[aarnetalman]

Motivation

Deterministic inference needs batch-invariant decode attention: a request's output must not depend on how many other requests share its batch, or how long they are. On gfx950, the AITER attention backend decodes non-MLA models with AITER's paged_attention_ragged. That kernel splits each request's KV cache into fixed 256-token partitions, so its result for one request does not depend on its batchmates. Nothing tests that, though. An AITER update that changed it, for example by choosing the number of partitions from the batch, would make outputs batch-dependent without any test failing.

This adds a small kernel-level test to the MI35x per-PR suite. It pins that property, so a regression shows up in CI at the AITER bump that causes it, instead of as non-reproducible outputs later. #44489 proposes allowing AITER decode attention in deterministic mode on gfx950 on the strength of this property. This test is useful with or without it and does not depend on it.

Impact: test only. No runtime code changes. Adds about 35 s to the stage-b-test-1-gpu-small-amd-mi35x suite. It is skipped when AITER is not installed or the GPU is not gfx950.

Modifications

• New test/registered/attention/test_aiter_decode_batch_invariance.py, registered with register_amd_ci(est_time=60, suite="stage-b-test-1-gpu-small-amd-mi35x"):
  • Inputs: decodes one token per request with paged_attention_ragged, called the way AiterAttnBackend calls it (same argument order, partition size _AITER_PARTITION_SIZE_ROCM, same workspace layout).
  • Comparison: a target request is decoded alone, then inside batches of 2, 8 and 33 at a random position, next to batchmates with random KV lengths of 1 to 16,383 tokens. Its output row must be bit-identical (torch.equal) to the alone result.
  • Target lengths: 1, 255, 256, 257, 4,000 and 12,000 KV tokens, covering a partial partition, partition edges and many partitions.
  • Cache types: run for a bf16 KV cache and an fp8 (fp8_e4m3) KV cache.
• No other files change.

Accuracy Tests

Test only, so no model outputs can change.

MI350X (gfx950), ROCm 7.2.4, images rocm/sgl-dev:v0.5.22-rocm724-mi35x-20261018 and rocm/sgl-dev:v0.5.22-rocm724-mi35x-20261021, with the AITER each image ships:

python3 -m pytest -v test/registered/attention/test_aiter_decode_batch_invariance.py

test_bf16_kv_cache PASSED
test_fp8_kv_cache PASSED
2 passed in 33.88s   (20261018)
2 passed in 34.98s   (20261021)

The bitwise comparison is sensitive enough to catch a batch-dependent kernel. The same alone-versus-batch comparison, applied to AITER's prefill kernel mha_batch_prefill_func with one query token per request (prefix 4,096 and 32,768 tokens, 1, 3 and 15 batchmates), finds a difference in 6 of 6 cases (max abs diff 3.05e-5, one bf16 ulp). So the paged decode kernel is batch-invariant where that kernel is not, and a loss of the property would fail this test.

Speed Tests and Profiling

Test only; no runtime code changes, so no performance effect. The test takes about 35 s on MI350X.

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 #40475355281
Latest PR Test (Extra): ❌ Run #40475355210
Latest PR Test (AMD ROCm 10): ❌ Run #40475355201