Prhub

#44359 [MoE] Share apply_moe_activation support metadata

原始 PR 作者 mgoin 合并时间 2026-08-06 03:05 文件变更 13 提交数 2 评论 7 代码增减 +425 / -241

执行摘要

MoE 激活配置集中化,统一各后端契约

PR body 明确指出要“Centralize the activation capability and configuration contract for MoE backends that delegate to apply_moe_activation()”。此前每个后端都复制一份激活列表(如 _supports_activation 中的大段 activation in [...]),且 clamp/alpha/beta/activation_situ_beta 等参数以多个标量形式在函数签名中传递,新增激活类型需要同步修改所有后端。本 PR 通过单一配置对象和共享支持查询函数,降低维护成本并避免各后端行为漂移。

值得精读。该 PR 是 MoE 激活层的一次架构性收敛,展示了如何将分散在各后端的重复元数据与参数传递统一为单一配置对象。重点关注 activation.pyApplyMoEActivationConfig.from_configs() 的优先级解析、modular_kernel.py 中实例持有配置的时机,以及 Marlin 中“standalone 显式传参、实例回调自持配置”的双轨设计。对需要扩展 MoE 激活类型的开发者,本 PR 提供了清晰的扩展点。

讨论亮点

Review 中 bnellnm 提出三个关键疑问,mgoin 逐一解释:

  • CUTLASS FP8/W4A8 的支持集为何比共享实现更受限:bnellnm 问“Do you know why the supported set here is more limited even though it uses apply_moe_activation?” mgoin 回应这是有意的,因为 FP8/W4A8 硬编码 2*N GEMM 形状并显式拒绝非 gated 激活,属于后端自身的形状约束。

  • Triton 的 activation() 是否应该显式接收配置:bnellnm 建议把 config 传入。mgoin 解释基类 modular activation 已读取 self.activation_config,显式传递会破坏实例持有的配置结构,故不传。

  • Marlin 的 activation_with_lora 路径是否也需要 config:bnellnm 担心 activation_func 分支遗漏配置。mgoin 说明 activation_with_lora 内部调用 self.activation(),而后者使用实例的 self.activation_config;只有 standalone Marlin 需要显式传 config。

所有讨论最终达成一致,bnellnm 给出 APPROVED。

实现拆解

本 PR 按以下步骤完成 MoE 激活配置的集中化:

  1. activation.py 引入能力元数据与配置对象:新增 _APPLY_MOE_ACTIVATIONS frozenset 与 apply_moe_activation_supported() 查询函数;新增 @dataclass(frozen=True)ApplyMoEActivationConfig,包含 clamp_limitalphabetaactivation_situ_betaactivation_situ_linear_beta 五个字段,并提供 from_configs() 类方法从 FusedMoEConfigFusedMoEQuantConfig 解析(量化配置优先,模型配置兜底)。apply_moe_activation() 签名从多个标量参数改为单个 activation_config 对象,内部使用 _DEFAULT_APPLY_MOE_ACTIVATION_CONFIG 兜底。

  2. modular_kernel.pyFusedMoEExpertsModular 中解析并持有配置:构造函数中一次调用 ApplyMoEActivationConfig.from_configs(moe_config, quant_config) 存入 self.activation_config;基类的 activation() 方法不再接收 clamp_limit/alpha/beta/situ beta 等参数,统一改用 self.activation_config。所有模块化专家实例自动获得自己的配置,无需各自解析。

  3. 各后端删除重复支持列表,改用共享查询与实例配置triton_moe.pyfused_humming_moe.pymarlin_moe.py_supports_activation() 简化为 return apply_moe_activation_supported(activation)cutlass_moe.py 中 FP8/W4A8 保留 activation.is_gated and apply_moe_activation_supported(activation)(因硬编码 2*N GEMM 形状),FP4/MXFP4 直接 return apply_moe_activation_supported(activation)TritonExperts.activation()HummingExpertsBase.apply_activation()MarlinExperts.apply() 均改为读取 self.activation_config

  4. CUTLASS FP4/MXFP4 与 NVFP4 oracle 行为调整run_cutlass_moe_fp4/run_cutlass_moe_mxfp4 中,只有当 SILU 且无 clamp 时才走 fused SiLU+quantization 快捷路径;配置了 clamp 时改走通用的 apply_moe_activation() 路径。oracle/nvfp4.py 允许 vLLM CUTLASS 和 Humming 在配置 clamp 时被选中。

  5. Marlin standalone helper 与实例回调拆分_fused_marlin_moe/fused_marlin_moe/batched_fused_marlin_moe 删除 5 个标量参数,改为接收 activation_config: ApplyMoEActivationConfig | None;当 activation_func is None 时内部构造默认配置并调用共享 apply_moe_activation,否则仅传 topk_ids/expert_map 给实例回调(第二个 commit 修复了 E2E 中 activation_config 关键字泄漏到实例回调的回归)。

  6. 测试配套:新增/更新 tests/kernels/moe/test_cutlass_moe.pytest_moe.pytest_triton_moe_no_act_mul.py,覆盖元数据追踪共享实现、配置对象构造与所有权、NVFP4 clamp 选择器、Marlin/Humming 委托行为,并在 NVIDIA B300 上跑通 165+ 项测试。

