[Bugfix] Fix RMSNorm kernels to multiply in weight's native dtype

[liulanze]

Purpose

Fix a regression introduced by #40860 (DeepSeek V4 support) where csrc/layernorm_kernels.cu was changed to always upcast the RMSNorm weight to float32 before the final multiplication. The correct behaviour, per the Python spec in vllm/ir/ops/layernorm.py:20, is to multiply in the weight's native dtype (BF16/FP16):

x = x.to(weight.dtype) * weight  # multiply in weight.dtype

The FP32 weight path introduces ~0.03 per-element error per layer. For models with many Q/K RMSNorm layers (e.g. LLaDA-2, DeepSeek with Q/K norm) this compounds catastrophically — 20-layer cumulative divergence exceeds 100, and only ~2% of output tokens match the reference.

The bug was reproduced locally on an A100-SXM4-40GB before implementing the fix, confirming the reported max diff of ~3e-02. The fix was built from source and verified on the same hardware.

Fixes #42325

Changes

csrc/layernorm_kernels.cu — 3 locations:

• rms_norm_kernel vectorized path
• fused_add_rms_norm_kernel FP16/BF16 optimised path
• fused_add_rms_norm_kernel generic scalar fallback

Pattern in each location:

// Before (regressed in #40860):
float w = static_cast<float>(weight_elem);
out = static_cast<scalar_t>(x * s_variance * w);

// After (restores v0.19.1 behaviour):
out = static_cast<scalar_t>(x * s_variance) * weight_elem;

csrc/layernorm_quant_kernels.cu — 3 matching locations:
The FP8 quant kernels contain an explicit comment explaining they were designed to match the non-fused composite path (rms_norm output → static_scaled_fp8_quant). After fixing the non-fused path, the quant kernels require the same update to stay consistent. Without it, test_fused_rms_norm_quant fails because the two paths diverge.

Test Plan

pytest tests/kernels/core/test_layernorm.py -x -q
pytest tests/kernels/ir/test_layernorm.py -x -q

Regression repro script (shows FIXED after patch):

import torch
from vllm import _custom_ops as ops

x = torch.randn(48, 128, dtype=torch.bfloat16, device='cuda')
weight = torch.randn(128, dtype=torch.bfloat16, device='cuda')
out = torch.empty_like(x)
ops.rms_norm(out, x, weight, 1e-6)

x_fp32 = x.to(torch.float32)
variance = x_fp32.pow(2).mean(-1, keepdim=True)
x_normed = x_fp32 * torch.rsqrt(variance + 1e-6)
ref = (x_normed.to(weight.dtype) * weight).to(x.dtype)

diff = (out.float() - ref.float()).abs().max().item()
print(f"Max diff: {diff:.6e}")
print("FIXED" if diff < 0.001 else "BUG CONFIRMED")

Test Result

Reproduced the bug and verified the fix locally on an A100-SXM4-40GB, building vllm from source on this branch:

	Before (bug reproduced)	After (fix verified)
rms_norm max diff (BF16)	3.125e-02 ✗	0.000e+00 ✓
fused_add_rms_norm max diff (BF16)	3.125e-02 ✗	0.000e+00 ✓

tests/kernels/core/test_layernorm.py: 865 passed
tests/kernels/ir/test_layernorm.py:   1442 passed, 361 skipped

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Note summary:

DeepSeek V4 support accidentally changed 6 CUDA kernel sites to upcast weights to float32 before multiplying, deviating from the spec (multiply in weight's native BF16 dtype); tiny per-layer rounding errors compounded across 20+ norm layers until the model output was nearly completely wrong — fixed by restoring the correct cast-then-multiply order.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

Hi @liulanze, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[gemini-code-assist]

Code Review

This pull request modifies the RMS normalization and fused FP8 quantization kernels in csrc/layernorm_kernels.cu and csrc/layernorm_quant_kernels.cu. The changes adjust the precision of the calculations by ensuring that the weight multiplication is performed in the weight's native data type after the normalized value has been cast to scalar_t. This update standardizes the calculation logic across various kernel implementations, including vectorized and non-vectorized versions. I have no feedback to provide as there were no review comments.

[liulanze]

Hi @AlonKellner-RedHat would appreciate your review when you get a chance. Happy to discuss the approach if anything is unclear. Thanks.

[liulanze]

Hi @yewentao256 / @mgoin, this touches layernorm_kernels.cu which has no CODEOWNER. Would appreciate a review when you have a moment!

[yewentao256]

Thanks for the work!
Please add a lm_eval metrics to show we don't hurt in e2e acc

[liulanze]

Thanks for the work! Please add a lm_eval metrics to show we don't hurt in e2e acc

Thanks @yewentao256, here are the lm_eval results comparing upstream/main (v0.21.0, unfixed kernels) against this fix branch, both using TinyLlama/TinyLlama-1.1B-Chat-v1.0 in BF16 on an H100 80GB:

Task	Metric	Baseline (v0.21.0)	Fix branch
arc_challenge	acc	0.3046 ± 0.0134	0.3046 ± 0.0134
arc_challenge	acc_norm	0.3276 ± 0.0137	0.3276 ± 0.0137
hellaswag	acc	0.4652 ± 0.0050	0.4652 ± 0.0050
hellaswag	acc_norm	0.6044 ± 0.0049	0.6044 ± 0.0049

Scores are identical with no regression. TinyLlama (RMSNorm, BF16) hits the fixed kernels, and I built from source to ensure the patched files were used.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @liulanze.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[yewentao256]

Please solve the conflicts

[liulanze]

Please solve the conflicts

Rebased onto latest main, the kernels have since moved to csrc/libtorch_stable/. Could you take another look? Thanks!

[yewentao256]

LGTM, thanks for the work!

[liulanze]

@yewentao256 CI checks failing on this PR are unrelated to the change.
kernels-fusedmoe-layer-test-2-b200s died on NVSHMEM/IBRC connection setup, and both MI300 multimodal jobs hit a simultaneous "GPU Hang" hardware event on node-2.
Could you please help to proceed with merge? Thanks!

[AndreasKaratzas]

I don't see why this PR would create a problem but given that the proposed changes are CU level and that one of the tests is failing with a seg fault, I would advise to investigate this a bit further. We are also transposing to MI325 cluster that should be more stable.

cc @yewentao256

[yewentao256]

CC @liulanze for the seg fault

[AndreasKaratzas]

Also we just migrated to different cluster, so probably rebase too so that it runs on mi325 instead of mi300.

[AndreasKaratzas]

The multi modal test on AMD is not a regression of this PR, this PR can safely be force merged, I am going to take a look into the AMD failure in a different PR.

[liulanze]

@AndreasKaratzas Thanks for the review!

[liulanze]

Friendly ping, this PR is approved by @yewentao256 and @AndreasKaratzas, and lm_eval on TinyLlama-1.1B/H100 shows identical scores. The remaining CI failures are unrelated to this change.

@mgoin @WoosukKwon - would you be able to force-merge when you get a chance? Happy to address anything else if needed. Thanks!

[AndreasKaratzas]

I posted the PR in slack to get it force merged.

[mgoin]

I'm not sure why accuracy on tiny llama is taken as a proxy here for deepseekv4, as that is the model that changed the behavior.. maybe @zyongye or @ivanium has more context possibly on the change

[zyongye]

Only check with python IR implementation is not ideal as it is written follow the CUDA implementation. And ideally all non Tensor core operation shall be operate on FP32 to maximize precision. If the score has no regression, then there's no point change the behavior? And I think we need to change the IR side instead of CUDA impl.