diff --git a/dnn-providers/hip-kernel-provider/descriptor-packaging/tests/rocke/test_desk_check_rocke_corpus.py b/dnn-providers/hip-kernel-provider/descriptor-packaging/tests/rocke/test_desk_check_rocke_corpus.py index 68a3d2353f32..29fbdb8ecd3e 100644 --- a/dnn-providers/hip-kernel-provider/descriptor-packaging/tests/rocke/test_desk_check_rocke_corpus.py +++ b/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/.yaml) and " + "regenerate, or teach the C++ geometry that order:\n " + + "\n ".join(sorted(mismatches)[:10]) + ) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_dense_prefill_live.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_dense_prefill_live.py index 42a8604736d9..cd7060696cb5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_dense_prefill_live.py +++ b/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) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/README.md b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/README.md index 60472b383e0e..2797c1afbf3b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/README.md +++ b/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 diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/attention_dense_prefill.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/attention_dense_prefill.py index f799dfb2219a..5d481f2483a2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/attention_dense_prefill.py +++ b/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 diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/ALGORITHM.md b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/ALGORITHM.md index aa13dcd9d5a8..311fb89119b0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/ALGORITHM.md +++ b/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 diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/README.md b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/README.md index 03485acb4244..90598211e642 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/README.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/README.md @@ -38,8 +38,9 @@ Non-persistent kernels still take the shape parameters; that ABI does not make a baked diagonal safe to reuse for a different sequence-length difference. Tile/resource knobs are `block_n`, `waves_per_eu`, and `lds_k_group_pad`; -persistent scheduling knobs are `num_persistent`, `persist_decode`, `interleave`, -and `wide_lds_dma`. +persistent scheduling knobs are `num_persistent`, `interleave`, and `wide_lds_dma`; +block-order knobs are `persist_decode`, `nonpersist_decode`, and +`chiplet_num_xcds` (see `../ALGORITHM.md` §8.2). ## Persistent (grid-stride) mode @@ -49,15 +50,12 @@ a 1-D grid of `num_persistent` long-lived CTAs grid-strides over the setup + K/V-prime cold-start is amortized once per CU instead of once per query-block. This closes the causal fixed-cost amortization gap. `num_persistent=256` = one 8-wave block per CU on MI355X (256 CUs) at 2 waves/SIMD; larger oversubscribes the CUs (tail -loss). The work-item decode is `persist_decode="auto"` by default. Auto selects: - -1. **gqa-pair** when its one-phase balance equation - (`NP == NQB*Hkv*B`, with even `NQB` and GQA ratio) holds; -2. **gqa-pair-2phase** when its two-phase balance equation - (`NP == NQB*Hkv*B*gqa/2`, with even `NQB`) holds; -3. **hkv-major** when its broader GQA balance condition - (`gqa*NQB*B >= 2*NP`) holds; or -4. **qb-major** otherwise. +loss). The work-item decode is `persist_decode="auto"` by default. For aligned +causal attention auto selects `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`. Query blocks are folded +or pair-folded under causal masking without a sliding window. The GQA-pair decodes below are explicit +only. ### Balanced GQA-pair decode @@ -69,16 +67,12 @@ through L2. The explicit mode requires persistent causal attention, aligned sequence lengths, even `NQB` and GQA ratio, and `NP == NQB*Hkv*B`. It composes numerically with -`interleave`, attention sinks, and aligned sliding-window attention. Auto uses it -for aligned non-windowed causal shapes whenever the balance equation holds. -Sliding-window requests fall back to qb-major because complementary query blocks -do not have constant combined work under a finite window. MHA (`Hq == Hkv`) also -falls back because it has no grouped query heads to reuse. +`interleave`, attention sinks, and aligned sliding-window attention. Auto does not +select it: on the shapes it fits, the default block orders measured ahead. `persist_decode="gqa_pair_2phase"` assigns one query head to each CTA and processes -complementary query blocks in two grid-stride phases. Auto selects it when -`NP == NQB*Hkv*B*gqa/2`; this covers, for example, the D128 Hq32/Hkv8 S4096 and -Hq64/Hkv8 S2048 dashboard shapes at NP=256. +complementary query blocks in two grid-stride phases. It requires +`NP == NQB*Hkv*B*gqa/2` and, like `gqa_pair`, is explicit only. ### Wide LDS DMA @@ -92,12 +86,16 @@ per-row 32-bit DMA layout into FlyDSL-compatible slabs: - IGLP-1 instead of the narrow path's manual scheduling directives. The dispatcher enables it for aligned persistent causal D128/BN64 fp16 and bf16 -shapes without sliding windows, sinks, or ragged lengths. GQA-pair selection is -independent: MHA and GQA shapes whose pair equations do not hold still use wide -DMA with qb-major or hkv-major work ordering. +shapes without sliding windows, sinks, or ragged lengths, whatever the block +order. ## Measured (MI355X, bf16, D=128, Hq=128, Hkv=8, causal, Sq=8192) +> These measurements predate the current block orders. Where they show +> `hkv-major` as the default or `gqa-pair` as auto-selected, auto now picks +> `bt_hkv_minor` / `qb_major`, which measured ahead of the pair decodes on +> their shapes. + Absolute MI355X TFLOPS swing **±25–30% with auto-clock**, so only **same-session ratios are load-bearing**; the table below is one representative session, with each number pinned to its exact config (grid / decode / V-pad / lazy): @@ -270,11 +268,12 @@ There are two algorithms, one per body, and three candidates: Nothing is chosen for the caller: an unpinned request keeps the unified 2D/3D path. The tile (`block_m` 128 or 256, `block_n`) and knobs such as -`num_persistent`, `persist_decode` (`gqa_pair` / `gqa_pair_2phase`), -`interleave` and `waves_per_eu` are spec fields of the candidate; set them with -`tuning_spec_with_knobs(req, spec_id, knobs)` or replay a swept `tuning_id`. -The kernel name exposes the persistent decisions through `wdma`, `gqapair`, or -`gqapair2` tokens. +`num_persistent`, `persist_decode` (including `gqa_pair` / `gqa_pair_2phase`), +`interleave`, `nonpersist_decode` (grid body) and `waves_per_eu` are spec +fields of the candidate; set them with `tuning_spec_with_knobs(req, spec_id, +knobs)` or replay a swept `tuning_id`. The kernel name exposes the decisions +through `wdma` and the block-order tokens (e.g. `bthkvmin`, `gqapair`, +`nphqminswz`). For bottom-right masking, set `mask_type=AttentionMaskType.BOTTOM_RIGHT_CAUSAL` (exported by `dispatch.attention`). With unequal Q/K lengths, only the grid diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/attention_dense_prefill.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/attention_dense_prefill.py index 105705f8573c..aa1edf3dece5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/attention_dense_prefill.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/attention_dense_prefill.py @@ -36,17 +36,20 @@ import torch # noqa: E402 from dispatch.attention import AttentionRequest, tuning_spec_with_knobs # noqa: E402 -from kernels.common.attention_dense_spec import DENSE_TILE_GEOMETRIES # noqa: E402 +from kernels.common.attention_dense_spec import ( # noqa: E402 + DENSE_TILE_GEOMETRIES, +) from kernels.gfx950.attention_dense import ( # noqa: E402 GFX950_DENSE_LAYOUTS, Gfx950AttentionDenseSpec, + _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, @@ -69,26 +72,11 @@ def _make_launcher(spec: Gfx950AttentionDenseSpec): 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") + return KernelLauncher( + hsaco=art.hsaco, + kernel_name=art.kernel_name, + signature=attention_dense_signature(spec), ) - if spec.runtime_shape: - sb = ( - sb.scalar("batch", "i32") - .scalar("seqlen_q", "i32") - .scalar("seqlen_kv", "i32") - ) - if spec.use_sinks: - sb = sb.ptr("sink_ptr", spec.dtype) - if spec.varlen: - sb = sb.ptr("cu_seqlens_q", "i32").ptr("cu_seqlens_kv", "i32") - sig = sb.build() - return KernelLauncher(hsaco=art.hsaco, kernel_name=art.kernel_name, signature=sig) def _launch_config(spec: Gfx950AttentionDenseSpec, stream) -> LaunchConfig: @@ -178,6 +166,7 @@ def make_spec_from_shape(shape: dict[str, Any]) -> Gfx950AttentionDenseSpec: "lazy_rescale": bool, "num_persistent": int, "persist_decode": str, + "nonpersist_decode": str, } knobs = { name: coerce(shape[name]) for name, coerce in coercers.items() if name in shape @@ -197,7 +186,11 @@ def run_benchmark( ok = (not check) or (err < _TOL) return { "kernel_name": spec.kernel_name(), - "persist_decode": spec.resolved_persist_decode, + # Only the decode the body uses; the other one is None. + "persist_decode": spec.resolved_persist_decode if spec.persistent else None, + "nonpersist_decode": ( + None if spec.persistent else spec.resolved_nonpersist_decode + ), "ms": ms, "tflops": tf, "max_abs": err, @@ -236,7 +229,7 @@ def run( stream = torch.cuda.current_stream().cuda_stream cfg = _launch_config(spec, stream) vals = {"q_ptr": q, "k_ptr": k, "v_ptr": v, "o_ptr": out, "scale": scale} - if spec.runtime_shape: + if _has_shape_params(spec): vals["batch"] = int(spec.batch) vals["seqlen_q"] = int(spec.seqlen_q) vals["seqlen_kv"] = int(spec.seqlen_kv) @@ -365,13 +358,12 @@ def main(): ap.add_argument( "--persist-decode", default="auto", - choices=[ - "auto", - "qb_major", - "hkv_major", - "gqa_pair", - "gqa_pair_2phase", - ], + choices=sorted(Gfx950AttentionDenseSpec.supported_persist_decodes()), + ) + ap.add_argument( + "--nonpersist-decode", + default="auto", + choices=sorted(Gfx950AttentionDenseSpec.supported_nonpersist_decodes()), ) ap.add_argument( "--sw", type=int, default=0, help="sliding_window (0=off; multiple of --bn)" @@ -415,6 +407,7 @@ def main(): "sliding_window": args.sw, "use_sinks": args.use_sinks, "persist_decode": args.persist_decode, + "nonpersist_decode": args.nonpersist_decode, "wide_lds_dma": args.wide_lds_dma, } spec = make_spec_from_shape(shape) @@ -450,6 +443,7 @@ def main(): sliding_window=args.sw, use_sinks=args.use_sinks, persist_decode=args.persist_decode, + nonpersist_decode=args.nonpersist_decode, wide_lds_dma=args.wide_lds_dma, ) run(spec) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/AGENTS.md b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/AGENTS.md index 8cd99c9162ba..c748008c8bfb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/AGENTS.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/AGENTS.md @@ -166,10 +166,13 @@ template and its DS-read count, each on its own axis), exp2 and PV-loop codegen, the tile (`block_m` 128/256, `block_n` 32/64/128; `ragged` follows it), the O store width, `num_persistent` (symbolic policies such as `gqa_pair` and CU multiples, resolved per problem and per `block_m`), -`persist_decode`, and `interleave`. `production` sets every applicable knob -to each legal value one at a time from the default spec, pairing a value with -its prerequisite when it is illegal alone, and repeats that pass at each -`block_m` (the query tile changes what every other knob does); `full` walks +`persist_decode`, `interleave`, and `nonpersist_decode` (the grid body's +block order; both decode axes take their values from the decode registries in +`kernels/common/attention_dense_decode.py`). `production` sets every +applicable knob to each legal value one at a time from the default spec, +pairing a value with its prerequisite when it is illegal alone, and repeats +that pass at each `block_m` (the query tile changes what every other knob +does); `full` walks or samples the pruned product. The kernel's spec validator is the only legality gate. An explicit value equal to its policy, or a knob the body does not read for that spec, re-emits the same IR @@ -183,8 +186,9 @@ gfx942 dense uses the same walk over `_GFX942_DENSE_AXES` knob there alongside `block_m` and `block_n`, and the persistent-only knobs are pruned per spec. Its knobs are the K row and V row pads, the D64 K group pad, the conflict-free-V store and its swizzle, -exp2, `iglp` / `iglp_mode`, the PV fence mask, `pv_priority`, -`pv_loop_order`, the bf16 O store width, and `causal_diag_split`. The +the persistent and non-persistent block orders, exp2, `iglp` / `iglp_mode`, +the PV fence mask, `pv_priority`, `pv_loop_order`, the bf16 O store width, +and `causal_diag_split`. The LDS-saving knobs lead as enablers so an over-budget tile can still reach the setting that fits. Untunable here: `lds_num_buffers` (only 1 is implemented) and `lazy_rescale` (the body never reads it). diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/axes.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/axes.py index 406592d2fb82..ca8ac6783bf2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/axes.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/axes.py @@ -12,7 +12,12 @@ from typing import Mapping, Tuple -from kernels.gfx950.attention_dense import gfx950_dense_tile_legality +from kernels.common.attention_dense_decode import NONPERSIST_DECODES, PERSIST_DECODES +from kernels.gfx942.attention_dense import Gfx942AttentionDenseSpec +from kernels.gfx950.attention_dense import ( + Gfx950AttentionDenseSpec, + gfx950_dense_tile_legality, +) from rocke.dispatch.tuning.axes import ( # noqa: F401 (re-exported for the rules) KnobAxis, @@ -26,6 +31,7 @@ values as _values, ) + # AMDGPU ``sched_barrier`` ABI: 0 is a full fence, and each bit lets one # instruction class cross it (ALU, VALU, SALU, MFMA, VMEM, VMEM rd, VMEM wr, # DS, DS rd, DS wr, transcendental). 0x108 is MFMA + DS-read. @@ -151,6 +157,9 @@ "num_kv_blocks", "use_sinks", "causal_bottom_right", + # A fact of the target GPU, not a tuning choice: the base spec keeps + # its default (8 XCDs). + "chiplet_num_xcds", } ) # Problem fields only some arches' dense specs declare (the band's right bound). @@ -193,7 +202,7 @@ "gqa_pair_2phase", "work", ) -# gfx942 implements only the qb_major / hkv_major decodes. +# gfx942 has no gqa_pair decodes, so none of their exact CTA counts. _GFX942_NUM_PERSISTENT_POLICIES = _NUM_PERSISTENT_MULTIPLES + ("work",) _GFX942_LDS_ROW_PADS = (0, 4, 8, 12, 16, 24, 32) _GFX942_V_ROW_PADS = (0, 8, 16, 32, 64) @@ -220,13 +229,20 @@ def _num_persistent_axis(policies: Tuple[str, ...]) -> KnobAxis: ) +def _decode_axis(name: str, registry: Mapping, supported) -> KnobAxis: + """``name`` over the registry's decodes (in registry order) that the arch's + spec accepts; ``auto`` is the default.""" + return _values(name, "auto", tuple(d for d in registry if d in supported)) + + # One axis per spec field. The PV scheduling knobs lead as enablers: IGLP needs # the manual fence and sched_group template off, and a manual one needs IGLP # off, so each can only be reached once the others are decided. Then # prerequisites-first: the tile sets the query-block count the CTA counts and # decodes are exact for, persist_decode's gqa_pair modes need an exact -# num_persistent, interleave is read only on the resolved qb_major decode, and -# the fence mask / DS-read count / rescale threshold follow their parents. +# num_persistent, interleave is read only on the resolved qb_major decode, +# nonpersist_decode only by the grid body, and the fence mask / DS-read count / +# rescale threshold follow their parents. _GFX950_DENSE_AXES: Tuple[KnobAxis, ...] = ( _values("pv_sched_fence", None, (True, False), enabler=True), _values("pv_sched_group_template", None, (True, False), enabler=True), @@ -234,12 +250,17 @@ def _num_persistent_axis(policies: Tuple[str, ...]) -> KnobAxis: _choices_axis("block_m", _GFX950_BLOCK_M), _choices_axis("block_n", _GFX950_BLOCK_N), _num_persistent_axis(_GFX950_NUM_PERSISTENT_POLICIES), - _values( + _decode_axis( "persist_decode", - "auto", - ("qb_major", "hkv_major", "gqa_pair", "gqa_pair_2phase"), + PERSIST_DECODES, + Gfx950AttentionDenseSpec.supported_persist_decodes(), ), _flag("interleave"), + _decode_axis( + "nonpersist_decode", + NONPERSIST_DECODES, + Gfx950AttentionDenseSpec.supported_nonpersist_decodes(), + ), _values("lazy_rescale", True, (False,)), _values("lazy_rescale_threshold", 8.0, _DENSE_LAZY_THRESHOLDS), _values("pv_sched_fence_mask", 0, _SCHED_BARRIER_MASKS), @@ -262,7 +283,8 @@ def _num_persistent_axis(policies: Tuple[str, ...]) -> KnobAxis: # limit (e.g. D64 block_n=256 fits only with lds_k_group_pad=0). Then # prerequisites-first: exp2 policy reads persistent, v_row_pad policy reads # block_n, a non-derived v_row_pad needs the swizzle off, interleave is read -# only on the resolved qb_major decode, and iglp_mode / the PV fence follow iglp. +# only on the resolved qb_major decode, nonpersist_decode only on the +# non-persistent grid, and iglp_mode / the PV fence follow iglp. _GFX942_DENSE_AXES: Tuple[KnobAxis, ...] = ( _values("use_cfvst", None, (True, False), enabler=True), _values("lds_row_pad", 8, _GFX942_LDS_ROW_PADS, enabler=True), @@ -271,8 +293,17 @@ def _num_persistent_axis(policies: Tuple[str, ...]) -> KnobAxis: _choices_axis("block_m", _GFX942_BLOCK_M), _choices_axis("block_n", _GFX942_BLOCK_N), _num_persistent_axis(_GFX942_NUM_PERSISTENT_POLICIES), - _values("persist_decode", "auto", ("qb_major", "hkv_major")), + _decode_axis( + "persist_decode", + PERSIST_DECODES, + Gfx942AttentionDenseSpec.supported_persist_decodes(), + ), _flag("interleave"), + _decode_axis( + "nonpersist_decode", + NONPERSIST_DECODES, + Gfx942AttentionDenseSpec.supported_nonpersist_decodes(), + ), _values("use_v_swizzle", None, (True, False)), _values("v_row_pad", None, _GFX942_V_ROW_PADS), _values("use_exp2_fast", None, (True, False)), diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/dense_rules.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/dense_rules.py index 89f48093ce68..9caa736068b6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/dense_rules.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/dense_rules.py @@ -98,6 +98,10 @@ class _DenseTuningArch: "num_persistent": lambda s: bool(s.persistent), "persist_decode": lambda s: bool(s.persistent), "interleave": lambda s: bool(s.persistent) and bool(s.causal), + # A knob only where the variant leaves the order to auto: a variant + # that sets it (qbminor) keeps it in sweeps and ignores the knob. + "nonpersist_decode": lambda s: not s.persistent + and s.nonpersist_decode == "auto", # Wide DMA locks the K/V slab pads. "lds_k_row_pad": lambda s: int(s.head_size) == 128 and not s.wide_lds_dma, "lds_v_row_pad": lambda s: int(s.head_size) == 128 and not s.wide_lds_dma, @@ -254,6 +258,17 @@ def _dense_redundant_knobs(spec, arch: str) -> dict: redundant["persist_decode"] = ( f"persist_decode={spec.persist_decode!r} is what auto resolves to" ) + if spec.persistent: + if spec.nonpersist_decode != "auto": + redundant["nonpersist_decode"] = "read only by the non-persistent body" + elif ( + spec.nonpersist_decode != "auto" + and replace(spec, nonpersist_decode="auto").resolved_nonpersist_decode + == spec.nonpersist_decode + ): + redundant["nonpersist_decode"] = ( + f"nonpersist_decode={spec.nonpersist_decode!r} is what auto resolves to" + ) if spec.interleave and "interleave" not in redundant: nqb = -(-int(spec.seqlen_q) // int(spec.block_m)) if not spec.causal or spec.resolved_persist_decode != "qb_major" or nqb < 2: diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/gfx942_dense.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/gfx942_dense.py index 612658c7cd31..e5e39fb885f3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/gfx942_dense.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/gfx942_dense.py @@ -44,8 +44,8 @@ def _base_spec(req: AttentionRequest): """The default ``Gfx942AttentionDenseSpec`` for ``req``. The gfx942 twin of ``gfx950_dense._base_spec``. The persistent grid turns on - once there is enough work to fill it (``nqb*Hq*B >= num_persistent``), and - non-tile-multiple self-attention lengths take the on-chip ragged path. + at D128, and non-tile-multiple self-attention lengths take the on-chip + ragged path. ``waves_per_eu`` comes from the kernel's own per-config policy (``_tuned_waves_per_eu``): any value the kernel bakes into its ``kernel_name`` must be resolved by the kernel's policy, or the name tag and @@ -74,8 +74,9 @@ def _base_spec(req: AttentionRequest): # on-chip ragged padding for ragged self-attention lengths (seqlen_q==seqlen_kv, # not a 256/block_n multiple). Cross-attention ragged is left to the validator. ragged = (sq == sk) and ((sq % bm != 0) or (sk % bn != 0)) - nqb = (sq + bm - 1) // bm - work = nqb * int(req.nhead_q) * int(req.batch) + # The persistent grid measured ahead at D128 for every batch and mask; at + # D64 the non-persistent grid did. + persistent = head_size == 128 return Gfx942AttentionDenseSpec( batch=int(req.batch), seqlen_q=sq, @@ -87,7 +88,7 @@ def _base_spec(req: AttentionRequest): dtype=dtype, block_m=bm, block_n=bn, - persistent=work >= _GFX942_NUM_PERSISTENT, + persistent=persistent, num_persistent=_GFX942_NUM_PERSISTENT, persist_decode="auto", ragged=ragged, @@ -134,7 +135,7 @@ def _make_gfx942_attention_dense_candidate() -> KernelCandidate: # LSE output; no moving bottom-right diagonal. Head size stays out: # D64/D128 coverage is ``supports_attention_dense``'s call. features=frozenset({"causal", "sliding_window", "sinks", "lse"}), - # _base_spec picks these from the work size and dtype/head size. + # _base_spec picks these from the dtype and head size. recorded=frozenset({"persistent", "waves_per_eu"}), ) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/gfx950_dense.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/gfx950_dense.py index bfc5ad665f81..49a204aecaa0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/gfx950_dense.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/gfx950_dense.py @@ -52,6 +52,7 @@ class Gfx950DenseVariant: persistent: bool wide_lds_dma: bool features: frozenset + nonpersist_decode: str = "auto" @property def candidate_name(self) -> str: @@ -70,6 +71,16 @@ def spec_id(self) -> str: wide_lds_dma=False, features=_FEATURES | {"causal_bottom_right"}, ), + # The grid body at the qb_minor block order: the grid the hipDNN kernel ingestor + # launches. + Gfx950DenseVariant( + "qbminor", + DENSE_GRID_ALGORITHM, + persistent=False, + wide_lds_dma=False, + features=_FEATURES | {"causal_bottom_right"}, + nonpersist_decode="qb_minor", + ), Gfx950DenseVariant( "persist", DENSE_PERSIST_ALGORITHM, @@ -190,6 +201,7 @@ def _base_spec(req: AttentionRequest, variant: Gfx950DenseVariant): sliding_window=int(req.sliding_window), use_sinks=bool(req.use_sinks), wide_lds_dma=variant.wide_lds_dma, + nonpersist_decode=variant.nonpersist_decode, causal_bottom_right=moving_bottom_right, emit_lse=bool(req.emit_lse), right_bound=right_bound, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_dense_decode.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_dense_decode.py new file mode 100644 index 000000000000..4b1da0f779f5 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_dense_decode.py @@ -0,0 +1,658 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Work-item decodes for the dense attention kernel. + +A decode maps a CTA's work index to the (query block, query head, batch) it +computes. Every decode is a bijection onto the work space, so the choice only +moves which items run together (L2 locality, causal load balance), never the +result. One class per ``persist_decode`` / ``nonpersist_decode`` value keeps +each decode's IR, grid, legality check and kernel-name tag together; both arch +builders call the same class. + +The shape reaches a decode through :class:`DecodeShape`: each extent is either +baked (emitted as a constant at its use) or a runtime value from the kernel's +params, so one decode body serves both, and a shape-dependent choice (causal +balancing under a window, the pair fold) is a compile-time branch when baked and +a ``select`` on a CTA-uniform flag when runtime, so one binary serves every shape +of a runtime-shape kernel. + +A decode takes only the shape and the block or work index: the shape carries the +spec and the IR builder ``b`` the decode emits through. +""" + +from __future__ import annotations + +from typing import Any, ClassVar, NamedTuple + + +class Decoded(NamedTuple): + """A decoded work item. ``hkv`` is set only where the decode has the kv + head as a digit; otherwise callers derive it as ``hq // gqa``.""" + + qb: Any + hq: Any + bt: Any + hkv: Any = None + + +def _spec_nqb(spec) -> int: + """Query blocks per head for the spec's shape (a ragged last block counts).""" + return (spec.seqlen_q + spec.block_m - 1) // spec.block_m + + +def _registry(*classes) -> dict: + return {c.name: c() for c in classes} + + +def balance_causal(spec) -> bool: + """Whether the decodes balance causal work (longest-first, fold, pair fold). + A sliding window shorter than the keys gives every query block about the + same work, so there it only moves blocks; a longer one cuts nothing off.""" + return spec.causal and not 0 < spec.sliding_window < spec.seqlen_kv + + +class DecodeShape: + """The problem shape as the decodes see it. + + ``values`` maps an extent (``batch``, ``seqlen_q``, ``seqlen_kv``, + ``num_query_heads``, ``num_kv_heads``, ``gqa``, ``nqb``, ``w``) to the runtime + IR value the kernel holds; any other extent is baked and emitted as a constant + at its use. The decodes are written once on these values, and LLVM folds the + baked ones. ``nqb`` and ``w``, when not passed, are derived from ``seqlen_q`` + (ceil) and the heads/batch at each use. + + A runtime divide goes through ``fastdiv(n, name)`` when the builder passes one + (the persistent loop: host-computed magic/shift pairs, no division in the + loop), else it is a plain division (once per CTA on the default grid). + ``scale`` divides an extent by a baked factor that divides it: ``n // (d/k)`` + is computed as ``(n*k) // d``, which needs only ``d``'s fastdiv pair. + """ + + def __init__(self, b, spec, values=None, fastdiv=None): + self.b, self.spec = b, spec + self.values = dict(values or {}) + self.fastdiv = fastdiv + + def _baked(self, name: str) -> int: + s = self.spec + nqb = _spec_nqb(s) + return { + "batch": s.batch, + "seqlen_q": s.seqlen_q, + "seqlen_kv": s.seqlen_kv, + "num_query_heads": s.num_query_heads, + "num_kv_heads": s.num_kv_heads, + "gqa": s.num_queries_per_kv, + "nqb": nqb, + "w": nqb * s.num_query_heads * s.batch, + }[name] + + def runtime(self, name: str) -> bool: + if name in self.values: + return True + if name == "nqb": + return self.runtime("seqlen_q") + if name == "w": + return any(self.runtime(n) for n in ("nqb", "num_query_heads", "batch")) + return False + + def static(self, name: str) -> int | None: + """The baked value, or None when ``name`` is runtime.""" + return None if self.runtime(name) else self._baked(name) + + def val(self, name: str, scale: int = 1): + """``name // scale`` as an IR value.""" + b = self.b + if not self.runtime(name): + return b.const_i32(self._baked(name) // scale) + if name in self.values: + v = self.values[name] + elif name == "nqb": + bm = self.spec.block_m + q = self.values["seqlen_q"] + v = b.div(b.add(q, b.const_i32(bm - 1)), b.const_i32(bm)) + else: # w + v = b.mul( + b.mul(self.val("nqb"), self.val("num_query_heads")), self.val("batch") + ) + return v if scale == 1 else b.div(v, b.const_i32(scale)) + + def div(self, n, name: str, scale: int = 1): + """``n // (name // scale)``.""" + b = self.b + if not self.runtime(name): + return b.div(n, b.const_i32(self._baked(name) // scale)) + if self.fastdiv is None: + return b.div(n, self.val(name, scale)) + if scale != 1: + n = b.mul(n, b.const_i32(scale)) + return self.fastdiv(n, name) + + def mod(self, n, name: str, scale: int = 1): + """``n % (name // scale)``.""" + b = self.b + if not self.runtime(name): + return b.mod(n, b.const_i32(self._baked(name) // scale)) + if self.fastdiv is None: + return b.mod(n, self.val(name, scale)) + return b.sub(n, b.mul(self.div(n, name, scale), self.val(name, scale))) + + def balance(self): + """:func:`balance_causal` for this shape: a bool when it is fixed at + compile time, else an i1 (a window is baked, the keys may not be).""" + spec = self.spec + if not spec.causal: + return False + W = spec.sliding_window + if W <= 0: + return True + if not self.runtime("seqlen_kv"): + return not W < spec.seqlen_kv + b = self.b + return b.cmp_ge(b.const_i32(W), self.values["seqlen_kv"]) + + def pair_fold(self, NP: int): + """Whether the pair fold applies: balancing, even NQB and more than one + round of work (``W > NP``). A bool or an i1, as :meth:`balance`.""" + bal = self.balance() + if bal is False: + return False + nqb = self.static("nqb") + if nqb is not None and not self.runtime("w"): + return bal is True and nqb % 2 == 0 and self._baked("w") > NP + b = self.b + flag = b.land( + b.cmp_eq(b.land(self.val("nqb"), b.const_i32(1)), b.const_i32(0)), + b.cmp_lt(b.const_i32(NP), self.val("w")), + ) + return flag if bal is True else b.land(flag, bal) + + def branch(self, flag, then, other): + """``then()`` if ``flag`` else ``other()``: only the taken side is + emitted when ``flag`` is a bool; with an i1 both are emitted and the + results (a value or a tuple of values) selected.""" + if isinstance(flag, bool): + return then() if flag else other() + t, o = then(), other() + if isinstance(t, tuple): + return tuple(self.b.select(flag, x, y) for x, y in zip(t, o)) + return self.b.select(flag, t, o) + + +# --------------------------------------------------------------------------- # +# Non-persistent grid: one CTA per work item. +# --------------------------------------------------------------------------- # + + +def emit_reverse(b, blk, shape: DecodeShape): + """Longest-first query-block order for causal cost: ``NQB-1-blk``. + + Under a runtime shape one binary serves every seqlen, so NQB must come from + the param; a baked NQB would reverse against the wrong extent. + """ + return b.sub(b.sub(shape.val("nqb"), b.const_i32(1)), blk) + + +class NonpersistDecode: + """Block ids -> work item, for the one-CTA-per-item grid. + + The hardware dispatches workgroups x-fastest and round-robins them over the + XCDs, so the grid shape and the decode together decide which items share an + XCD's L2 and which run last. + """ + + name: ClassVar[str] + tag: ClassVar[str] = "" # kernel-name part; "" = none + + def check(self, spec) -> str | None: + """Why ``spec`` cannot use this decode, or None.""" + return None + + def grid(self, spec) -> tuple[int, int, int]: + """Launch grid, from the spec's actual (host-side) shape.""" + raise NotImplementedError + + def emit_decode(self, shape: DecodeShape, bx, by, bz) -> Decoded: + raise NotImplementedError + + +class NonpersistQbMinor(NonpersistDecode): + """Grid ``(nqb, Hq, B)``: the query block is the fastest digit.""" + + name = "qb_minor" + + def grid(self, spec): + return (_spec_nqb(spec), spec.num_query_heads, spec.batch) + + def emit_decode(self, shape, bx, by, bz): + del shape + return Decoded(bx, by, bz) + + +class NonpersistBtHkvMinor(NonpersistDecode): + """Grid ``(B, Hq, nqb)`` with the kv head as the low digit of the head axis: + batch, then kv head, are the fastest digits, so the XCD round-robin spreads + distinct K/V streams over the XCDs. Query blocks run longest-first under + causal masking without a sliding window.""" + + name = "bt_hkv_minor" + tag = "npbthkvmin" + + def grid(self, spec): + return (spec.batch, spec.num_query_heads, _spec_nqb(spec)) + + def emit_decode(self, shape, bx, by, bz): + b = shape.b + hq = b.add( + b.mul(shape.mod(by, "num_kv_heads"), shape.val("gqa")), + shape.div(by, "num_kv_heads"), + ) + qb = shape.branch( + shape.balance(), lambda: emit_reverse(b, bz, shape), lambda: bz + ) + return Decoded(qb, hq, bx) + + +class NonpersistHqMinorSwz(NonpersistDecode): + """Swizzled Head-first (arXiv 2511.02132, Fig. 11): grid + ``(M*nqb, Hq/M, B)`` with M = ``chiplet_num_xcds``, so XCD ``a`` owns the + contiguous query-head band ``[a*Hq/M, (a+1)*Hq/M)`` and runs all its query + blocks before the next head. Batch is the slowest digit. Query blocks run + longest-first under causal masking without a sliding window.""" + + name = "hq_minor_swz" + tag = "nphqminswz" + + def check(self, spec): + if spec.num_query_heads % spec.chiplet_num_xcds != 0: + return ( + f"{self.name} needs num_query_heads ({spec.num_query_heads}) " + f"divisible by chiplet_num_xcds ({spec.chiplet_num_xcds})" + ) + return None + + def grid(self, spec): + M = spec.chiplet_num_xcds + return (M * _spec_nqb(spec), spec.num_query_heads // M, spec.batch) + + def emit_decode(self, shape, bx, by, bz): + b, spec = shape.b, shape.spec + M = spec.chiplet_num_xcds + # Head band of Hq/M heads per XCD. + hq = b.add( + b.mul(b.mod(bx, b.const_i32(M)), shape.val("num_query_heads", M)), by + ) + blk = b.div(bx, b.const_i32(M)) + qb = shape.branch( + shape.balance(), lambda: emit_reverse(b, blk, shape), lambda: blk + ) + return Decoded(qb, hq, bz) + + +NONPERSIST_DECODES = _registry( + NonpersistQbMinor, NonpersistBtHkvMinor, NonpersistHqMinorSwz +) + + +# --------------------------------------------------------------------------- # +# Persistent grid: NP CTAs grid-stride over the work index ``wi``. +# --------------------------------------------------------------------------- # + + +def _baked_only(shape: DecodeShape, name: str) -> int: + """NQB for the decodes that exist only on baked-shape bodies (gfx950).""" + nqb = shape.static("nqb") + if nqb is None: + raise NotImplementedError(f"persist_decode={name!r} needs a baked shape") + return nqb + + +def emit_fold(b, blk, shape: DecodeShape): + """Causal fold: ``blk < half -> blk``, else ``NQB-1-(blk-half)``, so a CTA + striding both halves pairs a cheap block with an expensive one.""" + nqb = shape.val("nqb") + half = b.lshr(nqb, b.const_i32(1)) + qb_hi = b.sub(b.add(b.sub(nqb, b.const_i32(1)), half), blk) # NQB-1-(blk-half) + return b.select(b.cmp_lt(blk, half), blk, qb_hi) + + +def emit_pair_unit(b, wi, shape: DecodeShape, NP: int): + """``(unit, half)`` for ``wi < W`` (``W`` even): unit ``u`` holds items + ``2u, 2u+1``, and one CTA runs both halves on consecutive grid-stride steps. + In the last round, CTAs that still get two items get both halves of a unit. + + Branch-free: every case is emitted and selected, each correction a no-op where + it does not apply (no whole double rounds, no tail, no paired tail); the only + divisors are NP, 2*NP and 2.""" + c = b.const_i32 + W = shape.val("w") + # Items before the last round: whole double rounds of 2*NP items. + rounds_end = b.mul(b.div(W, c(2 * NP)), c(2 * NP)) + re_half = b.lshr(rounds_end, c(1)) + paired_tail = b.sub(b.sub(W, rounds_end), c(NP)) # may be <= 0 + pt = b.smax(paired_tail, c(0)) + j = b.div(wi, c(NP)) + unit = b.add(b.sub(wi, b.mul(j, c(NP))), b.mul(b.div(j, c(2)), c(NP))) + half = b.mod(j, c(2)) + # Last round: CTAs below `paired_tail` still get both halves of a unit. + r = b.sub(wi, rounds_end) + s = b.sub(r, pt) + unit_t = b.add(b.add(re_half, pt), b.div(s, c(2))) + half_t = b.mod(s, c(2)) + lo = b.cmp_lt(r, paired_tail) # never when paired_tail <= 0 + hi = b.cmp_lt(c(NP - 1), r) # never when paired_tail <= 0 (tail <= NP) + unit_t = b.select( + lo, + b.add(re_half, r), + b.select(hi, b.add(b.sub(re_half, c(NP)), r), unit_t), + ) + half_t = b.select(lo, c(0), b.select(hi, c(1), half_t)) + in_rounds = b.cmp_lt(wi, rounds_end) + return b.select(in_rounds, unit, unit_t), b.select(in_rounds, half, half_t) + + +def emit_pair_block(b, p, half, shape: DecodeShape): + """Query block of fold pair ``p``: ``p`` for the first half, ``NQB-1-p`` for + the second, so every unit costs the same under causal masking.""" + first = b.cmp_lt(half, b.const_i32(1)) + last = b.sub(shape.val("nqb"), b.const_i32(1)) + return b.select(first, p, b.sub(last, p)) + + +class PersistDecode: + """Grid-stride work index ``wi`` -> work item, for the persistent grid. + + ``fastdiv_divisors``: the extents a runtime-shape body divides by in the + grid-stride loop, each needing a magic/shift kernarg pair; empty for the + decodes that only baked-shape bodies emit.""" + + name: ClassVar[str] + tag: ClassVar[str] = "" # kernel-name part; "" = none + fastdiv_divisors: ClassVar[tuple[str, ...]] = () + + def check(self, spec) -> str | None: + """Why ``spec`` cannot use this decode, or None.""" + return None + + def fastdiv_dividend_bound(self, spec) -> int: + """An upper bound on every runtime fastdiv dividend: the work count.""" + return _spec_nqb(spec) * spec.num_query_heads * spec.batch + + def emit_decode(self, shape: DecodeShape, wi) -> Decoded: + raise NotImplementedError + + +class PersistQbMajor(PersistDecode): + """``wi = qb*(Hq*B) + hq*B + bt``: the query block is the slowest digit, so + grid-stride spreads cheap and expensive blocks over each CTA. Under causal + masking without a sliding window the blocks are folded; ``interleave`` + instead reverses the block for odd ``hq`` (measured behind the fold, so auto + never sets it). Keying the reversal on ``hq`` keeps it a bijection: each head's + blocks are either in order or reversed.""" + + name = "qb_major" + fastdiv_divisors = ("batch", "num_query_heads", "gqa") + + def emit_decode(self, shape, wi): + b, spec = shape.b, shape.spec + bt = shape.mod(wi, "batch") + rem = shape.div(wi, "batch") + hq = shape.mod(rem, "num_query_heads") + qb0 = shape.div(rem, "num_query_heads") + if spec.interleave and spec.causal: + # With NQB == 1, qb0 is 0 and the flip is too. + odd = b.cmp_eq(b.mod(hq, b.const_i32(2)), b.const_i32(1)) + last = b.sub(shape.val("nqb"), b.const_i32(1)) + qb = b.select(odd, b.sub(last, qb0), qb0) + else: + qb = shape.branch( + shape.balance(), lambda: emit_fold(b, qb0, shape), lambda: qb0 + ) + return Decoded(qb, hq, bt) + + +class PersistHkvMajor(PersistDecode): + """``wi = hkv*(NQB*gqa*B) + blk*(gqa*B) + hql*B + bt``: each grid-stride + phase stays within about one kv head, so its K/V stays in L2 across the GQA + group. Under causal masking without a sliding window, with even NQB and + ``W > NP``, one CTA runs both blocks of a fold pair ``{p, NQB-1-p}`` on + consecutive steps (as ``hq_minor_swz``); otherwise the block is folded.""" + + name = "hkv_major" + tag = "hkvmaj" + fastdiv_divisors = ("batch", "gqa", "nqb") + + @staticmethod + def _digits(b, idx, shape: DecodeShape, nqb_scale: int): + """``idx = hkv*(nblk*gqa*B) + blk*(gqa*B) + hql*B + bt`` -> + ``(bt, hq, blk, hkv)``, with ``nblk = NQB // nqb_scale``.""" + bt = shape.mod(idx, "batch") + rem = shape.div(idx, "batch") # hkv*(nblk*gqa) + blk*gqa + hql + hql = shape.mod(rem, "gqa") + r2 = shape.div(rem, "gqa") # hkv*nblk + blk + blk = shape.mod(r2, "nqb", nqb_scale) + hkv = shape.div(r2, "nqb", nqb_scale) + hq = b.add(b.mul(hkv, shape.val("gqa")), hql) + return bt, hq, blk, hkv + + def emit_decode(self, shape, wi): + b, spec = shape.b, shape.spec + NP = spec.num_persistent + + def paired(): + unit, half = emit_pair_unit(b, wi, shape, NP) + bt, hq, p, hkv = self._digits(b, unit, shape, 2) + return emit_pair_block(b, p, half, shape), hq, bt, hkv + + def folded(): + bt, hq, blk, hkv = self._digits(b, wi, shape, 1) + return emit_fold(b, blk, shape), hq, bt, hkv + + return Decoded(*shape.branch(shape.pair_fold(NP), paired, folded)) + + +class PersistBtHkvMinor(PersistDecode): + """``wi = ((blk*gqa + hql)*Hkv + hkv)*B + bt``: batch, then kv head, are the + fastest digits, so each XCD's CTAs stride over few K/V streams; query + blocks are folded under causal masking without a sliding window.""" + + name = "bt_hkv_minor" + tag = "bthkvmin" + fastdiv_divisors = ("batch", "num_query_heads") + + def fastdiv_dividend_bound(self, spec) -> int: + # gqa * t with t < Hq, besides the work index itself. + return max( + super().fastdiv_dividend_bound(spec), + spec.num_queries_per_kv * spec.num_query_heads, + ) + + def emit_decode(self, shape, wi): + b = shape.b + # rem = blk*Hq + t with t = hql*Hkv + hkv < Hq, so the head digits split + # on Hq: hql = (gqa*t) // Hq exactly (gqa*hkv < gqa*Hkv = Hq), and no + # divide by Hkv is needed. + bt = shape.mod(wi, "batch") + rem = shape.div(wi, "batch") + t = shape.mod(rem, "num_query_heads") + blk = shape.div(rem, "num_query_heads") + hql = shape.div(b.mul(t, shape.val("gqa")), "num_query_heads") + hkv = b.sub(t, b.mul(hql, shape.val("num_kv_heads"))) + hq = b.add(b.mul(hkv, shape.val("gqa")), hql) + qb = shape.branch( + shape.balance(), lambda: emit_fold(b, blk, shape), lambda: blk + ) + return Decoded(qb, hq, bt, hkv) + + +class PersistHqMinorSwz(PersistDecode): + """Swizzled Head-first on the persistent grid: XCD ``a`` (``wi % M``, needs + ``M | NP``) owns head band ``a``. Under causal masking without a sliding + window one CTA runs both blocks of a fold pair ``{p, NQB-1-p}`` on consecutive steps, so per-CTA cost + is balanced however NP and NQB divide; the unit's XCD is ``wi % M``. Per XCD, + fastest first: pair, head in band, batch. With ``W <= NP`` the plain decode + is emitted. + + Odd NQB is not supported. It worked when tried: also pair the middle blocks + of two consecutive (head, batch) of the band into one unit (with + ``NQB = 2P+1`` it costs ``2·(P+1) = NQB+1``, as a fold pair), and run the + per-XCD leftover middle (odd heads x batch per band) as single items at the + end of the work index (``wi >= W - M``, XCD ``wi % M``).""" + + name = "hq_minor_swz" + tag = "hqminswz" + fastdiv_divisors = ("nqb", "num_query_heads", "gqa") + + def check(self, spec): + M = spec.chiplet_num_xcds + if spec.num_query_heads % M != 0 or spec.num_persistent % M != 0: + return ( + f"{self.name} needs num_query_heads and num_persistent divisible by {M}" + ) + if balance_causal(spec) and _spec_nqb(spec) % 2 != 0: + return ( + f"{self.name} needs an even number of query blocks under causal " + "masking without a sliding window" + ) + return None + + def emit_decode(self, shape, wi): + b, spec = shape.b, shape.spec + M, NP = spec.chiplet_num_xcds, spec.num_persistent + + def paired(): + unit, half = emit_pair_unit(b, wi, shape, NP) + a = b.mod(unit, b.const_i32(M)) + rest = b.div(unit, b.const_i32(M)) + p = shape.mod(rest, "nqb", 2) + hb = shape.div(rest, "nqb", 2) + return a, hb, emit_pair_block(b, p, half, shape) + + def plain(): + a = b.mod(wi, b.const_i32(M)) + rest = b.div(wi, b.const_i32(M)) + qb = shape.mod(rest, "nqb") # if balancing: W <= NP, one item per CTA + hb = shape.div(rest, "nqb") + return a, hb, qb + + a, hb, qb = shape.branch(shape.pair_fold(NP), paired, plain) + # Head band of Hq/M heads per XCD: hb // (Hq/M) == (M*hb) // Hq. + hq = b.add( + b.mul(a, shape.val("num_query_heads", M)), + shape.mod(hb, "num_query_heads", M), + ) + bt = shape.div(hb, "num_query_heads", M) + return Decoded(qb, hq, bt) + + +def _aligned_causal_error(spec, name: str): + # right_bound (a lookahead) exists only on gfx950. + if not spec.persistent or not spec.causal or getattr(spec, "right_bound", 0): + return f"{name} requires persistent causal attention" + if spec.ragged or spec.varlen or spec.paged: + return f"{name} is validated only for aligned dense attention" + return None + + +class PersistGqaPair(PersistDecode): + """NP = NQB*Hkv*B CTAs; two neighbouring CTAs cover one (qb pair, hkv, bt) + group, each half the local query heads at both complementary blocks, so the + two costs sum to a constant. Not auto-selected: behind the best order on + every shape auto used to pick it for.""" + + name = "gqa_pair" + tag = "gqapair" + + def check(self, spec): + nqb = _spec_nqb(spec) + expected_np = nqb * spec.num_kv_heads * spec.batch + why = _aligned_causal_error(spec, self.name) + if why: + return why + if nqb % 2 != 0 or spec.num_queries_per_kv % 2 != 0: + return f"{self.name} requires even NQB and even GQA ratio" + if spec.num_persistent != expected_np: + return ( + f"{self.name} requires num_persistent == NQB*Hkv*B " + f"({expected_np}), got {spec.num_persistent}" + ) + return None + + def emit_decode(self, shape, wi): + b, spec = shape.b, shape.spec + nqb = _baked_only(shape, self.name) + NP, B, Hkv = spec.num_persistent, spec.batch, spec.num_kv_heads + gqa = spec.num_queries_per_kv + cta = b.mod(wi, b.const_i32(NP)) + phase = b.div(wi, b.const_i32(NP)) + pair_lane = b.mod(cta, b.const_i32(2)) + rem = b.div(cta, b.const_i32(2)) + bt = b.mod(rem, b.const_i32(B)) + rem = b.div(rem, b.const_i32(B)) + hkv = b.mod(rem, b.const_i32(Hkv)) + qb_pair = b.div(rem, b.const_i32(Hkv)) + half_gqa = gqa // 2 + high = b.cmp_ge(phase, b.const_i32(half_gqa)) + phase_half = b.mod(phase, b.const_i32(half_gqa)) + hql = b.add(b.mul(pair_lane, b.const_i32(half_gqa)), phase_half) + hq = b.add(b.mul(hkv, b.const_i32(gqa)), hql) + qb = b.select(high, b.sub(b.const_i32(nqb - 1), qb_pair), qb_pair) + return Decoded(qb, hq, bt, hkv) + + +class PersistGqaPair2Phase(PersistDecode): + """NP = W/2 CTAs; gqa neighbouring CTAs cover all local query heads of one + (qb pair, hkv, bt), and phase 0/1 selects the complementary blocks. Not + auto-selected: behind the best order on every shape auto used to pick it + for.""" + + name = "gqa_pair_2phase" + tag = "gqapair2" + + def check(self, spec): + gqa = spec.num_queries_per_kv + nqb = _spec_nqb(spec) + expected_np = nqb * spec.num_kv_heads * spec.batch * gqa // 2 + why = _aligned_causal_error(spec, self.name) + if why: + return why + if nqb % 2 != 0 or gqa < 2: + return f"{self.name} requires even NQB and GQA ratio >= 2" + if spec.num_persistent != expected_np: + return ( + f"{self.name} requires num_persistent == W/2 " + f"({expected_np}), got {spec.num_persistent}" + ) + return None + + def emit_decode(self, shape, wi): + b, spec = shape.b, shape.spec + nqb = _baked_only(shape, self.name) + NP, B, Hkv = spec.num_persistent, spec.batch, spec.num_kv_heads + gqa = spec.num_queries_per_kv + cta = b.mod(wi, b.const_i32(NP)) + phase = b.div(wi, b.const_i32(NP)) + hql = b.mod(cta, b.const_i32(gqa)) + rem = b.div(cta, b.const_i32(gqa)) + bt = b.mod(rem, b.const_i32(B)) + rem = b.div(rem, b.const_i32(B)) + hkv = b.mod(rem, b.const_i32(Hkv)) + qb_pair = b.div(rem, b.const_i32(Hkv)) + hq = b.add(b.mul(hkv, b.const_i32(gqa)), hql) + qb = b.select( + b.cmp_ne(phase, b.const_i32(0)), + b.sub(b.const_i32(nqb - 1), qb_pair), + qb_pair, + ) + return Decoded(qb, hq, bt, hkv) + + +PERSIST_DECODES = _registry( + PersistQbMajor, + PersistHkvMajor, + PersistBtHkvMinor, + PersistHqMinorSwz, + PersistGqaPair, + PersistGqaPair2Phase, +) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_dense_spec.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_dense_spec.py index d7aa8509263a..24a052442b09 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_dense_spec.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_dense_spec.py @@ -16,6 +16,19 @@ from rocke.core.ir import BF16, F16 from rocke.helpers.spec import kernel_name_join +from kernels.common.attention_dense_decode import ( + NONPERSIST_DECODES, + PERSIST_DECODES, + NonpersistBtHkvMinor, + NonpersistHqMinorSwz, + NonpersistQbMinor, + PersistBtHkvMinor, + PersistHkvMajor, + PersistHqMinorSwz, + PersistQbMajor, + balance_causal, +) + _DTYPE_IR = {"bf16": BF16, "fp16": F16} @@ -28,7 +41,6 @@ ) DEFAULT_DENSE_TILE_GEOMETRY = DENSE_TILE_GEOMETRIES["default"] -_COMMON_PERSIST_DECODES = frozenset({"auto", "qb_major", "hkv_major"}) # Signed 32-bit ceiling for tensor extents. See ``check_dense_spec_preflight`` # check 4 for why the SIGNED bound binds even though the buffer-resource @@ -75,10 +87,23 @@ class AttentionDenseSpec: use_sinks: bool = False # Appended for positional compatibility with existing concrete specs. causal_bottom_right: bool = field(default=False, kw_only=True) - - def supported_persist_decodes(self) -> frozenset[str]: - """Decode values the concrete kernel type can actually emit.""" - return _COMMON_PERSIST_DECODES + # Block order of the non-persistent grid; ignored on the persistent one. + nonpersist_decode: str = field(default="auto", kw_only=True) + # XCDs the hardware round-robins workgroups over. A wrong value keeps results + # correct and only weakens the L2 locality of the XCD-aware block orders. + chiplet_num_xcds: int = field(default=8, kw_only=True) + + @classmethod + def supported_persist_decodes(cls) -> frozenset[str]: + """``persist_decode`` values the concrete kernel type can emit. The shared + spec accepts every known decode; each arch's spec lists its own.""" + return frozenset({"auto", *PERSIST_DECODES}) + + @classmethod + def supported_nonpersist_decodes(cls) -> frozenset[str]: + """``nonpersist_decode`` values the concrete kernel type can emit; every + arch emits every grid decode.""" + return frozenset({"auto", *NONPERSIST_DECODES}) def _supports_noncausal_window(self) -> bool: """Whether this arch's body implements ``causal=False`` + ``sliding_window`` @@ -157,16 +182,37 @@ def __post_init__(self) -> None: f"num_query_heads ({self.num_query_heads}) must be a positive " f"multiple of num_kv_heads ({self.num_kv_heads})" ) - if self.persistent and self.num_persistent <= 0: - raise ValueError( - f"num_persistent must be positive, got {self.num_persistent}" - ) + # Both decode fields always hold a known value; only the grid in use + # checks its decode against the shape. if self.persist_decode not in self.supported_persist_decodes(): raise ValueError( f"persist_decode must be one of " f"{sorted(self.supported_persist_decodes())}, " f"got {self.persist_decode!r}" ) + if self.nonpersist_decode not in self.supported_nonpersist_decodes(): + raise ValueError( + f"nonpersist_decode must be one of " + f"{sorted(self.supported_nonpersist_decodes())}, " + f"got {self.nonpersist_decode!r}" + ) + if self.chiplet_num_xcds <= 0: + raise ValueError( + f"chiplet_num_xcds must be positive, got {self.chiplet_num_xcds}" + ) + if self.persistent: + if self.num_persistent <= 0: + raise ValueError( + f"num_persistent must be positive, got {self.num_persistent}" + ) + if self.persist_decode != "auto": + why = PERSIST_DECODES[self.persist_decode].check(self) + if why: + raise ValueError(why) + elif self.nonpersist_decode != "auto": + why = NONPERSIST_DECODES[self.nonpersist_decode].check(self) + if why: + raise ValueError(why) if self.sliding_window < 0: raise ValueError(f"sliding_window must be >= 0, got {self.sliding_window}") if self.sliding_window > 0: @@ -251,17 +297,70 @@ def dtype_ir(self): def num_queries_per_kv(self) -> int: return self.num_query_heads // self.num_kv_heads + @property + def _aligned_causal(self) -> bool: + """Causal attention on the plain dense layout: where the auto block + orders were measured.""" + moving_diagonal = self.causal_bottom_right and self.seqlen_q != self.seqlen_kv + return ( + balance_causal(self) + and not moving_diagonal + and not self.ragged + and not self.varlen + and not self.paged + ) + @property def resolved_persist_decode(self) -> str: - """Resolve the common auto policy to hkv-major or qb-major.""" + """Resolve the persistent auto policy.""" if self.persist_decode != "auto": return self.persist_decode gqa = self.num_queries_per_kv nqb = (self.seqlen_q + self.block_m - 1) // self.block_m + work = nqb * self.num_query_heads * self.batch + if self._aligned_causal: + if self.interleave: + return PersistQbMajor.name # the only decode interleave applies to + # Swizzled Head-first measured ahead of every other order on causal + # MHA with more than about 8 grid-stride rounds of work. + if ( + gqa == 1 + and PERSIST_DECODES[PersistHqMinorSwz.name].check(self) is None + and work > 8 * self.num_persistent + ): + return PersistHqMinorSwz.name + # Batch is the fastest digit and xcd = wi % num_xcds, so with fewer + # batches than XCDs the second digit still picks the XCD: + # bt_hkv_minor then gives each XCD one kv head. From there on the + # rule below measured ahead. + if self.batch < self.chiplet_num_xcds: + return PersistBtHkvMinor.name per_hkv = gqa * nqb * self.batch if gqa > 1 and per_hkv >= 2 * self.num_persistent: - return "hkv_major" - return "qb_major" + return PersistHkvMajor.name + return PersistQbMajor.name + + @property + def resolved_nonpersist_decode(self) -> str: + """Resolve the non-persistent auto policy. The existing qb_minor grid + stays the best measured order without causal masking.""" + if self.nonpersist_decode != "auto": + return self.nonpersist_decode + if not self._aligned_causal: + return NonpersistQbMinor.name + # Swizzled Head-first measured ahead of every other order, persistent + # or not, on causal MHA with enough query blocks per head and enough + # work per CU; it loses on GQA and on small problems. + nqb = (self.seqlen_q + self.block_m - 1) // self.block_m + hq = self.num_query_heads + if ( + hq == self.num_kv_heads + and hq % self.chiplet_num_xcds == 0 + and nqb >= 16 + and self.batch * hq * nqb >= 8192 + ): + return NonpersistHqMinorSwz.name + return NonpersistBtHkvMinor.name @property def runtime_param_fields(self) -> tuple[str, ...]: @@ -312,7 +411,7 @@ def _algorithm_name_parts(self) -> tuple[str, ...]: return ("lazyrs",) if self.lazy_rescale else () def _persist_decode_name_part(self) -> str: - return "hkvmaj" if self.resolved_persist_decode == "hkv_major" else "" + return PERSIST_DECODES[self.resolved_persist_decode].tag def kernel_name(self) -> str: parts = [ @@ -350,6 +449,12 @@ def kernel_name(self) -> str: parts.append(decode) if self.interleave: parts.append("intl") + else: + decode = NONPERSIST_DECODES[self.resolved_nonpersist_decode].tag + if decode: + parts.append(decode) + if self.chiplet_num_xcds != 8: + parts.append(f"xcd{self.chiplet_num_xcds}") return kernel_name_join(*parts) @@ -370,10 +475,11 @@ class object rather than its name keeps subclass identity stringly-typing it; if a stable on-disk identity is ever needed, that wants ``__qualname__`` or a digest instead. - On the persistent grid ``persist_decode`` is keyed by its resolved order: - ``"auto"`` resolves from the shape, and the order picks the work-decode body, - so once the shape leaves the key the raw ``"auto"`` would let two shapes that - resolve differently share one slot. + The decode the body uses is keyed by its resolved order: ``persist_decode`` on + the persistent grid, ``nonpersist_decode`` on the default grid. ``"auto"`` + resolves from the shape, and the order picks the work-decode body, so once the + shape leaves the key the raw ``"auto"`` would let two shapes that resolve + differently share one slot. """ if not arch: raise ValueError("attention dense cache identity requires an explicit arch") @@ -382,6 +488,8 @@ class object rather than its name keeps subclass identity def _value(name: str): if name == "persist_decode" and spec.persistent: return spec.resolved_persist_decode + if name == "nonpersist_decode" and not spec.persistent: + return spec.resolved_nonpersist_decode return getattr(spec, name) rest = tuple( diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py index 266360cb8555..9aef59557672 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py @@ -190,6 +190,15 @@ from rocke.helpers.attention import mfma_32x32x8_for_dtype # Shared problem/geometry fields live in an architecture-neutral module. +from kernels.common.attention_dense_decode import ( + NONPERSIST_DECODES, + PERSIST_DECODES, + DecodeShape, + PersistBtHkvMinor, + PersistHkvMajor, + PersistHqMinorSwz, + PersistQbMajor, +) from kernels.common.attention_dense_spec import ( AttentionDenseSpec, DENSE_TILE_GEOMETRIES, @@ -456,6 +465,20 @@ class Gfx942AttentionDenseSpec(AttentionDenseSpec): # off emits no new IR, so the shipped kernel and its goldens are unchanged. emit_lse: bool = field(default=False, kw_only=True) + @classmethod + def supported_persist_decodes(cls) -> frozenset[str]: + # The persistent body takes the shape at runtime, so only decodes with a + # runtime-operand form (no gqa_pair*, whose NP equals the baked work count). + return frozenset( + { + "auto", + PersistQbMajor.name, + PersistHkvMajor.name, + PersistBtHkvMinor.name, + PersistHqMinorSwz.name, + } + ) + def __post_init__(self) -> None: super().__post_init__() if type(self.emit_lse) is not bool: @@ -562,18 +585,16 @@ def fastdiv_divisors(self) -> tuple[str, ...]: Each one gets a host-computed magic multiplier and shift (:func:`_fastdiv_kernargs`), so the in-loop divide stays a mul-hi, an add and - a shift instead of the backend's software integer division. ``gqa`` is in - both orders: qb_major derives ``hkv = hq / gqa``, and hkv_major splits its - work index by ``gqa`` into the query head within the kv group. + a shift instead of the backend's software integer division. Each decode + class declares its own (``PersistDecode.fastdiv_divisors``); a decode + without the kv head as a digit lists ``gqa`` for ``hkv = hq / gqa``. A function of the resolved decode order only, never of the shape, so it is fixed by the cache key: the key carries the resolved order on the persistent grid (``attention_dense_cache_key``).""" if not (self.persistent and self.runtime_shape): return () - if self.resolved_persist_decode == "hkv_major": - return ("batch", "gqa", "nqb") - return ("batch", "num_query_heads", "gqa") + return PERSIST_DECODES[self.resolved_persist_decode].fastdiv_divisors @property def runtime_kernarg_fields(self) -> tuple[str, ...]: @@ -1416,9 +1437,19 @@ def _build_attention_dense_single_buffer( gqa_p = b.div(hq_p, hkv_p) stride_q_p = b.mul(hq_p, b.const_i32(D)) stride_k_p = b.mul(hkv_p, b.const_i32(D)) + # The runtime shape as the work decodes read it (DecodeShape). + shape_values = { + "batch": batch_p, + "seqlen_q": seqlen_q_p, + "seqlen_kv": seqlen_kv_p, + "num_query_heads": hq_p, + "num_kv_heads": hkv_p, + "gqa": gqa_p, + } else: batch_p = seqlen_q_p = seqlen_kv_p = None hq_p = hkv_p = gqa_p = stride_q_p = stride_k_p = None + shape_values = None # Per-query-head sink logits [Hq], after every runtime kernarg field (the gfx950 # ABI order). Declared only when use_sinks, so the non-sink ABI and IR are # unchanged. @@ -2193,30 +2224,18 @@ def tile_body(j, carry, masked): # constant, so this per-CTA divide strength-reduces. nqb_p = b.div(seqlen_q_p, b.const_i32(BLOCK_M)) w_p = b.mul(b.mul(nqb_p, hq_p), batch_p) - divisor = { - "batch": batch_p, - "num_query_heads": hq_p, - "gqa": gqa_p, - "nqb": nqb_p, - } def _div(n, d): # (umulhi(n, magic) + n) >> shift, exact for n < 2**30 (the add is nsw); - # attention_dense_runtime_args holds W under that bound. + # attention_dense_runtime_args holds every decode dividend under that + # bound. return b.lshr( b.add(b.umul_hi_i32(n, rt[f"{d}_magic"]), n), rt[f"{d}_shift"] ) - def _divmod(n, d): - q = _div(n, d) - return q, b.sub(n, b.mul(q, divisor[d])) - - if spec.resolved_persist_decode == "hkv_major": - half_p = b.lshr(nqb_p, b.const_i32(1)) - fold_p = b.add(b.sub(nqb_p, b.const_i32(1)), half_p) # NQB-1+half - elif spec.interleave and causal: - # NQB == 1 needs no gate: qb0 is then 0 and the flip yields 0 too. - nqb_m1_p = b.sub(nqb_p, b.const_i32(1)) + work_shape = DecodeShape( + b, spec, {**shape_values, "nqb": nqb_p, "w": w_p}, fastdiv=_div + ) cta_id = b.block_id_x() outer = b.scf_for(cta_id, w_p, b.const_i32(NP), iv_name="wi") with outer as wi: @@ -2228,44 +2247,21 @@ def _divmod(n, d): # barrier suffices because there is no pending lgkm to sink here. b.s_waitcnt(vmcnt=0) b.s_barrier_bare() - if spec.resolved_persist_decode == "hkv_major": - # hkv-MAJOR + causal-balanced decode (gfx950 §): - # wi = hkv*(NQB*gqa*B) + blk*(gqa*B) + hql*B + bt - # hkv in the MSB keeps each grid-stride phase within ~1 kv-head so the - # shared GQA K/V stays L2-resident across its gqa query heads; blk is - # folded so a CTA striding both halves of a kv-head does qb=X and - # qb=NQB-1-X (constant causal cost) -- qb_major's balance + L2 win. - rem, bt_v = _divmod(wi, "batch") - r2, hql = _divmod(rem, "gqa") - hkv_v, blk = _divmod(r2, "nqb") - hq_v = b.add(b.mul(hkv_v, gqa_p), hql) - qb_hi = b.sub(fold_p, blk) # NQB-1-(blk-half) - qb_v = b.select(b.cmp_lt(blk, half_p), blk, qb_hi) - elif spec.resolved_persist_decode == "qb_major": - # qb-MAJOR decode: wi = qb*(Hq*B) + hq*B + bt. Putting qb (the - # triangular causal-cost index) in the MSB spreads cheap+expensive - # query blocks across each CTA under grid-stride. Optional interleave - # reverses qb for ODD hq to further balance the causal tail. Keying on - # hq (not rem = qb0*Hq + hq) keeps the map a bijection: each hq's qb - # range is either identity or reversed. For even Hq the two agree; - # for odd Hq rem parity aliases two qb0 onto one qb and skips another. - rem, bt_v = _divmod(wi, "batch") - qb0, hq_v = _divmod(rem, "num_query_heads") - hkv_v = _div(hq_v, "gqa") - if spec.interleave and causal: - odd = b.cmp_eq(b.mod(hq_v, b.const_i32(2)), b.const_i32(1)) - qb_v = b.select(odd, b.sub(nqb_m1_p, qb0), qb0) - else: - qb_v = qb0 - else: - raise ValueError( - "gfx942 attention_dense: persist_decode=" - f"{spec.resolved_persist_decode!r} is not implemented " - "by this builder" - ) - _run_work_item(qb_v, hq_v, bt_v, hkv=hkv_v) + d = PERSIST_DECODES[spec.resolved_persist_decode].emit_decode( + work_shape, wi + ) + # The kv head from the decode where it is a digit, else hq // gqa through + # the gqa fastdiv pair: no runtime division in the grid-stride loop. + hkv_v = d.hkv if d.hkv is not None else work_shape.div(d.hq, "gqa") + _run_work_item(d.qb, d.hq, d.bt, hkv=hkv_v) else: - _run_work_item(b.block_id_x(), b.block_id_y(), b.block_id_z()) + d = NONPERSIST_DECODES[spec.resolved_nonpersist_decode].emit_decode( + DecodeShape(b, spec, shape_values), + b.block_id_x(), + b.block_id_y(), + b.block_id_z(), + ) + _run_work_item(d.qb, d.hq, d.bt) b.ret() return b.kernel @@ -2275,7 +2271,8 @@ def _divmod(n, d): def attention_dense_grid(spec: AttentionDenseSpec) -> tuple[int, int, int]: """Launch grid: persistent = 1-D grid of ``num_persistent`` CTAs; default = - one CTA per (query-block, query-head, batch). + one CTA per (query-block, query-head, batch), laid out by the nonpersist + decode. Sized from ``spec.block_m`` (``_BLOCK_M`` at the default) so the grid and the body's query tiling cannot disagree -- a mismatch writes some rows twice and @@ -2285,11 +2282,7 @@ def attention_dense_grid(spec: AttentionDenseSpec) -> tuple[int, int, int]: if spec.persistent: return (spec.num_persistent, 1, 1) spec = _as_gfx942_spec(spec) - # ceil kept for parity with the gfx950 helper; on gfx942 it is always exact, - # because ragged is rejected and supports_attention_dense then enforces - # seqlen_q % block_m == 0. - nqb = (spec.seqlen_q + spec.block_m - 1) // spec.block_m - return (nqb, spec.num_query_heads, spec.batch) + return NONPERSIST_DECODES[spec.resolved_nonpersist_decode].grid(spec) def attention_dense_block(spec: AttentionDenseSpec) -> tuple[int, int, int]: @@ -2349,18 +2342,20 @@ def _fastdiv_kernargs(spec: Gfx942AttentionDenseSpec) -> dict[str, int]: is a uint32 passed as its two's-complement i32 bit pattern, which is what the kernel's unsigned mul-hi reads. - Raises when the work-item count reaches the fast-division dividend limit: every - decode dividend is below it, and past it the quotient is wrong with no fault.""" + Raises when a decode dividend can reach the fast-division dividend limit (the + work-item count, or the decode's own bound): past it the quotient is wrong with + no fault.""" if not spec.fastdiv_divisors: return {} from rocke.helpers.magic_division import calculate_magic_numbers nqb = spec.seqlen_q // spec.block_m - work = nqb * spec.num_query_heads * spec.batch - if work >= _FASTDIV_DIVIDEND_LIMIT: + bound = PERSIST_DECODES[spec.resolved_persist_decode].fastdiv_dividend_bound(spec) + if bound >= _FASTDIV_DIVIDEND_LIMIT: raise ValueError( - f"gfx942 attention_dense: persistent work-item count {work} " - f"(nqb*Hq*B) must be below 2**30 for the work decode's fast division" + f"gfx942 attention_dense: persistent decode dividend bound {bound} " + f"(the work-item count nqb*Hq*B, or the decode's own) must be below " + f"2**30 for the work decode's fast division" ) value = { "batch": spec.batch, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_dense.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_dense.py index 4477ef5065ea..949f503661ea 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_dense.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_dense.py @@ -97,6 +97,17 @@ attention_dense_cache_key, check_dense_spec_preflight, ) +from kernels.common.attention_dense_decode import ( + NONPERSIST_DECODES, + PERSIST_DECODES, + DecodeShape, + PersistBtHkvMinor, + PersistGqaPair, + PersistGqaPair2Phase, + PersistHkvMajor, + PersistHqMinorSwz, + PersistQbMajor, +) from kernels.gfx950.attention_tiled_2d import _mfma_32x32_c_row, _mfma_32x32_c_col from kernels.common._lse_store import ( dense_row_has_key, @@ -196,6 +207,11 @@ def _supports_noncausal_cross_length_ragged(self) -> bool: def _band_name_parts(self) -> tuple[str, ...]: return (f"rb{self.right_bound}",) if self.right_bound > 0 else () + @property + def _aligned_causal(self) -> bool: + # The auto block orders were measured without a lookahead. + return super()._aligned_causal and self.right_bound == 0 + # Performance-only codegen knobs. Every legal value computes the same # attention output; defaults reproduce the shipped kernel byte-for-byte. # ``None`` resolves through the shipped per-path policy (``resolved_*``). @@ -266,11 +282,19 @@ def resolved_pv_loop_order(self) -> str: return "k_major" if self.wide_lds_dma else "d_major" return self.pv_loop_order - def supported_persist_decodes(self) -> frozenset[str]: - return super().supported_persist_decodes() | { - "gqa_pair", - "gqa_pair_2phase", - } + @classmethod + def supported_persist_decodes(cls) -> frozenset[str]: + return frozenset( + { + "auto", + PersistQbMajor.name, + PersistHkvMajor.name, + PersistBtHkvMinor.name, + PersistHqMinorSwz.name, + PersistGqaPair.name, + PersistGqaPair2Phase.name, + } + ) def __post_init__(self) -> None: super().__post_init__() @@ -335,41 +359,6 @@ def __post_init__(self) -> None: "wide_lds_dma requires K/V slab padding of 8/32 elements" ) - if self.persist_decode == "gqa_pair": - gqa = self.num_queries_per_kv - nqb = (self.seqlen_q + self.block_m - 1) // self.block_m - expected_np = nqb * self.num_kv_heads * self.batch - if not self.persistent or not self.causal or self.right_bound: - raise ValueError("gqa_pair requires persistent causal attention") - if self.ragged or self.varlen or self.paged: - raise ValueError( - "gqa_pair is validated only for aligned dense attention" - ) - if nqb % 2 or gqa % 2: - raise ValueError("gqa_pair requires even NQB and even GQA ratio") - if self.num_persistent != expected_np: - raise ValueError( - "gqa_pair requires num_persistent == NQB*Hkv*B " - f"({expected_np}), got {self.num_persistent}" - ) - if self.persist_decode == "gqa_pair_2phase": - gqa = self.num_queries_per_kv - nqb = (self.seqlen_q + self.block_m - 1) // self.block_m - expected_np = nqb * self.num_kv_heads * self.batch * gqa // 2 - if not self.persistent or not self.causal or self.right_bound: - raise ValueError("gqa_pair_2phase requires persistent causal attention") - if self.ragged or self.varlen or self.paged: - raise ValueError( - "gqa_pair_2phase is validated only for aligned dense attention" - ) - if nqb % 2 or gqa < 2: - raise ValueError("gqa_pair_2phase requires even NQB and GQA ratio >= 2") - if self.num_persistent != expected_np: - raise ValueError( - "gqa_pair_2phase requires num_persistent == W/2 " - f"({expected_np}), got {self.num_persistent}" - ) - def _validate_codegen_knobs(self) -> None: if self.lds_k_row_pad < 0 or self.lds_k_row_pad % 8 != 0: raise ValueError( @@ -434,31 +423,6 @@ def _validate_codegen_knobs(self) -> None: "pv_sched_group_template when enabling it" ) - @property - def resolved_persist_decode(self) -> str: - if self.persist_decode != "auto": - return self.persist_decode - gqa = self.num_queries_per_kv - nqb = (self.seqlen_q + self.block_m - 1) // self.block_m - aligned_causal = ( - self.persistent - and self.causal - and not self.ragged - and not self.varlen - and not self.paged - and self.sliding_window == 0 - and self.right_bound == 0 - ) - if aligned_causal and nqb % 2 == 0 and gqa % 2 == 0: - pair_np = nqb * self.num_kv_heads * self.batch - if self.num_persistent == pair_np: - return "gqa_pair" - if aligned_causal and nqb % 2 == 0 and gqa >= 2: - two_phase_np = nqb * self.num_kv_heads * self.batch * gqa // 2 - if self.num_persistent == two_phase_np: - return "gqa_pair_2phase" - return super().resolved_persist_decode - @property def runtime_shape(self) -> bool: """Whether the body reads the problem shape *only* from its kernel params, @@ -542,13 +506,6 @@ def _codegen_name_part(self) -> str: return "" return "cg" + hashlib.sha256(repr(changed).encode()).hexdigest()[:8] - def _persist_decode_name_part(self) -> str: - return { - "hkv_major": "hkvmaj", - "gqa_pair": "gqapair", - "gqa_pair_2phase": "gqapair2", - }.get(self.resolved_persist_decode, "") - _CODEGEN_KNOB_DEFAULTS = MappingProxyType( { @@ -926,9 +883,11 @@ def build_attention_dense( rcp_ln2 = b.const_f32(LOG2E) one_f = b.const_f32(1.0) - qb = b.block_id_x() - hq = b.block_id_y() - bt = b.block_id_z() + # Heads are baked here; NQB comes from the seqlen_q param. + grid_shape = DecodeShape(b, spec, {"seqlen_q": seqlen_q_p}) + qb, hq, bt, _ = NONPERSIST_DECODES[spec.resolved_nonpersist_decode].emit_decode( + grid_shape, b.block_id_x(), b.block_id_y(), b.block_id_z() + ) hkv = b.div(hq, b.const_i32(gqa)) q_tok0 = b.add(b.mul(qb, b.const_i32(BLOCK_M)), b.mul(wave, b.const_i32(32))) @@ -1756,7 +1715,6 @@ def _build_attention_dense_persistent(spec: AttentionDenseSpec) -> KernelDef: NBUF = spec.lds_num_buffers PAD = spec.lds_k_row_pad NP = spec.num_persistent - INTERLEAVE = spec.interleave K_STEPS = D // 16 D_TILES = D // 32 @@ -1925,85 +1883,11 @@ def _build_attention_dense_persistent(spec: AttentionDenseSpec) -> KernelDef: b.s_waitcnt(vmcnt=0) b.s_barrier_bare() - if spec.resolved_persist_decode == "gqa_pair_2phase": - # NP=W/2 CTAs. gqa neighboring CTAs cover all local query heads - # for one (qb_pair,hkv,bt); phase 0/1 selects complementary qbs. - cta = b.mod(wi, b.const_i32(NP)) - phase = b.div(wi, b.const_i32(NP)) - hql = b.mod(cta, b.const_i32(gqa)) - rem = b.div(cta, b.const_i32(gqa)) - bt = b.mod(rem, b.const_i32(B)) - rem = b.div(rem, b.const_i32(B)) - hkv = b.mod(rem, b.const_i32(Hkv)) - qb_pair = b.div(rem, b.const_i32(Hkv)) - hq = b.add(b.mul(hkv, b.const_i32(gqa)), hql) - qb = b.select( - b.cmp_ne(phase, b.const_i32(0)), - b.sub(b.const_i32(NQB - 1), qb_pair), - qb_pair, - ) - elif spec.resolved_persist_decode == "gqa_pair": - # NP=NQB*Hkv*B CTAs. Two neighboring CTAs cover one - # (qb_pair,hkv,bt) group; each handles half the local query heads - # at both complementary qbs. The low/high costs sum to a constant. - cta = b.mod(wi, b.const_i32(NP)) - phase = b.div(wi, b.const_i32(NP)) - pair_lane = b.mod(cta, b.const_i32(2)) - rem = b.div(cta, b.const_i32(2)) - bt = b.mod(rem, b.const_i32(B)) - rem = b.div(rem, b.const_i32(B)) - hkv = b.mod(rem, b.const_i32(Hkv)) - qb_pair = b.div(rem, b.const_i32(Hkv)) - half_gqa = gqa // 2 - high = b.cmp_ge(phase, b.const_i32(half_gqa)) - phase_half = b.mod(phase, b.const_i32(half_gqa)) - hql = b.add(b.mul(pair_lane, b.const_i32(half_gqa)), phase_half) - hq = b.add(b.mul(hkv, b.const_i32(gqa)), hql) - qb = b.select( - high, - b.sub(b.const_i32(NQB - 1), qb_pair), - qb_pair, - ) - elif spec.resolved_persist_decode == "hkv_major": - # hkv-MAJOR + causal-balanced decode: - # wi = hkv*(NQB*gqa*B) + blk*(gqa*B) + hql*B + bt - # * hkv in the MSB -> each grid-stride phase (NP consecutive wi) stays - # within ~1 kv-head, so the shared GQA K/V is L2-resident across its - # gqa query heads (recovers the non-persistent grid's locality: - # measured L2 hit 57% -> ~90%+ vs qb_major). - # * `blk` (0..NQB-1) is folded to a query-block index that PAIRS a low - # and a high qb per CTA: blk qb=blk (cheap), blk>=half -> - # qb=NQB-1-(blk-half) (expensive), so a CTA that grid-strides over - # both halves of a kv-head does qb=X and qb=NQB-1-X (constant causal - # cost) -> keeps qb_major's load balance while gaining L2 locality. - half = NQB // 2 - bt = b.mod(wi, b.const_i32(B)) - rem = b.div(wi, b.const_i32(B)) # hkv*(NQB*gqa) + blk*gqa + hql - hql = b.mod(rem, b.const_i32(gqa)) - r2 = b.div(rem, b.const_i32(gqa)) # hkv*NQB + blk - blk = b.mod(r2, b.const_i32(NQB)) - hkv = b.div(r2, b.const_i32(NQB)) - hq = b.add(b.mul(hkv, b.const_i32(gqa)), hql) - # qb = blk 1: - # Reverse qb for ODD hq. Keying on rem (= qb0*Hq + hq) instead is - # not a bijection for odd Hq: two qb0 alias onto one qb. - odd = b.cmp_eq(b.mod(hq, b.const_i32(2)), b.const_i32(1)) - qb = b.select(odd, b.sub(b.const_i32(NQB - 1), qb0), qb0) - else: - qb = qb0 + # No shape params on this body: the decode reads the baked shape. + qb, hq, bt, hkv = PERSIST_DECODES[spec.resolved_persist_decode].emit_decode( + DecodeShape(b, spec), wi + ) + if hkv is None: hkv = b.div(hq, b.const_i32(gqa)) q_tok0 = b.add(b.mul(qb, b.const_i32(BLOCK_M)), b.mul(wave, b.const_i32(32))) @@ -2646,8 +2530,9 @@ def _clamp(x, lo, hi): o_acc = do_pv(o_acc, p_prev, last_vbuf) # Epilogue: recompute (bt, hq) from the live loop IV so they need not cross - # the KV loop (keeps the loop-carried live set minimal -> 0 spill). Must - # mirror the work-item decode used at the top of the loop. + # the KV loop (keeps the loop-carried live set minimal -> 0 spill). Same + # decode as the top of the loop; its unused outputs stay in the IR as dead + # ops, which LLVM drops, so the code object does not change. if spec.emit_lse: positive_l = b.fcmp("ogt", l_i, b.const_f32(0.0)) rcp_l = b.rcp(b.select(positive_l, l_i, b.const_f32(1.0))) @@ -2660,37 +2545,9 @@ def _clamp(x, lo, hi): rcp_l, b.const_f32(0.0), ) - if spec.resolved_persist_decode == "gqa_pair_2phase": - cta_e = b.mod(wi, b.const_i32(NP)) - hql_e = b.mod(cta_e, b.const_i32(gqa)) - rem_e = b.div(cta_e, b.const_i32(gqa)) - bt_e = b.mod(rem_e, b.const_i32(B)) - rem_e = b.div(rem_e, b.const_i32(B)) - hkv_e = b.mod(rem_e, b.const_i32(Hkv)) - hq_e = b.add(b.mul(hkv_e, b.const_i32(gqa)), hql_e) - elif spec.resolved_persist_decode == "gqa_pair": - cta_e = b.mod(wi, b.const_i32(NP)) - phase_e = b.div(wi, b.const_i32(NP)) - pair_lane_e = b.mod(cta_e, b.const_i32(2)) - rem_e = b.div(cta_e, b.const_i32(2)) - bt_e = b.mod(rem_e, b.const_i32(B)) - rem_e = b.div(rem_e, b.const_i32(B)) - hkv_e = b.mod(rem_e, b.const_i32(Hkv)) - phase_half_e = b.mod(phase_e, b.const_i32(gqa // 2)) - hql_e = b.add( - b.mul(pair_lane_e, b.const_i32(gqa // 2)), - phase_half_e, - ) - hq_e = b.add(b.mul(hkv_e, b.const_i32(gqa)), hql_e) - elif spec.resolved_persist_decode == "hkv_major": - bt_e = b.mod(wi, b.const_i32(B)) - rem_e = b.div(wi, b.const_i32(B)) - hql_e = b.mod(rem_e, b.const_i32(gqa)) - hkv_e = b.div(b.div(rem_e, b.const_i32(gqa)), b.const_i32(NQB)) - hq_e = b.add(b.mul(hkv_e, b.const_i32(gqa)), hql_e) - else: - bt_e = b.mod(wi, b.const_i32(B)) - hq_e = b.mod(b.div(wi, b.const_i32(B)), b.const_i32(Hq)) + _, hq_e, bt_e, _ = PERSIST_DECODES[spec.resolved_persist_decode].emit_decode( + DecodeShape(b, spec), wi + ) o_base = b.add( b.mul(b.mul(bt_e, b.const_i32(Sq)), b.const_i32(stride_q_tok)), b.mul(hq_e, b.const_i32(D)), @@ -2759,13 +2616,11 @@ def _clamp(x, lo, hi): def attention_dense_grid(spec: AttentionDenseSpec) -> Tuple[int, int, int]: """Launch grid for ``spec``. Persistent = 1-D grid of ``num_persistent`` CTAs; - default = one CTA per (query-block, query-head, batch).""" + default = one CTA per (query-block, query-head, batch), laid out by the + nonpersist decode.""" if spec.persistent: return (spec.num_persistent, 1, 1) - nqb = ( - spec.seqlen_q + spec.block_m - 1 - ) // spec.block_m # ceil: ragged partial block - return (nqb, spec.num_query_heads, spec.batch) + return NONPERSIST_DECODES[spec.resolved_nonpersist_decode].grid(spec) def attention_dense_block(spec: AttentionDenseSpec) -> Tuple[int, int, int]: diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_attention.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_attention.py index 08ae61e1a983..6e04adbb91d3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_attention.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_attention.py @@ -64,6 +64,7 @@ }, "attention_gfx942_dense": {"causal", "lse", "sinks", "sliding_window"}, "attention_gfx950_dense_grid": set(_GFX950_DENSE_GRID_FEATURES), + "attention_gfx950_dense_qbminor": set(_GFX950_DENSE_GRID_FEATURES), "attention_gfx950_dense_persist": set(_GFX950_DENSE_FEATURES), "attention_gfx950_dense_persist_widedma": set(_GFX950_DENSE_FEATURES), "attention_d256_decode": {"causal", "causal_bottom_right"}, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_gfx942_dense_wiring.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_gfx942_dense_wiring.py index 4ff01fda55fb..6c96abe02024 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_gfx942_dense_wiring.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_gfx942_dense_wiring.py @@ -9,8 +9,8 @@ - selection returns an ``AttentionTuningSpec`` (``path="dense"``) whose ``tuning_id`` replays; ``auto`` is the default spec - routing on gfx942, and rejection of every out-of-scope request - - the default spec turns the persistent grid on once there is enough work; - ``persistent`` and ``waves_per_eu`` are knobs of the candidate + - the default spec turns the persistent grid on at D128; ``persistent`` and + ``waves_per_eu`` are knobs of the candidate - both gfx942 dense grids read ``batch`` / ``seqlen_q`` / ``seqlen_kv`` and both head counts as runtime kernel params, so those fields drop out of ``kernel_name()`` and the cache key and the dispatched signature includes them. @@ -250,7 +250,7 @@ def test_concrete_support_rejects_shared_bottom_right_spec(self): Gfx942AttentionDenseSpec(**common) def test_equal_length_bottom_right_preserves_persistent_policy(self): - common = dict(seqlen_q=8192, seqlen_k=8192) + common = dict(batch=16, seqlen_q=2048, seqlen_k=2048) mask_pairs = ( (AttentionMaskType.TOP_LEFT_CAUSAL, 2), (1, AttentionMaskType.BOTTOM_RIGHT_CAUSAL), @@ -269,12 +269,13 @@ def test_equal_length_bottom_right_preserves_persistent_policy(self): class TestGfx942DensePersistent(unittest.TestCase): - def test_default_spec_turns_persistent_on_for_large_sq(self): - """The default spec turns the persistent grid-stride variant on once - there is enough work to fill the grid -- the large-Sq prefill regime.""" + def test_default_spec_turns_persistent_on_for_d128(self): + """The default spec turns the persistent grid-stride variant on at D128, + whatever the Sq and batch, and leaves it off at D64.""" with _Gfx942Arch(): + self.assertTrue(_spec(_req()).persistent) self.assertTrue(_spec(_req(seqlen_q=8192, seqlen_k=8192)).persistent) - self.assertFalse(_spec(_req(nhead_q=8, nhead_k=8)).persistent) + self.assertFalse(_spec(_req(hdim_q=64, hdim_v=64)).persistent) def test_persistent_is_a_knob_of_the_candidate(self): """Persistence is reached through the knob space (and its tuning ids), @@ -284,6 +285,39 @@ def test_persistent_is_a_knob_of_the_candidate(self): self.assertTrue(_spec(req, persistent=True).persistent) self.assertFalse(_spec(_req(), persistent=False).persistent) + def test_large_causal_mha_uses_persistent_swizzled_head_first(self): + with _Gfx942Arch(): + spec = _spec( + _req(batch=16, nhead_q=64, nhead_k=64, seqlen_q=4096, seqlen_k=4096) + ) + self.assertTrue(spec.persistent) + self.assertEqual(spec.resolved_persist_decode, "hq_minor_swz") + + def test_nonpersist_decode_is_a_knob_of_the_grid(self): + """The non-persistent block order is a knob: kept on the grid unless it + restates auto, dropped on the persistent body, which never reads it.""" + with _Gfx942Arch(): + req = _req() + grid = _spec(req, persistent=False, nonpersist_decode="hq_minor_swz") + self.assertFalse(grid.persistent) + self.assertEqual(grid.resolved_nonpersist_decode, "hq_minor_swz") + self.assertIn("nphqminswz", grid.kernel_name()) + auto = tuning_spec_with_knobs(req, _SPEC_ID, {"persistent": False}) + restated = tuning_spec_with_knobs( + req, + _SPEC_ID, + { + "persistent": False, + "nonpersist_decode": auto.kernel_spec.resolved_nonpersist_decode, + }, + ) + self.assertEqual(restated, auto) + persistent = tuning_spec_with_knobs( + req, _SPEC_ID, {"persistent": True, "nonpersist_decode": "hq_minor_swz"} + ) + self.assertNotIn("nonpersist_decode", dict(persistent.knobs)) + self.assertEqual(persistent.kernel_spec.nonpersist_decode, "auto") + class TestGfx942DenseWavesPerEu(unittest.TestCase): def test_shipped_policy_and_knob_override(self): @@ -343,7 +377,7 @@ def test_kernel_name_follows_the_runtime_shape_contract(self): with _Gfx942Arch(): for persistent, tail in ( (False, []), - # Every shape here resolves 'auto' to qb_major (asserted below). + # Every shape here resolves 'auto' to bt_hkv_minor (asserted below). ( True, [ @@ -351,8 +385,6 @@ def test_kernel_name_follows_the_runtime_shape_contract(self): "batch_shift", "num_query_heads_magic", "num_query_heads_shift", - "gqa_magic", - "gqa_shift", ], ), ): @@ -373,7 +405,8 @@ def test_kernel_name_follows_the_runtime_shape_contract(self): self.assertNotRegex(name, r"_s[qk]\d+") if persistent: self.assertEqual( - {s.resolved_persist_decode for s in specs}, {"qb_major"} + {s.resolved_persist_decode for s in specs}, + {"bt_hkv_minor"}, ) self.assertIn("persist304", name) names = [p["name"] for p in attention_dense_signature(specs[0])] @@ -441,7 +474,9 @@ class TestGfx942DensePersistentRuntimeShape(unittest.TestCase): # Per resolved decode order, the divisors in kernarg order. _DIVISORS = { + "bt_hkv_minor": ("batch", "num_query_heads"), "hkv_major": ("batch", "gqa", "nqb"), + "hq_minor_swz": ("nqb", "num_query_heads", "gqa"), "qb_major": ("batch", "num_query_heads", "gqa"), } @@ -526,17 +561,23 @@ def test_auto_decode_is_keyed_by_its_resolved_order(self): lands in the same slot as the 'auto' that resolves to it.""" common = dict(seqlen_q=8192, seqlen_k=8192) groups = { - # gqa*nqb*B < 2*NP (608), or gqa == 1 - "qb_major": [ + # GQA, batch below the XCD count + "bt_hkv_minor": [ dict(batch=1, nhead_q=128, nhead_k=8), dict(batch=1, nhead_q=40, nhead_k=8), - dict(batch=4, nhead_q=32, nhead_k=32), + dict(batch=4, nhead_q=32, nhead_k=8), ], - # gqa > 1 and gqa*nqb*B >= 2*NP + # gqa > 1, batch >= the XCD count and gqa*nqb*B >= 2*NP (608) "hkv_major": [ - dict(batch=2, nhead_q=128, nhead_k=8), - dict(batch=4, nhead_q=128, nhead_k=8), - dict(batch=4, nhead_q=40, nhead_k=8), + dict(batch=8, nhead_q=128, nhead_k=8), + dict(batch=8, nhead_q=40, nhead_k=8), + dict(batch=8, nhead_q=32, nhead_k=8), + ], + # MHA with more than 8 grid-stride rounds of work + "hq_minor_swz": [ + dict(batch=4, nhead_q=32, nhead_k=32), + dict(batch=2, nhead_q=64, nhead_k=64), + dict(batch=8, nhead_q=32, nhead_k=32), ], } keys, names = {}, {} @@ -563,10 +604,16 @@ def test_auto_decode_is_keyed_by_its_resolved_order(self): keys[decode], ) self.assertEqual(explicit.kernel_name(), names[decode]) - self.assertNotEqual(keys["qb_major"], keys["hkv_major"]) - self.assertNotEqual(names["qb_major"], names["hkv_major"]) - self.assertIn("hkvmaj", names["hkv_major"]) - self.assertNotIn("hkvmaj", names["qb_major"]) + self.assertEqual(len(set(keys.values())), len(groups)) + self.assertEqual(len(set(names.values())), len(groups)) + for decode, tag in ( + ("bt_hkv_minor", "bthkvmin"), + ("hkv_major", "hkvmaj"), + ("hq_minor_swz", "hqminswz"), + ): + self.assertEqual( + [n for n in names.values() if tag in n], [names[decode]] + ) def test_work_item_count_at_the_fastdiv_limit_raises(self): """The decode's fast division is exact only below 2**30; past it the diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_gfx950_dense_wiring.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_gfx950_dense_wiring.py index ebfcccab6b0b..e3c344179dc2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_gfx950_dense_wiring.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_gfx950_dense_wiring.py @@ -34,6 +34,7 @@ iter_registered_attention_combos, tuning_spec_with_knobs, ) +from dispatch.attention.gfx950_dense import GFX950_DENSE_VARIANTS from kernels.common.attention_dense_spec import ( DENSE_TILE_GEOMETRIES, attention_dense_cache_key, @@ -52,9 +53,11 @@ _GRID = "gfx950_dense_grid" _PERSIST = "gfx950_dense_persist" _WIDE = "gfx950_dense_persist_widedma" +_QBMINOR = "gfx950_dense_qbminor" _GRID_NAME = f"attention_{_GRID}" _PERSIST_NAME = f"attention_{_PERSIST}" _WIDE_NAME = f"attention_{_WIDE}" +_QBMINOR_NAME = f"attention_{_QBMINOR}" _DENSE_ALGORITHMS = (DENSE_GRID_ALGORITHM, DENSE_PERSIST_ALGORITHM) @@ -136,6 +139,29 @@ def test_unpinned_requests_route_to_unified(self): with self.subTest(**kw), self.assertRaises(ValueError): dispatch_attention(_gfx950_dense_req(**kw)) + def test_every_variant_needs_both_pins(self): + """Every variant is opt-in: unpinned or half-pinned requests are refused; + pinning its algorithm and spec_id routes to it.""" + req = _gfx950_dense_req(hdim_q=128, hdim_v=128, nhead_q=32, nhead_k=8) + by_name = {c.name: c for c in attention_candidates()} + for variant in GFX950_DENSE_VARIANTS: + candidate = by_name[variant.candidate_name] + with self.subTest(spec_id=variant.spec_id): + for kw in ( + {}, + dict(algorithm=variant.algorithm), + dict(spec_id=variant.spec_id), + ): + ok, why = candidate.admits(replace(req, **kw)) + self.assertFalse(ok, kw) + self.assertIn("opt-in", why) + pinned = replace( + req, algorithm=variant.algorithm, spec_id=variant.spec_id + ) + self.assertEqual( + dispatch_attention(pinned).candidate.name, variant.candidate_name + ) + def test_spec_id_under_the_other_body_is_refused(self): """The algorithm names the body: the gfx942 ``attention_dense`` name and a spec_id of the other gfx950 body select nothing.""" @@ -213,74 +239,80 @@ def test_production_and_full_sweeps_expand_wpe(self): class TestDenseGqaPairWiring(unittest.TestCase): - """Invariant-compatible shapes select a balanced GQA-local decode.""" + """Invariant-compatible shapes accept an explicit balanced GQA-local decode; + auto picks the L2-local one.""" - def test_exact_llama3_8b_prefill_auto_selects_gqa_pair(self): - spec = _spec(_gfx950_dense_req(**_LLAMA_8K), _WIDE) + def test_exact_llama3_8b_prefill_explicit_gqa_pair(self): + spec = _spec(_gfx950_dense_req(**_LLAMA_8K), _WIDE, persist_decode="gqa_pair") self.assertEqual(spec.resolved_persist_decode, "gqa_pair") self.assertIn("gqapair", spec.kernel_name()) self.assertTrue(spec.wide_lds_dma) self.assertIn("wdma", spec.kernel_name()) self.assertEqual(spec.num_persistent, 256) - def test_explicit_gqa_pair_knob_canonicalizes_to_the_default(self): + def test_explicit_auto_decode_knob_canonicalizes_to_the_default(self): """Pinning the decode auto already picks is the default kernel: the knob is dropped, so the id is the default's and no duplicate is minted.""" req = _gfx950_dense_req(**_LLAMA_8K) - pinned = tuning_spec_with_knobs(req, _WIDE, {"persist_decode": "gqa_pair"}) - self.assertEqual(pinned, attention_tuning_spec(req, _WIDE)) + default = attention_tuning_spec(req, _WIDE) + self.assertEqual(default.kernel_spec.resolved_persist_decode, "bt_hkv_minor") + pinned = tuning_spec_with_knobs(req, _WIDE, {"persist_decode": "bt_hkv_minor"}) + self.assertEqual(pinned, default) # Wide DMA records its problem-dependent default tile. self.assertEqual(pinned.knobs, (("block_m", 256),)) - self.assertEqual(pinned.kernel_spec.resolved_persist_decode, "gqa_pair") - self.assertIn("gqapair", pinned.kernel_name()) def test_invalid_dense_decode_is_rejected(self): with self.assertRaisesRegex(ValueError, "persist_decode"): _spec(_gfx950_dense_req(), _PERSIST, persist_decode="not-a-decode") - def test_s4096_shape_selects_two_phase_pair_and_wide_dma(self): + def test_s4096_shape_accepts_two_phase_pair_and_wide_dma(self): req = _gfx950_dense_req(**{**_LLAMA_8K, "seqlen_q": 4096, "seqlen_k": 4096}) - spec = _spec(req, _WIDE) + spec = _spec(req, _WIDE, persist_decode="gqa_pair_2phase") self.assertEqual(spec.resolved_persist_decode, "gqa_pair_2phase") self.assertTrue(spec.wide_lds_dma) self.assertIn("gqapair2", spec.kernel_name()) self.assertIn("wdma", spec.kernel_name()) def test_bf16_shape_uses_invariant_compatible_fast_path(self): - spec = _spec(_gfx950_dense_req(**{**_LLAMA_8K, "dtype": "bf16"}), _WIDE) + spec = _spec( + _gfx950_dense_req(**{**_LLAMA_8K, "dtype": "bf16"}), + _WIDE, + persist_decode="gqa_pair", + ) self.assertEqual(spec.resolved_persist_decode, "gqa_pair") self.assertTrue(spec.wide_lds_dma) - def test_s2048_h64_selects_two_phase_pair(self): + def test_s2048_h64_accepts_two_phase_pair(self): req = _gfx950_dense_req( **{**_LLAMA_8K, "nhead_q": 64, "seqlen_q": 2048, "seqlen_k": 2048} ) - spec = _spec(req, _WIDE) + spec = _spec(req, _WIDE, persist_decode="gqa_pair_2phase") self.assertEqual(spec.resolved_persist_decode, "gqa_pair_2phase") self.assertTrue(spec.wide_lds_dma) - def test_explicit_two_phase_pair_knob_canonicalizes_to_the_default(self): + def test_explicit_pair_decodes_are_kept_with_their_own_ids(self): + """A decode auto does not pick is kept, and gets its own id.""" req = _gfx950_dense_req(**{**_LLAMA_8K, "seqlen_q": 4096, "seqlen_k": 4096}) - pinned = tuning_spec_with_knobs( - req, _WIDE, {"persist_decode": "gqa_pair_2phase"} - ) - self.assertEqual(pinned, attention_tuning_spec(req, _WIDE)) - self.assertEqual(pinned.kernel_spec.resolved_persist_decode, "gqa_pair_2phase") - # A decode auto does not pick is kept, and gets its own id. - other = tuning_spec_with_knobs(req, _WIDE, {"persist_decode": "hkv_major"}) - self.assertEqual( - dict(other.knobs), {"block_m": 256, "persist_decode": "hkv_major"} - ) - self.assertNotEqual(other.tuning_id, pinned.tuning_id) + default = attention_tuning_spec(req, _WIDE) + ids = {default.tuning_id} + for decode in ("gqa_pair_2phase", "hkv_major"): + with self.subTest(decode=decode): + pinned = tuning_spec_with_knobs(req, _WIDE, {"persist_decode": decode}) + self.assertEqual( + dict(pinned.knobs), {"block_m": 256, "persist_decode": decode} + ) + self.assertEqual(pinned.kernel_spec.resolved_persist_decode, decode) + self.assertNotIn(pinned.tuning_id, ids) + ids.add(pinned.tuning_id) def test_exact_shape_with_sinks_keeps_gqa_pair(self): req = _gfx950_dense_req(**_LLAMA_8K, use_sinks=True) - spec = _spec(req, _PERSIST) + spec = _spec(req, _PERSIST, persist_decode="gqa_pair") self.assertEqual(spec.resolved_persist_decode, "gqa_pair") self.assertTrue(spec.use_sinks) self.assertFalse(spec.wide_lds_dma) # No dispatch gate stricter than the kernel: wide DMA admits sinks too. - self.assertTrue(_spec(req, _WIDE).use_sinks) + self.assertTrue(_spec(req, _WIDE, persist_decode="gqa_pair").use_sinks) def test_sliding_window_falls_back_to_qb_major(self): req = _gfx950_dense_req(**_LLAMA_8K, sliding_window=512) @@ -290,12 +322,35 @@ def test_sliding_window_falls_back_to_qb_major(self): _spec(req, spec_id).resolved_persist_decode, "qb_major" ) - def test_mha_falls_back_to_qb_major(self): + def test_mha_selects_bt_hkv_minor(self): spec = _spec(_gfx950_dense_req(**{**_LLAMA_8K, "nhead_k": 32}), _WIDE) - self.assertEqual(spec.resolved_persist_decode, "qb_major") + self.assertEqual(spec.resolved_persist_decode, "bt_hkv_minor") self.assertTrue(spec.wide_lds_dma) self.assertNotIn("gqapair", spec.kernel_name()) + def test_large_batch_selects_hkv_major_for_gqa(self): + gqa = _gfx950_dense_req(**{**_LLAMA_8K, "batch": 16}) + self.assertEqual(_spec(gqa, _PERSIST).resolved_persist_decode, "hkv_major") + + def test_large_batch_selects_hq_minor_swz_for_mha(self): + mha = dict(_LLAMA_8K, batch=16, nhead_k=32, seqlen_q=2048, seqlen_k=2048) + spec = _spec(_gfx950_dense_req(**mha), _PERSIST) + self.assertEqual(spec.resolved_persist_decode, "hq_minor_swz") + + def test_large_batch_mha_without_head_bands_selects_qb_major(self): + # head bands need a head count divisible by the XCD count + mha = dict( + _LLAMA_8K, batch=16, nhead_q=36, nhead_k=36, seqlen_q=2048, seqlen_k=2048 + ) + spec = _spec(_gfx950_dense_req(**mha), _PERSIST) + self.assertEqual(spec.resolved_persist_decode, "qb_major") + + def test_large_causal_mha_uses_persistent_swizzled_head_first(self): + mha = dict(_LLAMA_8K, batch=4, nhead_q=64, nhead_k=64) + spec = _spec(_gfx950_dense_req(**mha), _WIDE) + self.assertTrue(spec.wide_lds_dma) + self.assertEqual(spec.resolved_persist_decode, "hq_minor_swz") + class TestDenseBottomRightWiring(unittest.TestCase): @staticmethod @@ -441,7 +496,7 @@ def test_persistent_candidates_refuse_moving_bottom_right(self): with self.assertRaisesRegex(ValueError, "bottom.right"): _spec(self._moving_req(mask_type), spec_id) - def test_equal_length_bottom_right_preserves_gqa_pair_and_wide_dma(self): + def test_equal_length_bottom_right_preserves_decode_and_wide_dma(self): mask_pairs = ( (AttentionMaskType.TOP_LEFT_CAUSAL, 2), (1, AttentionMaskType.BOTTOM_RIGHT_CAUSAL), @@ -457,7 +512,9 @@ def test_equal_length_bottom_right_preserves_gqa_pair_and_wide_dma(self): self.assertEqual(bottom_right_spec, top_left_spec) self.assertFalse(bottom_right_spec.causal_bottom_right) self.assertTrue(bottom_right_spec.persistent) - self.assertEqual(bottom_right_spec.resolved_persist_decode, "gqa_pair") + self.assertEqual( + bottom_right_spec.resolved_persist_decode, "bt_hkv_minor" + ) self.assertTrue(bottom_right_spec.wide_lds_dma) @@ -497,7 +554,7 @@ def test_dispatch_captures_named_default_geometry(self): self.assertEqual(spec.lds_v_row_pad, layout["lds_v_row_pad"]) def test_block_m_controls_grid_and_kernel_identity(self): - default = self._spec() + default = replace(self._spec(), nonpersist_decode="qb_minor") bm128 = replace(default, block_m=128) self.assertEqual(attention_dense_grid(default), (2, 32, 1)) self.assertEqual(attention_dense_grid(bm128), (4, 32, 1)) @@ -998,10 +1055,11 @@ class TestGfx950DenseCandidates(unittest.TestCase): """Two algorithms (grid and persistent body); wide DMA is a persistent candidate; the tile is a knob of each.""" - def test_three_dense_candidates_are_registered(self): + def test_four_dense_candidates_are_registered(self): by_name = {c.name: c for c in attention_candidates()} expected = { _GRID_NAME: DENSE_GRID_ALGORITHM, + _QBMINOR_NAME: DENSE_GRID_ALGORITHM, _PERSIST_NAME: DENSE_PERSIST_ALGORITHM, _WIDE_NAME: DENSE_PERSIST_ALGORITHM, } @@ -1013,14 +1071,18 @@ def test_three_dense_candidates_are_registered(self): with self.subTest(name=name): self.assertEqual(by_name[name].algorithm, algorithm) - def test_d128_admits_all_three_dense_candidates(self): + def test_d128_admits_all_four_dense_candidates(self): req = _gfx950_dense_req( hdim_q=128, hdim_v=128, nhead_q=32, nhead_k=8, seqlen_q=2048, seqlen_k=2048 ) - self.assertEqual(_admitting(req), {_GRID_NAME, _PERSIST_NAME, _WIDE_NAME}) + self.assertEqual( + _admitting(req), {_GRID_NAME, _QBMINOR_NAME, _PERSIST_NAME, _WIDE_NAME} + ) def test_d64_refuses_wide_dma(self): - self.assertEqual(_admitting(_gfx950_dense_req()), {_GRID_NAME, _PERSIST_NAME}) + self.assertEqual( + _admitting(_gfx950_dense_req()), {_GRID_NAME, _QBMINOR_NAME, _PERSIST_NAME} + ) def test_tile_is_a_knob(self): req = _gfx950_dense_req(hdim_q=128, hdim_v=128, nhead_q=32, nhead_k=8) @@ -1036,6 +1098,49 @@ def test_tile_is_a_knob(self): with self.assertRaises(ValueError): _spec(req, _WIDE, block_n=128) + def test_nonpersist_decode_is_a_knob_of_the_grid_body(self): + """The non-persistent block order is a knob of the grid candidate, kept + unless it restates auto; the persistent bodies never read it and the + qbminor variant fixes it, so there it is dropped.""" + req = _gfx950_dense_req(hdim_q=128, hdim_v=128, nhead_q=32, nhead_k=8) + grid = _spec(req, _GRID, nonpersist_decode="hq_minor_swz") + self.assertEqual(grid.resolved_nonpersist_decode, "hq_minor_swz") + self.assertIn("nphqminswz", grid.kernel_name()) + auto = attention_tuning_spec(req, _GRID) + restated = tuning_spec_with_knobs( + req, + _GRID, + {"nonpersist_decode": auto.kernel_spec.resolved_nonpersist_decode}, + ) + self.assertEqual(restated, auto) + for spec_id, fixed in ( + (_PERSIST, "auto"), + (_WIDE, "auto"), + (_QBMINOR, "qb_minor"), + ): + with self.subTest(spec_id=spec_id): + pinned = tuning_spec_with_knobs( + req, spec_id, {"nonpersist_decode": "hq_minor_swz"} + ) + self.assertEqual(pinned, attention_tuning_spec(req, spec_id)) + self.assertEqual(pinned.kernel_spec.nonpersist_decode, fixed) + + def test_qbminor_sweep_keeps_the_ingestor_grid(self): + req = _gfx950_dense_req( + hdim_q=128, hdim_v=128, nhead_q=32, nhead_k=8, seqlen_q=2048, seqlen_k=2048 + ) + orders = { + name: { + spec.kernel_spec.resolved_nonpersist_decode + for _c, spec in iter_registered_attention_combos( + req, candidate_prefix=name + ) + } + for name in (_GRID_NAME, _QBMINOR_NAME) + } + self.assertEqual(orders[_QBMINOR_NAME], {"qb_minor"}) + self.assertGreater(len(orders[_GRID_NAME]), 1) + def test_tile_knob_recomputes_ragged(self): req = _gfx950_dense_req( hdim_q=128, hdim_v=128, nhead_q=32, nhead_k=8, seqlen_q=384, seqlen_k=384 @@ -1078,7 +1183,7 @@ def test_llama3_8b_s8192_needs_an_explicit_pin(self): self.assertEqual(spec.block_m, 256) self.assertTrue(spec.persistent) self.assertTrue(spec.wide_lds_dma) - self.assertEqual(spec.resolved_persist_decode, "gqa_pair") + self.assertEqual(spec.resolved_persist_decode, "bt_hkv_minor") def test_registered_combos_include_dense_not_unified_2d(self): from dispatch.attention import registered_attention_combos 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 9e023fb80ed9..f8bae6411308 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 @@ -828,7 +828,7 @@ def test_production_varies_every_declared_knob(self): def test_production_specs_are_legal_distinct_and_start_at_the_shipped_spec(self): req = _dense_request() grouped = _by_candidate(_dense_combos(req)) - self.assertEqual(len(grouped), 3, sorted(grouped)) + self.assertEqual(len(grouped), 4, sorted(grouped)) for name, (candidate, specs) in grouped.items(): with self.subTest(candidate=name): probe = replace( @@ -875,7 +875,7 @@ def test_num_persistent_policies_reach_every_persistent_decode(self): kernels = [spec.kernel_spec for spec in specs] self.assertEqual( {k.resolved_persist_decode for k in kernels}, - {"qb_major", "hkv_major", "gqa_pair", "gqa_pair_2phase"}, + Gfx950AttentionDenseSpec.supported_persist_decodes() - {"auto"}, ) swept_counts = {k.num_persistent for k in kernels} for policy in ("gqa_pair", "gqa_pair_2phase", "0.75x", "1.25x", "2x"): @@ -1079,18 +1079,26 @@ def test_redundant_settings_are_pruned(self): def canonical(knobs, r=req): return tuning_spec_with_knobs(r, "gfx942_dense", knobs) + self.assertTrue(base.persistent) # D128: the default spec is persistent redundant = ( dict(use_cfvst=True), # the policy already turns it on here dict(use_exp2_fast=base.resolved_use_exp2_fast()), dict(persistent=base.persistent), # base-relative no-op - dict(num_persistent=base.num_persistent * 2), # grid body ignores it - dict(persist_decode="hkv_major"), - dict(interleave=True), + dict(nonpersist_decode="hq_minor_swz"), # persistent body ignores it ) default = attention_tuning_spec(req, "gfx942_dense") for knobs in redundant: with self.subTest(**knobs): self.assertEqual(canonical(knobs), default) + # The grid body ignores the persistent-only knobs. + grid = dict(persistent=False) + for knobs in ( + dict(num_persistent=base.num_persistent * 2), + dict(persist_decode="hkv_major"), + dict(interleave=True), + ): + with self.subTest(grid=knobs): + self.assertEqual(canonical({**grid, **knobs}), canonical(grid)) # No cfvst path at D64: the V knobs are out of scope, so dropped, not refused. d64_default = attention_tuning_spec(d64_req, "gfx942_dense") self.assertEqual(d64_default.kernel_spec, d64) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_dense_gfx942_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_dense_gfx942_ir_sha256.json index 7980ad83b94f..823e1f2845f8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_dense_gfx942_ir_sha256.json +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_dense_gfx942_ir_sha256.json @@ -3,116 +3,116 @@ "llvm20": { "cases": { "attention_dense_gfx942/default_d128_bf16_causal": { - "bytes": 170850, - "sha256": "f114a2f8804581630ce87caf972d7aabd2e32ffa122bc7f94f9b46ba6e038167" + "bytes": 171590, + "sha256": "3a21804add31a71453e5f86874912285e3f66bb0cb8c3b481f44e6b55ad00812" }, "attention_dense_gfx942/default_d128_bf16_mha": { - "bytes": 170850, - "sha256": "f114a2f8804581630ce87caf972d7aabd2e32ffa122bc7f94f9b46ba6e038167" + "bytes": 171590, + "sha256": "3a21804add31a71453e5f86874912285e3f66bb0cb8c3b481f44e6b55ad00812" }, "attention_dense_gfx942/default_d128_fp16_causal": { - "bytes": 113206, - "sha256": "67333a65c6a6f4ccaa4b67c3eb59a48ff4c1cd49a1e4f2359b57ef336b53e16f" + "bytes": 113934, + "sha256": "445cb4dce51274a349a8e543bfee9bb59bbcf7b04d0a87eb3fef2a7ac9a48de7" }, "attention_dense_gfx942/default_d128_fp16_full": { "bytes": 97814, "sha256": "a183a50eaa5374d0f78a93c76821e92d8c871d4fd944cfc9d9e478f4f00efa61" }, "attention_dense_gfx942/default_d128_fp16_gqa16_8": { - "bytes": 113206, - "sha256": "67333a65c6a6f4ccaa4b67c3eb59a48ff4c1cd49a1e4f2359b57ef336b53e16f" + "bytes": 113934, + "sha256": "445cb4dce51274a349a8e543bfee9bb59bbcf7b04d0a87eb3fef2a7ac9a48de7" }, "attention_dense_gfx942/default_d64_bf16_bn128": { - "bytes": 185988, - "sha256": "de162ee0bd46b1dd0c2cf721f951fc39276cf14845e9f46495efb6a5da715b86" + "bytes": 186729, + "sha256": "772788965bdb222c3efe676757bd61982148d507ed217e30db100e0d3a768528" }, "attention_dense_gfx942/default_d64_bf16_causal": { - "bytes": 103197, - "sha256": "5df3dd8bc0d490daa515f51c39a5c7265c6ae9eb67cc5db34c30dbd2925bd900" + "bytes": 103760, + "sha256": "875b2984e0841ef186bf3b2fdb8b75c09da7e1bb65faf6e6c83863d32b88a9c3" }, "attention_dense_gfx942/default_d64_fp16_causal": { - "bytes": 97180, - "sha256": "29129acc4ed1a52a4c8c5c122a65ed546864018752ad15dd548134061fb655a2" + "bytes": 97743, + "sha256": "7e288d18dec9a9648035df2a29552ac71605d07b60e08368c53f3c58eb8b57ac" }, "attention_dense_gfx942/dispatch_d64_bf16_causal": { - "bytes": 104409, - "sha256": "dee9e70e44c2d724f07a5f6f3c400ad06d7340b48a0f526d606964a043d1b322" + "bytes": 104976, + "sha256": "ff3ce426327f0b87b412cf764c4a2ee1b570667605371d574eef31d5348a9600" }, "attention_dense_gfx942/dispatch_d64_fp16_causal": { - "bytes": 98392, - "sha256": "305409b65503241921220ccaa73b350d1011a727c352ae41e8f1ead88ed54fef" + "bytes": 98959, + "sha256": "7ae8ad27c77e9c0fbcf03b653d0baa3b02cf796b71976380c5b8c445ccecffca" }, "attention_dense_gfx942/dispatch_lse_d128_bf16_causal": { - "bytes": 177689, - "sha256": "e161b8a0de9e4164977f4ae8631aa266b009d55e37329b38aa0f124f36999930" + "bytes": 181452, + "sha256": "fcf6b7da60898c8062227a978235cae2ade574586db3d2ff380f76b160a4298c" }, "attention_dense_gfx942/dispatch_sink_d64_bf16_causal": { - "bytes": 104831, - "sha256": "e5bd846f9ab5dac5e7c1a56bcd7c5a3c127b70b4b84738090aed07574e0f6e0b" + "bytes": 105393, + "sha256": "200518cbffbb142e3ec56c951382e59674795147219980a3ebba436a393391da" }, "attention_dense_gfx942/dispatch_swa_d128_bf16_w128": { - "bytes": 175001, - "sha256": "443e2dc10e981be13c5946a9594db674ff39af026386dda7b6538373ae1ea418" + "bytes": 178454, + "sha256": "0e7dce4882d00b1cce04c36ba66f5ae749f061fa66d45fd21fce02e8d14e4925" }, "attention_dense_gfx942/lse_d128_bf16_causal": { - "bytes": 177689, - "sha256": "e161b8a0de9e4164977f4ae8631aa266b009d55e37329b38aa0f124f36999930" + "bytes": 178429, + "sha256": "15390d0c3bb2f04b83a10166f160c44a5553d8298057f68e64f820155e7add4a" }, "attention_dense_gfx942/lse_persist_empty_window_d128_bf16": { - "bytes": 184443, - "sha256": "5d6db8b2dc0fef4f94d618f6bf64861e87614fe31c590b357c6e3ac555a0556b" + "bytes": 185931, + "sha256": "0695fa68401f74c016c4b677a89bcb09582a8b2692d50dcc71c92fe44851f9c2" }, "attention_dense_gfx942/lse_sink_persist_empty_window_d128_bf16": { - "bytes": 184990, - "sha256": "d0fa692777ad8667736de1c417b698ca4efdc786ceb303daaff38f3bd91f2b15" + "bytes": 186480, + "sha256": "b6a7bbe360dff6bdc8f98c81f730736002d60c5d48129479df3e55900bf80f85" }, "attention_dense_gfx942/persist_d128_bf16_intl": { - "bytes": 173617, - "sha256": "c2d907f63da1a257cacae2efb9adf3a226d2cf3f2d9eec5b1827aaca6620fa23" + "bytes": 174216, + "sha256": "279d5d2fbfbeaf1e36529843e5958901109c561011497045684cfccd9dc5124d" }, "attention_dense_gfx942/persist_d128_bf16_qbmaj": { - "bytes": 173206, - "sha256": "21b2414ed8e74b100057607234503f17af6df1661f222a119db911384b8e42e3" + "bytes": 174079, + "sha256": "c80fa920c957c6fed01908f39bc1b0f5b32b0613f3e4fa98361006b0322c9c75" }, "attention_dense_gfx942/persist_d128_fp16_hkvmaj": { - "bytes": 115765, - "sha256": "29559af89c008d4b1f73a605c4c8a8645720aa5f55f3a03c5b2c683069f98c5c" + "bytes": 120777, + "sha256": "ddaa37fe0e2b96ba3535efa00704986efe999a84e92ae13d58656d26783f1a1f" }, "attention_dense_gfx942/persist_d64_bf16_hkvmaj": { - "bytes": 106693, - "sha256": "e6d25511f4fb1f1d4e7f57aa136bf72b18ac0ee06f01e76987f7ce9887db80af" + "bytes": 111676, + "sha256": "7391eacb7296dab59c509628019abb89d7adb928334534138d8d7eadfb2e5bc8" }, "attention_dense_gfx942/persist_d64_fp16_qbmaj": { - "bytes": 100228, - "sha256": "ce955ab86bbf411befc81b3eda72d7cd9ea67a8388c8cf5247f112d2c15ee992" + "bytes": 101084, + "sha256": "fb41aa30f46608819cb1a10327aaaf99613b4dd45c6b87bb399ab62ff85c9dae" }, "attention_dense_gfx942/persist_sink_d128_bf16_qbmaj": { - "bytes": 173736, - "sha256": "9d4fb5f841cc49e5f52f78f0d53a1aa8779d3c1debcccae6a3b28b103ab476ea" + "bytes": 174592, + "sha256": "4cacee0fc5604453e35fb33740b6e030a0768f851e825d578a030e46b01bef39" }, "attention_dense_gfx942/persist_sink_d128_fp16_hkvmaj": { - "bytes": 116269, - "sha256": "36038c6aeb695898708fd8070322038878d746a4caac49dc5329e6ebf5f80aeb" + "bytes": 121279, + "sha256": "ab4c06e5eb2f6360e89994169bf1c20748f766f6099b749810b96ff3acbbcdda" }, "attention_dense_gfx942/persist_swa_d128_bf16_w128": { - "bytes": 177498, - "sha256": "8a9ed70d9e9c9dc65e52e502908348437b42c45d012055988dceeaa555f30de3" + "bytes": 178454, + "sha256": "0e7dce4882d00b1cce04c36ba66f5ae749f061fa66d45fd21fce02e8d14e4925" }, "attention_dense_gfx942/persist_swa_sink_d128_bf16_w128": { - "bytes": 178016, - "sha256": "f4173b11417cc798bff1bb0b38350adb127fb60a0aebea15653f294ed2a8909b" + "bytes": 178969, + "sha256": "ffb28779bc551be8722dcc2ab8b141f02599b5cc7668b3a1e2b7b82a1d560c17" }, "attention_dense_gfx942/sink_d128_bf16_causal": { - "bytes": 171366, - "sha256": "143a9755f894cca55c75b268c2556537cc477bf41be55b5709c691ff1de9b888" + "bytes": 172103, + "sha256": "eff00bf729567a82bcd771814ca4e702b9420cb93585b55a61418bc5660822ec" }, "attention_dense_gfx942/sink_d128_fp16_full": { "bytes": 98321, "sha256": "5189bf6a4d982854b6095015b8c9991852b030790a37b264f0764f828cc14d77" }, "attention_dense_gfx942/sink_d64_fp16_causal": { - "bytes": 97591, - "sha256": "b11be08d9eca3e6cd963d9e657060385408d61310bc770ebf0aedc1982c34a3d" + "bytes": 98156, + "sha256": "2e125ff9d2da8b100d4a127948f5dca85f50c03e8595da7e1284d3f1e1f3a700" }, "attention_dense_gfx942/swa_d128_bf16_w128": { "bytes": 175001, @@ -135,116 +135,116 @@ "llvm22": { "cases": { "attention_dense_gfx942/default_d128_bf16_causal": { - "bytes": 170946, - "sha256": "24380550a9745e6d209ad497be3c3b81336a5708f163ffb4c5e96bf744132bd1" + "bytes": 171686, + "sha256": "2658a44f0016de96754d17f3dff885a6c6d018ef70629ee728410050218b3c24" }, "attention_dense_gfx942/default_d128_bf16_mha": { - "bytes": 170946, - "sha256": "24380550a9745e6d209ad497be3c3b81336a5708f163ffb4c5e96bf744132bd1" + "bytes": 171686, + "sha256": "2658a44f0016de96754d17f3dff885a6c6d018ef70629ee728410050218b3c24" }, "attention_dense_gfx942/default_d128_fp16_causal": { - "bytes": 113302, - "sha256": "17d7827f7e0f8c203c9529abfdedfed8f461e6b491c5fd70ac9287be87fbc8f1" + "bytes": 114030, + "sha256": "665bfbb89797f0981c4ea0b660a6c2398a723661816ce40f0a5a9a311f279eb5" }, "attention_dense_gfx942/default_d128_fp16_full": { "bytes": 97910, "sha256": "30a2bf4bf475f87ae02068198e459357bac3bd5074d21ba25c45e8deacd4454f" }, "attention_dense_gfx942/default_d128_fp16_gqa16_8": { - "bytes": 113302, - "sha256": "17d7827f7e0f8c203c9529abfdedfed8f461e6b491c5fd70ac9287be87fbc8f1" + "bytes": 114030, + "sha256": "665bfbb89797f0981c4ea0b660a6c2398a723661816ce40f0a5a9a311f279eb5" }, "attention_dense_gfx942/default_d64_bf16_bn128": { - "bytes": 186084, - "sha256": "64d991d862d20a3111d55a509080fb61151c4e45e9d7cca468c4a2305fc29513" + "bytes": 186825, + "sha256": "700cf25032f17ef16affc1d057115e203549f92ec72d05ba38aa9955040dbf25" }, "attention_dense_gfx942/default_d64_bf16_causal": { - "bytes": 103293, - "sha256": "0f6489269d39d83f116b77e599a6aa362192e70f6a2466449a5cb2a1ae656d14" + "bytes": 103856, + "sha256": "49f0e2ddea3d54f07b8c2750115c3b58af8a1b089c479ed2eb5e0f16d1f8b5c1" }, "attention_dense_gfx942/default_d64_fp16_causal": { - "bytes": 97276, - "sha256": "2af0ea0aa87cad7ccbd05774364f3ef4db629b67f74911f89c572e197098a5c1" + "bytes": 97839, + "sha256": "6aa77b6f76773c8bb8632f44965c508a7b3266a0d367eba0e42444e12a58f508" }, "attention_dense_gfx942/dispatch_d64_bf16_causal": { - "bytes": 104505, - "sha256": "36760a0a303f9cf4e4bb5b24ea921775e7b83b46d00079e0923f4e6e5a9e228b" + "bytes": 105072, + "sha256": "c916ad134243e659e040dcb76b73d4faf5ca7eec68ce675dee9bb52237d27cf3" }, "attention_dense_gfx942/dispatch_d64_fp16_causal": { - "bytes": 98488, - "sha256": "930b02fedbb935624accd81ff1fe561c6597a28ca11cf615becb9784d7522439" + "bytes": 99055, + "sha256": "c0573d6d737b64b02caf60ffa098e802c8cee9a4aa552649aa680af481121915" }, "attention_dense_gfx942/dispatch_lse_d128_bf16_causal": { - "bytes": 177785, - "sha256": "63490a820e7354947c259d49672c5ea01889f9c9864e02efe0ed259ca9ffe61d" + "bytes": 181548, + "sha256": "c67ca5774a53bc3f8cd23810b5f722d13ad96f07a88e2b2d2708a5535a6c6272" }, "attention_dense_gfx942/dispatch_sink_d64_bf16_causal": { - "bytes": 104927, - "sha256": "b812d4f95f9a0d54b145fd37adf33ebb5cd103b27904a1041c9cd95e7d1753a9" + "bytes": 105489, + "sha256": "88e42b5032a74a9ce36283af3017729737c743cdb72f94d551d84a564d339a80" }, "attention_dense_gfx942/dispatch_swa_d128_bf16_w128": { - "bytes": 175097, - "sha256": "831246c3f5a0e5f0509f926405cab8b5f30e29df9e6a7c6d5da46bc45db399f6" + "bytes": 178550, + "sha256": "5466a37a743495c049282c02cbe19eac1df6ce96955e8aec346f4906385918ec" }, "attention_dense_gfx942/lse_d128_bf16_causal": { - "bytes": 177785, - "sha256": "63490a820e7354947c259d49672c5ea01889f9c9864e02efe0ed259ca9ffe61d" + "bytes": 178525, + "sha256": "360f96a913ca34db8880643a4101efa9a1559272d347ee87f425cd500bb8b390" }, "attention_dense_gfx942/lse_persist_empty_window_d128_bf16": { - "bytes": 184539, - "sha256": "43c47ec48fdd586066beaa21f871a63ac389f84574af1152cccd612ab4fad1fb" + "bytes": 186027, + "sha256": "dd2d507d110839f77636a9705f2b2d4b064351f2f203a8ebc1177bed01d66ad5" }, "attention_dense_gfx942/lse_sink_persist_empty_window_d128_bf16": { - "bytes": 185086, - "sha256": "a55bdf94404df1fa02009d4429dcacef74fddd5e4dd4e6798dcda7c50a58e9b4" + "bytes": 186576, + "sha256": "85ba825b459ad10fd73ac96c80a7b5f3ef51d0d0a88f5c2d99e7b646b57c8068" }, "attention_dense_gfx942/persist_d128_bf16_intl": { - "bytes": 173713, - "sha256": "aac88b95ea00d399b5bd67e7094ead9659925d6abd91aedfad211e4a048c15ed" + "bytes": 174312, + "sha256": "8b187788a121447658b3a60dc296b659f2ac992140c8274e1107aee5e33acac5" }, "attention_dense_gfx942/persist_d128_bf16_qbmaj": { - "bytes": 173302, - "sha256": "36ad138df0b091a5399b29ad818fb03213d131aa94c6dda7a884867a7b2b5eca" + "bytes": 174175, + "sha256": "5f7ce04016f784a1bfcc82849945ea94ef9a86871c818ea2f7b308731835fa84" }, "attention_dense_gfx942/persist_d128_fp16_hkvmaj": { - "bytes": 115861, - "sha256": "0d9d315ad11fcac7d10dd7b87f0bf771b398d5f226cdb7ec1647adfe0244f2f1" + "bytes": 120873, + "sha256": "f221e7175ec2607aa41bc9ea94106408d0692cca7dfcd3ab2959405ededf326c" }, "attention_dense_gfx942/persist_d64_bf16_hkvmaj": { - "bytes": 106789, - "sha256": "b04e2a2482db1a690b246995a46430c65d3fe32420c16de8aaba3e1c58201bbc" + "bytes": 111772, + "sha256": "0ef3ff4aebc9c9e0374d95119a84ea76b6ac8db00ac7ceaebe8deca24dda3a40" }, "attention_dense_gfx942/persist_d64_fp16_qbmaj": { - "bytes": 100324, - "sha256": "e1ba9da81bc46b0934e2203e7f1d10bb4dcbec0eafeba1c892e9330465285e61" + "bytes": 101180, + "sha256": "9795b13b2179739e6504054478a30732b4fb5346cb9f5cab5926a49c5ba8c185" }, "attention_dense_gfx942/persist_sink_d128_bf16_qbmaj": { - "bytes": 173832, - "sha256": "566915921d6e8c9a9688af470d4c920933742e2d6be65b74892cc49c249b6843" + "bytes": 174688, + "sha256": "8c4f15cff00adee735dcec47d96c6b1917a8609182a624fd18c53b197b00beaf" }, "attention_dense_gfx942/persist_sink_d128_fp16_hkvmaj": { - "bytes": 116365, - "sha256": "c7b860e4c607f534228998bec0eb851791838ceb9f1c395e5cf05ab2ca82e207" + "bytes": 121375, + "sha256": "de006d6ce86a18a7ebb3d880db2993551fb4433f63a44fa8327d97a7d1b404a3" }, "attention_dense_gfx942/persist_swa_d128_bf16_w128": { - "bytes": 177594, - "sha256": "2750a53f8244bc7f61acf5d42abed57b1dd2aea0ae5c38dfa9f672dbe7273677" + "bytes": 178550, + "sha256": "5466a37a743495c049282c02cbe19eac1df6ce96955e8aec346f4906385918ec" }, "attention_dense_gfx942/persist_swa_sink_d128_bf16_w128": { - "bytes": 178112, - "sha256": "0330552760e69cb1fd54ba3fefdc501599621773085fc3dcd0a1d5669acdc102" + "bytes": 179065, + "sha256": "9218c9179f0b72d2c6eab8f32495e91c70520d737c21201a1fda0b4ce9c33cb6" }, "attention_dense_gfx942/sink_d128_bf16_causal": { - "bytes": 171462, - "sha256": "09c9c902d0b51946979ecc7c635cbf3edde6f1f66847b8e5838d28d94b4f9a73" + "bytes": 172199, + "sha256": "48fdf25db27ad6317f33e4d40c5a286f74377471fbfec608160680f809fd15d3" }, "attention_dense_gfx942/sink_d128_fp16_full": { "bytes": 98417, "sha256": "4ba5547683227ebb00e8620714e4ab189d36c3e15fd0070ea81dd13bbfc910de" }, "attention_dense_gfx942/sink_d64_fp16_causal": { - "bytes": 97687, - "sha256": "36b65514f928489a0fd5b2a91f2a1c5451dbf271706f7469cf71eb7355365880" + "bytes": 98252, + "sha256": "936454e6484a7a3bd3178b0475b5896ea306147870ed31011a66666c3db80fe6" }, "attention_dense_gfx942/swa_d128_bf16_w128": { "bytes": 175097, @@ -267,116 +267,116 @@ "llvm23": { "cases": { "attention_dense_gfx942/default_d128_bf16_causal": { - "bytes": 171010, - "sha256": "6dd00349c8c44cb04c8e96661dc302637d0ec4b433ae0551c42ae15429fe98bb" + "bytes": 171750, + "sha256": "bc12b9a1aba7d9bfae67d64c24339d006cec293c40e31a9f41e33aefbe7981eb" }, "attention_dense_gfx942/default_d128_bf16_mha": { - "bytes": 171010, - "sha256": "6dd00349c8c44cb04c8e96661dc302637d0ec4b433ae0551c42ae15429fe98bb" + "bytes": 171750, + "sha256": "bc12b9a1aba7d9bfae67d64c24339d006cec293c40e31a9f41e33aefbe7981eb" }, "attention_dense_gfx942/default_d128_fp16_causal": { - "bytes": 113366, - "sha256": "4a7395df00355cbe3b620074fea850cae8435b40107d30d6438aaee97a434863" + "bytes": 114094, + "sha256": "421c4f48f5c8499e82cd905ff2c848500062da48aa84b36df764e545a82ec377" }, "attention_dense_gfx942/default_d128_fp16_full": { "bytes": 97974, "sha256": "e6462fd928ba04e0305ee5e9507c43606a09af03d0f92a229c18b419213e820b" }, "attention_dense_gfx942/default_d128_fp16_gqa16_8": { - "bytes": 113366, - "sha256": "4a7395df00355cbe3b620074fea850cae8435b40107d30d6438aaee97a434863" + "bytes": 114094, + "sha256": "421c4f48f5c8499e82cd905ff2c848500062da48aa84b36df764e545a82ec377" }, "attention_dense_gfx942/default_d64_bf16_bn128": { - "bytes": 186148, - "sha256": "6efd7961192087168947cedf8f7624a06338ff8522152fe364c693c05252bd8e" + "bytes": 186889, + "sha256": "5a3fa8a6288c53506e6da383a5cd20ef7e3230c01f5685522599fa695eea0740" }, "attention_dense_gfx942/default_d64_bf16_causal": { - "bytes": 103357, - "sha256": "24bfb10905b13519781fc9ddb6297770bbce7da01d8a17d2b7cae38a70af9cf5" + "bytes": 103920, + "sha256": "8a9e4b4a0f981ba133720bb07e7355a6e2d4a563d2d1528bb1142ca6db7f1624" }, "attention_dense_gfx942/default_d64_fp16_causal": { - "bytes": 97340, - "sha256": "3112c90d6732d110d18ab377bed87bfbe6b7a2672e719800b2689ecd4539e09e" + "bytes": 97903, + "sha256": "70b2e04519dc783e54890edaf6f7d71cf8490baf41f625a687fa05de3ab3686a" }, "attention_dense_gfx942/dispatch_d64_bf16_causal": { - "bytes": 104569, - "sha256": "7b253bc717bf6149355a6db3d529fdfb469bb3ee440bc54ffa7e5aaff60cb293" + "bytes": 105136, + "sha256": "5fb581ff0c94f35dbf7ce1b889c750211cbf40620142575d90b87fe623592038" }, "attention_dense_gfx942/dispatch_d64_fp16_causal": { - "bytes": 98552, - "sha256": "d28c4327b2b18d3987f1574ae59f32d7107b327a6b9aed1c3b668462d17839a7" + "bytes": 99119, + "sha256": "081490b1ad69f7cce85652ef337adb2dea827b324f7420aa66948746544ca1ba" }, "attention_dense_gfx942/dispatch_lse_d128_bf16_causal": { - "bytes": 177849, - "sha256": "f3418840281d17d90e0434d94cfd621994943c5f8d18b7ad7fe941205ffc96ee" + "bytes": 181612, + "sha256": "beb0fc172ef17a23503d1ea8f9631f414887ba5abadb950c79717e3c2fb1a72a" }, "attention_dense_gfx942/dispatch_sink_d64_bf16_causal": { - "bytes": 104991, - "sha256": "848048dafc4dde8c3ee3d6b81e3147d0806d4c7887a2d31072eeceb8a1ec553a" + "bytes": 105553, + "sha256": "e02b1355cda68b1be2c3bad803a3dbdd778efde7ec78a04ea10e1e477586333b" }, "attention_dense_gfx942/dispatch_swa_d128_bf16_w128": { - "bytes": 175161, - "sha256": "d74be2c22e40bc578efe055fc2d6d0f16ae806efdadbfdea71b7e21c997af12b" + "bytes": 178614, + "sha256": "b21ca73b56adfb45d2b243192a77faaf672677b6fe92c05a95deef642e2b01e5" }, "attention_dense_gfx942/lse_d128_bf16_causal": { - "bytes": 177849, - "sha256": "f3418840281d17d90e0434d94cfd621994943c5f8d18b7ad7fe941205ffc96ee" + "bytes": 178589, + "sha256": "31d3c09d639960a9803fdc847515ec3dab0a852082091c90f9f13ef396de7c0b" }, "attention_dense_gfx942/lse_persist_empty_window_d128_bf16": { - "bytes": 184603, - "sha256": "f5308fb42fa20539fdd172b8ab3804fe3f94c43a07ee1111579d73df162d7dac" + "bytes": 186091, + "sha256": "a724043ad5acf38e53968e888c0f8e02019c5c858dd1bc2047941fddf34c5a87" }, "attention_dense_gfx942/lse_sink_persist_empty_window_d128_bf16": { - "bytes": 185150, - "sha256": "76e220469b933345536e1e33db30973b99d2a1016958c794418aba782676d1d4" + "bytes": 186640, + "sha256": "3631253c777bc46831d8b85fa421f81cffe0df9460a83a9a1742c3c0bfdad2df" }, "attention_dense_gfx942/persist_d128_bf16_intl": { - "bytes": 173777, - "sha256": "59f3c8b8d9c19d2f3381f27321647b6af88ef39341992b2788f9867f16a4d883" + "bytes": 174376, + "sha256": "6e231a1399c80e2918e18ca1e92080523df45163c09b5fd2aab621a8326fe11c" }, "attention_dense_gfx942/persist_d128_bf16_qbmaj": { - "bytes": 173366, - "sha256": "4db838378afd810de15ad71a24ba7e738409d21775fe1437798f9ce36d0e48a6" + "bytes": 174239, + "sha256": "36f758e6d65fca80952878584c0a515f2c4376793df33bfbd18c2899402daa33" }, "attention_dense_gfx942/persist_d128_fp16_hkvmaj": { - "bytes": 115925, - "sha256": "ffc03b3b9d67e7bb0efe684298d2c0a11cbb0022118b535d9d7a500ed6b95201" + "bytes": 120937, + "sha256": "a695a4b474328c98b094a24efea16d6b00a7022a41698fd99ad272d3e6da9924" }, "attention_dense_gfx942/persist_d64_bf16_hkvmaj": { - "bytes": 106853, - "sha256": "79acf7785e05757c4c0992554957132ad4a143ca80318744438be448c69b01a7" + "bytes": 111836, + "sha256": "6c5d1eedfb97f3db3fee2cbabde77f3ba6134684c0a8f03a37013eda9684dccb" }, "attention_dense_gfx942/persist_d64_fp16_qbmaj": { - "bytes": 100388, - "sha256": "9f7739d8474304475434d99bd9f2ac6b5b5982f1b26d06311f1c2c307974c681" + "bytes": 101244, + "sha256": "f480aeab86891c6110ff294f7997309ad4fec1bd62c3c75e1f12404a8a3c4ea7" }, "attention_dense_gfx942/persist_sink_d128_bf16_qbmaj": { - "bytes": 173896, - "sha256": "b27bcc336bd61564b218f13c7eaf4b003e8a04ef5d5e5dd0be09ef1ea9431c4d" + "bytes": 174752, + "sha256": "0848ea40055a8d4b810794512cec6c04917a1306636300d02ea4e28245d71283" }, "attention_dense_gfx942/persist_sink_d128_fp16_hkvmaj": { - "bytes": 116429, - "sha256": "c464ae67b79b294b7e0a718c3e263b350d602c6b4754a5cc7bcb5a13d1337075" + "bytes": 121439, + "sha256": "cce4d337e066f92856b53c55b7557eaab29d44f0a16a5a19536e217e613b9681" }, "attention_dense_gfx942/persist_swa_d128_bf16_w128": { - "bytes": 177658, - "sha256": "3c9a0d02f295d9c7bd0ac7ca42a7f6181f61d8454d01a799bc76882504d9c969" + "bytes": 178614, + "sha256": "b21ca73b56adfb45d2b243192a77faaf672677b6fe92c05a95deef642e2b01e5" }, "attention_dense_gfx942/persist_swa_sink_d128_bf16_w128": { - "bytes": 178176, - "sha256": "6590d638c43bdb648a1efc9ed02be20d3a65f2cefc8746d89c347f81ac18635d" + "bytes": 179129, + "sha256": "7499a1560311c9ecc61e06e8df7fda97ea285e5ad75da3625a58e8c49ae086db" }, "attention_dense_gfx942/sink_d128_bf16_causal": { - "bytes": 171526, - "sha256": "b23026700d7ed9ff8d6ec71ec8ff80967fd7e1c6123b59a0585b97225e8e3bdb" + "bytes": 172263, + "sha256": "31d29f2bfc8d35c82af0911a0ef50ffdeaca07cb4198ffc135bcaaa355da1916" }, "attention_dense_gfx942/sink_d128_fp16_full": { "bytes": 98481, "sha256": "5fd974a2134030513a7093b24a6bd38bb72f6070e65ea3f9a159f4a4fb2f2858" }, "attention_dense_gfx942/sink_d64_fp16_causal": { - "bytes": 97751, - "sha256": "78e398b78df8b4ad5e7e27bc0cdfa7551c3217aaa2461ab4c110bbc46e12e82d" + "bytes": 98316, + "sha256": "31db44a4a235ccb5e2f7d686fda39976f6cd77faf7af8403a0a3238772ae0435" }, "attention_dense_gfx942/swa_d128_bf16_w128": { "bytes": 175161, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py index c9806e4d88e7..26c7d1d35c7b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py @@ -129,12 +129,12 @@ def _compile_or_skip(kernel, *, arch: str): # gfx950 dense persistent prefill: the 512-thread (8-wave) workgroup pins VGPR # at the 256 ceiling independent of waves_per_eu; it needs a hair more -> 20 B. # Tuned for 2 waves/SIMD (512 VGPR/SIMD // 256). - "rocke_attention_dense_d128_hq32_kv8_bn64_bf16_sq2048_sk2048_causal_lazyrs_persist256": { + "rocke_attention_dense_d128_hq32_kv8_bn64_bf16_sq2048_sk2048_causal_lazyrs_persist256_bthkvmin": { "max_scratch_bytes": 32, "min_waves_per_simd": 2, }, # LSE retains the same persistent occupancy/scratch limits. - "rocke_attention_dense_d128_hq32_kv8_bn64_bf16_sq2048_sk2048_causal_lazyrs_lse_persist256": { + "rocke_attention_dense_d128_hq32_kv8_bn64_bf16_sq2048_sk2048_causal_lazyrs_lse_persist256_bthkvmin": { "max_scratch_bytes": 32, "min_waves_per_simd": 2, }, @@ -1538,6 +1538,189 @@ def test_gfx950_dense_paged_builds_and_signature(self): self.assertIn(req, names) self.assertIsNotNone(build_attention_dense(spec, arch="gfx950")) + def test_gfx950_signature_matches_the_built_kernels_params(self): + """gfx950 ``attention_dense_signature`` matches the built kernel's params. + + The arms straddle the case where ``runtime_shape`` and + ``_has_shape_params`` disagree (the shape is a param AND baked), which a + launcher gating on the wrong predicate mis-binds on GPU only. + """ + from rocke.core.ir import PtrType + from rocke.helpers.spec import ptr_type_str + from kernels.gfx950.attention_dense import ( + AttentionDenseSpec, + _has_shape_params, + attention_dense_signature, + build_attention_dense, + ) + + def type_str(param): + if isinstance(param.type, PtrType): + return ptr_type_str(param.type.pointee.name, param.type.space) + return param.type.name + + base = dict( + batch=2, + seqlen_q=2048, + seqlen_kv=2048, + num_query_heads=32, + num_kv_heads=8, + head_size=128, + causal=True, + dtype="fp16", + block_n=64, + ) + arms = { + "runtime-shape": AttentionDenseSpec(**base), + "sliding-window": AttentionDenseSpec(**base, sliding_window=512), + "bottom-right": AttentionDenseSpec( + **{**base, "seqlen_q": 1024}, causal_bottom_right=True + ), + "persistent": AttentionDenseSpec( + **base, persistent=True, num_persistent=256 + ), + } + for arm, spec in arms.items(): + with self.subTest(arm=arm): + params = build_attention_dense(spec, arch="gfx950").params + sig = attention_dense_signature(spec) + self.assertEqual([a["name"] for a in sig], [p.name for p in params]) + self.assertEqual( + [a["type"] for a in sig], [type_str(p) for p in params] + ) + self.assertEqual( + _has_shape_params(spec), "batch" in [p.name for p in params] + ) + # The predicates really diverge on these arms. + self.assertEqual( + [ + arms[a].runtime_shape + for a in ("runtime-shape", "sliding-window", "bottom-right") + ], + [True, False, False], + ) + + def test_persist_decodes_build_with_distinct_names_on_both_arches(self): + """Every persist decode an arch supports builds and has its own name, + so no two decodes can share a cached binary.""" + from kernels.gfx942 import attention_dense as k942 + from kernels.gfx950 import attention_dense as k950 + + base = dict( + batch=1, + seqlen_q=2048, + seqlen_kv=2048, + num_query_heads=32, + num_kv_heads=8, + head_size=128, + causal=True, + dtype="fp16", + block_n=64, + persistent=True, + ) + # gqa_pair* pin num_persistent to their own work split (NQB=8, Hkv=8, gqa=4). + num_persistent = {"gqa_pair": 64, "gqa_pair_2phase": 128} + for arch, K, Spec in ( + ("gfx942", k942, k942.Gfx942AttentionDenseSpec), + ("gfx950", k950, k950.Gfx950AttentionDenseSpec), + ): + with self.subTest(arch=arch): + names = {} + for decode in sorted( + Spec(**base).supported_persist_decodes() - {"auto"} + ): + spec = Spec( + **base, + persist_decode=decode, + num_persistent=num_persistent.get(decode, 256), + ) + K.build_attention_dense(spec, arch=arch) + names[decode] = spec.kernel_name() + self.assertIn("bt_hkv_minor", names) + self.assertEqual(len(set(names.values())), len(names), names) + + def test_nonpersist_auto_decode_by_shape(self): + """Non-persistent auto: Swizzled Head-first for large causal MHA, + bt_hkv_minor for other causal attention, qb_minor otherwise.""" + from kernels.gfx942.attention_dense import Gfx942AttentionDenseSpec + from kernels.gfx950.attention_dense import Gfx950AttentionDenseSpec + + cases = [ + # (Hq, Hkv, S, B, causal, extra) -> expected + ((64, 64, 8192, 4, True, {}), "hq_minor_swz"), # 32 blocks, 8192 items + ((128, 128, 4096, 4, True, {}), "hq_minor_swz"), # 16 blocks + ((64, 8, 8192, 4, True, {}), "bt_hkv_minor"), # GQA + ((128, 128, 2048, 8, True, {}), "bt_hkv_minor"), # only 8 blocks + ((32, 32, 8192, 4, True, {}), "bt_hkv_minor"), # 4096 items + ((64, 64, 8192, 4, True, {"chiplet_num_xcds": 128}), "bt_hkv_minor"), + ((64, 64, 8192, 4, True, {"sliding_window": 512}), "qb_minor"), + ((64, 64, 8192, 4, False, {}), "qb_minor"), + ] + for Spec in (Gfx942AttentionDenseSpec, Gfx950AttentionDenseSpec): + for (hq, hkv, s, b, causal, extra), want in cases: + with self.subTest( + spec=Spec.__name__, shape=(hq, hkv, s, b, causal), extra=extra + ): + spec = Spec( + batch=b, + seqlen_q=s, + seqlen_kv=s, + num_query_heads=hq, + num_kv_heads=hkv, + head_size=128, + causal=causal, + dtype="fp16", + block_n=64, + **extra, + ) + self.assertEqual(spec.resolved_nonpersist_decode, want) + + def test_nonpersist_decode_builds_names_and_rejects_on_both_arches(self): + """Every nonpersist decode builds, fills the grid and gets its own name; + the persistent grid ignores it, and an unknown value is rejected.""" + from kernels.common.attention_dense_decode import NONPERSIST_DECODES + from kernels.gfx942 import attention_dense as k942 + from kernels.gfx950 import attention_dense as k950 + + base = dict( + batch=2, + seqlen_q=2048, + seqlen_kv=2048, + num_query_heads=32, + num_kv_heads=8, + head_size=128, + causal=True, + dtype="fp16", + block_n=64, + ) + work = (2048 // 256) * 32 * 2 + for arch, K, Spec in ( + ("gfx942", k942, k942.Gfx942AttentionDenseSpec), + ("gfx950", k950, k950.Gfx950AttentionDenseSpec), + ): + with self.subTest(arch=arch): + names = {} + for decode, impl in NONPERSIST_DECODES.items(): + spec = Spec(**base, nonpersist_decode=decode) + K.build_attention_dense(spec, arch=arch) + gx, gy, gz = K.attention_dense_grid(spec) + self.assertEqual(gx * gy * gz, work, decode) + names[decode] = spec.kernel_name() + if impl.tag: + self.assertIn(impl.tag, names[decode]) + self.assertEqual(len(set(names.values())), len(names)) + auto = Spec(**base) + self.assertEqual( + auto.kernel_name(), names[auto.resolved_nonpersist_decode] + ) + with self.assertRaisesRegex(ValueError, "nonpersist_decode"): + Spec(**base, nonpersist_decode="nope") + persistent = dict(base, persistent=True, num_persistent=256) + self.assertEqual( + Spec(**persistent, nonpersist_decode="bt_hkv_minor").kernel_name(), + Spec(**persistent).kernel_name(), + ) + def test_gfx950_dense_paged_prefill_compiles_and_fits_budget(self): """comgr build + resource-budget net for the PAGED gfx950 dense prefill (fp16/bf16 D128 sliding-window, single-seq). Mirrors the non-paged dense @@ -3124,10 +3307,12 @@ def test_waves_per_eu_cache_key_isolation(self): # Sanity: the two waves_per_eu variants are otherwise indistinguishable, # so the split above is attributable to waves_per_eu alone. The name # tags a non-default waves_per_eu (2 keeps the shipped symbol), so the - # names differ by exactly that tag. + # names differ by exactly that tag (wherever the block-order tags put it). + wpe1 = specs[1].kernel_name().split("_") + wpe1.remove("wpe1") self.assertEqual( - specs[1].kernel_name(), - specs[2].kernel_name() + "_wpe1", + wpe1, + specs[2].kernel_name().split("_"), "kernel_name() differed beyond the waves_per_eu tag — test setup error", ) @@ -3467,12 +3652,20 @@ def _ir_sha(spec): ).hexdigest() def test_runtime_shape_specs_sharing_a_key_lower_to_identical_ir(self): - """Shapes that collapse to one cache key must emit one kernel.""" + """Shapes that collapse to one cache key must emit one kernel, for every + nonpersist decode (a reversed order must read NQB from the param).""" + from kernels.common.attention_dense_decode import NONPERSIST_DECODES + + for decode in ("auto", *NONPERSIST_DECODES): + with self.subTest(nonpersist_decode=decode): + self._check_runtime_shape_collision(nonpersist_decode=decode) + + def _check_runtime_shape_collision(self, **over): from dataclasses import replace from kernels.common.attention_dense_spec import attention_dense_cache_key from kernels.gfx950.attention_dense import AttentionDenseSpec - base = AttentionDenseSpec(**self._BASE_KWARGS) + base = AttentionDenseSpec(**self._BASE_KWARGS, **over) self.assertTrue( base.runtime_shape, "test setup error: the base spec is not on the runtime-shape path", @@ -3504,6 +3697,35 @@ def test_runtime_shape_specs_sharing_a_key_lower_to_identical_ir(self): "AttentionDenseSpec.runtime_param_fields so it splits the key again.", ) + def test_shape_dependent_auto_order_splits_the_key(self): + """The auto non-persistent order reads batch/seqlen (hq_minor_swz for + large causal MHA), so runtime-shape specs that resolve to different + orders must not share a key; the same order still shares one kernel.""" + from dataclasses import replace + + from kernels.common.attention_dense_spec import attention_dense_cache_key + from kernels.gfx950.attention_dense import AttentionDenseSpec + + base = AttentionDenseSpec(**{**self._BASE_KWARGS, "num_kv_heads": 32}) + small = replace(base, batch=1, seqlen_q=512, seqlen_kv=512) + large = replace(base, batch=16, seqlen_q=4096, seqlen_kv=4096) + larger = replace(base, batch=32, seqlen_q=4096, seqlen_kv=4096) + self.assertTrue(small.runtime_shape and large.runtime_shape) + self.assertNotEqual( + small.resolved_nonpersist_decode, large.resolved_nonpersist_decode + ) + self.assertEqual( + large.resolved_nonpersist_decode, larger.resolved_nonpersist_decode + ) + + def key(spec): + return attention_dense_cache_key(spec, arch="gfx950") + + self.assertNotEqual(key(small), key(large)) + self.assertEqual(key(large), key(larger)) + self.assertEqual(large.kernel_name(), larger.kernel_name()) + self.assertEqual(self._ir_sha(large), self._ir_sha(larger)) + def test_baked_shape_specs_split_both_key_and_ir(self): """Control: off the runtime path, each shape keeps its own key and IR. @@ -3621,11 +3843,19 @@ def _ir_sha(spec): def test_runtime_shape_specs_sharing_a_key_lower_to_identical_ir(self): """Shapes that collapse to one cache key must emit one kernel -- and, - on gfx942, one symbol name.""" + on gfx942, one symbol name -- for every nonpersist decode (a reversed + order must read NQB from the param).""" + from kernels.common.attention_dense_decode import NONPERSIST_DECODES + + for decode in ("auto", *NONPERSIST_DECODES): + with self.subTest(nonpersist_decode=decode): + self._check_runtime_shape_collision(nonpersist_decode=decode) + + def _check_runtime_shape_collision(self, **over): from dataclasses import replace from kernels.common.attention_dense_spec import attention_dense_cache_key - base = self._spec(**self._BASE_KWARGS) + base = self._spec(**self._BASE_KWARGS, **over) self.assertTrue( base.runtime_shape, "test setup error: the base gfx942 spec is not on the runtime-shape path", @@ -3680,6 +3910,34 @@ def test_runtime_shape_specs_sharing_a_key_lower_to_identical_ir(self): f"runtime-shape gfx942 name still carries {tok!r}: {name}", ) + def test_shape_dependent_auto_order_splits_the_key(self): + """The auto non-persistent order reads batch/seqlen (hq_minor_swz for + large causal MHA), so runtime-shape specs that resolve to different + orders must not share a key; the same order still shares one kernel.""" + from dataclasses import replace + + from kernels.common.attention_dense_spec import attention_dense_cache_key + + base = self._spec(**{**self._BASE_KWARGS, "num_kv_heads": 32}) + small = replace(base, batch=1, seqlen_q=512, seqlen_kv=512) + large = replace(base, batch=16, seqlen_q=4096, seqlen_kv=4096) + larger = replace(base, batch=32, seqlen_q=4096, seqlen_kv=4096) + self.assertTrue(small.runtime_shape and large.runtime_shape) + self.assertNotEqual( + small.resolved_nonpersist_decode, large.resolved_nonpersist_decode + ) + self.assertEqual( + large.resolved_nonpersist_decode, larger.resolved_nonpersist_decode + ) + + def key(spec): + return attention_dense_cache_key(spec, arch="gfx942") + + self.assertNotEqual(key(small), key(large)) + self.assertEqual(key(large), key(larger)) + self.assertEqual(large.kernel_name(), larger.kernel_name()) + self.assertEqual(self._ir_sha(large), self._ir_sha(larger)) + def test_grid_still_varies_per_shape_under_one_cache_key(self): """One binary, but a fresh launch grid for every shape. @@ -3729,10 +3987,18 @@ def test_grid_still_varies_per_shape_under_one_cache_key(self): # under test, or spec construction rejects it. Each maps to the name token its # order adds (qb_major adds none). _PERSIST_VARIANTS = { + "bt_hkv_minor": ( + dict(persistent=True, num_persistent=64, persist_decode="bt_hkv_minor"), + "_bthkvmin", + ), "hkv_major": ( dict(persistent=True, num_persistent=64, persist_decode="hkv_major"), "_hkvmaj", ), + "hq_minor_swz": ( + dict(persistent=True, num_persistent=64, persist_decode="hq_minor_swz"), + "_hqminswz", + ), "qb_major": ( dict(persistent=True, num_persistent=64, persist_decode="qb_major"), None, @@ -3925,9 +4191,10 @@ def test_auto_persist_decode_keys_on_the_resolved_order(self): base = self._spec(**self._BASE_KWARGS, persistent=True, num_persistent=64) self.assertEqual(base.persist_decode, "auto") - # gqa=4, NP=64: hkv_major iff gqa*nqb*B >= 2*NP (nqb = Sq/256 here). + # Batch below the XCD count: bt_hkv_minor. Else hkv_major iff + # gqa*nqb*B >= 2*NP (nqb = Sq/256 here). shapes = { - "qb_major": ((1, 512, 512), (1, 1024, 1024)), + "bt_hkv_minor": ((1, 512, 512), (1, 1024, 1024)), "hkv_major": ((8, 2048, 2048), (64, 4096, 8192)), } @@ -3961,9 +4228,9 @@ def artifact(spec): for i, part in enumerate(("cache key", "IR", "kernel name")): self.assertNotEqual( - by_order["qb_major"][i], + by_order["bt_hkv_minor"][i], by_order["hkv_major"][i], - f"auto resolving to qb_major and hkv_major share one {part}; the " + f"auto resolving to bt_hkv_minor and hkv_major share one {part}; the " "cache would serve one decode body to the other order's shapes", ) @@ -4017,6 +4284,42 @@ def artifact(spec): "baked into the prune and must split identity", ) + def test_window_balancing_is_runtime_for_every_decode(self): + """A window at least as long as the keys counts as plain causal, so the + decodes' causal balancing (longest-first, fold, pair fold) depends on the + launch's seqlen_kv. On the runtime path that is a select on a runtime flag, + not a compile-time branch: every pinned decode, on both grids, keeps ONE + (key, IR, name) across shapes on both sides of the window.""" + from dataclasses import replace + from kernels.common.attention_dense_decode import NONPERSIST_DECODES + from kernels.common.attention_dense_spec import attention_dense_cache_key + from kernels.gfx942.attention_dense import Gfx942AttentionDenseSpec + + window = 1024 + # (batch, seqlen): window >= keys for the first two, < keys for the rest. + shapes = ((2, 512), (1, 1024), (1, 2048), (3, 4096)) + grids = [dict(nonpersist_decode=n) for n in sorted(NONPERSIST_DECODES)] + [ + dict(persistent=True, num_persistent=64, persist_decode=n) + for n in sorted( + Gfx942AttentionDenseSpec.supported_persist_decodes() - {"auto"} + ) + ] + for kw in grids: + base = self._spec(**self._BASE_KWARGS, sliding_window=window, **kw) + arts = set() + for bt, sq in shapes: + with self.subTest(**kw, shape=(bt, sq)): + spec = replace(base, batch=bt, seqlen_q=sq, seqlen_kv=sq) + self.assertTrue(spec.runtime_shape) + arts.add( + ( + attention_dense_cache_key(spec, arch="gfx942"), + self._ir_sha(spec), + spec.kernel_name(), + ) + ) + self.assertEqual(len(arts), 1, f"{kw}: {len(arts)} distinct artifacts") + def test_bf16_d128_default_is_on_the_runtime_path(self): """bf16 D128 at the exp2_fast tri-state default is one binary per shape. @@ -4086,7 +4389,7 @@ def spec_for(batch, seqlen, nhead_q, nhead_k): "gfx942_dense", ).kernel_spec - specs = [spec_for(8, 8192, 32, 8), spec_for(2, 16384, 64, 8)] + specs = [spec_for(8, 8192, 32, 8), spec_for(16, 4096, 64, 8)] if not all(s.persistent for s in specs): self.skipTest( "shipped gfx942 dispatch no longer selects persistent for these " @@ -4256,6 +4559,9 @@ class TestAttentionDenseGfx942RuntimeHeads(unittest.TestCase): # smallest config under test. gfx950 dense is the baked-heads control # (test_gfx950_baked_heads_split_key_ir_and_name). _PERSIST_VARIANTS = { + "bt_hkv_minor": dict( + persistent=True, num_persistent=64, persist_decode="bt_hkv_minor" + ), "hkv_major": dict( persistent=True, num_persistent=64, persist_decode="hkv_major" ), @@ -4694,7 +5000,8 @@ def test_gfx950_and_base_names_keep_the_head_tokens(self): self.assertEqual( AttentionDenseSpec(**self._BASE_KWARGS).kernel_name(), - "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq2048_sk2048_causal_lazyrs", + "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq2048_sk2048_causal_lazyrs" + "_npbthkvmin", ) gfx950 = Gfx950AttentionDenseSpec(**self._BASE_KWARGS) self.assertIn("_hq128_kv8_", gfx950.kernel_name()) @@ -4884,13 +5191,17 @@ def test_runtime_shape_collision_does_not_smuggle_an_oversized_shape_past(self): from kernels.common.attention_dense_spec import attention_dense_cache_key from kernels.gfx950.attention_dense import supports_attention_dense - small = self._spec(batch=1, num_query_heads=8, num_kv_heads=8) + # One pinned block order: auto would pick a different one per shape. + small = self._spec( + batch=1, num_query_heads=8, num_kv_heads=8, nonpersist_decode="bt_hkv_minor" + ) huge = self._spec( batch=64, seqlen_q=16384, seqlen_kv=16384, num_query_heads=8, num_kv_heads=8, + nonpersist_decode="bt_hkv_minor", ) self.assertTrue(small.runtime_shape and huge.runtime_shape) self.assertEqual( diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_decode.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_decode.py new file mode 100644 index 000000000000..a7a491c6a5d5 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_decode.py @@ -0,0 +1,402 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Host-side checks of the dense-attention decode classes. + +A GPU numerical test cannot tell two decodes apart: every decode is a bijection +onto the work space, so any of them computes the right answer. These tests run +each class's real ``grid`` and ``emit_decode`` on Python ints instead, so a +decode that silently covers the work space in a different order, drops or +duplicates items, or bakes a shape value into a runtime-shape body is caught on +the host. + +Runtime shapes are checked the way a runtime-shape binary runs: the decode is +emitted for one spec (the compile-time fields) and fed another launch's shape +(and, on the persistent grid, that launch's fast-division kernargs); it must give +the launched spec's baked mapping. The fake builder's comparisons return +``numpy.bool_``, not ``bool``, so a runtime flag takes the ``select`` path. +""" + +from __future__ import annotations + +import dataclasses +import itertools + +import numpy as np +import pytest + +from kernels.common.attention_dense_decode import ( + NONPERSIST_DECODES, + PERSIST_DECODES, + DecodeShape, +) +from kernels.gfx942.attention_dense import ( + Gfx942AttentionDenseSpec, + attention_dense_runtime_args, +) +from kernels.gfx950.attention_dense import ( + Gfx950AttentionDenseSpec, +) + + +def _sdiv(a, b): + """i32 ``sdiv``: truncates toward zero (an unselected path may divide a + negative value; Python's ``//`` would floor it).""" + q = abs(a) // abs(b) + return q if (a >= 0) == (b >= 0) else -q + + +class _IntBuilder: + """Evaluates the IR-builder integer ops the decodes emit on Python ints, with + i32 semantics. Comparisons return ``numpy.bool_`` so a runtime flag is not a + ``bool`` and :meth:`DecodeShape.branch` takes its ``select`` path.""" + + const_i32 = staticmethod(int) + add = staticmethod(lambda a, b: a + b) + sub = staticmethod(lambda a, b: a - b) + mul = staticmethod(lambda a, b: a * b) + div = staticmethod(_sdiv) + mod = staticmethod(lambda a, b: a - _sdiv(a, b) * b) + lshr = staticmethod(lambda a, b: (a & 0xFFFFFFFF) >> b) + umul_hi_i32 = staticmethod(lambda a, b: ((a & 0xFFFFFFFF) * (b & 0xFFFFFFFF)) >> 32) + smax = staticmethod(max) + land = staticmethod(lambda a, b: a & b) + cmp_lt = staticmethod(lambda a, b: np.bool_(a < b)) + cmp_ge = staticmethod(lambda a, b: np.bool_(a >= b)) + cmp_eq = staticmethod(lambda a, b: np.bool_(a == b)) + cmp_ne = staticmethod(lambda a, b: np.bool_(a != b)) + select = staticmethod(lambda c, a, b: a if c else b) + + +B_ = _IntBuilder() + +_SHAPE_FIELDS = ("batch", "seqlen_q", "seqlen_kv", "num_query_heads", "num_kv_heads") + + +def _baked(spec): + return DecodeShape(B_, spec) + + +def _runtime(built, launched, fastdiv=None, extra=None): + """``built``'s decode fed ``launched``'s shape, as one runtime-shape binary + serves another launch.""" + values = {f: getattr(launched, f) for f in _SHAPE_FIELDS} + values["gqa"] = launched.num_queries_per_kv + values.update(extra or {}) + return DecodeShape(B_, built, values, fastdiv=fastdiv) + + +def _fastdiv(launched, dividends): + """The kernel's fast division, from the launch's real host kernargs.""" + args = attention_dense_runtime_args(launched) + + def div(n, d): + dividends.append(n) + q = (B_.umul_hi_i32(n, args[f"{d}_magic"]) + n) >> args[f"{d}_shift"] + return q + + return div + + +# (Hq, Hkv, B, S): GQA, MHA, and a head count that is not a power of two. +_SHAPES = ((32, 8, 2, 2048), (64, 64, 1, 1024), (40, 8, 3, 512)) + + +def _spec(Hq, Hkv, B, S, causal, **over): + return Gfx950AttentionDenseSpec( + batch=B, + seqlen_q=S, + seqlen_kv=S, + num_query_heads=Hq, + num_kv_heads=Hkv, + head_size=128, + causal=causal, + dtype="fp16", + **over, + ) + + +def _block_linear(spec, qb, hq, bt, nqb): + """Expected hardware dispatch index of the CTA for (qb, hq, bt), per decode. + Written from each order's definition, independently of the decode code; + ``blk`` is the query block before any causal reversal.""" + Hq, Hkv, B = spec.num_query_heads, spec.num_kv_heads, spec.batch + gqa = Hq // Hkv + name = spec.nonpersist_decode + if name == "qb_minor": # digits, fastest first: Q, head, B + return qb + nqb * (hq + Hq * bt) + if name == "bt_hkv_minor": # B, V, G, Q; Q longest-first under causal + blk = nqb - 1 - qb if spec.causal else qb + hkv, hql = divmod(hq, gqa) + return bt + B * (hkv + Hkv * (hql + gqa * blk)) + if name == "hq_minor_swz": + # arXiv 2511.02132 Fig. 11 per batch, with batch the slowest digit: + # head = (w % M)*(Hq/M) + w // (M*nqb), block = (w % (M*nqb)) // M + # inverted to w for (hq, blk); query blocks longest-first under causal. + M = spec.chiplet_num_xcds + blk = nqb - 1 - qb if spec.causal else qb + a, c = divmod(hq, Hq // M) + return a + M * blk + M * nqb * c + Hq * nqb * bt + raise AssertionError(f"no expected order for nonpersist decode {name!r}") + + +def _decode_all(decode, spec, make_shape): + gx, gy, gz = decode.grid(spec) + out = {} + for bz, by, bx in itertools.product(range(gz), range(gy), range(gx)): + d = decode.emit_decode(make_shape(spec), bx, by, bz) + out[bx + gx * (by + gy * bz)] = (d.qb, d.hq, d.bt) + return out + + +# Baked; the gfx950 grid (only seqlen_q is a param); the gfx942 runtime grid. +_GRID_SHAPES = { + "baked": _baked, + "seqlen_q": lambda s: DecodeShape(B_, s, {"seqlen_q": s.seqlen_q}), + "runtime": lambda s: _runtime(s, s), +} + + +@pytest.mark.parametrize("name", sorted(NONPERSIST_DECODES)) +@pytest.mark.parametrize("causal", (True, False)) +@pytest.mark.parametrize("shape", _SHAPES) +@pytest.mark.parametrize("xcds", (8, 4)) +def test_nonpersist_decode_linearizes_to_its_documented_order( + name, causal, shape, xcds +): + decode = NONPERSIST_DECODES[name] + spec = _spec(*shape, causal, nonpersist_decode=name, chiplet_num_xcds=xcds) + nqb = (spec.seqlen_q + spec.block_m - 1) // spec.block_m + work = nqb * spec.num_query_heads * spec.batch + for make_shape in _GRID_SHAPES.values(): + items = _decode_all(decode, spec, make_shape) + assert sorted(items) == list(range(work)) + assert len(set(items.values())) == work, "decode is not a bijection" + for wid, (qb, hq, bt) in items.items(): + assert _block_linear(spec, qb, hq, bt, nqb) == wid, (name, wid) + + +# Launches one runtime-shape grid binary serves: other batch, seqlen and heads, +# windows shorter and longer than the keys. +_GRID_LAUNCHES = ( + (32, 8, 2, 4096, 0), + (64, 64, 1, 1024, 0), + (40, 8, 3, 512, 512), + (40, 8, 1, 2048, 512), + (16, 4, 2, 1024, 0), +) + + +@pytest.mark.parametrize("name", sorted(NONPERSIST_DECODES)) +@pytest.mark.parametrize("window", (0, 512)) +def test_nonpersist_decode_reads_the_shape_from_the_runtime_params(name, window): + """Under a runtime shape one binary serves every shape: the decode built for + one spec and fed another launch's shape must give that launch's baked + mapping, not the built spec's (with a window, balancing follows the launch's + keys).""" + decode = NONPERSIST_DECODES[name] + built = _spec(32, 8, 2, 2048, True, sliding_window=window, nonpersist_decode=name) + for Hq, Hkv, B, S, w in _GRID_LAUNCHES: + if w != window: + continue + launched = dataclasses.replace( + built, + batch=B, + seqlen_q=S, + seqlen_kv=S, + num_query_heads=Hq, + num_kv_heads=Hkv, + ) + gx, gy, gz = decode.grid(launched) + for bz, by, bx in itertools.product(range(gz), range(gy), range(gx)): + got = decode.emit_decode(_runtime(built, launched), bx, by, bz) + want = decode.emit_decode(_baked(launched), bx, by, bz) + assert got == want, (name, (Hq, Hkv, B, S), (bx, by, bz)) + + +def _persist_specs(): + """Every persist decode on every shape where it is legal, with and without + causal masking, plus the interleave variant of qb_major.""" + return list(_iter_persist_specs()) + + +def _iter_persist_specs(): + for name in sorted(Gfx950AttentionDenseSpec.supported_persist_decodes() - {"auto"}): + for shape, causal in itertools.product(_SHAPES, (True, False)): + Hq, Hkv, B, S = shape + nqb = S // 256 + num_persistent = { + "gqa_pair": nqb * Hkv * B, + "gqa_pair_2phase": nqb * Hq * B // 2, + }.get(name, 256) + for interleave in ((False, True) if name == "qb_major" else (False,)): + try: + spec = _spec( + *shape, + causal, + persistent=True, + num_persistent=num_persistent, + persist_decode=name, + interleave=interleave, + ) + except ValueError: + continue # not legal for this shape (e.g. gqa_pair on MHA) + yield pytest.param( + spec, + id=f"{name}-{Hq}x{Hkv}-B{B}-S{S}-" + f"{'causal' if causal else 'full'}" + f"{'-intl' if interleave else ''}", + ) + + +@pytest.mark.parametrize("spec", _persist_specs()) +def test_persist_decode_is_a_bijection(spec): + decode = PERSIST_DECODES[spec.persist_decode] + gqa = spec.num_queries_per_kv + nqb = (spec.seqlen_q + spec.block_m - 1) // spec.block_m + work = nqb * spec.num_query_heads * spec.batch + items = set() + for wi in range(work): + d = decode.emit_decode(_baked(spec), wi) + assert 0 <= d.qb < nqb and 0 <= d.hq < spec.num_query_heads + assert 0 <= d.bt < spec.batch + assert d.hkv is None or d.hkv == d.hq // gqa + items.add((d.qb, d.hq, d.bt)) + assert len(items) == work, "decode drops or duplicates work items" + + +def test_every_persist_decode_is_covered(): + covered = {p.values[0].persist_decode for p in _persist_specs()} + assert covered == set(PERSIST_DECODES) + + +def test_bt_hkv_minor_follows_its_documented_order(): + """``wi = ((blk*gqa + hql)*Hkv + hkv)*B + bt`` (the class docstring), checked + independently of the decode's Hq-first split of the head digits.""" + for Hq, Hkv, B, S in _SHAPES + ((48, 6, 2, 1024),): + spec = _spec( + Hq, + Hkv, + B, + S, + False, + persistent=True, + num_persistent=256, + persist_decode="bt_hkv_minor", + ) + gqa, nqb = Hq // Hkv, S // 256 + for wi in range(nqb * Hq * B): + d = PERSIST_DECODES["bt_hkv_minor"].emit_decode(_baked(spec), wi) + hkv, hql = d.hq // gqa, d.hq % gqa + assert d.hkv == hkv + assert ((d.qb * gqa + hql) * Hkv + hkv) * B + d.bt == wi + + +def _gfx942_persist(Hq, Hkv, B, S, name, causal=True, window=0, np_=256, **over): + return Gfx942AttentionDenseSpec( + batch=B, + seqlen_q=S, + seqlen_kv=S, + num_query_heads=Hq, + num_kv_heads=Hkv, + head_size=128, + causal=causal, + dtype="fp16", + sliding_window=window, + persistent=True, + num_persistent=np_, + persist_decode=name, + **over, + ) + + +# (Hq, Hkv, B, S) launches of one persistent binary: NQB even and odd, W below, +# near and far above NP, GQA and MHA, a head count that is not a power of two. +_PERSIST_LAUNCHES = ( + (32, 8, 1, 512), + (32, 8, 2, 2048), + (40, 8, 1, 768), + (64, 64, 1, 1024), + (16, 4, 3, 1280), + (48, 8, 4, 4096), +) + + +@pytest.mark.parametrize( + "name", sorted(Gfx942AttentionDenseSpec.supported_persist_decodes() - {"auto"}) +) +@pytest.mark.parametrize("causal,window", ((True, 0), (False, 0), (True, 1024))) +@pytest.mark.parametrize("interleave", (False, True)) +def test_runtime_persist_decode_matches_the_launched_baked_decode( + name, causal, window, interleave +): + """One gfx942 persistent binary (built for the first launch) serves every + launch through the shape params and that launch's fast-division kernargs: it + must decode exactly as the launch's own baked decode, every fast-division + dividend must stay within the decode's bound, and the work must be covered + once.""" + if interleave and name != "qb_major": + pytest.skip("interleave applies to qb_major only") + decode = PERSIST_DECODES[name] + built = None + for Hq, Hkv, B, S in _PERSIST_LAUNCHES: + try: + launched = _gfx942_persist( + Hq, Hkv, B, S, name, causal, window, interleave=interleave + ) + except ValueError: + continue # not legal for this launch (e.g. hq_minor_swz, odd NQB) + built = built or launched + nqb = S // launched.block_m + work = nqb * Hq * B + dividends = [] + shape = _runtime( + built, launched, _fastdiv(launched, dividends), {"nqb": nqb, "w": work} + ) + got = [decode.emit_decode(shape, wi) for wi in range(work)] + want = [decode.emit_decode(_baked(launched), wi) for wi in range(work)] + assert got == want, (name, (Hq, Hkv, B, S)) + assert len({(d.qb, d.hq, d.bt) for d in got}) == work + assert max(dividends, default=0) <= decode.fastdiv_dividend_bound(launched) + assert built is not None, f"no legal launch for {name}" + + +# Decodes that require causal masking, so have no window-free causal-free order. +_CAUSAL_ONLY = {"gqa_pair", "gqa_pair_2phase"} + + +@pytest.mark.parametrize("persistent", (True, False)) +def test_sliding_window_orders_blocks_as_without_causal_masking(persistent): + """A sliding window gives every query block about the same work, so no + decode balances there (longest-first, fold, pair fold): each orders blocks + as without causal masking. For qb_major and hkv_major that is the previous + order (ascending, plain fold). A window at least as long as the keys cuts + nothing off, so there the orders are the causal ones. The window-free causal + specs are the control: some decode must still move blocks there.""" + Hq, Hkv, B, S = 32, 8, 2, 4096 + nqb = S // 256 + registry = PERSIST_DECODES if persistent else NONPERSIST_DECODES + moved = set() + for name in sorted(set(registry) - _CAUSAL_ONLY): + kw = ( + dict(persistent=True, num_persistent=256, persist_decode=name) + if persistent + else dict(nonpersist_decode=name) + ) + swa = _spec(Hq, Hkv, B, S, True, sliding_window=512, **kw) + plain = _spec(Hq, Hkv, B, S, False, **kw) + full = _spec(Hq, Hkv, B, S, True, **kw) + long = _spec(Hq, Hkv, B, S, True, sliding_window=S, **kw) + decode = registry[name] + if persistent: + orders = [ + [decode.emit_decode(_baked(s), wi) for wi in range(nqb * Hq * B)] + for s in (swa, plain, full, long) + ] + else: + orders = [_decode_all(decode, s, _baked) for s in (swa, plain, full, long)] + assert orders[0] == orders[1], name + assert orders[3] == orders[2], name + if orders[2] != orders[1]: + moved.add(name) + assert moved, "without a window some decode should balance causal work" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py index df84d57fc51e..9818d76d292c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py @@ -136,6 +136,10 @@ def test_bottom_right_field_is_keyword_only_without_shifting_concrete_signatures params["causal_bottom_right"].kind is inspect.Parameter.KEYWORD_ONLY ), spec_type.__name__ assert params["causal_bottom_right"].default is False + for kw_only in ("nonpersist_decode", "chiplet_num_xcds"): + assert ( + params[kw_only].kind is inspect.Parameter.KEYWORD_ONLY + ), spec_type.__name__ positionals = tuple( name for name, param in params.items() @@ -305,7 +309,9 @@ def test_kernel_name_pins_the_two_shipped_cache_collisions(): "persistent": (True, False), "num_persistent": (152, 304), "interleave": (True, False), - "persist_decode": ("qb_major", "hkv_major"), + "persist_decode": ("qb_major", "hkv_major", "bt_hkv_minor"), + "nonpersist_decode": ("bt_hkv_minor", "qb_minor", "hq_minor_swz"), + "chiplet_num_xcds": (1, 4), "lazy_rescale": (False, True), "use_sinks": (True,), } @@ -363,6 +369,21 @@ def test_kernel_name_pins_the_two_shipped_cache_collisions(): _RUNTIME_PARAM_FIELDS = frozenset(_spec().runtime_param_fields) +# Every persistent decode order the gfx942 body emits. +_PERSIST_DECODES = sorted( + Gfx942AttentionDenseSpec.supported_persist_decodes() - {"auto"} +) + + +def _decode_kw(decode: str) -> dict: + """Pin ``decode``; hq_minor_swz also needs num_persistent divisible by the XCD + count (the base's 228 is not).""" + kw = dict(persist_decode=decode) + if decode == "hq_minor_swz": + kw["num_persistent"] = 224 + return kw + + def _persistent_spec(**kw) -> Gfx942AttentionDenseSpec: """The persistent injectivity base, with overrides. @@ -560,7 +581,7 @@ def test_kernel_name_covers_every_baked_parameter(field): ) -@pytest.mark.parametrize("decode", ["qb_major", "hkv_major"]) +@pytest.mark.parametrize("decode", _PERSIST_DECODES) def test_persistent_grid_shares_one_kernel_across_batches(decode): """The persistent grid reads batch from a kernel param too: every batch is ONE name, ONE body and ONE cache key, with no ``_b{N}`` token. @@ -576,7 +597,7 @@ def test_persistent_grid_shares_one_kernel_across_batches(decode): The decode is pinned: ``"auto"`` would flip order with batch, which is a separate identity (:func:`test_auto_persist_decode_keys_on_the_resolved_order`). """ - specs = [_persistent_spec(batch=b, persist_decode=decode) for b in (1, 2, 4, 8)] + specs = [_persistent_spec(batch=b, **_decode_kw(decode)) for b in (1, 2, 4, 8)] assert all(s.runtime_shape for s in specs) names = {gfx942_kernel_name(s) for s in specs} assert len(names) == 1, f"persistent batches split the name: {names}" @@ -614,20 +635,27 @@ def test_build_emits_one_symbol_across_batches(make): _HEAD_CONFIGS = ((128, 8), (40, 8), (28, 4), (32, 32)) -@pytest.mark.parametrize("decode", ["qb_major", "hkv_major"]) +@pytest.mark.parametrize("decode", _PERSIST_DECODES) def test_persistent_grid_shares_one_kernel_across_head_configs(decode): """The persistent work decode divides by the runtime head counts (through the fast-division kernargs), so every head config -- pow2 and non-pow2 GQA, MHA -- is ONE name, ONE body and ONE cache key per decode order, with no hq/kv token. The head counts still have to reach the launch: each config gets distinct - kernarg values, ``gqa`` magic/shift included. That the IR probe can see baked + kernarg values, fast-division pairs included. That the IR probe can see baked head counts is the gfx950 control (:func:`test_ir_body_probe_detects_a_baked_shape_on_the_gfx950_control`), so the - single-body assertion here is not vacuous.""" - specs = [ - _persistent_spec(num_query_heads=hq, num_kv_heads=hkv, persist_decode=decode) + single-body assertion here is not vacuous. hq_minor_swz needs Hq divisible by + the XCD count, so it skips the configs that are not.""" + configs = [ + (hq, hkv) for hq, hkv in _HEAD_CONFIGS + if decode != "hq_minor_swz" or hq % 8 == 0 + ] + assert len(configs) >= 3 + specs = [ + _persistent_spec(num_query_heads=hq, num_kv_heads=hkv, **_decode_kw(decode)) + for hq, hkv in configs ] for s in specs: ok, why = supports_attention_dense(s, arch="gfx942") @@ -637,10 +665,14 @@ def test_persistent_grid_shares_one_kernel_across_head_configs(decode): assert not re.search(r"_(hq|kv)\d+", names.pop()) assert len({attention_dense_cache_key(s, arch="gfx942") for s in specs}) == 1 assert len({_ir_body_sha(s) for s in specs}) == 1 - # Every config has a distinct gqa (16, 5, 7, 1), so its fast-division pair must - # differ too; the exact values are pinned by the kernarg-ABI test below. + # Every config has a distinct Hq (128, 40, 28, 32) and gqa (16, 5, 7, 1), and + # each decode divides by one of them, so the fast-division pairs must differ + # too; the exact values are pinned by the kernarg-ABI test below. args = [attention_dense_runtime_args(s) for s in specs] - assert len({(a["gqa_magic"], a["gqa_shift"]) for a in args}) == len(specs) + pairs = [ + tuple(v for k, v in a.items() if k.endswith(("_magic", "_shift"))) for a in args + ] + assert len(set(pairs)) == len(specs) def test_kernel_name_drops_heads_on_the_runtime_shape_grid(): @@ -1043,8 +1075,8 @@ def test_gfx942_auto_decode_cannot_leak_to_gqa_pair(): num_persistent=304, persist_decode="auto", ) - assert spec.resolved_persist_decode == "hkv_major" - assert "hkvmaj" in spec.kernel_name() + assert spec.resolved_persist_decode == "bt_hkv_minor" + assert "bthkvmin" in spec.kernel_name() assert "gqapair" not in spec.kernel_name() @@ -1240,7 +1272,7 @@ def test_supports_true_implies_build_succeeds(row): # launch geometry # --------------------------------------------------------------------------- # def test_grid_and_block_geometry(): - s = _spec(seqlen_q=2048, num_query_heads=128, batch=1) + s = _spec(seqlen_q=2048, num_query_heads=128, batch=1, nonpersist_decode="qb_minor") assert attention_dense_grid(s) == (2048 // _BLOCK_M, 128, 1) assert attention_dense_block(s) == (s.num_waves * 64, 1, 1) assert attention_dense_block(s) == (_EXPECTED_WORKGROUP_SIZE, 1, 1) @@ -1249,7 +1281,12 @@ def test_grid_and_block_geometry(): def test_grid_covers_every_query_row_and_head(): for sq, hq, batch in ((2048, 128, 1), (512, 16, 4), (4096, 40, 2)): s = _spec( - seqlen_q=sq, seqlen_kv=sq, num_query_heads=hq, num_kv_heads=8, batch=batch + seqlen_q=sq, + seqlen_kv=sq, + num_query_heads=hq, + num_kv_heads=8, + batch=batch, + nonpersist_decode="qb_minor", ) nqb, ghq, gb = attention_dense_grid(s) assert ( @@ -2033,17 +2070,16 @@ def test_codegen_knob_rejections(overrides, reason): # --------------------------------------------------------------------------- # @pytest.mark.parametrize("dtype", ["fp16", "bf16"]) @pytest.mark.parametrize("d", [64, 128]) -@pytest.mark.parametrize("decode", ["qb_major", "hkv_major"]) +@pytest.mark.parametrize("decode", _PERSIST_DECODES) def test_persistent_builds_for_both_decodes(dtype, d, decode): - """P4: the persistent grid-stride body builds for every in-scope config and both - work-decode orders. The decode is a build-time (Python) branch, so both must emit; + """P4: the persistent grid-stride body builds for every in-scope config and every + work-decode order. The decode is a build-time (Python) branch, so each must emit; GPU numeric correctness (the decode bijection) is verified in the live cohort.""" spec = _spec( head_size=d, dtype=dtype, persistent=True, - num_persistent=228, - persist_decode=decode, + **{"num_persistent": 228, **_decode_kw(decode)}, ) ok, why = supports_attention_dense(spec, arch="gfx942") assert ok, why @@ -2088,7 +2124,9 @@ def test_persistent_and_default_share_one_inner_body(): # The divisors each resolved persistent decode order divides by in the grid-stride # loop, in kernarg order. Restated (not imported) on purpose: this IS the ABI. _EXPECTED_FASTDIV_DIVISORS = { + "bt_hkv_minor": ("batch", "num_query_heads"), "hkv_major": ("batch", "gqa", "nqb"), + "hq_minor_swz": ("nqb", "num_query_heads", "gqa"), "qb_major": ("batch", "num_query_heads", "gqa"), } @@ -2099,7 +2137,7 @@ def _fastdiv(n: int, magic_i32: int, shift: int) -> int: return (((n * (magic_i32 & 0xFFFFFFFF)) >> 32) + n) >> shift -@pytest.mark.parametrize("decode", ["qb_major", "hkv_major"]) +@pytest.mark.parametrize("decode", _PERSIST_DECODES) def test_persistent_kernarg_abi_carries_the_fastdiv_pairs(decode): """The persistent kernarg ABI is the five shape params, then one ``_magic`` / ``_shift`` i32 pair per divisor of the RESOLVED @@ -2112,17 +2150,19 @@ def test_persistent_kernarg_abi_carries_the_fastdiv_pairs(decode): packed as its two's-complement i32), and -- through the host model of the in-kernel divide -- exact for dividends up to just under the 2**30 bound. Non-pow2 divisors throughout (B=3, NQB=13, Hq=40, gqa=5), so the magic is not the trivial - pow2 one. + pow2 one; hq_minor_swz pairs query blocks under causal masking, so it takes + NQB=14. """ from rocke.helpers.magic_division import calculate_magic_numbers + nqb = 14 if decode == "hq_minor_swz" else 13 spec = _persistent_spec( batch=3, - seqlen_q=13 * _BLOCK_M, + seqlen_q=nqb * _BLOCK_M, seqlen_kv=4096, num_query_heads=40, num_kv_heads=8, - persist_decode=decode, + **_decode_kw(decode), ) assert supports_attention_dense(spec, arch="gfx942")[0] divisors = _EXPECTED_FASTDIV_DIVISORS[decode] @@ -2139,7 +2179,7 @@ def test_persistent_kernarg_abi_carries_the_fastdiv_pairs(decode): args = attention_dense_runtime_args(spec) assert tuple(args) == spec.runtime_kernarg_fields - value = {"batch": 3, "num_query_heads": 40, "gqa": 5, "nqb": 13} + value = {"batch": 3, "num_query_heads": 40, "gqa": 5, "nqb": nqb} for d in divisors: magic, shift = calculate_magic_numbers(value[d]) assert -(1 << 31) <= args[f"{d}_magic"] < (1 << 31), "magic must fit an i32" @@ -2166,16 +2206,22 @@ def test_auto_persist_decode_keys_on_the_resolved_order(): the raw ``"auto"`` would let two shapes that resolve differently share one slot and serve one decode's binary for the other. - gqa=4, NQB=16, NP=228: per-kv-head work 64*B crosses 2*NP between B=4 and B=8. - Same-order shapes share key, name and body; cross-order shapes differ in all - three; and an ``"auto"`` that resolves to an order is the SAME kernel as that - order spelled explicitly. + gqa=4, NQB=16, NP=228: batch below the XCD count resolves to bt_hkv_minor, and + at B=8 the per-kv-head work 64*B is past 2*NP, so hkv_major. Same-order shapes + share key, name and body; cross-order shapes differ in all three; and an + ``"auto"`` that resolves to an order is the SAME kernel as that order spelled + explicitly. """ by_batch = { b: _persistent_spec(batch=b, persist_decode="auto") for b in (1, 2, 8, 16) } order = {b: s.resolved_persist_decode for b, s in by_batch.items()} - assert order == {1: "qb_major", 2: "qb_major", 8: "hkv_major", 16: "hkv_major"} + assert order == { + 1: "bt_hkv_minor", + 2: "bt_hkv_minor", + 8: "hkv_major", + 16: "hkv_major", + } def ident(spec): return ( @@ -2190,11 +2236,11 @@ def ident(spec): lo, hi = ids[1], ids[8] assert all(a != b for a, b in zip(lo, hi)), "key, name and body must all split" assert "hkvmaj" in hi[1] and "hkvmaj" not in lo[1] - for b, resolved in ((1, "qb_major"), (8, "hkv_major")): + for b, resolved in ((1, "bt_hkv_minor"), (8, "hkv_major")): explicit = dataclasses.replace(by_batch[b], persist_decode=resolved) assert ident(explicit) == ids[b] # The fast-division ABI follows the resolved order as well. - assert by_batch[1].fastdiv_divisors == _EXPECTED_FASTDIV_DIVISORS["qb_major"] + assert by_batch[1].fastdiv_divisors == _EXPECTED_FASTDIV_DIVISORS["bt_hkv_minor"] assert by_batch[8].fastdiv_divisors == _EXPECTED_FASTDIV_DIVISORS["hkv_major"] @@ -2219,11 +2265,11 @@ def test_persistent_runtime_args_reject_work_past_the_fastdiv_bound(): assert attention_dense_runtime_args(default)["batch"] == at_limit -def test_dispatch_persistent_auto_turns_on_for_large_sq_only(): - """P4 dispatch: the default spec turns persistent on once the work - (nqb*Hq*B) fills the gfx942 persistent grid (num_persistent defaulted to 304), and - stays off for small Sq. The ``persistent`` knob is honored; gfx950 keeps its 256 default - and is otherwise untouched.""" +def test_dispatch_persistent_auto_turns_on_for_d128_only(): + """P4 dispatch: the default spec turns persistent on at D128 for any Sq and + batch (num_persistent defaulted to 304), and stays off at D64. The + ``persistent`` knob is honored; gfx950 keeps its 256 default and is otherwise + untouched.""" from dispatch.attention import AttentionRequest from dispatch.attention import attention_tuning_spec, tuning_spec_with_knobs @@ -2235,15 +2281,15 @@ def _dense_spec(req, **knobs): def _dense_spec_gfx950(req): return attention_tuning_spec(req, "gfx950_dense_grid").kernel_spec - def _req(sq, arch): + def _req(sq, arch, d=128): return AttentionRequest( batch=1, nhead_q=16, nhead_k=4, seqlen_q=sq, seqlen_k=sq, - hdim_q=128, - hdim_v=128, + hdim_q=d, + hdim_v=d, arch=arch, mask_type=1, dtype="fp16", @@ -2251,10 +2297,11 @@ def _req(sq, arch): # gfx942 num_persistent defaulted to the 304-CU part's CU count. assert _dense_spec(_req(8192, "gfx942")).num_persistent == 304 - # auto: on for large Sq (nqb*Hq = 32*16 = 512 >= 304), off for small. + # default: on at D128 whatever the Sq, off at D64. + assert _dense_spec(_req(2048, "gfx942")).persistent is True assert _dense_spec(_req(8192, "gfx942")).persistent is True - assert _dense_spec(_req(2048, "gfx942")).persistent is False # 8*16 = 128 < 304 - # explicit modes honored. + assert _dense_spec(_req(8192, "gfx942", d=64)).persistent is False + # the persistent knob is honored. assert _dense_spec(_req(8192, "gfx942"), persistent=False).persistent is False assert _dense_spec(_req(256, "gfx942"), persistent=True).persistent is True # gfx950 untouched: keeps the 256 default (not the gfx942 304 override). diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942_numeric.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942_numeric.py index 072a5826df14..813a033e587a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942_numeric.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942_numeric.py @@ -881,23 +881,57 @@ def test_one_sliding_window_binary_serves_every_shape_and_head_config(): (2, 1024, 15, 5), ) -# (dtype, persist_decode, interleave, sliding_window). Both decode orders pinned -# explicitly: "auto" resolves per shape, which would split the set across two keys. -# interleave only changes the qb_major body, so it rides that order alone. +# hq_minor_swz needs Hq divisible by the XCD count and an even NQB under causal +# masking, so it gets its own set, straddling NP the same way: +# W=80 < NP; W=320 and W=608 (pair fold, even NQB); MHA W=384; W=1536, W % NP = 16 +_SWZ_CONFIGS = ( + (1, 512, 40, 8), + (2, 1024, 40, 8), + (2, 1536, 32, 32), + (19, 512, 16, 4), + (3, 2048, 64, 8), +) + +# (dtype, persist_decode, interleave, sliding_window, configs). Every decode order +# pinned explicitly: "auto" resolves per shape, which would split the set across +# keys. interleave only changes the qb_major body, so it rides that order alone. A +# window of 1024 is at least as long as the keys for some configs and shorter for +# others, so one binary balances causal work for some launches and not for others. _PERSIST_ONE_BINARY = [ - pytest.param("fp16", "hkv_major", False, 0, id="hkv_major"), - pytest.param("fp16", "qb_major", False, 0, id="qb_major"), - pytest.param("fp16", "qb_major", True, 0, id="qb_major-interleave"), - pytest.param("bf16", "hkv_major", False, 128, id="hkv_major-swa128"), - pytest.param("bf16", "qb_major", False, 128, id="qb_major-swa128"), + pytest.param("fp16", "hkv_major", False, 0, _PERSIST_CONFIGS, id="hkv_major"), + pytest.param("fp16", "qb_major", False, 0, _PERSIST_CONFIGS, id="qb_major"), + pytest.param( + "fp16", "qb_major", True, 0, _PERSIST_CONFIGS, id="qb_major-interleave" + ), + pytest.param("fp16", "bt_hkv_minor", False, 0, _PERSIST_CONFIGS, id="bt_hkv_minor"), + pytest.param("fp16", "hq_minor_swz", False, 0, _SWZ_CONFIGS, id="hq_minor_swz"), + pytest.param( + "bf16", "hkv_major", False, 128, _PERSIST_CONFIGS, id="hkv_major-swa128" + ), + pytest.param( + "bf16", "qb_major", False, 128, _PERSIST_CONFIGS, id="qb_major-swa128" + ), + *( + pytest.param( + "bf16", + name, + False, + 1024, + _SWZ_CONFIGS if name == "hq_minor_swz" else _PERSIST_CONFIGS, + id=f"{name}-swa1024", + ) + for name in ("qb_major", "hkv_major", "bt_hkv_minor", "hq_minor_swz") + ), ] @requires_gfx942_gpu @pytest.mark.gpu -@pytest.mark.parametrize("dtype,persist_decode,interleave,window", _PERSIST_ONE_BINARY) +@pytest.mark.parametrize( + "dtype,persist_decode,interleave,window,configs", _PERSIST_ONE_BINARY +) def test_one_persistent_binary_serves_every_shape_and_head_config( - dtype, persist_decode, interleave, window + dtype, persist_decode, interleave, window, configs ): """The persistent-grid twin of the one-binary tests above. @@ -939,7 +973,7 @@ def test_one_persistent_binary_serves_every_shape_and_head_config( ), interleave=interleave, ) - for b, s, hq, hkv in _PERSIST_CONFIGS + for b, s, hq, hkv in configs ] # Preconditions: persistent, on the runtime path, the pinned order, one key -- @@ -956,13 +990,13 @@ def test_one_persistent_binary_serves_every_shape_and_head_config( keys = [attention_dense_cache_key(s, arch="gfx942") for s in specs] assert ( len(set(keys)) == 1 - ), f"persistent configs {_PERSIST_CONFIGS} did not share a cache key" + ), f"persistent configs {configs} did not share a cache key" _DENSE_LAUNCHER_CACHE.pop(keys[0], None) before = set(_DENSE_LAUNCHER_CACHE) launchers = [] - for (B, S, hq, hkv), spec in zip(_PERSIST_CONFIGS, specs): + for (B, S, hq, hkv), spec in zip(configs, specs): torch.manual_seed(0) q = torch.randn(B, S, hq, d, device="cuda", dtype=tdt) k = torch.randn(B, S, hkv, d, device="cuda", dtype=tdt) @@ -999,7 +1033,7 @@ def test_one_persistent_binary_serves_every_shape_and_head_config( "longer keyed by attention_dense_cache_key and this test is blind" ) assert all(lau is launchers[0] for lau in launchers), ( - f"persistent configs {_PERSIST_CONFIGS} share a cache key but were served " + f"persistent configs {configs} share a cache key but were served " "by DIFFERENT launcher objects -- the persistent kernel recompiled per config" ) assert set(_DENSE_LAUNCHER_CACHE) - before == {keys[0]}, ( diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_ir_cpp_parity.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_ir_cpp_parity.py index a84678e0d536..b685bbbb0e9e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_ir_cpp_parity.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_ir_cpp_parity.py @@ -133,7 +133,7 @@ def test_attention_ir_cpp_python_byte_identity(): ) # gfx942 has one shared epilogue for both grids; cover each grid with and - # without keyless rows (the harness has no gfx942 dense builder). + # without keyless rows, with LSE on (the harness's gfx942 cases run without it). from kernels.gfx942.attention_dense import ( Gfx942AttentionDenseSpec, build_attention_dense as build_gfx942_dense, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt index 717a89679250..8b6f2a73293c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt +++ b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt @@ -568,10 +568,12 @@ if(ROCKE_INSTALL_TESTS) " TIMEOUT 180)\n") # Library attention CPU tests (no GPU required): gfx950 LDS budget model, - # spec-wrapper parity gate, and the 3D decode split rule and K/V cache policy. + # dense block-order decodes, spec-wrapper parity gate, and the 3D decode split + # rule and K/V cache policy. # File-list style (same pattern as conv_pytest) so PYTHONPATH resolves rocke/ # correctly from the installed tree. set(_ROCKE_LIBRARY_ATTENTION_CPU_TEST_RELPATHS + "tests/library/tests/test_attention_dense_decode.py" "tests/library/tests/test_attention_dense_gfx950_lds.py" "tests/library/tests/test_attention_spec_predicate_parity.py" "tests/library/tests/test_attention_tiled_3d_decode_split.py" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/attention_thread_block_mapping_case_study.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/attention_thread_block_mapping_case_study.md new file mode 100644 index 000000000000..ba8a23378086 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/attention_thread_block_mapping_case_study.md @@ -0,0 +1,584 @@ + + +# Dense attention thread-block mapping — gfx942/gfx950 case study + +Measured scope: which thread block computes which work item in the dense attention +prefill kernel (`library/kernels/{gfx942,gfx950}/attention_dense.py`), and what that +costs. The kernel is self-attention (`Sq == Skv`) in bf16/fp16 with D 64/128, causal or +non-causal, on gfx942 and gfx950. It records the experiments behind the XCD- and +L2-aware block orders of these kernels. The previous version (the baseline) is `develop` +at `c81bca5ecf`, whose dense kernels are unchanged up to `5d7f9be53e`. The new version +is the first version of the orders (2026-09-30). The pair fold and the persistent +`hq_minor_swz` (§2.3) came later, on top of it. §6.1 measures the final version +(2026-10-13) against `develop` at `429ba1fe444`. + +Per [`platform/AGENTS.md`](../../AGENTS.md) §Compliance, this file records relative +results only: ratios between two code paths on the same device, with the correctness +gate, method and caveats stated. Absolute throughput is kept out of the repo. Results +marked **(experimental)** were measured on an unmerged experiments branch (based on +`develop` `7f09e16be8`) or on a prototype. The mapping each one uses is defined in +this document, and each passed the same correctness gate. Varlen, paged, +sliding-window and bottom-right-causal attention are not studied. They keep the +previous orders, apart from the automatic query-block fold described below (none +under a sliding window shorter than the keys). + +## 1. Shape sets + +| set | shapes | what it is for | +|---|---|---| +| end-to-end set | 99 (gfx942) / 114 (gfx950) dense-eligible prefill shapes of public models: MHA with 32–40 heads; GQA with 4 or 8 KV heads and 28–128 query heads; `S` 512–16K; `B = 1` (one `B = 2`); D128 (5–7 D64); bf16/fp16; causal, 11 with a sliding window shorter than the keys, and one non-causal on gfx950. gfx950 also runs 46 decode shapes (`Sq = 1`), and with tuned specs only, 6 more sliding-window shapes (windows of 128–256 tokens) | real-model shapes run end to end through the dispatcher, `develop` vs the final version, with default and with per-shape tuned specs (§6.1) | +| order-comparison grid | 48 prefill geometries plus a grid of MHA 8–128 heads and GQA `Hkv` ∈ {2, 4, 8, 16, 32} with `Hq` 32–128, `S` ∈ {8K, 32K}, `B` ∈ {1, 2, 4, 8}; causal and non-causal (about 250 shapes per arch) | the non-persistent orders compared with each other; the grid part is the Swizzled Head-first paper's grid, extended with more KV-head counts (§6.2) | +| grid-choice set | the 46 distinct geometries of an earlier, smaller end-to-end set (83 gfx942 / 105 gfx950 shapes, `S` up to 8K), each run causal and non-causal | whether auto picks the right grid, persistent or non-persistent (§6.4) | +| large-batch set | `B` ∈ {8, 16, 32} × MHA 64/64 and GQA 28/4, 32/8, 128/8 × `S` ∈ {2K, 8K}, causal (20 shapes within the kernel's 32-bit extents) | persistent orders and grid choice at large batch (§6.4) | +| profiling set | six representative shapes (§6.3) | every order and grid with L2 counters (§6.3) | + +## 2. Background + +### 2.1 How blocks reach the hardware + +gfx942 and gfx950 have 8 XCDs (compute dies). Each XCD has its own 4 MB L2 and 38 +(gfx942) or 32 (gfx950) CUs, and the XCDs share a last-level cache in front of HBM. +Workgroups are dispatched in linear-id order, x fastest, and **round-robin over the +XCDs: `xcd = linear_id % 8`**. So the low digits of the linear id decide which work +shares an L2. + +The kernel has two grids: + +- **Non-persistent:** one workgroup per work item. A CU takes the next workgroup when it + frees up, so the assignment is *dynamic* (a queue). +- **Persistent:** `NP` workgroups (the CU count) grid-stride over a work index `wi`. + Workgroup `c` runs `c, c+NP, c+2NP, …`, so the assignment is *static* (pinned). + +### 2.2 Work items and their cost + +A work item is `(query block, query head, batch)`. Its K/V is shared by every query +block and by the G query heads of one `(batch, kv head)`. Without a mask every item +costs the same. Under causal masking query block `qb` reads `qb+1` K/V tiles, so cost +grows linearly with `qb`. + +The order therefore decides three things, which turned out to matter in this order: + +1. **Causal balance and the tail.** What runs last, and how evenly the pinned + workgroups are loaded. This was the largest effect measured. +2. **Which XCD serves a K/V stream.** If a stream's consumers span several XCDs, each + XCD fetches the stream into its own L2. +3. **How many streams each XCD has live at once.** This matters mostly at long `S` + with many heads. + +### 2.3 Query-block traversal + +- **Longest first** (`NQB-1-qb`) is right for the dynamic queue. It is list + scheduling with the longest jobs first, and it leaves the cheapest blocks for the + tail. +- **Fold** (`qb < NQB/2 ? qb : NQB-1-(qb-NQB/2)`) is right for the pinned grid. Along a + workgroup's stride it pairs a cheap block with an expensive one, so per-workgroup + totals equalize. +- **Pair fold** (persistent `hq_minor_swz`, `hkv_major`; even `NQB`) runs both blocks + of `{p, NQB-1-p}` on one workgroup's consecutive steps, so every unit costs the + same whatever the digit order. The fold balances only when a stride crosses both + halves: automatic when the query block is the slowest digit (`qb_major`, + `bt_hkv_minor`, where pairing measured at most 0.4% either way, §7), not under XCD head bands + or with the kv head slowest. +- **Alternating directions** (`interleave`: reverse on odd heads) is weaker than both. + Every workgroup still sees the full cost range, and half of them end on the most + expensive blocks. +- **No traversal can rescue an order whose fastest digit is the query block.** Every run + of `NQB` consecutive items then holds each cost exactly once. A traversal only + permutes costs inside a run, so the most expensive blocks recur up to the end of the + dispatch. + +Ordering stops mattering in two regimes. With a single wave (`W ≤` CUs) there is no +tail to shape. When every work item is very long (`S ≥ 16K`), the tail is a small +fraction of the runtime. + +### 2.4 Capacity and temporal locality + +K plus V of one `(batch, kv head)` is `512·S` bytes at D128, so a whole head fills one +XCD's L2 at about `S = 8K`. To confine a stream to one XCD, the stream also needs +enough work to fill it: `G·NQB` items against 38 or 32 CUs. MHA never fills an XCD with +one head's blocks. That is why MHA needs a head *band* per XCD (Swizzled Head-first, +§3) rather than one stream per XCD. + +**Temporal locality beyond capacity.** A head that does not fit still gets L2 reuse. +Workgroups that consume the same stream and start together walk its K/V tiles from 0 +upward at nearly the same pace. They are not synchronized, but every KV iteration +costs the same except the diagonal tile, so they drift apart slowly. A tile the +leader loads is still resident when the followers reach it. The live set per XCD is +then *streams in flight × drift window*, not whole heads. + +- **It holds when:** + - a stream's consumers are adjacent in dispatch order, so they start in the same + wave on the same XCD; + - per-iteration cost is uniform, which is true for causal and non-causal alike over + the shared K/V prefix; + - there are few streams per XCD. +- **It weakens when:** + - starts are staggered (later waves, or the uneven finishes that longest-first + creates on purpose); + - a stream's consumers are spread over XCDs or waves; + - an XCD interleaves many streams; + - drift accumulates over very long K/V loops. + +Measured (§6.3): at `S = 32K`, where one head's K/V is four times an XCD's L2, the +non-persistent auto order (`bt_hkv_minor`) keeps an 87% (gfx942) / 95% (gfx950) L2 +hit rate. + +### 2.5 Notation + +Orders are written as digits, fastest first: `Q` query block, `B` batch, `V` kv head, +`G` query head within its kv group (`hq = hkv·G + g`). + +## 3. Prior art + +| work | mapping | relevance | +|---|---|---| +| [Swizzled Head-first](https://arxiv.org/abs/2511.02132) (arXiv 2511.02132) | each XCD owns a contiguous band of query heads and runs one head's blocks at a time | added here as `hq_minor_swz`. The paper has no causal traversal, its batch formula is inconsistent for `B > 1`, and its GQA tests are only at 8 KV heads | +| aiter `remap_xcd` (Triton MHA) | heads banded per XCD, query blocks ascending ("Swizzled Block-first") | the de-facto AMD convention | +| aiter `remap_workgroup_spatial` (opt-in) | KV head pinned to an XCD, block-first over its query heads; fewer than 8 KV heads are split over XCDs by contiguous block ranges | prototyped here (§7). Its split is unbalanced under causal masking, and its fallback for head counts not divisible by 8 is not a bijection | +| AOTriton causal forward | persistent, tiles claimed from an atomic counter, ascending | balances with a dynamic queue instead of an order | +| [FlashAttention-3/4](https://github.com/Dao-AILab/flash-attention/blob/main/hopper/tile_scheduler.hpp) tile scheduler | batch outermost, head sections sized to fit L2, longest-first query blocks, dynamic claiming | the closest analogue: L2-sized sections plus longest-first | +| [HipKittens](https://arxiv.org/abs/2511.08083) | notes that per-XCD L2 locality and shared-cache traffic trade off on 8-XCD parts | a caution against maximizing L2 locality alone | +| [Sawtooth](https://arxiv.org/abs/2601.16032) | alternates the K/V loop direction between a workgroup's consecutive items | only pays once K/V exceeds the cache | + +None of the AMD attention libraries read here reverses or folds causal query blocks; +they balance dynamically or not at all. Versions read from source: aiter `c8325e00` +and `21ae719`, AOTriton 0.14b (`b5e8cfb`), FlashAttention `e9cf2c1`. + +## 4. Previous status + +- **Non-persistent:** one fixed order, grid `(NQB, Hq, B)` (now `qb_minor`, `QGVB`). + The query block is the fastest digit and runs ascending, which is the worst class + for causal attention (§2.3). Adjacent workgroups also split one K/V stream across + XCDs by query block. +- **Persistent:** `qb_major` (ascending), `hkv_major` (folded), the `interleave` knob + and, on gfx950, `gqa_pair` / `gqa_pair_2phase`. Auto chose `hkv_major` or + `qb_major`, and on gfx950 it chose the pair decodes where their CTA-count condition + happened to hold. +- **Consequence:** on causal shapes the non-persistent path was 26–44% slower than the + persistent auto path at `S ≤ 8K`, and 14–20% slower at `S = 32K` (profiling set, + §6.3). So gfx942 dispatch went persistent as soon as the grid filled one wave. + +## 5. What this case study implemented + +| path | value | order | used for (auto) | +|---|---|---|---| +| non-persistent | `qb_minor` | `QGVB`, the previous grid | non-causal, and everything that is not aligned causal | +| non-persistent | `bt_hkv_minor` (new) | `BVGQ`, longest-first | aligned causal (default) | +| non-persistent | `hq_minor_swz` (new) | Swizzled Head-first: XCD `a` owns head band `a`, one head's blocks at a time, longest-first | causal MHA with many heads × batch × blocks, where dispatch picks the non-persistent grid | +| persistent | `hq_minor_swz` (new) | Swizzled Head-first, pair fold | aligned causal MHA with more than 8 grid-stride rounds of work (`Hq` divisible by 8, even `NQB`) | +| persistent | `bt_hkv_minor` (new) | `BVGQ`, folded | aligned causal, batch below the XCD count | +| persistent | `hkv_major` | `BGQV`, now pair fold under causal without a sliding window | GQA with enough work per kv head (the previous rule), including aligned causal from `chiplet_num_xcds` batches | +| persistent | `qb_major` | `BGVQ`, now folded under causal without a sliding window | everything else on the persistent grid | +| persistent | `gqa_pair*` | unchanged | explicit only | + +- **Knobs:** `nonpersist_decode` is new and mirrors `persist_decode`. + `chiplet_num_xcds` defaults to 8; another value stays correct and only weakens + locality. +- **Tuning knob and builder flag:** `nonpersist_decode` is an axis of the dense + tuning space (on gfx950, of the grid body), and the builders take + `--nonpersist-decode`. +- **Kernel name and cache key:** both include the resolved order. +- **Balance is automatic:** longest-first on the non-persistent grid, fold or pair + fold (§2.3) on the persistent grid. There is no knob. Only without a sliding + window shorter than the keys: such a window makes the work per query block + about uniform, and the `qb_major` fold measured 4-5% slower there. A window at + least as long counts as plain causal. +- **"Aligned causal"** means causal attention on the plain dense layout: no sliding + window shorter than the keys, ragged, varlen, paged, or moving bottom-right + diagonal. +- **Default grid (§6.4):** + - gfx942's default spec takes the persistent grid at D128 for every batch and + mask, and the non-persistent grid at D64; + - gfx950 has no default grid: dense is opt-in and the caller pins the grid or the + persistent body. The measurements used its previous rule, persistent once the + work fills the grid. +- **The persistent grid stays ahead at D128 on both arches, for different reasons.** + - gfx950's persistent builder is a different kernel body with extra optimizations: + wide LDS DMA (128-bit buffer→LDS loads) and a different scheduling template. So + non-persistent on gfx950 is still slower than persistent. + - On gfx942 one body serves both grids, yet the persistent grid with the new + decodes also measured ahead at D128, even when the work fits in one wave (§6.4). + At D64 the non-persistent grid is ahead. +- **Correctness:** every order is a bijection onto the work space, so outputs are + bit-identical across orders. Unit tests evaluate every order's decode on the host, + and every measured cell was checked on the GPU. + +## 6. Performance effect + +**Method, for every table below:** +- **correctness:** bf16/fp16 inputs; outputs checked against a PyTorch SDPA reference + (max abs error ≤ 2e-2) wherever the reference was affordable, and all orders of a + shape verified bit-identical. On shapes too large for the reference, correctness + rests on bit-identity with an order that passed SDPA at smaller shapes; +- **timing:** one process per kernel; median of 3–5 timed runs of back-to-back + launches (warm caches); two independent passes (§6.1: five rounds, below); +- **uncertainty:** every result carries a 95% interval, written `±x` (points of the + percentage) or `[lo, hi]`: + - §6.1 alternates develop and the new version in five rounds on one node, so each + round gives one paired ratio per shape. A shape's interval is a t-interval over + its five ratios. An aggregate's interval is a t-interval over the five per-round + geomeans, so offsets that move a whole round are included. An A/A comparison + (the same version in two rounds) shows what no change looks like. + - §6.2–§6.4 and §7 (sweeps of two passes) pool the noise per sweep and arch from + the pass-to-pass differences of the paired ratios: σ = 0.26–0.9% per paired + measurement. A shape's interval is `±1.96·σ/√passes` (±0.4–1.2%). A geomean over + `n` shapes gets `±(1.96·σ/√(passes·n) + half its pass-to-pass difference)` + (±0.1–0.3 points). Rows measured in one pass take σ from A/A pairs in the same + run (the same kernel on both sides): σ = 0.3–1.1%. + - The intervals cover run-to-run noise on one node. They do not cover + node-to-node variation or the choice of shapes (§6.1 adds a bootstrap over + shapes for the latter). +- **decision rule:** a delta counts as a gain or a loss only when its interval + excludes 0; otherwise it is reported as within noise. Every choice in §5 and §7 + holds under this rule; where a delta is within noise, the text says so. +- **configurations:** arms that pin an order or a grid are diagnostic configurations; + "auto" is the default spec; +- **dispatcher:** the runs of §6.2–§6.4 and §7 predate the tuning-knob dispatch + interface and pinned orders through the request fields it replaced; the + kernels are the same. + +Replay a single point with the builders +`library/builders/{gfx942,gfx950}/attention/prefill/attention_dense_prefill.py`, which +take `--nonpersist-decode` / `--persist-decode` and check parity against SDPA. + +### 6.1 End to end, `develop` vs the final version, end-to-end set + +Both versions ran through the dispatcher: `develop` at `429ba1fe444` against the final +version. Two configurations: +- **default:** the dispatcher's default dense spec; +- **tuned** (a diagnostic configuration): a per-shape dense spec from a tuner. `develop` + uses the specs tuned for it. The final version keeps each spec's tile and code + generation options but re-chooses its grid and block order per shape. + +gfx942 ran on two parts with 304 CUs, (a) and (b); tuned specs exist for (b) and gfx950 +only. Per part, five rounds ran on one node; each round ran every configuration once +per version, `develop` and the final version back to back, alternating which went +first. Each point is the median of 5 iterations of 250 timed executions after 75 +warmups. Every output of both versions in every round was checked against an fp64 +reference, max abs error ≤ 5e-2: all passed. Ratios are throughput, final version over +`develop`. `[lo, hi]` is the 95% interval over the rounds (§6). "Shape set" is a +bootstrap over shapes: how much the geomean depends on which shapes were measured. + +| prefill shapes | default, gfx942 (a) | default, gfx942 (b) | default, gfx950 | tuned, gfx942 (b) | tuned, gfx950 | +|---|---|---|---|---|---| +| shapes | 99 | 99 | 114 | 99 | 120 | +| **geomean** | **+8.5% [+8.4, +8.6]** | **+9.3% [+8.8, +9.8]** | **+2.7% [+2.5, +3.0]** | **+9.0% [+8.6, +9.3]** | **+1.4% [+1.1, +1.8]** | +| shape set | [+6.7, +10.5] | [+7.4, +11.3] | [+1.3, +4.1] | [+7.0, +11.1] | [+0.6, +2.3] | +| causal, `S ≤ 1K` | +5.4% [+5.1, +5.6] | +6.2% [+6.1, +6.4] | −1.2% [−1.8, −0.5] | +3.8% [+3.4, +4.2] | −1.0% [−1.8, −0.2] | +| causal, `S` 2K–4K | +13.5% [+13.4, +13.5] | +14.3% [+13.7, +14.9] | +6.2% [+6.1, +6.3] | +18.0% [+17.9, +18.2] | +3.6% [+3.6, +3.7] | +| causal, `S ≥ 8K` | +5.8% [+5.7, +6.0] | +6.2% [+5.7, +6.7] | +5.6% [+5.5, +5.8] | +5.9% [+5.8, +6.0] | +1.4% [+1.4, +1.5] | +| sliding window shorter than the keys | +0.1% [−0.4, +0.6] | +0.3% [−0.9, +1.6] | −0.8% [−0.9, −0.7] | −0.7% [−1.6, +0.2] | −0.4% [−0.8, −0.1] | +| D64 | +20.3% [+19.8, +20.8] | +22.9% [+21.7, +24.2] | +1.5% [+1.0, +2.0] | +8.1% [+7.0, +9.2] | +4.0% [+3.1, +4.9] | +| faster / within noise / slower | 74 / 14 / 11 | 75 / 16 / 8 | 65 / 24 / 25 | 66 / 20 / 13 | 62 / 29 / 29 | +| worst / best shape | −8.3% / +39.2% | −5.5% / +42.6% | −11.3% / +27.3% | −9.8% / +48.4% | −9.4% / +21.7% | +| shape interval, median half-width | ±0.4% | ±0.8% | ±0.3% | ±0.7% | ±0.3% | + +**A/A** (the same version in consecutive rounds): a whole round moves by at most 0.14% +(a), 0.94% (b) and 0.43% (gfx950) in geomean; 95% of single-shape A/A deltas are within +1.0%, 2.1% and 2.2%. + +- **gfx942 gains at every sequence length,** default and tuned. Most shapes now run + the persistent grid with the new decodes (§6.4). +- **The slower shapes fall into three groups,** plus four single shapes at −0.1% to + −1.3%: + - `S ≤ 1K` (7 / 5 / 22 of the default slower shapes; 11 / 26 tuned): at `S` = 512 a + launch is host-bound. The final version does more host work per launch (two more + spec fields in every validation, cache-key and name pass, and the order + resolution). A separate host-vs-device measurement on gfx950 at `S` = 512 put the + host time per launch up by about 8% and the device time up by only 0.4–6%. This + is why gfx950 `S ≤ 1K` reads −1%; + - sliding window: 1 / 1 / 3 default, 3 / 5 tuned. The worst is a 128-token window at + `S` = 512, tuned gfx942 (b) −9.8% [−11.5, −8.0]: its re-chosen spec runs the + persistent grid (`hq_minor_swz`) where `develop`'s tuned spec ran the + non-persistent one; + - 40 query heads at `S` = 4K on gfx942 (4 / 2 shapes, −1.0% to −4.7%): the work is + just over a whole number of persistent rounds, so the grid-stride tail is long; + per-shape tuned specs avoid it (§8). +- **gfx950 decode** (`Sq = 1`, 46 shapes, default only; its kernels are unchanged): + −0.6% [−0.9, −0.4], 38 within noise. This is the host cost above, with nothing to + gain. + +### 6.2 Non-persistent orders, order-comparison grid + +Non-persistent grid only. Geomean vs the previous non-persistent order (`qb_minor`), +order-comparison grid. "Auto" is the non-persistent auto order; which grid auto picks +is §6.4. + +| regime | mask | gfx942: `bt_hkv_minor` / `hq_minor_swz` / auto | auto vs best order | gfx950: `bt_hkv_minor` / `hq_minor_swz` / auto | auto vs best | +|---|---|---|---|---|---| +| MHA | causal | +24.1% / +27.5% / **+28.5%** | −0.8% | +26.8% / +26.1% / **+31.2%** | −0.1% | +| GQA `Hkv < 8` | causal | +29.4% / +26.2% / **+29.4%** | −0.1% | +35.6% / +32.8% / **+35.6%** | −0.3% | +| GQA `Hkv = 8` | causal | +32.3% / +23.6% / **+32.3%** | −0.2% | +38.1% / +28.6% / **+38.1%** | −0.7% | +| GQA `Hkv > 8` | causal | +28.1% / +28.4% / **+28.1%** | −1.8% | +35.5% / +38.4% / **+35.5%** | −3.3% | +| all | non-causal | −8.4% … +0.2% / −4.1% … −0.9% / **0** | −0.0 … −0.7% | −7.5% … −1.0% / −3.0% … −1.2% / **0** | −0.4 … −1.1% | + +Uncertainty (§6): every geomean in this table has a 95% interval of ±0.1–0.3 points; +one paired measurement has σ = 0.44% (gfx942) / 0.57% (gfx950). Within noise (the +interval contains 0): the zero entries, gfx950 MHA causal auto vs best (−0.1%), and +gfx942 `Hkv < 8` non-causal `bt_hkv_minor` (+0.2%). Every other entry excludes 0. +"Auto vs best" picks the best order per shape on the mean of both passes; picking it on +one pass and scoring on the other gives the same values within 0.1 point, so these gaps +are not a selection artefact. + +- **Causal:** the new orders are worth about +24% to +38% over the previous order, and + +45% to +67% at `S` 4K–8K. Auto is within 0.1–1.8% (gfx942) / 0.1–3.3% (gfx950) of + the best order per shape. +- **Where auto loses most:** large GQA (many heads, `B ≥ 4`, long `S`). There + `hq_minor_swz` beats `bt_hkv_minor` by up to 13.8% ±0.6 (gfx942) / 16.9% ±0.8 + (gfx950), both at GQA 128/32, `S` = 8K, `B` = 8. Auto selects it only for MHA (§8). +- **Non-causal:** the previous order is still the best non-persistent order. Step 0a + agrees: on gfx942 every alternative is 1.2–3.2% behind (±0.2–0.4); on gfx950 + `bt_hkv_minor` and Swizzled Head-first are 2.0% / 0.5% behind and two more orders tie + (−0.2% ±0.3). So auto keeps it. + +### 6.3 Representative shapes: throughput and L2 counters + +bf16, D128. "Auto" is the dispatcher's default choice: grid and order both automatic. +Throughput is relative to previous auto on the same arch. L2 hit rate is +`TCC_HIT / (TCC_HIT + TCC_MISS)`. "L2 misses" counts L2 read misses to the fabric +(`TCC_EA0_RDREQ`), relative to previous auto, median over 4 dispatches under +rocprofv3. The non-persistent orders were bit-identical on every shape. The SDPA +check ran on the `B = 1` shapes (the other two were too large). + +Uncertainty (§6): each throughput entry is the mean of two passes, with a 95% interval +of ±0.4% (gfx942) / ±0.6% (gfx950) (pass-to-pass σ = 0.26% / 0.45% per paired +measurement; same-kernel A/A pairs in these runs agree). Within noise: gfx942 32/8 +`S` = 32K non-persistent `hq_minor_swz` (−0.1%), gfx950 128/8 `B` = 16 new auto +(+0.5%), and gfx950 non-causal non-persistent `qb_minor` / `bt_hkv_minor` (−0.3% / ++0.1%). Every other entry excludes 0. + +**gfx942** + +| shape | arm | throughput vs previous auto | L2 hit rate | L2 misses vs previous auto | +|---|---|---:|---:|---:| +| GQA 32/8, S=4K, B=1, causal | previous auto (persistent `qb_major`) | +0.0% | 54% | 1.00× | +| | previous non-persistent `qb_minor` | -25.9% | 82% | 0.35× | +| | new non-persistent `bt_hkv_minor` | +23.1% | 86% | 0.24× | +| | new non-persistent `hq_minor_swz` | -8.0% | 81% | 0.35× | +| | new auto (persistent `bt_hkv_minor`) | +29.3% | 81% | 0.38× | +| GQA 64/8, S=8K, B=1, causal | previous auto (persistent `qb_major`) | +0.0% | 27% | 1.00× | +| | previous non-persistent `qb_minor` | -33.2% | 81% | 0.23× | +| | new non-persistent `bt_hkv_minor` | +2.1% | 91% | 0.10× | +| | new non-persistent `hq_minor_swz` | -0.8% | 77% | 0.29× | +| | new auto (persistent `bt_hkv_minor`) | +5.8% | 73% | 0.35× | +| GQA 128/8, S=4K, B=16, causal | previous auto (persistent `hkv_major`) | +0.0% | 56% | 1.00× | +| | previous non-persistent `qb_minor` | -35.5% | 87% | 0.24× | +| | new non-persistent `bt_hkv_minor` | -6.6% | 65% | 0.75× | +| | new non-persistent `hq_minor_swz` | -8.1% | 76% | 0.49× | +| | new auto (persistent `hkv_major`) | +5.9% | 87% | 0.24× | +| MHA 64/64, S=8K, B=4, causal | previous auto (persistent `qb_major`) | +0.0% | 25% | 1.00× | +| | previous non-persistent `qb_minor` | -28.7% | 63% | 0.48× | +| | new non-persistent `bt_hkv_minor` | -3.4% | 28% | 0.95× | +| | new non-persistent `hq_minor_swz` | +2.5% | 79% | 0.26× | +| | new auto (persistent `hq_minor_swz`) | +10.4% | 51% | 0.65× | +| GQA 32/8, S=32K, B=1, causal | previous auto (persistent `qb_major`) | +0.0% | 9% | 1.00× | +| | previous non-persistent `qb_minor` | -14.0% | 61% | 0.42× | +| | new non-persistent `bt_hkv_minor` | +1.3% | 87% | 0.13× | +| | new non-persistent `hq_minor_swz` | -0.1% | 79% | 0.22× | +| | new auto (persistent `bt_hkv_minor`) | +6.3% | 58% | 0.45× | +| GQA 32/8, S=4K, B=1, non-causal | previous auto (persistent `qb_major`) | +0.0% | 70% | 1.00× | +| | previous non-persistent `qb_minor` | -8.7% | 88% | 0.71× | +| | new non-persistent `bt_hkv_minor` | -12.5% | 90% | 0.38× | +| | new non-persistent `hq_minor_swz` | -11.8% | 83% | 0.63× | +| | new auto (persistent `qb_major`) | -0.0% | 70% | 1.00× | + +**gfx950** + +| shape | arm | throughput vs previous auto | L2 hit rate | L2 misses vs previous auto | +|---|---|---:|---:|---:| +| GQA 32/8, S=4K, B=1, causal | previous auto (persistent `gqa_pair_2phase`) | +0.0% | 72% | 1.00× | +| | previous non-persistent `qb_minor` | -38.1% | 80% | 0.67× | +| | new non-persistent `bt_hkv_minor` | -2.6% | 90% | 0.25× | +| | new non-persistent `hq_minor_swz` | -27.7% | 91% | 0.24× | +| | new auto (persistent `bt_hkv_minor`) | +4.3% | 91% | 0.23× | +| GQA 64/8, S=8K, B=1, causal | previous auto (persistent `gqa_pair`) | +0.0% | 67% | 1.00× | +| | previous non-persistent `qb_minor` | -43.7% | 88% | 0.30× | +| | new non-persistent `bt_hkv_minor` | -3.9% | 93% | 0.16× | +| | new non-persistent `hq_minor_swz` | -8.5% | 91% | 0.21× | +| | new auto (persistent `bt_hkv_minor`) | +2.5% | 93% | 0.14× | +| GQA 128/8, S=4K, B=16, causal | previous auto (persistent `hkv_major`) | +0.0% | 89% | 1.00× | +| | previous non-persistent `qb_minor` | -41.6% | 89% | 1.00× | +| | new non-persistent `bt_hkv_minor` | -8.4% | 58% | 4.90× | +| | new non-persistent `hq_minor_swz` | -9.6% | 90% | 0.82× | +| | new auto (persistent `hkv_major`) | +0.5% | 89% | 1.00× | +| MHA 64/64, S=8K, B=4, causal | previous auto (persistent `qb_major`) | +0.0% | 17% | 1.00× | +| | previous non-persistent `qb_minor` | -26.0% | 57% | 0.50× | +| | new non-persistent `bt_hkv_minor` | -2.2% | 18% | 1.00× | +| | new non-persistent `hq_minor_swz` | +7.9% | 86% | 0.14× | +| | new auto (persistent `hq_minor_swz`) | +24.2% | 59% | 0.48× | +| GQA 32/8, S=32K, B=1, causal | previous auto (persistent `hkv_major`) | +0.0% | 73% | 1.00× | +| | previous non-persistent `qb_minor` | -19.9% | 71% | 1.08× | +| | new non-persistent `bt_hkv_minor` | -4.2% | 95% | 0.15× | +| | new non-persistent `hq_minor_swz` | -5.3% | 93% | 0.23× | +| | new auto (persistent `bt_hkv_minor`) | +1.3% | 91% | 0.32× | +| GQA 32/8, S=4K, B=1, non-causal | previous auto (persistent `qb_major`) | +0.0% | 86% | 1.00× | +| | previous non-persistent `qb_minor` | -0.3% | 86% | 1.00× | +| | new non-persistent `bt_hkv_minor` | +0.1% | 93% | 0.40× | +| | new non-persistent `hq_minor_swz` | +1.1% | 93% | 0.40× | +| | new auto (persistent `qb_major`) | +0.7% | 86% | 1.00× | + +**Conclusions.** + +- **Balance explains most of the speedup, not hit rate.** On GQA `B = 1` causal shapes, + the previous non-persistent order already hit L2 80–88% of the time, yet was 26–44% + slower than previous auto. The new non-persistent `bt_hkv_minor` raises the hit rate + only to 86–93%, but removes the ascending tail. +- **The new auto cuts L2 misses substantially:** + - GQA `B = 1` causal: 2.2–2.9× fewer on gfx942 and 3–7× fewer on gfx950 than + previous auto, whose `qb_major` / `gqa_pair*` choices spread each K/V stream over + many XCDs; + - large MHA: the non-persistent `hq_minor_swz` needs 4–7× fewer (hit rate 79–86% + vs 17–28% for the other orders). The persistent `hq_minor_swz` that auto now + picks misses more (51–59% hit rate) but is the fastest, +10.4% (gfx942) / +24.2% + (gfx950): balance again outweighs hit rate. +- **Hit rate does not rank the grids.** On gfx942 the new auto (persistent + `bt_hkv_minor`) has a lower hit rate than the non-persistent `bt_hkv_minor` (58–81% vs + 86–91%) and is still 3–5% faster; see §6.4. +- **Temporal locality at `S = 32K`** (§2.4): the new non-persistent `bt_hkv_minor` keeps + 87% (gfx942) / 95% (gfx950) of reads in L2, against 9% / 73% for previous auto. +- **Unchanged where nothing better was found:** + - non-causal auto is still the previous persistent `qb_major` (the non-persistent + orders are 9–13% behind it on gfx942); + - large-batch GQA auto is still `hkv_major`, now with the pair fold: +5.9% on + gfx942, where its hit rate rises from 56% to 87% because each grid-stride phase + spans fewer kv heads, and unchanged within noise on gfx950 (+0.5% ±0.6), where + it was already 89%. + +### 6.4 Grid choice (dispatch) + +Auto vs previous auto, measured in the same session per arch. The shapes are the 46 +prefill geometries (mostly `B = 1`, a few `B` 2–4; D128, plus five D64), each run +causal and non-causal, plus 20 large-batch shapes (`B` 8–32, causal, D128). The +large-batch shapes were re-measured with the pair fold (§2.3); outputs were +bit-identical to the previous version. Every +run with an affordable SDPA reference matched it (max abs error 3.9e-3 in bf16). + +| slice | gfx942 | gfx950 | +|---|---|---| +| causal, `B ≤ 4` (46) | **+10.9% ±0.1** (worst −3.5% ±0.6) | **+6.5% ±0.3** (worst −1.3% ±2.1) | +| non-causal, `B ≤ 4` (46) | +5.5% ±0.1 (worst −3.1% ±0.6) | +0.2% ±0.3 (worst −1.6% ±2.1) | +| D64, both masks (10 of the 92 above) | +18.2% ±0.2 | +2.9% ±0.6 | +| causal, `B` 8–32 (20) | +4.3% ±0.2 (worst −1.0% ±0.5) | +3.7% ±0.1 (worst −2.8% ±0.5) | +| all (112) | **+7.4% ±0.1** (worst −3.5% ±0.6) | **+3.3% ±0.2** (worst −2.8% ±0.5) | + +Uncertainty (§6): ± is the 95% interval. The `B ≤ 4` rows are one pass; their noise +comes from A/A pairs in the same run (non-causal shapes whose kernel is the same on both +versions): σ = 0.31% (gfx942) / 1.04% (gfx950) per paired measurement. The `B` 8–32 +rows are two passes (σ = 0.38% / 0.33%). Within noise: gfx950 non-causal (+0.2%, so +unchanged) and both gfx950 `B ≤ 4` worst shapes. Every other entry excludes 0. + +- **gfx942 at D128: the persistent grid wins at every batch and mask.** With the new + decodes it is ahead of the non-persistent auto by 3.8% ±0.3 on causal and 10.4% ±0.3 + on non-causal shapes (`B = 1`). That holds even when the work fits in one wave, so part + of the gap is the grid's code path rather than the order. The cause is not + identified. +- **gfx942 at D64: the non-persistent grid wins,** by 10.8% ±0.9 (causal) and + 15.8% ±0.9 (non-causal). +- **Large batch:** MHA now runs the persistent `hq_minor_swz`: +10.5% ±0.3 (gfx942) / + +20.0% ±0.4 (gfx950). GQA keeps the previous rule, `hkv_major` where it applied, now + with the pair fold: +2.8% ±0.2 / +0.0% ±0.1 (unchanged), with gfx950 up to 2.8% + behind at `S = 8K`. +- **gfx950 at `S ≤ 1K`:** the unchanged previous rule picks the non-persistent grid + when the work does not fill the persistent grid. There the persistent grid measured + 7.9% ±0.6 ahead on causal shapes (13 of 13) and 2.4% ±0.6 on non-causal (§8). +- **Caveats:** + - the `B ≤ 4` set is mostly `B = 1`; + - D64 covers one geometry at five sequence lengths; + - sliding window was spot-checked on six gfx942 shapes only: persistent ahead on + five, behind on a 128-token window. + +## 7. What else was tried + +| idea | result | why | +|---|---|---| +| all 24 digit orders × {ascending, longest-first, fold} **(experimental)** | only the batch/kv-head-fastest class generalizes. Orders with the query block fastest lose 22–31%; `BVGQ` was within 1.6% ±0.1 geomean of the per-shape best on every arch and grid (1.4% when the best is picked on one pass and scored on the other) | §2.3: nothing rescues query-block-fastest orders; balance is a hard gate | +| locality-maximizing orders (one K/V stream per XCD, e.g. `VGQB`) **(experimental)** | the best locality measured *worst*; `BVGQ` won with 8–32× poorer locality | pinning one stream per XCD gives each workgroup one query block, which ruins causal balance. Locality only separates orders once balance is equal | +| kv-phase split (split the fused `(batch, kv head)` identity across the fastest and slowest digits) **(experimental)** | +1.3% ±0.2 on gfx942 persistent at `Hkv = 8` (large batch only); −0.5% to +0.4% (±0.2–0.4) elsewhere; about 12 extra integer ops per item | not worth a knob | +| `gqa_pair`, `gqa_pair_2phase` (gfx950) | very restrictive: GQA with even `G`, even `NQB`, a baked shape, and a CTA count that must equal `NQB·Hkv·B` (or half the work), so under the default CTA count they fire only by coincidence. At matched CTA count they are **2.2% ±0.2 (93 shapes) / 3.9% ±0.2 (64 shapes) behind** the new auto; at the default CTA count, far behind | the new persistent orders balance and localize at least as well without the constraints; now explicit only | +| pair fold for persistent `bt_hkv_minor` and `qb_major` (causal, 54 shapes) **(experimental)** | `bt_hkv_minor` −0.3% ±0.2 (gfx942) / −0.4% ±0.2 (gfx950), a small loss; `qb_major` −0.3% ±0.3 / +0.1% ±0.2, a tie | the query block is their slowest digit, so a workgroup's stride already crosses both halves of the fold | +| `interleave` vs fold (persistent `qb_major`, causal, 62 shapes) **(experimental)** | fold ahead by 3.3% ±0.2 (gfx942) / 2.8% ±0.2 (gfx950) geomean, and by 4.5% / 4.9% (±0.4) on the 30 shapes with `B < 8`; interleave only 0.7% / 0.4% (±0.1–0.2) better than ascending | §2.3; interleave is kept only as an explicit knob | +| Swizzled Head-first on the persistent grid, as-is **(experimental)** | 10.5% (gfx942) / 27% (gfx950) behind auto on causal shapes, whatever the traversal | grid-stride advances the block digit by `NP/8`; when that is a multiple of `NQB`, every workgroup keeps one query block for its whole life. The pair fold fixes this (§2.3) | +| aiter's KV-head-first swizzle (`hkv_minor_swz`: grid `(Hq·NQB, 1, B)`, XCD `a` owns kv heads `a·Hkv/8 …` one at a time, block-first over their query heads) **(prototype on the new version)** | vs non-persistent auto: `Hkv = 8` +0.1–0.5% on three of the four arch × shape-set cells, +2.8% ±0.2 on the fourth (gfx950, paper grid); `Hkv > 8` causal **+1.6% ±0.2 (gfx942) / +4.4% ±0.3 (gfx950)**, but gfx950 auto uses the faster persistent grid there anyway; `Hkv < 8` causal **13–35% behind**. MHA is identical to `hq_minor_swz` | `bt_hkv_minor` already pins each `(batch, kv head)` to one XCD whenever `B·Hkv % 8 == 0`. The `Hkv < 8` split hands one XCD the late, expensive blocks | +| first version of the new rules, gfx942: persistent only from 16 batches | behind the previous auto on non-causal D128 (−3.3% ±0.3 geomean, worst −14.9% ±1.9) and at `B = 8` (up to −7.3% ±1.0) | it relied on the non-persistent grid beating the *old* persistent decodes; with the new ones the persistent grid wins at D128 (§6.4). The rule now keys on head size | +| first version of the new rules: folded `qb_major` for aligned causal from `chiplet_num_xcds` batches | behind the previous `hkv_major` on gfx950 GQA (−2.9% ±0.3 geomean, worst −10.7% ±1.2) | the earlier study compared it only with `bt_hkv_minor`; auto now keeps the previous rule there | + +**Measurement lessons.** +- **The launcher cache is keyed by kernel name.** An order missing from the name + silently re-times another binary (it happened three times). Every order now has a + name tag and is part of the cache key. +- **Re-baseline against the current auto.** Quoting against an order nothing selects + overstates gains several-fold. +- **Re-check the grid choice after improving either grid.** Both first-version rule + errors above came from comparing a new candidate against a stale alternative. +- **Median over runs, then two passes.** Single bad runs invent regressions, and the + pass-to-pass differences give the noise for the intervals (§6). +- **Free noise probes.** Orders that coincide at `G = 1` or `B = 1` compile to the + same mapping, so any gap between them is noise. + +## 8. What can still be checked + +- **Pair longest and shortest block in one unit.** + - *Persistent:* done for even `NQB` (§2.3). Odd `NQB` still uses the fold; the + `hq_minor_swz` docstring has a recipe (pair the middle blocks of two heads). + - *Non-persistent:* one workgroup computes both blocks. That halves the grid, + equalizes workgroup cost and shares the K/V prefix, but needs `W/2 ≥` CUs and + makes each workgroup twice as long. Unlike `interleave`, it pairs within a + workgroup rather than alternating across heads. +- **`hq_minor_swz` for large GQA** (§6.2): the largest remaining auto losses. +- **The prototype KV-head-first swizzle for `Hkv > 8`** (small gain on gfx942), and + uneven head bands for head counts not divisible by 8. +- **Dynamic persistent scheduling:** an atomic work counter plus longest-first order, + like AOTriton and FlashAttention. +- **`num_persistent = CUs × occupancy`** instead of the CU count. On gfx942 D64 the + smaller tile runs two workgroups per CU while the grid has one per CU, which may be + why the persistent grid loses there (§6.4). +- **gfx950 grid choice for small problems:** persistent measured about 8% ahead at + `S ≤ 1K` (causal), where the unchanged rule picks the non-persistent grid (§6.4). +- **gfx942 GQA at large batch:** the folded `qb_major` measured 1.5% ahead of the + folded `hkv_major` the rule picked. The pair fold has since gained `hkv_major` + 3.1% there, so recheck before adding a gfx942-only rule. +- **gfx942 D128 shapes where the non-persistent grid still wins:** 40 heads at `S ≥ 4K` + (up to 18% ahead of persistent; end to end these shapes gave back 7–17% of the first + version's gain, landing within 3% of the previous version), and a 128-token sliding + window, where persistent is 11.5% slower. The first may be fold imbalance when the + work does not divide the 304-workgroup grid. +- **Exploit temporal locality deliberately at long `S`:** + - keep one stream's consumers in the same wave and XCD, and bound the streams per + XCD (FlashAttention's L2-sized sections); + - K/V-direction reversal (sawtooth) once one head exceeds L2 (`S > 8K`). +- **Coverage gaps,** each needing its own benchmarks: + - `B ≥ 16` on the non-persistent grid; + - `S ≥ 64K`, and non-power-of-two `S`; + - non-causal at large batch; + - varlen, paged, sliding window and bottom-right causal. + + On the order-comparison grid, auto is within about 1–2% of the best order on + average up to `B = 8` and `S = 32K`, apart from the large-GQA case above. + +## 9. Other possible improvements (not orderings) + +- **Port wide LDS DMA and the persistent scheduling template to gfx950's non-persistent + body.** Disabling `wide_lds_dma` on the persistent grid costs about 6% + **(experimental)**, so porting it could give the non-persistent grid a similar gain. + That would close most of the gap in §5 and reach varlen, which only has a + non-persistent path. +- **Unify gfx950's persistent and non-persistent bodies if possible,** as gfx942 already does. +- **Runtime-shape persistent decodes.** Today they bake `NQB`. +- **GQA packing:** put a kv head's G query heads in one tile, so K/V reuse no longer + depends on the cache. +- **Retire `interleave`, and consider retiring `gqa_pair*`** if broader benchmarks + confirm they are not ahead of the other orders. + +## 10. Where the code lives + +| concern | file | +|---|---| +| decode classes (grid + decode per order) | `library/kernels/common/attention_dense_decode.py` | +| knobs, legality, auto rules, name tags, cache key | `library/kernels/common/attention_dense_spec.py` | +| call sites (grid, decode, persistent loop) | `library/kernels/{gfx942,gfx950}/attention_dense.py` | +| persistent vs non-persistent dispatch | `library/dispatch/attention/{gfx942,gfx950}.py` | +| builders (`--persist-decode`, `--nonpersist-decode`) | `library/builders/{gfx942,gfx950}/attention/prefill/attention_dense_prefill.py` | +| host tests of every decode | `library/tests/test_attention_dense_decode.py` | diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/attention.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/attention.md index 41cae2b68ecb..a100a16e09d1 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/attention.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/attention.md @@ -10,7 +10,8 @@ Production tiled matrix-core, paged-KV decode: - `library/kernels/gfx950/attention_tiled_2d.py`, `library/kernels/gfx942/attention_tiled_2d.py`, `library/kernels/gfx1250/attention_tiled_2d.py` - `library/kernels/gfx950/attention_tiled_3d.py`, `library/kernels/gfx942/attention_tiled_3d.py`, `library/kernels/gfx1250/attention_tiled_3d.py` - `library/kernels/gfx950/attention_tiled_2d_fastkv_regp.py` -- fast paged-KV + register-resident P (gfx950) -- `library/kernels/gfx950/attention_dense.py` -- dense (non-paged) flash prefill +- `library/kernels/gfx950/attention_dense.py` -- dense (non-paged) flash prefill; block orders and + grid choice: [case study](../architecture/attention_thread_block_mapping_case_study.md) - `library/kernels/common/fmha_mfma.py` (unified MFMA/WMMA forward, MFMA on CDNA / WMMA on RDNA) - `platform/python/rocke/helpers/attention.py` @@ -565,9 +566,11 @@ Step 1 above builds the key from the spec. By default *every* spec field partici A **runtime param field** is a spec field the kernel body reads as a kernel argument instead of baking into the body. The spec declares them by overriding `runtime_param_fields`, and `attention_dense_cache_key` excludes exactly those fields from the key, so one compiled binary serves every value of them: ```text -attention_dense_cache_key(spec, arch) -> (arch, type(spec), ) +attention_dense_cache_key(spec, arch) -> (arch, type(spec), + ) ``` +The block order the kernel uses is keyed explicitly because auto may pick it from shape fields the key drops. + The declaration lives on the spec that owns the body, not in the shared key function, because the two must agree: declaring a field the body still bakes is a cache collision — different problems served by the wrong binary. `library/tests/test_attention_builds.py::TestAttentionDenseRuntimeShapeCollision` guards that direction by asserting specs sharing a key lower to identical IR. Today only `library/kernels/gfx950/attention_dense.py` opts in, declaring `("batch", "seqlen_q", "seqlen_kv")` on its aligned, unshifted dense path. Sub-modes that still bake seqlen into the body — persistent, ragged, varlen, paged, sliding-window, and bottom-right causal with unequal Q/K lengths — declare nothing and keep per-shape identity. Equal-length bottom-right adds no extra shape specialization. `library/kernels/gfx942/attention_dense.py` opts in further on both its default and persistent grids (sliding-window included), declaring `("batch", "seqlen_q", "seqlen_kv", "num_query_heads", "num_kv_heads")`; only ragged, varlen and paged declare nothing. The persistent body also takes a host-computed fast-division magic/shift pair per work-decode divisor, after those five (`runtime_kernarg_fields`); the pairs are derived from the declared fields, so they do not enter the key. Every other family also declares nothing, so its identity is unchanged. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md index be33108a59a5..f77b3809d998 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md @@ -2108,7 +2108,7 @@ them by hand only through an explicit tuning spec. |---|---|---|---| | `chiplet_swizzle` | GEMM / conv | False | Remap WGIDs so contiguous stripes land on the same XCD (L2 reuse). Rejected by `is_valid_dgrad_spec`: the dgrad builder does not implement it (dgrad has `xcd_tile_order`) | | `chiplet_wgm` | GEMM / conv | 8 | Super-tile WGM grouping | -| `chiplet_num_xcds` | GEMM / conv | 8 | MI300X / MI325X / MI350X have 8 XCDs | +| `chiplet_num_xcds` | GEMM / conv / dense attention | 8 | MI300X / MI325X / MI350X have 8 XCDs; dense attention's XCD-aware block orders read it | | `chiplet_chunk_size` | GEMM / conv | 64 | XCD round-robin chunk size | | `helpers/grid.py::chiplet_transform_chunked` | helper | — | Pure helper if you author your own kernel | | Constants `NUM_XCDS_MI300X / MI325X / MI350X` | `helpers/grid.py` | — | All 8 today | diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json index 0cb88103c848..3f7756e99bc1 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json @@ -330,12 +330,52 @@ "sha256": "cd4307bce99bb62a7f06701560e1b8d398295b7597a4b6701e27380694469e84", "status": "ok" }, + "attention_dense/gfx942/causal_sq512": { + "arch": "gfx942", + "bytes": 171590, + "family": "attention_dense", + "kernel_name": "rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_npbthkvmin_gfx942_wpe2", + "sha256": "3a21804add31a71453e5f86874912285e3f66bb0cb8c3b481f44e6b55ad00812", + "status": "ok" + }, + "attention_dense/gfx942/hqswz_sq512": { + "arch": "gfx942", + "bytes": 171616, + "family": "attention_dense", + "kernel_name": "rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_nphqminswz_gfx942_wpe2", + "sha256": "e866945043e182753a9ce9087fd65a9ed07d2e50cb795a5ff9600a5b8f5295e7", + "status": "ok" + }, + "attention_dense/gfx942/persist_causal_sq512": { + "arch": "gfx942", + "bytes": 174611, + "family": "attention_dense", + "kernel_name": "rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe2", + "sha256": "b10368fbc7e749be0f4146cde5dee5fc09be047a52dcd343282695ed01d8bad4", + "status": "ok" + }, + "attention_dense/gfx942/persist_hkvmaj_sq512": { + "arch": "gfx942", + "bytes": 178853, + "family": "attention_dense", + "kernel_name": "rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_hkvmaj_gfx942_wpe2", + "sha256": "4d09d83e48d8115affc8e493cda04e7ea3df86e5f2e9aaa57f0f9462105bf769", + "status": "ok" + }, + "attention_dense/gfx942/persist_hqswz_sq512": { + "arch": "gfx942", + "bytes": 177201, + "family": "attention_dense", + "kernel_name": "rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_hqminswz_gfx942_wpe2", + "sha256": "e245eb2bfbc2ca772a6b20d9d8a714978fb536f03cd115d2ee728bef64c73d5c", + "status": "ok" + }, "attention_dense/gfx950/bn128_sq512": { "arch": "gfx950", - "bytes": 590394, + "bytes": 591989, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn128_bf16_causal_lazyrs", - "sha256": "c7cc9fd2df45728a32c5431e0db09e69e71a8409f68b9c1a24ad7f48ded92340", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn128_bf16_causal_lazyrs_npbthkvmin", + "sha256": "75cdfe6461c367d35e3a590b63b050090a39b69191e531ee46c3d435ba23cfe0", "status": "ok" }, "attention_dense/gfx950/bottom_right_bm128_sq512": { @@ -364,18 +404,18 @@ }, "attention_dense/gfx950/default_causal_sq512": { "arch": "gfx950", - "bytes": 324104, + "bytes": 325096, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_lazyrs", - "sha256": "dfbb93589f31f072247ce8d0c41eb8f3a02f621f558479cadbc7aae68bd07917", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_lazyrs_npbthkvmin", + "sha256": "2b86578ddd20fd7756a0af7cebbc13aab55aad1abcf10db73155ae5b9f095293", "status": "ok" }, "attention_dense/gfx950/fp16_h64_sq512": { "arch": "gfx950", - "bytes": 207716, + "bytes": 208367, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_fp16_kpad8_causal_lazyrs", - "sha256": "e7e23f95703cd126968435cf8cd35999cb68af88002f083ae3583b0187bec516", + "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_fp16_kpad8_causal_lazyrs_npbthkvmin", + "sha256": "1f5d46356371ac7b1ad56a15510c1fb7066cd8d1ef491ccb38ed701a949d1ac9", "status": "ok" }, "attention_dense/gfx950/full_ragged_cross_197x400": { @@ -386,12 +426,20 @@ "sha256": "ad8ffd227e6aa5d3dfe9181bbde83f5b82a93d7f796e5e256017b644825e29ed", "status": "ok" }, + "attention_dense/gfx950/hqswz_sq512": { + "arch": "gfx950", + "bytes": 325096, + "family": "attention_dense", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_lazyrs_nphqminswz", + "sha256": "091599f50697dda3b9f9b6dd0cf94f8336ed9848d4c36ad9cfbadf0dfb63895e", + "status": "ok" + }, "attention_dense/gfx950/lazy_off_sq512": { "arch": "gfx950", - "bytes": 319738, + "bytes": 320730, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal", - "sha256": "d23fe8ec952f05a2fcdb7f62df6b51782802255360691ce2e48192d7ebb834a3", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_npbthkvmin", + "sha256": "2adb7791889a716456ebc4fbd816072fba3f2c589f78b118bf1cc03cce0af1d0", "status": "ok" }, "attention_dense/gfx950/lookahead_sq512": { @@ -404,18 +452,18 @@ }, "attention_dense/gfx950/lse_default_sq512": { "arch": "gfx950", - "bytes": 331898, + "bytes": 332890, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_lazyrs_lse", - "sha256": "41ac66f0799a777fdca1b0b564756fe72d21ee665720761549db48db42818881", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_lazyrs_lse_npbthkvmin", + "sha256": "07113b73657ddeed4a2d22b69f15bfe50c8b2fa275a7d991b76e4c4a5607c832", "status": "ok" }, "attention_dense/gfx950/lse_persist_empty_window_sinks_sq256": { "arch": "gfx950", - "bytes": 432925, + "bytes": 434564, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq256_sk64_causal_swa64_sinks_lazyrs_lse_persist256", - "sha256": "6cbd5f2ce7f1bcee8e9c5b493257cea15015246c680259595eed70ab016a0fac", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq256_sk64_causal_swa64_sinks_lazyrs_lse_persist256_bthkvmin", + "sha256": "bfb996b984c7d37d80124736338f9765ac084cdaadd28e189c0b4d064fadc614", "status": "ok" }, "attention_dense/gfx950/lse_varlen_sq512": { @@ -452,130 +500,138 @@ }, "attention_dense/gfx950/persist_h64_sq512": { "arch": "gfx950", - "bytes": 211415, + "bytes": 212669, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_bf16_kpad8_sq512_sk512_causal_lazyrs_persist256", - "sha256": "32aa0388437da6635dff712aa5647a1d7c1f282f772667b4b09a8db80eb3e639", + "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_bf16_kpad8_sq512_sk512_causal_lazyrs_persist256_bthkvmin", + "sha256": "9cd1250fb4e975b1d172ca87db0d759ccc3c487376ccbbe3d98aeba77ff0b250", "status": "ok" }, "attention_dense/gfx950/persist_hkvmaj_sq512": { "arch": "gfx950", - "bytes": 308263, + "bytes": 308669, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_lazyrs_persist256_hkvmaj", - "sha256": "9efff76c40d673fe64b2bd94b7c7941855bc0871aa8eb90cd7126c6ac5ccb847", + "sha256": "8f96ead8fc44baf050a380486d8e6aa731a7c453ca4413440d425294c8fe118b", + "status": "ok" + }, + "attention_dense/gfx950/persist_hqswz_sq512": { + "arch": "gfx950", + "bytes": 308319, + "family": "attention_dense", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_lazyrs_persist256_hqminswz", + "sha256": "758151a7bcac901e2e62005a53997d23504fdadf0b7aaeb78cf028b6b25a7d5b", "status": "ok" }, "attention_dense/gfx950/persist_intl_sq512": { "arch": "gfx950", - "bytes": 307931, + "bytes": 308205, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_lazyrs_persist256_intl", - "sha256": "48b6b78b737c5784f63f68a1dedd9aa97ead1e109761b4c0c110f300990a07e4", + "sha256": "be820045fd3eb798c96015a312d1b1be89c69cbb7f318eb239c0e8b6fa74531f", "status": "ok" }, "attention_dense/gfx950/persist_lazy_off_sq512": { "arch": "gfx950", - "bytes": 303176, + "bytes": 304647, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_persist256", - "sha256": "894f7329f1ad785344b8de7a5e27c48d18d0d7a40a60d797be963bb27812ee9d", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_persist256_bthkvmin", + "sha256": "93bf2ddb35f17495a2fe3901be8fe88e6c0fc84884316a8fec9d183ac15edc5f", "status": "ok" }, "attention_dense/gfx950/persist_lookahead_sq512": { "arch": "gfx950", - "bytes": 307875, + "bytes": 308408, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_rb64_lazyrs_persist256", - "sha256": "d1dd21117969afe8a4a1b14e2bd04a207cf58eaae741664978ecdfd763957620", + "sha256": "31df9271c1cf19b3def385aa33afe5e068cae003d28b79ccf2513607d4904806", "status": "ok" }, "attention_dense/gfx950/persist_noncausal_window_sq512": { "arch": "gfx950", - "bytes": 307955, + "bytes": 307991, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_full_swa128_lazyrs_persist256", - "sha256": "1af5c8ef1247792196428012fabe39015e40ff56ba94f0625afd5b621dd0a78d", + "sha256": "2c084455e04a01b6209ca87de94725783f3ee1831030d720b87be6f6e4cec759", "status": "ok" }, "attention_dense/gfx950/persist_qbmaj_sq512": { "arch": "gfx950", - "bytes": 307485, + "bytes": 308020, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_lazyrs_persist256", - "sha256": "fe9d3a6ea2d25b23583e587bfc18e0994a53f371681e56d7fad88c59f73d4383", + "sha256": "a611f3d8e6053f0889081c6c3fb7ea30c6e50062cb9830dfbdcee53b568e510e", "status": "ok" }, "attention_dense/gfx950/persist_ragged_sq500": { "arch": "gfx950", - "bytes": 308810, + "bytes": 309344, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq500_sk500_causal_ragged_lazyrs_persist256", - "sha256": "b083a398a4068739c67238d5df62869531c97dd08e06cddbaede40c31a43e6c9", + "sha256": "2cdbd067203a873b8380bf1d0228d9ff3978f17bc406b5564fa6dc712d575818", "status": "ok" }, "attention_dense/gfx950/persist_sinks_d128_sq512": { "arch": "gfx950", - "bytes": 308193, + "bytes": 309668, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_sinks_lazyrs_persist256", - "sha256": "e585a9e07a2e3cc2e63f114a62d1041571d92178c7d69ba029962221c6563885", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_sinks_lazyrs_persist256_bthkvmin", + "sha256": "1fe15ecce526876534f488bd43ed6f882d1abf9fd14fa8277ff58bd946c13df6", "status": "ok" }, "attention_dense/gfx950/persist_sinks_d64_sq512": { "arch": "gfx950", - "bytes": 211977, + "bytes": 213237, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_bf16_kpad8_sq512_sk512_causal_sinks_lazyrs_persist256", - "sha256": "11cc1ee108b4d383997f51b31456b60f6b7922eae6f622a0d6028327183b775e", + "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_bf16_kpad8_sq512_sk512_causal_sinks_lazyrs_persist256_bthkvmin", + "sha256": "4f5aff0dc031912a84c07267deaff53d2b1c84b0dab9260bb6a7b60ef912093b", "status": "ok" }, "attention_dense/gfx950/persist_swa_sinks_d128_sq512": { "arch": "gfx950", - "bytes": 424823, + "bytes": 424859, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_swa128_sinks_lazyrs_persist256", - "sha256": "ae76f122d197b01b5e9b408801705d9dbe19366ca0ee0994fc51e91c5670403c", + "sha256": "cd0822141c1c5424c7e638ae73012ad1ba07389ef2d0462884fb271f35629b92", "status": "ok" }, "attention_dense/gfx950/persist_swa_sinks_d64_sq512": { "arch": "gfx950", - "bytes": 296264, + "bytes": 296300, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_bf16_kpad8_sq512_sk512_causal_swa128_sinks_lazyrs_persist256", - "sha256": "a7f430db80dfa72188523bac0d4107f90bee4be69b400d1afc562af04f741283", + "sha256": "d32450c52a4fbb78818d2c34b3fa73dc4e6f4b63b8b96530d7cf3f3c1f7c6565", "status": "ok" }, "attention_dense/gfx950/persist_swa_w128_sq512": { "arch": "gfx950", - "bytes": 424015, + "bytes": 424051, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_swa128_lazyrs_persist256", - "sha256": "6d13af976b197069501fdb7f839666bd2888d2940f25de36e47e394d8b55c8da", + "sha256": "e1185c19e50e38b23aff242b0934e1b71efede8605e000b27e1b8f88353c8316", "status": "ok" }, "attention_dense/gfx950/persist_wdma_gqapair_fp16_sq512": { "arch": "gfx950", - "bytes": 263509, + "bytes": 263676, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq32_kv8_bn64_fp16_sq512_sk512_causal_lazyrs_wdma_persist16_gqapair", - "sha256": "1e991860dee3cbe1aa341dfc24b23a23684f246545cdb7fd9ab15f53fbb6f962", + "sha256": "2f08c12157db7f8639aec57753167aedf4eca3632583e699eabc75c3b1f41218", "status": "ok" }, "attention_dense/gfx950/persist_window_empty_rows_1024x512": { "arch": "gfx950", - "bytes": 424311, + "bytes": 424347, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq1024_sk512_causal_swa128_lazyrs_persist256", - "sha256": "5d8a5d6e3b890fa481e1e09493bbeedf14b9124e5a48ce31a02703a0ba8199fd", + "sha256": "e2f3aca3da4edd5ac03377f2a3c117b3bb0be0f01946d8264523c7f49f1ca782", "status": "ok" }, "attention_dense/gfx950/persistent_causal_sq512": { "arch": "gfx950", - "bytes": 307485, + "bytes": 308956, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_lazyrs_persist256", - "sha256": "fe9d3a6ea2d25b23583e587bfc18e0994a53f371681e56d7fad88c59f73d4383", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_lazyrs_persist256_bthkvmin", + "sha256": "01a99c4e15a586806fc7e7757fdc697be213df04014679022ac2ce8ac6f1d64e", "status": "ok" }, "attention_dense/gfx950/ragged_bottom_right_sq500": { @@ -612,18 +668,18 @@ }, "attention_dense/gfx950/sinks_d128_sq512": { "arch": "gfx950", - "bytes": 324843, + "bytes": 325802, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_sinks_lazyrs", - "sha256": "a2054abc28fa55351b4839a2b0d6583618aafa6d2fee012acc4447e6fa3ee681", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_sinks_lazyrs_npbthkvmin", + "sha256": "92a3eb6c678fe656c99945deeaf084bba40a0e6d8d48a1b7b2d447b9cc0d9191", "status": "ok" }, "attention_dense/gfx950/sinks_d64_sq512": { "arch": "gfx950", - "bytes": 211288, + "bytes": 211955, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_bf16_kpad8_causal_sinks_lazyrs", - "sha256": "59ad975aec0d4c6a0547d1e10583c145ffe54d87fcdabd7279a923ba20d0791b", + "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_bf16_kpad8_causal_sinks_lazyrs_npbthkvmin", + "sha256": "ff4b65d857d38521acbba1832c28585eda6150978ec4bf68a60ee6fc01c34634", "status": "ok" }, "attention_dense/gfx950/swa_sinks_d128_sq512": { @@ -2173,15 +2229,15 @@ "gfx1201": 8, "gfx1250": 12, "gfx90a": 5, - "gfx942": 52, - "gfx950": 178 + "gfx942": 57, + "gfx950": 180 }, "expected_failures": 6, "families": { "attention": 38, "attention_d128_swa": 1, "attention_d256": 2, - "attention_dense": 44, + "attention_dense": 51, "conv": 20, "conv_dgrad": 13, "conv_direct": 26, @@ -2201,7 +2257,7 @@ "moe": 9, "nontemporal": 2 }, - "ok_cases": 267 + "ok_cases": 274 } }, "llvm22": { @@ -2534,12 +2590,52 @@ "sha256": "cf3163aba7b796909e402cf81ca203af1f865ba4343aec58d9bdeb228df731ab", "status": "ok" }, + "attention_dense/gfx942/causal_sq512": { + "arch": "gfx942", + "bytes": 171686, + "family": "attention_dense", + "kernel_name": "rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_npbthkvmin_gfx942_wpe2", + "sha256": "2658a44f0016de96754d17f3dff885a6c6d018ef70629ee728410050218b3c24", + "status": "ok" + }, + "attention_dense/gfx942/hqswz_sq512": { + "arch": "gfx942", + "bytes": 171712, + "family": "attention_dense", + "kernel_name": "rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_nphqminswz_gfx942_wpe2", + "sha256": "506a3887ffd5f612d9165126be44d70a52a4d2a1f188cfd3d0b273d3e84495b6", + "status": "ok" + }, + "attention_dense/gfx942/persist_causal_sq512": { + "arch": "gfx942", + "bytes": 174707, + "family": "attention_dense", + "kernel_name": "rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe2", + "sha256": "a394730c2571e12aac91bed2f7d2dea724c5c9a976af38550eacb72141a0bf6d", + "status": "ok" + }, + "attention_dense/gfx942/persist_hkvmaj_sq512": { + "arch": "gfx942", + "bytes": 178949, + "family": "attention_dense", + "kernel_name": "rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_hkvmaj_gfx942_wpe2", + "sha256": "7648f077ef297d2a2638d8cd4702a815dd6e24f4695ac24f6f3706f48d7b7e26", + "status": "ok" + }, + "attention_dense/gfx942/persist_hqswz_sq512": { + "arch": "gfx942", + "bytes": 177297, + "family": "attention_dense", + "kernel_name": "rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_hqminswz_gfx942_wpe2", + "sha256": "d69fffc1ae57a98d2fbbbaa96a3a1e20b59181ccf68972324b8ebf772aa57415", + "status": "ok" + }, "attention_dense/gfx950/bn128_sq512": { "arch": "gfx950", - "bytes": 590494, + "bytes": 592089, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn128_bf16_causal_lazyrs", - "sha256": "43af47065664c7b0871b6ffa41bec596184187d5986c37e4c30ec6df3044a088", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn128_bf16_causal_lazyrs_npbthkvmin", + "sha256": "c464c39a5b20a7b5f2f4e6e0d65cee4a5d41160fe90b0b8e785552405127ba16", "status": "ok" }, "attention_dense/gfx950/bottom_right_bm128_sq512": { @@ -2568,18 +2664,18 @@ }, "attention_dense/gfx950/default_causal_sq512": { "arch": "gfx950", - "bytes": 324204, + "bytes": 325196, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_lazyrs", - "sha256": "dfb6146f3ec997437f6cdc68058cf6e8d17cf08cd6342d8d9b8ac306964f4e32", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_lazyrs_npbthkvmin", + "sha256": "19bb0e66a369af1f3a5b926e2925db7e4af1cbb1abf5a80510c9bfdac34bce48", "status": "ok" }, "attention_dense/gfx950/fp16_h64_sq512": { "arch": "gfx950", - "bytes": 207812, + "bytes": 208463, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_fp16_kpad8_causal_lazyrs", - "sha256": "0e1ea2912e4171f10f9e31094cb3089085ef560e5b7031e86b21f225e332d088", + "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_fp16_kpad8_causal_lazyrs_npbthkvmin", + "sha256": "dbc1f781072770d36c8e0fafbcb9e7206a987499bc0249a51965567f5cede493", "status": "ok" }, "attention_dense/gfx950/full_ragged_cross_197x400": { @@ -2590,12 +2686,20 @@ "sha256": "fc9f47a1dbd2433baf4da66be2bb098b85c768cdc60b3a9937f5b2e9fd3a7d3d", "status": "ok" }, + "attention_dense/gfx950/hqswz_sq512": { + "arch": "gfx950", + "bytes": 325196, + "family": "attention_dense", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_lazyrs_nphqminswz", + "sha256": "b3800a0d94c38ebd1c6e3c1ef5f131f59f2d5740b5a7ab072c9805a12219e9e8", + "status": "ok" + }, "attention_dense/gfx950/lazy_off_sq512": { "arch": "gfx950", - "bytes": 319838, + "bytes": 320830, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal", - "sha256": "fc9d1e8365218673e8bddc3b5d608b86bd1f5d56b7cec8b6baccf307c4423b8a", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_npbthkvmin", + "sha256": "9374159fe9732562a3652587016f808415936c79cb8230e453e9666c830a34e8", "status": "ok" }, "attention_dense/gfx950/lookahead_sq512": { @@ -2608,18 +2712,18 @@ }, "attention_dense/gfx950/lse_default_sq512": { "arch": "gfx950", - "bytes": 331998, + "bytes": 332990, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_lazyrs_lse", - "sha256": "4138bc34d657a3d50c83b8a11c46c2c51ef15b02c5a4f91ffefb93d3ea5db36c", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_lazyrs_lse_npbthkvmin", + "sha256": "e43160487487c229bbc18e74c79a607ea9168e7d674860fde26f6af0e8ab34e0", "status": "ok" }, "attention_dense/gfx950/lse_persist_empty_window_sinks_sq256": { "arch": "gfx950", - "bytes": 433025, + "bytes": 434664, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq256_sk64_causal_swa64_sinks_lazyrs_lse_persist256", - "sha256": "d378d3cb3a85ca2fd8ab0e13480907d3316f903c42d478c8e901d46587b45236", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq256_sk64_causal_swa64_sinks_lazyrs_lse_persist256_bthkvmin", + "sha256": "3130138ebdc1190ed56e5003efb67020ff237e614359e55ceefc84c63d2b90e9", "status": "ok" }, "attention_dense/gfx950/lse_varlen_sq512": { @@ -2656,130 +2760,138 @@ }, "attention_dense/gfx950/persist_h64_sq512": { "arch": "gfx950", - "bytes": 211511, + "bytes": 212765, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_bf16_kpad8_sq512_sk512_causal_lazyrs_persist256", - "sha256": "68116b7a779cf880b2df951277a600b5cd7160176f34bb96a09be824bc9f1283", + "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_bf16_kpad8_sq512_sk512_causal_lazyrs_persist256_bthkvmin", + "sha256": "2be04f91348450a4f378deace520bcaebb3875f7ad91a45416cef73985b7d871", "status": "ok" }, "attention_dense/gfx950/persist_hkvmaj_sq512": { "arch": "gfx950", - "bytes": 308363, + "bytes": 308769, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_lazyrs_persist256_hkvmaj", - "sha256": "0f4f04a065849838a3eb4cffa97cca3d337c483253595b6e117e40b3190477d8", + "sha256": "64039c6e1113c45b83a1f4655552afc09e5c35b05438fabf4dad0958ef4f0fb1", + "status": "ok" + }, + "attention_dense/gfx950/persist_hqswz_sq512": { + "arch": "gfx950", + "bytes": 308419, + "family": "attention_dense", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_lazyrs_persist256_hqminswz", + "sha256": "e9459964e2c614ee0f1f09f1f73e6898c3cc0784147d513926ed77e95acff78b", "status": "ok" }, "attention_dense/gfx950/persist_intl_sq512": { "arch": "gfx950", - "bytes": 308031, + "bytes": 308305, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_lazyrs_persist256_intl", - "sha256": "863f766c7131a0ea8cecfc0a5458e462505f5492d92e885b2f76063c8a6a05f5", + "sha256": "ad5788f3268c4732582a9719d91dcca8d5cc40fb78d411fa7515eb4c96993370", "status": "ok" }, "attention_dense/gfx950/persist_lazy_off_sq512": { "arch": "gfx950", - "bytes": 303276, + "bytes": 304747, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_persist256", - "sha256": "f10fe9efca043e4c045ba546b6334aff2a4529183c64eae5309cde79e0315e28", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_persist256_bthkvmin", + "sha256": "e43edd13ee400222dd66ce4bd75ebf34661200b60b1ff6600d3c793f5c625132", "status": "ok" }, "attention_dense/gfx950/persist_lookahead_sq512": { "arch": "gfx950", - "bytes": 307975, + "bytes": 308508, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_rb64_lazyrs_persist256", - "sha256": "b89eedd970eeb353ca2c9b5edfc2db179b745a68767cf870aa35cba74df6978b", + "sha256": "06f256af69f842a1dc9dd610c68392ba70ef31b7bfdbf9a2ccdb2c05f58cf003", "status": "ok" }, "attention_dense/gfx950/persist_noncausal_window_sq512": { "arch": "gfx950", - "bytes": 308055, + "bytes": 308091, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_full_swa128_lazyrs_persist256", - "sha256": "6124395ebe1ef4727b327417d75bde327558eb4877599610fd68c20281643e26", + "sha256": "56e699e9e5402a46d071b30c512e57b608b18b68c6a7607e6f5fe9ffbd0897f0", "status": "ok" }, "attention_dense/gfx950/persist_qbmaj_sq512": { "arch": "gfx950", - "bytes": 307585, + "bytes": 308120, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_lazyrs_persist256", - "sha256": "056626a53d1457ca7bcbc11949e3ea85aa4a2114fde8e212ea6beeb7a3c54348", + "sha256": "044b538645d60264264bb5158ab9f30278e1ed5e1b5f1b4276ee54eacacfc94a", "status": "ok" }, "attention_dense/gfx950/persist_ragged_sq500": { "arch": "gfx950", - "bytes": 308955, + "bytes": 309489, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq500_sk500_causal_ragged_lazyrs_persist256", - "sha256": "689e8ea1b0aee258399bcb919e0fcc0cbf16f7244e2893bcc77e674946723651", + "sha256": "34e94aac5052f3588a8231e28f142694d0ebb94b1bbcd1490ef2831418493f8a", "status": "ok" }, "attention_dense/gfx950/persist_sinks_d128_sq512": { "arch": "gfx950", - "bytes": 308293, + "bytes": 309768, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_sinks_lazyrs_persist256", - "sha256": "cc1d21882f1a2fcf5aa9ecb15ab9a10d633ef30fc9e3fa13874d3306beb6e987", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_sinks_lazyrs_persist256_bthkvmin", + "sha256": "9f9d2ff0b7b4277d947023b17eb5b254869ce447b16584b365356cd8572230cc", "status": "ok" }, "attention_dense/gfx950/persist_sinks_d64_sq512": { "arch": "gfx950", - "bytes": 212073, + "bytes": 213333, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_bf16_kpad8_sq512_sk512_causal_sinks_lazyrs_persist256", - "sha256": "0b76cbc48849557311851be0befae7639e4a4809f45a9b104907582be3f30e22", + "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_bf16_kpad8_sq512_sk512_causal_sinks_lazyrs_persist256_bthkvmin", + "sha256": "c0434f12837856998a160887a98bdf170abf6a0f6f1b112fcfe3e47b98ca7e25", "status": "ok" }, "attention_dense/gfx950/persist_swa_sinks_d128_sq512": { "arch": "gfx950", - "bytes": 424923, + "bytes": 424959, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_swa128_sinks_lazyrs_persist256", - "sha256": "c1cb8d8b671aa02192d48e6089aeb3fe1faedec846e96638b9833f3c499153df", + "sha256": "636f64458986d83717a08ca1cb54138055fe54b802c2de3040bc59f330ab48ce", "status": "ok" }, "attention_dense/gfx950/persist_swa_sinks_d64_sq512": { "arch": "gfx950", - "bytes": 296360, + "bytes": 296396, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_bf16_kpad8_sq512_sk512_causal_swa128_sinks_lazyrs_persist256", - "sha256": "445f98cec6971306906bb1ed01a88bc4f63cd0b3049abc2965e264c2f000596d", + "sha256": "de5e0e013aa79665d563b2b294cf4535ed47215051ae2d919b2a34dd5749f4b3", "status": "ok" }, "attention_dense/gfx950/persist_swa_w128_sq512": { "arch": "gfx950", - "bytes": 424115, + "bytes": 424151, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_swa128_lazyrs_persist256", - "sha256": "3b2f060e7ab35446ae3bfc4c17e042d82700eb4cfaf007f7e09a9c51c3e043d4", + "sha256": "b483815a935e28dbf8616293ded03c322cf859d1792964523ddf6e450b22ed06", "status": "ok" }, "attention_dense/gfx950/persist_wdma_gqapair_fp16_sq512": { "arch": "gfx950", - "bytes": 263609, + "bytes": 263776, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq32_kv8_bn64_fp16_sq512_sk512_causal_lazyrs_wdma_persist16_gqapair", - "sha256": "cad7d10ee1dacc38b8d0995654eb01674288ff17c6b264a2d5ef9ddd02449182", + "sha256": "a341916b990513a4c2a30ea06658bc06c6eb428d7350ba92ed9ed3aaf3eaed03", "status": "ok" }, "attention_dense/gfx950/persist_window_empty_rows_1024x512": { "arch": "gfx950", - "bytes": 424411, + "bytes": 424447, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq1024_sk512_causal_swa128_lazyrs_persist256", - "sha256": "8770edfc25915fd3c4714b4a2197469ca1946906f23db7e4b3f9d610d01d4375", + "sha256": "04c006d28344e7cdaca55c90312ca671bf443d98987d8163ada60b88ef7692e6", "status": "ok" }, "attention_dense/gfx950/persistent_causal_sq512": { "arch": "gfx950", - "bytes": 307585, + "bytes": 309056, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_lazyrs_persist256", - "sha256": "056626a53d1457ca7bcbc11949e3ea85aa4a2114fde8e212ea6beeb7a3c54348", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_lazyrs_persist256_bthkvmin", + "sha256": "443122b2417d4d876d081c91b6496c2b6a94b578dcd09ab9896fb1c7ff2ccd10", "status": "ok" }, "attention_dense/gfx950/ragged_bottom_right_sq500": { @@ -2816,18 +2928,18 @@ }, "attention_dense/gfx950/sinks_d128_sq512": { "arch": "gfx950", - "bytes": 324943, + "bytes": 325902, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_sinks_lazyrs", - "sha256": "8a8518f03051013fa1e2e36117d5178b1682108f9073d1abfffc441162c90fcf", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_sinks_lazyrs_npbthkvmin", + "sha256": "d92d8c2fd919540a44df58ac829aea7d77c689afe4175d0a3331016a01ed8ed5", "status": "ok" }, "attention_dense/gfx950/sinks_d64_sq512": { "arch": "gfx950", - "bytes": 211384, + "bytes": 212051, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_bf16_kpad8_causal_sinks_lazyrs", - "sha256": "b19c4a35746652e4c9c68c2c024e2b210dc5b2f65d02922854f40a7ca862482b", + "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_bf16_kpad8_causal_sinks_lazyrs_npbthkvmin", + "sha256": "89a7d49b71784c6995d47a67184a238a3a64c72283fa9be0409bc20496c5b11c", "status": "ok" }, "attention_dense/gfx950/swa_sinks_d128_sq512": { @@ -4377,15 +4489,15 @@ "gfx1201": 8, "gfx1250": 12, "gfx90a": 5, - "gfx942": 52, - "gfx950": 178 + "gfx942": 57, + "gfx950": 180 }, "expected_failures": 6, "families": { "attention": 38, "attention_d128_swa": 1, "attention_d256": 2, - "attention_dense": 44, + "attention_dense": 51, "conv": 20, "conv_dgrad": 13, "conv_direct": 26, @@ -4405,7 +4517,7 @@ "moe": 9, "nontemporal": 2 }, - "ok_cases": 267 + "ok_cases": 274 } }, "llvm23": { @@ -4738,12 +4850,52 @@ "sha256": "e20ccb61ab8ca4ba8413c69ba4dec0f50f9ec7025db28b15e608a7db0672729f", "status": "ok" }, + "attention_dense/gfx942/causal_sq512": { + "arch": "gfx942", + "bytes": 171750, + "family": "attention_dense", + "kernel_name": "rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_npbthkvmin_gfx942_wpe2", + "sha256": "bc12b9a1aba7d9bfae67d64c24339d006cec293c40e31a9f41e33aefbe7981eb", + "status": "ok" + }, + "attention_dense/gfx942/hqswz_sq512": { + "arch": "gfx942", + "bytes": 171776, + "family": "attention_dense", + "kernel_name": "rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_nphqminswz_gfx942_wpe2", + "sha256": "4a1a78f21b2b8db16bd7033c211ad36cbbe479fe7a29db3e7ad7ba5a0adf8bfb", + "status": "ok" + }, + "attention_dense/gfx942/persist_causal_sq512": { + "arch": "gfx942", + "bytes": 174771, + "family": "attention_dense", + "kernel_name": "rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe2", + "sha256": "22798644d75f8d43cf7ed1e0db84408b8aff160a4597771c849f1932878398bb", + "status": "ok" + }, + "attention_dense/gfx942/persist_hkvmaj_sq512": { + "arch": "gfx942", + "bytes": 179013, + "family": "attention_dense", + "kernel_name": "rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_hkvmaj_gfx942_wpe2", + "sha256": "9d95fa2e4cca8c46fda343e891bc1b177d1ea86a61fdfe8c0ab0fdfc36565475", + "status": "ok" + }, + "attention_dense/gfx942/persist_hqswz_sq512": { + "arch": "gfx942", + "bytes": 177361, + "family": "attention_dense", + "kernel_name": "rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_hqminswz_gfx942_wpe2", + "sha256": "9ac2689a1c2b0572551c4c87597a97b11f0984a64ddbbaa4063d6862fce0c5b2", + "status": "ok" + }, "attention_dense/gfx950/bn128_sq512": { "arch": "gfx950", - "bytes": 590558, + "bytes": 592153, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn128_bf16_causal_lazyrs", - "sha256": "610e90920049fc8d5d63c78dcd9f82bae5eb40e18ff9df985a80e262d2315c3f", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn128_bf16_causal_lazyrs_npbthkvmin", + "sha256": "42b53154e1a096666d99691df4edd655e0cf2c15fc350f4006bb50fb8a48906f", "status": "ok" }, "attention_dense/gfx950/bottom_right_bm128_sq512": { @@ -4772,18 +4924,18 @@ }, "attention_dense/gfx950/default_causal_sq512": { "arch": "gfx950", - "bytes": 324268, + "bytes": 325260, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_lazyrs", - "sha256": "f348b2ad956842abaaea08f498fda04fd1cb647d292098446f1d8fc5852a898d", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_lazyrs_npbthkvmin", + "sha256": "d6e9a5050722f3187e0394c2aefa80385ef12406054656e6ca2dd3d4aa582429", "status": "ok" }, "attention_dense/gfx950/fp16_h64_sq512": { "arch": "gfx950", - "bytes": 207876, + "bytes": 208527, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_fp16_kpad8_causal_lazyrs", - "sha256": "513d286add701de6801e9e6d64c4ca29706cb673e8d3fa1794f6fff2e10b2a55", + "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_fp16_kpad8_causal_lazyrs_npbthkvmin", + "sha256": "60626bc3b241a9db7636d574304ce1b0d55b3f1fa1d28b4920ea4c1138d11b63", "status": "ok" }, "attention_dense/gfx950/full_ragged_cross_197x400": { @@ -4794,12 +4946,20 @@ "sha256": "304595cfd13fbc1ad44ac20986974ec7e6442d88f126ac007ddc720c6e7091e8", "status": "ok" }, + "attention_dense/gfx950/hqswz_sq512": { + "arch": "gfx950", + "bytes": 325260, + "family": "attention_dense", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_lazyrs_nphqminswz", + "sha256": "31d6b442f10c18fc1f5833be9ee32ac60a02688683a847697a5b33d266566def", + "status": "ok" + }, "attention_dense/gfx950/lazy_off_sq512": { "arch": "gfx950", - "bytes": 319902, + "bytes": 320894, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal", - "sha256": "85c0f328bdbeeab4e630aebada79aa2413af2718b4b153c0d62e3f9012344456", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_npbthkvmin", + "sha256": "9361d758bee30253f176ac382ef74a5ff03d34a51e3853cb658ab95440a8ff23", "status": "ok" }, "attention_dense/gfx950/lookahead_sq512": { @@ -4812,18 +4972,18 @@ }, "attention_dense/gfx950/lse_default_sq512": { "arch": "gfx950", - "bytes": 332062, + "bytes": 333054, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_lazyrs_lse", - "sha256": "2a9393101e236a4518c2974f5ec2809baa743b71fd136b672730e57d2add9b6e", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_lazyrs_lse_npbthkvmin", + "sha256": "9e1cb463f9a1dd64eef7ea61072c5bd4812217b0c41d2717576766eecdc458ed", "status": "ok" }, "attention_dense/gfx950/lse_persist_empty_window_sinks_sq256": { "arch": "gfx950", - "bytes": 433089, + "bytes": 434728, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq256_sk64_causal_swa64_sinks_lazyrs_lse_persist256", - "sha256": "086f1582d0037d3760c0a21eb60bf87c71da4692d0061ecbac027098efac6461", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq256_sk64_causal_swa64_sinks_lazyrs_lse_persist256_bthkvmin", + "sha256": "fe6df270dfb9fa15ef18ad7c6e7357fcd7ce1c5be0541f0a0d3aa5979f1d5776", "status": "ok" }, "attention_dense/gfx950/lse_varlen_sq512": { @@ -4860,130 +5020,138 @@ }, "attention_dense/gfx950/persist_h64_sq512": { "arch": "gfx950", - "bytes": 211575, + "bytes": 212829, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_bf16_kpad8_sq512_sk512_causal_lazyrs_persist256", - "sha256": "b6ce349ead2af7fc181ddecc3837ea53dd78f3977929adba4c75c46d3ca578c5", + "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_bf16_kpad8_sq512_sk512_causal_lazyrs_persist256_bthkvmin", + "sha256": "765db5314152cc1dae544474adb236ba92b2114467f803d1d6bb887ae5da2ee3", "status": "ok" }, "attention_dense/gfx950/persist_hkvmaj_sq512": { "arch": "gfx950", - "bytes": 308427, + "bytes": 308833, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_lazyrs_persist256_hkvmaj", - "sha256": "49a8caf64d2f6344e3e1139365a5098ad5914c94be9f29f285adb21a5dfbe3ab", + "sha256": "80d77c90ffd776ef8404b22437b91a21dcc4357fe814486e034d981aec2d0017", + "status": "ok" + }, + "attention_dense/gfx950/persist_hqswz_sq512": { + "arch": "gfx950", + "bytes": 308483, + "family": "attention_dense", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_lazyrs_persist256_hqminswz", + "sha256": "08402632a00c4a97cf0d9f3551c4a03560949da70854aaf206777feb08f45eba", "status": "ok" }, "attention_dense/gfx950/persist_intl_sq512": { "arch": "gfx950", - "bytes": 308095, + "bytes": 308369, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_lazyrs_persist256_intl", - "sha256": "013de7a93659c6389b50d6f76913a7de67ccfb11a015e586e10a61150d41c321", + "sha256": "b7312386b83704e63dd293e6cdd3d69491efe7c70f9034ae7ef52e2b7d150e5a", "status": "ok" }, "attention_dense/gfx950/persist_lazy_off_sq512": { "arch": "gfx950", - "bytes": 303340, + "bytes": 304811, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_persist256", - "sha256": "d5a172f72f9a26388558a14203823793ac4adc8a637692c3eb4635ce3db0bb82", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_persist256_bthkvmin", + "sha256": "add24eee4294effb24c40c241174b0b9db3a3cfa65f588dc9c61c2dae32bef2a", "status": "ok" }, "attention_dense/gfx950/persist_lookahead_sq512": { "arch": "gfx950", - "bytes": 308039, + "bytes": 308572, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_rb64_lazyrs_persist256", - "sha256": "bb0836bbdfea8f44c61f6aebb1b6119d932e087ce681d770d15e55b3f6bfcd02", + "sha256": "84433767cbd7541fa00d1e3f75eb7cff4a4a15d739190eaae13d8857521aa5bf", "status": "ok" }, "attention_dense/gfx950/persist_noncausal_window_sq512": { "arch": "gfx950", - "bytes": 308119, + "bytes": 308155, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_full_swa128_lazyrs_persist256", - "sha256": "e357edac22b3f6dc076cbfc68c12bffb0a1ecdb7285e094996bab7f3c47205cc", + "sha256": "e0ff11bc3886852a3c8781157cc2ce2c17b1e423348c0329597a6a8c5cc4e308", "status": "ok" }, "attention_dense/gfx950/persist_qbmaj_sq512": { "arch": "gfx950", - "bytes": 307649, + "bytes": 308184, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_lazyrs_persist256", - "sha256": "5929ed47c8d4f0e7fbc6efdca2769b3f74732737640ec1478a5c25784101a4ae", + "sha256": "9fefcc6985f6c70aaeaf27f6a851c5f18e5400300f5900faa6741918b290626c", "status": "ok" }, "attention_dense/gfx950/persist_ragged_sq500": { "arch": "gfx950", - "bytes": 309019, + "bytes": 309553, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq500_sk500_causal_ragged_lazyrs_persist256", - "sha256": "f49911c945348de95f915d48fe36b85ee3ff57c1601061fb7caf014fde7f3960", + "sha256": "06af111eda447f2767162cadb2159b4dc5440927b644a57960bb388df2c6880e", "status": "ok" }, "attention_dense/gfx950/persist_sinks_d128_sq512": { "arch": "gfx950", - "bytes": 308357, + "bytes": 309832, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_sinks_lazyrs_persist256", - "sha256": "2703e73ce331ce14b7322c4e1150172651640342809f02706750d21ca3ce3ee6", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_sinks_lazyrs_persist256_bthkvmin", + "sha256": "db0df262331f7d8205d1cdf995804990562f6a27fd18100d78efe596781e885e", "status": "ok" }, "attention_dense/gfx950/persist_sinks_d64_sq512": { "arch": "gfx950", - "bytes": 212137, + "bytes": 213397, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_bf16_kpad8_sq512_sk512_causal_sinks_lazyrs_persist256", - "sha256": "724fa59c12322b7e3ec575459063a629e9098a3f2c7e84919e71d481a48192f9", + "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_bf16_kpad8_sq512_sk512_causal_sinks_lazyrs_persist256_bthkvmin", + "sha256": "b5d9fd0ac457b205080d1cd23167ed9c3eee472d79b22b6cbc56b5a285e99839", "status": "ok" }, "attention_dense/gfx950/persist_swa_sinks_d128_sq512": { "arch": "gfx950", - "bytes": 424987, + "bytes": 425023, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_swa128_sinks_lazyrs_persist256", - "sha256": "080e4171e5839df40394e2234cb83a93f219fda3fe0f715d3e81233ce547569a", + "sha256": "4f755f21b05f25dc8d398b38b94907838e699042a8c15929a4e71f8a89c9d486", "status": "ok" }, "attention_dense/gfx950/persist_swa_sinks_d64_sq512": { "arch": "gfx950", - "bytes": 296424, + "bytes": 296460, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_bf16_kpad8_sq512_sk512_causal_swa128_sinks_lazyrs_persist256", - "sha256": "4a27c2966ce3d9da52ce5bba6e536a2d09e67e6fcd4a6bb8185a35e8b000dc4f", + "sha256": "5e14421578949d5747c345903dde6fdcb694fe2678a6236060365c71e4999883", "status": "ok" }, "attention_dense/gfx950/persist_swa_w128_sq512": { "arch": "gfx950", - "bytes": 424179, + "bytes": 424215, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_swa128_lazyrs_persist256", - "sha256": "e957084c75b6ba162816dbcee74184da74f822ab4622ab9d922f927e271c0278", + "sha256": "f3a12d195c8c3d4d2b67e434812cebaa336bf33c242d18899386c465f215f615", "status": "ok" }, "attention_dense/gfx950/persist_wdma_gqapair_fp16_sq512": { "arch": "gfx950", - "bytes": 263673, + "bytes": 263840, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq32_kv8_bn64_fp16_sq512_sk512_causal_lazyrs_wdma_persist16_gqapair", - "sha256": "9333f5fa7f31cac10dd19f75916e2d7e6279f8c09711409a168dd52a1b7d4dab", + "sha256": "d39bbc78962eed76af2b5050e0768d99957825d6744aedde0d8a33ebb02a6e3c", "status": "ok" }, "attention_dense/gfx950/persist_window_empty_rows_1024x512": { "arch": "gfx950", - "bytes": 424475, + "bytes": 424511, "family": "attention_dense", "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq1024_sk512_causal_swa128_lazyrs_persist256", - "sha256": "0800784ea22338cd886217e18a74dce8dd5badc214c80990c6006bd5baa3817a", + "sha256": "c1fda97dc2d8bb79c2e3276074b24b197969d68e9ac76cd0fe849d276c6f83ab", "status": "ok" }, "attention_dense/gfx950/persistent_causal_sq512": { "arch": "gfx950", - "bytes": 307649, + "bytes": 309120, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_lazyrs_persist256", - "sha256": "5929ed47c8d4f0e7fbc6efdca2769b3f74732737640ec1478a5c25784101a4ae", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_sq512_sk512_causal_lazyrs_persist256_bthkvmin", + "sha256": "0ff123ec789a95786922b47219f4c933430510218bdc4fc41ac2e9672a246525", "status": "ok" }, "attention_dense/gfx950/ragged_bottom_right_sq500": { @@ -5020,18 +5188,18 @@ }, "attention_dense/gfx950/sinks_d128_sq512": { "arch": "gfx950", - "bytes": 325007, + "bytes": 325966, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_sinks_lazyrs", - "sha256": "c3cde27d4e436085540aec401b6bfbd89cb71f5190a74366109e6f9516692e4a", + "kernel_name": "rocke_attention_dense_d128_hq128_kv8_bn64_bf16_causal_sinks_lazyrs_npbthkvmin", + "sha256": "918ec3e4aa184255bb90eca7b161e0f839c94a67a5c20b3df088d3c862dfbe9b", "status": "ok" }, "attention_dense/gfx950/sinks_d64_sq512": { "arch": "gfx950", - "bytes": 211448, + "bytes": 212115, "family": "attention_dense", - "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_bf16_kpad8_causal_sinks_lazyrs", - "sha256": "0d90cd11620570a936b41fa09e1ab1f44cc0067be8038d27f9f6a9f20513922c", + "kernel_name": "rocke_attention_dense_d64_hq128_kv8_bn64_bf16_kpad8_causal_sinks_lazyrs_npbthkvmin", + "sha256": "ecf4b1c7db03322f30c7f141551c2f23061d9521f0e5c61fdb5949fc4801b6f7", "status": "ok" }, "attention_dense/gfx950/swa_sinks_d128_sq512": { @@ -6604,15 +6772,15 @@ "gfx1201": 8, "gfx1250": 18, "gfx90a": 5, - "gfx942": 52, - "gfx950": 178 + "gfx942": 57, + "gfx950": 180 }, "expected_failures": 0, "families": { "attention": 38, "attention_d128_swa": 1, "attention_d256": 2, - "attention_dense": 44, + "attention_dense": 51, "conv": 20, "conv_dgrad": 13, "conv_direct": 26, @@ -6633,7 +6801,7 @@ "moe": 9, "nontemporal": 2 }, - "ok_cases": 273 + "ok_cases": 280 } } }, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py index 7529d7b5b757..5dbc6bfda184 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py @@ -548,6 +548,31 @@ def _build(): return _build +def build_attention_dense_gfx942(**over): + """Dense flash-attn prefill spec (library ``kernels/gfx942/attention_dense``).""" + + def _build(): + from kernels.gfx942.attention_dense import ( + Gfx942AttentionDenseSpec, + build_attention_dense as _build_dense, + ) + + spec = dict( + batch=1, + seqlen_q=512, + seqlen_kv=512, + num_query_heads=32, + num_kv_heads=8, + head_size=128, + causal=True, + dtype="bf16", + ) + spec.update(over) + return _build_dense(Gfx942AttentionDenseSpec(**spec), arch="gfx942") + + return _build + + def build_kda_chunkwise_gfx950(kind, arch, **over): """gfx950 chunkwise KDA kernel from a representative spec. @@ -3399,6 +3424,7 @@ def add(family, case_id, arch, build): ("fp16_h64_sq512", {"dtype": "fp16", "head_size": 64}), ("bn128_sq512", {"block_n": 128}), ("noncausal_sq512", {"causal": False}), + ("hqswz_sq512", {"nonpersist_decode": "hq_minor_swz"}), # --- bottom-right diagonal. Four cases pin the aligned and arbitrary # shifted-diagonal routes, the sink composition, and the BM128 geometry. ( @@ -3440,6 +3466,14 @@ def add(family, case_id, arch, build): "persist_hkvmaj_sq512", {"persistent": True, "num_persistent": 256, "persist_decode": "hkv_major"}, ), + ( + "persist_hqswz_sq512", + { + "persistent": True, + "num_persistent": 256, + "persist_decode": "hq_minor_swz", + }, + ), ( "persist_intl_sq512", {"persistent": True, "num_persistent": 256, "interleave": True}, @@ -3622,6 +3656,23 @@ def add(family, case_id, arch, build): build_attention_dense("gfx950", **_over), ) + # Dense flash-attn prefill on gfx942: the block orders of both grids (the + # library's gfx942 dense golden pins more variants, but outside this gate). + _persist942 = {"persistent": True, "num_persistent": 304} + for _variant, _over in ( + ("causal_sq512", {}), + ("hqswz_sq512", {"nonpersist_decode": "hq_minor_swz"}), + ("persist_causal_sq512", _persist942), + ("persist_hqswz_sq512", {**_persist942, "persist_decode": "hq_minor_swz"}), + ("persist_hkvmaj_sq512", {**_persist942, "persist_decode": "hkv_major"}), + ): + add( + "attention_dense", + f"attention_dense/gfx942/{_variant}", + "gfx942", + build_attention_dense_gfx942(**_over), + ) + # Attention band on the unified kernels: top-left alignment, full attention # (right_bound=-1), lookahead and non-causal / two-sided windows. These bodies have # no C++ mirror yet, so this golden is the only pin of their lowered IR. diff --git a/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/descriptors/rocKE/gfx950_attention_dense/artifacts.dvc b/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/descriptors/rocKE/gfx950_attention_dense/artifacts.dvc index 8fb02fcce76a..2e3eeb0f7373 100644 --- a/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/descriptors/rocKE/gfx950_attention_dense/artifacts.dvc +++ b/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/descriptors/rocKE/gfx950_attention_dense/artifacts.dvc @@ -1,6 +1,6 @@ outs: -- md5: 106ebd4caeba7b508fa33fcd8383c130.dir - size: 1221062 +- md5: 6ad8e5d93e5fa3c1523dea593bffb980.dir + size: 1257182 nfiles: 1 hash: md5 path: artifacts diff --git a/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/Gfx950AttentionDenseGeometry.hpp b/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/Gfx950AttentionDenseGeometry.hpp index 7ac97648f8fa..39056b03a96b 100644 --- a/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/Gfx950AttentionDenseGeometry.hpp +++ b/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/Gfx950AttentionDenseGeometry.hpp @@ -140,7 +140,7 @@ struct Gfx950AttentionDenseGeometry /** * @brief The launch geometry for one selected candidate over one graph. * - * Mirrors `attention_dense_grid`: + * Mirrors `attention_dense_grid` for the qb_minor non-persistent block order: * * nqb = (spec.seqlen_q + spec.block_m - 1) // spec.block_m * return (nqb, spec.num_query_heads, spec.batch) diff --git a/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionDenseGeometry.cpp b/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionDenseGeometry.cpp index 1bd688da2f84..3e4157c4a193 100644 --- a/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionDenseGeometry.cpp +++ b/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionDenseGeometry.cpp @@ -109,8 +109,9 @@ TEST(TestGfx950AttentionDenseGeometry, SingleHeadSingleBatchIsAOneDeepGrid) } /// Representative launches, read off attention_dense_grid/_block for each -/// block_m. The same graph launches differently per selected tile, so a prepare() that -/// used one tile's block_m for another's binary produces exactly one of these mismatches. +/// block_m, at the qb_minor non-persistent block order. The same graph launches +/// differently per selected tile, so a prepare() that used one tile's block_m for +/// another's binary produces exactly one of these mismatches. TEST(TestGfx950AttentionDenseGeometry, WitnessLaunchesDependOnTheSelectedBlockM) { // BF16/D128/H32/8 causal, B3, Sq=Skv=1536. diff --git a/projects/hipdnn/tools/IngestorGenerator/configs/gfx950_attention_dense.engine.yaml b/projects/hipdnn/tools/IngestorGenerator/configs/gfx950_attention_dense.engine.yaml index af1014345a4b..4feedcad2f42 100644 --- a/projects/hipdnn/tools/IngestorGenerator/configs/gfx950_attention_dense.engine.yaml +++ b/projects/hipdnn/tools/IngestorGenerator/configs/gfx950_attention_dense.engine.yaml @@ -120,7 +120,7 @@ enumerate: candidates: dispatch.attention:attention_candidates request: dispatch.attention.common:AttentionRequest # Selects the one non-persistent grid candidate. - pin: {algorithm: attention_dense_grid, spec_id: gfx950_dense_grid} + pin: {algorithm: attention_dense_grid, spec_id: gfx950_dense_qbminor} # Request fields every cell shares. fixed: {arch: gfx950, batch: 1, seqlen_q: 512, seqlen_k: 512} # Crossed in order, first group outermost (the catalog order). @@ -177,7 +177,7 @@ enumerate: # lds_k_row_pad too: `space: slice` leaves it at its default while `space: full` # varies it. tests/test_registry_regeneration.py asserts the defaults (g1) and # diffs the enumeration against the shipped catalog on the shipped fields (g3). - spec_fields: [dtype, head_size, num_query_heads, num_kv_heads, causal, batch, seqlen_q, seqlen_kv, sliding_window, ragged, varlen, paged, persistent, block_m, block_n, waves_per_eu, lds_k_group_pad, lds_v_row_pad, lds_k_row_pad, lazy_rescale, wide_lds_dma, interleave, persist_decode, num_persistent, block_size, num_kv_blocks, use_sinks] + spec_fields: [dtype, head_size, num_query_heads, num_kv_heads, causal, batch, seqlen_q, seqlen_kv, sliding_window, ragged, varlen, paged, persistent, block_m, block_n, waves_per_eu, lds_k_group_pad, lds_v_row_pad, lds_k_row_pad, lazy_rescale, wide_lds_dma, interleave, persist_decode, nonpersist_decode, num_persistent, block_size, num_kv_blocks, use_sinks] # Frozen build cost: a registry that grows past this fails here, not at pack time. max_kernels: 840 diff --git a/projects/hipdnn/tools/IngestorGenerator/configs/gfx950_attention_dense.yaml b/projects/hipdnn/tools/IngestorGenerator/configs/gfx950_attention_dense.yaml index d12cc7f667cc..99efb393dcba 100644 --- a/projects/hipdnn/tools/IngestorGenerator/configs/gfx950_attention_dense.yaml +++ b/projects/hipdnn/tools/IngestorGenerator/configs/gfx950_attention_dense.yaml @@ -193,7 +193,7 @@ packs: metadata: [dtype, head_size, num_query_heads, num_kv_heads, causal, ragged, sliding_window, batch, seqlen_q, seqlen_kv, block_m, block_n] vocabulary: dtype: {bf16: BF16, fp16: FP16} - spec_order: [dtype, head_size, num_query_heads, num_kv_heads, causal, batch, seqlen_q, seqlen_kv, sliding_window, ragged, varlen, paged, persistent, block_m, block_n, waves_per_eu, lds_k_group_pad, lds_v_row_pad, lazy_rescale, wide_lds_dma, interleave, persist_decode, num_persistent, block_size, num_kv_blocks, use_sinks] + spec_order: [dtype, head_size, num_query_heads, num_kv_heads, causal, batch, seqlen_q, seqlen_kv, sliding_window, ragged, varlen, paged, persistent, block_m, block_n, waves_per_eu, lds_k_group_pad, lds_v_row_pad, lazy_rescale, wide_lds_dma, interleave, persist_decode, nonpersist_decode, num_persistent, block_size, num_kv_blocks, use_sinks] spec_defaults: head_size: 64 batch: 1 @@ -211,6 +211,7 @@ packs: wide_lds_dma: false interleave: false persist_decode: auto + nonpersist_decode: qb_minor num_persistent: 256 block_size: 0 num_kv_blocks: 0 @@ -289,7 +290,7 @@ packs: metadata: [dtype, head_size, num_query_heads, num_kv_heads, causal, ragged, sliding_window, batch, seqlen_q, seqlen_kv, block_m, block_n] vocabulary: dtype: {bf16: BF16, fp16: FP16} - spec_order: [dtype, head_size, num_query_heads, num_kv_heads, causal, batch, seqlen_q, seqlen_kv, sliding_window, ragged, varlen, paged, persistent, block_m, block_n, waves_per_eu, lds_k_group_pad, lds_v_row_pad, lazy_rescale, wide_lds_dma, interleave, persist_decode, num_persistent, block_size, num_kv_blocks, use_sinks] + spec_order: [dtype, head_size, num_query_heads, num_kv_heads, causal, batch, seqlen_q, seqlen_kv, sliding_window, ragged, varlen, paged, persistent, block_m, block_n, waves_per_eu, lds_k_group_pad, lds_v_row_pad, lazy_rescale, wide_lds_dma, interleave, persist_decode, nonpersist_decode, num_persistent, block_size, num_kv_blocks, use_sinks] spec_defaults: head_size: 128 batch: 1 @@ -307,6 +308,7 @@ packs: wide_lds_dma: false interleave: false persist_decode: auto + nonpersist_decode: qb_minor num_persistent: 256 block_size: 0 num_kv_blocks: 0