文件 模块 状态 重要度
vllm/model_executor/layers/fused_moe/activation.py 激活层 modified 8.53
vllm/model_executor/layers/fused_moe/modular_kernel.py 模块内核 modified 6.82
vllm/model_executor/layers/fused_moe/experts/marlin_moe.py Marlin 后端 modified 7.49
vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py CUTLASS 后端 modified 7.23
vllm/model_executor/layers/fused_moe/experts/triton_moe.py Triton 后端 modified 7.17
vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py Humming 后端 modified 7.06
tests/kernels/moe/test_cutlass_moe.py MoE 测试 modified 7.28
tests/kernels/moe/test_moe.py MoE 测试 modified 6.84
tests/kernels/moe/test_triton_moe_no_act_mul.py MoE 测试 modified 6.68

关键符号

apply_moe_activation_supported ApplyMoEActivationConfig.from_configs apply_moe_activation FusedMoEExpertsModular.activation _fused_marlin_moe fused_marlin_moe batched_fused_marlin_moe run_cutlass_moe_fp8 run_cutlass_moe_fp4 run_cutlass_moe_mxfp4 run_cutlass_moe_w4a8_fp8 TritonExperts.activation HummingExpertsBase.apply_activation

关键源码片段

vllm/model_executor/layers/fused_moe/activation.py data-contract

核心变更文件:新增 `apply_moe_activation_supported()` 能力查询与不可变配置对象 `ApplyMoEActivationConfig`,并将 `apply_moe_activation()` 的签名从多个标量参数收敛为单一配置对象,是全 PR 的契约中心。

# vllm/model_executor/layers/fused_moe/activation.py
# 共享支持元数据:apply_moe_activation() 能处理的全部激活类型,单一事实来源。
_APPLY_MOE_ACTIVATIONS = frozenset(
    {
        MoEActivation.SILU,
        MoEActivation.GELU,
        MoEActivation.GELU_TANH,
        MoEActivation.SITU,
        MoEActivation.SWIGLUOAI,
        MoEActivation.SWIGLUOAI_UNINTERLEAVE,
        MoEActivation.SWIGLUSTEP,
        MoEActivation.SILU_NO_MUL,
        MoEActivation.GELU_NO_MUL,
        MoEActivation.GELU_TANH_NO_MUL,
        MoEActivation.RELU2_NO_MUL,
    }
)
​
​
def apply_moe_activation_supported(activation: MoEActivation) -> bool:
    """Whether ``apply_moe_activation`` supports an activation."""
    return activation in _APPLY_MOE_ACTIVATIONS
​
​
@dataclass(frozen=True)
class ApplyMoEActivationConfig:
    """Configuration forwarded to ``apply_moe_activation``."""
​
    clamp_limit: float | None = None
    alpha: float = 1.0
    beta: float = 0.0
    activation_situ_beta: float | None = None
    activation_situ_linear_beta: float | None = None
