Prhub

#40372 [Kernel] Batch invariant NVFP4 MoE using cutlass

原始 PR 作者 jzakrzew 合并时间 2026-08-05 12:31 文件变更 5 提交数 14 评论 5 代码增减 +316 / -102

执行摘要

固化 CUTLASS NVFP4 MoE 内核批不变性契约并补测试

PR body 明确说明:"Much like in case of nvfp4 linear, the current VLLM_CUTLASS MoE backend happens to be already batch invariant in practice. This PR adds a few comments and static asserts in the .cu file so that no one breaks it accidentally, along with a specific test case for batch invariance." 即目标是替代需要大改实现的 #39520,用轻量方式确认现有实现的批不变性属性,并通过编译期断言与测试把它固化为不可破坏的契约。

值得精读。对从事内核调度、量化后端或确定性保证的工程师尤其有参考价值,值得关注的三点设计决策:一是用 static_assert 把实践属性固化为编译期契约,配套错误信息引导开发者走专门路径;二是三层批不变性验证策略(full-M 逐行 vs M=1、整批行置换、单行独立执行);三是通过测试文件合并解决进程隔离约束的 CI 组织方式。另外可与 #39520 对比,理解"发现已支持 → 转固化契约"的 PR 替代决策。

讨论亮点

核心讨论有三处:

  1. yewentao256 在 CI 配置 review 中要求:"Please add the test to the linear test file and combine as test cutlass",即把新增 MoE 测试合并进 linear 批不变测试文件。作者后续通过 commit "Consolidate Cutlass batch-invariance tests" 响应,删除了独立的 test_nvfp4_batch_invariant_scaled_mm.py,将其逻辑并入 test_cutlass_batch_invariance.py

  2. gemini-code-assist 高优先级建议:PR body 声明 e=40 用例可能因 bug #40351 失败,建议用 pytest.mark.xfail 标记避免阻塞 CI。最终合入代码未见 xfail 标记(_NVFP4_MOE_NUM_EXPERTS 仍为 (40, 64)),推测 #40351 已修复或测试实际通过;PR body 的 note 未同步删除,存在轻微文档不一致。

  3. mergify[bot] 在合入前提示 merge conflict 与 pre-commit 失败,作者通过多轮 merge main 解决(14 个 commit 中 11 个为 merge),并由 yewentao256 协助重试 CI,作者在评论中请求"take another look"。

实现拆解

实现分为 4 步:

  1. 宣告能力契约vllm/model_executor/layers/fused_moe/experts/cutlass_moe.pyCutlassExpertsFp4 新增静态方法 _supports_batch_invariance(),返回 True,明确声明该后端支持批不变性,供上层框架与测试查询。

  2. 编译期固话内核约束csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu 中,run_fp4_blockwise_scaled_group_mm_sm100cute::is_same_v 断言 TileScheduler 必须为 PersistentTileSchedulerSm100Grouprun_fp4_blockwise_scaled_group_mm_sm120cute::is_base_of_v 断言主循环调度必须源自 KernelPtrArrayTmaWarpSpecializedCooperative。错误信息明确要求未来改动走独立 batch invariant 路径而非放宽检查。

  3. 补充批不变性测试tests/v1/determinism/test_cutlass_batch_invariance.py 新增 294 行。test_cutlass_nvfp4_scaled_mm_batch_invariant 覆盖 7 组形状 × 2 种 dtype,逐行将 full-M 结果与 M=1 结果做 torch.equal 严格对比;test_cutlass_nvfp4_moe_batch_invariant 构造 FusedMoEKernel + CutlassExpertsFp4 真实调用链,覆盖 e∈{40,64}、topk∈{1,4}、SILU/SWIGLUSTEP 与 3 组 batch 形状,验证行置换后输出不变以及单行独立执行与整批对应行一致。文件 docstring 要求以 VLLM_BATCH_INVARIANT=1 在 fresh pytest 进程中运行。

  4. 合并重复测试与清理 CI:删除 tests/v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py(其 test_nvfp4_gemm_batch_invariance 逻辑并入上述新文件),并在 .buildkite/test_areas/misc.yaml 的 batch-invariance-b200 任务中移除对已删除文件的引用,避免同一进程多次隔离执行的要求造成 CI 配置复杂。

