perf(rocke): runtime heads, window and batch for gfx942/gfx950 tiled attention

dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml

@@ -74,13 +74,15 @@ test_categories:
       - "rocke_mx_scale_decode"
       - "rocke_buffer_rsrc_hip"
       - "rocke_tiled_attention_2d_reentrancy"
+      - "rocke_tiled_attention_entry_guards"
       - "rocke_unified_attention_scalar_lower_to_llvm"
       - "rocke_python_import_smoke"
       - "rocke_golden_static"
       - "rocke_differential_golden_ll"
       - "rocke_differential_golden_canonical"
       - "rocke_library_conv_pytest"
       - "rocke_library_attention_cpu_pytest"
+      - "rocke_library_attention_cache_key_pytest"
       - "rocke_library_attention_pytest"
       # Not registered on Windows (see rocke/platform/CMakeLists.txt); the
       # parser ignores a listed name that has no registered test.
@@ -132,13 +134,15 @@ test_categories:
       - "rocke_mx_scale_decode"
       - "rocke_buffer_rsrc_hip"
       - "rocke_tiled_attention_2d_reentrancy"
+      - "rocke_tiled_attention_entry_guards"
       - "rocke_unified_attention_scalar_lower_to_llvm"
       - "rocke_python_import_smoke"
       - "rocke_golden_static"
       - "rocke_differential_golden_ll"
       - "rocke_differential_golden_canonical"
       - "rocke_library_conv_pytest"
       - "rocke_library_attention_cpu_pytest"
+      - "rocke_library_attention_cache_key_pytest"
       - "rocke_library_attention_pytest"
       # Not registered on Windows (see rocke/platform/CMakeLists.txt); the
       # parser ignores a listed name that has no registered test.
@@ -159,6 +163,7 @@ test_categories:
     description: "rocKE generated-artifact reproduction checks"
     test_patterns:
       - "rocke_arch_domain_reproducibility"
+      - "rocke_library_attention_cache_key_sweep_pytest"
     labels:
       - "host"
       - "toolchain"

dnn-providers/hip-kernel-provider/descriptor-packaging/python/hkp_pack/rocke_compile.py

@@ -323,15 +323,15 @@ def _module_from_source(source):
 
 
 # Builders whose validation predicate is not derivable from the builder name.
-# The tiled ``supports_tiled_2d``/``3d`` are keyword-only, so these builders use the
-# spec-first ``supports_tiled_*_spec`` wrappers in the gfx942/gfx950 modules. The
-# fastkv-register-P tiled builders, the gfx942 reduce builder and all gfx1250 tiled
-# builders have no spec-first predicate, so ``_check_support_predicate`` refuses to pack
-# them.
+# The tiled ``supports_tiled_2d``/``3d`` and ``supports_fastkv_register_p_2d`` are
+# keyword-only, so these builders use the spec-first ``supports_*_spec`` wrappers in
+# the gfx942/gfx950 modules. The gfx942 reduce builder and all gfx1250 tiled builders
+# have no spec-first predicate, so ``_check_support_predicate`` refuses to pack them.
 _PREDICATE_ALIASES = {
     "build_unified_attention_2d_tiled": "supports_tiled_2d_spec",
     "build_gfx942_4warp_gqa": "supports_gfx942_4warp_gqa_spec",
     "build_unified_attention_3d_tiled": "supports_tiled_3d_spec",
+    "build_unified_attention_2d_fastkv_register_p": "supports_fastkv_register_p_2d_spec",
     "build_unified_attention_reduce_tiled": "supports_tiled_reduce_spec",
 }
 

dnn-providers/hip-kernel-provider/descriptor-packaging/tests/rocke/test_hkp_pack_producer_guards.py

@@ -23,6 +23,7 @@
     _module_from_source,
     _require_spec_arch_signature,
     _resolve_spec_class,
+    _resolve_support_predicate,
     build_spec,
     rocke_compile_arch,
 )
@@ -386,6 +387,36 @@ def test_every_shipped_rocke_builder_passes_the_pack_time_predicate_check():
             pytest.fail(f"{source}:{builder} @ {arch} ({target}): {exc}")
 
 
