[aarnetalman]

Motivation

In deterministic mode, SGLang runs every bf16 linear layer through a batch-invariant Triton matmul. At decode sizes (a few to a few hundred tokens) that kernel launches only a few programs and leaves most of an MI350X idle, so the decode GEMMs are one of the largest remaining costs of deterministic mode on gfx950.

This PR adds a segmented split-K matmul for bf16 on gfx950: K is split into 4 fixed segments that are summed in a fixed order at every batch size, and at small batch sizes the segments run in separate programs, so decode GEMMs use 4 times as many programs. A token's result still does not depend on its batch.

Impact

• Deterministic-mode throughput on MI350X: Qwen3-8B +12% (1k/1k), +12% (4k/512), +3% (8k/1k); Qwen3.8-27B +12% / +14%; Phi-4 +14% / +14% / +5%. Batch invariance and accuracy unchanged (tables below).
• Only --enable-deterministic-inference on gfx950 with bf16 models changes. Normal mode, MI300 (gfx942), NVIDIA, and fp16 / fp32 matmuls are unchanged.
• The reduction order changes, so deterministic-mode outputs on gfx950 differ in the last bits from before this PR, as with any change of the batch-invariant kernel.

Design note: the kernel never software-pipelines the loop over segments within one program (num_stages 1 there). With pipelining, some K values were not run-to-run deterministic on MI350X (seen only with 2 segments; cause not understood), and a unit test keeps it off.

Modifications

• batch_invariant_ops.py:
  • _segmented_matmul_kernel (Triton): per output tile, sums 4 fixed K segments in order; SPLIT runs one segment per program and stores fp32 partials, otherwise one program loops over all 4. Bias is added after the segment sum. int64 offsets.
  • _segment_reduce_kernel: adds the 4 partials in order (and the bias) and casts to the output dtype.
  • _segmented_matmul_launch(M, N): tile, warps, num_stages and split by M and N, measured on MI350X over the bf16 linear shapes of Qwen3-8B, Qwen3.8-27B, Qwen3-30B-A3B and gpt-oss-20b (M = 1-8192). Split up to M = 1,024 (M = 128 when N > 8,192); unsplit launches use num_stages 1.
  • matmul_persistent: on gfx950, bf16 inputs with bf16 output and K above 256 (_SEGMENT_MIN_K) go to _matmul_segmented_triton.
• Tests:
  • test/registered/unit/batch_invariant_ops/test_segmented_matmul.py (gfx950, registered on the MI35x per-PR suite): bitwise equality of each row across batch sizes 1-4096 (with and without bias), run-to-run equality, the same bits for every launch configuration in the table and for split vs unsplit, and relative error against fp32 below 1%; shapes include K = 2,880 (not a multiple of the segment size) and N = 32. A short-K test checks that rows at K = 96-512 do not depend on the batch size through matmul_persistent.
  • test_matmul_persistent_configs.py (CPU): the segment constants, and that unsplit launches are never pipelined.

Accuracy Tests

MI350X, rocm/sgl-dev:v0.5.22-rocm724-mi35x-20261018 (ROCm 7.2.4), this PR's batch_invariant_ops.py over the image's SGLang, with the deterministic-mode fixes MI350X needed then (all merged since). Server:

SGLANG_USE_AITER=1 python3 -m sglang.launch_server --model-path <model> --tp <N> \
  --attention-backend triton --prefill-attention-backend triton --decode-attention-backend aiter \
  --enable-deterministic-inference

Checks: python3 -m sglang.test.test_deterministic --test-mode {single,prefix,radix_cache} --n-trials 50 (prefix also with --return-logprob), and a bitwise comparison of 1,275 responses (the same prompt in batches of 1-50, token IDs and log-probabilities) against batch size 1. Long context (Qwen3-8B with YaRN x4): a 64k and a 120k prompt alone vs next to other requests; a 100k prompt split at different chunk boundaries by other requests; 30k-token cached prefixes.