文件 模块 状态 重要度
tests/v1/determinism/test_cutlass_batch_invariance.py 确定性测试 modified 7.52
csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu 量化内核 modified 4.55
vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py MoE 专家 modified 5.5
tests/v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py 确定性测试 removed 5.74
.buildkite/test_areas/misc.yaml CI 配置 modified 2.36

关键符号

_supports_batch_invariance test_cutlass_nvfp4_scaled_mm_batch_invariant test_cutlass_nvfp4_moe_batch_invariant _make_cutlass_fp4_moe_batch_invariant_case _run_cutlass_fp4_moe run_fp4_blockwise_scaled_group_mm_sm100 run_fp4_blockwise_scaled_group_mm_sm120

关键源码片段

tests/v1/determinism/test_cutlass_batch_invariance.py test-coverage

PR 主体:新增 294 行批不变性测试,含 NVFP4 scaled-MM 逐行对比与 MoE 行置换 / 单行验证,并明示 VLLM_BATCH_INVARIANT=1 与 fresh 进程要求。

@_NVFP4_REQUIRES_SM100
@pytest.mark.parametrize("case_config", _NVFP4_MOE_BATCH_INVARIANT_CASES)
@pytest.mark.parametrize("activation", [MoEActivation.SILU, MoEActivation.SWIGLUSTEP])
@pytest.mark.parametrize("e", _NVFP4_MOE_NUM_EXPERTS)
@pytest.mark.parametrize("topk", _NVFP4_MOE_TOPKS)
@torch.inference_mode()
def test_cutlass_nvfp4_moe_batch_invariant(
    case_config, activation, e, topk, dtype, default_vllm_config, workspace_init
) -> None:
    case = _make_cutlass_fp4_moe_batch_invariant_case(
        case_config, activation, e, topk, dtype
    )
    # 先跑一次完整 batch 得到 baseline,后续所有重放都对照它的行切片
    batch_output = _run_cutlass_fp4_moe(case, case["hidden_states"], case["score"])
​
    # 断言后端显式声明支持批不变性,防止契约失效
    assert CutlassExpertsFp4._supports_batch_invariance()
​
    # 对整批行做两种置换重跑:若结果与 baseline 按原顺序逐行完全一致,
    # 说明 grouped GEMM 输出不依赖 token 被打包进 expert 任务的顺序
    indices = torch.arange(case["hidden_states"].size(0), device=case["hidden_states"].device)
    for perm_name, perm in (
        ("reversed", torch.flip(indices, dims=(0,))),
        ("evens_then_odds", torch.cat((indices[::2], indices[1::2]))),
    ):
        permuted_output = _run_cutlass_fp4_moe(
            case, case["hidden_states"][perm], case["score"][perm]
        )
        torch.testing.assert_close(
            batch_output[perm],
            permuted_output,
            atol=0,
            rtol=0,
            msg=f"{case_id}: permutation '{perm_name}' changed outputs",
        )
    # 再对每一行以 batch-size-1 独立执行并逐行对比,进一步验证 M=1 一致性
csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu core-logic

批不变性的编译期保护机制:SM100/SM120 两处 static_assert 锁定调度器与主循环 schedule,防止未来改动无意破坏确定性。

// SM100 NVFP4 分组 GEMM:编译期锁定 persistent tile scheduler。
// 批不变性要求单个 token 行的结果不依赖 batch 内其他行的调度顺序,
// PersistentTileSchedulerSm100Group 按行分组持久调度,恰好满足该性质。
using Gemm1SM = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
using Gemm = Gemm1SM;static_assert(
    cute::is_same_v<
        typename Gemm::GemmKernel::TileScheduler,
        cutlass::gemm::kernel::detail::PersistentTileSchedulerSm100Group<
            ProblemShape, Gemm::GemmKernel::SchedulerPipelineStageCount>>,
    "SM100 NVFP4 grouped GEMM must use PersistentTileSchedulerSm100Group "
    "for batch invariance (VLLM_BATCH_INVARIANT=1). "
    "If you want to change this, add a dedicated config/code path "
    "for batch invariant mode, instead of relaxing this check.");// SM120 分支:要求主循环采用 cooperative ptr-array 调度,
