执行摘要

• 一句话：修复 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。

// 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 一致性。

# 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。