@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 8fe9317f2e

[Model][MoE] DeepSeek-V4: add opt-in FlashInfer moe_ep expert backend (#49636)

Signed-off-by: Md Saidul Hoque Anik <mhoqueanik@nvidia.com> Signed-off-by: Anerudhan Gopal <agopal@nvidia.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Anerudhan Gopal <agopal@nvidia.com>

Md Saidul Hoque Anik committed Aug 24, 2026 at 20:35 UTC 8fe9317f2e401aff6e13044098ac7f59e95dce97
5 files changed +1083 -12
tests/models/test_deepseek_v4_fi_moe_ep.py new
+259
@@ -0,0 +1,259 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for the flashinfer moe_ep backend plumbing.
4 +
5 +Everything here runs without a GPU or a flashinfer install: the flashinfer
6 +modules the helpers import lazily are replaced with capture fakes.
7 +"""
8 +
9 +import sys
10 +from dataclasses import dataclass, field
11 +from types import ModuleType, SimpleNamespace
12 +from typing import Any
13 +
14 +import pytest
15 +import torch
16 +
17 +from vllm.config.kernel import (
18 + FLASHINFER_MOE_EP_BACKENDS,
19 + MEGA_MOE_BACKENDS,
20 + validate_flashinfer_moe_ep_model,
21 +)
22 +from vllm.utils.flashinfer_moe_ep import (
23 + _E2M1_LUT,
24 + FI_MOE_EP_BACKEND_SPECS,
25 + _dequant_fp4_ue8m0_gran32,
26 + build_fi_mega_config,
27 + fi_moe_ep_backend_spec,
28 + make_fi_moe_ep_bootstrap,
29 + megakernel_runtime_requirements,
30 +)
31 +
32 +
33 +@dataclass
34 +class _FakeBootstrapConfig:
35 + world_size: int
36 + rank: int
37 + process_group: Any = None
38 + auto_bootstrap: bool = True
39 + device: int | None = field(default=None, kw_only=True)
40 +
41 +
42 +@dataclass
43 +class _FakeDeepGemmMegaMoeConfig:
44 + intermediate_size: int
45 + top_k: int
46 + activation_clamp: float | None
47 + fast_math: bool
48 +
49 +
50 +@dataclass
51 +class _FakeNvfp4CutedslMegaMoeConfig:
52 + intermediate_size: int
53 + top_k: int
54 + activation_clamp: float | None
55 + fast_math: bool
56 +
57 +
58 +@dataclass
59 +class _FakeMegaConfig:
60 + megakernel: Any
61 + preprocess_weights: bool
62 + quantize_input: bool
63 +
64 +
65 +@pytest.fixture
66 +def fake_flashinfer(monkeypatch):
67 + """Install a minimal fake flashinfer.moe_ep for the lazy imports."""
68 + moe_ep = ModuleType("flashinfer.moe_ep")
69 + core = ModuleType("flashinfer.moe_ep.core")
70 + runtime = ModuleType("flashinfer.moe_ep.core.runtime")
71 + flashinfer = ModuleType("flashinfer")
72 + fake_attrs: dict[ModuleType, dict[str, Any]] = {
73 + moe_ep: {
74 + "BootstrapConfig": _FakeBootstrapConfig,
75 + "DeepGemmMegaMoeConfig": _FakeDeepGemmMegaMoeConfig,
76 + "Nvfp4CutedslMegaMoeConfig": _FakeNvfp4CutedslMegaMoeConfig,
77 + "MegaConfig": _FakeMegaConfig,
78 + "core": core,
79 + },
80 + runtime: {"TORCH_DIST": "torch_dist", "NVSHMEM": "nvshmem"},
81 + flashinfer: {"moe_ep": moe_ep},
82 + core: {"runtime": runtime},
83 + }
84 + for mod, attrs in fake_attrs.items():
85 + for attr, value in attrs.items():
86 + setattr(mod, attr, value)
87 +
88 + for name, mod in {
89 + "flashinfer": flashinfer,
90 + "flashinfer.moe_ep": moe_ep,
91 + "flashinfer.moe_ep.core": core,
92 + "flashinfer.moe_ep.core.runtime": runtime,
93 + }.items():
94 + monkeypatch.setitem(sys.modules, name, mod)
95 + return moe_ep
96 +
97 +
98 +def test_fi_backend_strings_are_registered_mega_moe_backends():
99 + assert set(FI_MOE_EP_BACKEND_SPECS) == FLASHINFER_MOE_EP_BACKENDS
100 + assert FLASHINFER_MOE_EP_BACKENDS < MEGA_MOE_BACKENDS
101 +
102 +
103 +@pytest.mark.parametrize("moe_backend", sorted(FLASHINFER_MOE_EP_BACKENDS))
104 +def test_fi_moe_ep_backend_rejected_for_non_dsv4(moe_backend):
105 + """An FI moe_ep backend with a non-DSv4 model must fail at config time
106 + instead of silently falling through to the generic FusedMoE path."""
107 + with pytest.raises(ValueError, match="only supported for DeepSeek-V4"):
108 + validate_flashinfer_moe_ep_model(moe_backend, ["MixtralForCausalLM"])
109 +
110 +
111 +@pytest.mark.parametrize("moe_backend", sorted(FLASHINFER_MOE_EP_BACKENDS))
112 +def test_fi_moe_ep_backend_accepted_for_dsv4(moe_backend):
113 + validate_flashinfer_moe_ep_model(moe_backend, ["DeepseekV4ForCausalLM"])
114 +
115 +
116 +@pytest.mark.parametrize(
117 + "architectures",
118 + [["KimiK3ForConditionalGeneration"], ["MixtralForCausalLM"]],
119 +)
120 +def test_native_deep_gemm_mega_moe_not_arch_gated(architectures):
121 + """vLLM's own deep_gemm mega path is not DSv4-only (Kimi K3 uses it);
122 + models validate their own constraints at construction time."""
123 + validate_flashinfer_moe_ep_model("deep_gemm_mega_moe", architectures)
124 +
125 +
126 +def test_non_fi_backend_ignores_architectures():
127 + validate_flashinfer_moe_ep_model("auto", ["MixtralForCausalLM"])
128 +
129 +
130 +@pytest.mark.parametrize("moe_backend", sorted(MEGA_MOE_BACKENDS))
131 +def test_all_mega_backends_get_sequence_parallel_moe(moe_backend):
132 + """Every mega backend must qualify for sequence-parallel MoE at
133 + TP>1/EP: the predicate once matched only the native backend string,
134 + which silently ran the fi backends full-batch with an all-reduce on
135 + every rank — 0.42-0.65x native e2e at TP8."""
136 + from vllm.models.deepseek_v4.nvidia.model import _use_sequence_parallel
137 +
138 + vllm_config = SimpleNamespace(
139 + parallel_config=SimpleNamespace(
140 + pipeline_parallel_size=1,
141 + enable_expert_parallel=True,
142 + tensor_parallel_size=8,
143 + data_parallel_size=1,
144 + ),
145 + kernel_config=SimpleNamespace(moe_backend=moe_backend),
146 + )
147 + assert _use_sequence_parallel(vllm_config)
148 +
149 +
150 +def test_fi_moe_ep_backend_spec_kernel_and_nvshmem_contract():
151 + dg = fi_moe_ep_backend_spec("flashinfer_moe_ep_mega_deep_gemm")
152 + assert dg.megakernel == "deep_gemm_mega"
153 + assert not dg.needs_nvshmem
154 +
155 + cd = fi_moe_ep_backend_spec("flashinfer_moe_ep_mega_cutedsl")
156 + assert cd.megakernel == "nvfp4_cutedsl"
157 + assert cd.needs_nvshmem
158 +
159 + with pytest.raises(ValueError, match="not a flashinfer moe_ep backend"):
160 + fi_moe_ep_backend_spec("deep_gemm_mega_moe")
161 +
162 +
163 +def test_megakernel_runtime_requirements(fake_flashinfer):
164 + dg = megakernel_runtime_requirements(
165 + fi_moe_ep_backend_spec("flashinfer_moe_ep_mega_deep_gemm")
166 + )
167 + assert dg == frozenset({"torch_dist"})
168 +
169 + cd = megakernel_runtime_requirements(
170 + fi_moe_ep_backend_spec("flashinfer_moe_ep_mega_cutedsl")
171 + )
172 + assert cd == frozenset({"torch_dist", "nvshmem"})
173 +
174 +
175 +def test_bootstrap_pins_the_device_vllm_bound(fake_flashinfer, monkeypatch):
176 + """The runtime must not rederive the device from LOCAL_RANK/rank: under a
177 + remapped CUDA_VISIBLE_DEVICES that ordinal points at the wrong GPU
178 + (CUDA_ERROR_ILLEGAL_ADDRESS in the weight transforms). vLLM passes the
179 + device it already bound via BootstrapConfig.device."""
180 + import vllm.utils.flashinfer_moe_ep as mod
181 +
182 + pg = object()
183 + monkeypatch.setattr(
184 + mod,
185 + "get_ep_group",
186 + lambda: SimpleNamespace(world_size=4, rank_in_group=2, device_group=pg),
187 + )
188 + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 3)
189 +
190 + bootstrap = make_fi_moe_ep_bootstrap()
191 +
192 + assert bootstrap.world_size == 4
193 + assert bootstrap.rank == 2
194 + assert bootstrap.process_group is pg
195 + assert bootstrap.auto_bootstrap is False
196 + assert bootstrap.device == 3
197 +
198 +
199 +def test_build_fi_mega_config_selects_kernel_config(fake_flashinfer):
200 + dg = build_fi_mega_config(
201 + intermediate_size=2048,
202 + top_k=8,
203 + activation_clamp=7.0,
204 + megakernel="deep_gemm_mega",
205 + )
206 + assert isinstance(dg.megakernel, _FakeDeepGemmMegaMoeConfig)
207 + assert dg.megakernel.intermediate_size == 2048
208 + assert dg.megakernel.top_k == 8
209 + assert dg.megakernel.activation_clamp == 7.0
210 + assert dg.preprocess_weights and dg.quantize_input
211 +
212 + cd = build_fi_mega_config(
213 + intermediate_size=2048,
214 + top_k=8,
215 + activation_clamp=None,
216 + megakernel="nvfp4_cutedsl",
217 + )
218 + assert isinstance(cd.megakernel, _FakeNvfp4CutedslMegaMoeConfig)
219 +
220 + with pytest.raises(ValueError, match="Unsupported fi_moe_ep megakernel"):
221 + build_fi_mega_config(
222 + intermediate_size=2048,
223 + top_k=8,
224 + activation_clamp=None,
225 + megakernel="deep_gemm",
226 + )
227 +
228 +
229 +def test_ckpt_uses_nvfp4_experts_reads_moe_quant_algo():
230 + from vllm.models.deepseek_v4.nvidia.fi_moe import ckpt_uses_nvfp4_experts
231 +
232 + nvfp4 = SimpleNamespace(quant_config=SimpleNamespace(moe_quant_algo="NVFP4"))
233 + assert ckpt_uses_nvfp4_experts(nvfp4)
234 +
235 + mxfp4 = SimpleNamespace(quant_config=SimpleNamespace(moe_quant_algo=None))
236 + assert not ckpt_uses_nvfp4_experts(mxfp4)
237 +
238 + no_algo = SimpleNamespace(quant_config=SimpleNamespace())
239 + assert not ckpt_uses_nvfp4_experts(no_algo)
240 +
241 +
242 +def test_dequant_fp4_ue8m0_gran32_decodes_lut_and_scales():
243 + """One 32-element scale group per row: low nibble is the even element,
244 + high nibble the odd one, ue8m0 scale applies to the whole group."""
245 + packed = torch.arange(32, dtype=torch.uint8).reshape(2, 16)
246 + sf = torch.tensor([[127], [128]], dtype=torch.uint8) # 2**0, 2**1
247 +
248 + out = _dequant_fp4_ue8m0_gran32(packed, sf)
249 +
250 + assert out.shape == (2, 32)
251 + assert out.dtype == torch.bfloat16
252 + expected = torch.empty(2, 32)
253 + for row in range(2):
254 + for col in range(16):
255 + byte = int(packed[row, col])
256 + expected[row, 2 * col] = _E2M1_LUT[byte & 0x0F]
257 + expected[row, 2 * col + 1] = _E2M1_LUT[byte >> 4]
258 + expected[row] *= 2.0**row
259 + assert torch.equal(out, expected.to(torch.bfloat16))
vllm/config/kernel.py
+56 -1
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import contextlib
4 -from collections.abc import Callable
4 +from collections.abc import Callable, Iterable
5 from dataclasses import asdict, fields
6 from typing import TYPE_CHECKING, Any, Literal
7
@@ -131,6 +131,8 @@ MoEBackend = Literal[
131 "flashinfer_cutedsl",
132 "flashinfer_b12x",
133 "b12x",
134 + "flashinfer_moe_ep_mega_deep_gemm",
135 + "flashinfer_moe_ep_mega_cutedsl",
136 "marlin",
137 "humming",
138 "triton_unfused",
@@ -140,6 +142,46 @@ MoEBackend = Literal[
142 "emulation",
143 ]
144
145 +# Backends that run the mega-MoE model path through the flashinfer moe_ep
146 +# runtime. Only architectures in FLASHINFER_MOE_EP_ARCHITECTURES wire up
147 +# these experts.
148 +FLASHINFER_MOE_EP_BACKENDS = frozenset(
149 + {
150 + "flashinfer_moe_ep_mega_deep_gemm",
151 + "flashinfer_moe_ep_mega_cutedsl",
152 + }
153 +)
154 +
155 +# Backends that run a mega-MoE model path (fused expert module plus
156 +# prepare_megamoe routing): vLLM's native deep_gemm path, which any model
157 +# with a mega-MoE module may use (DeepSeek-V4, Kimi K3), plus the flashinfer
158 +# moe_ep variants.
159 +MEGA_MOE_BACKENDS = frozenset({"deep_gemm_mega_moe"}) | FLASHINFER_MOE_EP_BACKENDS
160 +
161 +# Architectures whose model code wires up the flashinfer moe_ep experts. MTP
162 +# and DSpark draft variants inherit the setting from these target models.
163 +FLASHINFER_MOE_EP_ARCHITECTURES = frozenset(
164 + {
165 + "DeepseekV4ForCausalLM",
166 + "DeepSeekV4MTPModel",
167 + }
168 +)
169 +
170 +
171 +def validate_flashinfer_moe_ep_model(
172 + moe_backend: str, architectures: Iterable[str]
173 +) -> None:
174 + """Reject flashinfer moe_ep backends for models that lack the FI path."""
175 + if moe_backend not in FLASHINFER_MOE_EP_BACKENDS:
176 + return
177 + if not any(arch in FLASHINFER_MOE_EP_ARCHITECTURES for arch in architectures):
178 + raise ValueError(
179 + f"moe_backend={moe_backend!r} is only supported for DeepSeek-V4 "
180 + f"models ({sorted(FLASHINFER_MOE_EP_ARCHITECTURES)}), but the "
181 + f"model is {list(architectures)}."
182 + )
183 +
184 +
185 LinearBackend = Literal[
186 "auto",
187 "cutlass",
@@ -206,6 +248,14 @@ class KernelConfig:
248 - "flashinfer_b12x": Use FlashInfer CuteDSL fused MoE for SM12x
249 (RTX Pro 6000 / DGX Spark)
250 - "b12x": Use b12x FP4 MoE kernels on SM12x
251 + - "flashinfer_moe_ep_mega_deep_gemm": Use the FlashInfer moe_ep
252 + expert-parallel mega-MoE with the DeepGEMM megakernel, which consumes an
253 + MXFP4 checkpoint verbatim (Blackwell, requires expert parallel;
254 + DeepSeek-V4 only)
255 + - "flashinfer_moe_ep_mega_cutedsl": Same, with the CuteDSL megakernel
256 + (additionally requires NVSHMEM). The checkpoint selects the weight path:
257 + an NVFP4 checkpoint is consumed prequantized, MXFP4 weights are
258 + requantized at load
259 - "marlin": Use Marlin kernels (weight-only quantization)
260 - "humming": Use Humming Mixed Precision kernels
261 - "triton_unfused": Use Triton unfused MoE kernels
@@ -288,6 +338,11 @@ class KernelConfig:
338 """Set platform-specific defaults for the kernel config."""
339 from vllm.platforms import current_platform
340
341 + if vllm_config.model_config is not None:
342 + validate_flashinfer_moe_ep_model(
343 + self.moe_backend, vllm_config.model_config.architectures
344 + )
345 +
346 platform_op_priority = current_platform.get_default_ir_op_priority(vllm_config)
347 logger.debug(
348 "Setting platform-specific IR op priority defaults: %s, user-defined: %s",
vllm/models/deepseek_v4/nvidia/fi_moe.py new
+397
@@ -0,0 +1,397 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""FlashInfer ``moe_ep`` expert module for DeepSeek V4."""
4 +
5 +from __future__ import annotations
6 +
7 +from typing import TYPE_CHECKING, Any
8 +
9 +import torch
10 +import torch.nn as nn
11 +
12 +from vllm.model_executor.utils import set_weight_attrs
13 +from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MegaMoEExperts
14 +from vllm.utils.flashinfer_moe_ep import (
15 + build_fi_mega_layer,
16 + ensure_fi_moe_ep_runtime,
17 + fi_moe_ep_backend_spec,
18 + make_fi_moe_ep_bootstrap,
19 + mega_moe_weight_pack_from_params,
20 +)
21 +
22 +if TYPE_CHECKING:
23 + from flashinfer.moe_ep import MoEEpMegaLayer
24 +
25 + from vllm.config import VllmConfig
26 + from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MLP
27 +
28 +_MOE_SKIP_PADDING: bool | None = None
29 +
30 +
31 +def ckpt_uses_nvfp4_experts(vllm_config: VllmConfig) -> bool:
32 + """True when the loaded checkpoint quantizes experts with modelopt NVFP4
33 + (e2m1 + fp8-e4m3 per-16 block scales + per-tensor weight_scale_2), i.e.
34 + the recipe the nvfp4_cutedsl prequantized-weights path consumes verbatim.
35 +
36 + DeepSeek-V4 NVFP4-expert checkpoints declare ``moe_quant_algo: NVFP4`` in
37 + their ``quantization_config``, which DeepseekV4FP8Config surfaces as
38 + ``moe_quant_algo``.
39 + """
40 + return getattr(vllm_config.quant_config, "moe_quant_algo", "") == "NVFP4"
41 +
42 +
43 +def resolve_mega_moe_is_padding(num_tokens: int) -> torch.Tensor | None:
44 + from vllm.forward_context import get_forward_context, is_forward_context_available
45 +
46 + global _MOE_SKIP_PADDING
47 + if _MOE_SKIP_PADDING is None:
48 + import vllm.envs as envs
49 +
50 + _MOE_SKIP_PADDING = bool(envs.VLLM_MOE_SKIP_PADDING)
51 + if not _MOE_SKIP_PADDING or not is_forward_context_available():
52 + return None
53 + is_padding = get_forward_context().is_padding
54 + if is_padding is None:
55 + return None
56 + return is_padding[:num_tokens]
57 +
58 +
59 +def apply_mega_moe_routing_preprocess(
60 + topk_ids: torch.Tensor,
61 + *,
62 + is_padding: torch.Tensor | None = None,
63 +) -> torch.Tensor:
64 + """Padding-only routing preprocess (EPLB hooks go here later)."""
65 + if is_padding is not None:
66 + topk_ids = torch.where(is_padding.unsqueeze(1), -1, topk_ids)
67 + return topk_ids
68 +
69 +
70 +def nvfp4_prequant_pack_and_alphas(
71 + w13_weight: torch.Tensor,
72 + w13_weight_scale: torch.Tensor,
73 + w13_weight_scale_2: torch.Tensor, # (E_local, 2) fp32: [:,0]=gate(w1), [:,1]=up(w3)
74 + w2_weight: torch.Tensor,
75 + w2_weight_scale: torch.Tensor,
76 + w2_weight_scale_2: torch.Tensor, # (E_local,) fp32
77 + *,
78 + intermediate_size: int,
79 +):
80 + """NVFP4 checkpoint params -> (MoEWeightPack, fc1_alpha, fc2_alpha).
81 +
82 + Activation quantization is fully dynamic, so the checkpoint's static
83 + ``input_scale`` drops out and each GEMM's epilogue alpha reduces to the
84 + weight's per-tensor ``weight_scale_2``.
85 +
86 + fc1_alpha is one scalar per expert, but gate (w1) and up (w3) carry their
87 + own ``weight_scale_2``. When they differ, the ratio is folded into the up
88 + half's e4m3 block scales — only if the fold round-trips exactly, since a
89 + lossy rescale would silently change the model.
90 + """
91 + from flashinfer.moe_ep import MoEWeightPack
92 +
93 + inter = intermediate_size
94 + gate_s2 = w13_weight_scale_2[:, 0].float()
95 + up_s2 = w13_weight_scale_2[:, 1].float()
96 + if (gate_s2 <= 0).any() or (up_s2 <= 0).any() or (w2_weight_scale_2 <= 0).any():
97 + raise ValueError(
98 + "nvfp4 prequant: non-positive weight_scale_2 loaded — checkpoint "
99 + "scale tensors missing or loader routed them wrong."
100 + )
101 +
102 + w13_scale = w13_weight_scale
103 + if not torch.equal(gate_s2, up_s2):
104 + ratio = up_s2 / gate_s2 # (E_local,)
105 + up_sf = w13_scale[:, inter:, :].float()
106 + folded = up_sf * ratio[:, None, None]
107 + folded_e4m3 = folded.to(torch.float8_e4m3fn)
108 + if not torch.equal(folded_e4m3.float(), folded):
109 + raise ValueError(
110 + "NVFP4 MoE cannot merge the gate and up weight_scale_2 values "
111 + "because their ratio is not exactly representable in "
112 + "float8_e4m3fn. Use a checkpoint quantized with a shared "
113 + "gate/up weight_scale_2."
114 + )
115 + w13_scale = w13_scale.clone()
116 + w13_scale[:, inter:, :] = folded_e4m3
117 +
118 + fc1_alpha = gate_s2.clone().contiguous()
119 + fc2_alpha = w2_weight_scale_2.float().clone().contiguous()
120 + pack = MoEWeightPack(
121 + w13=w13_weight,
122 + w2=w2_weight,
123 + w13_scale=w13_scale,
124 + w2_scale=w2_weight_scale,
125 + )
126 + return pack, fc1_alpha, fc2_alpha
127 +
128 +
129 +class DeepseekV4MegaMoEExpertsFI(DeepseekV4MegaMoEExperts):
130 + """Same weight layout/loader as the native mega experts, FI compute path."""
131 +
132 + def __init__(
133 + self,
134 + vllm_config: VllmConfig,
135 + *,
136 + activation_clamp: float | None = None,
137 + **kwargs: Any,
138 + ) -> None:
139 + super().__init__(vllm_config, **kwargs)
140 + self._vllm_config = vllm_config
141 + self._activation_clamp = activation_clamp
142 + self._mega_layer: MoEEpMegaLayer | None = None
143 + self._fast_ctx: tuple[Any, Any, Any, int, bool] | None = None
144 + self._epilogue_alphas: tuple[torch.Tensor, torch.Tensor] | None = None
145 + self._nvfp4_prequant = ckpt_uses_nvfp4_experts(vllm_config)
146 + if self._nvfp4_prequant:
147 + megakernel = fi_moe_ep_backend_spec(
148 + vllm_config.kernel_config.moe_backend
149 + ).megakernel
150 + if megakernel != "nvfp4_cutedsl":
151 + raise ValueError(
152 + "NVFP4-quantized expert checkpoint requires "
153 + "moe_backend=flashinfer_moe_ep_mega_cutedsl, got a "
154 + f"backend using megakernel {megakernel!r} "
155 + "(deep_gemm consumes the MXFP4 checkpoint instead)."
156 + )
157 + self._realloc_nvfp4_params()
158 +
159 + def _realloc_nvfp4_params(self) -> None:
160 + """Swap the mx-recipe scale params for the NVFP4 checkpoint's:
161 + fp8-e4m3 per-16 block scales plus the per-tensor second-level
162 + scales (weight_scale_2, input_scale) modelopt exports."""
163 + n_e = self.num_local_experts
164 + inter = self.intermediate_size
165 + hidden = self.hidden_size
166 + attrs = {"weight_loader": self.weight_loader}
167 +
168 + def _param(shape: tuple, dtype: torch.dtype) -> nn.Parameter:
169 + p = nn.Parameter(torch.zeros(*shape, dtype=dtype), requires_grad=False)
170 + set_weight_attrs(p, attrs)
171 + return p
172 +
173 + self.w13_weight_scale = _param(
174 + (n_e, 2 * inter, hidden // 16), torch.float8_e4m3fn
175 + )
176 + self.w13_weight_scale.quant_method = "block"
177 + self.w2_weight_scale = _param((n_e, hidden, inter // 16), torch.float8_e4m3fn)
178 + self.w2_weight_scale.quant_method = "block"
179 + # (E, 2): column 0 = gate (w1), column 1 = up (w3).
180 + self.w13_weight_scale_2 = _param((n_e, 2), torch.float32)
181 + self.w2_weight_scale_2 = _param((n_e,), torch.float32)
182 + # Loaded for completeness; unused — activation quant is dynamic
183 + # (see nvfp4_prequant_pack_and_alphas).
184 + self.w13_input_scale = _param((n_e, 2), torch.float32)
185 + self.w2_input_scale = _param((n_e,), torch.float32)
186 +
187 + def weight_loader(
188 + self,
189 + param: nn.Parameter,
190 + loaded_weight: torch.Tensor,
191 + weight_name: str,
192 + shard_id: str,
193 + expert_id: int,
194 + return_success: bool = False,
195 + ) -> bool | None:
196 + # NVFP4 checkpoint second-level scalars route here; everything
197 + # else (packed weights, block scales) matches the base layout.
198 + if "weight_scale_2" in weight_name or "input_scale" in weight_name:
199 + local_expert_ids = self._map_global_expert_id(expert_id)
200 + if not local_expert_ids:
201 + return False if return_success else None
202 + value = loaded_weight.reshape(()).to(torch.float32)
203 + for local_expert_id in local_expert_ids:
204 + if shard_id in ("w1", "w3"):
205 + if "w13_" not in weight_name:
206 + return False if return_success else None
207 + param.data[local_expert_id, 0 if shard_id == "w1" else 1] = value
208 + elif shard_id == "w2":
209 + if "w2_" not in weight_name:
210 + return False if return_success else None
211 + param.data[local_expert_id] = value
212 + else:
213 + raise ValueError(f"Unsupported expert shard id: {shard_id}")
214 + return True if return_success else None
215 + return super().weight_loader(
216 + param,
217 + loaded_weight,
218 + weight_name,
219 + shard_id,
220 + expert_id,
221 + return_success,
222 + )
223 +
224 + def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None:
225 + # The FlashInfer megakernel has no shared-expert fusion; the caller's
226 + # serial shared MLP path handles shared_experts.
227 + if self._mega_layer is not None:
228 + return
229 + if self.w13_weight is None:
230 + return
231 +
232 + self._check_runtime_supported()
233 + ensure_fi_moe_ep_runtime(self._vllm_config)
234 +
235 + if self._nvfp4_prequant:
236 + # NVFP4 checkpoint: hand the packed weights + both scale
237 + # planes straight to the backend (no dequant->requant);
238 + # per-expert globals become fc1/fc2 epilogue alphas staged
239 + # at every forward via MoEEpTensors.
240 + weights, fc1_alpha, fc2_alpha = nvfp4_prequant_pack_and_alphas(
241 + self.w13_weight.data,
242 + self.w13_weight_scale.data,
243 + self.w13_weight_scale_2.data,
244 + self.w2_weight.data,
245 + self.w2_weight_scale.data,
246 + self.w2_weight_scale_2.data,
247 + intermediate_size=self.intermediate_size,
248 + )
249 + self._epilogue_alphas = (fc1_alpha, fc2_alpha)
250 + else:
251 + weights = mega_moe_weight_pack_from_params(
252 + self.w13_weight,
253 + self.w13_weight_scale,
254 + self.w2_weight,
255 + self.w2_weight_scale,
256 + megakernel=fi_moe_ep_backend_spec(
257 + self._vllm_config.kernel_config.moe_backend
258 + ).megakernel,
259 + )
260 + self._mega_layer = build_fi_mega_layer(
261 + make_fi_moe_ep_bootstrap(),
262 + vllm_config=self._vllm_config,
263 + num_experts=self.num_experts,
264 + max_tokens_per_rank=self.max_num_tokens,
265 + hidden_size=self.hidden_size,
266 + intermediate_size=self.intermediate_size,
267 + top_k=self.top_k,
268 + activation_clamp=self._activation_clamp,
269 + weights=weights,
270 + )
271 + del weights
272 + # Allocate (or attach to) the pooled workspace before first
273 + # forward so warmup/capture never hits the lazy path.
274 + self._mega_layer._ensure_workspace()
275 + self.w13_weight = None
276 + self.w13_weight_scale = None
277 + self.w2_weight = None
278 + self.w2_weight_scale = None
279 + if self._nvfp4_prequant:
280 + self.w13_weight_scale_2 = None
281 + self.w2_weight_scale_2 = None
282 + self.w13_input_scale = None
283 + self.w2_input_scale = None
284 +
285 + # EPLB is rejected for these backends in validate_fi_moe_ep_config, so
286 + # nothing constructs an EplbState and none of these run. They are kept
287 + # as explicit errors rather than deleted because inheriting the native
288 + # implementations would be worse: set_eplb_state would record a map
289 + # this forward path never consults, silently splitting routing from
290 + # weights, and get_expert_weights reads transformed-weight attributes
291 + # the flashinfer path releases after preprocess.
292 + _EPLB_UNSUPPORTED = (
293 + "EPLB is not supported with the flashinfer moe_ep backends; "
294 + "validate_fi_moe_ep_config should have rejected this "
295 + "configuration at startup."
296 + )
297 +
298 + def set_eplb_state(
299 + self,
300 + moe_layer_idx: int,
301 + expert_load_view: torch.Tensor,
302 + logical_to_physical_map: torch.Tensor,
303 + logical_replica_count: torch.Tensor,
304 + ) -> None:
305 + raise NotImplementedError(self._EPLB_UNSUPPORTED)
306 +
307 + def get_expert_weights(self) -> list[torch.Tensor]:
308 + raise NotImplementedError(self._EPLB_UNSUPPORTED)
309 +
310 + def update_expert_map(self) -> None:
311 + raise NotImplementedError(self._EPLB_UNSUPPORTED)
312 +
313 + def forward(
314 + self,
315 + hidden_states: torch.Tensor,
316 + topk_weights: torch.Tensor,
317 + topk_ids: torch.Tensor,
318 + *,
319 + activation_clamp: float | None,
320 + fast_math: bool = True,
321 + ) -> torch.Tensor:
322 + # fast_math is a native deep_gemm knob; the FI kernels have no
323 + # equivalent toggle, so it is accepted for signature parity only.
324 + if hidden_states.shape[0] > self.max_num_tokens:
325 + raise ValueError(
326 + f"DeepSeek V4 MegaMoE got {hidden_states.shape[0]} tokens, "
327 + f"but the symmetric buffer was sized for {self.max_num_tokens}."
328 + )
329 +
330 + from flashinfer.moe_ep import MoEEpTensors
331 +
332 + num_tokens = hidden_states.shape[0]
333 + is_padding = resolve_mega_moe_is_padding(num_tokens)
334 + topk_ids = apply_mega_moe_routing_preprocess(
335 + topk_ids,
336 + is_padding=is_padding,
337 + )
338 +
339 + # Fast path: after the first successful full forward the layer is
340 + # immutable, so skip MoEEpMegaLayer.forward()'s per-call validation
341 + # and go straight to the kernel backend's stage_inputs + compute.
342 + alphas = self._epilogue_alphas
343 + fc1_alpha = alphas[0] if alphas is not None else None
344 + fc2_alpha = alphas[1] if alphas is not None else None
345 +
346 + fast = self._fast_ctx
347 + if fast is not None:
348 + kernel, workspace, transformed, hidden_size, zero_copy = fast
349 + t = MoEEpTensors(
350 + hidden_states=hidden_states,
351 + topk_ids=topk_ids,
352 + topk_weights=topk_weights,
353 + fc1_alpha=fc1_alpha,
354 + fc2_alpha=fc2_alpha,
355 + )
356 + kernel.stage_inputs(t, workspace, quantize_input=True)
357 + if zero_copy:
358 + # Zero-copy (cutedsl backends): consume the workspace [:n]
359 + # view directly — valid under stream ordering until the next
360 + # MoE layer's launch on the shared workspace.
361 + return kernel.compute(workspace, transformed, output=None)
362 + # deep_gemm_mega's compute() requires a real output tensor.
363 + out = torch.empty(
364 + num_tokens,
365 + hidden_size,
366 + dtype=torch.bfloat16,
367 + device=hidden_states.device,
368 + )
369 + return kernel.compute(workspace, transformed, output=out)
370 +
371 + ensure_fi_moe_ep_runtime(self._vllm_config)
372 + self.finalize_weights()
373 + assert self._mega_layer is not None
374 +
375 + y = self._mega_layer.forward(
376 + MoEEpTensors(
377 + hidden_states=hidden_states,
378 + topk_ids=topk_ids,
379 + topk_weights=topk_weights,
380 + fc1_alpha=fc1_alpha,
381 + fc2_alpha=fc2_alpha,
382 + )
383 + )
384 + layer = self._mega_layer
385 + if hidden_states.dtype == torch.bfloat16:
386 + self._fast_ctx = (
387 + layer._kernel,
388 + layer._ensure_workspace(),
389 + layer._transformed,
390 + layer._fleet_params.token_hidden_size,
391 + # zero-copy output views are a cutedsl-backend contract
392 + layer._kernel.kernel_name() != "deep_gemm_mega",
393 + )
394 + return y
395 +
396 +
397 +DeepseekV4MegaMoEExpertsFI.weight_loader.supports_moe_loading = True # type: ignore[attr-defined]
vllm/models/deepseek_v4/nvidia/model.py
+32 -11
@@ -11,6 +11,7 @@ import torch.nn as nn
11
12 import vllm.envs as envs
13 from vllm.config import VllmConfig
14 +from vllm.config.kernel import MEGA_MOE_BACKENDS
15 from vllm.distributed import (
16 get_ep_group,
17 get_pp_group,
@@ -83,6 +84,10 @@ from vllm.models.deepseek_v4.nvidia.flashmla import DeepseekV4FlashMLAAttention
84 from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
85 from vllm.platforms import current_platform
86 from vllm.sequence import IntermediateTensors
87 +from vllm.utils.flashinfer_moe_ep import (
88 + is_fi_moe_ep_backend,
89 + validate_fi_moe_ep_config,
90 +)
91 from vllm.utils.math_utils import cdiv
92 from vllm.v1.attention.backends.registry import AttentionBackendEnum
93 from vllm.v1.worker.ubatching import dbo_current_ubatch_id
@@ -752,9 +757,10 @@ class DeepseekV4MoE(nn.Module):
757 quant_config = vllm_config.quant_config
758 self.prefix = prefix
759 self.use_sequence_parallel = use_sequence_parallel
755 - self.use_mega_moe = (
756 - vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
757 - )
760 + moe_backend = vllm_config.kernel_config.moe_backend
761 + validate_fi_moe_ep_config(vllm_config)
762 + self.use_mega_moe = moe_backend in MEGA_MOE_BACKENDS
763 + self.use_fi_mega_moe = is_fi_moe_ep_backend(moe_backend)
764 if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:
765 raise NotImplementedError(
766 "DeepSeek V4 MegaMoE currently requires expert parallel. "
@@ -779,7 +785,7 @@ class DeepseekV4MoE(nn.Module):
785 raise NotImplementedError(
786 "DeepSeek V4 MegaMoE only supports fp4 experts; got expert_dtype="
787 f"{config.expert_dtype!r}. Drop --kernel-config moe_backend="
782 - "deep_gemm_mega_moe for this checkpoint."
788 + f"{moe_backend} for this checkpoint."
789 )
790
791 self.gate = GateLinear(
@@ -867,15 +873,29 @@ class DeepseekV4MoE(nn.Module):
873 # Native DeepGEMM fusion requires each EP rank to own the complete
874 # shared MLP. Sequence parallel replicates those weights while sharding
875 # tokens. TP=1 is also naturally replicated. With PP+TP the shared MLP
870 - # remains tensor-sharded, so retain the serial path.
876 + # remains tensor-sharded, so retain the serial path. The FlashInfer
877 + # megakernel has no shared-expert fusion, so it keeps the serial path.
878 fuse_shared_experts = bool(
879 self.shared_experts is not None
880 and not envs.VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION
881 and (self.use_sequence_parallel or self.tp_size == 1)
882 + and not self.use_fi_mega_moe
883 )
884
877 - self.experts = DeepseekV4MegaMoEExperts(
878 - vllm_config,
885 + activation_clamp = (
886 + float(self.swiglu_limit) if self.swiglu_limit is not None else None
887 + )
888 + if self.use_fi_mega_moe:
889 + # Deferred: fi_moe subclasses DeepseekV4MegaMoEExperts, so a
890 + # module-level import here would be circular.
891 + from vllm.models.deepseek_v4.nvidia.fi_moe import (
892 + DeepseekV4MegaMoEExpertsFI,
893 + )
894 +
895 + experts_cls: type[DeepseekV4MegaMoEExperts] = DeepseekV4MegaMoEExpertsFI
896 + else:
897 + experts_cls = DeepseekV4MegaMoEExperts
898 + expert_kwargs: dict[str, typing.Any] = dict(
899 num_experts=self.n_physical_experts,
900 num_local_experts=self.n_local_physical_experts,
901 experts_start_idx=self.physical_expert_start,
@@ -886,6 +906,9 @@ class DeepseekV4MoE(nn.Module):
906 num_shared_experts=(self.n_shared_experts if fuse_shared_experts else 0),
907 prefix=f"{prefix}.experts",
908 )
909 + if self.use_fi_mega_moe:
910 + expert_kwargs["activation_clamp"] = activation_clamp
911 + self.experts = experts_cls(vllm_config, **expert_kwargs)
912
913 def _init_fused_moe_experts(
914 self,
@@ -1031,7 +1054,7 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]:
1054
1055 def _use_sequence_parallel(vllm_config: VllmConfig) -> bool:
1056 parallel_config = vllm_config.parallel_config
1034 - use_mega_moe = vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
1057 + use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS
1058 return (
1059 parallel_config.pipeline_parallel_size == 1
1060 and parallel_config.enable_expert_parallel
@@ -1226,9 +1249,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
1249 self.config = config
1250 self.quant_config = quant_config
1251 self.parallel_config = vllm_config.parallel_config
1229 - self.use_mega_moe = (
1230 - vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
1231 - )
1252 + self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS
1253 self.use_sequence_parallel = _use_sequence_parallel(vllm_config)
1254 if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:
1255 raise NotImplementedError(
vllm/utils/flashinfer_moe_ep.py new
+339
@@ -0,0 +1,339 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""FlashInfer ``moe_ep`` helpers for DeepSeek V4 vLLM integration."""
4 +
5 +from __future__ import annotations
6 +
7 +from dataclasses import dataclass
8 +from typing import TYPE_CHECKING, Any
9 +
10 +import torch
11 +import torch.nn as nn
12 +
13 +from vllm.config.kernel import FLASHINFER_MOE_EP_BACKENDS
14 +from vllm.distributed import get_ep_group
15 +from vllm.platforms import current_platform
16 +
17 +if TYPE_CHECKING:
18 + from flashinfer.moe_ep import BootstrapConfig, MoEEpMegaLayer
19 +
20 + from vllm.config import VllmConfig
21 +
22 +
23 +# Every mega kernel is Blackwell-only, so the arch is a property of the family
24 +# rather than of a particular backend: validate it against the live device
25 +# instead of encoding it in the backend name. An explicit allowlist so new
26 +# archs (SM110/SM120, Rubin) fail loudly until flashinfer supports them.
27 +FI_MOE_EP_SUPPORTED_CAPABILITIES = frozenset({(10, 0), (10, 3)})
28 +
29 +
30 +@dataclass(frozen=True)
31 +class FiMoeEpBackendSpec:
32 + """Static properties of one ``flashinfer_moe_ep_*`` backend string.
33 +
34 + The backend names the kernel *family*; the arch comes from the device and
35 + the weight handling from the checkpoint, so this only has to carry which
36 + megakernel to build and whether it needs NVSHMEM in the runtime set.
37 + """
38 +
39 + megakernel: str
40 + needs_nvshmem: bool
41 +
42 +
43 +FI_MOE_EP_BACKEND_SPECS: dict[str, FiMoeEpBackendSpec] = {
44 + # Consumes an MXFP4 checkpoint verbatim (e2m1 weights + E8M0 per-32
45 + # scales) -- the same recipe the native deep_gemm mega path uses.
46 + "flashinfer_moe_ep_mega_deep_gemm": FiMoeEpBackendSpec(
47 + megakernel="deep_gemm_mega",
48 + needs_nvshmem=False,
49 + ),
50 + # The checkpoint picks the weight path, not the kernel: an NVFP4
51 + # checkpoint is consumed prequantized (no round trip), while MXFP4 weights
52 + # are dequantized to bf16 and requantized. See
53 + # models/deepseek_v4/nvidia/fi_moe.py:ckpt_uses_nvfp4_experts.
54 + "flashinfer_moe_ep_mega_cutedsl": FiMoeEpBackendSpec(
55 + megakernel="nvfp4_cutedsl",
56 + needs_nvshmem=True,
57 + ),
58 +}
59 +
60 +assert set(FI_MOE_EP_BACKEND_SPECS) == FLASHINFER_MOE_EP_BACKENDS
61 +
62 +_FI_RUNTIME_HANDLE: Any = None
63 +
64 +
65 +def is_fi_moe_ep_backend(moe_backend: str) -> bool:
66 + return moe_backend in FI_MOE_EP_BACKEND_SPECS
67 +
68 +
69 +def fi_moe_ep_backend_spec(moe_backend: str) -> FiMoeEpBackendSpec:
70 + try:
71 + return FI_MOE_EP_BACKEND_SPECS[moe_backend]
72 + except KeyError:
73 + raise ValueError(
74 + f"{moe_backend!r} is not a flashinfer moe_ep backend; expected "
75 + f"one of {sorted(FI_MOE_EP_BACKEND_SPECS)}"
76 + ) from None
77 +
78 +
79 +def validate_fi_moe_ep_config(vllm_config: VllmConfig) -> None:
80 + """Config-time checks for the mega-MoE backends, native and flashinfer."""
81 + moe_backend = vllm_config.kernel_config.moe_backend
82 + if not is_fi_moe_ep_backend(moe_backend):
83 + return
84 +
85 + # flashinfer validates the arch too, but not until the layer constructor
86 + # runs during weight load; check here so the error names the flag the user
87 + # actually typed.
88 + capability = current_platform.get_device_capability()
89 + if capability is not None:
90 + cc = (capability.major, capability.minor)
91 + if cc not in FI_MOE_EP_SUPPORTED_CAPABILITIES:
92 + supported = ", ".join(
93 + f"{m}.{n}" for m, n in sorted(FI_MOE_EP_SUPPORTED_CAPABILITIES)
94 + )
95 + raise ValueError(
96 + f"moe_backend={moe_backend!r} is only supported on compute "
97 + f"capability {supported} (SM100/SM103), but this device is "
98 + f"{cc[0]}.{cc[1]}."
99 + )
100 +
101 + if vllm_config.parallel_config.enable_eplb:
102 + raise NotImplementedError(
103 + f"EPLB is not supported with moe_backend={moe_backend!r}: the "
104 + "flashinfer moe_ep experts neither apply the logical-to-physical "
105 + "expert map nor report per-expert load, so rebalancing would move "
106 + "weights without moving routing. Use "
107 + "moe_backend=deep_gemm_mega_moe to run the mega path with EPLB."
108 + )
109 +
110 +
111 +def make_fi_moe_ep_bootstrap() -> BootstrapConfig:
112 + from flashinfer.moe_ep import BootstrapConfig
113 +
114 + ep = get_ep_group()
115 + return BootstrapConfig(
116 + world_size=ep.world_size,
117 + rank=ep.rank_in_group,
118 + process_group=ep.device_group,
119 + auto_bootstrap=False,
120 + # vLLM already bound this worker's (possibly remapped) device;
121 + # without this the runtime would rebind to cuda:LOCAL_RANK|rank and
122 + # launch weight transforms against another device's pointers.
123 + device=torch.accelerator.current_device_index(),
124 + )
125 +
126 +
127 +def megakernel_runtime_requirements(spec: FiMoeEpBackendSpec) -> frozenset[str]:
128 + from flashinfer.moe_ep.core.runtime import NVSHMEM, TORCH_DIST
129 +
130 + if spec.needs_nvshmem:
131 + return frozenset({TORCH_DIST, NVSHMEM})
132 + return frozenset({TORCH_DIST})
133 +
134 +
135 +def ensure_fi_moe_ep_runtime(vllm_config: VllmConfig) -> None:
136 + """Acquire the process-wide flashinfer moe_ep runtime once per worker."""
137 + global _FI_RUNTIME_HANDLE
138 + if _FI_RUNTIME_HANDLE is not None:
139 + return
140 +
141 + from flashinfer.moe_ep import bootstrap_moe_ep_runtime
142 +
143 + bootstrap = make_fi_moe_ep_bootstrap()
144 + spec = fi_moe_ep_backend_spec(vllm_config.kernel_config.moe_backend)
145 + _FI_RUNTIME_HANDLE = bootstrap_moe_ep_runtime(
146 + bootstrap,
147 + megakernel_runtime_requirements(spec),
148 + )
149 +
150 +
151 +def finalize_fi_moe_ep_runtime() -> None:
152 + """Release the process-wide flashinfer moe_ep runtime."""
153 + global _FI_RUNTIME_HANDLE
154 + if _FI_RUNTIME_HANDLE is None:
155 + return
156 +
157 + from flashinfer.moe_ep import finalize_moe_ep_runtime
158 +
159 + finalize_moe_ep_runtime(_FI_RUNTIME_HANDLE)
160 + _FI_RUNTIME_HANDLE = None
161 +
162 +
163 +_E2M1_LUT = (
164 + 0.0,
165 + 0.5,
166 + 1.0,
167 + 1.5,
168 + 2.0,
169 + 3.0,
170 + 4.0,
171 + 6.0,
172 + -0.0,
173 + -0.5,
174 + -1.0,
175 + -1.5,
176 + -2.0,
177 + -3.0,
178 + -4.0,
179 + -6.0,
180 +)
181 +
182 +
183 +def _dequant_fp4_ue8m0_gran32(
184 + packed: torch.Tensor, sf_ue8m0: torch.Tensor
185 +) -> torch.Tensor:
186 + """[rows, K//2] packed e2m1 + [rows, K//32] ue8m0-uint8 scales -> bf16 [rows, K]."""
187 + raw = packed.view(torch.uint8)
188 + lut = torch.tensor(_E2M1_LUT, dtype=torch.float32, device=raw.device)
189 + vals = torch.empty(
190 + raw.shape[0], raw.shape[1] * 2, dtype=torch.float32, device=raw.device
191 + )
192 + vals[:, ::2] = lut[(raw & 0x0F).to(torch.int64)]
193 + vals[:, 1::2] = lut[(raw >> 4).to(torch.int64)]
194 + sf = (sf_ue8m0.to(torch.int32) << 23).view(torch.float32)
195 + return (vals * sf.repeat_interleave(32, dim=-1)).to(torch.bfloat16)
196 +
197 +
198 +def _dequant_expert_weights_to_bf16(
199 + weight: torch.Tensor, scale: torch.Tensor
200 +) -> torch.Tensor:
201 + """[E, N, K//2] fp4 + [E, N, K//32] ue8m0 -> [E, N, K] bf16 (expert loop)."""
202 + num_experts, n, k_half = weight.shape
203 + out = torch.empty(
204 + num_experts, n, k_half * 2, dtype=torch.bfloat16, device=weight.device
205 + )
206 + for e in range(num_experts):
207 + out[e] = _dequant_fp4_ue8m0_gran32(weight[e], scale[e])
208 + return out
209 +
210 +
211 +def mega_moe_weight_pack_from_params(
212 + w13_weight: nn.Parameter,
213 + w13_weight_scale: nn.Parameter,
214 + w2_weight: nn.Parameter,
215 + w2_weight_scale: nn.Parameter,
216 + *,
217 + megakernel: str = "deep_gemm_mega",
218 +):
219 + from flashinfer.moe_ep import MoEWeightPack
220 +
221 + if megakernel == "deep_gemm_mega":
222 + # Same fp4-e2m1 + ue8m0-per-32 recipe as the native path: pass verbatim,
223 + # flashinfer runs the identical deep_gemm transform.
224 + return MoEWeightPack(
225 + w13=w13_weight.data,
226 + w2=w2_weight.data,
227 + w13_scale=w13_weight_scale.data,
228 + w2_scale=w2_weight_scale.data,
229 + )
230 + # The cutedsl kernel quantizes with its own recipe (nvfp4
231 + # e2m1+e4m3-per-16): dequantize the checkpoint fp4 to bf16 and let the
232 + # backend preprocess requantize. Double quantization: outputs are close to
233 + # but not bit-identical with the native path.
234 + return MoEWeightPack(
235 + w13=_dequant_expert_weights_to_bf16(w13_weight.data, w13_weight_scale.data),
236 + w2=_dequant_expert_weights_to_bf16(w2_weight.data, w2_weight_scale.data),
237 + )
238 +
239 +
240 +def build_fi_mega_config(
241 + *,
242 + intermediate_size: int,
243 + top_k: int,
244 + activation_clamp: float | None,
245 + megakernel: str,
246 +):
247 + from flashinfer.moe_ep import (
248 + DeepGemmMegaMoeConfig,
249 + MegaConfig,
250 + Nvfp4CutedslMegaMoeConfig,
251 + )
252 +
253 + # fast_math selects approximate exp/rcp in DeepGEMM's fused SwiGLU
254 + # epilogue; the cutedsl kernels accept it for API parity only.
255 + if megakernel == "deep_gemm_mega":
256 + mk = DeepGemmMegaMoeConfig(
257 + intermediate_size=intermediate_size,
258 + top_k=top_k,
259 + activation_clamp=activation_clamp,
260 + fast_math=True,
261 + )
262 + elif megakernel == "nvfp4_cutedsl":
263 + mk = Nvfp4CutedslMegaMoeConfig(
264 + intermediate_size=intermediate_size,
265 + top_k=top_k,
266 + activation_clamp=activation_clamp,
267 + fast_math=True,
268 + )
269 + else:
270 + raise ValueError(f"Unsupported fi_moe_ep megakernel {megakernel!r}")
271 +
272 + return MegaConfig(
273 + megakernel=mk,
274 + preprocess_weights=True,
275 + quantize_input=True,
276 + )
277 +
278 +
279 +# All MoE layers share one symmetric workspace, like the native path's
280 +# class-level DeepseekV4MegaMoEExperts._symm_buffer_cache. Without this the
281 +# fi path allocates one symm buffer PER LAYER (43x memory + cold working
282 +# sets); the workspace is stateless across forwards (kernel tail-cleans) and
283 +# layers execute sequentially on one stream, so sharing is safe.
284 +def build_fi_mega_layer(
285 + bootstrap: BootstrapConfig,
286 + *,
287 + vllm_config: VllmConfig,
288 + num_experts: int,
289 + max_tokens_per_rank: int,
290 + hidden_size: int,
291 + intermediate_size: int,
292 + top_k: int,
293 + activation_clamp: float | None,
294 + weights,
295 +) -> MoEEpMegaLayer:
296 + from flashinfer.moe_ep import FleetParams, MoEEpLayer
297 +
298 + megakernel = fi_moe_ep_backend_spec(
299 + vllm_config.kernel_config.moe_backend
300 + ).megakernel
301 + mega_config = build_fi_mega_config(
302 + intermediate_size=intermediate_size,
303 + top_k=top_k,
304 + activation_clamp=activation_clamp,
305 + megakernel=megakernel,
306 + )
307 + layer = MoEEpLayer(
308 + bootstrap=bootstrap,
309 + fleet_params=FleetParams(
310 + num_experts=num_experts,
311 + max_tokens_per_rank=max_tokens_per_rank,
312 + token_hidden_size=hidden_size,
313 + ),
314 + weights=weights,
315 + backend=mega_config,
316 + )
317 + from flashinfer.moe_ep import MoEEpMegaLayer
318 +
319 + if not isinstance(layer, MoEEpMegaLayer):
320 + raise TypeError(
321 + f"fi_moe_ep expected MoEEpMegaLayer, got {type(layer).__name__}"
322 + )
323 + return layer
324 +
325 +
326 +__all__ = [
327 + "FI_MOE_EP_BACKEND_SPECS",
328 + "FiMoeEpBackendSpec",
329 + "build_fi_mega_config",
330 + "build_fi_mega_layer",
331 + "ensure_fi_moe_ep_runtime",
332 + "fi_moe_ep_backend_spec",
333 + "finalize_fi_moe_ep_runtime",
334 + "is_fi_moe_ep_backend",
335 + "make_fi_moe_ep_bootstrap",
336 + "mega_moe_weight_pack_from_params",
337 + "megakernel_runtime_requirements",
338 + "validate_fi_moe_ep_config",
339 +]