# PR #48391 完整报告

- 仓库：`vllm-project/vllm`
- 标题：[Bugfix][Kernel] Fix batch invariance in RMSNorm kernels by pinning block size
- 合并时间：2026-07-28 22:24
- 原文链接：http://prhub.com.cn/vllm-project/vllm/pull/48391

---

# 执行摘要

- 一句话：修复 RMSNorm kernel 的 batch invariance 漏洞：固定 block size
- 推荐动作：值得精读。虽然改动量仅 190 行，但 PR 展示了如何通过深思熟虑的测试参数化暴露隐藏的浮点并行问题，并系统性地修复同类 kernel。对于关注推理确定性的团队是重要参考。

# 功能与动机

此 PR 是 #27433（batch invariance 追踪 issue）和 #40413（将残差路径路由到 fused_add_rms_norm）的 follow-up。#40413 的测试仅比较了 num_tokens=1 和 4（均小于 256），未发现 block size 分歧。作者发现 `fused_add_rms_norm` kernel 的 block size 选择依赖 `num_tokens`，但 `batch_invariant_launch` 标志仅禁用向量化，未固定 block size，导致 batch 间的不一致。

# 实现拆解

1. **CUDA kernel 逻辑修复（5 个函数）**: 在 `csrc/libtorch_stable/layernorm_kernels.cu`、`layernorm_quant_kernels.cu` 和 `fused_layernorm_dynamic_per_token_quant.cu` 中，分别修改 `rms_norm`、`fused_add_rms_norm`、`rms_norm_static_fp8_quant`、`fused_add_rms_norm_static_fp8_quant`、`rms_norm_per_block_quant_dispatch` 的 block size 计算：若 `vllm_is_batch_invariant()` 为真则固定为 1024（per_block_quant 为 512），否则保持原动态逻辑。同时移除了原本在 vectorization 分支前重复的 `batch_invariant_launch` 检查。
2. **单元测试大幅增强**: `tests/v1/determinism/test_rms_norm_batch_invariant.py` 中，为已有测试增加 `n_extra=299`（使总 token 数跨过 256 阈值）和种子遍历（16 个种子）。新增 `_assert_rows_bit_identical` 辅助函数和 4 个 regression test，分别覆盖非残差 `rms_norm`、`rms_norm_static_fp8_quant`、`fused_add_rms_norm_static_fp8_quant`、`rms_norm_per_block_quant_dispatch`。这些测试通过比较小 batch（255 tokens）和大 batch（300 tokens）的前 255 行是否 bit-identical 来验证 batch invariance。
3. **端到端确定性测试增强**: `tests/v1/determinism/test_batch_invariance.py` 中的 `test_v1_generation_is_deterministic_across_batch_sizes_with_needle` 新增 `rms_norm_impl` 参数化（`default` 和 `vllm_c`），当为 `vllm_c` 时通过 `kernel_config` 强制走 C++ 实现，以覆盖 block size 依赖路径。
4. **CI 配置调整**: `.buildkite/test_areas/misc.yaml` 中增加 A100/H100/B200 上 batch-invariance 测试的超时时间（40→60、35→45 分钟），并调整 `-k` 过滤器以适配新参数化后的节点名变更。

关键文件：
- `csrc/libtorch_stable/layernorm_kernels.cu`（模块 内核层；类别 source；类型 core-logic；符号 rms_norm, fused_add_rms_norm）: 核心修复文件：修改 `rms_norm` 和 `fused_add_rms_norm` 函数中的 block size 选择逻辑，添加 batch-invariant 检查以固定 block size 为 1024。
- `csrc/libtorch_stable/layernorm_quant_kernels.cu`（模块 内核层；类别 source；类型 core-logic；符号 rms_norm_static_fp8_quant, fused_add_rms_norm_static_fp8_quant）: 修复 `rms_norm_static_fp8_quant` 和 `fused_add_rms_norm_static_fp8_quant` 的 block size 选择，与 layernorm_kernels.cu 同理。
- `csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu`（模块 内核层；类别 source；类型 core-logic；符号 rms_norm_per_block_quant_dispatch）: 修复 `rms_norm_per_block_quant_dispatch` 的 block size 选择，per_block_quant 使用不同的默认值（512/256），batch-invariant 下固定为 512。
- `tests/v1/determinism/test_rms_norm_batch_invariant.py`（模块 测试；类别 test；类型 test-coverage；符号 _assert_rows_bit_identical, test_rms_norm_batch_invariant_nonresidual_kernel, test_rms_norm_static_fp8_quant_batch_invariant, test_fused_add_rms_norm_static_fp8_quant_batch_invariant）: 大幅增强的单元测试：增加跨 256 阈值的参数和种子遍历，新增 4 个针对不同 kernel 的 batch invariance regression test，使用 `_assert_rows_bit_identical` 验证 bit-exact 一致性。
- `tests/v1/determinism/test_batch_invariance.py`（模块 测试；类别 test；类型 test-coverage；符号 test_v1_generation_is_deterministic_across_batch_sizes_with_needle）: 端到端确定性测试增强：增加 `rms_norm_impl` 参数化以强制走 C++ kernel 路径，覆盖 block size 依赖。
- `.buildkite/test_areas/misc.yaml`（模块 CI 配置；类别 config；类型 configuration）: CI 配置调整：增加 batch-invariance 测试超时时间，并更新 `-k` 过滤器以适配新参数化。

关键符号：rms_norm, fused_add_rms_norm, rms_norm_static_fp8_quant, fused_add_rms_norm_static_fp8_quant, rms_norm_per_block_quant_dispatch, _assert_rows_bit_identical, test_v1_generation_is_deterministic_across_batch_sizes_with_needle

