[aarnetalman]

Motivation

Deterministic inference (--enable-deterministic-inference) runs prefill with the unified extend-attention kernel (extend_attention_fwd_unified), whose reduction order does not depend on the batch. On gfx950 it launches every head_dim <= 128 with a 64-row query tile and 4 warps. Each workgroup streams the whole prefix for its query rows, so with long prompts the kernel is bound by KV reads: at 64k-token inputs it takes 82% of deterministic mode's GPU time (Qwen3-8B on MI350X), and deterministic throughput falls to about 0.2x of normal mode.

A 256-row tile with 8 warps reads each KV block once for 4x as many query rows. The KV block size (BLOCK_N), and so the order in which each row's softmax and sums are accumulated, is unchanged, so the output is bitwise identical to today's.

Modifications

• extend_attention.py (extend_attention_fwd_unified): on gfx950, for head_dim <= 128 with the default BLOCK_N of 64, launch with BLOCK_M = 256 and 8 warps. Other GPUs, head dims, the shared block-size table and the non-deterministic extend kernels are unchanged.
• test_triton_attention_kernels.py: on gfx950, the unified kernel's output with the new tile equals the default tile's output bit for bit (head_dim 128 and 64, three sequences with 0-4k prefixes).

A pipelined variant (128-row tile, 2 stages) was faster in isolation but aborted in ROCm Triton's compiler (LLVM iota_range assertion) for the server's kernel specialisation, so it is not used.

Accuracy Tests

The output is bitwise identical to the current launch, so accuracy cannot change. Checked three ways on MI350X (rocm/sgl-dev:v0.5.22-rocm724-mi35x-20261018, SGLANG_USE_AITER=1, our open deterministic-mode fixes for MI350X applied):

• Kernel level: 500 launch configurations timed on 8k-token chunks after 0-56k-token prefixes and a 4-sequence mixed batch, head_dim 128 (Qwen3-8B: 32/8 heads) and 64 (gpt-oss: 64/8 heads, sinks). Every configuration with BLOCK_N = 64 gives output equal to the current launch (torch.equal); the new one is 2.0-2.6x faster.
• New unit test test_extend_attention_unified_gfx950_tile_bitwise passes on MI350X. Two existing tests in the same file fail on gfx950 with and without this change (same results on the unchanged image): test_extend_attention_block_sizes (expects the tile [ROCm] Extend the gfx950 extend-attention tile to head_dim <= 128: -43% kernel, -14% TTFT, bit-identical #34461 added, which [AMD] Optimize Kimi-K3 Triton MLA prefill on gfx950 #35770 narrowed) and the head_dim-80 case of test_extend_attention_unified_vs_regular (max diff 0.167, tolerance 0.15/0.17). Both are pre-existing; the first is fixed by [Perf] [AMD] attention: Restore the gfx950 extend tile for head_dim <= 128 #45099, which restores that tile and is independent of this PR (the two apply together cleanly and were tested together).
• End to end, deterministic mode with the change, python3 -m sglang.test.test_deterministic --test-mode {single,prefix,radix_cache} --n-trials 50, --test-mode prefix --return-logprob, and a bitwise batch-size check (the same prompt in batches of 1 to 50, 1,275 responses, token IDs and logprobs compared exactly): all pass for Qwen3-8B (TP1, TP2, TP8), Qwen3-30B-A3B, Phi-4, gpt-oss-20b (TP1, TP2) and gpt-oss-120b. Long prompts (16k-120k tokens alone vs with other requests in the batch, and with chunked-prefill boundaries moved by other requests) give identical token IDs and logprobs. Rechecked on the 2026-10-19 nightly (v0.5.22-rocm724-mi35x-20261019): Qwen3-8B and Qwen3-30B-A3B pass every check.

Speed Tests and Profiling

python3 -m sglang.benchmark.serving --backend sglang --dataset-name random --random-range-ratio 1.0 --request-rate inf, deterministic mode, before and after alternating twice on the same GPU (every pair within 0.5%), output throughput:

Model, attention	1k in / 1k out (256 prompts, conc. 256)	8k / 1k (256, conc. 128)	15k-32k / 256 (64, conc. 16)	64k / 256 (32, conc. 8)	120k-126k / 256 (16, conc. 4)
Qwen3-8B, Triton prefill + AITER decode	7,070 -> 7,234 (+2.3%)	1,211 -> 1,469 (+21%)	47.9 -> 86.2 (+80%, 32k)	17.8 -> 35.2 (+98%)*	5.9 -> 13.7 (+130%)*
Qwen3-30B-A3B, Triton prefill + AITER decode	6,635 -> 6,822 (+2.8%)	1,262 -> 1,664 (+32%)
Phi-4, Triton prefill + AITER decode		778 -> 924 (+19%)	95.4 -> 139.8 (+47%, 15k)
gpt-oss-20b, Triton		1,644 -> 1,877 (+14%)	90.6 -> 160.5 (+77%, 32k)		13.0 -> 27.5 (+112%)
gpt-oss-20b TP2, Triton		2,610 -> 2,893 (+11%)	154.5 -> 246.3 (+59%, 32k)
gpt-oss-120b, Triton		1,031 -> 1,169 (+13%)	59.6 -> 105.3 (+77%, 32k)

* Qwen3-8B with YaRN (factor 4, 128k context). Time to first token drops by the same factors (Qwen3-8B 32k: 41.4 -> 22.3 s); inter-token latency is unchanged (decode does not use this kernel). gpt-oss ran with ROCm/aiter#6583's deterministic MoE.

Profile at 64k input (Qwen3-8B, 8 prompts, 40 forward steps, deterministic mode): _fwd_kernel_unified 79.0 -> 30.7 s of GPU time; total GPU time 96 -> 48 s.

No change without --enable-deterministic-inference or off gfx950.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• 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.