[aarnetalman]

Motivation

This PR speeds up deterministic inference on MI350X/MI355X (gfx950): about 6% more throughput with 1k-token inputs and about 1.5% with 8k-token inputs. It gives the batch-invariant bf16 matmul a launch configuration tuned for gfx950.

With --enable-deterministic-inference, every linear layer goes through the persistent batch-invariant Triton matmul (_matmul_persistent_triton). That kernel uses one bf16 configuration on every GPU (128x128x64 tiles, 3 stages, 8 warps), with no AMD-specific options. On gfx950, a 128x64x128 tile with 16-wide MFMA instructions is up to 2x faster at decode batch sizes.

The configuration is still fixed for every M, so each output element keeps the same reduction order whatever the batch: batch invariance is unchanged.

Modifications

• batch_invariant_ops.py: on gfx950 only, the bf16 entry becomes BLOCK_SIZE_M=128, BLOCK_SIZE_N=64, BLOCK_SIZE_K=128, GROUP_SIZE_M=8, num_stages=3, num_warps=8, matrix_instr_nonkdim=16.
  • It needs about 147 KB of LDS, so it is limited to gfx950 (160 KB).
  • Other GPUs, fp16/fp32 and the non-deterministic path are unchanged.
  • In deterministic mode every unquantized bf16 linear layer (aten::mm/addmm) runs this kernel on ROCm, so the change applies to any bf16 model. Quantized layers, MoE experts (Triton MoE runner) and attention don't use it.
• Chosen from a sweep of about 2,900 configurations on MI350X, each checked bitwise for batch invariance.
• The launch configurations are now built by a small helper, _matmul_persistent_configs(), and a new CPU test (test/registered/unit/batch_invariant_ops/test_matmul_persistent_configs.py, base-a-test-cpu) guards the preset: on gfx950 the bf16 entry must be exactly this configuration and fit gfx950's 160 KB of LDS, and on other GPUs it must stay the original one.

Accuracy Tests

MI350X, rocm/sgl-dev:v0.5.21-rocm724-mi35x-20261012 with this change applied, Qwen3-8B BF16, TP1, --attention-backend triton --enable-deterministic-inference, SGLANG_USE_AITER=0:

• python3 -m sglang.test.test_deterministic --test-mode prefix --return-logprob and --test-mode radix_cache (50 trials): pass.
• Bitwise comparison of 1,275 responses (batch sizes 1-50) against batch size 1: 1275/1275 token IDs and log-probabilities identical.
• test/registered/unit/batch_invariant_ops/test_batch_invariant_ops.py: 6 passed, 54 subtests, before and after (also with the helper refactor).
• test_matmul_persistent_configs.py: 2 passed (2 failed on the code before this PR, which has no helper).

Speed Tests and Profiling

End to end, deterministic mode as above, python3 -m sglang.benchmark.serving --backend sglang --dataset-name random --random-range-ratio 1.0 --request-rate inf --seed 1, median of 3 runs each:

Workload		Output tok/s	ITL p50	TTFT p50	Request latency p50
ISL/OSL 1024/1024, 256 prompts, concurrency 256	main	3,370	66.79 ms	5.03 s	77.7 s
	this PR	3,585 (1.064x)	61.98 ms (-7.2%)	5.28 s (+4.8%)	73.0 s (-6.0%)
ISL/OSL 8192/1024, 256 prompts, concurrency 128, GPU A	main	686.9	121.61 ms	34.5 s	190.8 s
	this PR, GPU B	695.4 (1.012x)	118.34 ms (-2.7%)	35.1 s (+1.7%)	188.4 s (-1.3%)
same, GPUs swapped	main, GPU B	686.0	121.36 ms	34.8 s	191.0 s
	this PR, GPU A	697.6 (1.017x)	118.27 ms (-2.5%)	34.8 s (0.0%)	187.9 s (-1.6%)

• 1k/1k: same GPU, before and after; an earlier run with the GPUs swapped gave the same gain (1.059x on both GPUs).
• With 8k inputs most of the time goes to prefill, where the new tile is about 5% slower at the GEMM level (table below), so the gain is smaller. Decode is still faster, and request latency still improves.

Kernel, speed-up over the current configuration (triton.testing.do_bench, weights passed transposed as F.linear does):

