TheRock/rocm-libraries/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/ALGORITHM.md

gfx942 unified attention, from the math up

A derivation of what the gfx942 (CDNA3 / MI300X) tiled SDPA-forward kernels compute, and why they are shaped the way they are — the online-softmax recurrence, the paged-KV layout the example feeds them, and the CDNA3-specific matmul choice (the transposed-x8 flash regime) that the harness exercises.

This is written for a reader with a math / CS background who has not seen flash attention before. We start from the definition of attention, derive the streaming-softmax recurrence that makes it tile-able, and then map that recurrence onto what the two shipped paths — the narrow 16x16x16 default and the transposed-x8 flash regime — actually do.

If you just want to run the example, see README.md. The two kernels themselves live in kernels.gfx942.attention_tiled_2d (prefill) and ..._tiled_3d (the split-KV decode path); this folder is a torch-reference parity + latency harness over them.

---

0. Notation

symbol	shape	meaning
Q	S_q \times d	query matrix, one (\text{batch}, \text{head}) slice
K	S_k \times d	key matrix
V	S_k \times d	value matrix
O	S_q \times d	output (same shape as Q)
S_q, S_k	scalar	sequence lengths (query, key)
d	scalar	head dimension (head_size, here 64 or 128)
\tau	scalar	softmax scale, \tau = 1/\sqrt{d}
T	scalar	the KV tile width (block_size, 64 in this example)

A subscript like Q_i denotes row i; S_{i,:} denotes row i of a score matrix. \exp, \max act rowwise.

---

1. What attention computes (the specification)

Scaled dot-product attention, for one head of one batch element, is

O = \operatorname{softmax}\!\left(\tau\, Q K^{\top}\right) V,
\qquad \tau = \frac{1}{\sqrt{d}} .

Written one query row i at a time, with the score vector s = \tau\, Q_i K^{\top} \in \mathbb{R}^{S_k} (one score per key):

O_i \;=\; \sum_{j=1}^{S_k} \underbrace{\frac{e^{\,s_j}}{\sum_{j'} e^{\,s_{j'}}}}_{\text{attention weight } a_j}\, V_j .

So O_i is a convex combination of the value rows, weighted by how much query i attends to key j. That is the entire mathematical object. The rest of this document is about computing it without ever materializing the S_q \times S_k score matrix — which for long sequences (the long_prefill group reaches S=8192) does not fit in registers or LDS.

---

2. The softmax stability trick

Softmax is invariant to subtracting any constant c from every score:

\frac{e^{\,s_j}}{\sum_{j'} e^{\,s_{j'}}}
= \frac{e^{\,s_j - c}}{\sum_{j'} e^{\,s_{j'} - c}}
\qquad\text{for any } c .

In finite precision e^{s_j} overflows for large s_j; choosing c = \max_j s_j keeps the largest exponent at e^0 = 1. The whole difficulty of flash attention is that we want this max, but we process keys in tiles and do not know the global max until we have seen every tile. Section 3 resolves that.

---

3. Online (streaming) softmax — the heart of the algorithm

Process the S_k keys in tiles of width T. For a fixed query row, write the score sub-vector of tile t as s^{(t)} \in \mathbb{R}^{T}. Maintain three running quantities, updated tile by tile:

• m^{(t)} — the running max of all scores seen through tile t,
• \ell^{(t)} — the running denominator \sum e^{s - m^{(t)}},
• \mathbf{o}^{(t)} \in \mathbb{R}^{d} — the running un-normalized output.

Initial state: m^{(0)} = -\infty, \ell^{(0)} = 0, \mathbf{o}^{(0)} = \mathbf{0}.

Update for tile t. Extend the running max, then form the rescale factor

m^{(t)} = \max\!\big(m^{(t-1)},\; \max_j s^{(t)}_j\big), \qquad
\boxed{\;\alpha^{(t)} = e^{\,m^{(t-1)} - m^{(t)}}\;}\quad (0 < \alpha^{(t)} \le 1),

and apply it before adding this tile's contribution:

\begin{aligned}
p^{(t)}_j &= e^{\,s^{(t)}_j - m^{(t)}} & &\text{(tile weights, current max)}\\[2pt]
\ell^{(t)} &= \alpha^{(t)}\,\ell^{(t-1)} + \textstyle\sum_j p^{(t)}_j & &\text{(rescale old sum, add new)}\\[2pt]
\mathbf{o}^{(t)} &= \alpha^{(t)}\,\mathbf{o}^{(t-1)} + \textstyle\sum_j p^{(t)}_j\, V^{(t)}_j & &\text{(rescale old output, add new)}
\end{aligned}

Final normalization after the last tile: O_i = \mathbf{o}^{(N)} / \ell^{(N)}.

