[Bugfix][Kernel] Fix batch invariance in RMSNorm kernels by pinning block size - #48391
Merged
DarkLight1337 merged 33 commits intoJul 28, 2026
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
Contributor
Author
|
@yewentao256 pls help to review this |
yewentao256
reviewed
Jul 12, 2026
yewentao256
left a comment
yewentao256
left a comment
Member
There was a problem hiding this comment.
Thanks for the work!
Please attach with command line in main with full log output that can reproduce the issue, and in your branch this issue got solved.
Contributor
Author
|
@yewentao256 I've updated the PR with the reproduction commands and full logs as requested. Let me know if you need anything else. |
yewentao256
reviewed
Jul 13, 2026
yewentao256
left a comment
yewentao256
left a comment
Member
There was a problem hiding this comment.
VLLM_BATCH_INVARIANT=1 pytest tests/v1/determinism/test_rms_norm_batch_invariant.py
======================== 93 passed, 16 warnings in 23.75s =========================I can't reproduce it in main. What's device are you currently using?
Contributor
Author
@yewentao256 Thanks for checking! The 93 passed means the run used main's old version of this test, not the one in this PR — on main, test_fused_add_rms_norm_batch_invariant_residual_path compares fixed shapes (1 row vs 4 rows, both < 256 tokens), so max_block_size stays 1024 in both and it can never trigger the bug. Could you pull this PR's test file and re-run? The updated test adds n_extra=299 (num_tokens=300 ≥ 256, flipping block size to 256) plus a seed sweep, so that single test is 128 cases: git checkout main
git fetch origin pull/48391/head
git checkout FETCH_HEAD -- tests/v1/determinism/test_rms_norm_batch_invariant.py
VLLM_BATCH_INVARIANT=1 pytest tests/v1/determinism/test_rms_norm_batch_invariant.py::test_fused_add_rms_norm_batch_invariant_residual_path -qThis reproduces 5 failed, 123 passed on the current main kernel. My environment: GPU: NVIDIA L40S (compute capability 8.9), driver 580.159.03
OS: Ubuntu 22.04.5 LTS, kernel 6.8.0-1055-aws
Python: 3.10.12
PyTorch: 2.11.0+cu130 (CUDA 13.0, cuDNN 91900) |
yewentao256
reviewed
Jul 14, 2026
yewentao256
left a comment
yewentao256
left a comment
Member
There was a problem hiding this comment.
Thanks for the work! Could you also fix rms_norm_static_fp8_quant as there are similar issues?
oops-oom
changed the title
oops-oom
changed the title
oops-oom
changed the title
oops-oom
changed the title
Contributor
Author
|
@yewentao256 Thanks for the review! I've addressed the feedback — the block-size pinning is now applied consistently across all five kernels (fused_add_rms_norm, fused_add_rms_norm_static_fp8_quant, rms_norm_static_fp8_quant, rms_norm, rms_norm_per_block_quant). Would you mind taking another look when you have time? Happy to make any further changes. FYI, here's the full map of where an RMSNorm layer actually lands and which paths reach the fused kernels.
Note In eager mode (rms_norm custom op enabled) under Batch Invariance, torch.ops._C.rms_norm is not called on either branch: no-residual short-circuits to the Triton _rms_norm_kernel, and the residual path goes to torch.ops._C.fused_add_rms_norm (a different CUDA op — which does run, and is what #48391 pins). The plain _C.rms_norm kernel is reachable only via forward_native + --ir-op-priority.rms_norm=vllm_c. vllm serve meta-llama/Llama-3.2-1B \
--ir-op-priority.rms_norm=vllm_c \
--ir-op-priority.fused_add_rms_norm=vllm_cIn this override, the non-residual C++ kernels execute with batch invariance, making the fixed block size a strict dependency.
|
…lock size The RMSNorm-family CUDA kernels pick their block size from the token count (`max_block_size = (num_tokens < 256) ? 1024 : 256`, and `(num_tokens <= 256) ? 512 : 256` for the per-block quant kernel), independent of whether batch-invariant mode is enabled. Under `VLLM_BATCH_INVARIANT=1` the residual RMSNorm path (`RMSNorm.forward_cuda` with a residual -> `ops.fused_add_rms_norm`) routes to these kernels, so the *same* token reduces with a wide block when processed in a small batch but a narrow block inside a batch of >= 256 tokens. The two block sizes give different `cub::BlockReduce` partitions and therefore a different fp32 sum-of-squares order, so the normalized output is not bit-exact across batch sizes whenever `hidden_size` exceeds the block threshold. Lock the block size to a `num_tokens`-independent constant when `vllm_is_batch_invariant()` is true, across all affected kernels: - `rms_norm` / `fused_add_rms_norm` (layernorm_kernels.cu) - `rms_norm_static_fp8_quant` / `fused_add_rms_norm_static_fp8_quant` (layernorm_quant_kernels.cu) - `rms_norm_per_block_quant` (fused_layernorm_dynamic_per_token_quant.cu), pinned to 512 (its existing small-batch value, already an exercised launch config) to stay valid for the per-group/per-warp reduction math Extend the batch-invariant RMSNorm determinism test to sweep seeds and cross the `num_tokens=256` block-size threshold, which the previous single-seed, below-threshold test could not catch. Signed-off-by: oops-oom <73481342@qq.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
oops-oom
force-pushed
the
fix/fused-add-rms-norm-batch-invariant-block-size
branch
from
July 20, 2026 02:54
c14f2f2 to
a927cd1
Compare
Contributor
Author
|
@yewentao256 pls help to review this |
yewentao256
reviewed
Jul 20, 2026
yewentao256
left a comment
yewentao256
left a comment
Member
There was a problem hiding this comment.
Thanks for the work!
Please benchmark performance before/after your change and also possibly add unit tests to cover other kernel changes.
BTW, I am curious why e2e acc doesn't break without your PR in main.
Contributor
|
Tested on NVIDIA H100 NVL (95830 MiB). Note: this run used the test file from vLLM main, not the PR branch — which is why all 95 pass. As @oops-oom explained above, main's Command: VLLM_BATCH_INVARIANT=1 pytest tests/v1/determinism/test_rms_norm_batch_invariant.py -vFull outputHardware: |
Signed-off-by: oops-oom <73481342@qq.com>
Sweep seeds and add a >=256-token batch case to the Triton vs standard RMSNorm comparison, and add batch-invariance regression tests for the C++ rms_norm, rms_norm_static_fp8_quant, fused_add_rms_norm_static_fp8_quant, and rms_norm_per_block_quant kernels. Each feeds the same rows through small launches (num_tokens<256) and one large launch (num_tokens=300) that crosses the block-size threshold, asserting every row is bit-for-bit identical so the ~0.1%-of-rows reduction-width divergence cannot hide. Signed-off-by: oops-oom <73481342@qq.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Contributor
Author
|
Thanks for the review! benchmark performancePinning to the larger block does not regress . Note: I expected a small regression here, since pinning to 1024 overrides the batch-size-tuned 256 for large launches.The measured delta is within noise instead (−0.009%, see table).
add unit tests to cover other kernel changesDone. why e2e acc doesn't break in mainThe buggy kernel isn't on the default path. Under torch.compile (the default), RMSNorm lowers through the native IR-op priority — RMSNorm.forward_native → ir.ops.rms_norm → native ATen, which Inductor then codegens into Triton. The num_tokens-dependent block-size bug lives only in the C++ vllm_c kernels, which this default path never invokes — so the existing e2e determinism test, So to exercise the fix end-to-end I extended test_batch_invariance.py with a rms_norm_impl parameter: the "vllm_c" variant sets that kernel_config, routing the model's RMSNorm to the C++ kernels so the needle test actually crosses the block-size threshold. |
yewentao256
reviewed
Jul 22, 2026
yewentao256
left a comment
yewentao256
left a comment
Member
There was a problem hiding this comment.
Generally LGTM, several small updates
Contributor
Author
|
CI won't pass until this PR #50060 is merged. |
29 tasks
aoshen02
added a commit
to aoshen02/vllm
that referenced
this pull request
Aug 18, 2026
compute_num_split derives the K-split count from n_sms // cdiv(num_tokens, 64), so the reduction tree changes with the batch (same defect class as the RMSNorm fix in vllm-project#48391); use_small_fma switches to a second implementation at num_tokens <= 16. Under VLLM_BATCH_INVARIANT=1 the split count is pinned (min(cap, n_sms//4) rounded down to a power of two = 32 on GB200; divides the 256 K-blocks, single wave up to M=256) and the small-fma fork is disabled. Cross-split merge is already a T.serial ordered loop, so pinning the count pins the reduction tree. CUDA-graph cost after pinning: parity with baseline for n<=128 and 193-256, faster at 129-192; only n=1 keeps +2us (fused kernel split into post+GEMM). Tests (10): bitwise stability across batch boundaries for mhc_pre / mhc_fused_post_pre / broadcast / fused-RMSNorm variants, mhc_post and hc_head regressions, negative controls with checkpoint-realistic magnitudes (synthetic small weights wash out real defects), correctness vs the torch reference. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Consolidation pass folded in (codex r10/r10b reviewed): shared test helpers, pure Triton key fn, tl.constexpr-instantiated constants (plain global ints fail to compile under Triton 3.7), repo-pinned ruff format. Container suite 39/39 green; topk equivalence probe 240/240 bitwise.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Fix a batch-invariance gap in the
fused_add_rms_normCUDA kernel.The kernel selects its block size from the token count:
max_block_size = (num_tokens < 256) ? 1024 : 256, independent of whether batch-invariant mode is enabled. UnderVLLM_BATCH_INVARIANT=1, the residual RMSNorm path (RMSNorm.forward_cudawith a residual →ops.fused_add_rms_norm) routes to this kernel. As a result the same token reduces withblock=1024when processed in a small batch (num_tokens < 256) but withblock=256inside a batch of>= 256tokens. The two block sizes give differentcub::BlockReducepartitions and therefore a different fp32 sum-of-squares order, so the normalized output is not bit-exact across batch sizes wheneverhidden_size > 256.The existing
batch_invariant_launchflag only disabled vectorization (width 8 → 0); it did not touch the block size. This PR locksmax_block_sizeto 1024 whenvllm_is_batch_invariant()is true, making the reduction width independent ofnum_tokens.This is a follow-up to #40413 (which routed the residual path to
fused_add_rms_normon the assumption it is already batch-invariant) and is tracked under #27433. It is not a duplicate: no open PR addresses the block-size-vs-num_tokensgap; #40413's test only comparednum_tokens=1vs4, both below the 256 threshold, so it could not catch this.Test
Both runs use this PR's updated test (
test_fused_add_rms_norm_batch_invariant_residual_path, the seed-swept version). The previous single-seed test (seed=42) passes on the buggy kernel, which is why it never caught this.The decisive parameter is
n_extra:n_extra=299makesnum_tokens=300 >= 256(crossing the block-size threshold and exposing the bug), whilen_extra=3(num_tokens=4) stays below it. That is why every failure below is ann_extra=299case, and the seed is swept so a single lucky seed can't hide it.n_extra=299→num_tokens = 300 >= 256, so that same token now reduces withblock=256— crossing the threshold and exposing the bug.n_extra=3→num_tokens = 4 < 256, stays atblock=1024, so no divergence.That is why every failure below is an
n_extra=299case (athidden_size=4096),and why the seed is swept — a single seed can pass by luck even at 299.
On
main— issue reproducedOn this branch — issue solved