执行摘要

• 一句话：新增 SM100 BF16x3 路由器 GEMM 加速
• 推荐动作：推荐关注此 PR 的设计思路：利用数值分解在低精度硬件上模拟高精度计算，以及 CuteDSL 与 TMA 的使用方式。若团队涉及 SM100 MoE 优化，应仔细阅读 kernel 实现和精度分析。

功能与动机

当前 MoE 路由器权重常以 FP32 存储，而输入为 BF16，现有方案使用 FP32 GEMM 或 CuBLAS BF16x9 方案。本 PR 受 CuBLAS BF16x9 启发，针对路由器场景优化：权重分解为 3xBF16 并在此精度下计算，以利用 SM100 Tensor Core 的 BF16 吞吐，同时避免 FP32 GEMM 的性能开销。

实现拆解

1. 核心 kernel 实现（bf16x3_router_gemm_cutedsl.py）：定义 _decompose_fp32x2_to_3xbf16x2 内联 PTX 函数，将两个 FP32 值分解为 3 对 BF16 值；实现 Sm100BF16x3RouterGemm 类，包含 TMA 拷贝、主 GEMM kernel 及 split-K 规约 kernel。
2. 配置项添加（KernelConfig）：新增 enable_bf16x3_router_gemm 字段，默认 False。
3. CLI 参数（EngineArgs）：使用 kwargs 风格添加 --enable-bf16x3-router-gemm，并通过 create_engine_config 写入 KernelConfig。
4. 调度集成（GateLinear）：在多级 GEMM 调度链中插入 bf16x3_router_gemm 作为第 4 级（Tier 4），当启用且权重为 FP32 时优先于 cuBLAS 及 fallback。
5. 测试覆盖（test_bf16x3_router_gemm_cutedsl.py）：参数化测试多种形状，验证 kernel 输出与 FP64 参考的 MAE 低于 5e-6。

关键文件：

• vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py（模块 MoE 路由；类别 source；类型 core-logic；符号 _decompose_fp32x2_to_3xbf16x2, Sm100BF16x3RouterGemm, init, _make_tma）: 核心 kernel 实现，包括 FP32 到 3xBF16 的分解 PTX 和完整的 CuteDSL GEMM kernel 及 split-K 规约。
• vllm/model_executor/layers/fused_moe/router/gate_linear.py（模块 MoE 路由；类别 source；类型 core-logic）: MoE 路由器线性层的入口，在此添加了 bf16x3_router_gemm 的调度判断与调用。
• tests/kernels/test_bf16x3_router_gemm_cutedsl.py（模块 测试；类别 test；类型 test-coverage；符号 _requires_sm100_cutedsl, test_bf16x3_router_gemm_matches_reference）: 单元测试，验证 kernel 结果与 FP64 参考的 MAE 低于阈值。
• vllm/engine/arg_utils.py（模块 引擎参数；类别 source；类型 configuration）: 添加了 --enable-bf16x3-router-gemm CLI 参数和配置传递。
• vllm/config/kernel.py（模块 Kernel 配置；类别 source；类型 configuration）: 在 KernelConfig 中添加 enable_bf16x3_router_gemm 配置项。

关键符号：bf16x3_router_gemm, _decompose_fp32x2_to_3xbf16x2, Sm100BF16x3RouterGemm.init, Sm100BF16x3RouterGemm.kernel, Sm100BF16x3RouterGemm._splitk_reduce_kernel, Sm100BF16x3RouterGemm.compile, GateLinear.init, GateLinear.forward, test_bf16x3_router_gemm_matches_reference

关键源码片段

vllm/model_executor/layers/fused_moe/router/gate_linear.py

MoE 路由器线性层的入口，在此添加了 bf16x3_router_gemm 的调度判断与调用。

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, None]:
        # 多级调度链：
        # Tier 1-3: 专用 kernel (DSV3, fp32 等 )
        # Tier 4: experimental bf16x3 CuteDSL kernel
        # Tier 5: cuBLAS bf16×bf16→fp32
        # Tier 6: F.linear (fallback)
        if self.allow_bf16x3_router_gemm and x.dtype == torch.bfloat16:
            from vllm.model_executor.layers.fused_moe.router.bf16x3_router_gemm_cutedsl import ( # noqa: E501
                bf16x3_router_gemm,
            )
            output = bf16x3_router_gemm(x, self.weight)
            return output, None
        # ... 后续 tier

评论区精华

simon-veitner-redhat 提出了多个风格与设计建议：使用 cute.arch.warp_idx()/lane_idx() 代替手动计算、将常量作为 constexpr 传递给 kernel、将部分计算移到 host 端、为 PTX 分解添加注释等。gau-nernst 已采纳部分建议（注释、warp/lane ID 改用内置函数、移除 output_dim % 8 == 0 约束）。hmellor 要求 CLI 参数使用 kwargs 风格，gau-nernst 已更新。另一个关于与 PR#44343 性能对比的评论已被作者划掉。

• CLI 参数使用 kwargs 风格 (style): gau-nernst 已改为使用 **kernel_kwargs["enable_bf16x3_router_gemm"]。
• 使用 cute.arch 内置函数获取 warp/lane ID (design): gau-nernst 已修改为使用内置函数。
• PTX 分解添加注释 (documentation): gau-nernst 已在 _decompose_fp32x2_to_3xbf16x2 中添加注释。

风险与影响

• 风险：此 PR 默认关闭，对现有行为无影响。启用后仅对 SM100 GPU 生效，且仅当权重为 FP32 且 input_size % 8 == 0 时触发。主要风险包括：
  1) 新 kernel 的数值精度在特定形状下可能恶化（作者通过 2 个 MMA 累加器缓解）；
  2) CuteDSL 编译时间可能影响启动；
  3) 目前仅 Blackwell 架构可用，无 fallback 兼容。
  • 影响：对用户：提供可选的性能提升路径，需 SM100 GPU 并显式开启。对系统：新增 ~450 行 CuteDSL kernel 代码，增加编译与维护负担。对团队：需要维护 CuteDSL 和 PTX 内联汇编，对 CUDA 编译链依赖加深。
  • 风险标记：需 SM100 GPU, 默认关闭, 实验性功能, CuteDSL 编译依赖

关联脉络

• 暂无明显关联 PR