[jiaqiang-dot-liu]

Motivation

On gfx950 the extend-attention kernel only takes the 128-row query tile when 128 < head_dim <= 256:

elif _is_gfx95 and 128 < Lq <= 256:
    BLOCK_M, BLOCK_N = (128, 64)
    num_warps = 8
else:
    BLOCK_M, BLOCK_N = (64, 64)
    num_warps = 4

Every workgroup in this kernel streams the whole prefix, so the KV bytes read per call scale with the number of query tiles: doubling BLOCK_M halves them. And BLOCK_M / num_warps = 128 / 8 = 16 rows per warp is exactly one MFMA tile at matrix_instr_nonkdim=16, so the wider tile maps cleanly onto the hardware rather than leaving a partial tile.

Neither part of that reasoning depends on the head dim being above 128. Head dims 64 and 128 — which cover most dense models, including the Llama family — are left on the narrow tile for no stated reason.

Modifications

Drop the lower bound on the condition:

elif _is_gfx95 and Lq <= 256:

BLOCK_N stays at 64 in both branches, so the KV sweep per query row is unchanged and each row's online-softmax accumulates in the same order. BLOCK_M only changes how many query rows a program handles, and rows are independent. The output is therefore expected to be bit-identical, and it is — see below. Other AMD architectures, CUDA, and head dims above 256 are untouched.

Accuracy Tests

MI355X (gfx950), 256 CU, ROCm 7.2, torch 2.9.1, Triton 3.7.0. B=32, H_Q=64, H_KV=8, bf16, causal, tensors from the repo's _build_batch.

12 configurations, head_dim {64, 128, 256} × N_CTX {1024, 2048, 4096, 8192}, run before and after. Each run hashes its inputs as well as its o_extend, so the output comparison is only accepted when the inputs provably matched.

12/12 inputs identical, 12/12 outputs bitwise identical.

test/registered/attention/unittests/dense/test_triton.py and .../swa/test_triton.py: 14 passed, 2 skipped, 55 subtests passed.

Speed Tests

Minimum over 7 repetitions of triton.testing.do_bench(warmup=100ms, rep=500ms, return_mode="median"), one idle GPU. The baseline sweep was run twice, before and after the patched sweep; maximum run-to-run drift was 0.33%.

head_dim 64 — BLOCK_M 64 → 128

N_CTX	before (ms)	after (ms)	delta
1024	0.4712	0.2588	−45.1%
2048	1.8496	0.8606	−53.5%
4096	5.9978	2.8041	−53.2%
8192	16.8052	8.0892	−51.9%

head_dim 128 — BLOCK_M 64 → 128

N_CTX	before (ms)	after (ms)	delta
1024	0.6572	0.5349	−18.6%
2048	2.4469	1.8698	−23.6%
4096	8.0976	6.2074	−23.3%
8192	22.8645	17.4134	−23.8%

head_dim 256 — control, tile choice unchanged

N_CTX	before (ms)	after (ms)	delta
1024	0.8084	0.8083	−0.0%
2048	2.8070	2.8060	−0.0%
4096	9.4010	9.3854	−0.2%
8192	26.3511	26.3371	−0.1%

All four control rows are inside the 0.33% drift, as they must be — the condition change cannot reach this head dim.

Register pressure

A wider query tile costs registers, and this file has a history of Triton 3.7 gfx950 spill regressions, so wall-clock alone is not enough. From the compiled AMDGCN at N_CTX=4096:

head_dim	sgpr_count	sgpr_spill	vgpr_count	vgpr_spill
64	105 → 106	4 → 2	107 → 114	0 → 0
128	106 → 106	4 → 6	169 → 151	0 → 0
256	106 → 106	14 → 14	227 → 227	0 → 0

No VGPR spills appear at any head dim, the SGPR spill counts stay in the single digits and move in both directions, and head_dim 128 actually uses 18 fewer VGPRs on the wider tile. Nothing here resembles the scratch-spill regime that motivated the narrower tile for Lq == 576.

What this does not cover

gfx950 only, and kernel-level only — no end-to-end serving numbers. The change is gated on _is_gfx95, so no other architecture is affected by construction.

---

CI States

Latest PR Test (Base): ❌ Run #35090178559
Latest PR Test (Extra): ❌ Run #35090178192
Latest PR Test (AMD ROCm 10): ❌ Run #35090178552

[jiaqiang-dot-liu]

Superseded by the upstream PR sgl-project#39808 (same branch).