perf(rocke): prepare saved paged decode bindings once

dnn-providers/hip-kernel-provider/rocke/docs/decode_host_overhead_case_study.md

@@ -0,0 +1,131 @@
+# Saved paged split-KV decode binding host overhead case study
+
+This case study concerns host submission from saved explicit gfx942/gfx950
+non-FP8, LSE-off paged one-token attention bindings. It describes preparation
+outside the timed loop and validation of equivalent segment/reducer kernels.
+It does not measure the native provider or change heuristic attention routing.
+
+## Caller and preparation
+
+The measured caller is an executable tuning candidate's saved Torch binding:
+`attention_execution_candidates` → `bind_tuning_attention_torch` →
+`TorchBinding.launch`. The first ordinary eager call goes through the existing
+public runner, validating the launch and initializing an existing
+`Attention3DExecution`. Binding creation neither launches nor drains.
+
+The binding then holds its compiled pipeline, launch geometry, stream wrapper,
+scalar defaults and three private scratch tensors. Compatible warm calls use
+the existing ordinary pipeline directly. They avoid allocation, repeated
+immutable admission/spec/cache work and redundant Torch context construction.
+Kernargs are still packed and retained on every invocation; live pointers and
+scalar overrides are read on every call.
+
+The new binding path uses `None` for Torch-current selection. It resolves that
+selection before passing an integer to the existing runner and pipeline.
+Legacy zero inputs and explicit integer handles preserve the current runtime
+API. An explicit stream different from the caller's current stream is entered
+using the cached stream object and restored afterwards.
+
+## Ownership and compatibility
+
+Scratch belongs to one saved binding. Separate bindings, including compatible
+maximum-KV bindings, share existing immutable compiled caches but own separate
+scratch. Each binding uses
+`4 * total_q * num_query_heads * num_segments * (head_size + 2)` scratch bytes.
+There is no new global cache, eviction policy, until-drain registry or cleanup
+API. Generic HIP runtime and launcher code are unchanged.
+
+LSE-on bindings use the general runner on every invocation, preserving live LSE
+buffer overrides and the runner's output validation. They keep the existing
+integer-stream interface and default; they do not prepare private eager scratch.
+
+Only the binding's creating host thread reuses its scratch under the existing
+execution owner's whole-pair guard. Different streams/devices/layouts, foreign
+host threads, reentrant calls, internal graphs and external capture use the
+general runner with its existing ownership. This path does not introduce a
+general concurrent runtime guarantee.
+
+Ordinary per-launch retention keeps submitted tensors and scratch alive after
+the binding is discarded. Fenced calls complete through the ordinary pipeline;
+unfenced calls require the caller's existing finite-batch completion boundary.
+Drains remain outside host enqueue timing and must happen before reading
+results. No context remains entered between binding invocations.
+
+## Validation method
+
+[Host regressions](../library/tests/test_attention_decode_binding.py) cover
+allocation/context reuse, independent bindings, executable-candidate reachability,
+live pointer repacking, scalars, stream selection/overrides, changed layouts,
+foreign/reentrant calls, capture/graph fallback, enqueue failure and ordinary
+retention after discarding a binding. The focused test is installed in the
+attention CPU lane. Tests for the withdrawn global cache and runtime changes
+are removed; pre-existing suites retain their baseline source.
+
+The LSE dispatch host regression separately checks repeated calls of one saved
+binding on gfx942/gfx950, including LSE-buffer and integer-stream overrides.
+It verifies that each call reaches the general runner without preparation.
+
+Fresh hardware validation compares the final binding implementation against the
+pinned base using identical executable specs, LLVM/HSACO hashes and geometry.
+The synthetic BF16 cases have B/KV = 1/2048 and 4/4096, Hq/Hkv = 32/4, D = 128
+and page size 16. Stable bindings and a rotating KV/KV+1/KV+2 maximum-length
+series are warmed separately; device lengths stay at KV. Internal and outer
+graphs are disabled for measurements.
+
+Binding, setup, compilation warmup and drains are excluded from
+`perf_counter_ns` host enqueue timing under `no_fence()`. HIP-event intervals
+are measured separately and include host pacing. Three paired trials alternate
+arm order on the same GPU and pinned CPU, using 20 warmups, 100 iterations per
+finite batch and seven samples; the first sample is excluded from medians.
+
+Every measured arm checks outputs against an independent NumPy FP64 reference
+over quantized BF16 inputs before and after timing, with `atol=0.01` and
+`rtol=0.02`. Post-checks poison output before a fresh launch. The validation
+harness and private raw measurements remain outside this PR. gfx1250 hardware
+and a migration of the generic stream API are outside scope.
+
+## Results
+
+The saved binding reduced host enqueue time at every measured configuration
+and in every paired trial. Each value below is
+`100 * (1 - median(final/base))` over the three paired trials. These are
+per-configuration results, not confidence intervals.
+
+| Target | B / KV | Binding mode | Host enqueue reduction | Host-paced HIP-event interval reduction |
+| --- | --- | --- | ---: | ---: |
+| gfx942 | 1 / 2048 | Stable | 57.1% | 56.8% |
+| gfx942 | 1 / 2048 | Maximum-length series | 57.5% | 57.1% |
+| gfx942 | 4 / 4096 | Stable | 57.0% | 56.1% |
+| gfx942 | 4 / 4096 | Maximum-length series | 57.6% | 57.0% |
+| gfx950 | 1 / 2048 | Stable | 53.7% | 54.0% |
+| gfx950 | 1 / 2048 | Maximum-length series | 54.1% | 53.3% |
+| gfx950 | 4 / 4096 | Stable | 54.2% | 53.9% |
+| gfx950 | 4 / 4096 | Maximum-length series | 54.3% | 54.2% |
+
+All compared segment/reducer LLVM and HSACO hashes, executable specs and
+grid/block geometry matched across arms and trials. Each target passed 48
+independent before/after checks in the measured runs, 64 existing runtime
+numerical cases and nine stream/capture/retention/completion checks. Ten
+existing LDS-admission numerical cases were explicitly deselected. No hardware
+skip or failed check is counted as a pass.
+
+The baseline was `41317132e3f7c081bb36fd228efb59145fe6368d`; the measured
+implementation was `b106aeda4da3265c35612b99c69f2a99a2c8c0de`. The explicit
+specs were `gfx942_u3d_splitkv_seg128_t1xb` and
+`gfx950_u3d_splitkv_seg128_t1xb`. The environment used Python 3.12, the Python
+engine, LLVM23, ROCm 7.14.0, Torch `2.13.0+rocm7.14.0`, HIP `7.14.60850` and
+NumPy 2.2.6. Compilation ran in Torch-free COMGR workers; launch and compilation
+used release-matched wheel HIP/COMGR libraries.
+
+The source-and-adapter archive SHA-256 was
+`df8c036a8b55370c134fe0ddeb9db3793ac81d82d8bcd6bf792e426b7c6deb31`;
+the separately hashed timing harness was
+`bdc02b9d13d80d3c370b15e0276bc331494ae6b801f92a91e03922590b3bfa3e`.
+The validation record retains per-file provenance, commands and paired trials.
+
+These results establish savings at the saved Python binding's host submission
+boundary for these configurations. HIP-event intervals include host pacing and
+do not establish a device-kernel speedup. The runs do not measure a new binding
+on every invocation, native provider performance, two-device concurrent
+submission or gfx1250 hardware. Earlier measurements of the withdrawn generic
+runtime/cache design are not used for these results.

dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/README.md

@@ -218,3 +218,5 @@ absolute timings depend on the GPU, ROCm/torch versions, clocks, and thermal
 state of the box, and are **not benchmarked into this folder**. Run
 `final_shapes_check.py` to produce the CK-vs-Torch eager/graph table for your own
 machine; the script prints, and gates on, the correctness result for every shape.
+
+See the [saved paged decode binding case study](../../../../docs/decode_host_overhead_case_study.md) for private scratch, prepared launch context and host validation.

dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/README.md

@@ -959,3 +959,5 @@ scenario (paths relative to `library/`):
 Measured conditions and numbers are intentionally omitted from this repository
 (per `platform/AGENTS.md` §Compliance); they are recorded only in the internal
 perf record.
+
+See the [saved paged decode binding case study](../../../../docs/decode_host_overhead_case_study.md) for private scratch, prepared launch context and host validation.

dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/bindings.py

@@ -362,6 +362,162 @@ def _reject_additive_bias_inputs(inputs: Mapping[str, Any], where: str) -> None:
         )
 
 
+def _bind_paged_decode_launch(problem, spec, run):
+    """Prepare one binding's eager pair after the public runner validates it.
+
+    Only the creating host thread reuses scratch. Different layouts, streams,
+    reentrant calls and captures keep the ordinary runner's private allocation.
+    Kernargs are packed and retained by the ordinary pipeline on every call.
+    """
+    from threading import current_thread
+
+    from kernels.common import attention_unified as au
+    from rocke.runtime.torch_interop import resolve_stream
+
+    thread = current_thread()
+    owner = None
+    prepared = None
+    stream_object = None
+    layouts = None
+    device = None
+    segment_defaults = None
+    reduce_defaults = None
+    names = (
+        "q",
+        "k",
+        "v",
+        "out",
+        "cu_seqlens_q",
+        "seqused_k",
+        "block_table",
+        "sinks",
+        "alibi_slopes",
+        "qq_bias",
+    )
+
+    def launch(call):
+        nonlocal owner, prepared, stream_object, layouts, device
+        nonlocal segment_defaults, reduce_defaults
+        requested_stream = call["stream"]
+        call["stream"] = resolve_stream(requested_stream, device=call["q"].device)
+        if current_thread() is not thread or (
+            owner is not None and owner._lock.locked()
+        ):
+            return run(**call)
+        import torch
+
+        if (
+            au._torch_stream_capturing()
+            or au._enable_3d_graph_replay(problem, spec.arch)
+            or torch.cuda.current_device() != call["q"].device.index
+        ):
+            return run(**call)
+        live_layouts = tuple(au._attention_tensor_layout(call[name]) for name in names)
+        if prepared is None:
+            # No extra warmup launch or synchronization at binding creation.
+            # A failed first submission stays on the public path for retry.
+            execution = au.Attention3DExecution()
+            owner = execution
+            try:
+                stream_object = torch.cuda.ExternalStream(
+                    call["stream"], device=call["q"].device
+                )
+                # Capture may belong to an explicit non-current stream.
+                # Check on that stream before giving the public runner an owner.
+                with torch.cuda.stream(stream_object):
+                    if au._torch_stream_capturing():
+                        owner = None
+                        return run(**call)
+                    result = run(**call, execution=execution)
+                prepared = au._get_3d_pipeline(
+                    problem,
+                    spec.cache_key(),
+                    int(spec.kernel_spec.num_segments),
+                    arch=spec.arch,
+                    tuning_spec=spec,
+                )
+                device = call["q"].device
+                layouts = live_layouts
+                output, maximum, expsum = execution._workspace
+                segment_defaults = {
+                    "segm_output_ptr": output,
+                    "segm_max_ptr": maximum,
+                    "segm_expsum_ptr": expsum,
+                    "num_seqs": int(problem.num_seqs),
+                    "block_table_stride": int(call["block_table"].stride(0)),
+                    "qq_bias_stride_0": 0,
+                    "additive_bias_ptr": 0,
+                    "additive_bias_batch_stride": 0,
+                    "additive_bias_head_stride": 0,
+                    "additive_bias_sq_stride": 0,
+                    "k_scale": 1.0,
+                    "v_scale": 1.0,
+                }
+                reduce_defaults = {
+                    "segm_output_ptr": output,
+                    "segm_max_ptr": maximum,
+                    "segm_expsum_ptr": expsum,
+                }
+            except BaseException:
+                owner = None
+                prepared = None
+                raise
+            return result
+        if (
+            call["q"].device != device
+            or call["stream"] != owner._stream
+            or live_layouts != layouts
+        ):
+            return run(**call)
+        # This is the existing execution owner's whole-pair guard, not a
+        # generic runtime submission lock. Reentrant work gets fresh scratch.
+        if not owner._lock.acquire(blocking=False):
+            return run(**call)
+        try:
+            segment = dict(
+                segment_defaults,
+                query_ptr=call["q"],
+                key_cache_ptr=call["k"],
+                value_cache_ptr=call["v"],
+                query_start_len_ptr=call["cu_seqlens_q"],
+                seq_lens_ptr=call["seqused_k"],
+                block_tables_ptr=call["block_table"],
+                sink_ptr=call["sinks"],
+                alibi_slopes_ptr=(
+                    call["alibi_slopes"] if call["alibi_slopes"] is not None else 0
+                ),
+                qq_bias_ptr=call["qq_bias"] if call["qq_bias"] is not None else 0,
+                scale=call["softmax_scale"],
+                softcap=call["softcap"],
+            )
+            reduce = dict(
+                reduce_defaults,
+                output_ptr=call["out"],
+                seq_lens_ptr=call["seqused_k"],
+            )
+            # None/legacy-zero selected the live Torch stream above. An
+            # explicit handle may require entering the cached stream object.
+            matching_stream = requested_stream is None or int(requested_stream) == 0
+            if not matching_stream:
+                matching_stream = (
+                    int(torch.cuda.current_stream(device).cuda_stream) == call["stream"]
+                )
+            if matching_stream:
+                return au._launch_3d_pipeline(
+                    prepared, segment, reduce, call["stream"], capturing=False
+                )
+            with torch.cuda.stream(stream_object):
+                if au._torch_stream_capturing():
+                    return run(**call)
+                return au._launch_3d_pipeline(
+                    prepared, segment, reduce, call["stream"], capturing=False
+                )
+        finally:
+            owner._lock.release()
+
+    return launch
+
+
 def bind_tuning_attention_torch(
     request, spec, tensors: Mapping[str, Any], **kwargs
 ) -> TorchBinding:
@@ -372,6 +528,12 @@ def bind_tuning_attention_torch(
     invariants externally may pass
     ``unsafe_skip_paged_value_validation=True`` to skip the synchronized value
     check; structural ABI validation is never skipped.
+
+    Saved gfx942/gfx950 non-FP8, LSE-off one-token 3D bindings prepare private
+    eager scratch and fixed launch context on their first ordinary launch.
+    LSE-on bindings, other execution contexts and graph calls use the general
+    runner. None selects Torch's current stream for prepared bindings; existing
+    integer stream behavior is preserved.
     """
     from kernels.common.attention_unified import (
         UnifiedAttentionProblem,
@@ -413,16 +575,29 @@ def bind_tuning_attention_torch(
     problem = replace(problem, num_kv_blocks=num_blocks)
     if hasattr(spec, "with_num_kv_blocks"):
         spec = spec.with_num_kv_blocks(num_blocks)
-    stream = kwargs.get("stream", 0)
     path = str(getattr(spec, "path", "2d"))
     backend = "tiled" if path == "2d" else path
+    decode = (
+        path == "3d"
+        and spec.arch in ("gfx942", "gfx950")
+        and problem.all_decode
+        and not problem.use_fp8
+        and not problem.emit_lse
+    )
+    stream = kwargs.get("stream", None if decode else 0)
     scale = kwargs.get("softmax_scale")
     if scale is None:
         scale = 1.0 / math.sqrt(int(problem.head_size))
 
+    prepared_launch = (
+        _bind_paged_decode_launch(problem, spec, run_unified_attention_torch)
+        if decode
+        else None
+    )
+
     def launch(**_kw):
         _reject_additive_bias_inputs(_kw, "launch kwargs")
-        return run_unified_attention_torch(
+        call = dict(
             problem=problem,
             q=tensors["q"],
             k=tensors["k"],
@@ -437,10 +612,14 @@ def launch(**_kw):
             alibi_slopes=_kw.get("alibi_slopes", tensors.get("alibi_slopes")),
             qq_bias=_kw.get("qq_bias", tensors.get("qq_bias")),
             backend=backend,
-            stream=int(_kw.get("stream", stream)),
+            stream=_kw.get("stream", stream),
             tuning_spec=spec,
             lse=_kw.get("lse", tensors.get("lse")),
         )
+        if prepared_launch is not None:
+            return prepared_launch(call)
+        call["stream"] = int(call["stream"])
+        return run_unified_attention_torch(**call)
 
     return TorchBinding(
         launch=launch,