执行摘要

• 一句话：修复 mHC block-M prenorm GEMM 跨行归约累加器污染
• 推荐动作：建议尽快合入 main 分支。该 PR 修复了一个严重的计算错误，改动量小且验证充分。值得关注的要点是 tilelang 中 T.alloc_var 在循环内的 hoist 行为需要开发者谨慎使用，建议在相关文档或代码注释中提醒。

功能与动机

hc_prenorm_gemm_block_m_tilelang 内核在处理 block-M 时，由于 T.alloc_var(init=0.0) 声明在 unrolled i_m 循环内部，但编译器将其提升至循环外仅初始化一次，导致跨行归约累加器 carry-over。这造成每个 block 的奇数行输出被错误地累加了偶数行的值。该 bug 在 DeepSeek-V4 mHC prefill 中产生 corrupted mixes，影响 num_tokens >= 1024 的场景。PR 描述指出“every odd output row (and its sqrsum) is the sum of both rows in its block”。

实现拆解

1. 识别根因：在 hc_prenorm_gemm_block_m_tilelang 内核的 warp_id == 0 分支中，原先在每个 i_m 迭代中通过 T.alloc_var 声明归约累加器，但 tilelang 编译器将其 hoist 到循环外仅初始化一次。
2. 修复方案：提前在 warp_id == 0 分支开始处分配 block_m 大小的局部数组 reduced_acc 和 reduced_sqr，并使用 T.clear 清零，然后在 i_m 循环中通过数组索引 reduced_acc[i_m] 和 reduced_sqr[i_m] 进行累加，确保每个 token 的归约结果独立。
3. 输出修正：将原本直接赋值的 out[0, token_idx, out_idx] = reduced_acc 改为 out[0, token_idx, out_idx] = reduced_acc[i_m]，sqrsum[0, token_idx] = reduced_sqr 改为 sqrsum[0, token_idx] = reduced_sqr[i_m]，输出正确的行独立结果。
4. 测试验证：PR 提及已在 B200 上验证，修复前奇数行输出等于 ref(even)+ref(odd) 且 sqrsum 为 2 倍，修复后所有 tests/kernels/test_mhc_kernels.py 用例通过，包括之前失败的 num_tokens=2048 参数化测试。该变更仅修改 vllm/model_executor/kernels/mhc/tilelang_kernels.py 一个文件，+8/-6 行。

关键文件：

• vllm/model_executor/kernels/mhc/tilelang_kernels.py（模块 内核；类别 source；类型 data-contract；符号 hc_prenorm_gemm_block_m_tilelang）: 包含所有变更，修复 hc_prenorm_gemm_block_m_tilelang 内核中跨行归约累加器 carry-over 的严重 bug。

关键符号：hc_prenorm_gemm_block_m_tilelang

关键源码片段

vllm/model_executor/kernels/mhc/tilelang_kernels.py

包含所有变更，修复 hc_prenorm_gemm_block_m_tilelang 内核中跨行归约累加器 carry-over 的严重 bug。

# 修复前：T.alloc_var 在 unrolled i_m 循环内被 hoist 导致跨行 carry-over
# 修复后：提前分配 block_m 大小的数组并清零，确保每行独立累加
if warp_id == 0:
    # 分配每行独立的归约累加器数组
    reduced_acc = T.alloc_local((block_m,), T.float32)
    reduced_sqr = T.alloc_local((block_m,), T.float32)
    T.clear(reduced_acc)
    T.clear(reduced_sqr)
    for i_m in T.unroll(block_m):
        token_idx = i_mt * block_m + i_m
        if token_idx < num_tokens:
            if lane < tile_n:
                # 使用数组索引累加，避免跨行污染
                for i_w in T.unroll(num_warps):
                    reduced_acc[i_m] += warp_acc[i_w, i_m, lane]
                out_idx = i_t * tile_n + lane
                if out_idx < n_out:
                    out[0, token_idx, out_idx] = reduced_acc[i_m]
            if lane == 0 and i_t == 0:
                for i_w in T.unroll(num_warps):
                    reduced_sqr[i_m] += warp_sqr[i_w, i_m]
                sqrsum[0, token_idx] = reduced_sqr[i_m]

评论区精华

该 PR 获得 tjtanaa 的批准，且在 MI355X 上验证了修复前后单元测试均通过（因为该 bug 仅影响 block-M 路径，而 MI355X 可能不使用 block-M 内核或测试用例覆盖不全）。claude[bot] 自动评论表示来自 fork 的 PR 无法自动审查。没有其他 review 讨论。

• 暂无高价值评论线程

风险与影响

• 风险：
  1. 回归风险低：修复仅针对 block-M prenorm GEMM 内核的特定 bug，改动范围小（+8/-6 行），且已在 B200 上验证所有测试用例通过。
  2. 性能影响极小：将单变量替换为 block_m 大小的数组，额外内存开销很小，且仅在 warp_id == 0 分支内，对整体性能影响可忽略。
  3. 未覆盖的测试：虽然所有测试用例通过，但测试可能未覆盖所有 block-M 大小和 token 数量的组合，存在边缘情况未被验证。
• 影响：
  1. 用户影响：修复了 DeepSeek-V4 在 mHC prefill 路径中 num_tokens >= 1024 时的推理错误，确保奇数 token 的输出正确，对使用 DeepSeek-V4 的用户是重要修复。
  2. 系统影响：仅修改一个内核文件，对 vLLM 其他组件无影响。
  3. 团队影响：小范围、低风险的 bugfix，易于 review 和合并。
     - 风险标记：核心路径变更, 依赖编译器行为（hoist）

关联脉络

• 暂无明显关联 PR