perf(rocke): build one K/V resource per page on the gfx942 2D i64 path

dnn-providers/hip-kernel-provider/rocke/docs/gpu-attention-test-coverage.md

@@ -49,6 +49,7 @@ listed separately in the inventory, not added to the workload rows.
 | gfx942 forced exp2 A/B | 2 | No A/B assertion | Skipped |
 | gfx942 paged D128 ring regression | 12 | No: different path | Skipped |
 | gfx942 paged GQA head-fold | 4 | No: different path | Skipped |
+| gfx942 paged decode with i64 K/V addressing | 4 | No: different path | Skipped |
 | gfx950 dense attention | Historical inventory: 51 shapes in 50 configurations | Base causal configurations on gfx950 | Device-gated on gfx942; see gfx950 validation below |
 | gfx1151 standalone WMMA attention | Pinned cohort | Yes: FP16 D64/D128 × MHA/GQA × full/causal; source and installed validation passed | Device-gated on gfx942; local artifact routing checked |
 | gfx1200 standalone WMMA attention compiled for gfx12-generic | Adapter cohort | Pending GPU qualification; host selection, ABI, and contract tests only | No installed pinned CTest |
@@ -129,6 +130,7 @@ The historical source inventory below is unchanged by this new cohort.
 | Forced exp2 variants | BF16 D128 Hq16/Hkv4 B1/S512, default/persistent | The reference does not force fast and plain paths and compare them bitwise |
 | Paged D128 ring regression | FP16/BF16 × Hq/Hkv=(32,8)/(16,16) × S=512/1024/4096, B1 | Unified/paged builder and cache layout differ from dense SDPA |
 | Paged head-fold | BF16 D128, Hq32/Hkv8, `(S,W)=(2048,512),(8192,4096)` × page size16/32, B1 | No head-fold kernel, paging, long sequence, or clipped window in reference cohort |
+| Paged i64 K/V addressing | BF16 D128 decode, page size 16: B8 Sk1024 Hq/Hkv=(32,8),(64,4) and B8 ragged Sk≤1000 with a block-table tail (`num_kv_blocks` overstated to force i64); B4 Sk1024 Hq/Hkv=(32,8) on a cache just over 2 GiB with pages straddling 2^31 bytes | No reference case pages the cache or uses i64 K/V addressing |
 | Seeds/distributions | One fixed normally distributed recipe per configuration in required suite | Multiple runs test repeatability, not additional input coverage |
 | Additional dense combinations | No BF16 full-mask case; no D64 full-mask case; not every dtype × grid × head-ratio combination | The selected rows are not a Cartesian-product sweep |
 | Unequal query/KV lengths, bottom-right masks | Existing detailed numeric cohort is gfx950: see 11 rows in shape table | All required Sq=Sk=512; no bottom-right or ragged reference |
