feat(rocke): L2-locality-aware block orders for dense attention

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

@@ -14,8 +14,12 @@
 fixture engineered to violate it), so no check only ever sees valid data. The
 producer-agnostic halves of invariants 2 and 3 are also held on a real hip pack
 in ``tests/test_desk_check_invariants.py``.
+
+The real-bundle regressions at the end also hold the shipped dense attention specs to
+the C++ ingestor's restated launch geometry.
 """
 
+import dataclasses
 import json
 import shutil
 
@@ -27,6 +31,7 @@
     symbol_distinctness,
     toc_key_uniqueness,
 )
+from hkp_pack.rocke_compile import build_spec
 from pack_helpers import (
     _EXAMPLES,
     _ROOT_IDS,
@@ -323,3 +328,65 @@ def test_real_rocke_example_passes_out_of_the_box(rocke_root):
         assert proc.returncode == 0, str(kdp) + proc.stdout + proc.stderr
         assert "metadata/authored-spec drift: none" in proc.stdout, kdp
         assert "duplicate matcher tuples: none" in proc.stdout, kdp
+
+
+def _ingestor_geometry(spec):
+    """The C++ ingestor's launch geometry, restated: `gfx950AttentionDenseGeometry()`
+    (src/engines/kernel_ingestor_engine/packs/Gfx950AttentionDenseGeometry.hpp),
+    grid `(ceil(Sq / block_m), Hq, B)` and block `(block_m / 32 * 64, 1, 1)`."""
+    nqb = (spec.seqlen_q + spec.block_m - 1) // spec.block_m
+    grid = (nqb, spec.num_query_heads, spec.batch)
+    return grid, (spec.block_m // 32 * 64, 1, 1)
+
+
+@pytest.mark.quick
+@pytest.mark.parametrize("rocke_root", _ROCKE_EXAMPLE, ids=_ROOT_IDS)
+def test_real_attention_dense_bundles_launch_with_the_ingestor_geometry(
+    rocke_root, rocke_importable
+):
+    """The C++ ingestor relaunches each packed dense attention binary with its own
+    restated geometry, while the binary's block order is whatever the packed spec
+    resolves to in rocKE; an order with another grid runs out of bounds under that
+    launch, with no error. Nothing at pack time compares the two, and the C++
+    geometry tests pin literals, so every shipped `build_attention_dense` spec, built
+    the way the packer builds it, must give the restated grid and block: at its
+    canonical shape and at one with 4 query blocks and batch 3, so an order whose
+    grid coincides with the restated one at the canonical shape cannot pass."""
+    import importlib
+
+    from hkp_pack.rocke_compile import _module_from_source, _resolve_spec_class
+
+    kernels = [
+        (kdp, k)
+        for kdp in _require_bundles(rocke_root)
+        for k in _kernels(_read(kdp))
+        if k["kernel_source"].get("builder") == "build_attention_dense"
+    ]
+    if not kernels:
+        pytest.skip(f"{rocke_root.parent} ships no build_attention_dense kernel")
+    mismatches = set()
+    for kdp, kernel in kernels:
+        source = kernel["kernel_source"]
+        module = importlib.import_module(_module_from_source(source["source"]))
+        spec_cls = _resolve_spec_class(module, getattr(module, source["builder"]))
+        spec = build_spec(spec_cls, source["spec"])
+        seqlen = 4 * spec.block_m
+        for shape in ({}, {"batch": 3, "seqlen_q": seqlen, "seqlen_kv": seqlen}):
+            launched = dataclasses.replace(spec, **shape)
+            got = (
+                module.attention_dense_grid(launched),
+                module.attention_dense_block(launched),
+            )
+            if got != _ingestor_geometry(launched):
+                mismatches.add(
+                    f"{kernel['name']} ({launched.resolved_nonpersist_decode}): "
+                    f"B={launched.batch} Sq={launched.seqlen_q} gives {got}, the "
+                    f"ingestor launches {_ingestor_geometry(launched)}"
+                )
+    assert not mismatches, (
+        f"{len(mismatches)} shipped launch(es) disagree with the ingestor geometry; pin "
+        "the block order in the generator config "
+        "(projects/hipdnn/tools/IngestorGenerator/configs/<bundle>.yaml) and "
+        "regenerate, or teach the C++ geometry that order:\n  "
+        + "\n  ".join(sorted(mismatches)[:10])
+    )

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

@@ -55,13 +55,14 @@
 
 from kernels.gfx950.attention_dense import (  # noqa: E402
     AttentionDenseSpec,
+    _has_shape_params,
     attention_dense_block,
     attention_dense_grid,
+    attention_dense_signature,
     build_attention_dense,
     supports_attention_dense,
 )
 from rocke.helpers.compile import compile_kernel  # noqa: E402
-from rocke.helpers.spec import SignatureBuilder  # noqa: E402
 from rocke.runtime import (  # noqa: E402
     KernelLauncher,
     LaunchConfig,
@@ -126,24 +127,13 @@ def _dense_launcher(spec: AttentionDenseSpec) -> KernelLauncher:
         backend="python",
         capture_ir_text=False,
     )
-    sb = (
-        SignatureBuilder()
-        .ptr("q_ptr", spec.dtype)
-        .ptr("k_ptr", spec.dtype)
-        .ptr("v_ptr", spec.dtype)
-        .ptr("o_ptr", spec.dtype)
-        .scalar("scale", "f32")
-    )
-    if spec.runtime_shape:
-        sb = (
-            sb.scalar("batch", "i32")
-            .scalar("seqlen_q", "i32")
-            .scalar("seqlen_kv", "i32")
-        )
-    if spec.varlen:
-        sb = sb.ptr("cu_seqlens_q", "i32").ptr("cu_seqlens_kv", "i32")
+    # The kernel module owns the ABI. A hand-rolled copy here gated the shape
+    # scalars on runtime_shape, which is narrower than _has_shape_params
+    # (varlen/ragged/paged/SWA), so those launches packed misaligned kernargs.
     lch = KernelLauncher(
-        hsaco=art.hsaco, kernel_name=art.kernel_name, signature=sb.build()
+        hsaco=art.hsaco,
+        kernel_name=art.kernel_name,
+        signature=attention_dense_signature(spec),
     )
     _LAUNCHER_CACHE[key] = lch
     return lch
@@ -251,10 +241,10 @@ def bench_dense(
             stream=stream,
         )
         vals = {"q_ptr": q, "k_ptr": k, "v_ptr": v, "o_ptr": out, "scale": scale}
-        if spec.runtime_shape:
-            vals["batch"] = int(spec.batch)
-            vals["seqlen_q"] = int(spec.seqlen_q)
-            vals["seqlen_kv"] = int(spec.seqlen_kv)
+    if _has_shape_params(spec):
+        vals["batch"] = int(spec.batch)
+        vals["seqlen_q"] = int(spec.seqlen_q)
+        vals["seqlen_kv"] = int(spec.seqlen_kv)
 
     def call():
         lch(vals, config=cfg)

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

@@ -62,7 +62,8 @@ decision), and the baseline takes `wide4`, `narrow` or `narrow_d64` as
 
 Every config wins at every measured length, and the margin grows with sequence length
 — the persistent grid turns on at the long end, and the causal work per tile rises
-faster than the fixed overhead. Two things worth knowing about how to read the table:
+faster than the fixed overhead. (Measured before the current block orders; the default
+spec now goes persistent at D128.) Two things worth knowing about how to read the table:
 
 - **The D64 rows moved a long way, and the D64 K-bank-conflict pad is why.** An
   earlier comparison, run before the pad was adopted, measured D64 at −17 % (fp16,
@@ -181,7 +182,8 @@ magnitudes live outside the repo per `AGENTS.md`.
 | Per-config `waves_per_eu` | **bf16 D64** → 4 | forces the allocator low enough that a second workgroup co-resides (1 → 2 WG/CU) | **shipped** — large at long sequences |
 | D64 K-bank-conflict pad | **D64 both dtypes** | 2-row-group boundary pad takes the `do_qk` K reads from 32-way to 4-way | **shipped** — large, cross-part confirmed |
 | `lds_k_group_pad` default (8) confirmed | **D64 both dtypes** | gfx942 reuses the shared `AttentionDenseSpec.lds_k_group_pad` field; default 8 inherits from the gfx950 sweep (840 configs, pad ∈ {0,8,16,24,32} × bn × GQA × mode × seqlen). Whole-wave bank model and decision record: `library/builders/gfx950/attention/prefill/README.md §Tuning`. | **adopted** — default unchanged |
-| Persistent grid-stride | all | `num_persistent` CTAs grid-stride over decoded work items; qb-major and hkv-major decodes | **shipped** — large at long sequences; auto-on when work fills the grid |
+| Persistent grid-stride | all | `num_persistent` CTAs grid-stride over decoded work items; `persist_decode` orders them (auto: `hq_minor_swz` for large causal MHA, `bt_hkv_minor` for causal below `chiplet_num_xcds` batches) | **shipped** — default at D128 |
+| Default-grid block order | causal | `nonpersist_decode`: batch, then kv head, fastest with longest-first query blocks (`bt_hkv_minor`); Swizzled Head-first (`hq_minor_swz`) for large causal MHA | **shipped** — auto; see the [case study](../../../../../platform/dsl_docs/architecture/attention_thread_block_mapping_case_study.md) |
 
 ### Evaluated and rejected
 
@@ -207,6 +209,7 @@ python attention_dense_prefill.py                        # parity + bench, defau
 python attention_dense_prefill.py --dtype fp16 --d 64
 python attention_dense_prefill.py --persistent --np 304  # persistent grid
 python attention_dense_prefill.py --use-sinks --sw 256   # SWA + sinks
+python attention_dense_prefill.py --nonpersist-decode hq_minor_swz  # pin the default-grid order
 ```
 
 Full-cohort parity and perf are driven by the live harness at

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

@@ -106,6 +106,7 @@
     "lds_k_group_pad",
     "num_persistent",
     "persist_decode",
+    "nonpersist_decode",
 )
 
 
@@ -156,10 +157,17 @@ def add_dense_tuning_args(ap: argparse.ArgumentParser) -> None:
     ap.add_argument(
         "--persist-decode",
         dest="persist_decode",
-        choices=["auto", "qb_major", "hkv_major"],
+        choices=sorted(Gfx942AttentionDenseSpec.supported_persist_decodes()),
         default=None,
         help="persistent work-item decode; default = dispatch's 'auto'",
     )
+    ap.add_argument(
+        "--nonpersist-decode",
+        dest="nonpersist_decode",
+        choices=sorted(Gfx942AttentionDenseSpec.supported_nonpersist_decodes()),
+        default=None,
+        help="non-persistent block order; default = dispatch's 'auto'",
+    )
     ap.add_argument(
         "--interleave",
         dest="interleave",
@@ -515,10 +523,15 @@ def main():
         )
         spec = resolve_dense_spec(req, overrides)
         if args.dry_run:
+            decode = (
+                spec.resolved_persist_decode
+                if spec.persistent
+                else spec.resolved_nonpersist_decode
+            )
             print(
                 f"Sq={sq:<6d} {describe_dense_spec(spec, overrides)}  "
                 f"(persistent={spec.persistent} np={spec.num_persistent} "
-                f"decode={spec.persist_decode} wpe={spec.waves_per_eu} "
+                f"decode={decode} wpe={spec.waves_per_eu} "
                 f"bn={spec.block_n} kpad={spec.lds_k_group_pad})"
             )
             continue

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

@@ -386,9 +386,10 @@ byte-identically for D=128, so D=64 support costs the D=128 schedule nothing.
 
 Same tiling and MFMA/LDS schedule, different outer work assignment:
 
-- **Default** — one CTA per $(\text{q-block}, \text{query-head}, \text{batch})$.
-  Grid `(⌈S_q/256⌉, H_q, B)`. Simple; the per-CTA launch/dispatch + scalar setup
-  + K/V-prime cold-start (~4.5 tile-equivalents) is paid once **per query block**.
+- **Default** — one CTA per $(\text{q-block}, \text{query-head}, \text{batch})$,
+  laid out by `nonpersist_decode` (below). Simple; the per-CTA launch/dispatch +
+  scalar setup + K/V-prime cold-start (~4.5 tile-equivalents) is paid once **per
+  query block**.
 - **Persistent (grid-stride)** — a 1-D grid of `num_persistent` long-lived CTAs
   (256 = exactly one 8-wave block per CU on MI355X's 256 CUs at 2 waves/SIMD).
   Each CTA grid-strides over the flattened work-item space
@@ -407,17 +408,49 @@ $(\text{qb}, h_q, \text{bt})$ decides load balance *and* L2 locality:
 
 - **`qb_major`** — `wi = qb·(H_q·B) + h_q·B + bt`. Putting the triangular causal
   cost index `qb` in the MSB spreads cheap + expensive query blocks across each
-  CTA under grid-stride. But every 256-CTA grid-stride phase spans *all* KV
-  heads at once → large L2 footprint (57% L2 hit at GQA-8).
-- **`hkv_major`** — `wi = hkv·(NQB·g·B) + blk·(g·B) + h_ql·B + bt`, with `blk`
-  folded to a **low/high-paired** query-block index (`blk < half → qb = blk`;
-  else `qb = NQB−1−(blk−half)`). Putting `hkv` in the MSB keeps each grid-stride
-  phase within ~1 KV head so the shared GQA K/V stays L2-resident across its $g$
-  query heads (**L2 hit 57% → ~93%, HBM misses 5.9× lower**); the low/high qb
-  pairing preserves `qb_major`'s causal-triangle balance. Valid only when the
-  CTA grid-strides across both halves of a KV head ($g\cdot NQB\cdot B \ge 2\,NP$).
-- **`auto`** (default) — `hkv_major` when it is balance-safe **and** GQA
-  ($g>1$), else `qb_major`. Strictly ≥ `qb_major`.
+  CTA under grid-stride; under causal masking `qb` is also folded (as in
+  `hkv_major` below; `interleave` is an explicit alternative). But every 256-CTA
+  grid-stride phase spans *all* KV heads at once → large L2 footprint (57% L2 hit
+  at GQA-8).
+- **`hkv_major`** — `wi = hkv·(NQB·g·B) + blk·(g·B) + h_ql·B + bt`. Putting
+  `hkv` in the MSB keeps each grid-stride phase within ~1 KV head so the shared
+  GQA K/V stays L2-resident across its $g$ query heads (**L2 hit 57% → ~93%, HBM
+  misses 5.9× lower**). Under causal masking with even NQB, `blk` indexes a pair
+  `{p, NQB−1−p}` whose two blocks one CTA runs on consecutive grid-stride steps
+  (pair fold), so every CTA gets the same cost; otherwise `blk` is folded
+  (`blk < half → qb = blk`, else `qb = NQB−1−(blk−half)`).
+- **`bt_hkv_minor`** — `wi = ((blk·g + h_ql)·H_{kv} + hkv)·B + bt`, `blk`
+  folded under causal. Batch, then KV head, are the fastest digits; the hardware
+  assigns XCDs round-robin (`xcd = wi mod 8`), so with fewer batches than XCDs
+  each XCD streams few KV heads' K/V.
+- **`hq_minor_swz`** — Swizzled Head-first: XCD `wi mod M` owns a band of query
+  heads; under causal masking the pair fold as in `hkv_major`.
+- **`gqa_pair`, `gqa_pair_2phase`** — explicit only; see the prefill README.
+- **`auto`** (default) — for aligned causal attention: `hq_minor_swz` for large
+  MHA, else `bt_hkv_minor` with fewer batches than XCDs; otherwise `hkv_major`
+  for GQA with enough work per KV head, else `qb_major`.
+
+**Default-grid block order (`nonpersist_decode`).** The grid shape and the
+block-id → work-item map:
+
+- **`qb_minor`** — grid `(⌈S_q/256⌉, H_q, B)`; the query block is the fastest
+  digit.
+- **`bt_hkv_minor`** — grid `(B, H_q, ⌈S_q/256⌉)` with the KV head as the low
+  digit of the head axis; query blocks run longest-first under causal masking.
+- **`hq_minor_swz`** — Swizzled Head-first (arXiv 2511.02132): grid
+  `(M·⌈S_q/256⌉, H_q/M, B)`, $M$ = `chiplet_num_xcds`, so each XCD owns a
+  contiguous band of query heads; longest-first under causal masking.
+- **`auto`** (default) — `hq_minor_swz` for large causal MHA, `bt_hkv_minor`
+  for other aligned causal attention, else `qb_minor`.
+
+Both grids balance causal work (longest-first, fold, pair fold) only without a
+sliding window shorter than the keys: such a window gives every query block
+about the same work, so there the orders keep their blocks in place (`qb_major`
+ascending, `hkv_major` plain fold, as before these orders were added). A window
+at least as long cuts nothing off and counts as plain causal.
+
+Measurements, L2 counters and the orders ruled out are in the
+[dense attention thread-block mapping case study](../../../../platform/dsl_docs/architecture/attention_thread_block_mapping_case_study.md).
 
 **Measured (MI355X, bf16, D=128, causal, $S = 8192$, 128/8 GQA, 0 spill, err
 ≈1.46e-3).** Absolute MI355X TFLOPS swing **±25–30% with auto-clock**, so only
@@ -429,14 +462,16 @@ session, each pinned to its config (grid / decode / V-pad / lazy):
 | default grid | one-CTA/q-block | — | 32 | on | ≈543 |
 | persistent baseline | persistent NP=256 | qb-major | 0 | off | ≈877 |
 | persistent + V-pad | persistent NP=256 | qb-major | 32 | off | ≈912 |
-| **persistent (shipped default)** | persistent NP=256 | **hkv-major** | 32 | on | **≈948** |
+| **persistent (then default)** | persistent NP=256 | **hkv-major** | 32 | on | **≈948** |
 
 The clock-invariant deltas are the load-bearing part: **hkv/qb ≈ 1.04×** (L2 hit
-57%→~93%), **V-pad 0→32 ≈ +5%**, **lazy ≈ +2%** — and the shipped default
+57%→~93%), **V-pad 0→32 ≈ +5%**, **lazy ≈ +2%** — and the then-shipped default
 (`persistent=True`, `persist_decode="auto"`, `lazy_rescale=True`,
 `lds_v_row_pad=32`, the last row) is byte-identical IR to a re-measure, so the
-ratio reproduces regardless of the absolute clock. The persistent variant is the
-production choice for dense prefill.
+ratio reproduces regardless of the absolute clock. These measurements predate
+the current block orders. The persistent variant measured ahead for dense
+prefill, except large causal MHA, where the default grid with `hq_minor_swz`
+did.
 
 ### 8.3 Ragged sequence lengths — on-chip boundary padding