## 关键源码片段

### `csrc/libtorch_stable/layernorm_kernels.cu`

核心修复文件：修改 `rms_norm` 和 `fused_add_rms_norm` 函数中的 block size 选择逻辑，添加 batch-invariant 检查以固定 block size 为 1024。

```cpp
// csrc/libtorch_stable/layernorm_kernels.cu

void rms_norm(torch::stable::Tensor& out, /* ... */) {
    // ...
    // 在 batch-invariant 模式下，必须固定 block size 1024
    // 否则同一 token 在不同 batch 大小下会因 block 大小不同而产生不同的 fp32 求和顺序
    const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
    const int max_block_size =
        batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);
    // ...
}

void fused_add_rms_norm(torch::stable::Tensor& input, /* ... */) {
    // ...
    const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
    const int max_block_size =
        batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);
    // 移除了 vectorization 前重复的 batch_invariant_launch 检查
    // ...
}

```

### `tests/v1/determinism/test_rms_norm_batch_invariant.py`

大幅增强的单元测试：增加跨 256 阈值的参数和种子遍历，新增 4 个针对不同 kernel 的 batch invariance regression test，使用 `_assert_rows_bit_identical` 验证 bit-exact 一致性。

```python
# tests/v1/determinism/test_rms_norm_batch_invariant.py

# 关键辅助函数：验证小 batch 与大 batch 的前 255 行 bit-identical
# 255 是保证两个 launch 都使用相同 block size 的最大 token 数
# （300 tokens 会触发 block size 切换，而 255 不会）
def _assert_rows_bit_identical(small, large, msg):
    """Assert that the first 255 rows of small and large are bit-identical."""
    n = min(small.shape[0], 255, large.shape[0])
    torch.testing.assert_close(
        small[:n], large[:n], rtol=0.0, atol=0.0, msg=msg
    )

# 新增非残差路径测试
def test_rms_norm_batch_invariant_nonresidual_kernel(...):
    # 构造 small (255 tokens) 和 large (300 tokens)，调用 rms_norm
    # 通过 _assert_rows_bit_identical 验证前 255 行一致
    ...

# 类似地，增加 static_fp8_quant、fused_add_rms_norm_static_fp8_quant、per_block_quant 的测试

```

# 评论区精华

- **复现问题讨论**: yewentao256 在 main 上无法复现，作者解释 main 的测试只用 batch_size<256 所以不会触发 bug。作者展示在 n_extra=299 下稳定复现。
- **扩展到其他 kernel**: yewentao256 要求同时修复 `rms_norm_static_fp8_quant` 等类似 kernel，作者扩展到 5 个函数，获得认可。
- **性能影响**: 作者进行 benchmark（Llama-3.1-8B-Instruct-FP8），显示延迟变化约 -0.009%，在噪声范围内，无显著回归。
- **为何 e2e acc 不坏**: 作者推测多数用户不使用 `VLLM_BATCH_INVARIANT`，且即使使用，小 batch 也可能不触发 block size 切换。
- **参数调整**: reviewer 建议 num_trials 从 10 改为 5，作者采纳。

- 无法复现 bug (question): 作者解释 main 的测试只使用 batch_size<256，不触发 block size 切换。展示在 n_extra=299 时稳定复现。
- 扩展到其他类似 kernel (design): 作者将修复扩展到 5 个函数（rms_norm, fused_add_rms_norm, rms_norm_static_fp8_quant, fused_add_rms_norm_static_fp8_quant, rms_norm_per_block_quant_dispatch）。
- 性能影响评估 (performance): 作者 benchmark Llama-3.1-8B-Instruct-FP8，结果显示延迟变化约 -0.009%，在噪声范围内，无显著回归。
- 测试参数调整 (testing): 作者采纳建议并修正。

# 风险与影响

- 风险：
 - **核心路径变更**: 修改了 RMSNorm 系列 5 个 CUDA kernel 的 block size 逻辑，影响所有使用这些 kernel 的模型在 batch-invariant 模式下的行为。固定到较大 block size（1024）可能减少 SM 并发度，但 benchmark 显示无显著性能影响。此变更仅影响 `VLLM_BATCH_INVARIANT=1` 模式，默认模式不受影响。
 - **测试覆盖**: 新增的单元测试覆盖了各个 kernel 的 batch invariance，但可能仍有其他类似 kernel（如 l2_norm）未包含。CI 超时增加表明测试耗时增加。
 - **兼容性**: 行为变化仅限于 batch-invariant 模式。用户若依赖之前的不一致行为（不推荐）会观察到变化。
- 影响：
 - **用户**: 启用 `VLLM_BATCH_INVARIANT=1` 的用户将获得 bit-exact 一致的推理输出，无论 batch 大小如何。对于需要确定性结果（如测试、调试、可重复实验）的用户至关重要。
 - **系统**: 性能无显著变化，GPU 内存使用不变。CI 测试耗时略有增加。
 - **团队**: 维护了 batch invariance 承诺的正确性，减少相关 issue。新增的测试为未来重构提供了安全网。
 - 风险标记：核心路径变更 , 潜在性能影响（已验证无）

# 关联脉络

- PR #40413 残差路径路由到 fused_add_rms_norm（假设已 batch-invariant）: 本 PR 是该 PR 的 follow-up，修复了其未发现的 block size 不一致问题。
- PR #50060 CI 修复，本 PR 依赖该修复才能通过 CI: PR 评论指出 CI 在 #50060 合并前无法通过。
- PR #27433 Batch invariance 追踪 issue: 本 PR 是该 issue 的一部分，解决了其中记录的一个具体 bug。