Skip to content

vllm.model_executor.layers.fused_moe.router.norm_gate_linear

Fused RMSNorm + GateLinear for DeepSeek V4 MoE routing.

NormGateLinear

Bases: Module

RMSNorm + GateLinear, fused on DSV4-Pro only.

Source code in vllm/model_executor/layers/fused_moe/router/norm_gate_linear.py 
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
@PluggableLayer.register("norm_gated_linear")
class NormGateLinear(nn.Module):
    """RMSNorm + GateLinear, fused on DSV4-Pro only."""

    def __init__(
        self,
        hidden_size: int,
        num_experts: int,
        rms_eps: float = 1e-6,
        params_dtype: torch.dtype | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        self.num_experts = num_experts
        self.rms_eps = rms_eps

        self.norm = RMSNorm(hidden_size, eps=rms_eps, dtype=params_dtype)
        self.gate = GateLinear(
            hidden_size,
            num_experts,
            bias=False,
            out_dtype=torch.float32,  # DSV4 router output is fp32
            params_dtype=params_dtype,
            prefix=f"{prefix}.gate" if prefix else "gate",
        )

        self.e_score_correction_bias = None
        self.tid2eid = None

        self._fused_kernel_supported = (
            hidden_size == DSV4_PRO_HIDDEN_SIZE
            and num_experts == DSV4_PRO_NUM_EXPERTS
            and self.gate.allow_dsv3_router_gemm  # cuda platform
        )

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        if self._fused_kernel_supported:
            assert x.shape[1] == DSV4_PRO_HIDDEN_SIZE
            assert self.gate.weight.shape == (
                DSV4_PRO_NUM_EXPERTS,
                DSV4_PRO_HIDDEN_SIZE,
            )
            # This must be wrapped in a custom op because our torch.compile integration
            # does not support runtime dispatching on num_tokens.
            return torch.ops.vllm.dsv4_pro_norm_gate(
                x, self.norm.weight, self.gate.weight, self.rms_eps
            )

        # Non-Pro fallback (e.g. DSV4-Flash with hidden_size=4096):

        normed_x = self.norm(x)
        logits, _ = self.gate(normed_x)
        return normed_x, logits

_dsv4_pro_norm_gate 

_dsv4_pro_norm_gate(
    x: Tensor,
    norm_weight: Tensor,
    gate_weight: Tensor,
    rms_eps: float,
) -> tuple[Tensor, Tensor]

Runtime dispatcher: fused dsv4_norm_router_gemm (M<=16) vs the unfused rms_norm + dsv3_router_gemm fallback (M>16).

Source code in vllm/model_executor/layers/fused_moe/router/norm_gate_linear.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
def _dsv4_pro_norm_gate(
    x: torch.Tensor,
    norm_weight: torch.Tensor,
    gate_weight: torch.Tensor,
    rms_eps: float,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Runtime dispatcher: fused ``dsv4_norm_router_gemm`` (M<=16) vs the
    unfused ``rms_norm + dsv3_router_gemm`` fallback (M>16).

    """
    if x.shape[0] <= DSV4_PRO_MAX_NUM_TOKENS:
        return ops.dsv4_norm_router_gemm(x, norm_weight, gate_weight, rms_eps)

    normed = torch.empty_like(x)
    # Call `_C::rms_norm` here to avoid select the path of native rms
    torch.ops._C.rms_norm(normed, x, norm_weight, rms_eps)
    logits = torch.mm(normed, gate_weight.t(), out_dtype=torch.float32)
    return normed, logits