[aarnetalman]

Motivation

In deterministic mode, RMSNorm with a residual add (both norms in every decoder layer) falls back to eager PyTorch, which takes about ten kernel launches instead of a single fused kernel. This PR adds a fused, batch-invariant add+RMSNorm kernel for that path. It applies to any model whose decoder layers use RMSNorm with a residual (Llama, Mistral, Qwen2/3, Qwen3-MoE, Phi, DeepSeek, gpt-oss and most others). Tested on MI350X with Qwen3-8B, Qwen3-30B-A3B and Phi-4: single-request decode throughput with the fused kernel is 56-73% higher and 12-22% higher at concurrency 16, and every determinism check passes on Qwen3-8B (TP1, TP2) and Qwen3-30B-A3B.

The fallback is a large share of deterministic decode time. Profile of a single-request decode (1k-token prompt, 256 output tokens, CUDA graphs off) in deterministic mode on MI350X:

model	GPU time, deterministic	GPU time, normal mode	eager add+RMSNorm fallback
Qwen3-8B	603 ms	323 ms	about 205 ms (34%)
Qwen3-30B-A3B	678 ms	298 ms	about 215 ms (32%)
Phi-4	902 ms	474 ms	about 205 ms (22%)

In normal mode the same norms take 12.6 ms on Qwen3-8B as one fused kernel.

Scope: deterministic mode only, on every platform that runs the batch-invariant Triton kernels. The new kernel uses the same float32 math as forward_native. Models that use GemmaRMSNorm (Gemma, Qwen3-Next/3.5/3.8) take a different path and are not affected. Normal mode is unchanged.

Modifications

• batch_invariant_ops.py: new _fused_add_rms_norm_kernel and fused_add_rms_norm_batch_invariant, next to the existing rms_norm kernel and in the same style. One program per row:

  • adds the residual (and the optional post-residual addend) to the input in float32;
  • stores the sum as the new residual, in the input's dtype or float32 (fp32_residual);
  • RMS-normalizes the float32 sum and applies the weight.

  The reduction runs over one row with a fixed block size, so a row's result does not depend on the batch.

• layernorm.py: in the batch-invariant branches of RMSNorm.forward_cuda, forward_aiter, forward_hip and forward_xpu, calls with a residual use the new kernel where forward_native's math applies. These cases keep falling back to forward_native as before:

  • cast_x_before_out_mul;
  • a variance-size override;
  • a padded output;
  • the FSDP RL target.

• New test/registered/unit/batch_invariant_ops/test_fused_add_rms_norm_batch_invariant.py (per-PR, CUDA and MI35x):

  • every row's output and new residual are bit-identical (torch.equal) when computed in batches of 1, 3, 17 and 64;
  • the new residual equals the eager float32 sum exactly, and the output matches the eager path within rounding;
  • covered: hidden sizes 2048/4096/5120/2880, bf16 and fp16, with and without the addend, fp32_residual on and off.

• Unchanged: normal mode; batch-invariant calls without a residual; the fallback cases above.

Accuracy Tests

MI350X, ROCm 7.2.4, rocm/sgl-dev:v0.5.22-rocm724-mi35x-20261021.

Unit tests:

python3 -m pytest -q test/registered/unit/batch_invariant_ops/test_fused_add_rms_norm_batch_invariant.py test/registered/unit/batch_invariant_ops/test_batch_invariant_ops.py

7 passed, 1 skipped, 86 subtests passed in 16.00s