@@ -192,6 +194,7 @@ See [the SDPA guide](../library/tests/sdpa_reference/docs/sdpa-test-reference.md
 | [test_gdn_decode_gfx950_numeric.py](../library/tests/test_gdn_decode_gfx950_numeric.py) | 13 | Selected by library CTest file list |
 | [test_gfx1250_attention.py](../library/tests/test_gfx1250_attention.py) | 2 | Selected by library CTest file list |
 | [test_gfx942_gqa_head_fold_numeric.py](../library/tests/test_gfx942_gqa_head_fold_numeric.py) | 1 | Selected by library CTest file list |
+| [test_attention_gfx942_kv_pages.py](../library/tests/test_attention_gfx942_kv_pages.py) | 3 | Selected by library CTest file list |
 | [test_kda_chunkwise_gfx942_numeric.py](../library/tests/test_kda_chunkwise_gfx942_numeric.py) | 6 | Selected by library CTest file list |
 | [test_kda_chunkwise_gfx950_numeric.py](../library/tests/test_kda_chunkwise_gfx950_numeric.py) | 14 | Selected by library CTest file list |
 | [test_kda_decode_gfx950_numeric.py](../library/tests/test_kda_decode_gfx950_numeric.py) | 8 | Selected by library CTest file list |
@@ -293,6 +296,10 @@ not executed GPU-kernel counts. All filenames below are relative to
 | [test_gfx942_gqa_head_fold_numeric.py](../library/tests/test_gfx942_gqa_head_fold_numeric.py) :: test_fold_numeric_vs_fp32_windowed_oracle | gfx942 | bf16 | 1 | 8192 | 8192 | 32 | 8 | 128 | paged head-fold; window=4096; page size=16 | Missing |
 | [test_gfx942_gqa_head_fold_numeric.py](../library/tests/test_gfx942_gqa_head_fold_numeric.py) :: test_fold_numeric_vs_fp32_windowed_oracle | gfx942 | bf16 | 1 | 2048 | 2048 | 32 | 8 | 128 | paged head-fold; window=512; page size=32 | Missing |
 | [test_gfx942_gqa_head_fold_numeric.py](../library/tests/test_gfx942_gqa_head_fold_numeric.py) :: test_fold_numeric_vs_fp32_windowed_oracle | gfx942 | bf16 | 1 | 8192 | 8192 | 32 | 8 | 128 | paged head-fold; window=4096; page size=32 | Missing |
+| [test_attention_gfx942_kv_pages.py](../library/tests/test_attention_gfx942_kv_pages.py) :: test_i64_decode_matches_reference | gfx942 | bf16 | 8 | 1 | 1024 | 32 | 8 | 128 | paged tiled 2D; i64 K/V addressing (num_kv_blocks overstated); page size=16 | Missing |
+| [test_attention_gfx942_kv_pages.py](../library/tests/test_attention_gfx942_kv_pages.py) :: test_i64_decode_matches_reference | gfx942 | bf16 | 8 | 1 | 1024 | 64 | 4 | 128 | paged tiled 2D; i64 K/V addressing (num_kv_blocks overstated); page size=16 | Missing |
+| [test_attention_gfx942_kv_pages.py](../library/tests/test_attention_gfx942_kv_pages.py) :: test_i64_ragged_lengths_and_block_table_tail | gfx942 | bf16 | 8 | 1 | 17-1000 (ragged) | 32 | 8 | 128 | paged tiled 2D; i64 K/V addressing (num_kv_blocks overstated); last tile past the block table; page size=16 | Missing |
+| [test_attention_gfx942_kv_pages.py](../library/tests/test_attention_gfx942_kv_pages.py) :: test_i64_decode_reads_past_2gib | gfx942 | bf16 | 4 | 1 | 1024 | 32 | 8 | 128 | paged tiled 2D; i64 K/V addressing; cache just over 2 GiB, pages straddle 2^31 bytes; skips under ~5 GiB free; page size=16 | Missing |
 | [test_sdpa_pinned_reference.py](../library/tests/test_sdpa_pinned_reference.py) :: test_sdpa_correctness_against_qualified_rocke | gfx942 | fp16 | 1 | 512 | 512 | 16 | 4 | 128 | persistent=False; causal=True | Required GPU |
 | [test_sdpa_pinned_reference.py](../library/tests/test_sdpa_pinned_reference.py) :: test_sdpa_correctness_against_qualified_rocke | gfx942 | fp16 | 1 | 512 | 512 | 16 | 4 | 128 | persistent=True; causal=True | Required GPU |
 | [test_sdpa_pinned_reference.py](../library/tests/test_sdpa_pinned_reference.py) :: test_sdpa_correctness_against_qualified_rocke | gfx942 | bf16 | 1 | 512 | 512 | 16 | 4 | 128 | persistent=True; causal=True | Required GPU |

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

@@ -3000,11 +3000,11 @@ def _enable_i64_kv_addr(problem: UnifiedAttentionProblem) -> bool:
     The default load path puts the full byte offset (incl.
     ``physical_block * block_stride``) in a 32-bit buffer voffset, bounded at
     ``_I32_KV_BUFFER_BYTES``; past that, loads silently return zeros. When the
-    cache is that large we fold the per-block offset into a
-    64-bit buffer base (a tiny, wave-uniform ``make_buffer_rsrc`` per
-    block; measured within ~1-2% of the i32 path). Below the cap we keep
-    the fast i32 path. ``num_kv_blocks == 0`` means the cache size is
-    unknown (caller did not supply it) -> assume small / fast path.
+    cache is that large we fold the per-block offset into a 64-bit buffer
+    base (a wave-uniform ``make_buffer_rsrc`` per block; the gfx942 tiled 2D
+    K/V copies build one per page of a tile, not one per copy). Below the cap
+    we keep the i32 path. ``num_kv_blocks == 0`` means the cache size is
+    unknown (caller did not supply it) -> assume small / i32 path.
     """
     return kv_cache_needs_i64_addr(problem.num_kv_blocks, _kv_block_bytes(problem))
 

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

@@ -334,9 +334,10 @@ class UnifiedAttention2DTiledSpec:
     # voffset, which overflows for paged caches > 2 GiB (~65 K bf16 / ~131 K
     # fp8 blocks) and silently corrupts loads. When True, the per-block
     # ``physical_block`` offset is folded into a 64-bit buffer *base* (small
-    # within-block voffset) so any cache size is addressable. It costs a
-    # per-call ``make_buffer_rsrc`` so the dispatcher only enables it when the
-    # cache actually exceeds the 2 GiB cap (see _enable_i64_kv_addr).
+    # within-block voffset) so any cache size is addressable. The K/V loaders
+    # build one ``make_buffer_rsrc`` per page of a tile; the dispatcher only
+    # enables it when the cache exceeds the i32 buffer range (see
+    # _enable_i64_kv_addr).
     use_i64_kv_addr: bool = False
     # Experimental schedule: issue the current V async copy immediately after
     # the iter-start K drain/barrier, before QK. This gives V the whole QK plus
@@ -2675,6 +2676,61 @@ def _fast_paged_kv_voff(call: int, block0: Value, block1: Value):
                 ),
             )
 
+    # i64 K/V addressing: load the tile's page ids once and build one buffer
+    # resource per page. A resource per copy would reload its page id and wait
+    # on it, and since vmcnt drains in order, those waits also drain the copies
+    # issued before it. The non-fast K/V loaders below are i32-only otherwise.
+    PAGE_RSRC_I64 = I64_KV_ADDR and not FAST_PAGED_KV_DESC
+    if PAGE_RSRC_I64:
+        # The spec and supports_tiled_2d guarantee both: tile_size is a multiple
+        # of block_size, and a copy (at most 8 warps x 64 lanes x 2 halves) never
+        # exceeds a page (at least 16 x 64 halves), all powers of two. Raised,
+        # not asserted, so python -O cannot emit copies that straddle pages.
+        if T % BS != 0 or (BS * HD) % KV_HALVES_PER_CALL != 0:
+            raise ValueError(
+                f"i64 K/V copies need whole pages per tile and per copy "
+                f"(T={T}, BS={BS}, HD={HD}, halves per copy={KV_HALVES_PER_CALL})"
+            )
+        CALLS_PER_PAGE = (BS * HD) // KV_HALVES_PER_CALL
+
+        def _tile_pages(
+            src: Value, kv_tile_idx: Value
+        ) -> tuple[list[Value], list[Value], Value]:
+            """Base pointers and buffer resources of the tile's pages, and the
+            kv head's byte offset within a page."""
+            first = b.add(seq_base, b.mul(kv_tile_idx, b.const_i32(N_BLOCKS_PER_TILE)))
+            ptrs = []
+            for page in range(N_BLOCKS_PER_TILE):
+                idx = b.add(first, b.const_i32(page))
+                physical = b.masked_global_load(
+                    block_tables,
+                    idx,
+                    b.cmp_lt(idx, block_table_max_idx),
+                    b.const_i32(0),
+                    dtype=I32,
+                    align=4,
+                )
+                base = b.mul(
+                    b.zext(b.to_sgpr_u32(physical), I64), b.const_i64(kv_stride_blk_b)
+                )
+                ptrs.append(b.global_ptr_add(src, base))
+            rsrcs = [b.buffer_rsrc(ptr, kv_block_bytes_c) for ptr in ptrs]
+            kv_head_b = b.mul(kv_head_idx, b.const_i32(kv_stride_h_b))
+            return ptrs, rsrcs, kv_head_b
+
+        def _page_voff(call: int, kv_head_b: Value) -> Value:
+            # Byte offset of this lane's halves within the call's page.
+            half = b.add(
+                b.const_i32((call % CALLS_PER_PAGE) * KV_HALVES_PER_CALL),
+                lane_half_base,
+            )
+            token = b.div(half, b.const_i32(HD))
+            dim = b.mod(half, b.const_i32(HD))
+            return b.add(
+                b.add(b.mul(token, b.const_i32(kv_stride_tok_b)), kv_head_b),
+                b.mul(dim, b.const_i32(KV_BYTES)),
+            )
+
     def _issue_k_load_runtime(kv_tile_idx: Value, buf_idx: Value) -> None:
         """Issue async K loads for one tile into K_lds[buf_idx].
 
@@ -2696,6 +2752,8 @@ def _issue_k_load_runtime(kv_tile_idx: Value, buf_idx: Value) -> None:
         K_wave_base = b.smem_ptr_add(K_buf_base, wave_lds_offset_i64)
         if FAST_PAGED_KV_DESC:
             fast_block0, fast_block1 = _fast_paged_kv_blocks(kv_tile_idx)
+        if PAGE_RSRC_I64:
+            page_ptrs, page_rsrcs, kv_head_b = _tile_pages(key, kv_tile_idx)
         for call in range(kv_calls_per_tile):
             k_rsrc = key_rsrc
             k_ptr = key
@@ -2704,27 +2762,20 @@ def _issue_k_load_runtime(kv_tile_idx: Value, buf_idx: Value) -> None:
                 if base_i64 is not None:  # I64_KV_ADDR: per-block 64-bit base
                     k_ptr = b.global_ptr_add(key, base_i64)
                     k_rsrc = b.buffer_rsrc(k_ptr, kv_block_bytes_c)
+            elif PAGE_RSRC_I64:
+                k_ptr = page_ptrs[call // CALLS_PER_PAGE]
+                k_rsrc = page_rsrcs[call // CALLS_PER_PAGE]
+                voff = _page_voff(call, kv_head_b)
             else:
                 linear_half = b.add(
                     b.const_i32(call * KV_HALVES_PER_CALL), lane_half_base
                 )
-                if I64_KV_ADDR:
-                    base_i64, voff, _ = paged_kv_desc.offset_i64_split(
-                        b,
-                        "physical_block",
-                        tile_idx=kv_tile_idx,
-                        linear_half=linear_half,
-                        kv_head=kv_head_idx,
-                    )
-                    k_ptr = b.global_ptr_add(key, base_i64)
-                    k_rsrc = b.buffer_rsrc(k_ptr, kv_block_bytes_c)
-                else:
-                    voff, _ = paged_kv_desc.offset(
-                        b,
-                        tile_idx=kv_tile_idx,
-                        linear_half=linear_half,
-                        kv_head=kv_head_idx,
-                    )
+                voff, _ = paged_kv_desc.offset(
+                    b,
+                    tile_idx=kv_tile_idx,
+                    linear_half=linear_half,
+                    kv_head=kv_head_idx,
+                )
             k_dst = b.smem_ptr_add(K_wave_base, b.const_i64(call * bytes_per_call))
             if USE_GLOBAL_LOAD_LDS_K:
                 b.global_load_lds(
@@ -2836,6 +2887,8 @@ def _issue_v_load_runtime(kv_tile_idx: Value, buf_idx: Value) -> None:
         V_wave_base = b.smem_ptr_add(V_lds_addr, v_wave_lds_offset_i64)
         if FAST_PAGED_KV_DESC:
             fast_block0, fast_block1 = _fast_paged_kv_blocks(kv_tile_idx)
+        if PAGE_RSRC_I64:
+            _, page_rsrcs, kv_head_b = _tile_pages(value, kv_tile_idx)
         for call in range(kv_calls_per_tile):
             v_rsrc = value_rsrc
             if FAST_PAGED_KV_DESC:
@@ -2844,28 +2897,19 @@ def _issue_v_load_runtime(kv_tile_idx: Value, buf_idx: Value) -> None:
                     v_rsrc = b.buffer_rsrc(
                         b.global_ptr_add(value, base_i64), kv_block_bytes_c
                     )
+            elif PAGE_RSRC_I64:
+                v_rsrc = page_rsrcs[call // CALLS_PER_PAGE]
+                voff = _page_voff(call, kv_head_b)
             else:
                 linear_half = b.add(
                     b.const_i32(call * KV_HALVES_PER_CALL), lane_half_base
                 )
-                if I64_KV_ADDR:
-                    base_i64, voff, _ = paged_kv_desc.offset_i64_split(
-                        b,
-                        "physical_block",
-                        tile_idx=kv_tile_idx,
-                        linear_half=linear_half,
-                        kv_head=kv_head_idx,
-                    )
-                    v_rsrc = b.buffer_rsrc(
-                        b.global_ptr_add(value, base_i64), kv_block_bytes_c
-                    )
-                else:
-                    voff, _ = paged_kv_desc.offset(
-                        b,
-                        tile_idx=kv_tile_idx,
-                        linear_half=linear_half,
-                        kv_head=kv_head_idx,
-                    )
+                voff, _ = paged_kv_desc.offset(
+                    b,
+                    tile_idx=kv_tile_idx,
+                    linear_half=linear_half,
+                    kv_head=kv_head_idx,
+                )
             v_dst = b.smem_ptr_add(
                 V_wave_base, b.const_i64(call * V_BYTES_PER_CALL_SWZ)
             )

dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_2d_emit.c

@@ -112,6 +112,76 @@ static int make_spec(int idx, rocke_attention_tiled_2d_spec_t* s)
         s->use_k_single_buffer = false;
         s->use_agpr_alloc_zero = true;
         break;
+    case 7: /* Selector spec for bf16 D128 decode with i64 K/V addressing. */
+        s->head_size = 128;
+        s->block_size = 16;
+        s->num_query_heads = 32;
+        s->num_kv_heads = 8;
+        s->dtype = "bf16";
+        s->use_sinks = false;
+        s->sliding_window = 0;
+        s->has_softcap = false;
+        s->num_seqs = 256;
+        s->has_waves_per_eu = true;
+        s->waves_per_eu = 2;
+        s->use_register_pv = true;
+        s->has_tile_size = true;
+        s->tile_size = 64;
+        s->use_i64_kv_addr = true;
+        break;
+    case 8: /* i64 K/V, four warps, two pages per tile. */
+        s->head_size = 128;
+        s->block_size = 32;
+        s->num_query_heads = 32;
+        s->num_kv_heads = 8;
+        s->dtype = "bf16";
+        s->use_sinks = false;
+        s->sliding_window = 0;
+        s->has_softcap = false;
+        s->num_seqs = 256;
+        s->num_warps = 4;
+        s->has_waves_per_eu = true;
+        s->waves_per_eu = 2;
+        s->use_register_pv = true;
+        s->has_tile_size = true;
+        s->tile_size = 64;
+        s->use_i64_kv_addr = true;
+        break;
+    case 9: /* i64 K/V, one page per tile. */
+        s->head_size = 128;
+        s->block_size = 64;
+        s->num_query_heads = 32;
+        s->num_kv_heads = 8;
+        s->dtype = "bf16";
+        s->use_sinks = false;
+        s->sliding_window = 0;
+        s->has_softcap = false;
+        s->num_seqs = 256;
+        s->has_waves_per_eu = true;
+        s->waves_per_eu = 2;
+        s->use_register_pv = true;
+        s->has_tile_size = true;
+        s->tile_size = 64;
+        s->use_i64_kv_addr = true;
+        break;
+    case 10: /* cfg7 with K copied by global_load_lds. */
+        s->head_size = 128;
+        s->block_size = 16;
+        s->num_query_heads = 32;
+        s->num_kv_heads = 8;
+        s->dtype = "bf16";
+        s->use_sinks = false;
+        s->sliding_window = 0;
+        s->has_softcap = false;
+        s->num_seqs = 256;
+        s->has_waves_per_eu = true;
+        s->waves_per_eu = 2;
+        s->use_register_pv = true;
+        s->has_tile_size = true;
+        s->tile_size = 64;
+        s->use_i64_kv_addr = true;
+        s->use_global_load_lds_k = true;
+        break;
     default:
         return -1;
     }

dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_2d_emit.py

@@ -107,6 +107,74 @@
         use_k_single_buffer=False,
         use_agpr_alloc_zero=True,
     ),
+    # cfg7: the selector's spec for bf16 D128 decode (B=256, Skv=4096, GQA 32/8,
+    # block 16) on a paged cache past the i32 buffer range: i64 K/V addressing
+    # with four pages per tile, so the K/V copies take per-page resources.
+    7: dict(
+        head_size=128,
+        block_size=16,
+        num_query_heads=32,
+        num_kv_heads=8,
+        dtype="bf16",
+        use_sinks=False,
+        sliding_window=0,
+        has_softcap=False,
+        num_seqs=256,
+        waves_per_eu=2,
+        use_register_pv=True,
+        tile_size=64,
+        use_i64_kv_addr=True,
+    ),
+    # cfg8-cfg10: other i64 geometries of the per-page K/V loader. cfg8 has four
+    # warps and two 32-token pages per tile, cfg9 one 64-token page per tile,
+    # cfg10 is cfg7 with K copied by global_load_lds.
+    8: dict(
+        head_size=128,
+        block_size=32,
+        num_query_heads=32,
+        num_kv_heads=8,
+        dtype="bf16",
+        use_sinks=False,
+        sliding_window=0,
+        has_softcap=False,
+        num_seqs=256,
+        num_warps=4,
+        waves_per_eu=2,
+        use_register_pv=True,
+        tile_size=64,
+        use_i64_kv_addr=True,
+    ),
+    9: dict(
+        head_size=128,
+        block_size=64,
+        num_query_heads=32,
+        num_kv_heads=8,
+        dtype="bf16",
+        use_sinks=False,
+        sliding_window=0,
+        has_softcap=False,
+        num_seqs=256,
+        waves_per_eu=2,
+        use_register_pv=True,
+        tile_size=64,
+        use_i64_kv_addr=True,
+    ),
+    10: dict(
+        head_size=128,
+        block_size=16,
+        num_query_heads=32,
+        num_kv_heads=8,
+        dtype="bf16",
+        use_sinks=False,
+        sliding_window=0,
+        has_softcap=False,
+        num_seqs=256,
+        waves_per_eu=2,
+        use_register_pv=True,
+        tile_size=64,
+        use_i64_kv_addr=True,
+        use_global_load_lds_k=True,
+    ),
 }