// 保证 grouped GEMM 各 expert 分组的执行与 batch 行排列解耦。
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
static_assert(
    cute::is_base_of_v<
        cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperative,
        typename Gemm::GemmKernel::CollectiveMainloop::DispatchPolicy::Schedule>,
    "SM120 NVFP4 grouped GEMM must use a cooperative ptr-array mainloop "
    "schedule for batch invariance (VLLM_BATCH_INVARIANT=1). "
    "If you want to change this, add a dedicated config/code path "
    "for batch invariant mode, instead of relaxing this check.");// 今后若为性能更换 scheduler 或主循环 schedule,必须先为 batch invariant
// 模式提供独立代码路径,否则编译期直接报错,防止确定性被无意破坏
vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py data-contract

新增 _supports_batch_invariance() 契约方法,把批不变性从隐式实践变为显式能力声明,供测试与框架查询。

class CutlassExpertsFp4(mk.ExpertOp):
    # ... 其他能力声明省略 ...
​
    @staticmethod
    def _supports_batch_invariance() -> bool:
        # CUTLASS NVFP4 grouped GEMM 使用 persistent 调度且行间无依赖,
        # 输出天然与 batch 行排列无关;该声明使上层可以静态识别能力,
        # 并由 tests/v1/determinism/test_cutlass_batch_invariance.py 持续验证。
        return True
​
    # 若未来实现不再满足该性质,必须在此返回 False 并补充非批不变路径

评论区精华

测试文件合并进 test_cutlass_batch_invariance.py 设计

yewentao256 在 .buildkite/test_areas/misc.yaml 的 review 中要求:"Please add the test to the linear test file and combine as test cutlass",主张新增 MoE 测试并入现有 linear 批不变测试文件统一管理。

结论:作者通过 commit "Consolidate Cutlass batch-invariance tests" 将 test_nvfp4_batch_invariant_scaled_mm.py 删除并合入 test_cutlass_batch_invariance.py,同时移除 CI 中对应条目,避免同一进程多次隔离执行的要求。 · 已解决

e=40 用例是否需要 xfail 测试

gemini-code-assist 高优先级反馈:PR body 声明 e=40 测试可能因 bug #40351 失败,建议用 pytest.mark.xfail 标记避免阻塞 CI,并给出了具体 suggestion 代码。

结论:最终合入代码未包含 xfail 标记(_NVFP4_MOE_NUM_EXPERTS 仍为 (40, 64)),推测 #40351 已修复或测试实际通过;PR body 中的 note 未同步删除,存在轻微文档不一致。 · unresolved

风险与影响

主要风险有四类:

  1. 编译期断言限制内核演进.cu 中的 static_assert 硬编码了当前 scheduler 类型。未来若为性能引入新调度器(如 streaming scheduler 或非 cooperative 主循环),将直接编译失败。这是有意设计,但会抬高内核改动成本;同时若 CUTLASS 版本升级导致类型命名空间变化,可能误报失败。

  2. 测试强依赖进程隔离:native 代码在首个 GEMM 调用时缓存 VLLM_BATCH_INVARIANT 状态,一旦同一 pytest 进程先执行非批不变内核,后续该文件的测试结果不可信。文件 docstring 已警告,但 CI 中 batch-invariance-b200 任务仍在同一进程串行跑多个测试文件,顺序敏感。

  3. 覆盖不均衡:测试仅覆盖 bfloat16、SM100+(B200)、e∈{40,64}、topk∈{1,4};f16、SM120、更大专家数未覆盖。且 e=40 受 #40351 影响可能失败,存在 CI 波动风险。

  4. 删除文件连锁影响test_nvfp4_batch_invariant_scaled_mm.py 被删除,仓库内 CI 配置已同步清理;但若外部脚本或本地命令仍引用该路径会失效。

对开发者:修改 NVFP4 MoE CUTLASS 内核调度配置会立即在编译期被拦截,必须显式提供 batch invariant 专用路径,显著降低无意回归概率。对用户:使用 VLLM_CUTLASS + NVFP4(Blackwell)的用户获得更明确的确定性承诺,_supports_batch_invariance() 契约可被上层框架查询。对 CI:B200 批不变任务新增 294 行参数化测试(组合数达数百),增加运行时间;同时消除一个独立测试文件的进程隔离负担。对测试基建:确立了"批不变测试集中存放且必须 fresh 进程执行"的组织模式。整体影响面集中在 NVFP4/CUTLASS 后端与确定性测试体系,不影响其他模型路径。

编译期断言限制内核演进 测试强依赖进程隔离 e=40 已知缺陷影响 CI 覆盖仅限 SM100+ / bfloat16

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论