N x K	M=1	M=16	M=64	M=256	M=1024	M=4096
6144 x 4096	1.77x	1.92x	1.84x	1.60x	1.16x	0.96x
4096 x 4096	1.98x	1.94x	1.85x	1.57x	0.95x	0.96x
24576 x 4096	1.02x	1.01x	1.00x	1.14x	0.95x	0.96x
4096 x 12288	2.05x	2.02x	1.87x	1.56x	0.96x	0.95x
16384 x 5120	1.16x	1.15x	1.15x	0.88x	0.94x	0.94x
5120 x 6144	2.03x	1.99x	1.87x	1.58x	1.15x	0.98x
34816 x 5120	1.21x	1.20x	1.21x	1.04x	0.99x	0.96x
5120 x 17408	2.13x	2.11x	1.96x	1.57x	1.11x	0.95x

The shapes are the linear layers of Qwen3-8B and Qwen3.8-27B. The geomean is 1.59x at M <= 64 and 0.99x at M >= 1024. A per-(N, K) configuration could recover the prefill shapes later; the configuration must never depend on M.

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 #40028266705
Latest PR Test (Extra): ❌ Run #40028266638
Latest PR Test (AMD ROCm 10): ❌ Run #40028266701

[HaiShaw]

How to justify the tuning pertain to Qwen3-8B and Qwen3.8-27B to be generic in a common code path?

[aarnetalman]

@HaiShaw Good question. I checked the preset on 19 layer shapes it was not tuned on (Llama 3.1 8B and 70B TP8, Phi-4, gpt-oss-20b, Qwen3-30B-A3B, Qwen3-8B TP8, three LM heads), M = 1-8192, against the default config on MI350X:

	Decode (M <= 256)	Prefill (M >= 512)
Tuned-on shapes (Qwen3-8B)	1.81x	1.03x
19 untuned shapes	1.54x	1.03x

The gain is structural: at decode the output has N / BLOCK_N tiles, and the default 128-wide N tile leaves most of MI350X's 256 CUs idle; the preset's 64-wide tile doubles them. So it holds for any layer up to N ~ 8k (1.58-1.89x at decode). It does not help very wide layers, which already fill the GPU: Llama 3.1 8B gate_up (N = 28,672) 0.94x, LM heads 0.97-1.02x. Every shape stayed batch-invariant, and the preset only applies to bf16 on gfx950.

[aarnetalman]

@HaiShaw I also updated the code comment on the preset (961702b). It no longer names Qwen3-8B and Qwen3.8-27B; it now describes the shapes it was measured on: faster than the default on average for linear layers up to N ~ 8k, neutral on wider layers.

[aarnetalman]

/rerun-failed-ci

[aarnetalman]

/rerun-failed-ci

[aarnetalman]

/rerun-failed-ci

[1am9trash]

All pr-test-base passed.
Link: https://github.com/sgl-project/sglang/actions/runs/40028266705/job/120852732553?pr=44045

[1am9trash]

Hi, @aarnetalman
bmm_batch_invariant still hardcodes the same config. Should it use _matmul_persistent_configs() as well?

[aarnetalman]

Good question. I measured it, and the matmul preset isn't a straight win for bmm. Its main deterministic-mode user is MLA attention, whose two batched matmuls have very different shapes from the linear layers the preset was tuned on. On MI350X, here is the preset against the current bmm config:

bmm shape (heads x K x N)	M 1-128	M 1,024	M 4,096
16 x 512 x 128	1.44-1.51x faster	1.39x	0.97x
128 x 512 x 128	1.09-1.36x faster	0.99x	0.87x
16 x 128 x 512	0.99-1.01x	0.86x	0.80x
128 x 128 x 512	0.75-0.83x (slower)	0.78x	0.75x

Both configs are batch-invariant: every row is bitwise the same at every M. But the gains and losses depend on K and N, so switching bmm wholesale would slow the K=128, N=512 matmul.

I'd keep this PR to the 2D matmul it was tuned and measured on, and do bmm as a follow-up. There, the config would be picked by the bmm's own shape (K and N, never M, so it stays batch-invariant), with end-to-end numbers on an MLA model. Does that work for you?

[1am9trash]

Sounds great.
Thanks for clarification.

[1am9trash]

LGTM.
Only config change on gfx950.