Files
TheRock/rocm-libraries/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/ALGORITHM.md
T
derek.holloway 7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
Flatten subprojects into a single base monorepo
2026-09-02 16:34:47 -07:00

27 KiB
Raw Blame History

Unified paged attention on gfx950, from the math up

This file explains what the CK DSL unified_attention kernels compute and why they are shaped the way they are on CDNA (gfx950, MI355X, wave64, MFMA). The README.md is the parity + optimization history; this file is the specification, the data layout, and the two kernel strategies (2D and 3D split-KV) the production dispatcher chooses between. §8 additionally documents the dense prefill kernel — a non-paged, compile-time-sized peak-TFLOPS sibling (persistent grid-stride + on-chip ragged support) for the contiguous [B, S, H, d] prefill regime.

The harness in this folder benchmarks against AITER's Triton unified_attention; the CK DSL kernels mirror that Triton reference 1:1 in semantics (kernel_unified_attention_2d / _3d), which is why every output is bit-identical to Triton once cast back to the working dtype.


0. Notation

symbol shape meaning
Q S_q \times H_q \times d query, one batch slice; H_q query heads
K, V paged blocks of [\,b,\, H_k,\, d\,] key / value cache, H_k KV heads
O S_q \times H_q \times d output (same shape as Q)
S_q, S_k scalar query / key sequence length (per sequence)
d scalar head dimension (head_size, 64/128/256)
\tau scalar softmax scale, \tau = 1/\sqrt{d}
b scalar paged-cache block_size (16, 32, or 64)
g scalar GQA group size, g = H_q / H_k (queries per KV head)

Attention is paged and varlen: the KV cache is stored as fixed-size blocks scattered in HBM, a per-sequence block table maps logical KV positions to physical blocks, and cu_seqlens_q packs many sequences of different lengths into one ragged batch. This is the vLLM / serving layout, not a dense [B, H, S, d] tensor.


1. What unified attention computes (the specification)

For one query row i of one (\text{sequence}, \text{head}), with score vector s = \tau\, Q_i K^{\top} \in \mathbb{R}^{S_k}:


O_i \;=\; \sum_{j=1}^{S_k} \frac{e^{\,s_j}}{\sum_{j'} e^{\,s_{j'}}}\, V_j .

"Unified" means one kernel family serves the whole serving workload — decode (S_q = 1), prefill (S_q = S_k), and chunked prefill (1 < S_q < S_k) — plus every bias/mask variant a real deployment needs. The reference (ref_paged_attn in parity_unified_attention.py) and the kernels both apply, in this exact order, before the softmax row-reduction:

  1. scale s \leftarrow \tau\, Q_i K^\top;
  2. softcap (optional): s \leftarrow c\,\tanh(s/c) — bounds logits to \pm c;
  3. ALiBi (optional): s_j \mathrel{+}= \text{slope}_h \cdot (k_j - \text{ctx}) — a linear positional bias;
  4. QQ-bias (optional): s_j \mathrel{+}= \text{bias}[q_{\text{local}}, k_j - \text{ctx}] inside the query section;
  5. causal mask s_j \leftarrow -\infty for k_j > q_i + \text{ctx};
  6. sliding window (optional): also mask k_j < q_i + \text{ctx} - W;
  7. attention sinks (optional): append per-head learned logits as extra softmax columns (they soak up probability mass but contribute no value rows).

Each of these is a flag on UnifiedAttentionProblem / UnifiedAttention2DTiledSpec, so the dispatcher builds exactly the kernel the shape needs.


2. The streaming-softmax core (shared by both paths)

Both kernels are flash-attention: they tile the S_k keys, never materialize the S_q \times S_k score matrix, and maintain three running quantities per query row updated tile by tile — the running max m, the running denominator \ell, and the un-normalized output accumulator \mathbf{o} \in \mathbb{R}^{d}:


\begin{aligned}
m' &= \max\!\big(m, \operatorname{rowmax} s^{(t)}\big), &\quad \alpha &= e^{\,m - m'} \\
p^{(t)}_j &= e^{\,s^{(t)}_j - m'}, &\quad \ell' &= \alpha\,\ell + \textstyle\sum_j p^{(t)}_j \\
\mathbf{o}' &= \alpha\,\mathbf{o} + P^{(t)} V^{(t)}, &\quad O_i &= \mathbf{o}^{(N)} / \ell^{(N)}.
\end{aligned}

The rescale \alpha re-bases everything accumulated under the old max to the new max; the result is exact, not approximate. As in all of CK DSL, exponentials are computed in base 2 — the GPU has a hardware 2^x but no e^x. The tiled 2D/3D kernels take the raw softmax scale \tau as a runtime param and fold \log_2 e into it once per kernel on the device (qk_scale = scale * 1.4426950408889634, i.e. \tau \log_2 e), so every exp2(s - m) is one instruction. Softcap, ALiBi, and QQ-bias are likewise applied in the \log_2 domain (e.g. apply_softcap_log2).

Where the two paths differ is who owns which keys, which is dictated entirely by occupancy — how many CTAs the shape can keep the device busy with.


3. CDNA mapping: wave64 + MFMA

The matrix unit on gfx950 is v_mfma_* over wave64 (64 lanes), with native A \cdot B accumulation (no transpose, unlike RDNA's WMMA A B^\top). The default (non-transposed) path runs QK/PV on the wide-K 16x16x32 atom (mfma_f32_16x16x32_{f16,bf16}), falling back to 16x16x16 only for a K=16 head-dim tail step in PV. The transposed-32x32 prefill "combo" path runs QK/PV on gfx950's wide-K 32x32x16 atom (mfma_f32_32x32x16_{f16,bf16}, the gfx950-only 32x32 family). The flash-attention building blocks ported from CK Tile's BlockFmhaPipelineQRKSVSAsync are:

  • Q staged in LDS once per CTA, reused across the whole K-loop. (The transposed combo additionally hoists Q into VGPRs — Q32_reg — to drop Q's permanent LDS allocation.)
  • K and V streamed cache → LDS each tile, with the global load issued early so the QK MFMA starts the moment the LDS write retires (async DMA, see §7.1).
  • m, l in registers; the per-row max/sum reductions are XOR butterflies (CK Tile's block_tile_reduce_xor_sync pattern) — no LDS round-trip. The intra-16-lane / intra-32-lane stages lower to ds_swizzle_b32 SWAP mode (not ds_bpermute); only a cross-half (mask 32) stage — e.g. the split-KV reduce's wave64 fold, or the transposed scalar state's alpha broadcast — uses ds_bpermute.
  • o_acc in MFMA-accumulator distribution (per-lane <4×f32> per N-tile of the head dim), truncated to the output dtype through an LDS-staged shuffle epilogue with 16-byte stores.

4. The 2D path — one CTA owns a query block

The 2D kernel (attention_tiled_2d.py) assigns one CTA per $(\text{q-block}, \text{kv-head})$ and walks the entire KV sequence for that q-block in the K-loop. BLOCK_M = 16 packs the GQA group: with g queries per KV head, the 16 MFMA rows hold (\text{query-position}, \text{query-head}) pairs so one QK MFMA serves the whole group sharing a KV head.

  • Grid: (num_kv_heads, total_num_q_blocks, 1), block = 64 · num_warps.
  • Best for: chunked prefill and short / sliding-window contexts, where the q-block × kv-head grid is already large enough to fill the device. It is the path the production dispatcher selects for the d64/b32 GQA-8 serving traces.

The boundary masks (causal, sliding-window) are applied per element via v_cndmask. A causal early-exit clamps the K-loop so fully-future tiles are skipped; the boundary tile straddling the diagonal is still masked elementwise.

The transposed "combo" 2D variant (the prefill win)

The plain 2D kernel is single-warp-ish per CTA and leaves the device under-occupied on long prefill. The combo rewrite — selected automatically by _enable_combo_2d for the validated bf16 / fp8 d64/b32 GQA-8 family — restructures the softmax onto a transposed 32×32 MFMA layout (block_m_per_warp=32, use_mfma_32x32, use_transposed_qk_32x32) and layers several co-operating optimizations, each a spec flag:

flag what it does
use_transposed_scalar_state keep m/l in the transposed lane layout (no re-shuffle)
use_transposed_mask_once apply the causal mask once per tile, not per MFMA (full-attn only)
use_transposed_half_local_pv halve the PV operand staging via the transposed layout
use_mfma32_skip_legacy_qreg drop the legacy per-MFMA Q-register reload
use_transposed_mask_limit fold the causal row limit into the threshold (full-attn only)
use_fast_paged_kv_desc the 64/8-head fast paged-KV descriptor (one i64 block base)

These flags are mutually load-bearing — most are perf-neutral alone and only pay off together — which is why the harness sweeps them as a bundle (combo) and the dispatcher ships them as one validated policy. mask_once / mask_limit are valid only for full attention; under a sliding window the spec validator rejects them, so the SW combo runs with them off.

Occupancy is the 2D lever, not instruction count

Static ISA inspection shows the combo kernel is VALU/SALU-bound (~800 VALU + ~650 SALU vs only ~16 MFMA per kernel), dominated by the per-element causal-mask v_cndmask. But the binding constraint on this kernel is how many workgroups fit per CU — it is occupancy/latency-bound, and the wins all come from raising WG/CU, never from cutting instructions. Three levers shaped the tuning (all detailed in the README):

  1. Raise waves_per_eu (the d64/b32 trace family). That combo is VGPR-limited (~137 VGPR → 3 WG/CU at the default waves_per_eu=2); waves_per_eu=3 reaches 4 WG/CU (+15%), and waves_per_eu=4 adds ~5% more on full-attention shapes.
  2. K single-buffer for single-batch d128 prefill (the LDS lever). The single-batch d128 combo at num_warps=2 is not VGPR-limited but LDS-limited: at tile_size T = 2·block_size = 64 the K double-buffer + V single-buffer LDS (K_lds[2,T,HD] + V_lds[1,T,HD] = 48 KB) admits only 1 WG/CU, while the register file already admits two. Keeping the larger T = 64 tile (good long-context per-iter amortization) but halving K_lds via K single-buffer (use_k_single_buffer: K_lds[1] 16 KB + V_lds[1] 16 KB = 32 KB → 2 WG/CU, VGPR=215 AGPR=0) doubles occupancy and hides the per-iter latency. V double-buffer is OFF on d128 — a V[i+1] prefetch is a net drag there and would re-inflate LDS back to 1 WG/CU. The single K slot re-issues its next-K prefetch after the PV-wait barrier (all QK reads drained) so it cannot WAR-race. This took the single-batch d128 prefill cohort from below flash to ≈1.10x over torch SDPA flash (geomean; 1.36x at S1024, 1.02x at S2048, but 0.95x at the long S4096 holdout — a small honest loss to flash there, re-measured on verified llvm22). The same cohort still trails Triton's multi-warp 2D kernel (~0.55-0.60x); see README. d64 single-batch prefill instead keeps num_warps=4 with a wider tile_size = 128 to feed its KV loop.
  3. Lighten the prelude for sliding window. SW prunes the K-loop to a handful of tiles, so the per-CTA prelude (Q→LDS, binary search, sink init) dominates. For bf16 SW the combo drops to num_warps=2 (BLOCK_M=64, half the prelude, 2× the CTAs) but keeps tile_size = 2·block_size; fp8 SW is dequant-bound rather than prelude-bound, so it stays at num_warps=4 and is the only SW case that also shrinks tile_size = block_size (one paged-KV block per iter). The compute-bound no-SW combo keeps num_warps=4 / tile_size = 2·block_size for both dtypes.

Reducing instruction count by splitting the loop into a no-mask phase + a masked boundary phase was byte-identical but ~7% slower (I-cache / code-size cost on a latency-bound kernel), and was reverted — a representative "fewer instructions ≠ faster" result on this kernel.


5. The 3D split-KV path — many CTAs share a query block

When a single (\text{q-block}, \text{kv-head}) has so much KV that one CTA walking it serially leaves the device idle (long-context decode), the 3D kernel (attention_tiled_3d.py) splits the KV sequence into segments and gives each segment its own CTA, then reduces:

  • Grid: (total_num_q_blocks, num_kv_heads, NUM_SEGMENTS). Each segment CTA runs the §2 recurrence over only tile_start..tile_end of the KV sequence and writes a partial (m, \ell, \mathbf{o}) to a workspace segm_output[total_q, num_qh, num_segments, head_size] (plus segm_max and segm_expsum, all fp32).
  • reduce_segments then combines the per-segment partials into the final output. Because flash-attention's (m, \ell, \mathbf{o}) merge is associative (re-base each segment's accumulator to the global max, sum the denominators), the split is exact — the final acc /= L and the output-dtype cast happen in the reduce kernel.

The 3D path turns a serial KV walk into parallel segments + a cheap reduction, which is the only way to saturate the device on long-context single-query decode.


6. The dispatcher — which path runs (select_path)

UnifiedAttentionProblem.select_path() mirrors AITER's own use_2d_kernel selector so CK DSL and Triton make the same algorithmic choice. The rule, in spirit:

  • 2D when the q-block × kv-head grid already saturates the device (target = num_sms · 4), or for short context (S_k \le 512), or under a sliding window — the split-KV segments would only add launch overhead.
  • 3D split-KV otherwise — long, full-context sequences where the 2D grid is too small to fill the device.

Once a shape lands on the 2D path, a second tier of routing picks the kernel geometry for it, driven entirely by the occupancy (WG/CU) the shape can sustain:

  • Single-batch (num_seqs == 1) d128 long prefill (bf16/fp16, GQA, no bias/SW, S_q > 256): the full transposed-32×32 combo with num_warps=2, tile_size T = 2·block_size = 64, K single-buffer on, V double-buffer off. The K single-buffer halves K_lds so the larger T=64 tile still fits in the 32 KB → 2 WG/CU budget (see §4 lever 2). This is the LDS-bound occupancy story that reversed the old d128-prefill loss vs flash; it now wins ≈1.10x over torch SDPA flash (1.36x S1024 → 0.95x at the S4096 holdout) but still trails Triton's 2D kernel (~0.55-0.60x) — see the README cohort section. Only same-session ratios are load-bearing (±25-30% auto-clock).
  • Single-batch d64 long prefill: same combo but num_warps=4, tile_size = 128 (8·block_size) — d64 is wide-KV-loop-bound, so it wants the wider tile and more warps rather than the d128 LDS lever.
  • The d64/b32 GQA-8 serving traces (multi-seq chunked prefill, sinks): the combo with the waves_per_eu lever (§4 lever 1) and the prelude-light num_warps=2 SW geometry (§4 lever 3).
  • Decode / long full-context shapes route to 3D split-KV before the 2D geometry tier even runs — the segment grid is the only way to fill the device on single-query decode, and that lane is a clean ≈1.20x win over Triton (3D vs 3D, verified llvm22; see README).

This is also why the parity harness reports three tables (auto/2d/3d): auto is what production launches, while forcing both backends to the same path (2d-vs-2d, 3d-vs-3d) is the algorithmically-fair comparison.


7. Implementation details worth the math

7.1 Async DMA K/V (current-V-first, next-K-second)

K and V are streamed to LDS with async DMA. The issue order is deliberately current-V then next-K, so the PV matmul of the current tile only has to wait on the V stream while next-K loads in the shadow — the QK of the next tile then finds K already in flight. This overlap is what keeps the kernel HBM-bandwidth- fed on the long-context, large-cache regime where it wins (§ README).

7.2 64-bit paged-KV addressing

A paged-KV byte offset is physical_block · (block_size · num_kv_heads · head_size · dtype_bytes). That product overflows a 32-bit hardware buffer voffset once the cache exceeds 2 GiB (~65 K bf16 / ~131 K fp8 blocks) — and production caches are far larger (captured traces ~350 K blocks ≈ 11 GiB). _enable_i64_kv_addr switches the load paths to fold block · stride into a 64-bit base (only a small within-block offset stays in the 32-bit field) when num_kv_blocks · block_stride > 2³¹, so small caches keep the exact fast i32 path and only large caches pay the tiny per-block-base cost. Without it, large caches silently read garbage (verified max_abs ≈ 1.4).

7.3 fp8 KV cache (bf16-Q + fp8-KV)

The fp8 path stores K/V as fp8 e4m3 (half the HBM bytes of bf16). The sync-dequant loader writes bf16 into the K/V LDS (folding k_scale in) before the MFMA, so the transposed bf16 combo runs unchanged on fp8 inputs. A native fp8×fp8 QK MFMA (no dequant) was implemented and measured but is lose-lose on the compute-bound full-attention shape — slower and less accurate (quantizing Q to fp8 costs ~1e-2) — because the dequant is already hidden behind the K/V load latency. The accurate sync-dequant path is the production choice.

7.4 Magic division

Per-tile index arithmetic (pos // block_size, grid-axis decode) uses CK Tile's mul-hi magic division so a compile-time constant divisor folds to a v_mul_hi_u32 + add + shift instead of the ~20-cycle hardware integer divider.


8. The dense prefill kernel — peak-TFLOPS sibling (attention_dense.py)

The 2D/3D kernels above serve the paged, varlen deployment layout. Dense prefill (a contiguous [B, S, H, d] batch, no page table, self-attention S_q = S_k) is a different regime: the shape is known up front and there is no paging indirection to hide, so a dedicated kernel (kernels/gfx950/attention_dense.py) trades generality for peak MFMA throughput. It computes the exact same streaming-softmax core (§2) but bakes the shape into the kernel and reshapes the schedule around a 256-row query tile.

  • Compile-time-sized ABI. batch / seqlen / heads / head_size / causal / dtype are constants baked at build time (dense, statically-sized). The only runtime args are the q/k/v/o pointers and the f32 softmax scale. block_n (KV tile), waves_per_eu, and the persistent knobs are the tuning parameters; every algorithmic lever below is always-on.
  • 256-row query tile, 32×32×16 MFMA. _BLOCK_M = 256 query rows per CTA (8 wave64s = 512 threads) run QK/PV on gfx950's wide-K mfma_f32_32x32x16 atom. QK produces S^\top = K Q^\top so the key lands on the per-lane accumulator regs — the layout that keeps the softmax a cheap in-lane reduce + one lane^32 exchange and lets PV consume P with no relayout shuffle.
  • head_size 64 and 128 (bf16 / fp16), MHA and GQA including non-power-of-2 group sizes (e.g. 40/8, 28/4).

8.1 The winning levers (all always-on)

lever what it does measured
CK-1 transposed PV P feeds the PV MFMA in its native QK-output layout via a half-local V load (pv32_v_load_paired); the cross-half P-relayout shuffle is gone (~96 ds_bpermute removed). +35%
LDS bank-conflict pad on K ([NBUF, BN, D+8]) kills the 8-way conflict on the QK K-reads. +80% (base win)
LDS V pad (+32) the transposed PV read (ds_read_b64_tr_b16) has a stricter bank pattern than K; a +32 V-row pad fully clears its conflicts. +~5%
native exp2_fast (v_exp_f32, no overflow guard — softmax arg ≤ 0) one instruction per exp. +11.5%
depth-1 cluster fuses exp2(s − m) into the PV-MFMA loop so the softmax VALU/TRANS co-executes in the MFMA shadow (sched_group_barrier names the full DS_READ/MFMA/VALU/TRANS population per step). —
partial-vmcnt software prefetch per-tile K/V DMA drains to a partial vmcnt (keeps the freshest V prefetch in flight across the barrier) instead of a full vmcnt(0) serialize. raises MfmaUtil
PV-only s_setprio the PV MFMA cluster is bracketed at raised priority so it wins issue slots; paired with the prefetch. +3.5%
lazy online rescale keep the running max as a lazy max that only re-anchors when a tile exceeds it by >8 (log2); when every lane is within 8 (a wave_all vote) skip the O/ℓ rescale entirely (a 0/1-trip scf.for → a wave-uniform scalar branch), cutting the VALU between the QK and PV clusters. Numerically approximate (P bounded by 2^8) but parity-identical at bf16/fp16 tolerance. +~2%

Diagonal-only causal masking. For causal, below-diagonal KV tiles need no mask (~94% of tiles at S = 8192); the KV loop is split into a mask-free body [1, diag) and a masked diagonal tail [diag, n_upper), and the causal row limit clamps n_upper so fully-future tiles are never visited.

head_size and the DMA loader. One async_buffer_load_lds instruction moves 64 lanes × 2 bf16 = 128 elements. At D=128 that is exactly one padded K/V row per instruction (the fast path). At D=64 the loader packs 128//D = 2 rows per instruction into an unpadded contiguous LDS tile (lane l → row l/(D/2), col 2\,(l \bmod D/2)); the padded fast path is emitted byte-identically for D=128, so D=64 support costs the D=128 schedule nothing.

8.2 Two grids: default and persistent

Same inner pipeline, 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.
  • 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 W = ⌈S_q/256⌉ \cdot H_q \cdot B, so the fixed per-CTA cold-start is amortized once per CU instead of once per query block. The inner compute is byte-identical to the default path; the difference is the outer work loop, the work-item decode, and per-item state reset.

Persistent work-item decode (persist_decode). How W is unflattened to (\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.

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 same-session ratios are load-bearing; the numbers below are one representative session, each pinned to its config (grid / decode / V-pad / lazy):

config grid decode V-pad lazy TFLOPS
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

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 (persistent=True, persist_decode="auto", lazy_rescale=True, ROCKE_DENSE_VPAD=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.

8.3 Ragged sequence lengths — on-chip boundary padding

The 256-row / block_n-key tile geometry assumes aligned lengths. Real prompts are arbitrary, and this layout cannot pad on the host (no room to grow the dense buffers). Instead a separate kernel path (ragged=True, its own kernel_name) pads the boundary tiles on-chip — everything else is the §8.1 pipeline unchanged. Because seqlen is compile-time, the raggedness is fully known at build time:

  • Grid ceil'd — ⌈S_q/256⌉ query blocks (persistent: the work-item count is ceil'd), so the partial last query block is covered.
  • OOB query rows → register-zero pad. Q is loaded through a bounds-checked buffer_load_vN (buffer-resource extent = the real Q buffer); rows past S_q return 0 instead of faulting.
  • OOB keys → LDS-zero pad. The K/V async DMA already goes through buffer resources, so keys past S_k load as 0 into LDS.
  • Masking is nearly free for causal. A padded key has token index \ge S_k > every real query index, so the causal mask (k_j \le q_i) already drops it — causal ragged needs no extra key mask. Non-causal adds a compile-time k_j < S_k key mask on the partial tile.
  • Partial output rows dropped. The epilogue store is wrapped in a per-lane q_i < S_q guard, so padded rows never write (and never clobber a neighbouring batch's real rows — correct for any B).

Self-attention only (S_q = S_k); not combined with varlen or the sliding window (the validator rejects those). The aligned path is emitted byte-identically when ragged=False, so ragged support costs the aligned schedule nothing; the ragged kernel itself runs at ~890 TFLOPS (≈94% of the 948 aligned headline), the delta being the bounds-checked load + guarded store.


9. Where the algorithm ends and tuning begins

The math above is fixed and exact: online softmax, the bias/mask order of §1, the associative split-KV merge. Everything the README tunes — the combo flag bundle, waves_per_eu, num_warps, tile_size, async-DMA issue order, i64 addressing, fp8 dequant — changes only how these steps are scheduled onto gfx950, never what is computed. Correctness is pinned by bit-exact parity against Triton and ULP-level agreement with ref_paged_attn (unified paths) / torch SDPA (dense); for the dense kernel a golden LLVM-IR SHA gate additionally pins that every tuning change to one variant leaves the others' machine code byte-identical. Performance is the per-CTA schedule, the 2D-vs-3D occupancy choice, and the dense default-vs-persistent grid.


10. Where to go next

  • README.md — the parity harness, the three-table methodology, the prefill-2D optimization history (combo, waves_per_eu, i64 addressing, fp8), and the substantiated speedups.
  • parity_unified_attention.py — the canonical parity + benchmark harness (default / creative / fmha scenario sets).
  • The kernels themselves live in rocke (not in this folder): the unified paged kernels are kernels/gfx950/attention_tiled_2d.py, kernels/gfx950/attention_tiled_3d.py, and kernels/common/attention_unified.py (the dispatcher and shared spec); the dense prefill kernel (§8) is kernels/gfx950/attention_dense.py, gated by tests/test_attention_dense_golden.py (IR-SHA + C++/Python byte-identity).