[aarnetalman]

Motivation

This PR lets deterministic inference on MI350X/MI355X (gfx950) use AITER for decode attention: --decode-attention-backend aiter, with prefill attention still on a deterministic backend such as Triton. Decode attention is a large part of the cost of deterministic mode on these GPUs, and with AITER decode, deterministic throughput is 1.27-1.91x higher for Qwen3-8B, 1.53-1.61x for Qwen3-30B-A3B and 1.09-1.13x for Qwen3.8-27B, with every determinism check still passing and unchanged GSM8K.

AITER's paged decode attention (paged_attention_ragged) splits each request's KV cache into fixed 256-token partitions and reduces them per request. The split depends only on the request's own context length, never on the batch, so a request gets the same result whatever it is batched with.

This builds on #44438, which makes deterministic mode check --prefill-attention-backend and --decode-attention-backend. Decode has to be set explicitly: --attention-backend aiter --prefill-attention-backend triton stays refused, because --attention-backend itself must be a deterministic backend.

Refused in deterministic mode, and why

Configuration	Why
AITER as --attention-backend or --prefill-attention-backend	AITER prefill attention is not batch-invariant (in our tests its results changed with the batch from about 480 tokens, and with mixed batches)
AITER decode for MLA models	MLA decode sizes its KV splits by the batch size
AITER decode on other GPUs than gfx950	not tested
AITER decode with speculative decoding	target verify (--speculative-attention-mode decode) and some draft workers would run AITER's verify and extend kernels, which are not verified
AITER decode with sliding-window models, SGLANG_USE_AITER_UNIFIED_ATTN or the vectorized 5D KV-cache layout	these run AITER's unified-attention or 5D-pool decode kernels instead of paged_attention_ragged; not verified. Refused when the AITER backend starts
--enable-mixed-chunk (any backend)	a mixed batch runs its decode tokens through the prefill kernels; handled separately in #44701, which switches it off in deterministic mode

Modifications

• arg_groups/overrides.py (_deterministic_attention_backend): the allowed backends are built per phase first; aiter is added for decode on gfx950, for non-MLA models and without speculative decoding. Both phases are then checked in one loop. When AITER decode is refused, the error says why. The per-phase fields are read directly.
• layers/attention/aiter_backend.py: in deterministic mode, the AITER backend refuses to start when decode would not run paged_attention_ragged (unified attention or the vectorized 5D KV-cache layout), so an accepted configuration always runs the kernel that was verified.
• test_model_overrides.py: a table test (gfx950 or not, MLA or not, AITER prefill or not, speculative decoding or not) for which configurations are accepted. [Fix] deterministic: Validate the per-phase attention backends #44438's refusal test pins is_gfx95_supported to False and sets both per-phase fields, so it does not depend on the machine it runs on.
• Unchanged: deterministic mode off; every configuration that does not set --decode-attention-backend aiter; the radix cache (it follows --attention-backend, which stays a deterministic backend).

Accuracy Tests

Environment: 8x MI350X, rocm/sgl-dev:v0.5.21-rocm724-mi35x-20261013 (SGLang 27b68303d0, AITER e7d2453f2), SGLANG_USE_AITER=1 (the image default); aiter_backend.py is identical to upstream main. Deterministic mode on MI350X also needs the open fixes #43788 (AITER GEMM and MoE bypass; not merged yet) and #34543 (all-reduce for oversized inputs), which were applied in all runs, together with our other open PRs. These runs used the image without #44438, where --decode-attention-backend aiter was not checked yet; the server arguments are the ones this PR accepts:

python3 -m sglang.launch_server --model-path Qwen/Qwen3-8B --tp $TP \
  --attention-backend triton --prefill-attention-backend triton --decode-attention-backend aiter \
  --enable-deterministic-inference
# Qwen3.8-27B additionally: --linear-attn-prefill-backend triton --linear-attn-decode-backend triton

Checks, each on a fresh server:

• python3 -m sglang.test.test_deterministic --test-mode {single,prefix,radix_cache} --n-trials 50, and --test-mode prefix --return-logprob;
• a bitwise batch-size check: single's batches (the same prompt repeated 1 to 50 times, 1,275 responses), requesting logprobs and comparing every response's token IDs and selected-token logprobs with the batch-size-1 response, exact equality.