Model	TP	single	prefix (logprobs)	radix_cache	bitwise
Qwen3-8B	1, 2, 4, 8	1 unique	identical	pass	1,275/1,275
Qwen3-30B-A3B	1	1 unique	identical	pass	1,275/1,275
Qwen3.8-27B	2, 8	1 unique	identical	pass	1,275/1,275
Phi-4	1	1 unique	identical	pass	1,275/1,275
gpt-oss-20b (with ROCm/aiter#6583's deterministic MoE)	1	1 unique	identical	pass	1,275/1,275

The kernel's split and unsplit paths compute the same segment sums in the same order, so moving the split threshold does not change any output bit (checked on 200 shape and batch-size points, and for every launch in the table by the new unit test); the final launch table was re-checked end to end on Qwen3-8B TP1 (all of the above, and the long-context checks below).

Long context, Qwen3-8B TP1: 64k/120k 8/8 identical, chunk boundaries 12/12, cached 30k prefix 12/12.

GSM8K (1,319 questions, greedy), Qwen3-8B deterministic mode: 93.18% before, 93.33% after.

Unit tests on MI350X (this branch's tree on rocm/sgl-dev:v0.5.22-rocm724-mi35x-20261020): test_segmented_matmul.py and test_matmul_persistent_configs.py 10 passed (259 subtests), existing test_batch_invariant_ops.py 6 passed, 1 skipped.

Short K: before the fix in the latest commit, torch.mm in batch-invariant mode gave rows that depended on the batch size at K = 96, 128, 192 and 256 (N = 2,048-5,120, rows of batches of 1-1,024 vs the same rows in a batch of 4,096); this showed up end to end as a TP4 difference in Kimi-Linear, whose shared-expert down projection has K = 256 per rank. With the fix every audited shape (K = 96-2,304) is batch-invariant, the same as the persistent kernel on main.

Speed Tests and Profiling

Same node, image and stack as above, before and after this PR, interleaved in one session, 2 rounds each (both shown); deterministic mode, TP1; python3 -m sglang.benchmark.serving --backend sglang --dataset-name random --random-range-ratio 1.0 --request-rate inf --random-input-len <in> --random-output-len <out> --num-prompts <n> --max-concurrency <c>. Output throughput (tok/s):

Model	in/out, prompts, concurrency	Before	After	Change
Qwen3-8B	1024/1024, 256, 64	4,446 / 4,481	5,015 / 5,003	+12.2%
Qwen3-8B	4096/512, 128, 32	1,415 / 1,423	1,586 / 1,583	+11.7%
Qwen3-8B	8192/1024, 256, 128	1,506 / 1,510	1,559 / 1,557	+3.3%
Qwen3.8-27B	1024/1024, 64, 64	1,746 / 1,755	1,959 / 1,960	+11.9%
Qwen3.8-27B	8192/1024, 64, 32	552 / 556	630 / 635	+14.3%
Phi-4	1024/1024, 256, 64	2,912 / 2,923	3,307 / 3,316	+13.5%
Phi-4	4096/512, 128, 32	887 / 891	1,011 / 1,012	+13.8%
Phi-4	8192/1024, 256, 128	954 / 956	1,000 / 998	+4.6%

MoE models gain less, since their expert GEMMs do not use this path: Qwen3-30B-A3B about +2%, gpt-oss-20b +1-2% (earlier version of the launch table).

Kernel time (µs, CUDA graphs, MI350X, bf16), Qwen3-8B layers: current batch-invariant matmul / this PR (time ratio, above 1 = faster):

Layer (N x K)	M = 1	64	256	512	1,024	4,096
QKV (6144 x 4096)	25 / 12 (2.10x)	26 / 15 (1.76x)	28 / 33 (0.87x)	52 / 54 (0.98x)	86 / 102 (0.84x)	304 / 321 (0.95x)
O (4096 x 4096)	24 / 10 (2.37x)	25 / 13 (1.92x)	26 / 20 (1.31x)	29 / 35 (0.83x)	57 / 67 (0.85x)	202 / 212 (0.95x)
gate/up (24576 x 4096)	48 / 33 (1.47x)	49 / 41 (1.20x)	88 / 102 (0.86x)	161 / 191 (0.85x)	303 / 305 (0.99x)	1,503 / 1,307 (1.15x)
down (4096 x 12288)	67 / 25 (2.69x)	69 / 28 (2.43x)	69 / 41 (1.69x)	79 / 74 (1.07x)	153 / 147 (1.04x)	614 / 582 (1.06x)

The gain is at decode sizes, where most of a decode step's GEMM time goes; some mid-size prefill GEMMs are 1-17% slower, which the end-to-end numbers include.

No change without --enable-deterministic-inference.

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.

[aarnetalman]

I found a batch-invariance bug in this PR and pushed a fix (bb91ae0, together with a merge of current main).

When K is 256 or less, each of the 4 segments is just one 64-wide block. In that case the unsplit launch (used for large batches) rounds differently from the split launch (used for small batches), so a row's result could change with the batch size. I ran into it with Kimi-Linear at TP4: its shared-expert down projection has K = 256 per GPU, and the prefix determinism test failed.

The fix keeps matmuls with K <= 256 on the persistent kernel, the same as without this PR. Larger K is untouched, so the numbers in the description still hold. I also added a test for short K (96-512).

Checked on MI350X: no batch-size differences at any shape I tested (K from 96 to 2,304), the unit tests pass, and Kimi-Linear TP4 now passes every determinism check (prefix gives one output per prompt, before 1/2/1/2; 1,275/1,275 bitwise identical).