执行摘要

本PR修复了由PR #40860（DeepSeek V4支持）引入的RMSNorm CUDA内核回归问题。该回归导致内核在最终乘法时始终将权重转换为float32，违反Python参考规范。修复方法是在 layernorm_kernels.cu 和 layernorm_quant_kernels.cu 中恢复权重原生dtype乘法，并同步更新FP8量化路径以保持一致性。经lm_eval验证，端到端精度无损。

功能与动机

为什么做：PR #40860 合并后，RMSNorm内核将所有权重转换为float32进行乘法，与 vllm/ir/ops/layernorm.py 中的 x = x.to(weight.dtype) * weight 规范不符。该问题导致每层约0.03的误差，在多层（如DeepSeek V4的Q/K norm层）中累积至输出完全错误（Issue #42325）。

要解决的问题：确保RMSNorm CUDA内核在权重为BF16/FP16时，乘法在对应精度下进行，而非一律使用FP32。

实现拆解

1. 修复 rms_norm_kernel 及其fused变体：在 csrc/libtorch_stable/layernorm_kernels.cu 的三处位置，将 x * s_variance * w（w已转为float）改为 static_cast<scalar_t>(x * s_variance) * weight_elem，即先缩放到标量类型（如bfloat16），再乘以原生dtype的权重。

2. 修复 rms_norm_static_fp8_quant_kernel 及其fused变体：在 csrc/libtorch_stable/layernorm_quant_kernels.cu 的三处对应位置进行相同修改，以保持fused路径与非fused复合路径的数值一致性（否则 test_fused_rms_norm_quant 会失败）。

3. 验证：

   • 通过 tests/kernels/core/test_layernorm.py（865 pass）和 tests/kernels/ir/test_layernorm.py（1442 pass, 361 skipped）。
   • 在TinyLlama上运行lm_eval，各指标与v0.21.0基线一致。

关键源码片段

csrc/libtorch_stable/layernorm_kernels.cu

核心修复文件，包含rms_norm_kernel和fused_add_rms_norm_kernel的三个实现路径，是本次bugfix的主体。

// csrc/libtorch_stable/layernorm_kernels.cu - 修复后的 rms_norm_kernel (vectorized 路径 )
#pragma unroll
for (int j = 0; j < VEC_SIZE; j++) {
    float x = static_cast<float>(src1.val[j]);
    // 之前 (regressed) 的代码：
    // float w = static_cast<float>(src2.val[j]);
    // dst.val[j] = static_cast<scalar_t>(x * s_variance * w);
    // 修复后：先缩放至 scalar_t（如 bfloat16），再乘以权重原生 dtype 值
    // 这样可以保证乘法结果精确匹配 Python 参考实现 `x.to(weight.dtype) * weight`
    dst.val[j] = static_cast<scalar_t>(x * s_variance) * src2.val[j];
}

评论区精华

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

作者补充了TinyLlama-1.1B上的lm_eval结果，显示各指标完全一致，随后获得批准。

@zyongye："...ideally all non Tensor core operation shall be operate on FP32 to maximize precision. And I think we need to change the IR side instead of CUDA impl."

作者以测试无回归为由坚持当前修复方案。讨论未改代码方向，但提出了关于精度与规范之间取舍的深刻问题。

@AndreasKaratzas："The multi modal test on AMD is not a regression of this PR, this PR can safely be force merged"

协助确认CI失败与PR无关，并推动合并。

风险与影响

• 回归风险：低。变更仅在标量化后改变了乘法精度顺序，对主流模型无退化，lm_eval验证一致。
• 兼容性风险：FP8量化内核同步修改维持了fused与非fused路径的一致性，现有测试全部通过。
  • 数值精度敏感领域（如训练后量化，后训练校准）可能需要关注，但推理路径无影响。
• 影响范围：所有使用BF16/FP16权重的RMSNorm模型（DeepSeek V4、LLaDA-2等）将恢复到v0.19.1的正确行为。其他模型无影响。

关联脉络

本PR是Issue #42325的修复，其诱因是PR #40860（DeepSeek V4支持）。该PR在DeepSeek V4支持中加入了将权重转为float32的修改，忽略了通用规范。本PR还原了v0.19.1的正确行为，并同步修复了量化内核的对应路径。

与历史PR #45478（多模态模型权重映射修复）、#45252（M-RoPE prompt_embeds DoS修复）等一起，反映了v0.20.0以来由大规模重构引入的数值回滚和稳定性修复趋势。团队应加强对核心内核变更的回归测试和精度基准。