Model	TP	single	prefix	prefix logprobs	radix_cache	bitwise (token IDs and logprobs)
Qwen3-8B	1	1 output	1 per prompt	identical	pass	1,275 / 1,275
Qwen3-8B	2	1 output	1 per prompt	identical	pass	1,275 / 1,275
Qwen3-8B	4	1 output	1 per prompt	identical	pass	1,275 / 1,275
Qwen3-8B	8	1 output	1 per prompt	identical	pass	1,275 / 1,275
Qwen3.8-27B	1	1 output	1 per prompt	identical	pass	1,275 / 1,275
Qwen3.8-27B	2	1 output	1 per prompt	identical	pass	1,275 / 1,275
Qwen3-30B-A3B (MoE)	1	1 output	1 per prompt	identical	pass	1,275 / 1,275
Qwen3-8B, --kv-cache-dtype fp8_e4m3	1	1 output	1 per prompt	identical	pass	1,275 / 1,275
Qwen3.8-27B, --kv-cache-dtype fp8_e4m3	1	1 output	1 per prompt	identical	pass	1,275 / 1,275

The TP1 Qwen3-8B row was repeated on two more fresh servers with this PR's final code, with the same result.

Long contexts: test_deterministic's prompts stay near 1k tokens (at most about 4 KV partitions per request). So we also sent 8k, 12k and 16k-token prompts (token IDs), each alone twice and in one batch with 12 other requests of 64-4,096 tokens, with 4 different sets of batchmates, and compared the generated token IDs and logprobs exactly with the first batch-of-1 result: identical in 15 / 15 comparisons for Qwen3-8B at TP1 and TP2 and for Qwen3.8-27B at TP1.

GSM8K (sgl-eval, all 1,319 questions, temperature 0, thinking off), Triton decode / AITER decode: Qwen3-8B 93.56% / 93.48%; Qwen3-30B-A3B 94.77% / 94.39%; Qwen3.8-27B 96.06% / 95.83%.

