iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
Md Saidul Hoque Anik
2026-08-25 03:35:01 +00:00
committed by GitHub
co-authored by Claude Fable 5 Anerudhan Gopal
parent d3e2888c75
commit 8fe9317f2e
5 changed files with 1083 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/models/test_deepseek_v4_fi_moe_ep.py
+259
View File
@@ -0,0 +1,259 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the flashinfer moe_ep backend plumbing.
Everything here runs without a GPU or a flashinfer install: the flashinfer
modules the helpers import lazily are replaced with capture fakes.
"""
import sys
from dataclasses import dataclass, field
from types import ModuleType, SimpleNamespace
from typing import Any
import pytest
import torch
from vllm.config.kernel import (
FLASHINFER_MOE_EP_BACKENDS,
MEGA_MOE_BACKENDS,
validate_flashinfer_moe_ep_model,
)
from vllm.utils.flashinfer_moe_ep import (
_E2M1_LUT,
FI_MOE_EP_BACKEND_SPECS,
_dequant_fp4_ue8m0_gran32,
build_fi_mega_config,
fi_moe_ep_backend_spec,
make_fi_moe_ep_bootstrap,
megakernel_runtime_requirements,
)
@dataclass
class _FakeBootstrapConfig:
world_size: int
rank: int
process_group: Any = None
auto_bootstrap: bool = True
device: int | None = field(default=None, kw_only=True)
@dataclass
class _FakeDeepGemmMegaMoeConfig:
intermediate_size: int
top_k: int
activation_clamp: float | None
fast_math: bool
@dataclass
class _FakeNvfp4CutedslMegaMoeConfig:
intermediate_size: int
top_k: int
activation_clamp: float | None
fast_math: bool
@dataclass
class _FakeMegaConfig:
megakernel: Any
preprocess_weights: bool
quantize_input: bool
@pytest.fixture
def fake_flashinfer(monkeypatch):
"""Install a minimal fake flashinfer.moe_ep for the lazy imports."""
moe_ep = ModuleType("flashinfer.moe_ep")
core = ModuleType("flashinfer.moe_ep.core")
runtime = ModuleType("flashinfer.moe_ep.core.runtime")
flashinfer = ModuleType("flashinfer")
fake_attrs: dict[ModuleType, dict[str, Any]] = {
moe_ep: {
"BootstrapConfig": _FakeBootstrapConfig,
"DeepGemmMegaMoeConfig": _FakeDeepGemmMegaMoeConfig,
"Nvfp4CutedslMegaMoeConfig": _FakeNvfp4CutedslMegaMoeConfig,
"MegaConfig": _FakeMegaConfig,
"core": core,
},
runtime: {"TORCH_DIST": "torch_dist", "NVSHMEM": "nvshmem"},
flashinfer: {"moe_ep": moe_ep},
core: {"runtime": runtime},
}
for mod, attrs in fake_attrs.items():
for attr, value in attrs.items():
setattr(mod, attr, value)
for name, mod in {
"flashinfer": flashinfer,
"flashinfer.moe_ep": moe_ep,
"flashinfer.moe_ep.core": core,
"flashinfer.moe_ep.core.runtime": runtime,
}.items():
monkeypatch.setitem(sys.modules, name, mod)
return moe_ep
def test_fi_backend_strings_are_registered_mega_moe_backends():
assert set(FI_MOE_EP_BACKEND_SPECS) == FLASHINFER_MOE_EP_BACKENDS
assert FLASHINFER_MOE_EP_BACKENDS < MEGA_MOE_BACKENDS
@pytest.mark.parametrize("moe_backend", sorted(FLASHINFER_MOE_EP_BACKENDS))
def test_fi_moe_ep_backend_rejected_for_non_dsv4(moe_backend):
"""An FI moe_ep backend with a non-DSv4 model must fail at config time
instead of silently falling through to the generic FusedMoE path."""
with pytest.raises(ValueError, match="only supported for DeepSeek-V4"):
validate_flashinfer_moe_ep_model(moe_backend, ["MixtralForCausalLM"])
@pytest.mark.parametrize("moe_backend", sorted(FLASHINFER_MOE_EP_BACKENDS))
def test_fi_moe_ep_backend_accepted_for_dsv4(moe_backend):
validate_flashinfer_moe_ep_model(moe_backend, ["DeepseekV4ForCausalLM"])
@pytest.mark.parametrize(
"architectures",
[["KimiK3ForConditionalGeneration"], ["MixtralForCausalLM"]],
)
def test_native_deep_gemm_mega_moe_not_arch_gated(architectures):
"""vLLM's own deep_gemm mega path is not DSv4-only (Kimi K3 uses it);
models validate their own constraints at construction time."""
validate_flashinfer_moe_ep_model("deep_gemm_mega_moe", architectures)
def test_non_fi_backend_ignores_architectures():
validate_flashinfer_moe_ep_model("auto", ["MixtralForCausalLM"])
@pytest.mark.parametrize("moe_backend", sorted(MEGA_MOE_BACKENDS))
def test_all_mega_backends_get_sequence_parallel_moe(moe_backend):
"""Every mega backend must qualify for sequence-parallel MoE at
TP>1/EP: the predicate once matched only the native backend string,
which silently ran the fi backends full-batch with an all-reduce on
every rank — 0.42-0.65x native e2e at TP8."""
from vllm.models.deepseek_v4.nvidia.model import _use_sequence_parallel
vllm_config = SimpleNamespace(
parallel_config=SimpleNamespace(
pipeline_parallel_size=1,
enable_expert_parallel=True,
tensor_parallel_size=8,
data_parallel_size=1,
),
kernel_config=SimpleNamespace(moe_backend=moe_backend),
)
assert _use_sequence_parallel(vllm_config)
def test_fi_moe_ep_backend_spec_kernel_and_nvshmem_contract():
dg = fi_moe_ep_backend_spec("flashinfer_moe_ep_mega_deep_gemm")
assert dg.megakernel == "deep_gemm_mega"
assert not dg.needs_nvshmem
cd = fi_moe_ep_backend_spec("flashinfer_moe_ep_mega_cutedsl")
assert cd.megakernel == "nvfp4_cutedsl"
assert cd.needs_nvshmem
with pytest.raises(ValueError, match="not a flashinfer moe_ep backend"):
fi_moe_ep_backend_spec("deep_gemm_mega_moe")
def test_megakernel_runtime_requirements(fake_flashinfer):
dg = megakernel_runtime_requirements(
fi_moe_ep_backend_spec("flashinfer_moe_ep_mega_deep_gemm")
)
assert dg == frozenset({"torch_dist"})
cd = megakernel_runtime_requirements(
fi_moe_ep_backend_spec("flashinfer_moe_ep_mega_cutedsl")
)
assert cd == frozenset({"torch_dist", "nvshmem"})
def test_bootstrap_pins_the_device_vllm_bound(fake_flashinfer, monkeypatch):
"""The runtime must not rederive the device from LOCAL_RANK/rank: under a
remapped CUDA_VISIBLE_DEVICES that ordinal points at the wrong GPU
(CUDA_ERROR_ILLEGAL_ADDRESS in the weight transforms). vLLM passes the
device it already bound via BootstrapConfig.device."""
import vllm.utils.flashinfer_moe_ep as mod
pg = object()
monkeypatch.setattr(
mod,
"get_ep_group",
lambda: SimpleNamespace(world_size=4, rank_in_group=2, device_group=pg),
)
monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 3)
bootstrap = make_fi_moe_ep_bootstrap()
assert bootstrap.world_size == 4
assert bootstrap.rank == 2
assert bootstrap.process_group is pg
assert bootstrap.auto_bootstrap is False
assert bootstrap.device == 3
def test_build_fi_mega_config_selects_kernel_config(fake_flashinfer):
dg = build_fi_mega_config(
intermediate_size=2048,
top_k=8,
activation_clamp=7.0,
megakernel="deep_gemm_mega",
)
assert isinstance(dg.megakernel, _FakeDeepGemmMegaMoeConfig)
assert dg.megakernel.intermediate_size == 2048
assert dg.megakernel.top_k == 8
assert dg.megakernel.activation_clamp == 7.0
assert dg.preprocess_weights and dg.quantize_input
cd = build_fi_mega_config(
intermediate_size=2048,
top_k=8,
activation_clamp=None,
megakernel="nvfp4_cutedsl",
)
assert isinstance(cd.megakernel, _FakeNvfp4CutedslMegaMoeConfig)
with pytest.raises(ValueError, match="Unsupported fi_moe_ep megakernel"):
build_fi_mega_config(
intermediate_size=2048,
top_k=8,
activation_clamp=None,
megakernel="deep_gemm",
)
def test_ckpt_uses_nvfp4_experts_reads_moe_quant_algo():
from vllm.models.deepseek_v4.nvidia.fi_moe import ckpt_uses_nvfp4_experts
nvfp4 = SimpleNamespace(quant_config=SimpleNamespace(moe_quant_algo="NVFP4"))
assert ckpt_uses_nvfp4_experts(nvfp4)
mxfp4 = SimpleNamespace(quant_config=SimpleNamespace(moe_quant_algo=None))
assert not ckpt_uses_nvfp4_experts(mxfp4)
no_algo = SimpleNamespace(quant_config=SimpleNamespace())
assert not ckpt_uses_nvfp4_experts(no_algo)
def test_dequant_fp4_ue8m0_gran32_decodes_lut_and_scales():
"""One 32-element scale group per row: low nibble is the even element,
high nibble the odd one, ue8m0 scale applies to the whole group."""
packed = torch.arange(32, dtype=torch.uint8).reshape(2, 16)
sf = torch.tensor([[127], [128]], dtype=torch.uint8) # 2**0, 2**1
out = _dequant_fp4_ue8m0_gran32(packed, sf)
assert out.shape == (2, 32)
assert out.dtype == torch.bfloat16
expected = torch.empty(2, 32)
for row in range(2):
for col in range(16):
byte = int(packed[row, col])
expected[row, 2 * col] = _E2M1_LUT[byte & 0x0F]
expected[row, 2 * col + 1] = _E2M1_LUT[byte >> 4]
expected[row] *= 2.0**row
assert torch.equal(out, expected.to(torch.bfloat16))
vllm/config/kernel.py
+56 -1
View File
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import contextlib
from collections.abc import Callable
from collections.abc import Callable, Iterable
from dataclasses import asdict, fields
from typing import TYPE_CHECKING, Any, Literal
@@ -131,6 +131,8 @@ MoEBackend = Literal[
"flashinfer_cutedsl",
"flashinfer_b12x",
"b12x",
"flashinfer_moe_ep_mega_deep_gemm",
"flashinfer_moe_ep_mega_cutedsl",
"marlin",
"humming",
"triton_unfused",
@@ -140,6 +142,46 @@ MoEBackend = Literal[
"emulation",
]
# Backends that run the mega-MoE model path through the flashinfer moe_ep
# runtime. Only architectures in FLASHINFER_MOE_EP_ARCHITECTURES wire up
# these experts.
FLASHINFER_MOE_EP_BACKENDS = frozenset(
{
"flashinfer_moe_ep_mega_deep_gemm",
"flashinfer_moe_ep_mega_cutedsl",
}
)
# Backends that run a mega-MoE model path (fused expert module plus
# prepare_megamoe routing): vLLM's native deep_gemm path, which any model
# with a mega-MoE module may use (DeepSeek-V4, Kimi K3), plus the flashinfer
# moe_ep variants.
MEGA_MOE_BACKENDS = frozenset({"deep_gemm_mega_moe"}) | FLASHINFER_MOE_EP_BACKENDS
# Architectures whose model code wires up the flashinfer moe_ep experts. MTP
# and DSpark draft variants inherit the setting from these target models.
FLASHINFER_MOE_EP_ARCHITECTURES = frozenset(
{
"DeepseekV4ForCausalLM",
"DeepSeekV4MTPModel",
}
)
def validate_flashinfer_moe_ep_model(
moe_backend: str, architectures: Iterable[str]
) -> None:
"""Reject flashinfer moe_ep backends for models that lack the FI path."""
if moe_backend not in FLASHINFER_MOE_EP_BACKENDS:
return
if not any(arch in FLASHINFER_MOE_EP_ARCHITECTURES for arch in architectures):
raise ValueError(
f"moe_backend={moe_backend!r} is only supported for DeepSeek-V4 "
f"models ({sorted(FLASHINFER_MOE_EP_ARCHITECTURES)}), but the "
f"model is {list(architectures)}."
)
LinearBackend = Literal[
"auto",
"cutlass",
@@ -206,6 +248,14 @@ class KernelConfig:
- "flashinfer_b12x": Use FlashInfer CuteDSL fused MoE for SM12x
(RTX Pro 6000 / DGX Spark)
- "b12x": Use b12x FP4 MoE kernels on SM12x
- "flashinfer_moe_ep_mega_deep_gemm": Use the FlashInfer moe_ep
expert-parallel mega-MoE with the DeepGEMM megakernel, which consumes an
MXFP4 checkpoint verbatim (Blackwell, requires expert parallel;
DeepSeek-V4 only)
- "flashinfer_moe_ep_mega_cutedsl": Same, with the CuteDSL megakernel
(additionally requires NVSHMEM). The checkpoint selects the weight path:
an NVFP4 checkpoint is consumed prequantized, MXFP4 weights are
requantized at load
- "marlin": Use Marlin kernels (weight-only quantization)
- "humming": Use Humming Mixed Precision kernels
- "triton_unfused": Use Triton unfused MoE kernels
@@ -288,6 +338,11 @@ class KernelConfig:
"""Set platform-specific defaults for the kernel config."""
from vllm.platforms import current_platform
if vllm_config.model_config is not None:
validate_flashinfer_moe_ep_model(
self.moe_backend, vllm_config.model_config.architectures
)
platform_op_priority = current_platform.get_default_ir_op_priority(vllm_config)
logger.debug(
"Setting platform-specific IR op priority defaults: %s, user-defined: %s",
vllm/models/deepseek_v4/nvidia/fi_moe.py
+397
View File
@@ -0,0 +1,397 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""FlashInfer ``moe_ep`` expert module for DeepSeek V4."""
from __future__ import annotations
from typing import TYPE_CHECKING, Any
import torch
import torch.nn as nn
from vllm.model_executor.utils import set_weight_attrs
from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MegaMoEExperts
from vllm.utils.flashinfer_moe_ep import (
build_fi_mega_layer,
ensure_fi_moe_ep_runtime,
fi_moe_ep_backend_spec,
make_fi_moe_ep_bootstrap,
mega_moe_weight_pack_from_params,
)
if TYPE_CHECKING:
from flashinfer.moe_ep import MoEEpMegaLayer
from vllm.config import VllmConfig
from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MLP
_MOE_SKIP_PADDING: bool | None = None
def ckpt_uses_nvfp4_experts(vllm_config: VllmConfig) -> bool:
"""True when the loaded checkpoint quantizes experts with modelopt NVFP4
(e2m1 + fp8-e4m3 per-16 block scales + per-tensor weight_scale_2), i.e.
the recipe the nvfp4_cutedsl prequantized-weights path consumes verbatim.
DeepSeek-V4 NVFP4-expert checkpoints declare ``moe_quant_algo: NVFP4`` in
their ``quantization_config``, which DeepseekV4FP8Config surfaces as
``moe_quant_algo``.
"""
return getattr(vllm_config.quant_config, "moe_quant_algo", "") == "NVFP4"
def resolve_mega_moe_is_padding(num_tokens: int) -> torch.Tensor | None:
from vllm.forward_context import get_forward_context, is_forward_context_available
global _MOE_SKIP_PADDING
if _MOE_SKIP_PADDING is None:
import vllm.envs as envs
_MOE_SKIP_PADDING = bool(envs.VLLM_MOE_SKIP_PADDING)
if not _MOE_SKIP_PADDING or not is_forward_context_available():
return None
is_padding = get_forward_context().is_padding
if is_padding is None:
return None
return is_padding[:num_tokens]
def apply_mega_moe_routing_preprocess(
topk_ids: torch.Tensor,
*,
is_padding: torch.Tensor | None = None,
) -> torch.Tensor:
"""Padding-only routing preprocess (EPLB hooks go here later)."""
if is_padding is not None:
topk_ids = torch.where(is_padding.unsqueeze(1), -1, topk_ids)
return topk_ids
def nvfp4_prequant_pack_and_alphas(
w13_weight: torch.Tensor,
w13_weight_scale: torch.Tensor,
w13_weight_scale_2: torch.Tensor, # (E_local, 2) fp32: [:,0]=gate(w1), [:,1]=up(w3)
w2_weight: torch.Tensor,
w2_weight_scale: torch.Tensor,
w2_weight_scale_2: torch.Tensor, # (E_local,) fp32
*,
intermediate_size: int,
):
"""NVFP4 checkpoint params -> (MoEWeightPack, fc1_alpha, fc2_alpha).
Activation quantization is fully dynamic, so the checkpoint's static
``input_scale`` drops out and each GEMM's epilogue alpha reduces to the
weight's per-tensor ``weight_scale_2``.
fc1_alpha is one scalar per expert, but gate (w1) and up (w3) carry their
own ``weight_scale_2``. When they differ, the ratio is folded into the up
half's e4m3 block scales — only if the fold round-trips exactly, since a
lossy rescale would silently change the model.
"""
from flashinfer.moe_ep import MoEWeightPack
inter = intermediate_size
gate_s2 = w13_weight_scale_2[:, 0].float()
up_s2 = w13_weight_scale_2[:, 1].float()
if (gate_s2 <= 0).any() or (up_s2 <= 0).any() or (w2_weight_scale_2 <= 0).any():
raise ValueError(
"nvfp4 prequant: non-positive weight_scale_2 loaded — checkpoint "
"scale tensors missing or loader routed them wrong."
)
w13_scale = w13_weight_scale
if not torch.equal(gate_s2, up_s2):
ratio = up_s2 / gate_s2 # (E_local,)
up_sf = w13_scale[:, inter:, :].float()
folded = up_sf * ratio[:, None, None]
folded_e4m3 = folded.to(torch.float8_e4m3fn)
if not torch.equal(folded_e4m3.float(), folded):
raise ValueError(
"NVFP4 MoE cannot merge the gate and up weight_scale_2 values "
"because their ratio is not exactly representable in "
"float8_e4m3fn. Use a checkpoint quantized with a shared "
"gate/up weight_scale_2."
)
w13_scale = w13_scale.clone()
w13_scale[:, inter:, :] = folded_e4m3
fc1_alpha = gate_s2.clone().contiguous()
fc2_alpha = w2_weight_scale_2.float().clone().contiguous()
pack = MoEWeightPack(
w13=w13_weight,
w2=w2_weight,
w13_scale=w13_scale,
w2_scale=w2_weight_scale,
)
return pack, fc1_alpha, fc2_alpha
class DeepseekV4MegaMoEExpertsFI(DeepseekV4MegaMoEExperts):
"""Same weight layout/loader as the native mega experts, FI compute path."""
def __init__(
self,
vllm_config: VllmConfig,
*,
activation_clamp: float | None = None,
**kwargs: Any,
) -> None:
super().__init__(vllm_config, **kwargs)
self._vllm_config = vllm_config
self._activation_clamp = activation_clamp
self._mega_layer: MoEEpMegaLayer | None = None
self._fast_ctx: tuple[Any, Any, Any, int, bool] | None = None
self._epilogue_alphas: tuple[torch.Tensor, torch.Tensor] | None = None
self._nvfp4_prequant = ckpt_uses_nvfp4_experts(vllm_config)
if self._nvfp4_prequant:
megakernel = fi_moe_ep_backend_spec(
vllm_config.kernel_config.moe_backend
).megakernel
if megakernel != "nvfp4_cutedsl":
raise ValueError(
"NVFP4-quantized expert checkpoint requires "
"moe_backend=flashinfer_moe_ep_mega_cutedsl, got a "
f"backend using megakernel {megakernel!r} "
"(deep_gemm consumes the MXFP4 checkpoint instead)."
)
self._realloc_nvfp4_params()
def _realloc_nvfp4_params(self) -> None:
"""Swap the mx-recipe scale params for the NVFP4 checkpoint's:
fp8-e4m3 per-16 block scales plus the per-tensor second-level
scales (weight_scale_2, input_scale) modelopt exports."""
n_e = self.num_local_experts
inter = self.intermediate_size
hidden = self.hidden_size
attrs = {"weight_loader": self.weight_loader}
def _param(shape: tuple, dtype: torch.dtype) -> nn.Parameter:
p = nn.Parameter(torch.zeros(*shape, dtype=dtype), requires_grad=False)
set_weight_attrs(p, attrs)
return p
self.w13_weight_scale = _param(
(n_e, 2 * inter, hidden // 16), torch.float8_e4m3fn
)
self.w13_weight_scale.quant_method = "block"
self.w2_weight_scale = _param((n_e, hidden, inter // 16), torch.float8_e4m3fn)
self.w2_weight_scale.quant_method = "block"
# (E, 2): column 0 = gate (w1), column 1 = up (w3).
self.w13_weight_scale_2 = _param((n_e, 2), torch.float32)
self.w2_weight_scale_2 = _param((n_e,), torch.float32)
# Loaded for completeness; unused — activation quant is dynamic
# (see nvfp4_prequant_pack_and_alphas).
self.w13_input_scale = _param((n_e, 2), torch.float32)
self.w2_input_scale = _param((n_e,), torch.float32)
def weight_loader(
self,
param: nn.Parameter,
loaded_weight: torch.Tensor,
weight_name: str,
shard_id: str,
expert_id: int,
return_success: bool = False,
) -> bool | None:
# NVFP4 checkpoint second-level scalars route here; everything
# else (packed weights, block scales) matches the base layout.
if "weight_scale_2" in weight_name or "input_scale" in weight_name:
local_expert_ids = self._map_global_expert_id(expert_id)
if not local_expert_ids:
return False if return_success else None
value = loaded_weight.reshape(()).to(torch.float32)
for local_expert_id in local_expert_ids:
if shard_id in ("w1", "w3"):
if "w13_" not in weight_name:
return False if return_success else None
param.data[local_expert_id, 0 if shard_id == "w1" else 1] = value
elif shard_id == "w2":
if "w2_" not in weight_name:
return False if return_success else None
param.data[local_expert_id] = value
else:
raise ValueError(f"Unsupported expert shard id: {shard_id}")
return True if return_success else None
return super().weight_loader(
param,
loaded_weight,
weight_name,
shard_id,
expert_id,
return_success,
)
def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None:
# The FlashInfer megakernel has no shared-expert fusion; the caller's
# serial shared MLP path handles shared_experts.
if self._mega_layer is not None:
return
if self.w13_weight is None:
return
self._check_runtime_supported()
ensure_fi_moe_ep_runtime(self._vllm_config)
if self._nvfp4_prequant:
# NVFP4 checkpoint: hand the packed weights + both scale
# planes straight to the backend (no dequant->requant);
# per-expert globals become fc1/fc2 epilogue alphas staged
# at every forward via MoEEpTensors.
weights, fc1_alpha, fc2_alpha = nvfp4_prequant_pack_and_alphas(
self.w13_weight.data,
self.w13_weight_scale.data,
self.w13_weight_scale_2.data,
self.w2_weight.data,
self.w2_weight_scale.data,
self.w2_weight_scale_2.data,
intermediate_size=self.intermediate_size,
)
self._epilogue_alphas = (fc1_alpha, fc2_alpha)
else:
weights = mega_moe_weight_pack_from_params(
self.w13_weight,
self.w13_weight_scale,
self.w2_weight,
self.w2_weight_scale,
megakernel=fi_moe_ep_backend_spec(
self._vllm_config.kernel_config.moe_backend
).megakernel,
)
self._mega_layer = build_fi_mega_layer(
make_fi_moe_ep_bootstrap(),
vllm_config=self._vllm_config,
num_experts=self.num_experts,
max_tokens_per_rank=self.max_num_tokens,
hidden_size=self.hidden_size,
intermediate_size=self.intermediate_size,
top_k=self.top_k,
activation_clamp=self._activation_clamp,
weights=weights,
)
del weights
# Allocate (or attach to) the pooled workspace before first
# forward so warmup/capture never hits the lazy path.
self._mega_layer._ensure_workspace()
self.w13_weight = None
self.w13_weight_scale = None
self.w2_weight = None
self.w2_weight_scale = None
if self._nvfp4_prequant:
self.w13_weight_scale_2 = None
self.w2_weight_scale_2 = None
self.w13_input_scale = None
self.w2_input_scale = None
# EPLB is rejected for these backends in validate_fi_moe_ep_config, so
# nothing constructs an EplbState and none of these run. They are kept
# as explicit errors rather than deleted because inheriting the native
# implementations would be worse: set_eplb_state would record a map
# this forward path never consults, silently splitting routing from
# weights, and get_expert_weights reads transformed-weight attributes
# the flashinfer path releases after preprocess.
_EPLB_UNSUPPORTED = (
"EPLB is not supported with the flashinfer moe_ep backends; "
"validate_fi_moe_ep_config should have rejected this "
"configuration at startup."
)
def set_eplb_state(
self,
moe_layer_idx: int,
expert_load_view: torch.Tensor,
logical_to_physical_map: torch.Tensor,
logical_replica_count: torch.Tensor,
) -> None:
raise NotImplementedError(self._EPLB_UNSUPPORTED)
def get_expert_weights(self) -> list[torch.Tensor]:
raise NotImplementedError(self._EPLB_UNSUPPORTED)
def update_expert_map(self) -> None:
raise NotImplementedError(self._EPLB_UNSUPPORTED)
def forward(
self,
hidden_states: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
*,
activation_clamp: float | None,
fast_math: bool = True,
) -> torch.Tensor:
# fast_math is a native deep_gemm knob; the FI kernels have no
# equivalent toggle, so it is accepted for signature parity only.
if hidden_states.shape[0] > self.max_num_tokens:
raise ValueError(
f"DeepSeek V4 MegaMoE got {hidden_states.shape[0]} tokens, "
f"but the symmetric buffer was sized for {self.max_num_tokens}."
)
from flashinfer.moe_ep import MoEEpTensors
num_tokens = hidden_states.shape[0]
is_padding = resolve_mega_moe_is_padding(num_tokens)
topk_ids = apply_mega_moe_routing_preprocess(
topk_ids,
is_padding=is_padding,
)
# Fast path: after the first successful full forward the layer is
# immutable, so skip MoEEpMegaLayer.forward()'s per-call validation
# and go straight to the kernel backend's stage_inputs + compute.
alphas = self._epilogue_alphas
fc1_alpha = alphas[0] if alphas is not None else None
fc2_alpha = alphas[1] if alphas is not None else None
fast = self._fast_ctx
if fast is not None:
kernel, workspace, transformed, hidden_size, zero_copy = fast
t = MoEEpTensors(
hidden_states=hidden_states,
topk_ids=topk_ids,
topk_weights=topk_weights,
fc1_alpha=fc1_alpha,
fc2_alpha=fc2_alpha,
)
kernel.stage_inputs(t, workspace, quantize_input=True)
if zero_copy:
# Zero-copy (cutedsl backends): consume the workspace [:n]
# view directly — valid under stream ordering until the next
# MoE layer's launch on the shared workspace.
return kernel.compute(workspace, transformed, output=None)
# deep_gemm_mega's compute() requires a real output tensor.
out = torch.empty(
num_tokens,
hidden_size,
dtype=torch.bfloat16,
device=hidden_states.device,
)
return kernel.compute(workspace, transformed, output=out)
ensure_fi_moe_ep_runtime(self._vllm_config)
self.finalize_weights()
assert self._mega_layer is not None
y = self._mega_layer.forward(
MoEEpTensors(
hidden_states=hidden_states,
topk_ids=topk_ids,
topk_weights=topk_weights,
fc1_alpha=fc1_alpha,
fc2_alpha=fc2_alpha,
)
)
layer = self._mega_layer
if hidden_states.dtype == torch.bfloat16:
self._fast_ctx = (
layer._kernel,
layer._ensure_workspace(),
layer._transformed,
layer._fleet_params.token_hidden_size,
# zero-copy output views are a cutedsl-backend contract
layer._kernel.kernel_name() != "deep_gemm_mega",
)
return y
DeepseekV4MegaMoEExpertsFI.weight_loader.supports_moe_loading = True # type: ignore[attr-defined]
vllm/models/deepseek_v4/nvidia/model.py
+32 -11
View File
@@ -11,6 +11,7 @@ import torch.nn as nn
import vllm.envs as envs
from vllm.config import VllmConfig
from vllm.config.kernel import MEGA_MOE_BACKENDS
from vllm.distributed import (
get_ep_group,
get_pp_group,
@@ -83,6 +84,10 @@ from vllm.models.deepseek_v4.nvidia.flashmla import DeepseekV4FlashMLAAttention
from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
from vllm.platforms import current_platform
from vllm.sequence import IntermediateTensors
from vllm.utils.flashinfer_moe_ep import (
is_fi_moe_ep_backend,
validate_fi_moe_ep_config,
)
from vllm.utils.math_utils import cdiv
from vllm.v1.attention.backends.registry import AttentionBackendEnum
from vllm.v1.worker.ubatching import dbo_current_ubatch_id
@@ -752,9 +757,10 @@ class DeepseekV4MoE(nn.Module):
quant_config = vllm_config.quant_config
self.prefix = prefix
self.use_sequence_parallel = use_sequence_parallel
self.use_mega_moe = (
vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
)
moe_backend = vllm_config.kernel_config.moe_backend
validate_fi_moe_ep_config(vllm_config)
self.use_mega_moe = moe_backend in MEGA_MOE_BACKENDS
self.use_fi_mega_moe = is_fi_moe_ep_backend(moe_backend)
if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:
raise NotImplementedError(
"DeepSeek V4 MegaMoE currently requires expert parallel. "
@@ -779,7 +785,7 @@ class DeepseekV4MoE(nn.Module):
raise NotImplementedError(
"DeepSeek V4 MegaMoE only supports fp4 experts; got expert_dtype="
f"{config.expert_dtype!r}. Drop --kernel-config moe_backend="
"deep_gemm_mega_moe for this checkpoint."
f"{moe_backend} for this checkpoint."
)
self.gate = GateLinear(
@@ -867,15 +873,29 @@ class DeepseekV4MoE(nn.Module):
# Native DeepGEMM fusion requires each EP rank to own the complete
# shared MLP. Sequence parallel replicates those weights while sharding
# tokens. TP=1 is also naturally replicated. With PP+TP the shared MLP
# remains tensor-sharded, so retain the serial path.
# remains tensor-sharded, so retain the serial path. The FlashInfer
# megakernel has no shared-expert fusion, so it keeps the serial path.
fuse_shared_experts = bool(
self.shared_experts is not None
and not envs.VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION
and (self.use_sequence_parallel or self.tp_size == 1)
and not self.use_fi_mega_moe
)
self.experts = DeepseekV4MegaMoEExperts(
vllm_config,
activation_clamp = (
float(self.swiglu_limit) if self.swiglu_limit is not None else None
)
if self.use_fi_mega_moe:
# Deferred: fi_moe subclasses DeepseekV4MegaMoEExperts, so a
# module-level import here would be circular.
from vllm.models.deepseek_v4.nvidia.fi_moe import (
DeepseekV4MegaMoEExpertsFI,
)
experts_cls: type[DeepseekV4MegaMoEExperts] = DeepseekV4MegaMoEExpertsFI
else:
experts_cls = DeepseekV4MegaMoEExperts
expert_kwargs: dict[str, typing.Any] = dict(
num_experts=self.n_physical_experts,
num_local_experts=self.n_local_physical_experts,
experts_start_idx=self.physical_expert_start,
@@ -886,6 +906,9 @@ class DeepseekV4MoE(nn.Module):
num_shared_experts=(self.n_shared_experts if fuse_shared_experts else 0),
prefix=f"{prefix}.experts",
)
if self.use_fi_mega_moe:
expert_kwargs["activation_clamp"] = activation_clamp
self.experts = experts_cls(vllm_config, **expert_kwargs)
def _init_fused_moe_experts(
self,
@@ -1031,7 +1054,7 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]:
def _use_sequence_parallel(vllm_config: VllmConfig) -> bool:
parallel_config = vllm_config.parallel_config
use_mega_moe = vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS
return (
parallel_config.pipeline_parallel_size == 1
and parallel_config.enable_expert_parallel
@@ -1226,9 +1249,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
self.config = config
self.quant_config = quant_config
self.parallel_config = vllm_config.parallel_config
self.use_mega_moe = (
vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
)
self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS
self.use_sequence_parallel = _use_sequence_parallel(vllm_config)
if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:
raise NotImplementedError(
vllm/utils/flashinfer_moe_ep.py
+339
View File
@@ -0,0 +1,339 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""FlashInfer ``moe_ep`` helpers for DeepSeek V4 vLLM integration."""
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
import torch
import torch.nn as nn
from vllm.config.kernel import FLASHINFER_MOE_EP_BACKENDS
from vllm.distributed import get_ep_group
from vllm.platforms import current_platform
if TYPE_CHECKING:
from flashinfer.moe_ep import BootstrapConfig, MoEEpMegaLayer
from vllm.config import VllmConfig
# Every mega kernel is Blackwell-only, so the arch is a property of the family
# rather than of a particular backend: validate it against the live device
# instead of encoding it in the backend name. An explicit allowlist so new
# archs (SM110/SM120, Rubin) fail loudly until flashinfer supports them.
FI_MOE_EP_SUPPORTED_CAPABILITIES = frozenset({(10, 0), (10, 3)})
@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
FI_MOE_EP_BACKEND_SPECS: dict[str, FiMoeEpBackendSpec] = {
# Consumes an MXFP4 checkpoint verbatim (e2m1 weights + E8M0 per-32
# scales) -- the same recipe the native deep_gemm mega path uses.
"flashinfer_moe_ep_mega_deep_gemm": FiMoeEpBackendSpec(
megakernel="deep_gemm_mega",
needs_nvshmem=False,
),
# The checkpoint picks the weight path, not the kernel: an NVFP4
# checkpoint is consumed prequantized (no round trip), while MXFP4 weights
# are dequantized to bf16 and requantized. See
# models/deepseek_v4/nvidia/fi_moe.py:ckpt_uses_nvfp4_experts.
"flashinfer_moe_ep_mega_cutedsl": FiMoeEpBackendSpec(
megakernel="nvfp4_cutedsl",
needs_nvshmem=True,
),
}
assert set(FI_MOE_EP_BACKEND_SPECS) == FLASHINFER_MOE_EP_BACKENDS
_FI_RUNTIME_HANDLE: Any = None
def is_fi_moe_ep_backend(moe_backend: str) -> bool:
return moe_backend in FI_MOE_EP_BACKEND_SPECS
def fi_moe_ep_backend_spec(moe_backend: str) -> FiMoeEpBackendSpec:
try:
return FI_MOE_EP_BACKEND_SPECS[moe_backend]
except KeyError:
raise ValueError(
f"{moe_backend!r} is not a flashinfer moe_ep backend; expected "
f"one of {sorted(FI_MOE_EP_BACKEND_SPECS)}"
) from None
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."
)
def make_fi_moe_ep_bootstrap() -> BootstrapConfig:
from flashinfer.moe_ep import BootstrapConfig
ep = get_ep_group()
return BootstrapConfig(
world_size=ep.world_size,
rank=ep.rank_in_group,
process_group=ep.device_group,
auto_bootstrap=False,
# vLLM already bound this worker's (possibly remapped) device;
# without this the runtime would rebind to cuda:LOCAL_RANK|rank and
# launch weight transforms against another device's pointers.
device=torch.accelerator.current_device_index(),
)
def megakernel_runtime_requirements(spec: FiMoeEpBackendSpec) -> frozenset[str]:
from flashinfer.moe_ep.core.runtime import NVSHMEM, TORCH_DIST
if spec.needs_nvshmem:
return frozenset({TORCH_DIST, NVSHMEM})
return frozenset({TORCH_DIST})
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),
)
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
_E2M1_LUT = (
0.0,
0.5,
1.0,
1.5,
2.0,
3.0,
4.0,
6.0,
-0.0,
-0.5,
-1.0,
-1.5,
-2.0,
-3.0,
-4.0,
-6.0,
)
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)
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
def mega_moe_weight_pack_from_params(
w13_weight: nn.Parameter,
w13_weight_scale: nn.Parameter,
w2_weight: nn.Parameter,
w2_weight_scale: nn.Parameter,
*,
megakernel: str = "deep_gemm_mega",
):
from flashinfer.moe_ep import MoEWeightPack
if megakernel == "deep_gemm_mega":
# Same fp4-e2m1 + ue8m0-per-32 recipe as the native path: pass verbatim,
# flashinfer runs the identical deep_gemm transform.
return MoEWeightPack(
w13=w13_weight.data,
w2=w2_weight.data,
w13_scale=w13_weight_scale.data,
w2_scale=w2_weight_scale.data,
)
# The cutedsl kernel quantizes with its own recipe (nvfp4
# e2m1+e4m3-per-16): dequantize the checkpoint fp4 to bf16 and let the
# backend preprocess requantize. Double quantization: outputs are close to
# but not bit-identical with the native path.
return MoEWeightPack(
w13=_dequant_expert_weights_to_bf16(w13_weight.data, w13_weight_scale.data),
w2=_dequant_expert_weights_to_bf16(w2_weight.data, w2_weight_scale.data),
)
def build_fi_mega_config(
*,
intermediate_size: int,
top_k: int,
activation_clamp: float | None,
megakernel: str,
):
from flashinfer.moe_ep import (
DeepGemmMegaMoeConfig,
MegaConfig,
Nvfp4CutedslMegaMoeConfig,
)
# fast_math selects approximate exp/rcp in DeepGEMM's fused SwiGLU
# epilogue; the cutedsl kernels accept it for API parity only.
if megakernel == "deep_gemm_mega":
mk = DeepGemmMegaMoeConfig(
intermediate_size=intermediate_size,
top_k=top_k,
activation_clamp=activation_clamp,
fast_math=True,
)
elif megakernel == "nvfp4_cutedsl":
mk = Nvfp4CutedslMegaMoeConfig(
intermediate_size=intermediate_size,
top_k=top_k,
activation_clamp=activation_clamp,
fast_math=True,
)
else:
raise ValueError(f"Unsupported fi_moe_ep megakernel {megakernel!r}")
return MegaConfig(
megakernel=mk,
preprocess_weights=True,
quantize_input=True,
)
# All MoE layers share one symmetric workspace, like the native path's
# class-level DeepseekV4MegaMoEExperts._symm_buffer_cache. Without this the
# fi path allocates one symm buffer PER LAYER (43x memory + cold working
# sets); the workspace is stateless across forwards (kernel tail-cleans) and
# layers execute sequentially on one stream, so sharing is safe.
def build_fi_mega_layer(
bootstrap: BootstrapConfig,
*,
vllm_config: VllmConfig,
num_experts: int,
max_tokens_per_rank: int,
hidden_size: int,
intermediate_size: int,
top_k: int,
activation_clamp: float | None,
weights,
) -> MoEEpMegaLayer:
from flashinfer.moe_ep import FleetParams, MoEEpLayer
megakernel = fi_moe_ep_backend_spec(
vllm_config.kernel_config.moe_backend
).megakernel
mega_config = build_fi_mega_config(
intermediate_size=intermediate_size,
top_k=top_k,
activation_clamp=activation_clamp,
megakernel=megakernel,
)
layer = MoEEpLayer(
bootstrap=bootstrap,
fleet_params=FleetParams(
num_experts=num_experts,
max_tokens_per_rank=max_tokens_per_rank,
token_hidden_size=hidden_size,
),
weights=weights,
backend=mega_config,
)
from flashinfer.moe_ep import MoEEpMegaLayer
if not isinstance(layer, MoEEpMegaLayer):
raise TypeError(
f"fi_moe_ep expected MoEEpMegaLayer, got {type(layer).__name__}"
)
return layer
__all__ = [
"FI_MOE_EP_BACKEND_SPECS",
"FiMoeEpBackendSpec",
"build_fi_mega_config",
"build_fi_mega_layer",
"ensure_fi_moe_ep_runtime",
"fi_moe_ep_backend_spec",
"finalize_fi_moe_ep_runtime",
"is_fi_moe_ep_backend",
"make_fi_moe_ep_bootstrap",
"mega_moe_weight_pack_from_params",
"megakernel_runtime_requirements",
"validate_fi_moe_ep_config",
]