[aarnetalman]

Motivation

On MI350X/MI355X (gfx950), the Triton prefill attention kernel used to run faster for models with head size 64 or 128, such as gpt-oss, Llama and Qwen. #34461 made that change, with identical outputs. A later PR (#35770, for Kimi-K3) appears to have undone it by accident, so these models are back on the slower setting. The test #34461 updated still expects the faster setting and fails on MI350X; it runs only on MI300 in CI, so nobody noticed. This PR restores #34461's setting.

@clintg6, was narrowing it intended for the MLA kernels? MLA head dims (576/512, 192/128) take their own branches and are not affected by this change.

Modifications

• extend_attention.py (_get_block_sizes_for_extend_attention): restore [ROCm] Extend the gfx950 extend-attention tile to head_dim <= 128: -43% kernel, -14% TTFT, bit-identical #34461's condition and comment, _is_gfx95 and Lq <= 256. BLOCK_N is unchanged, so outputs stay bitwise identical.
• No test change: the existing test_extend_attention_block_sizes already expects this.

Accuracy Tests

The reduction order (BLOCK_N) is unchanged, so outputs are bitwise identical; checked at kernel level on MI350X for head_dim 64 and 128 (8k-token chunks after 0-56k-token prefixes, and a mixed 4-sequence batch): output equal to the 64-row tile's with torch.equal.

python3 -m pytest test/registered/attention/test_triton_attention_kernels.py -k extend on MI350X (rocm/sgl-dev:v0.5.22-rocm724-mi35x-20261018): before 3 failed / 5 passed, after 2 failed / 6 passed: test_extend_attention_block_sizes fails before and passes after. test_extend_attention and the head_dim-80 case of test_extend_attention_unified_vs_regular fail on MI350X with and without the change (tolerance checks, unrelated).

Speed Tests and Profiling

MI350X, Qwen3-8B TP1, python3 -m sglang.benchmark.serving --backend sglang --dataset-name random --random-range-ratio 1.0 --request-rate inf, before and after alternating twice (pairs within 0.3%), output throughput:

Server	8k in / 1k out (256 prompts, conc. 128)	32k in / 256 out (64 prompts, conc. 16)
--attention-backend triton	1,468 -> 1,509 tok/s (+2.8%)	101.7 -> 130.0 tok/s (+28%; TTFT 17.4 -> 12.6 s)
deterministic mode, Triton prefill (the deterministic kernel takes its tile from the same table; measured with AITER decode, #44489)	1,203 -> 1,376 (+14%)	47.9 -> 69.2 (+44%)

No change off gfx950, for head_dim > 256, or with other attention backends.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests. (existing test)
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #39546508413
Latest PR Test (Extra): ❌ Run #39546508460
Latest PR Test (AMD ROCm 10): ❌ Run #39546508338