End to end, Qwen3-8B TP1, deterministic mode with AITER prefill and decode attention. These runs use the open MI350X deterministic-mode PRs (#43788, #44045, #44438, #44489, #45599, #45601, #34543), which is the configuration that deterministic mode runs in on MI350X today; "before" is the same tree without this change:

python3 -m sglang.launch_server --model-path Qwen/Qwen3-8B --enable-deterministic-inference \
  --attention-backend triton --prefill-attention-backend aiter --decode-attention-backend aiter
python3 -m sglang.test.test_deterministic --test-mode {single,prefix,radix_cache} --n-trials 50

With this change, Qwen3-8B: single 1 output, prefix 1/1/1/1 (twice), prefix --return-logprob identical, radix_cache pass, and a bitwise batch-size check (the same prompt in batches of 1 to 50, 1,275 responses, token IDs and logprobs compared exactly with the batch-size-1 response) 1,275 / 1,275. Qwen3-30B-A3B (same setup): single 1, prefix 1/1/1/1, bitwise 1,275 / 1,275. Qwen3-8B at TP2: single 1, prefix 1/1/1/1, bitwise 1,275 / 1,275.

GSM8K (sgl-eval 0.1.2, --temperature 0 --max-tokens 2048, all 1,319 questions, Qwen3-8B, deterministic): 93.18% with this change (before: 93.25% deterministic, 93.48% normal mode, same image).

Speed Tests and Profiling

Deterministic mode (setup above, CUDA graphs on), before vs after on the same node and image, python3 -m sglang.benchmark.serving --backend sglang --dataset-name random --random-range-ratio 1.0 --request-rate inf --max-concurrency C (output throughput, median inter-token latency):

model	workload	before	after	change
Qwen3-8B	1k in / 256 out, concurrency 1	105 tok/s, 9.42 ms	173 tok/s, 5.64 ms	+65%, ITL -40%
Qwen3-8B	1k / 256, concurrency 16	1,663 tok/s, 7.70 ms	2,033 tok/s, 6.14 ms	+22%, ITL -20%
Qwen3-8B	8k / 256, concurrency 16	582 tok/s, 10.71 ms	653 tok/s, 9.11 ms	+12%, ITL -15%
Qwen3-30B-A3B	1k / 256, concurrency 1	94 tok/s, 10.5 ms	163 tok/s, 6.0 ms	+73%, ITL -43%
Qwen3-30B-A3B	1k / 256, concurrency 16	1,070 tok/s, 13.55 ms	1,282 tok/s, 11.18 ms	+20%, ITL -17%
Phi-4	1k / 256, concurrency 1	72 tok/s, 13.74 ms	112 tok/s, 8.70 ms	+56%, ITL -37%

A second alternating round for Qwen3-8B gives the same numbers within 1% (before 105 tok/s / 9.40 ms and 1,655 tok/s / 7.75 ms; after 174 tok/s / 5.64 ms and 2,036 tok/s / 6.11 ms).

For reference, normal mode on the same setup has an inter-token latency of 4.8 ms (Qwen3-8B) and 4.4 ms (Qwen3-30B-A3B) at concurrency 1, so this closes most of the deterministic-mode gap for single-request decode. No change outside deterministic mode.

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 #40025448735
Latest PR Test (Extra): ❌ Run #40025448730
Latest PR Test (AMD ROCm 10): ❌ Run #40025448697

[HaiShaw]

/tag-and-rerun-ci

[1am9trash]

All pr-test-base passed.
Link: https://github.com/sgl-project/sglang/actions/runs/40025448735/job/120733431113?pr=46159

[1am9trash]

Hi @aarnetalman
This PR replaces forward_native on the residual path, and the new kernel is used from all platform.
Should it be guarded (for example with is_hip)?

[aarnetalman]

Thanks @1am9trash . I'd keep it platform-neutral, for these reasons:

• It fills a gap that exists on every platform. Batch-invariant mode already runs RMSNorm without a residual through the Triton rms_norm_batch_invariant on every platform, including forward_cuda (since Add Batch‑Invariant RMSNorm #12144). Only the residual case fell back to the eager forward_native. This PR handles that case in the same three places (forward_cuda, forward_aiter, forward_hip). An is_hip guard would leave CUDA with a Triton kernel for one case and the eager path for the other.
• Nothing in the kernel is ROCm-specific. It is plain Triton and uses the same float32 math as forward_native: the new residual is bitwise equal to the eager path, and the output is within bf16 rounding of it. Each row is reduced in one program, so it is batch-invariant by construction.
• It's tested on NVIDIA in CI. The unit test (bitwise batch invariance, plus a comparison with the eager math) is registered on the CUDA base-b 1-GPU suite as well as on MI35x, and the base tests passed.
• Only batch-invariant mode changes. Normal mode is unchanged on every platform.

I haven't measured it end to end on NVIDIA. If you'd rather limit it to ROCm until someone has, I can add the guard and leave CUDA for a follow-up.

[1am9trash]

Okay, that works for me.
I was just worried about breaking servers on other platforms. Since the pr-test-base CI covers it, I'm fine keeping it platform-neutral.

[1am9trash]

LGTM.