Skip to content

vllm.utils.flashinfer_moe_ep

FlashInfer moe_ep helpers for DeepSeek V4 vLLM integration.

Classes:

Functions:

FiMoeEpBackendSpec dataclass

Static properties of one flashinfer_moe_ep_* backend string.

The backend names the kernel family; the arch comes from the device and the weight handling from the checkpoint, so this only has to carry which megakernel to build and whether it needs NVSHMEM in the runtime set.

Source code in vllm/utils/flashinfer_moe_ep.py 
30
31
32
33
34
35
36
37
38
39
40
@dataclass(frozen=True)
class FiMoeEpBackendSpec:
    """Static properties of one ``flashinfer_moe_ep_*`` backend string.

    The backend names the kernel *family*; the arch comes from the device and
    the weight handling from the checkpoint, so this only has to carry which
    megakernel to build and whether it needs NVSHMEM in the runtime set.
    """

    megakernel: str
    needs_nvshmem: bool

_dequant_expert_weights_to_bf16(weight, scale)

[E, N, K//2] fp4 + [E, N, K//32] ue8m0 -> [E, N, K] bf16 (expert loop).

Source code in vllm/utils/flashinfer_moe_ep.py 
198
199
200
201
202
203
204
205
206
207
208
def _dequant_expert_weights_to_bf16(
    weight: torch.Tensor, scale: torch.Tensor
) -> torch.Tensor:
    """[E, N, K//2] fp4 + [E, N, K//32] ue8m0 -> [E, N, K] bf16 (expert loop)."""
    num_experts, n, k_half = weight.shape
    out = torch.empty(
        num_experts, n, k_half * 2, dtype=torch.bfloat16, device=weight.device
    )
    for e in range(num_experts):
        out[e] = _dequant_fp4_ue8m0_gran32(weight[e], scale[e])
    return out

_dequant_fp4_ue8m0_gran32(packed, sf_ue8m0)

[rows, K//2] packed e2m1 + [rows, K//32] ue8m0-uint8 scales -> bf16 [rows, K].

Source code in vllm/utils/flashinfer_moe_ep.py 
183
184
185
186
187
188
189
190
191
192
193
194
195
def _dequant_fp4_ue8m0_gran32(
    packed: torch.Tensor, sf_ue8m0: torch.Tensor
) -> torch.Tensor:
    """[rows, K//2] packed e2m1 + [rows, K//32] ue8m0-uint8 scales -> bf16 [rows, K]."""
    raw = packed.view(torch.uint8)
    lut = torch.tensor(_E2M1_LUT, dtype=torch.float32, device=raw.device)
    vals = torch.empty(
        raw.shape[0], raw.shape[1] * 2, dtype=torch.float32, device=raw.device
    )
    vals[:, ::2] = lut[(raw & 0x0F).to(torch.int64)]
    vals[:, 1::2] = lut[(raw >> 4).to(torch.int64)]
    sf = (sf_ue8m0.to(torch.int32) << 23).view(torch.float32)
    return (vals * sf.repeat_interleave(32, dim=-1)).to(torch.bfloat16)

ensure_fi_moe_ep_runtime(vllm_config)

Acquire the process-wide flashinfer moe_ep runtime once per worker.

Source code in vllm/utils/flashinfer_moe_ep.py 
135
136
137
138
139
140
141
142
143
144
145
146
147
148
def ensure_fi_moe_ep_runtime(vllm_config: VllmConfig) -> None:
    """Acquire the process-wide flashinfer moe_ep runtime once per worker."""
    global _FI_RUNTIME_HANDLE
    if _FI_RUNTIME_HANDLE is not None:
        return

    from flashinfer.moe_ep import bootstrap_moe_ep_runtime

    bootstrap = make_fi_moe_ep_bootstrap()
    spec = fi_moe_ep_backend_spec(vllm_config.kernel_config.moe_backend)
    _FI_RUNTIME_HANDLE = bootstrap_moe_ep_runtime(
        bootstrap,
        megakernel_runtime_requirements(spec),
    )

finalize_fi_moe_ep_runtime()

Release the process-wide flashinfer moe_ep runtime.

Source code in vllm/utils/flashinfer_moe_ep.py 
151
152
153
154
155
156
157
158
159
160
def finalize_fi_moe_ep_runtime() -> None:
    """Release the process-wide flashinfer moe_ep runtime."""
    global _FI_RUNTIME_HANDLE
    if _FI_RUNTIME_HANDLE is None:
        return

    from flashinfer.moe_ep import finalize_moe_ep_runtime

    finalize_moe_ep_runtime(_FI_RUNTIME_HANDLE)
    _FI_RUNTIME_HANDLE = None

validate_fi_moe_ep_config(vllm_config)

Config-time checks for the mega-MoE backends, native and flashinfer.

Source code in vllm/utils/flashinfer_moe_ep.py 
 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
def validate_fi_moe_ep_config(vllm_config: VllmConfig) -> None:
    """Config-time checks for the mega-MoE backends, native and flashinfer."""
    moe_backend = vllm_config.kernel_config.moe_backend
    if not is_fi_moe_ep_backend(moe_backend):
        return

    # flashinfer validates the arch too, but not until the layer constructor
    # runs during weight load; check here so the error names the flag the user
    # actually typed.
    capability = current_platform.get_device_capability()
    if capability is not None:
        cc = (capability.major, capability.minor)
        if cc not in FI_MOE_EP_SUPPORTED_CAPABILITIES:
            supported = ", ".join(
                f"{m}.{n}" for m, n in sorted(FI_MOE_EP_SUPPORTED_CAPABILITIES)
            )
            raise ValueError(
                f"moe_backend={moe_backend!r} is only supported on compute "
                f"capability {supported} (SM100/SM103), but this device is "
                f"{cc[0]}.{cc[1]}."
            )

    if vllm_config.parallel_config.enable_eplb:
        raise NotImplementedError(
            f"EPLB is not supported with moe_backend={moe_backend!r}: the "
            "flashinfer moe_ep experts neither apply the logical-to-physical "
            "expert map nor report per-expert load, so rebalancing would move "
            "weights without moving routing. Use "
            "moe_backend=deep_gemm_mega_moe to run the mega path with EPLB."
        )