fix(rocke): bound gfx942 3D KV cache by its K/V loader

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

@@ -32,6 +32,7 @@
     _enable_gfx942_flash_q_direct,
     _enable_gfx942_fp16_flash,
     _enable_i64_kv_addr,
+    _enable_i64_kv_addr_3d,
     _enable_k_single_buffer,
     _enable_mfma_32x32,
     _enable_register_pv,
@@ -396,7 +397,7 @@ def _spec_generic_3d(problem: UnifiedAttentionProblem, arch: str):
         kv_storage_dtype=_kv_storage_dtype(problem),
         tile_size_override=tile_size_override,
         use_wide_kv_load=_enable_gfx942_3d_wide_kv_load(problem, arch),
-        use_i64_kv_addr=_enable_i64_kv_addr(problem),
+        use_i64_kv_addr=_enable_i64_kv_addr_3d(problem, arch),
         use_decode_grid=_use_decode_grid(problem, arch),
         # Additive bias is gfx950-only for now; gfx942's spec has no such field.
         **(

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

@@ -56,6 +56,7 @@
 
 from kernels.common.attention_unified import (
     UnifiedAttentionProblem,
+    gfx942_3d_kv_needs_i64_addr,
     kv_addr_width_for_builder,
     # Re-exported, never redeclared. The kernel owns what it covers; dispatch's
     # job is to state that coverage as a Capability so it can be filtered and
@@ -1078,13 +1079,26 @@ def with_num_kv_blocks(self, num_kv_blocks: int) -> "AttentionTuningSpec":
             * int(kernel_spec.head_size)
             * elem_bytes
         )
-        # A builder with no i64 path (gfx942 4-warp GQA) shares its spec class
-        # with the default builder, so no __post_init__ can refuse the flag the
-        # way the gfx942 3D spec does. It stays i32 up to its element limit and
-        # is refused at bind time past that, not at compile.
-        use_i64 = kv_addr_width_for_builder(
-            self.builder_kind, count, block_stride, elem_bytes
-        )
+        if self.path == "3d" and self.arch == "gfx942":
+            # No i64 path: the flag marks a cache past what the spec's K/V
+            # loader reads, and the spec's __post_init__ refuses it.
+            use_i64 = gfx942_3d_kv_needs_i64_addr(
+                count,
+                block_stride,
+                kv_fp8=kernel_spec.kv_storage_dtype == "fp8e4m3",
+                use_wide_kv_load=kernel_spec.use_wide_kv_load,
+                tile_size=kernel_spec.tile_size,
+                head_size=kernel_spec.head_size,
+            )
+        else:
+            # A builder with no i64 path (gfx942 4-warp GQA) shares its spec
+            # class with the default builder, so no __post_init__ can refuse the
+            # flag the way the gfx942 3D spec does. It stays i32 up to its
+            # element limit and is refused at bind time past that, not at
+            # compile.
+            use_i64 = kv_addr_width_for_builder(
+                self.builder_kind, count, block_stride, elem_bytes
+            )
         return replace(
             self,
             kernel_spec=replace(kernel_spec, use_i64_kv_addr=use_i64),

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

@@ -14,10 +14,12 @@
 
 from kernels.common.attention_unified import (
     UnifiedAttentionProblem,
+    _kv_block_bytes,
     _kv_cache_policy_3d,
     _num_segments,
     _tiled_2d_impl,
     _tiled_3d_impl,
+    gfx942_3d_kv_needs_i64_addr,
     kv_addr_width_for_builder,
 )
 
@@ -135,14 +137,8 @@ def _semantic_fields(
 ) -> dict:
     kv_storage_dtype = "fp8e4m3" if problem.use_fp8 else None
     elem_bytes = 1 if problem.use_fp8 else 2
-    block_stride = (
-        int(problem.block_size)
-        * int(problem.num_kv_heads)
-        * int(problem.head_size)
-        * elem_bytes
-    )
     use_i64 = kv_addr_width_for_builder(
-        builder_kind, int(problem.num_kv_blocks), block_stride, elem_bytes
+        builder_kind, int(problem.num_kv_blocks), _kv_block_bytes(problem), elem_bytes
     )
     return {
         "head_size": int(problem.head_size),
@@ -312,8 +308,18 @@ def make_explicit_attention_3d_specs(
         raise ValueError("gfx950 3D does not implement hoist/wide-KV tuning knobs")
     # Follows the problem unless the config sets it (``_kv_cache_policy_3d``).
     knobs.setdefault("kv_cache_policy", _kv_cache_policy_3d(problem, arch=arch))
+    semantic = _semantic_fields(problem)
+    if arch == "gfx942":
+        semantic["use_i64_kv_addr"] = gfx942_3d_kv_needs_i64_addr(
+            int(problem.num_kv_blocks),
+            _kv_block_bytes(problem),
+            kv_fp8=semantic["kv_storage_dtype"] is not None,
+            use_wide_kv_load=bool(knobs.get("use_wide_kv_load", False)),
+            tile_size=int(tile),
+            head_size=int(problem.head_size),
+        )
     segment = spec_type(
-        **_semantic_fields(problem),
+        **semantic,
         **_mask_fields(problem, arch),
         num_segments=num_segments,
         waves_per_eu=config.waves_per_eu,

dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py

@@ -2906,6 +2906,48 @@ def kv_addr_width_for_builder(
     return False
 
 
+def gfx942_3d_kv_needs_i64_addr(
+    num_kv_blocks: int,
+    block_bytes: int,
+    *,
+    kv_fp8: bool,
+    use_wide_kv_load: bool,
+    tile_size: int,
+    head_size: int,
+) -> bool:
+    """``use_i64_kv_addr`` for the gfx942 tiled 3D kernel on this paged cache.
+
+    The kernel has no i64 path, so its spec refuses the flag: True means the
+    cache is past what the selected K/V loader reads. The single rule for the
+    spec builder, the cache key and the dispatcher's tuning specs.
+    """
+    from ..gfx942.attention_tiled_3d import i32_kv_cache_bytes
+
+    limit = i32_kv_cache_bytes(
+        kv_fp8=kv_fp8,
+        use_wide_kv_load=use_wide_kv_load,
+        tile_size=tile_size,
+        head_size=head_size,
+    )
+    return num_kv_blocks > 0 and num_kv_blocks * block_bytes > limit
+
+
+def _enable_i64_kv_addr_3d(problem: UnifiedAttentionProblem, arch: str) -> bool:
+    """3D ``use_i64_kv_addr``: ``_enable_i64_kv_addr``, except on gfx942, whose
+    limit depends on the K/V loader (``gfx942_3d_kv_needs_i64_addr``)."""
+    # Below the buffer range every loader is i32, so skip the selector walk.
+    if arch != "gfx942" or not _enable_i64_kv_addr(problem):
+        return _enable_i64_kv_addr(problem)
+    return gfx942_3d_kv_needs_i64_addr(
+        problem.num_kv_blocks,
+        _kv_block_bytes(problem),
+        kv_fp8=_kv_storage_dtype(problem) is not None,
+        use_wide_kv_load=_enable_gfx942_3d_wide_kv_load(problem, arch),
+        tile_size=_gfx942_3d_tile_size_override(problem, arch) or problem.block_size,
+        head_size=problem.head_size,
+    )
+
+
 def _kv_block_bytes(problem: UnifiedAttentionProblem) -> int:
     """Byte size of one paged K (or V) block for ``problem``."""
     elem_bytes = 1 if problem.use_fp8 else 2
@@ -3004,15 +3046,35 @@ def _spec_kv_addr_limit(
     return min(limits, key=lambda limit: limit.nbytes, default=None)
 
 
+def _3d_spec_kv_addr_limit(
+    problem: UnifiedAttentionProblem, spec, arch: str
+) -> _KvAddrLimit | None:
+    """Largest paged K cache, in bytes, the 3D segment kernel built from ``spec``
+    reads. gfx942 has no i64 path, so its K/V loader sets the limit."""
+    if arch == "gfx942":
+        from ..gfx942.attention_tiled_3d import i32_kv_cache_bytes
+
+        return _KvAddrLimit(
+            i32_kv_cache_bytes(
+                kv_fp8=spec.kv_storage_dtype is not None,
+                use_wide_kv_load=spec.use_wide_kv_load,
+                tile_size=spec.tile_size,
+                head_size=spec.head_size,
+            ),
+            _KV_LIMIT_NO_I64_PATH,
+        )
+    return _spec_kv_addr_limit(problem, spec)
+
+
 def _kv_addr_fix(
     problem: UnifiedAttentionProblem, num_kv_blocks: int, kv_addr_limit: _KvAddrLimit
 ) -> str:
     """What a caller can do about a cache past ``kv_addr_limit``."""
     if kv_addr_limit.cause == _KV_LIMIT_NO_I64_PATH:
         return (
-            "This kernel has no i64 KV path (the scalar, gfx942 4-warp GQA and "
-            "gfx1250 kernels index at most 2^31 elements), so fixing "
-            "num_kv_blocks will not help. Use a smaller cache."
+            "This kernel has no i64 KV path (the scalar, gfx942 4-warp GQA, "
+            "gfx942 tiled 3D and gfx1250 kernels), so fixing num_kv_blocks "
+            "will not help. Use a smaller cache."
         )
     if kv_addr_limit.cause != _KV_LIMIT_I32_BUFFER:
         return (
@@ -3985,7 +4047,7 @@ def _tiled_3d_cache_key(problem: UnifiedAttentionProblem, arch: str) -> Tuple:
         _select_3d_waves_per_eu(problem, arch),
         _enable_gfx942_3d_wide_kv_load(problem, arch),
         _kv_storage_dtype(problem),
-        _enable_i64_kv_addr(problem),
+        _enable_i64_kv_addr_3d(problem, arch),
         decode_grid,
         _kv_cache_policy_3d(problem, arch),
         _gfx950_3d_tile_size_override(problem, arch),
@@ -4145,9 +4207,9 @@ def _attn_values(
     _require_additive_bias(problem, additive_bias)
     # Production 2D and scalar launches and the torch direct-launch harnesses
     # pack kernargs here, so check the real K cache against what the compiled
-    # kernel addresses (``kv_addr_limit``, from ``_spec_kv_addr_limit``). Not
-    # covered: the 3D path and the gfx1250 tiled_2d_verify harness, which pack
-    # their own.
+    # kernel addresses (``kv_addr_limit``, from ``_spec_kv_addr_limit``). The 3D
+    # path checks in ``_run_3d_tiled_on_stream``; the gfx1250 tiled_2d_verify
+    # harness packs its own and is not covered.
     _check_kv_addr_width(problem, k, kv_addr_limit)
     vals = {
         "output_ptr": out,
@@ -4370,6 +4432,19 @@ def _run_3d_tiled_on_stream(
     else:
         num_segments = _num_segments(problem, _launch_arch)
         cache_key = _tiled_3d_cache_key(problem, _launch_arch)
+    # The spec saw problem.num_kv_blocks, which a caller can leave stale, so
+    # check the real paged cache. gfx1250's 3D spec makes no i64 decision.
+    if kv_cache_layout is None and _launch_arch in ("gfx942", "gfx950"):
+        if cache_key not in _3D_KV_ADDR_LIMITS:
+            segment_spec = (
+                tuning_spec.kernel_spec
+                if tuning_spec is not None
+                else _tiled_3d_spec_from_problem(problem, _launch_arch)
+            )
+            _3D_KV_ADDR_LIMITS[cache_key] = _3d_spec_kv_addr_limit(
+                problem, segment_spec, _launch_arch
+            )
+        _check_kv_addr_width(problem, k, _3D_KV_ADDR_LIMITS[cache_key])
     if kv_cache_layout is not None:
         cache_key = cache_key + ("strided_kv",)
     capturing = _torch_stream_capturing()
@@ -4751,6 +4826,8 @@ class _Attention2DLaunchMeta:
 # Per-cache-key ``_spec_kv_addr_limit`` of the default-route spec. The spec is a
 # function of the cache key, so it is built once per kernel, not per total_q.
 _2D_KV_ADDR_LIMITS: dict[tuple, _KvAddrLimit | None] = {}
+# The same for the 3D segment kernel, from ``_3d_spec_kv_addr_limit``.
+_3D_KV_ADDR_LIMITS: dict[tuple, _KvAddrLimit | None] = {}
 _2D_GRAPHS: Dict[Tuple, Any] = {}
 _2D_GRAPH_REFS: Dict[Tuple, Tuple[Any, ...]] = {}
 _SCALAR_LAUNCHERS: Dict[Tuple, KernelLauncher] = {}

dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_3d.py

@@ -126,12 +126,10 @@ class UnifiedAttention3DTiledSpec:
     # instructions -> less VMEM issue pressure on the memory-bound decode segment.
     # fp16/bf16 only (the fp8 path already loads wide). Signature-tracked.
     use_wide_kv_load: bool = False
-    # 64-bit paged-KV addressing for caches > 2 GiB. Accepted for signature
-    # parity with the shared dispatch spec builder and the gfx950 spec; the
-    # gfx942 narrow segment loaders do not key on it yet (the gfx942 >2 GiB
-    # decode path is a separate follow-on), so a True value here would be a
-    # no-op on gfx942 -- guard it so the dispatcher cannot silently build an
-    # uncorrected gfx942 kernel for an oversized cache.
+    # 64-bit paged-KV addressing. Accepted for signature parity with the gfx950
+    # spec, but this kernel has no i64 path: the builders set it only for a
+    # cache past what the K/V loader reads (``i32_kv_cache_bytes``), and it is
+    # refused below so an oversized cache cannot build a kernel that misreads it.
     use_i64_kv_addr: bool = False
     # Direct non-paged [B,H,S,D] addressing; D remains contiguous.
     kv_layout: str = "paged"
@@ -319,6 +317,30 @@ def supports_tiled_3d(
     return True, "supported"
 
 
+# The wide K/V loader moves 8 halves per lane over the 64-lane wave; a tile it
+# cannot split evenly falls back to the async copies.
+_WIDE_KV_ELEMS = 8
+
+
+def wide_kv_load_fits(tile_size: int, head_size: int) -> bool:
+    """True when ``use_wide_kv_load`` takes effect for this tile."""
+    return (tile_size * head_size) % (64 * _WIDE_KV_ELEMS) == 0
+
+
+def i32_kv_cache_bytes(
+    *, kv_fp8: bool, use_wide_kv_load: bool, tile_size: int, head_size: int
+) -> int:
+    """Largest paged K/V cache, in bytes, the selected K/V loader reads.
+
+    This kernel has no i64 KV path. The fp8 dequant and wide loaders are flat
+    global loads indexed by an i32 byte offset, correct below 2^31 bytes. The
+    async copies read through a buffer resource bounded at 0x7FFF0000 bytes.
+    """
+    if kv_fp8 or (use_wide_kv_load and wide_kv_load_fits(tile_size, head_size)):
+        return 0x8000_0000
+    return 0x7FFF_0000
+
+
 def build_unified_attention_3d_tiled(
     spec: UnifiedAttention3DTiledSpec, *, arch: str = "gfx942"
 ) -> KernelDef:
@@ -800,8 +822,8 @@ def _issue_v_load(
     # 1 DWORD to 4 DWORDs, cutting load-instruction count 4x. The dim axis is
     # contiguous in the paged cache (stride = KV_BYTES) so 8 consecutive halves
     # are a legal coalesced b128 load.
-    WIDE_ELEMS = 8
-    WIDE_OK = (T * HD) % (THREADS * WIDE_ELEMS) == 0
+    WIDE_ELEMS = _WIDE_KV_ELEMS
+    WIDE_OK = wide_kv_load_fits(T, HD)
     wide_chunks_per_thread = (T * HD) // (THREADS * WIDE_ELEMS) if WIDE_OK else 0
 
     def _kv_col(row: Value, col: Value) -> Value:

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

@@ -314,8 +314,20 @@ def test_runtime_cache_size_specializes_i64_but_keeps_the_id(self):
             seqlen_q=1,
             seqlen_k=4096,
         )
+        # The gfx942 3D kernel has no i64 path; its K/V loader sets the cache it
+        # reads. The wide loader reads 2^31 bytes (65536 of these 32 KiB
+        # blocks), the async copies stop at the 0x7FFF0000 buffer range.
+        wide = next(s for s in gfx942_split_specs if s.kernel_spec.use_wide_kv_load)
+        not_wide = next(
+            s for s in gfx942_split_specs if not s.kernel_spec.use_wide_kv_load
+        )
+        for n in (65535, 65536):
+            self.assertFalse(wide.with_num_kv_blocks(n).kernel_spec.use_i64_kv_addr)
+        with self.assertRaisesRegex(NotImplementedError, "does not support"):
+            wide.with_num_kv_blocks(65537)
+        self.assertFalse(not_wide.with_num_kv_blocks(65534).kernel_spec.use_i64_kv_addr)
         with self.assertRaisesRegex(NotImplementedError, "does not support"):
-            gfx942_split_specs[0].with_num_kv_blocks(65535)
+            not_wide.with_num_kv_blocks(65535)
 
         # The 4-warp GQA builder has no i64 path either, but shares its spec
         # class with the default builder. Its i32 element offsets read 2^31