[alexnails]

Summary

extend_attention.py already carries a gfx950-tuned tile — (BLOCK_M, BLOCK_N) = (128, 64) with num_warps=8 — but gates it behind 128 < Lq <= 256, which excludes head_dim 64 and 128 for no stated reason. This changes the predicate to Lq <= 256.

One functional line. Output is bit-identical.

Before / after — isolated kernel

MI350X (gfx950, 256 CU), head_dim 64, measured per layer type because the model alternates full-attention and sliding-window layers:

layer type	(64,64)/4 (default)	(128,64)/8 (this PR)	delta
full attention	baseline		−47%
sliding window	baseline		−51%

head_dim 128 also gains −32%. head_dim 256 already took this path and is unchanged.

(128,64)/8 is the only configuration that wins both layer types. (64,64)/8 looks attractive on sliding-window layers (+126%) but loses 12% on full-attention layers, making it unusable for a model that alternates them.

Before / after — in a served profile

gpt-oss-120b, TP4, ISL 8192, cc8, SGLANG_USE_AITER=1 --attention-backend triton. Same-window A/B, four minutes apart on the same server script:

measure	before	after	delta
_fwd_kernel µs/call	1308.9	741.8	−43.3%
_fwd_kernel share of prefill GPU	36.0%	23.7%
total prefill GPU	654.6 ms	564.1 ms	−13.8%

Every other kernel's call count and per-call time is unchanged, so the reduction is attributable to the retile alone. Traces were checked for the roctracer corrupt-timestamp defect (4987 events, 1 outlier of 0.1 ms, sum/union 1.003).

Before / after — served

benchmark_serving.py, paired probes run in both A/B orders:

concurrency	Δ TTFT	Δ output throughput
8	−13.6%	+8.9%
16	−15.9%

Warm 1k is unchanged (105.1 vs 106.7 ms), matching the microbenchmark's prediction of a small short-prefill effect.

Numerics

Bit-identical: 14/14 shapes, max_abs exactly 0.0.

This is structural rather than incidental. BLOCK_N governs the online-softmax rescale sequence and the extent of the p@v reduction, so changing it genuinely alters reduction order — but BLOCK_M cannot. Holding BLOCK_N = 64 while raising only BLOCK_M therefore preserves the arithmetic exactly. The sweep confirms the rule in both directions: every BLOCK_N=64 configuration is bit-exact, and every BLOCK_N ∈ {32, 128} configuration differs by ~1e-3.

Also checked, no change made

The other HIP launch knobs on this path are already optimal and were left alone: waves_per_eu=1, matrix_instr_nonkdim=16, num_stages=1. kpack is a literal no-op on gfx950 — the compiler forces 2→1 — so it was excluded from the search.

Hardware: MI350X (gfx950), ROCm 7.2 userspace. Measured on MI350X only; an MI355X confirmation would be welcome.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #32696876435
Latest PR Test (Extra): ⏳ Run #32698281924
Latest PR Test (AMD ROCm 7.2): ⏳ Run #32696876681

[alexnails]

/tag-and-rerun-ci