执行摘要

• 一句话：为CPU/CUTLASS/WNA16 MoE后端添加GELU_TANH激活支持
• 推荐动作：值得精读，尤其是 WNA16 中从硬编码断言到动态传递 activation 的设计决策，以及 CPU C++ 中添加新激活的完整流程——枚举、解析、kernel 实现、分发函数。可作为后续扩展其他激活的参考模板。

功能与动机

核心GELU_TANH MoE激活已在主路径支持，但CPU fused MoE、CUTLASS FP8/FP4和WNA16后端仍拒绝该激活：CPU fused MoE在激活映射中缺少GELU_TANH条目会引发KeyError；CUTLASS未在_supports_activation中列出导致不必要的fallback；WNA16 MoE硬断言仅支持SiLU，遇到非SiLU模型会崩溃。

实现拆解

1. CPU C++ 核心：在 csrc/cpu/cpu_fused_moe.cpp 的 FusedMOEAct 枚举中添加 GeluTanhAndMul，并在 get_act_type 中增加 "gelu_tanh" 到该枚举的映射。实现向量化 gelu_tanh_and_mul 模板函数，使用 tanh 近似计算 GELU_TANH 激活。在 apply_gated_act 中增加对应的 case 分支。

2. CPU Python 层：在 vllm/model_executor/layers/fused_moe/cpu_fused_moe.py 的 _CPU_MOE_ACT_FN 字典中添加 MoEActivation.GELU_TANH 条目，使用 F.gelu(approximate="tanh") 实现。

3. CUTLASS 后端：在 vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py 的 _supports_activation 静态方法中（针对 FP8 和 FP4 两个类）添加 MoEActivation.GELU_TANH 和 MoEActivation.GELU_TANH_NO_MUL。

4. WNA16 量化层：在 vllm/model_executor/layers/quantization/moe_wna16.py 的 MoeWNA16Method.apply 中移除 assert layer.activation == MoEActivation.SILU 的硬限制，并在调用 fused_experts 时新增 activation=layer.activation 参数传递。

5. 测试配套：新增 tests/quantization/test_moe_wna16.py 测试 WNA16 方法能正确传递 GELU_TANH 激活；在 tests/kernels/moe/test_cutlass_moe.py 添加 test_cutlass_moe_supports_gelu_tanh_activation_metadata 验证 CUTLASS 后端的元数据支持；在 tests/kernels/moe/test_cpu_fused_moe.py 扩展 CPU 后端的测试覆盖。

关键文件：

• csrc/cpu/cpu_fused_moe.cpp（模块 CPU内核；类别 source；类型 core-logic；符号 FusedMOEAct, get_act_type, gelu_tanh_and_mul, apply_gated_act）: 核心 C++ 实现，新增 GELU_TANH 枚举、名称解析和向量化计算函数。
• tests/quantization/test_moe_wna16.py（模块 WNA16量化；类别 test；类型 test-coverage；符号 test_moe_wna16_apply_passes_layer_activation, fake_fused_experts）: 新增测试验证 WNA16 方法能正确传递 GELU_TANH 激活到 fused_experts。
• vllm/model_executor/layers/quantization/moe_wna16.py（模块 量化层；类别 source；类型 data-contract；符号 MoeWNA16Method.apply）: 移除 SiLU-only 断言，改为传递 layer.activation，使 WNA16 支持任何激活。
• vllm/model_executor/layers/fused_moe/cpu_fused_moe.py（模块 MoE层；类别 source；类型 data-contract；符号 _CPU_MOE_ACT_FN）: 在激活映射中添加 GELU_TANH 的 PyTorch 实现。
• vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py（模块 CUTLASS专家；类别 source；类型 data-contract；符号 _supports_activation）: 为 CUTLASS FP8 和 FP4 后端添加 GELU_TANH 及 GELU_TANH_NO_MUL 支持。
• tests/kernels/moe/test_cutlass_moe.py（模块 CUTLASS测试；类别 test；类型 test-coverage；符号 test_cutlass_moe_supports_gelu_tanh_activation_metadata）: 添加测试验证 CUTLASS 后端支持 GELU_TANH 激活元数据。
• tests/kernels/moe/test_cpu_fused_moe.py（模块 CPU测试；类别 test；类型 test-coverage）: 扩展 CPU fused MoE 的 GELU_TANH 测试覆盖。

关键符号：gelu_tanh_and_mul (C++ template), get_act_type (C++), apply_gated_act (C++), MoeWNA16Method.apply (Python), _CPU_MOE_ACT_FN (Python dict), _supports_activation (CUTLASS method), test_moe_wna16_apply_passes_layer_activation, test_cutlass_moe_supports_gelu_tanh_activation_metadata

关键源码片段

csrc/cpu/cpu_fused_moe.cpp

核心 C++ 实现，新增 GELU_TANH 枚举、名称解析和向量化计算函数。