​
    @classmethod
    def from_configs(
        cls,
        moe_config: "FusedMoEConfig",
        quant_config: "FusedMoEQuantConfig",
    ) -> "ApplyMoEActivationConfig":
        """从模型配置与量化配置构建,量化配置优先级更高,缺失时回退到模型配置。"""
        clamp_limit = quant_config.gemm1_clamp_limit
        if clamp_limit is None:
            clamp_limit = moe_config.swiglu_limit
        alpha = quant_config.gemm1_alpha
        if alpha is None:
            alpha = moe_config.swiglu_alpha
        beta = quant_config.gemm1_beta
        if beta is None:
            beta = moe_config.swiglu_beta
        return cls(
            clamp_limit=clamp_limit,
            alpha=1.0 if alpha is None else alpha,
            beta=0.0 if beta is None else beta,
            activation_situ_beta=moe_config.activation_situ_beta,
            activation_situ_linear_beta=moe_config.activation_situ_linear_beta,
        )# 全局默认配置,避免每次调用重复构造。
_DEFAULT_APPLY_MOE_ACTIVATION_CONFIG = ApplyMoEActivationConfig()
​
​
def apply_moe_activation(
    activation: MoEActivation,
    output: torch.Tensor,
    input: torch.Tensor,
    *,
    activation_config: ApplyMoEActivationConfig | None = None,
    topk_ids: torch.Tensor | None = None,
    expert_map: torch.Tensor | None = None,
) -> torch.Tensor:
    """统一激活入口:配置对象驱动 clamp/alpha/beta,路由张量仍为每次调用传入。"""
    config = (
        _DEFAULT_APPLY_MOE_ACTIVATION_CONFIG
        if activation_config is None
        else activation_config
    )
    # ... 形状断言与各激活分支(SILU/SITU/SWIGLUOAI_UNINTERLEAVE 等)
    if activation == MoEActivation.SILU:
        if config.clamp_limit is not None:
            silu_and_mul_with_clamp(output, input, config.clamp_limit, topk_ids, expert_map)
        else:
            torch.ops._C.silu_and_mul(output, input)
    # ...
    return output
vllm/model_executor/layers/fused_moe/modular_kernel.py data-contract

基类 `FusedMoEExpertsModular` 在构造时统一解析 `activation_config`,并让 `activation()` 方法使用实例持有的配置,是所有模块化专家获取配置的入口。

# vllm/model_executor/layers/fused_moe/modular_kernel.py
class FusedMoEExpertsModular(...):
    def __init__(self, moe_config, quant_config, ...):
        self.moe_config = moe_config
        self.quant_config = quant_config
        # 每个专家实例在构造时只解析一次配置,后续所有激活调用共享同一不可变对象。
        self.activation_config = ApplyMoEActivationConfig.from_configs(
            moe_config, quant_config
        )
        # ...
​
    def activation(
        self,
        activation: MoEActivation,
        output: torch.Tensor,
        input: torch.Tensor,
        *,
        topk_ids: torch.Tensor | None = None,
        expert_map: torch.Tensor | None = None,
    ) -> None:
        # 不再接收 clamp_limit/alpha/beta 等标量,统一从实例配置读取,
        # 子类只需覆写此方法并在需要时读取 self.activation_config。
        apply_moe_activation(
            activation,
            output,
            input,
            activation_config=self.activation_config,
            topk_ids=topk_ids,
            expert_map=expert_map,
        )
vllm/model_executor/layers/fused_moe/experts/marlin_moe.py data-contract

Marlin 后端改动最复杂:standalone helper 改用 `activation_config` 对象,实例回调只传路由参数;第二个 commit 修复了实例回调的 `activation_config` 关键字泄漏回归,是本 PR 中风险最高的后端。

