diff --git a/tests/models/test_deepseek_v4_fi_moe_ep.py b/tests/models/test_deepseek_v4_fi_moe_ep.py new file mode 100644 index 0000000000..77b05101d6 --- /dev/null +++ b/tests/models/test_deepseek_v4_fi_moe_ep.py @@ -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)) diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 8bd8830570..2716319e39 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -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", diff --git a/vllm/models/deepseek_v4/nvidia/fi_moe.py b/vllm/models/deepseek_v4/nvidia/fi_moe.py new file mode 100644 index 0000000000..2e94e75e7d --- /dev/null +++ b/vllm/models/deepseek_v4/nvidia/fi_moe.py @@ -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] diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 095d8bc741..a469526316 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -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( diff --git a/vllm/utils/flashinfer_moe_ep.py b/vllm/utils/flashinfer_moe_ep.py new file mode 100644 index 0000000000..5bdd5cbf01 --- /dev/null +++ b/vllm/utils/flashinfer_moe_ep.py @@ -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", +]