By induction over t, \alpha^{(t)} re-bases every previously-stored term from "$-m^{(t-1)}$" to "$-m^{(t)}$", so at t = N the max is global and \mathbf{o}^{(N)}/\ell^{(N)} equals the true softmax output of Section 1. The streaming computation is exact, not an approximation — it never holds more than one $T$-wide tile of scores at a time. This is what makes attention O(S_k) memory instead of O(S_k^2), and it is the recurrence both gfx942 kernels implement.

---

4. Causal masking and the early-exit

This example is causal-only for prefill (the kernel always applies a causal mask and has no non-causal prefill mode — see final_shapes_check.py). Causal attention forbids a query at position q from attending to keys k > q: in the tile update, s_j \leftarrow -\infty whenever key j is in the future of the query row, which makes its weight e^{-\infty} = 0.

Naively you still compute every score and mask the upper triangle — ~2× the matmul the math needs. The standard win is to stop the K-loop early: a query tile only needs KV tiles whose lowest key index is \le its largest query position, so fully-future tiles are skipped entirely. The final_shapes_check.py reference mirrors the kernel's masking with a bottom-right-aligned causal mask (torch.triu(..., diagonal = kv_len - query_len + 1)), which is the correct alignment when S_q \ne S_k — and a known point of divergence from Torch's is_causal=True (top-left aligned), called out in that script's docstring.

Decode (seqlen_q == 1) is run non-causal: a single query attends to all keys, which is identical to causal at q = 1.

---

5. The data layout the example feeds: paged KV

The harness does not hand the kernel a dense [B, H_q, S_q, d] tensor; it maps each dense SDPA problem onto the paged-KV layout the production dispatcher uses (parity_unified_attention.py::make_inputs):

• Each batch element becomes one sequence with (\text{query\_len}, \text{kv\_len}) = (\text{seqlen\_q}, \text{seqlen\_k}).
• The KV cache is [num_blocks, block_size, kv_heads, head_size] with block_size = 64; a per-sequence block table is a contiguous, non-overlapping run of block_size-token blocks, so each sequence's KV working set is genuinely distinct.
• Queries are packed [total_q, heads, head_size] with a cu_seqlens_q prefix sum; seqused_k carries each sequence's KV length.

block_size = 64 is chosen to match the gfx942 flash-regime tile (T = 64), so the wide / narrow geometries below apply cleanly. Grouped-query attention (GQA) is expressed by heads > kv_heads: a group of heads / kv_heads query heads shares one KV head (the reference repeat_interleaves K/V to match).

The launch grid is recomputed exactly as the production dispatcher does — grid = (kv_heads, total_num_q_blocks, 1) (or transposed under the use_q_major_grid flag), block = (64 * num_warps, 1, 1) — so the example exercises the same build + launch plumbing as the provider.

---

6. Two matmul orientations: narrow vs transposed-x8

The recurrence of Section 3 is fixed; the two shipped gfx942 paths differ only in how the two matmuls (QK^\top and PV) are mapped onto the CDNA3 matrix unit. The harness selects between them in _build_spec.

6.1 The narrow default (16x16x16)

In this harness's _build_spec, every shape that is not D128 fp16 uses the default narrow path: the mfma_16x16x16 atom, with num_warps keyed on block_size (D64 ships num_warps=4 + block_m_per_warp=32; D128 bf16 stays num_warps=2). Scores S = \tau Q K^\top are produced as the MFMA accumulator, the softmax recurrence runs over them, the probabilities P are staged to LDS (P_lds), and the PV matmul reads them back as the A-operand. This is the general-purpose path and the one whose LDS footprint the supports_tiled_2d gate models.

Note — harness vs production dispatcher. This narrow/flash split is the harness's explicit choice (parity_unified_attention.py::_build_spec routes only D128 fp16 to the flash regime). The production dispatcher's gfx942 fp16 flash family is broader: _enable_gfx942_fp16_flash now admits D64 fp16 too (head_size in (64, 128)), running the same gfx942-legal 32x32x8 atom plus a sliced-K / cfvst stack. So when final_shapes_check.py runs through the production dispatcher (Sections 8–9), a D64 fp16 prefill shape is not on the narrow path described here — it takes the dispatcher's D64 flash path. The bf16 fallback to narrow (Section 6.2) holds in both the harness and the dispatcher.

6.2 The transposed-x8 flash regime (D128 fp16)

For D128 fp16, the harness routes to the flash regime (use_mfma_32x32x8 = True

• use_transposed_qk_32x32 = True), which is the central CDNA3 trick. Rather than computing S = Q K^\top, the kernel computes the transpose

S^{\top} = K\,Q^{\top}

via the gfx942-legal mfma_f32_32x32x8_f16 atom (a wider-M MFMA: M=32, N=32, K-step 8, <16 x f32> accumulator). The payoff: the probabilities come out already in the orientation the PV matmul needs as its B-operand, so P^{\top} stays register-resident and feeds the PV MFMA directly — no P_lds round-trip. The output accumulator's LDS region (Acc_lds) is epilogue-only and is backend-aliased into the loop-dead K/V LDS region. Consequences the harness and provider account for:

• use_k_single_buffer is set to num_warps * 32 <= block_size: when the M-tile (\text{num\_warps} \times 32) fits the tile width T, K is single-buffered; wider tiles double-buffer K.
• supports_tiled_2d's LDS gate over-counts this path (it assumes P_lds present + K double-buffered + a full Acc_lds), so the harness skips the generic gate for the flash configs and lets the spec's own __post_init__ + comgr be the arbiters — mirroring the provider's accurate transposed-x8 LDS model. The generic gate is applied only to the narrow path.

use_mfma_32x32x8 is fp16-only on this kernel path — the gfx942 spec's __post_init__ rejects it for bf16 ("use_mfma_32x32x8 is fp16-only"). Note this is a kernel-path restriction, not a hardware one: gfx942 does have the bf16 32x32x8 atom (mfma_f32_32x32x8_bf16 is in the gfx942 MMA catalog and lowers through the ..._1k builtin), but the transposed-x8 attention path is only wired and validated for fp16. That is why D128 bf16 falls back to the narrow 16x16x16 path (Section 6.1) rather than the flash regime.

---

7. The wide4 default vs the L4 contrast

Within the D128-fp16 flash regime, the harness reproduces the two geometries the provider ships and demonstrates the difference explicitly (this is the point of parity_unified_attention.py, per its docstring):

config	num_warps	workgroup	K buffering	when
wide4	4	256 threads	double-buffered (32·4 > 64)	the provider's analytic default for qualifying D128 fp16
L4	1	64 threads	single-buffered (32·1 ≤ 64)	the HIPDNN_GFX942_FLASH_WIDE=0 contrast

The important subtlety the docstring records: wide4 is the provider's analytic default, not the DSL spec's default. A spec built with no flash knobs lands on L4 (WG=64). So parity_unified_attention.py sets num_warps=4 (+ use_mfma_32x32x8 / use_transposed_qk_32x32 / use_k_single_buffer=False) explicitly to reproduce the shipped peak; HIPDNN_GFX942_FLASH_WIDE=0 selects the L4 geometry for an A/B. HIPDNN_GFX942_FLASH_WIDE=2 selects an intermediate wide2 (num_warps=2, WG=128, BLOCK_M = 32·2 = 64 ≤ T, so K single-buffers).

---

8. The decode path (3D split-KV)

Decode (seqlen_q == 1) is structurally different: there is one query row but a long KV sequence, so parallelism comes from splitting the KV dimension across threadgroups (the attention_tiled_3d path) and combining the partial softmax-weighted outputs. The harness routes decode through the production dispatcher run_unified_attention_torch(..., backend="auto") (final_shapes_check.py), which selects the 3D path and — for the overhead-dominated decode regime — can engage an internal CUDA-graph replay (Section 9).

---

9. Host overhead and the graph heuristic

For tiny kernels (decode, short prefill) the per-launch host cost (Python dispatch + kernarg pack + hipModuleLaunchKernel) is a large fraction of total time, so CUDA-graph replay — which removes exactly that host cost — pays off. The dispatcher encodes this in _recommend_graph_replay(problem) (attention_unified.py), verified to return:

regime	rule	graph?
decode	max_seqlen_q == 1	yes (3D)
short prefill	max_seqlen_q <= 768	yes (2D, or 3D if routed there)
long prefill	max_seqlen_q > 768	no (kernel-bound; overhead is noise)

Feature-flagged problems (sinks / alibi / qq-bias / softcap / sliding-window / fp8) are excluded. The internal graph is engaged only when the caller is not already capturing (a framework that graphs the whole forward takes precedence), and is toggled by HIPDNN_GFX942_2D_GRAPH / HIPDNN_GFX942_3D_GRAPH (both default-on). final_shapes_check.py measures each shape in two modes — eager (host + kernel) and graph (kernel-only, host overhead removed symmetrically) — so a CK-vs-Torch gap can be attributed to kernel time vs launch overhead.

---

10. Numerical correctness gate

Both scripts gate the kernel output against an fp32 paged reference (ref_paged_attn): gather the per-sequence K/V from the block table, repeat_interleave for GQA, score in fp32, apply the bottom-right causal mask, softmax, and weight V. The tolerance is max_abs <= 2e-2 (fp16) / 4e-2 (bf16) — wide enough to absorb the fp16/bf16 MFMA accumulation order, tight enough to catch a real layout or masking bug. No latency number is reported for a kernel that fails the gate.

---

11. Where the algorithm ends and tuning begins

The recurrence (Sections 3–4) and the paged layout (Section 5) are fixed; the choices in Sections 6–9 — the matmul orientation, the wide4/L4 geometry, K buffering, the grid transpose, and graphing — change only how the fixed computation is mapped onto CDNA3, never what is computed. Correctness is pinned by the fp32 gate (Section 10); the rest is the per-shape schedule the harness lets you A/B through the documented HIPDNN_GFX942_* environment variables.