// 扩展枚举，添加 GeluTanhAndMul
enum class FusedMOEAct {
  SiluAndMul,
  SwigluOAIAndMul,
  GeluAndMul,
  GeluTanhAndMul, // 新增 : GELU_TANH 激活
};
​
// 在 get_act_type 中添加 "gelu_tanh" 分支
FusedMOEAct get_act_type(const std::string& act) {
  if (act == "silu") return FusedMOEAct::SiluAndMul;
  if (act == "swigluoai") return FusedMOEAct::SwigluOAIAndMul;
  if (act == "gelu") return FusedMOEAct::GeluAndMul;
  if (act == "gelu_tanh") return FusedMOEAct::GeluTanhAndMul; // 新增分支
  TORCH_CHECK(false, "Invalid act type: " + act);
}
​
// apply_gated_act 中新增 case
FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
                                  float* __restrict__ input,
                                  scalar_t* __restrict__ output,
                                  const int32_t m, const int32_t n,
                                  const int32_t input_stride,
                                  const int32_t output_stride) {
  switch (act) {
    // ... 其他 case ...
    case FusedMOEAct::GeluTanhAndMul:
      gelu_tanh_and_mul(input, output, m, n, input_stride, output_stride);
      return;
    default:
      TORCH_CHECK(false, "Unsupported act type.");
  }
}

vllm/model_executor/layers/quantization/moe_wna16.py

移除 SiLU-only 断言，改为传递 layer.activation，使 WNA16 支持任何激活。

def apply(
    self,
    layer: RoutedExperts,
    x: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    shared_experts: SharedExperts | None,
    shared_experts_input: torch.Tensor | None,
) -> torch.Tensor:
    from vllm.model_executor.layers.fused_moe import fused_experts
​
    # 移除了 assert layer.activation == MoEActivation.SILU 的限制
    return fused_experts(
        x,
        layer.w13_qweight,
        layer.w2_qweight,
        topk_weights=topk_weights,
        topk_ids=topk_ids,
        activation=layer.activation, # 新增：传递层配置的激活类型
        apply_router_weight_on_input=layer.apply_router_weight_on_input,
        global_num_experts=layer.global_num_experts,
        expert_map=layer.expert_map,
        quant_config=self.moe_quant_config,
    )

vllm/model_executor/layers/fused_moe/cpu_fused_moe.py

在激活映射中添加 GELU_TANH 的 PyTorch 实现。

# 激活名称到原生前向函数的映射
_CPU_MOE_ACT_FN: dict[MoEActivation, Callable[[torch.Tensor], torch.Tensor]] = {
    MoEActivation.SILU: lambda x: SiluAndMul(compile_native=False).forward_native(x),
    MoEActivation.SWIGLUOAI: _swigluoai_forward_native,
    MoEActivation.GELU: _gelu_and_mul,
    MoEActivation.GELU_TANH: ( # 新增：GELU_TANH 支持
        lambda x: F.gelu(x[..., : x.shape[-1] // 2], approximate="tanh")
        * x[..., x.shape[-1] // 2 :]
    ),
}

评论区精华

• 测试平台兼容性：审查者 AndreasKaratzas 指出 WNA16 测试应添加 CUDA-only 跳过的 pytest mark，因为 WNA16 是 CUDA 特定的量化路径。作者 lesj0610 响应并添加了 @pytest.mark.skipif(not current_platform.is_cuda())。

  • WNA16 测试应添加 CUDA-only 跳过标记 (testing): 作者 lesj0610 采纳建议，在测试函数上添加了 CUDA-only 跳过装饰器。

风险与影响

• 风险：
  • WNA16 断言移除：移除 SiLU 硬断言后，若 fused_experts 不支持的激活被传入，运行时错误将替代早期报错，可能使故障定位更困难。但当前 fused_experts 已具备完整的激活路由能力。
  • C++ kernel 新实现：gelu_tanh_and_mul 使用标量 tanh 回退循环，在未充分测试的平台上可能有性能隐忧，但与现有 gelu_and_mul 模式一致。
  • CUTLASS 元数据扩展：FP8/FP4 路径新增激活元数据后，若底层 kernel 实际不支持，将静默选择 fallback 路径，不影响正确性但可能改变性能特征。
• 影响：
  • 用户：使用 GELU_TANH 激活的模型（如 Gemma 4）现在可以在 CPU、CUTLASS (FP8/FP4) 和 WNA16 量化后端上正常推理，不再报错或 crash。
  • 系统：无性能回退预期，测试覆盖率显著提升，统一了各后端的激活支持模式。
  • 团队：降低了维护多后端激活映射的认知负担，后续添加新激活时有可参考的范式。
  • 风险标记：多后端同步变更, 断言移除可能导致运行时错误, 新 C++ kernel 未在非 x86 平台验证

关联脉络

• PR #41050 [Kernel][MoE] Add GELU_TANH to MoE activations: 该 PR 为主路径（TRT-LLM NvFP4 等）添加了 GELU_TANH 支持，本 PR 是对剩余后端的补全，作者在 PR body 中明确提及两者关联。