A per-layer check at TP1 and TP2 (Qwen3-8B; batch sizes 2-50 and mixed batches; forward hooks hashing each op's inputs and outputs) found no op whose output changed while its input was bitwise equal; attention was compared 5,940 times in decode per rank.

--test-mode p_vs_d (prefill vs decode logprobs of the same tokens) differs, as it does with Triton for both phases; it is not part of the AMD deterministic CI.

Unit tests, python3 -m pytest test/registered/unit/test_model_overrides.py -k deterministic (MI350X, same image): with #44438 only, the new test fails (1 failed, 6 passed); with this PR, 6 passed.

With this PR applied, a Qwen3-8B server with the command above starts with AITER decode and Triton prefill (radix cache on) and passes a short bitwise batch check; with SGLANG_USE_AITER_UNIFIED_ATTN=1 it refuses to start with the new error; with --speculative-algorithm NGRAM it is refused: "... but --decode-attention-backend is 'aiter' (aiter decode is accepted only on gfx950, for non-MLA models and without speculative decoding)". With only --prefill-attention-backend triton --decode-attention-backend aiter (no --attention-backend) the server starts with AITER decode and Triton prefill; --attention-backend itself then resolves to the existing deterministic default (fa3), which only the per-phase backends override.

Speed Tests and Profiling

Deterministic mode with AITER decode (this PR) vs Triton for both phases, same server otherwise, MI350X, same image and fixes as above. python3 -m sglang.benchmark.serving --backend sglang --dataset-name random --random-input-len {1024,8192} --random-output-len 1024 --random-range-ratio 1.0 --num-prompts 256 --max-concurrency {256 for 1k/1k, 128 for 8k/1k} --request-rate inf --seed 1 --warmup-requests 8 --flush-cache; fresh server per run, 3 rounds interleaved, medians (rounds within 0.5%):

Model	TP	Workload	Triton decode, tok/s	AITER decode, tok/s	Throughput	ITL p50 (Triton / AITER)
Qwen3-8B	1	1k/1k	3,691	7,064	1.91x	60.6 / 27.1 ms
Qwen3-8B	1	8k/1k	737	1,284	1.74x	116.2 / 42.2 ms
Qwen3-8B	2	1k/1k	5,783.6	9,269.9	1.60x	37.0 / 20.3 ms
Qwen3-8B	2	8k/1k	1,224	1,933	1.58x	63.5 / 25.0 ms
Qwen3-8B	8	1k/1k	12,356.0	15,627.0	1.27x	16.0 / 11.6 ms
Qwen3-8B	8	8k/1k	2,907.8	3,795.3	1.31x	22.3 / 12.0 ms
Qwen3.8-27B	1	1k/1k	1,811	2,032	1.12x	46.6 / 39.2 ms
Qwen3.8-27B	2	1k/1k	3,833	4,332	1.13x	49.7 / 42.0 ms
Qwen3.8-27B	1	8k/1k	653	726	1.11x	72.0 / 52.8 ms
Qwen3.8-27B	2	8k/1k	1,085	1,186	1.09x	46.1 / 36.0 ms
Qwen3-30B-A3B	1	1k/1k	4,159	6,700	1.61x	55.4 / 31.7 ms
Qwen3-30B-A3B	1	8k/1k	851	1,303	1.53x	91.4 / 39.1 ms

Qwen3.8-27B gains less because only a quarter of its layers use attention (the rest are Gated DeltaNet). TTFT p50 is unchanged where measured (prefill still uses Triton), except 12% lower for Qwen3.8-27B at TP1 1k/1k, from faster decode steps under load. The Qwen3-8B TP1, Qwen3-30B-A3B and Qwen3.8-27B 8k/1k rows were measured with this PR's final code; the others with its first version (same decode path). Each row compares both configurations on one node.

For scale, against normal mode with AITER attention on the same node: Qwen3-8B TP1 deterministic mode goes from 0.382x to 0.731x (1k/1k) and from 0.329x to 0.572x (8k/1k); at TP8 from 0.377x to 0.476x (1k/1k) and from 0.301x to 0.393x (8k/1k); Qwen3.8-27B TP2 8k/1k from 0.587x to 0.641x.

No change with deterministic mode off.

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.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ✅ Run #41101544435
Latest PR Test (Extra): ❌ Run #41101544353
Latest PR Test (AMD ROCm 10): ✅ Run #41101544477

[rkarhila-amd]

Logic is terrible.

Agents have the capability to shift through these if-elses, for humans please make this easier to follow.

[aarnetalman]

Logic is terrible.

Agents have the capability to shift through these if-elses, for humans please make this easier to follow.

Thanks, agreed. Updated: the check now builds the allowed backends for each phase first (aiter added for decode only on gfx950 for non-MLA models), then checks both phases in one loop, with no special case inside it. It also reads the per phase fields directly instead of using getattr(..., None).

[rkarhila-amd]

Much better now.

[rkarhila-amd]

On a second thought (what luxury to have a little time for thinking every now and then) we're approaching the problem from the wrong side. Aiter is just a library of kernels, and some of those kernels are deterministic and some not.

It depends on the model definitions and aiter version which kernels are actually used for a workload.
A blanket yes for using decode attention kernels works at this point of time for Qwen model, but does not guarantee decode determinism for other models or future releases of aiter - maybe someone comes up with a faster "accurate enough" kernel and applies it to qwen and then this check is not sufficient.

We need to get the information on determinism from the actual aiter implementation. I have not looked at all if that is kept track of, if it is possible to get this information from the aiter library.

[aarnetalman]

@rkarhila-amd Thanks for the feedback. Agreed, that's the right long-term shape, and it's how I'm doing the MoE: ROCm/aiter#6583 adds fused_moe(..., deterministic=True), so AITER itself picks a batch-invariant kernel and raises for configurations it can't make deterministic. I checked AITER main for the same for decode attention: there isn't one today, no flag and no recorded property.

Until there is, this PR is narrower than a blanket yes. It allows AITER only for decode, only on gfx950, and only on the code path we validated: non-MLA, the default paged KV layout (not unified attention or the 5D layout), no sliding-window models, no speculative decoding. Everything else is refused with the reason. That path was checked bitwise with Qwen3-8B (TP1-TP8), Qwen3-30B-A3B, Qwen3.8-27B, FP8 KV cache and 8k-16k contexts.

Your point about a future AITER release swapping in a faster kernel is the real gap, since nothing would notice. Two follow-ups I'd propose:

1. A kernel-level test in the AMD CI that runs AITER's paged decode on the same requests alone and in batches of different sizes and compares bitwise, so an AITER update that loses batch invariance fails CI instead of silently breaking determinism.
2. Create a deterministic option in AITER's paged decode attention, like the MoE one, that selects the batch-invariant kernel and refuses configurations it can't make deterministic, and switch SGLang to request it.

[aarnetalman]

/rerun-failed-ci

[aarnetalman]

/rerun-failed-ci

[RolaoDenthu]

LGTM. All ci passed. This PR mainly updates the argument checks to allow AITER decode in deterministic mode. The common-path changes should not affect other configurations.