# vllm/model_executor/layers/fused_moe/experts/marlin_moe.py
# 该分支处理两种调用方式:
# 1. standalone 模式(activation_func=None):显式接收并传入 activation_config;
# 2. 实例回调模式(activation_func 为实例方法,如 self.activation):
# 只传路由张量,配置由实例内部的 self.activation_config 提供,
# 避免把配置对象泄漏给不期望它的回调。
activation_input = intermediate_cache1.view(-1, w13_num_shards * N)
if activation_func is None:
    config = (
        ApplyMoEActivationConfig()
        if activation_config is None
        else activation_config
    )
    apply_moe_activation(
        activation,
        intermediate_cache2,
        activation_input,
        activation_config=config,
        topk_ids=topk_ids,
        expert_map=expert_map,
    )
else:
    activation_func(
        activation,
        intermediate_cache2,
        activation_input,
        topk_ids=topk_ids,
        expert_map=expert_map,
    )

评论区精华

CUTLASS FP8/W4A8 支持集为何比共享实现更受限 设计

bnellnm 在 cutlass_moe.py 的 `_supports_activation` 处提问,既然这些类也使用 `apply_moe_activation`,为什么支持集合比共享函数小。

结论:mgoin 回应这是有意的:CUTLASS FP8/W4A8 硬编码 `2*N` GEMM 形状并显式拒绝非 gated 激活,属于后端自身形状约束,后续可扩展。 · 已解决

Triton activation 是否应该显式接收配置对象 设计

bnellnm 在 triton_moe.py 的 `activation()` 中提问,是否应把 config 作为参数传入。

结论:mgoin 解释基础 modular activation 已读取 `self.activation_config`,显式传递会破坏实例持有的配置结构,因此不传。 · 已解决

Marlin activation_with_lora 路径是否需要 config 正确性

bnellnm 担心 `_fused_marlin_moe` 中 `activation_func` 分支(例如 `activation_with_lora`)不会收到配置,导致 clamp 参数丢失。

结论:mgoin 说明 `activation_with_lora` 内部调用 `self.activation()`,而后者使用实例的 `self.activation_config`;只有 standalone Marlin 需要显式传 config。 · 已解决

风险与影响

主要风险集中在契约重构对多个后端的影响:

  • 回调签名不兼容风险activation_func 从接收 8 个标量参数改为只接收 routing-only 参数,外部自定义回调或子类若仍按旧签名传参将直接报错。PR 中已通过 E2E Marlin 暴露并修复了实例回调回归(activation_config 关键字意外传入),但其他第三方扩展点仍可能受影响。

  • CUTLASS FP4/MXFP4 clamped SiLU 路径变化:配置 clamp 后从 fused SiLU+量化快捷路径切换到通用 apply_moe_activation() + 独立量化,性能可能下降且数值结果可能有微小差异。测试验证了正确性但未做性能对比。

  • 元数据契约漂移apply_moe_activation_supported() 用 frozenset 硬编码支持集合,若未来在 apply_moe_activation() 中新增分支而忘记同步 frozenset,元数据可能不一致。测试 test_apply_moe_activation_supported_contract 部分缓解了该风险,但无法完全杜绝。

  • 平台覆盖不足:测试仅在 NVIDIA B300(CUDA)上运行,ROCm/XPU 等平台未覆盖;TritonExperts._supports_activation 的行为变化可能影响 ROCm 上的 Triton MoE。

影响范围集中在 vLLM 的 MoE 推理核心路径:

  • 对用户:MoE 模型的激活行为数学上保持一致,clamp/alpha/beta 等参数解析顺序(量化配置优先于模型配置)不变;唯一可见变化是 clamped SiLU 在 CUTLASS FP4/MXFP4 上可能走性能略低的通用路径,以及 NVFP4 选择器在 clamp 配置下会允许 CUTLASS/Humming 后端。

  • 对系统apply_moe_activation() 的公共签名发生破坏性变更,所有调用方(内部后端、可能的第三方定制层)需同步迁移到 ApplyMoEActivationConfig。未来新增激活类型时只需修改 activation.py 单点,显著降低维护成本。

  • 对团队:确立了“实例持有配置 + 共享能力元数据”的设计模式,为后续 MoE 激活扩展提供了清晰契约;测试矩阵覆盖了全部 11 种激活的执行与元数据一致性。

核心路径变更 多后端契约统一 回调签名变更风险 clamped SiLU 路径行为变化 缺少 ROCm/XPU 覆盖

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论