From c9c73c2a8e85b3af809914a9df1d9830cbbcb6e3 Mon Sep 17 00:00:00 2001 From: dlangbe Date: Thu, 22 Oct 2026 14:57:35 -0400 Subject: [PATCH 1/2] fix(rocke/attention): bound gfx942 3D KV cache by its K/V loader The gfx942 tiled 3D kernel has no i64 path, and #12446 made it refuse any paged cache above 0x7FFF0000 bytes, the async buffer bound. Its wide and fp8 loaders are flat i32-byte-offset loads that read up to 2^31 bytes, so 2 GiB decode caches that ran correctly before are now refused. Decide the gfx942 3D i64 flag from the selected loader's limit, through one rule shared by the spec builder, the 3D cache key and the dispatcher's tuning specs. Emitted IR is unchanged. --- .../builders/common/attention_spec_builder.py | 3 +- .../library/dispatch/attention/common.py | 28 +++-- .../dispatch/attention/tuning_specs.py | 29 +++-- .../kernels/common/attention_unified.py | 44 +++++++- .../kernels/gfx942/attention_tiled_3d.py | 38 +++++-- .../dispatch/attention/test_tuning_space.py | 14 ++- .../library/tests/test_selector_parity.py | 105 +++++++++++++++++- 7 files changed, 235 insertions(+), 26 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py index c50ecbb76dee..ef9b99611332 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py +++ b/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. **( diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/common.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/common.py index c8251d77d527..da154bf483cb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/common.py +++ b/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), diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/tuning_specs.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/tuning_specs.py index 3ad88a35234c..2d734418223c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/tuning_specs.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/tuning_specs.py @@ -18,6 +18,7 @@ _num_segments, _tiled_2d_impl, _tiled_3d_impl, + gfx942_3d_kv_needs_i64_addr, kv_addr_width_for_builder, ) @@ -130,19 +131,23 @@ def _mask_fields(problem: UnifiedAttentionProblem, arch: str) -> dict: return fields -def _semantic_fields( - problem: UnifiedAttentionProblem, builder_kind: str = "tiled" -) -> dict: - kv_storage_dtype = "fp8e4m3" if problem.use_fp8 else None +def _kv_block_stride(problem: UnifiedAttentionProblem) -> int: elem_bytes = 1 if problem.use_fp8 else 2 - block_stride = ( + return ( int(problem.block_size) * int(problem.num_kv_heads) * int(problem.head_size) * elem_bytes ) + + +def _semantic_fields( + problem: UnifiedAttentionProblem, builder_kind: str = "tiled" +) -> dict: + kv_storage_dtype = "fp8e4m3" if problem.use_fp8 else None + elem_bytes = 1 if problem.use_fp8 else 2 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_stride(problem), elem_bytes ) return { "head_size": int(problem.head_size), @@ -312,8 +317,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_stride(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, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index 351480e4d083..a3df58e9d6cb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/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 @@ -3985,7 +4027,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), diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_3d.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_3d.py index 0a63602220f4..11d737730a37 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_3d.py +++ b/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: diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_tuning_space.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_tuning_space.py index 0e020115ffe1..9e023fb80ed9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_tuning_space.py +++ b/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 diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_selector_parity.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_selector_parity.py index f051c227c741..60ac85e15393 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_selector_parity.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_selector_parity.py @@ -24,6 +24,7 @@ from __future__ import annotations +import os import unittest from contextlib import contextmanager from dataclasses import asdict, fields, replace @@ -33,13 +34,16 @@ import builders.common.attention_spec_builder as _asb import kernels.common.attention_unified as _au from dispatch.attention import AttentionRequest, attention_candidates -from dispatch.attention.common import _tuning_problem +from dispatch.attention.common import AttentionTuningSpec, _tuning_problem from dispatch.attention.tuning_specs import ( ExplicitAttention2DConfig, + ExplicitAttention3DConfig, _semantic_fields, make_explicit_attention_2d_spec, + make_explicit_attention_3d_specs, ) from kernels import UnifiedAttentionProblem +from kernels.gfx942 import attention_tiled_3d as _gfx942_3d from kernels.gfx942.attention_tiled_2d import ( UnifiedAttention2DTiledSpec as Gfx942TiledSpec, build_gfx942_4warp_gqa, @@ -492,6 +496,105 @@ def test_spec_meta_and_guard_agree_at_every_count(self): _au._check_kv_addr_width(p, _kv_cache(n, p), limit) +# d128 bs16 4-kv-head bf16 blocks are 16 KiB, so 131072 blocks are exactly 2^31 +# bytes: 64 sequences of 32768 tokens. +_AT_2GIB = 131072 + + +def _gfx942_decode_prob(num_kv_blocks: int, **kw) -> UnifiedAttentionProblem: + fields = dict( + total_q=64, + num_seqs=64, + num_query_heads=64, + num_kv_heads=4, + head_size=128, + block_size=16, + max_seqlen_q=1, + max_seqlen_k=32768, + dtype="bf16", + num_kv_blocks=num_kv_blocks, + ) + fields.update(kw) + return UnifiedAttentionProblem(**fields) + + +class TestGfx942Tiled3DKvAddr(unittest.TestCase): + """The gfx942 tiled 3D kernel has no i64 path; its K/V loader sets the cache + it reads. The flat loaders read to 2^31 bytes, the async copies to the + 0x7FFF0000 buffer range.""" + + def setUp(self): + env = mock.patch.dict(os.environ) + env.start() + self.addCleanup(env.stop) + os.environ.pop("HIPDNN_GFX942_3D_WKV", None) + + def test_loader_limits(self): + limit = _gfx942_3d.i32_kv_cache_bytes + common = dict(tile_size=16, head_size=128) + self.assertEqual(limit(kv_fp8=False, use_wide_kv_load=True, **common), 2**31) + self.assertEqual(limit(kv_fp8=True, use_wide_kv_load=False, **common), 2**31) + self.assertEqual( + limit(kv_fp8=False, use_wide_kv_load=False, **common), 0x7FFF0000 + ) + # 8 x 32 halves do not split over 64 lanes x 8, so the async copies run. + self.assertEqual( + limit(kv_fp8=False, use_wide_kv_load=True, tile_size=8, head_size=32), + 0x7FFF0000, + ) + + def test_exactly_2gib_decode_builds_on_wide_loader(self): + for hq in (64, 32): + p = _gfx942_decode_prob(_AT_2GIB, num_query_heads=hq) + with self.subTest(num_query_heads=hq), _patch_arch("gfx942"): + spec = _asb._tiled_3d_spec_from_problem(p, "gfx942") + self.assertTrue(spec.use_wide_kv_load) + self.assertFalse(spec.use_i64_kv_addr) + + def test_past_2gib_wide_loader_refused(self): + p = _gfx942_decode_prob(_AT_2GIB + 1) + with _patch_arch("gfx942"), self.assertRaises(NotImplementedError): + _asb._tiled_3d_spec_from_problem(p, "gfx942") + + def test_exactly_2gib_async_loader_refused(self): + os.environ["HIPDNN_GFX942_3D_WKV"] = "0" + p = _gfx942_decode_prob(_AT_2GIB) + with _patch_arch("gfx942"), self.assertRaises(NotImplementedError): + _asb._tiled_3d_spec_from_problem(p, "gfx942") + + def test_3d_key_follows_loader_limit(self): + # No i64 binary exists, so caches the loader reads share one key, and a + # cache past it gets its own key and reaches the spec's refusal. + small = _gfx942_decode_prob(_AT_I32_LIMIT) + at_2gib = _gfx942_decode_prob(_AT_2GIB) + over = _gfx942_decode_prob(_AT_2GIB + 1) + with _patch_arch("gfx942"): + key = _au._tiled_3d_cache_key(small, "gfx942") + self.assertEqual(key, _au._tiled_3d_cache_key(at_2gib, "gfx942")) + self.assertNotEqual(key, _au._tiled_3d_cache_key(over, "gfx942")) + + def test_explicit_3d_spec_and_bind_follow_loader_limit(self): + p = _gfx942_decode_prob(_AT_2GIB) + config = ExplicitAttention3DConfig( + num_segments=0, knobs=(("use_wide_kv_load", True),) + ) + segment, reduce = make_explicit_attention_3d_specs(p, config, arch="gfx942") + self.assertFalse(segment.use_i64_kv_addr) + tuned = AttentionTuningSpec( + path="3d", + arch="gfx942", + builder_kind="tiled", + compile_backend="llvm", + candidate_name="test", + tuning_id="test", + kernel_spec=segment, + reduce_spec=reduce, + ) + self.assertFalse(tuned.with_num_kv_blocks(_AT_2GIB).kernel_spec.use_i64_kv_addr) + with self.assertRaises(NotImplementedError): + tuned.with_num_kv_blocks(_AT_2GIB + 1) + + @contextmanager def _clear_launch_memos(): """Run with empty launch-meta and per-key KV address limit memos.""" From 6120fae1db06dab54dadef5ec643cb814c82ef71 Mon Sep 17 00:00:00 2001 From: dlangbe Date: Fri, 23 Oct 2026 15:34:33 -0400 Subject: [PATCH 2/2] fix(rocke): check the real paged cache on 3D attention launches The 3D segment spec picks its KV addressing from problem.num_kv_blocks, which a caller can leave stale. Now that the gfx942 wide and fp8 loaders read caches up to 2^31 bytes, a stale count on a larger cache would read a wrapped address instead of zeros. _run_3d_tiled_on_stream now checks the real K cache against the segment kernel's limit before the eager and graph launches, as the 2D path already does. Strided K/V and gfx1250 are not checked. Also from review: use _kv_block_bytes in the explicit 3D tuning spec, and drop the environment override from the async-loader test (head size 256 selects the async copies). Co-Authored-By: Claude --- .../dispatch/attention/tuning_specs.py | 15 +-- .../kernels/common/attention_unified.py | 47 +++++++-- .../library/tests/test_selector_parity.py | 97 +++++++++++++++---- 3 files changed, 122 insertions(+), 37 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/tuning_specs.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/tuning_specs.py index 2d734418223c..4be5480c6f37 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/tuning_specs.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/tuning_specs.py @@ -14,6 +14,7 @@ from kernels.common.attention_unified import ( UnifiedAttentionProblem, + _kv_block_bytes, _kv_cache_policy_3d, _num_segments, _tiled_2d_impl, @@ -131,23 +132,13 @@ def _mask_fields(problem: UnifiedAttentionProblem, arch: str) -> dict: return fields -def _kv_block_stride(problem: UnifiedAttentionProblem) -> int: - elem_bytes = 1 if problem.use_fp8 else 2 - return ( - int(problem.block_size) - * int(problem.num_kv_heads) - * int(problem.head_size) - * elem_bytes - ) - - def _semantic_fields( problem: UnifiedAttentionProblem, builder_kind: str = "tiled" ) -> dict: kv_storage_dtype = "fp8e4m3" if problem.use_fp8 else None elem_bytes = 1 if problem.use_fp8 else 2 use_i64 = kv_addr_width_for_builder( - builder_kind, int(problem.num_kv_blocks), _kv_block_stride(problem), elem_bytes + builder_kind, int(problem.num_kv_blocks), _kv_block_bytes(problem), elem_bytes ) return { "head_size": int(problem.head_size), @@ -321,7 +312,7 @@ def make_explicit_attention_3d_specs( if arch == "gfx942": semantic["use_i64_kv_addr"] = gfx942_3d_kv_needs_i64_addr( int(problem.num_kv_blocks), - _kv_block_stride(problem), + _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), diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index a3df58e9d6cb..e871c0c04136 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py @@ -3046,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 ( @@ -4187,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, @@ -4412,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() @@ -4793,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] = {} diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_selector_parity.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_selector_parity.py index 60ac85e15393..7254f0ea0ce6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_selector_parity.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_selector_parity.py @@ -24,7 +24,6 @@ from __future__ import annotations -import os import unittest from contextlib import contextmanager from dataclasses import asdict, fields, replace @@ -501,7 +500,7 @@ def test_spec_meta_and_guard_agree_at_every_count(self): _AT_2GIB = 131072 -def _gfx942_decode_prob(num_kv_blocks: int, **kw) -> UnifiedAttentionProblem: +def _decode_prob(num_kv_blocks: int, **kw) -> UnifiedAttentionProblem: fields = dict( total_q=64, num_seqs=64, @@ -518,17 +517,40 @@ def _gfx942_decode_prob(num_kv_blocks: int, **kw) -> UnifiedAttentionProblem: return UnifiedAttentionProblem(**fields) +class _ReachedPipeline(Exception): + pass + + +def _launch_3d(problem: UnifiedAttentionProblem, num_blocks: int, arch: str) -> None: + """Run the 3D launch on a shape-only cache of ``num_blocks`` up to the + pipeline build, which raises ``_ReachedPipeline``.""" + with _patch_arch(arch), mock.patch.dict( + _au._3D_KV_ADDR_LIMITS, clear=True + ), mock.patch.object(_au, "_get_3d_pipeline", side_effect=_ReachedPipeline): + _au._run_3d_tiled_on_stream( + problem=problem, + q=None, + k=_kv_cache(num_blocks, problem), + v=None, + out=None, + cu_seqlens_q=None, + seqused_k=None, + softmax_scale=1.0, + block_table=None, + softcap=0.0, + sinks=None, + bt_stride=0, + warmup=0, + attempts=1, + use_graph=False, + ) + + class TestGfx942Tiled3DKvAddr(unittest.TestCase): """The gfx942 tiled 3D kernel has no i64 path; its K/V loader sets the cache it reads. The flat loaders read to 2^31 bytes, the async copies to the 0x7FFF0000 buffer range.""" - def setUp(self): - env = mock.patch.dict(os.environ) - env.start() - self.addCleanup(env.stop) - os.environ.pop("HIPDNN_GFX942_3D_WKV", None) - def test_loader_limits(self): limit = _gfx942_3d.i32_kv_cache_bytes common = dict(tile_size=16, head_size=128) @@ -545,36 +567,73 @@ def test_loader_limits(self): def test_exactly_2gib_decode_builds_on_wide_loader(self): for hq in (64, 32): - p = _gfx942_decode_prob(_AT_2GIB, num_query_heads=hq) + p = _decode_prob(_AT_2GIB, num_query_heads=hq) with self.subTest(num_query_heads=hq), _patch_arch("gfx942"): spec = _asb._tiled_3d_spec_from_problem(p, "gfx942") self.assertTrue(spec.use_wide_kv_load) self.assertFalse(spec.use_i64_kv_addr) def test_past_2gib_wide_loader_refused(self): - p = _gfx942_decode_prob(_AT_2GIB + 1) + p = _decode_prob(_AT_2GIB + 1) with _patch_arch("gfx942"), self.assertRaises(NotImplementedError): _asb._tiled_3d_spec_from_problem(p, "gfx942") - def test_exactly_2gib_async_loader_refused(self): - os.environ["HIPDNN_GFX942_3D_WKV"] = "0" - p = _gfx942_decode_prob(_AT_2GIB) - with _patch_arch("gfx942"), self.assertRaises(NotImplementedError): - _asb._tiled_3d_spec_from_problem(p, "gfx942") + def test_async_loader_bounded_by_buffer_range(self): + # Head size 256 takes the async copies. Its 32 KiB blocks fill the + # buffer range at 65534 and reach 2^31 bytes at 65536. + with _patch_arch("gfx942"): + spec = _asb._tiled_3d_spec_from_problem( + _decode_prob(65534, head_size=256), "gfx942" + ) + self.assertFalse(spec.use_wide_kv_load) + self.assertFalse(spec.use_i64_kv_addr) + for n in (65535, 65536): + with self.subTest(num_kv_blocks=n), self.assertRaises( + NotImplementedError + ): + _asb._tiled_3d_spec_from_problem( + _decode_prob(n, head_size=256), "gfx942" + ) def test_3d_key_follows_loader_limit(self): # No i64 binary exists, so caches the loader reads share one key, and a # cache past it gets its own key and reaches the spec's refusal. - small = _gfx942_decode_prob(_AT_I32_LIMIT) - at_2gib = _gfx942_decode_prob(_AT_2GIB) - over = _gfx942_decode_prob(_AT_2GIB + 1) + small = _decode_prob(_AT_I32_LIMIT) + at_2gib = _decode_prob(_AT_2GIB) + over = _decode_prob(_AT_2GIB + 1) with _patch_arch("gfx942"): key = _au._tiled_3d_cache_key(small, "gfx942") self.assertEqual(key, _au._tiled_3d_cache_key(at_2gib, "gfx942")) self.assertNotEqual(key, _au._tiled_3d_cache_key(over, "gfx942")) + def test_3d_launch_checks_the_real_cache(self): + # A stale num_kv_blocks builds the spec for a small cache, so the launch + # checks the real cache against what the K/V loader reads. + wide = _decode_prob(1024) + with self.assertRaises(_ReachedPipeline): + _launch_3d(wide, _AT_2GIB, "gfx942") + with self.assertRaisesRegex(ValueError, "fixing num_kv_blocks will not help"): + _launch_3d(wide, _AT_2GIB + 1, "gfx942") + not_wide = _decode_prob(1024, head_size=256) + with self.assertRaises(_ReachedPipeline): + _launch_3d(not_wide, 65534, "gfx942") + with self.assertRaisesRegex(ValueError, "fixing num_kv_blocks will not help"): + _launch_3d(not_wide, 65535, "gfx942") + + def test_gfx950_3d_launch_checks_the_real_cache(self): + # 131068 of these 16 KiB blocks fill the i32 buffer range. A stale count + # builds an i32 spec, which the launch refuses past that; the real count + # builds an i64 spec, which reads the whole cache. + stale = _decode_prob(1024) + with self.assertRaises(_ReachedPipeline): + _launch_3d(stale, 131068, "gfx950") + with self.assertRaisesRegex(ValueError, "below the cache's block count"): + _launch_3d(stale, 131069, "gfx950") + with self.assertRaises(_ReachedPipeline): + _launch_3d(_decode_prob(131069), 131069, "gfx950") + def test_explicit_3d_spec_and_bind_follow_loader_limit(self): - p = _gfx942_decode_prob(_AT_2GIB) + p = _decode_prob(_AT_2GIB) config = ExplicitAttention3DConfig( num_segments=0, knobs=(("use_wide_kv_load", True),) )