+@pytest.mark.skipif(not _HAS_ROCKE_LIBRARY, reason="rocke/library not in this tree")
+def test_fastkv_register_p_builder_packs_through_its_spec_predicate():
+    """The gfx950 fastKV register-P builder resolves its spec-first predicate by
+    alias; the predicate accepts the experiment's spec and refuses one the
+    builder would raise on."""
+    module = import_module("kernels.gfx950.attention_tiled_2d_fastkv_regp")
+    builder = "build_unified_attention_2d_fastkv_register_p"
+    name, _ = _resolve_support_predicate(module, builder)
+    assert name == "supports_fastkv_register_p_2d_spec"
+    spec = module.make_fastkv_register_p_spec(
+        module.UnifiedAttention2DTiledSpec(
+            head_size=64,
+            block_size=32,
+            num_query_heads=64,
+            num_kv_heads=8,
+            dtype="bf16",
+            use_sinks=False,
+            sliding_window=0,
+            has_softcap=False,
+        )
+    )
+    _check_support_predicate(module, builder, spec, "gfx950")
+    no_fast_desc = dataclasses.replace(spec, use_fast_paged_kv_desc=False)
+    with pytest.raises(HkpPackError, match="use_fast_paged_kv_desc"):
+        _check_support_predicate(module, builder, no_fast_desc, "gfx950")
+    # The shared tiled gate admits gfx942 for its narrow builder; this one must not.
+    with pytest.raises(HkpPackError, match="gfx950 kernel"):
+        _check_support_predicate(module, builder, spec, "gfx942")
+
+
 # --- C. Literal validation --------------------------------------------------
 @dataclasses.dataclass
 class _LiteralSpec:

dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/prefill/benchmark_prefill2d_live.py

@@ -464,9 +464,9 @@ def build(
             supports_tiled_2d,
         )
         from kernels.common.attention_unified import (
-            _attn_signature,
             _enable_i64_kv_addr,
             _tiled_2d_impl,
+            _tiled_2d_signature,
         )
         from rocke.runtime import KernelLauncher
 
@@ -525,12 +525,12 @@ def build(
             use_i64_kv_addr=flags["use_i64_kv_addr"] or _enable_i64_kv_addr(problem),
             use_register_pv=flags["use_register_pv"],
         )
-        # kernel_name() does not encode the KV addressing width, so key on it too.
+        # The name carries i64 addressing; the signature keys the per-shape
+        # launch values the cached spec feeds.
         key = (
             shape.signature,
             variant,
             spec.kernel_name(),
-            spec.use_i64_kv_addr,
             self.compile_backend,
         )
         if key not in self._launchers:
@@ -540,9 +540,7 @@ def build(
                 KernelLauncher(
                     hsaco=artifact.hsaco,
                     kernel_name=artifact.kernel_name,
-                    signature=_attn_signature(
-                        dtype, include_bt_stride=True, include_qq_bias_stride=True
-                    ),
+                    signature=_tiled_2d_signature(ARCH, dtype),
                     cache_key=("prefill2d_live_gfx942", key),
                 ),
                 spec,

dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_fastkv_regp.py

@@ -8,11 +8,10 @@
 ``attention_tiled_2d_fastkv_regp.py`` and optionally joins against a Triton CSV
 by ``shape_signature``.
 
-The emitted 2D kernel name omits specialization constants such as ``num_seqs``
-and the derived binary-search trip count.  This harness appends a short shape
-hash to each built ``KernelDef.name`` before compiling, so the ROCm module/HSACO
-cache never sees two different shape-specialized code objects with the same
-kernel symbol.
+Head counts, window and batch are kernargs, so shapes that share a kernel name
+share a code object. This harness still appends a short shape hash to each built
+``KernelDef.name`` before compiling, so every trace record loads its own module
+and the timings stay independent of the order records run in.
 
 Example:
 
@@ -308,7 +307,7 @@ def _variant_spec_and_builder(
     def _launcher(self, shape, problem, variant: str, sliding_window: int):
         from rocke import compile_kernel
         from kernels.common.attention_unified import (
-            _attn_signature,
+            _tiled_2d_signature,
             _enable_i64_kv_addr,
         )
         from rocke.runtime import KernelLauncher
@@ -317,28 +316,25 @@ def _launcher(self, shape, problem, variant: str, sliding_window: int):
             shape, problem, variant, sliding_window
         )
         spec = replace(spec, use_i64_kv_addr=_enable_i64_kv_addr(problem))
-        # kernel_name() does not encode the KV addressing width, so key on it too.
+        # The name carries i64 addressing; the signature keeps one module per
+        # trace record (see the module docstring).
         key = (
             variant,
             shape.signature,
             spec.kernel_name(),
-            spec.use_i64_kv_addr,
             self.compile_backend,
         )
         if key not in self._launchers:
             kernel = builder(spec)
-            # The display name omits compile-time shape constants.  Keep the
-            # human-readable spec name in CSV, but compile a shape-unique symbol
-            # to avoid ROCm HSACO/module-cache aliasing across trace records.
+            # Keep the human-readable spec name in CSV, but compile a
+            # shape-unique symbol so each trace record loads its own module.
             kernel.name = _shape_kernel_name(kernel.name, variant, shape.signature)
             artifact = compile_kernel(kernel, capture_ir_text=False)
             self._launchers[key] = (
                 KernelLauncher(
                     hsaco=artifact.hsaco,
                     kernel_name=artifact.kernel_name,
-                    signature=_attn_signature(
-                        spec.dtype, include_bt_stride=True, include_qq_bias_stride=True
-                    ),
+                    signature=_tiled_2d_signature("gfx950", spec.dtype),
                     cache_key=("prefill2d_fastkv_regp", key),
                 ),
                 spec,

dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_live.py

@@ -436,7 +436,7 @@ def build(
             supports_tiled_2d,
         )
         from kernels.common.attention_unified import (
-            _attn_signature,
+            _tiled_2d_signature,
             _enable_i64_kv_addr,
             _select_2d_compile_backend,
         )
@@ -498,12 +498,12 @@ def build(
             # A cache past the i32 buffer range needs i64 whatever the variant asks.
             use_i64_kv_addr=flags["use_i64_kv_addr"] or _enable_i64_kv_addr(problem),
         )
-        # kernel_name() does not encode the KV addressing width, so key on it too.
+        # The name carries i64 addressing; the signature keys the per-shape
+        # launch values the cached spec feeds.
         key = (
             shape.signature,
             variant,
             spec.kernel_name(),
-            spec.use_i64_kv_addr,
             self.compile_backend,
         )
         if key not in self._launchers:
@@ -521,9 +521,7 @@ def build(
                 KernelLauncher(
                     hsaco=artifact.hsaco,
                     kernel_name=artifact.kernel_name,
-                    signature=_attn_signature(
-                        dtype, include_bt_stride=True, include_qq_bias_stride=True
-                    ),
+                    signature=_tiled_2d_signature("gfx950", dtype),
                     cache_key=("prefill2d_live", key),
                 ),
                 spec,

dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_traces.py

@@ -128,7 +128,7 @@ def _launcher(self, shape, problem, sliding_window: int):
             supports_tiled_2d,
         )
         from kernels.common.attention_unified import (
-            _attn_signature,
+            _tiled_2d_signature,
             _enable_i64_kv_addr,
         )
         from rocke.runtime import KernelLauncher
@@ -181,15 +181,13 @@ def _launcher(self, shape, problem, sliding_window: int):
             use_fast_paged_kv_desc=use_combo,
             use_i64_kv_addr=_enable_i64_kv_addr(problem),
         )
-        # The display name intentionally omits several compile-time constants
-        # (notably num_seqs/binary_search_iters). Cache on the trace signature
-        # as well so shapes with the same kernel_name never reuse an HSACO that
-        # was specialized for a different batch geometry.
-        # kernel_name() does not encode the KV addressing width either.
+        # The kernel name identifies the binary (head counts, window and batch
+        # are kernargs; i64 addressing is in the name). Key on the trace
+        # signature as well because the cached spec also feeds per-shape launch
+        # values such as the K/V address limit.
         key = (
             shape.signature,
             spec.kernel_name(),
-            spec.use_i64_kv_addr,
             self.compile_backend,
         )
         if key not in self._launchers:
@@ -199,9 +197,7 @@ def _launcher(self, shape, problem, sliding_window: int):
                 KernelLauncher(
                     hsaco=artifact.hsaco,
                     kernel_name=artifact.kernel_name,
-                    signature=_attn_signature(
-                        dtype, include_bt_stride=True, include_qq_bias_stride=True
-                    ),
+                    signature=_tiled_2d_signature("gfx950", dtype),
                     cache_key=("prefill2d_trace_combo", key),
                 ),
                 spec,

dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py

@@ -21,6 +21,7 @@
     _has_score_bias,
     _enable_combo_2d,
     _enable_early_v_schedule,
+    _enable_fast_paged_kv_desc,
     _enable_fp8_mfma_qk,
     _enable_gfx942_3d_wide_kv_load,
     _enable_gfx942_bf16_flash,
@@ -240,15 +241,7 @@ def _base_2d_generic_fields(problem: UnifiedAttentionProblem, arch: str) -> dict
         use_transposed_mask_limit=mask_opts,
         use_mfma32_skip_legacy_qreg=skip_legacy_qreg,
         use_early_v_schedule=_enable_early_v_schedule(problem, arch),
-        use_fast_paged_kv_desc=(
-            combo_no_sw
-            and problem.dtype
-            == "bf16"  # fast descriptor is bf16-only (spec validator enforces)
-            and not problem.use_fp8
-            and problem.num_query_heads == 64
-            and problem.num_kv_heads == 8
-            and _select_2d_tile_size(problem, arch) == 64
-        ),
+        use_fast_paged_kv_desc=_enable_fast_paged_kv_desc(problem, arch),
         use_register_pv=_enable_register_pv(problem, arch),
         use_fp8_mfma_qk=_enable_fp8_mfma_qk(problem, arch),
         use_i64_kv_addr=_enable_i64_kv_addr(problem),

dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/bf16_wide_check.py

@@ -39,7 +39,7 @@
 from kernels.gfx942.attention_tiled_2d import (  # noqa: E402
     UnifiedAttention2DTiledSpec,
 )
-from kernels.common.attention_unified import _attn_signature  # noqa: E402
+from kernels.common.attention_unified import _tiled_2d_signature  # noqa: E402
 from rocke.runtime import KernelLauncher  # noqa: E402
 
 BS = H.BLOCK_SIZE
@@ -136,9 +136,7 @@ def build_launcher(s, spec, kind):
     launcher = KernelLauncher(
         hsaco=artifact.hsaco,
         kernel_name=artifact.kernel_name,
-        signature=_attn_signature(
-            s.dtype, include_bt_stride=True, include_qq_bias_stride=True
-        ),
+        signature=_tiled_2d_signature("gfx942", s.dtype),
         cache_key=("bf16wide", spec.kernel_name(), kind),
     )
     return launcher

dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/parity_unified_attention.py

@@ -774,7 +774,7 @@ def _build_kernel(s: Shape):
     """
     from rocke import compile_kernel
     from kernels import build_unified_attention_2d_tiled, supports_tiled_2d
-    from kernels.common.attention_unified import _attn_signature
+    from kernels.common.attention_unified import _tiled_2d_signature
     from rocke.runtime import KernelLauncher
 
     spec, config = _build_spec(s)
@@ -810,9 +810,7 @@ def _build_kernel(s: Shape):
     launcher = KernelLauncher(
         hsaco=artifact.hsaco,
         kernel_name=artifact.kernel_name,
-        signature=_attn_signature(
-            s.dtype, include_bt_stride=True, include_qq_bias_stride=True
-        ),
+        signature=_tiled_2d_signature("gfx942", s.dtype),
         cache_key=("gfx942_attn_example", spec.kernel_name(), config),
     )
     return launcher, spec, config

dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py

@@ -1735,7 +1735,7 @@ def _run_rocke(
             supports_tiled_2d,
         )
         from kernels.common.attention_unified import (
-            _attn_signature,
+            _tiled_2d_signature,
             _attn_values,
             _spec_kv_addr_limit,
             _select_2d_compile_backend,
@@ -1775,9 +1775,7 @@ def _run_rocke(
             launcher = KernelLauncher(
                 hsaco=artifact.hsaco,
                 kernel_name=artifact.kernel_name,
-                signature=_attn_signature(
-                    dtype_str, include_bt_stride=True, include_qq_bias_stride=True
-                ),
+                signature=_tiled_2d_signature(_ARCH, dtype_str),
                 cache_key=("d256_fast_direct", spec.kernel_name()),
             )
             if probe_occupancy:
@@ -1848,9 +1846,7 @@ def call_once():
         launcher = KernelLauncher(
             hsaco=artifact.hsaco,
             kernel_name=artifact.kernel_name,
-            signature=_attn_signature(
-                dtype_str, include_bt_stride=True, include_qq_bias_stride=True
-            ),
+            signature=_tiled_2d_signature(_ARCH, dtype_str),
             cache_key=("prod_2d_direct", spec.kernel_name()),
         )
         vals = _attn_values(