diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index 0b1a9151ec..6c71aa1f46 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -37,6 +37,7 @@ th { | deepep_low_latency | batched | fp8 | G(128),A,T3 | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] | | flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] | | flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8,fp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] | +| passthrough | standard | nvfp4,mxfp4 | G | N | Y | [`PassThroughPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.passthrough.PassThroughPrepareAndFinalize] | !!! info "Table key" 1. All types: mxfp4, nvfp4, int4, int8, fp8 diff --git a/docs/serving/expert_parallel_deployment.md b/docs/serving/expert_parallel_deployment.md index ea60233645..f3b3296102 100644 --- a/docs/serving/expert_parallel_deployment.md +++ b/docs/serving/expert_parallel_deployment.md @@ -27,6 +27,7 @@ vLLM provides multiple communication backends for EP. Use `--all2all-backend` to | `deepep_high_throughput` | Multi-node prefill | Grouped GEMM with continuous layout, optimized for prefill | Prefill-dominated workloads, high-throughput scenarios | | `deepep_low_latency` | Multi-node decode | CUDA graph support, masked layout, optimized for decode | Decode-dominated workloads, low-latency scenarios | | `flashinfer_nvlink_one_sided` | MNNVL systems | FlashInfer's one-sided A2A strategy for multi-node NVLink | High-throughput workloads | +| `passthrough` | Any | Internal, not accepted on the CLI: bound automatically for MoE backends that dispatch and combine themselves, so no all-to-all runs in the framework | `--moe-backend flashinfer_moe_ep_*` (FlashInfer MoE-EP megakernels) | | `flashinfer_nvlink_two_sided` | MNNVL systems | FlashInfer's two-sided A2A strategy for multi-node NVLink | Systems with NVLink across nodes | ## Single Node Deployment diff --git a/tests/distributed/test_engram_dp_shard.py b/tests/distributed/test_engram_dp_shard.py index 0158b342f3..ac73188534 100644 --- a/tests/distributed/test_engram_dp_shard.py +++ b/tests/distributed/test_engram_dp_shard.py @@ -511,7 +511,7 @@ def _check_dummy_hash_model_forward( model = model_ops.DeepseekV4Model.__new__(model_ops.DeepseekV4Model) torch.nn.Module.__init__(model) - model.use_mega_moe = False + model.use_native_mega_moe = False model.use_sequence_parallel = False model.fuse_mhc_all_reduce = False model.engram_hash = state diff --git a/tests/kernels/mhc/test_mhc_kernels.py b/tests/kernels/mhc/test_mhc_kernels.py index 0554e9ca66..00e9834d62 100644 --- a/tests/kernels/mhc/test_mhc_kernels.py +++ b/tests/kernels/mhc/test_mhc_kernels.py @@ -167,7 +167,7 @@ def test_v41_dspark_head_collapses_with_last_ffn_mix(num_tokens, monkeypatch): monkeypatch.setattr(dspark, "mhc_post_tilelang", lambda *args: streams) draft = SimpleNamespace( - use_mega_moe=False, + use_native_mega_moe=False, use_sequence_parallel=False, hc_mult=hc_mult, layers=[make_layer(mix) for mix in mixes], diff --git a/tests/kernels/moe/test_flashinfer_moe_ep.py b/tests/kernels/moe/test_flashinfer_moe_ep.py new file mode 100644 index 0000000000..c3223b7d98 --- /dev/null +++ b/tests/kernels/moe/test_flashinfer_moe_ep.py @@ -0,0 +1,307 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end tests for the FlashInfer MoE-EP adapter on real Blackwell GPUs. + +Each test spawns one process per GPU, bootstraps vLLM's expert-parallel +groups, builds a ``FlashInferMoeEp`` adapter from bf16 expert weights at the +MoE geometry of a real model, warms it up and runs a forward whose inputs are +chosen so that every quantization step in the megakernel is exact, for the +CuTeDSL kernel (NVFP4 activations and weights, E4M3 scales per 16) as well as +the DeepGEMM kernel (FP8 activations and FP4 weights, power-of-two scales per +32). The expected output is then known in closed form and compared bit for +bit. +""" + +from dataclasses import dataclass + +import pytest +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from tests.distributed.eplb_utils import distributed_run, set_env_vars_and_device +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.config.kernel import FLASHINFER_MOE_EP_CUTEDSL, FLASHINFER_MOE_EP_DEEP_GEMM +from vllm.distributed.parallel_state import ensure_model_parallel_initialized +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, +) +from vllm.model_executor.layers.fused_moe.experts.flashinfer_moe_ep import ( + FlashInferMoeEpExperts, +) +from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep import ( + FlashInferMoeEp, + FlashInferMoeEpWeights, + apply_topk_in_fc1, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Static +from vllm.platforms import current_platform +from vllm.utils.import_utils import has_deep_gemm + + +@dataclass(frozen=True) +class Geometry: + name: str + hidden_size: int + intermediate_size: int + num_experts: int + top_k: int + + +# Routed-expert geometry of the models the branch was evaluated on. DeepGEMM +# keeps one UE8M0 scale byte per 32 elements and requires the per-token scale +# rows to be TMA aligned, so hidden and intermediate sizes must be multiples +# of 512 (FlashInfer's own validation only checks 128); all four qualify. +GEOMETRIES = ( + Geometry("dsv4-pro", 7168, 3072, 384, 6), + Geometry("dsv4-flash", 4096, 2048, 256, 6), + Geometry("mistral-small-4", 4096, 2048, 128, 4), + Geometry("mistral-large-3", 7168, 4096, 128, 4), +) +BACKENDS = (FLASHINFER_MOE_EP_CUTEDSL, FLASHINFER_MOE_EP_DEEP_GEMM) + +NUM_TOKENS_PER_RANK = 4 +MAX_TOKENS_PER_RANK = 128 + +# Every block scale below is either 1 or a power of two, so the FP4 codes +# and the FP8 values survive quantization exactly: NVFP4 scales a 16-block by +# absmax / 6 in E4M3, DeepGEMM scales a 32-block by absmax / 6 (FP4) or +# absmax / 448 (FP8) rounded up to a power of two. A 6.0 in a block pins the +# FP4 scale to 1; 2.0, 6.0 and 144 are exact FP8 values at any such scale. +# +# Token t carries X_SIGNAL at hidden column 16 * t and a 6.0 sentinel next to +# it; every weight column the sentinel meets is zero. +X_SIGNAL = 2.0 +X_SENTINEL = 6.0 +# Gate and up rows are identical, so h = silu(12) * 12 = 144 = 6 * 24 +# whichever half the kernel treats as the gate, and 144 requantizes exactly. +W13_DIAG = 6.0 +H_SIGNAL = 144.0 +# Expert e writes E2M1_CODES[e % 7] * H_SIGNAL into output channel e. The 6.0 +# sentinel sits next to a zero intermediate channel and pins the block scale. +W2_SENTINEL = 6.0 +E2M1_CODES = (0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + + +def _expert_out(expert: int) -> float: + return E2M1_CODES[expert % len(E2M1_CODES)] + + +def _topk_weights(top_k: int) -> tuple[float, ...]: + # The DeepSeek-V4 routing policy applies the router weight before fc2, so + # weight * 144 must requantize exactly too: powers of two do. + return tuple(2.0 ** -(k + 1) for k in range(top_k)) + + +def _routing_method(backend: str) -> RoutingMethodType: + if backend == FLASHINFER_MOE_EP_DEEP_GEMM: + return RoutingMethodType.DeepseekV4 + return RoutingMethodType.Default + + +def _routing( + geometry: Geometry, rank: int, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor]: + """top_k distinct experts per token, strided across ranks so dispatch goes + remote.""" + base = ( + torch.arange(NUM_TOKENS_PER_RANK, device=device) + rank * NUM_TOKENS_PER_RANK + ) % geometry.num_experts + stride = geometry.num_experts // geometry.top_k + offsets = torch.arange(geometry.top_k, device=device) * stride + topk_ids = ((base[:, None] + offsets[None, :]) % geometry.num_experts).to( + torch.int32 + ) + topk_weights = torch.tensor( + [_topk_weights(geometry.top_k)] * NUM_TOKENS_PER_RANK, + dtype=torch.float32, + device=device, + ) + return topk_ids, topk_weights + + +def _hidden_states(geometry: Geometry, device: torch.device) -> torch.Tensor: + x = torch.zeros( + NUM_TOKENS_PER_RANK, geometry.hidden_size, dtype=torch.bfloat16, device=device + ) + for t in range(NUM_TOKENS_PER_RANK): + x[t, 16 * t] = X_SIGNAL + x[t, 16 * t + 1] = X_SENTINEL + return x + + +def _local_expert_weights( + geometry: Geometry, rank: int, num_local_experts: int, device: torch.device +) -> FlashInferMoeEpWeights: + w13 = torch.zeros( + num_local_experts, + 2 * geometry.intermediate_size, + geometry.hidden_size, + dtype=torch.bfloat16, + device=device, + ) + w2 = torch.zeros( + num_local_experts, + geometry.hidden_size, + geometry.intermediate_size, + dtype=torch.bfloat16, + device=device, + ) + token_columns = [16 * t for t in range(NUM_TOKENS_PER_RANK)] + for local in range(num_local_experts): + expert = rank * num_local_experts + local + w13[local, 0, token_columns] = W13_DIAG + w13[local, geometry.intermediate_size, token_columns] = W13_DIAG + w2[local, expert, 0] = _expert_out(expert) + w2[local, expert, 2] = W2_SENTINEL + return FlashInferMoeEpWeights(w13=w13, w2=w2) + + +def _expected_output( + geometry: Geometry, topk_ids: torch.Tensor, topk_weights: torch.Tensor +) -> torch.Tensor: + expected = torch.zeros( + NUM_TOKENS_PER_RANK, + geometry.hidden_size, + dtype=torch.float32, + device=topk_ids.device, + ) + for t in range(NUM_TOKENS_PER_RANK): + for k in range(geometry.top_k): + expert = int(topk_ids[t, k]) + weight = float(topk_weights[t, k]) + expected[t, expert] += weight * _expert_out(expert) * H_SIGNAL + return expected.to(torch.bfloat16) + + +def _moe_config( + vllm_config: VllmConfig, + geometry: Geometry, + world_size: int, + device: torch.device, + backend: str, +) -> FusedMoEConfig: + parallel = FusedMoEParallelConfig.make( + tp_size_=1, + pcp_size_=1, + dp_size_=world_size, + sp_size_=1, + vllm_parallel_config=vllm_config.parallel_config, + ) + return FusedMoEConfig( + num_experts=geometry.num_experts, + experts_per_token=geometry.top_k, + hidden_dim=geometry.hidden_size, + intermediate_size=geometry.intermediate_size, + num_local_experts=geometry.num_experts // world_size, + num_logical_experts=geometry.num_experts, + activation=MoEActivation.SILU, + device=device, + routing_method=_routing_method(backend), + moe_parallel_config=parallel, + in_dtype=torch.bfloat16, + moe_backend=backend, + max_num_tokens=MAX_TOKENS_PER_RANK, + ) + + +def _check_expert_views_alias_kernel_memory( + adapter: FlashInferMoeEp, + geometry: Geometry, + rank: int, + num_local_experts: int, + device: torch.device, +) -> None: + """EPLB moves experts through ``expert_weight_views``; the kernel must see it. + + Swap local experts 0 and 1 in place on every rank, then route through the + swapped physical slots: the output must be the one the unswapped logical + experts would have produced. + """ + for view in adapter.expert_weight_views(): + first = view[0].clone() + view[0].copy_(view[1]) + view[1].copy_(first) + torch.distributed.barrier() + + topk_ids, topk_weights = _routing(geometry, rank, device) + slot = topk_ids % num_local_experts + physical_ids = torch.where(slot == 0, topk_ids + 1, topk_ids) + physical_ids = torch.where(slot == 1, topk_ids - 1, physical_ids) + output = adapter(_hidden_states(geometry, device), physical_ids, topk_weights) + torch.accelerator.synchronize() + assert torch.equal(output, _expected_output(geometry, topk_ids, topk_weights)) + + +def _init_warmup_forward( + env: dict[str, str], world_size: int, geometry: Geometry, backend: str +) -> None: + set_env_vars_and_device(env) + vllm_config = VllmConfig() + vllm_config.parallel_config.data_parallel_size = world_size + vllm_config.parallel_config.enable_expert_parallel = True + vllm_config.kernel_config.moe_backend = backend + + with set_current_vllm_config(vllm_config): + ensure_model_parallel_initialized( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + rank = torch.distributed.get_rank() + device = torch.device(f"cuda:{rank}") + moe = _moe_config(vllm_config, geometry, world_size, device, backend) + supported, reason = FlashInferMoeEpExperts.is_supported_config( + FlashInferMoeEpExperts, + moe, + kMxfp4Static, + None, + mk.FusedMoEActivationFormat.Standard, + ) + assert supported, reason + + adapter = FlashInferMoeEp( + moe, + _local_expert_weights(geometry, rank, moe.num_local_experts, device), + apply_topk_in_fc1=apply_topk_in_fc1(moe), + ) + try: + adapter.warmup() + + topk_ids, topk_weights = _routing(geometry, rank, device) + output = adapter(_hidden_states(geometry, device), topk_ids, topk_weights) + torch.accelerator.synchronize() + + expected = _expected_output(geometry, topk_ids, topk_weights) + assert output.dtype == expected.dtype and output.shape == expected.shape + assert torch.equal(output, expected), ( + f"rank {rank}: max abs diff " + f"{(output.float() - expected.float()).abs().max().item()}, " + f"nonzero output channels per token " + f"{[torch.nonzero(row).flatten().tolist() for row in output]}" + ) + + if backend == FLASHINFER_MOE_EP_CUTEDSL: + _check_expert_views_alias_kernel_memory( + adapter, geometry, rank, moe.num_local_experts, device + ) + finally: + adapter.destroy() + + +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="The FlashInfer MoE-EP megakernels require SM100", +) +@pytest.mark.parametrize("geometry", GEOMETRIES, ids=[g.name for g in GEOMETRIES]) +@pytest.mark.parametrize("backend", BACKENDS, ids=("cutedsl", "deep_gemm")) +@pytest.mark.parametrize("world_size", [2, 4]) +def test_init_warmup_forward_is_exact_on_grid_inputs( + world_size: int, backend: str, geometry: Geometry +): + pytest.importorskip("flashinfer.moe_ep") + if backend == FLASHINFER_MOE_EP_DEEP_GEMM and not has_deep_gemm(): + pytest.skip("DeepGEMM is not available") + if torch.accelerator.device_count() < world_size: + pytest.skip(f"Need at least {world_size} GPUs") + distributed_run(_init_warmup_forward, world_size, geometry, backend) diff --git a/tests/models/test_deepseek_v4_fi_moe_ep.py b/tests/models/test_deepseek_v4_fi_moe_ep.py index 99d661e755..7687a6786f 100644 --- a/tests/models/test_deepseek_v4_fi_moe_ep.py +++ b/tests/models/test_deepseek_v4_fi_moe_ep.py @@ -1,116 +1,90 @@ # 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 +from dataclasses import replace +from types import SimpleNamespace import pytest import torch +from tests.kernels.moe.utils import make_dummy_moe_config from vllm.config.kernel import ( FLASHINFER_MOE_EP_BACKENDS, + FLASHINFER_MOE_EP_CUTEDSL, + FLASHINFER_MOE_EP_DEEP_GEMM, MEGA_MOE_BACKENDS, + NATIVE_MEGA_MOE_BACKENDS, + PASSTHROUGH_ALL2ALL_BACKEND, + bind_passthrough_all2all_backend, 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, +from vllm.engine.arg_utils import EngineArgs +from vllm.model_executor.layers.fused_moe import flashinfer_moe_ep as fi_ep +from vllm.model_executor.layers.fused_moe import modular_kernel as mk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + maybe_make_prepare_finalize, ) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEQuantConfig, + RoutingMethodType, + nvfp4_moe_quant_config, +) +from vllm.model_executor.layers.fused_moe.experts.flashinfer_moe_ep import ( + FlashInferMoeEpExperts, + epilogue_from_quant_config, +) +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + FLASHINFER_MOE_EP_MXFP4_BACKENDS, + Mxfp4MoeBackend, + make_mxfp4_moe_quant_config, + map_mxfp4_backend, + mxfp4_round_up_hidden_size_and_intermediate_size, +) +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + backend_to_kernel_cls as mxfp4_backend_to_kernel_cls, +) +from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( + NvFp4MoeBackend, + convert_to_nvfp4_moe_kernel_format, + map_nvfp4_backend, +) +from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( + backend_to_kernel_cls as nvfp4_backend_to_kernel_cls, +) +from vllm.model_executor.layers.fused_moe.prepare_finalize.passthrough import ( + PassThroughPrepareAndFinalize, +) +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp4Static, + kNvfp4Dynamic, + kNvfp4Static, +) +from vllm.utils.argparse_utils import FlexibleArgumentParser -@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) +def test_flashinfer_backends_are_megakernels_outside_the_native_model_path(): + assert FLASHINFER_MOE_EP_BACKENDS <= MEGA_MOE_BACKENDS + assert FLASHINFER_MOE_EP_BACKENDS.isdisjoint(NATIVE_MEGA_MOE_BACKENDS) -@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.""" +def test_only_deep_gemm_backend_is_dsv4_specific(): + validate_flashinfer_moe_ep_model( + FLASHINFER_MOE_EP_CUTEDSL, + ["MixtralForCausalLM"], + ) 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"]) + validate_flashinfer_moe_ep_model( + FLASHINFER_MOE_EP_DEEP_GEMM, + ["MixtralForCausalLM"], + ) + validate_flashinfer_moe_ep_model( + FLASHINFER_MOE_EP_DEEP_GEMM, + ["DeepseekV4ForCausalLM"], + ) @pytest.mark.parametrize( @@ -118,8 +92,7 @@ def test_fi_moe_ep_backend_accepted_for_dsv4(moe_backend): [["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.""" + """Models validate native deep_gemm mega constraints at construction time.""" validate_flashinfer_moe_ep_model("deep_gemm_mega_moe", architectures) @@ -128,11 +101,7 @@ def test_non_fi_backend_ignores_architectures(): @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.""" +def test_token_sharding_backends_enable_dsv4_sequence_parallel(moe_backend: str): from vllm.models.deepseek_v4.nvidia.model import _use_sequence_parallel vllm_config = SimpleNamespace( @@ -147,113 +116,262 @@ def test_all_mega_backends_get_sequence_parallel_moe(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_dsv4_requests_pre_fc2_router_weight_placement(): + dsv4 = SimpleNamespace(routing_method=RoutingMethodType.DeepseekV4) + default = SimpleNamespace(routing_method=RoutingMethodType.Default) + assert fi_ep.apply_topk_in_fc1(dsv4) + assert not fi_ep.apply_topk_in_fc1(default) -def test_megakernel_runtime_requirements(fake_flashinfer): - dg = megakernel_runtime_requirements( - fi_moe_ep_backend_spec("flashinfer_moe_ep_mega_deep_gemm") +def test_backend_specs_preserve_weight_format_contracts(): + cutedsl = fi_ep.flashinfer_moe_ep_backend_spec(FLASHINFER_MOE_EP_CUTEDSL) + assert cutedsl.kernel == "cutedsl" + assert cutedsl.weight_formats == frozenset({"nvfp4", "mxfp4"}) + + deep_gemm = fi_ep.flashinfer_moe_ep_backend_spec(FLASHINFER_MOE_EP_DEEP_GEMM) + assert deep_gemm.kernel == "deep_gemm" + assert deep_gemm.weight_formats == frozenset({"mxfp4"}) + + +def _megakernel_moe(moe_backend: str, **overrides) -> SimpleNamespace: + fields = dict( + moe_backend=moe_backend, + moe_parallel_config=SimpleNamespace( + use_ep=True, use_batched_activation_format=False + ), + is_act_and_mul=True, + activation=MoEActivation.SILU, + has_hash_routing=False, + routing_method=RoutingMethodType.DeepseekV4, + router_logits_dtype=torch.float32, + hidden_dim=256, + is_lora_enabled=False, + skip_final_all_reduce=False, + in_dtype=torch.bfloat16, + has_bias=False, + swiglu_alpha=None, + swiglu_beta=None, ) - assert dg == frozenset({"torch_dist"}) + fields.update(overrides) + return SimpleNamespace(**fields) - cd = megakernel_runtime_requirements( - fi_moe_ep_backend_spec("flashinfer_moe_ep_mega_cutedsl") + +def _is_supported(moe, weight_key, activation_key) -> tuple[bool, str]: + supported, reason = FlashInferMoeEpExperts.is_supported_config( + FlashInferMoeEpExperts, + moe, + weight_key, + activation_key, + mk.FusedMoEActivationFormat.Standard, ) - assert cd == frozenset({"torch_dist", "nvshmem"}) + return supported, reason or "" -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() +@pytest.fixture +def megakernel_host(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: + """A vLLM config without DBO/EPLB/weight transfer, on a supported device.""" + config = SimpleNamespace( + weight_transfer_config=None, + parallel_config=SimpleNamespace(enable_dbo=False, enable_eplb=False), + ) + monkeypatch.setattr(fi_ep, "get_current_vllm_config", lambda: config) monkeypatch.setattr( - mod, - "get_ep_group", - lambda: SimpleNamespace(world_size=4, rank_in_group=2, device_group=pg), + FlashInferMoeEpExperts, "_supports_current_device", staticmethod(lambda: True) ) - 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 + return config -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", +def test_megakernel_requires_expert_parallel(megakernel_host): + """The oracle rejects the megakernel through the experts' support check.""" + moe = _megakernel_moe(FLASHINFER_MOE_EP_CUTEDSL) + assert _is_supported(moe, kNvfp4Static, kNvfp4Dynamic) == (True, "") + + moe.moe_parallel_config.use_ep = False + supported, reason = _is_supported(moe, kNvfp4Static, kNvfp4Dynamic) + assert not supported and "parallel config" in reason + + +def test_megakernel_reports_its_own_constraints(megakernel_host): + """Constraints the framework does not know about surface as reasons.""" + _, reason = _is_supported( + _megakernel_moe(FLASHINFER_MOE_EP_CUTEDSL), kNvfp4Static, None ) - 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 + assert "A16 activations" in reason - cd = build_fi_mega_config( - intermediate_size=2048, - top_k=8, - activation_clamp=None, - megakernel="nvfp4_cutedsl", + deep_gemm = _megakernel_moe( + FLASHINFER_MOE_EP_DEEP_GEMM, routing_method=RoutingMethodType.Renormalize ) - assert isinstance(cd.megakernel, _FakeNvfp4CutedslMegaMoeConfig) + _, reason = _is_supported(deep_gemm, kMxfp4Static, None) + assert "routing method" in reason - 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", + megakernel_host.parallel_config.enable_eplb = True + _, reason = _is_supported( + _megakernel_moe(FLASHINFER_MOE_EP_DEEP_GEMM), kMxfp4Static, None + ) + assert reason.endswith("EPLB") + + +def test_megakernels_are_oracle_backends(): + """Quant methods reach the megakernel through the NVFP4/MXFP4 oracles.""" + nvfp4 = map_nvfp4_backend(FLASHINFER_MOE_EP_CUTEDSL) + assert nvfp4 is NvFp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL + assert nvfp4_backend_to_kernel_cls(nvfp4) == [FlashInferMoeEpExperts] + + assert map_mxfp4_backend(FLASHINFER_MOE_EP_DEEP_GEMM) == [ + Mxfp4MoeBackend.FLASHINFER_MOE_EP_DEEP_GEMM + ] + for backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS: + assert mxfp4_backend_to_kernel_cls(backend) == [FlashInferMoeEpExperts] + # The megakernel takes the model's own dimensions. + assert mxfp4_round_up_hidden_size_and_intermediate_size( + backend, 2880, 2880 + ) == (2880, 2880) + + quant_config = make_mxfp4_moe_quant_config( + Mxfp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL, w1_scale=None, w2_scale=None + ) + assert quant_config is not None and quant_config.use_nvfp4_w4a4 + + +def test_megakernel_is_a_modular_kernel_with_pass_through_stages(): + """The megakernel plugs into ``FusedMoEKernel`` like any other experts impl. + + Routing stays outside (top-k ids in), dispatch and combine happen inside, + so prepare passes tokens through untouched and finalize has nothing to do. + """ + moe = make_dummy_moe_config( + num_experts=8, experts_per_token=2, hidden_dim=256, intermediate_size=128 + ) + quant_config = FusedMoEQuantConfig.make("nvfp4", weight_dtype="nvfp4") + kernel = mk.FusedMoEKernel( + PassThroughPrepareAndFinalize(), + FlashInferMoeEpExperts(moe, quant_config), + ) + + assert not kernel.is_monolithic + assert kernel.prepare_finalize.output_is_reduced() + # Generic stages: the int32 top-k ids the megakernel wants are its concern. + assert kernel.prepare_finalize.topk_indices_dtype() is None + assert kernel.fused_experts.expects_unquantized_inputs + assert isinstance( + kernel.fused_experts.finalize_weight_and_reduce_impl(), TopKWeightAndReduceNoOP + ) + assert kernel.fused_experts.workspace_shapes( + 16, 256, 256, 2, 8, 8, None, MoEActivation.SILU + ) == ((0,), (0,), (16, 256)) + + tokens = torch.zeros(4, 256, dtype=torch.bfloat16) + prepared = kernel.prepare_finalize.prepare( + tokens, + torch.ones(4, 2), + torch.zeros(4, 2, dtype=torch.int32), + 8, + None, + False, + quant_config, + ) + assert prepared[0] is tokens and all(item is None for item in prepared[1:]) + + +def test_epilogue_takes_the_weight_global_scales_from_the_quant_config(): + """The canonical ``g1/g2_alphas`` become the per-expert fc1/fc2 alphas and + keep aliasing the layer's tensors, so EPLB permutations reach the kernel.""" + g1_alphas = torch.tensor([1.0, 2.0]) + g2_alphas = torch.tensor([3.0, 4.0]) + quant_config = nvfp4_moe_quant_config( + g1_alphas=g1_alphas, + g2_alphas=g2_alphas, + a1_gscale=torch.ones(2), + a2_gscale=torch.ones(2), + w1_scale=torch.ones(2, 4, 1), + w2_scale=torch.ones(2, 2, 1), + ) + + epilogue = epilogue_from_quant_config(quant_config) + + assert epilogue.fc1_alpha is not None and epilogue.fc2_alpha is not None + assert epilogue.fc1_alpha.data_ptr() == g1_alphas.data_ptr() + assert epilogue.fc2_alpha.data_ptr() == g2_alphas.data_ptr() + # Activations are quantized dynamically: no norm constants. + assert epilogue.input_norm_const == 1.0 and epilogue.fc1_norm_const is None + + +def test_epilogue_rejects_unloaded_weight_global_scales(): + quant_config = nvfp4_moe_quant_config( + g1_alphas=torch.tensor([1.0, float("nan")]), + g2_alphas=torch.ones(2), + a1_gscale=torch.ones(2), + a2_gscale=torch.ones(2), + w1_scale=torch.ones(2, 4, 1), + w2_scale=torch.ones(2, 2, 1), + ) + with pytest.raises(ValueError, match="g1_alphas"): + epilogue_from_quant_config(quant_config) + + +def test_mxfp4_checkpoints_keep_the_default_epilogue(): + """MXFP4 has no per-tensor global scales; the megakernel needs no alphas.""" + quant_config = FusedMoEQuantConfig.make("nvfp4", weight_dtype="mxfp4") + assert epilogue_from_quant_config(quant_config) == fi_ep.FlashInferMoeEpEpilogue() + + +def test_nvfp4_kernel_format_passes_megakernel_weights_through(): + """The megakernel takes the canonical tensors as they are; the oracle must + not reject or reshape them.""" + tensors = tuple(torch.arange(4.0).reshape(2, 2) + i for i in range(8)) + converted = convert_to_nvfp4_moe_kernel_format( + NvFp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL, + SimpleNamespace(), + *tensors, + is_act_and_mul=True, + ) + assert all(out is src for out, src in zip(converted, tensors)) + + +def test_passthrough_all2all_backend_selects_the_pass_through_stages(): + """``all2all_backend="passthrough"`` is how the oracles get the + pass-through prepare/finalize.""" + moe = make_dummy_moe_config( + num_experts=8, experts_per_token=2, hidden_dim=256, intermediate_size=128 + ) + moe.moe_parallel_config = replace( + moe.moe_parallel_config, all2all_backend=PASSTHROUGH_ALL2ALL_BACKEND + ) + assert not moe.moe_parallel_config.use_all2all_kernels + + prepare_finalize = maybe_make_prepare_finalize( + moe, FusedMoEQuantConfig.make("nvfp4", weight_dtype="nvfp4") + ) + + assert isinstance(prepare_finalize, PassThroughPrepareAndFinalize) + + +def test_passthrough_all2all_backend_is_not_a_cli_choice(): + """Users cannot pick the passthrough all2all backend explicitly.""" + parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) + with pytest.raises(SystemExit): + parser.parse_args(["--all2all-backend", PASSTHROUGH_ALL2ALL_BACKEND]) + parser.parse_args(["--all2all-backend", "deepep_low_latency"]) + + +def test_megakernel_moe_backend_binds_the_all2all_backend(): + """A megakernel deployment gets its all2all backend by default.""" + parallel = SimpleNamespace(all2all_backend="allgather_reducescatter") + bind_passthrough_all2all_backend(FLASHINFER_MOE_EP_CUTEDSL, parallel) + assert parallel.all2all_backend == PASSTHROUGH_ALL2ALL_BACKEND + bind_passthrough_all2all_backend(FLASHINFER_MOE_EP_CUTEDSL, parallel) + + with pytest.raises(ValueError, match="cannot be paired"): + bind_passthrough_all2all_backend( + FLASHINFER_MOE_EP_DEEP_GEMM, + SimpleNamespace(all2all_backend="deepep_low_latency"), + ) + with pytest.raises(ValueError, match="only valid with"): + bind_passthrough_all2all_backend( + "flashinfer_trtllm", + SimpleNamespace(all2all_backend=PASSTHROUGH_ALL2ALL_BACKEND), ) - -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)) + untouched = SimpleNamespace(all2all_backend="deepep_low_latency") + bind_passthrough_all2all_backend("flashinfer_trtllm", untouched) + assert untouched.all2all_backend == "deepep_low_latency" diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index f87dd75fb2..9ee0212fba 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -657,7 +657,7 @@ def test_deepseek_v4_mega_moe_does_not_double_add_fused_shared_expert( moe = DeepseekV4MoE.__new__(DeepseekV4MoE) torch.nn.Module.__init__(moe) - moe.use_mega_moe = True + moe.use_native_mega_moe = True moe.gate = FakeGate() moe.experts = FakeExperts() moe.shared_experts = FakeSharedExperts() diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index 3e70450f2c..f77832d608 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -241,7 +241,7 @@ def test_v41_dspark_loads_linear_scales( linear_scale_name=scale_name, pad_shared_expert=False, model=SimpleNamespace( - layers=[SimpleNamespace(ffn=SimpleNamespace(use_mega_moe=False))], + layers=[SimpleNamespace(ffn=SimpleNamespace(use_native_mega_moe=False))], confidence_head=None, ), named_parameters=lambda: [(runtime_name, param)], diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 7fe868ae5a..df69f17149 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -3,7 +3,7 @@ import contextlib from collections.abc import Callable, Iterable from dataclasses import asdict, fields -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, Final, Literal from pydantic import Field, field_validator @@ -12,6 +12,7 @@ from vllm.logger import init_logger if TYPE_CHECKING: from vllm.config import VllmConfig + from vllm.config.parallel import ParallelConfig logger = init_logger(__name__) @@ -128,10 +129,10 @@ MoEBackend = Literal[ "flashinfer_trtllm", "flashinfer_cutlass", "flashinfer_cutedsl", + "flashinfer_moe_ep_cutedsl", "flashinfer_b12x", "b12x", "flashinfer_moe_ep_mega_deep_gemm", - "flashinfer_moe_ep_mega_cutedsl", "marlin", "humming", "triton_unfused", @@ -143,21 +144,28 @@ MoEBackend = Literal[ "rdna3", ] -# 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_CUTEDSL = "flashinfer_moe_ep_cutedsl" +FLASHINFER_MOE_EP_DEEP_GEMM = "flashinfer_moe_ep_mega_deep_gemm" + +# Backends implemented by the shared FlashInfer MoE-EP adapter. FLASHINFER_MOE_EP_BACKENDS = frozenset( { - "flashinfer_moe_ep_mega_deep_gemm", - "flashinfer_moe_ep_mega_cutedsl", + FLASHINFER_MOE_EP_CUTEDSL, + FLASHINFER_MOE_EP_DEEP_GEMM, } ) +# Those backends dispatch and combine themselves, so the framework must not +# run a separate all2all around them. +PASSTHROUGH_ALL2ALL_BACKEND: Final = "passthrough" -# 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 +# "Native" here meaning that the backend does not use the fused MoE path. +# Currently only relevant for the vLLM provided DeepGEMM MegaMoE., +# TODO: migrate to the fused MoE path ultimately. +NATIVE_MEGA_MOE_BACKENDS = frozenset({"deep_gemm_mega_moe"}) + +# Every megakernel backend. The FlashInfer variants run through the standard +# fused-MoE model path; the quantization method binds the kernel. +MEGA_MOE_BACKENDS = NATIVE_MEGA_MOE_BACKENDS | FLASHINFER_MOE_EP_BACKENDS SparseIndexerTopkBackend = Literal[ "auto", @@ -170,8 +178,8 @@ SparseIndexerTopkBackend = Literal[ "aiter", ] -# Architectures whose model code wires up the flashinfer moe_ep experts. MTP -# and DSpark draft variants inherit the setting from these target models. +# DeepGEMM consumes the DeepSeek-V4 MXFP4 checkpoint recipe. CuTeDSL is +# selected by quantization methods and is not architecture-specific. FLASHINFER_MOE_EP_ARCHITECTURES = frozenset( { "DeepseekV4ForCausalLM", @@ -181,11 +189,48 @@ FLASHINFER_MOE_EP_ARCHITECTURES = frozenset( ) +def bind_passthrough_all2all_backend( + moe_backend: str, parallel_config: "ParallelConfig" +) -> None: + """Bind ``all2all_backend="passthrough"`` to MoE backends that own their + expert-parallel communication (currently the FlashInfer MoE-EP megakernels). + + Such a deployment gets ``passthrough`` when the all2all backend is left at + its default and rejects any other explicit choice; the value is meaningless + with a MoE backend that relies on the framework to communicate. The CLI does + not offer the value; the early return keeps repeated + ``set_platform_defaults`` calls idempotent. + """ + a2a = parallel_config.all2all_backend + if moe_backend in FLASHINFER_MOE_EP_BACKENDS: + if a2a == PASSTHROUGH_ALL2ALL_BACKEND: + return + if a2a != "allgather_reducescatter": + raise ValueError( + f"moe_backend={moe_backend!r} dispatches and combines itself and " + f"cannot be paired with all2all_backend={a2a!r}; leave " + f"--all2all-backend unset." + ) + logger.info_once( + "moe_backend=%r: using all2all_backend=%r (dispatch and combine run " + "inside the MoE backend).", + moe_backend, + PASSTHROUGH_ALL2ALL_BACKEND, + ) + parallel_config.all2all_backend = PASSTHROUGH_ALL2ALL_BACKEND + elif a2a == PASSTHROUGH_ALL2ALL_BACKEND: + raise ValueError( + f"all2all_backend={a2a!r} is only valid with a moe_backend that " + f"dispatches and combines itself ({sorted(FLASHINFER_MOE_EP_BACKENDS)})," + f" got moe_backend={moe_backend!r}." + ) + + 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: + """Reject model-specific FlashInfer MoE-EP backends.""" + if moe_backend != FLASHINFER_MOE_EP_DEEP_GEMM: return if not any(arch in FLASHINFER_MOE_EP_ARCHITECTURES for arch in architectures): raise ValueError( @@ -255,6 +300,9 @@ class KernelConfig: - "flashinfer_trtllm": Use FlashInfer with TRTLLM-GEN kernels - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels - "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels (FP4 only) + - "flashinfer_moe_ep_cutedsl": Use FlashInfer's CuTeDSL MoE-EP + mega-kernel with NVFP4 weights (MXFP4 checkpoints are requantized at + load); requires Blackwell, expert parallelism, and NVSHMEM - "flashinfer_b12x": Use FlashInfer CuteDSL fused MoE for SM12x (RTX Pro 6000 / DGX Spark) - "b12x": Use b12x FP4 MoE kernels on SM12x @@ -262,10 +310,6 @@ class KernelConfig: 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 @@ -391,6 +435,7 @@ class KernelConfig: validate_flashinfer_moe_ep_model( self.moe_backend, vllm_config.model_config.architectures ) + bind_passthrough_all2all_backend(self.moe_backend, vllm_config.parallel_config) platform_op_priority = current_platform.get_default_ir_op_priority(vllm_config) logger.debug( diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 7adc48254b..a5e01ee6f8 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -56,6 +56,7 @@ All2AllBackend = Literal[ "flashinfer_all2allv", # temporary alias for flashinfer_nvlink_two_sided "flashinfer_nvlink_two_sided", "flashinfer_nvlink_one_sided", + "passthrough", ] @@ -211,7 +212,10 @@ class ParallelConfig: - "moonep": MoonEP balanced EP with dynamic redundant experts (NVLink) - "nixl_ep": Use nixl-ep kernels - "flashinfer_nvlink_one_sided": Use flashinfer high-throughput a2a kernels - - "flashinfer_nvlink_two_sided": Use flashinfer two-sided kernels for mnnvl""" + - "flashinfer_nvlink_two_sided": Use flashinfer two-sided kernels for mnnvl + - "passthrough": No all2all at all: the MoE backend dispatches and combines + itself (e.g. the FlashInfer MoE-EP megakernels). Bound automatically for + such MoE backends; not selectable from the CLI.""" max_parallel_loading_workers: int | None = Field(default=None, ge=1) """Maximum number of parallel loading workers when loading model @@ -742,6 +746,7 @@ class ParallelConfig: "deepep_low_latency", "deepep_v2", "flashinfer_nvlink_one_sided", + "passthrough", "mori_high_throughput", "mori_low_latency", "nixl_ep", diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 5f339540b4..8658d8663b 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -49,6 +49,20 @@ if has_flashinfer_nvlink_one_sided(): logger = init_logger(__name__) +class PassThroughAll2AllManager(All2AllManagerBase): + """Placeholder for ``all2all_backend="passthrough"``. + + The MoE backend dispatches and combines itself, so there is no all2all to + manage. + """ + + def get_handle(self, kwargs): + raise RuntimeError( + "passthrough has no all2all handle: dispatch and combine run " + "inside the MoE backend." + ) + + class AgRsAll2AllManager(All2AllManagerBase): """An implementation of all2all communication based on all-gather (dispatch) and reduce-scatter (combine). diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index ef790f7732..0a203956f0 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -255,6 +255,12 @@ class CudaCommunicator(DeviceCommunicatorBase): from .all2all import FlashInferNVLinkOneSidedManager self.all2all_manager = FlashInferNVLinkOneSidedManager(self.cpu_group) + elif self.all2all_backend == "passthrough": + from .all2all import PassThroughAll2AllManager + + self.all2all_manager = PassThroughAll2AllManager( + self.cpu_group, tcp_store_group + ) else: raise ValueError(f"Unknown all2all backend: {self.all2all_backend}") diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index ce24ac28f2..f4cbcf9012 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -80,6 +80,7 @@ from vllm.config.cache import ( ) from vllm.config.device import Device from vllm.config.kernel import ( + PASSTHROUGH_ALL2ALL_BACKEND, IrOpPriorityConfig, LinearBackend, MoEBackend, @@ -1104,6 +1105,10 @@ class EngineArgs: # Parallel arguments parallel_kwargs = get_kwargs(ParallelConfig) + # Bound from --moe-backend in KernelConfig.set_platform_defaults(). + parallel_kwargs["all2all_backend"]["choices"].remove( + PASSTHROUGH_ALL2ALL_BACKEND + ) parallel_group = parser.add_argument_group( title="ParallelConfig", description=ParallelConfig.__doc__, diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 3186b5d3c4..6a0a2b2358 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -31,6 +31,9 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided import ( # noqa: E501 FlashInferNVLinkTwoSidedPrepareAndFinalize, ) +from vllm.model_executor.layers.fused_moe.prepare_finalize.passthrough import ( + PassThroughPrepareAndFinalize, +) from vllm.platforms import current_platform from vllm.utils.import_utils import ( has_deep_ep, @@ -185,6 +188,10 @@ def maybe_make_prepare_finalize( all2all_manager: Any | None = None, input_dtype: torch.dtype | None = None, ) -> FusedMoEPrepareAndFinalize | None: + if moe.moe_parallel_config.use_passthrough_all2all: + # The experts dispatch and combine themselves, in every EP topology. + return PassThroughPrepareAndFinalize() + if not moe.moe_parallel_config.use_all2all_kernels: if not allow_new_interface: return None diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 302ed121fe..158bc90769 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1068,6 +1068,12 @@ class FusedMoEParallelConfig: and self.all2all_backend == "allgather_reducescatter" ) + @property + def use_passthrough_all2all(self): + # Not gated on use_all2all_kernels: the experts dispatch and combine + # internally in every EP topology, including TP-only. + return self.all2all_backend == "passthrough" + @property def use_mori_kernels(self): return self.use_all2all_kernels and self.all2all_backend in ( @@ -1530,6 +1536,10 @@ class FusedMoEConfig: def use_deepep_ll_kernels(self): return self.moe_parallel_config.use_deepep_ll_kernels + @property + def use_passthrough_all2all(self): + return self.moe_parallel_config.use_passthrough_all2all + @property def use_mori_kernels(self): return self.moe_parallel_config.use_mori_kernels diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_moe_ep.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_moe_ep.py new file mode 100644 index 0000000000..8fbfcd9178 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_moe_ep.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FlashInfer MoE-EP megakernels as a modular-kernel experts implementation. + +The megakernel routes on vLLM's top-k output and then dispatches, computes and +combines in one launch, so it is paired with a pass-through prepare/finalize. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, +) +from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep import ( + FlashInferMoeEp, + FlashInferMoeEpEpilogue, + FlashInferMoeEpWeights, + apply_topk_in_fc1, + flashinfer_moe_ep_unsupported_reasons, + supports_current_device, + validate_flashinfer_moe_ep_layer, +) +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp4Static, + kNvfp4Static, + require_finite_positive_scales, +) +from vllm.model_executor.utils import replace_parameter + +if TYPE_CHECKING: + pass + +# Layer parameters the megakernel reads. After ``process_weights_after_loading`` +# they alias the megakernel's own tensors, so EPLB permutes exactly what it uses. +KERNEL_WEIGHT_NAMES = ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale") + + +_supported_weight_quant_schemes: frozenset[QuantKey] = frozenset( + (kNvfp4Static, kMxfp4Static) +) + +_supported_activations: frozenset[MoEActivation] = frozenset((MoEActivation.SILU,)) + + +class FlashInferMoeEpExperts(mk.FusedMoEExpertsModular): + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ) -> None: + super().__init__(moe_config, quant_config) + self._adapter: FlashInferMoeEp | None = None + + @staticmethod + def is_supported_config( + cls: type[mk.FusedMoEExperts], + moe_config: FusedMoEConfig, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + activation_format: mk.FusedMoEActivationFormat, + ) -> tuple[bool, str | None]: + supported, reason = mk.FusedMoEExperts.is_supported_config( + cls, moe_config, weight_key, activation_key, activation_format + ) + if not supported: + return False, reason + reasons = flashinfer_moe_ep_unsupported_reasons( + moe_config, weight_key, activation_key + ) + if reasons: + return False, f"kernel does not support {', '.join(reasons)}" + return True, None + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def _supports_current_device() -> bool: + return supports_current_device() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return weight_key in _supported_weight_quant_schemes + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation in _supported_activations + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + return ( + moe_parallel_config.use_ep + and not moe_parallel_config.use_batched_activation_format + ) + + @property + def expects_unquantized_inputs(self) -> bool: + return True + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + return (0,), (0,), (M, K) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """Build the megakernel from the layer's canonical weights. + + The oracle has already put the weights in kernel format; the per-expert + weight global scales arrive through the quant config and become the + epilogue alphas. + """ + validate_flashinfer_moe_ep_layer(layer) # type: ignore[arg-type] + adapter = FlashInferMoeEp( + self.moe_config, + FlashInferMoeEpWeights( + w13=layer.w13_weight, + w2=layer.w2_weight, + w13_scale=getattr(layer, "w13_weight_scale", None), + w2_scale=getattr(layer, "w2_weight_scale", None), + ), + epilogue_from_quant_config(self.quant_config), + apply_topk_in_fc1=apply_topk_in_fc1( + self.moe_config, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + ), + ) + adapter.warmup() + for name, tensor in zip(KERNEL_WEIGHT_NAMES, adapter.kernel_weights()): + replace_parameter(layer, name, tensor) + self._adapter = adapter + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ) -> None: + if self._adapter is None: + raise RuntimeError("FlashInfer MoE-EP weights have not been processed") + output.copy_( + self._adapter(hidden_states, topk_ids.to(torch.int32), topk_weights) + ) + + +def epilogue_from_quant_config( + quant_config: FusedMoEQuantConfig, +) -> FlashInferMoeEpEpilogue: + """Per-expert epilogue constants from the canonical quant config. + + The weight global scales (``g1_alphas``/``g2_alphas``) become the fc1/fc2 + alphas. Activations are quantized dynamically inside the megakernel, so the + checkpoint's input scales drop out. MXFP4 checkpoints carry no global scales + and keep the default epilogue. The alphas alias the layer's parameters, so + EPLB permutations reach the kernel. + """ + fc1_alpha, fc2_alpha = quant_config.g1_alphas, quant_config.g2_alphas + if fc1_alpha is None and fc2_alpha is None: + return FlashInferMoeEpEpilogue() + assert fc1_alpha is not None and fc2_alpha is not None + require_finite_positive_scales("g1_alphas", fc1_alpha) + require_finite_positive_scales("g2_alphas", fc2_alpha) + if fc1_alpha.dim() != 1 or fc2_alpha.dim() != 1: + raise ValueError("the megakernel needs one weight global scale per expert") + return FlashInferMoeEpEpilogue( + fc1_alpha=fc1_alpha.float().contiguous(), + fc2_alpha=fc2_alpha.float().contiguous(), + ) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py b/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py new file mode 100644 index 0000000000..bac5345255 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py @@ -0,0 +1,481 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import contextlib +import sys +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal, cast + +import torch + +from vllm.config import get_current_vllm_config +from vllm.config.kernel import ( + FLASHINFER_MOE_EP_BACKENDS, + FLASHINFER_MOE_EP_CUTEDSL, + FLASHINFER_MOE_EP_DEEP_GEMM, +) +from vllm.distributed import get_ep_group +from vllm.model_executor.layers.fused_moe.config import RoutingMethodType +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp4Static, + kNvfp4Static, +) +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import _import_deep_gemm + +if TYPE_CHECKING: + from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts + +FlashInferMoeEpKernel = Literal["cutedsl", "deep_gemm"] +FlashInferMoeEpWeightFormat = Literal["nvfp4", "mxfp4"] +_WEIGHT_FORMATS: dict[QuantKey | None, FlashInferMoeEpWeightFormat] = { + kNvfp4Static: "nvfp4", + kMxfp4Static: "mxfp4", +} + +_SUPPORTED_CAPABILITIES = frozenset({(10, 0), (10, 3)}) + + +def supports_current_device() -> bool: + capability = current_platform.get_device_capability() + return capability is not None and ( + (capability.major, capability.minor) in _SUPPORTED_CAPABILITIES + ) + + +_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, +) + + +@dataclass(frozen=True) +class FlashInferMoeEpBackendSpec: + kernel: FlashInferMoeEpKernel + weight_formats: frozenset[FlashInferMoeEpWeightFormat] + + +@dataclass(frozen=True) +class FlashInferMoeEpWeights: + w13: torch.Tensor + w2: torch.Tensor + w13_scale: torch.Tensor | None = None + w2_scale: torch.Tensor | None = None + + +@dataclass(frozen=True) +class FlashInferMoeEpEpilogue: + input_norm_const: float = 1.0 + fc1_alpha: torch.Tensor | None = None + fc2_alpha: torch.Tensor | None = None + fc1_norm_const: torch.Tensor | None = None + + +_BACKEND_SPECS = { + FLASHINFER_MOE_EP_CUTEDSL: FlashInferMoeEpBackendSpec( + kernel="cutedsl", + weight_formats=frozenset({"nvfp4", "mxfp4"}), + ), + FLASHINFER_MOE_EP_DEEP_GEMM: FlashInferMoeEpBackendSpec( + kernel="deep_gemm", + weight_formats=frozenset({"mxfp4"}), + ), +} + +assert frozenset(_BACKEND_SPECS) == FLASHINFER_MOE_EP_BACKENDS + + +def is_flashinfer_moe_ep_backend(moe_backend: str) -> bool: + return moe_backend in _BACKEND_SPECS + + +def flashinfer_moe_ep_backend_spec( + moe_backend: str, +) -> FlashInferMoeEpBackendSpec: + try: + return _BACKEND_SPECS[moe_backend] + except KeyError: + raise ValueError( + f"{moe_backend!r} is not a FlashInfer MoE-EP backend; expected " + f"one of {sorted(_BACKEND_SPECS)}" + ) from None + + +def flashinfer_moe_ep_unsupported_reasons( + moe: FusedMoEConfig, + weight_key: QuantKey | None, + activation_key: QuantKey | None, +) -> tuple[str, ...]: + """Deployment constraints of the megakernel beyond the framework's checks. + + ``activation_key`` is ``None`` for MXFP4 checkpoints (the kernel quantizes + bf16 activations itself) but marks a W4A16 recipe for NVFP4 weights. + """ + spec = flashinfer_moe_ep_backend_spec(moe.moe_backend) + unsupported: list[str] = [] + weight_format = _WEIGHT_FORMATS.get(weight_key) + if weight_format is None: + unsupported.append(f"weight scheme {weight_key}") + elif weight_format not in spec.weight_formats: + unsupported.append(f"{weight_format.upper()} weights") + if weight_key == kNvfp4Static and activation_key is None: + unsupported.append("A16 activations with NVFP4 weights") + if moe.skip_final_all_reduce: + unsupported.append("skip_final_all_reduce") + if moe.in_dtype != torch.bfloat16: + unsupported.append(f"activation dtype {moe.in_dtype}") + if moe.has_bias: + unsupported.append("expert bias") + if moe.swiglu_alpha is not None or moe.swiglu_beta is not None: + unsupported.append("custom SwiGLU alpha or beta") + if ( + spec.kernel == "deep_gemm" + and moe.routing_method is not RoutingMethodType.DeepseekV4 + ): + unsupported.append(f"routing method {moe.routing_method.name}") + + vllm_config = get_current_vllm_config() + if vllm_config.weight_transfer_config is not None: + unsupported.append("runtime weight transfer") + if vllm_config.parallel_config.enable_dbo: + unsupported.append("dual batch overlap") + if spec.kernel == "deep_gemm" and vllm_config.parallel_config.enable_eplb: + unsupported.append("EPLB") + return tuple(unsupported) + + +def validate_flashinfer_moe_ep_layer(layer: RoutedExperts) -> None: + if layer.expert_map_manager.placement_strategy != "linear": + raise ValueError( + f"{layer.moe_config.moe_backend} requires contiguous linear " + "expert placement" + ) + + +def apply_topk_in_fc1( + moe: FusedMoEConfig, + *, + apply_router_weight_on_input: bool = False, +) -> bool: + return ( + apply_router_weight_on_input + or moe.routing_method is RoutingMethodType.DeepseekV4 + ) + + +def _dequant_mxfp4_ue8m0_gran32( + packed: torch.Tensor, + scale: torch.Tensor, +) -> torch.Tensor: + raw = packed.view(torch.uint8) + lut = torch.tensor(_E2M1_LUT, dtype=torch.float32, device=raw.device) + values = torch.empty( + raw.shape[0], + raw.shape[1] * 2, + dtype=torch.float32, + device=raw.device, + ) + values[:, ::2] = lut[(raw & 0x0F).to(torch.int64)] + values[:, 1::2] = lut[(raw >> 4).to(torch.int64)] + encoded_scale = ( + scale + if scale.dtype == torch.float8_e8m0fnu + else scale.view(torch.float8_e8m0fnu) + ) + decoded_scale = encoded_scale.float() + return (values * decoded_scale.repeat_interleave(32, dim=-1)).to(torch.bfloat16) + + +def _dequant_mxfp4_experts( + weight: torch.Tensor, + scale: torch.Tensor, +) -> torch.Tensor: + num_experts, output_size, packed_input_size = weight.shape + result = torch.empty( + num_experts, + output_size, + packed_input_size * 2, + dtype=torch.bfloat16, + device=weight.device, + ) + for expert_id in range(num_experts): + result[expert_id] = _dequant_mxfp4_ue8m0_gran32( + weight[expert_id], scale[expert_id] + ) + return result + + +def mxfp4_moe_ep_weights( + moe_backend: str, + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, +) -> FlashInferMoeEpWeights: + spec = flashinfer_moe_ep_backend_spec(moe_backend) + if spec.kernel == "deep_gemm": + return FlashInferMoeEpWeights(w13, w2, w13_scale, w2_scale) + return FlashInferMoeEpWeights( + w13=_dequant_mxfp4_experts(w13, w13_scale), + w2=_dequant_mxfp4_experts(w2, w2_scale), + ) + + +class _FlashInferMoeEpApi: + def __init__(self) -> None: + try: + from flashinfer.moe_ep import ( + BootstrapConfig, + DeepGemmMegaMoeConfig, + FleetParams, + MegaConfig, + MoEEpMegaLayer, + MoEEpTensors, + MoEWeightPack, + Nvfp4CutedslMegaMoeConfig, + ) + except ImportError as err: + raise RuntimeError( + "FlashInfer MoE-EP requires the public API shipped with " + "FlashInfer 0.7.0 or newer" + ) from err + + self.BootstrapConfig = BootstrapConfig + self.DeepGemmMegaMoeConfig = DeepGemmMegaMoeConfig + self.FleetParams = FleetParams + self.MegaConfig = MegaConfig + self.MoEEpMegaLayer = MoEEpMegaLayer + self.MoEEpTensors = MoEEpTensors + self.MoEWeightPack = MoEWeightPack + self.Nvfp4CutedslMegaMoeConfig = Nvfp4CutedslMegaMoeConfig + + +def _load_flashinfer_moe_ep_api() -> _FlashInferMoeEpApi: + return _FlashInferMoeEpApi() + + +def _expose_deep_gemm_to_flashinfer() -> None: + deep_gemm = _import_deep_gemm() + if deep_gemm is None: + raise RuntimeError( + "FlashInfer MoE-EP DeepGEMM requires a usable DeepGEMM installation" + ) + sys.modules.setdefault("deep_gemm", deep_gemm) + + +class FlashInferMoeEp: + def __init__( + self, + moe: FusedMoEConfig, + weights: FlashInferMoeEpWeights, + epilogue: FlashInferMoeEpEpilogue | None = None, + *, + apply_topk_in_fc1: bool, + ) -> None: + spec = flashinfer_moe_ep_backend_spec(moe.moe_backend) + if spec.kernel == "deep_gemm": + _expose_deep_gemm_to_flashinfer() + api = _load_flashinfer_moe_ep_api() + if epilogue is None: + epilogue = FlashInferMoeEpEpilogue() + ep_group = get_ep_group() + bootstrap = api.BootstrapConfig( + world_size=ep_group.world_size, + rank=ep_group.rank_in_group, + process_group=ep_group.device_group, + device=torch.accelerator.current_device_index(), + ) + fleet_params = api.FleetParams( + num_experts=moe.num_experts, + max_tokens_per_rank=moe.max_num_tokens, + token_hidden_size=moe.hidden_dim, + ) + weight_pack = api.MoEWeightPack( + w13=weights.w13, + w2=weights.w2, + w13_scale=weights.w13_scale, + w2_scale=weights.w2_scale, + ) + if spec.kernel == "cutedsl": + megakernel = api.Nvfp4CutedslMegaMoeConfig( + intermediate_size=moe.intermediate_size, + top_k=moe.experts_per_token, + gate_up_clamp=moe.swiglu_limit, + fast_math=True, + apply_topk_in_fc1=apply_topk_in_fc1, + enable_in_kernel_fc2_reduce=False, + combine_dtype="bf16", + input_norm_const=epilogue.input_norm_const, + fc1_alpha=None, + fc2_alpha=None, + fc1_norm_const=None, + knobs=None, + ) + else: + megakernel = api.DeepGemmMegaMoeConfig( + intermediate_size=moe.intermediate_size, + top_k=moe.experts_per_token, + activation_clamp=moe.swiglu_limit, + fast_math=True, + ) + backend = api.MegaConfig( + megakernel=megakernel, + quantize_input=True, + preprocess_weights=True, + ) + + self._moe_ep_tensors_cls = api.MoEEpTensors + self._mega_layer: Any | None = api.MoEEpMegaLayer( + bootstrap, + fleet_params, + weight_pack, + backend, + ) + self._device = weights.w13.device + self._hidden_size = moe.hidden_dim + self._max_num_tokens = moe.max_num_tokens + self._num_experts = moe.num_experts + self._top_k = moe.experts_per_token + self._epilogue: FlashInferMoeEpEpilogue | None = epilogue + self._return_workspace_view = moe.routing_method is RoutingMethodType.DeepseekV4 + + @property + def can_overlap_shared_experts(self) -> bool: + return False + + @property + def output_is_reduced(self) -> bool: + return True + + @property + def topk_indices_dtype(self) -> torch.dtype: + return torch.int32 + + @property + def is_monolithic(self) -> bool: + return False + + def _tensors( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + ) -> Any: + epilogue = cast(FlashInferMoeEpEpilogue, self._epilogue) + return self._moe_ep_tensors_cls( + hidden_states=hidden_states, + topk_ids=topk_ids, + topk_weights=topk_weights, + fc1_alpha=epilogue.fc1_alpha, + fc2_alpha=epilogue.fc2_alpha, + fc1_norm_const=epilogue.fc1_norm_const, + ) + + def __call__( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + ) -> torch.Tensor: + if self._mega_layer is None: + raise RuntimeError("FlashInfer MoE-EP layer was destroyed") + if hidden_states.shape[0] > self._max_num_tokens: + raise ValueError( + f"FlashInfer MoE-EP got {hidden_states.shape[0]} tokens, but " + f"the workspace supports at most {self._max_num_tokens}" + ) + tensors = self._tensors(hidden_states, topk_ids, topk_weights) + if self._return_workspace_view and getattr( + self._mega_layer, "supports_output_view", False + ): + return self._mega_layer.forward(tensors, return_workspace_view=True) + return self._mega_layer(tensors) + + @torch.inference_mode() + def warmup(self) -> None: + if self._mega_layer is None: + return + hidden_states = torch.zeros( + 1, + self._hidden_size, + dtype=torch.bfloat16, + device=self._device, + ) + topk_ids = ( + torch.arange(self._top_k, dtype=torch.int32, device=self._device) + .mul_(self._num_experts // self._top_k) + .view(1, self._top_k) + ) + topk_weights = torch.full( + (1, self._top_k), + 1.0 / self._top_k, + dtype=torch.float32, + device=self._device, + ) + self._mega_layer.warmup(self._tensors(hidden_states, topk_ids, topk_weights)) + + def kernel_weights( + self, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Kernel-resident fc1 weight, fc1 scales, fc2 weight, fc2 scales. + + Contiguous per-expert byte tensors that alias the memory the megakernel + reads, so a caller may permute experts in place (EPLB). + """ + if self._mega_layer is None: + raise RuntimeError("FlashInfer MoE-EP layer was destroyed") + (fc1_weight, fc1_sf), (fc2_weight, fc2_sf) = self._mega_layer._transformed + return ( + _expert_major_bytes(fc1_weight), + _expert_major_bytes(fc1_sf), + _expert_major_bytes(fc2_weight), + _expert_major_bytes(fc2_sf), + ) + + def expert_weight_views(self) -> tuple[torch.Tensor, ...]: + """Every per-expert tensor the kernel reads: weights plus epilogue.""" + epilogue = cast(FlashInferMoeEpEpilogue, self._epilogue) + vectors = (epilogue.fc1_alpha, epilogue.fc2_alpha, epilogue.fc1_norm_const) + return (*self.kernel_weights(), *(v for v in vectors if v is not None)) + + def destroy(self) -> None: + mega_layer = self._mega_layer + if mega_layer is None: + return + mega_layer.destroy() + self._mega_layer = None + self._epilogue = None + + def __del__(self) -> None: + with contextlib.suppress(Exception): + self.destroy() + + +def _expert_major_bytes(tensor: torch.Tensor) -> torch.Tensor: + """Contiguous storage of a per-expert kernel tensor, 1-byte dtypes as uint8. + + The CuTeDSL kernel keeps (E, K, N) transpose views over (E, N, K) storage. + """ + if not tensor.is_contiguous(): + tensor = tensor.transpose(1, 2) + if not tensor.is_contiguous(): + raise ValueError(f"expert tensor is not expert-major: {tensor.stride()}") + return tensor.view(torch.uint8) if tensor.element_size() == 1 else tensor diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 27d59ba2e0..a3c65dc4a5 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -8,7 +8,11 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm import envs from vllm.config import get_current_vllm_config -from vllm.config.kernel import MoEBackend +from vllm.config.kernel import ( + FLASHINFER_MOE_EP_CUTEDSL, + FLASHINFER_MOE_EP_DEEP_GEMM, + MoEBackend, +) from vllm.config.quantization import QuantizationConfigArgs from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( @@ -145,6 +149,9 @@ class Mxfp4MoeBackend(Enum): EMULATION = "EMULATION" # Humming HUMMING = "HUMMING" + # FlashInfer MoE-EP megakernels: routing outside, dispatch/GEMMs/combine fused. + FLASHINFER_MOE_EP_CUTEDSL = "FLASHINFER_MOE_EP_CUTEDSL" + FLASHINFER_MOE_EP_DEEP_GEMM = "FLASHINFER_MOE_EP_DEEP_GEMM" # Backends that share the same TRTLLM weight format @@ -165,6 +172,13 @@ B12X_BACKENDS = ( Mxfp4MoeBackend.B12X_MXFP4_BF16, ) +# Megakernels consume the checkpoint's MXFP4 tensors through the shared +# FlashInfer MoE-EP adapter; vLLM does not convert their weight layout. +FLASHINFER_MOE_EP_MXFP4_BACKENDS = ( + Mxfp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL, + Mxfp4MoeBackend.FLASHINFER_MOE_EP_DEEP_GEMM, +) + def backend_to_kernel_cls( backend: Mxfp4MoeBackend, @@ -291,6 +305,13 @@ def backend_to_kernel_cls( return [OCP_MXQuantizationEmulationTritonExperts] + elif backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS: + from vllm.model_executor.layers.fused_moe.experts.flashinfer_moe_ep import ( # noqa: E501 + FlashInferMoeEpExperts, + ) + + return [FlashInferMoeEpExperts] + else: raise ValueError(f"Unknown MXFP4 MoE backend: {backend.value}") @@ -314,6 +335,10 @@ def map_mxfp4_backend(runner_backend: MoEBackend) -> list[Mxfp4MoeBackend]: Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8, ], "flashinfer_cutlass_afp8": [Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8], + "flashinfer_moe_ep_cutedsl": [Mxfp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL], + "flashinfer_moe_ep_mega_deep_gemm": [ + Mxfp4MoeBackend.FLASHINFER_MOE_EP_DEEP_GEMM + ], "triton": [Mxfp4MoeBackend.TRITON], "triton_unfused": [Mxfp4MoeBackend.TRITON_UNFUSED], "humming": [Mxfp4MoeBackend.HUMMING], @@ -735,6 +760,9 @@ def mxfp4_round_up_hidden_size_and_intermediate_size( activation: MoEActivation | None = None, ) -> tuple[int, int]: """Round up hidden_size and intermediate_size based on backend requirements.""" + if backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS: + # The megakernel takes the model's dimensions and validates them itself. + return hidden_size, intermediate_size if backend in B12X_BACKENDS: # b12x plans for the exact model dimensions. B12xExperts validates the # required MXFP4 block alignment before selecting the backend. @@ -1827,6 +1855,29 @@ def convert_weight_to_mxfp4_moe_kernel_format( w2_bias=w2_bias, _cache_permute_indices=_cache_permute_indices, ) + elif mxfp4_backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS: + from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep import ( + mxfp4_moe_ep_weights, + ) + + # DeepGEMM consumes MXFP4 directly; CuTeDSL gets bf16 and requantizes. + weights = mxfp4_moe_ep_weights( + FLASHINFER_MOE_EP_CUTEDSL + if mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL + else FLASHINFER_MOE_EP_DEEP_GEMM, + w13_weight, + w2_weight, + w13_weight_scale, + w2_weight_scale, + ) + return ( + weights.w13, + weights.w2, + weights.w13_scale, + weights.w2_scale, + w13_bias, + w2_bias, + ) elif mxfp4_backend == Mxfp4MoeBackend.CPU: from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( prepare_mxfp4_moe_layer_for_cpu, @@ -1874,6 +1925,9 @@ def make_mxfp4_moe_quant_config( layer: "RoutedExperts | None" = None, ) -> FusedMoEQuantConfig | None: """Create a FusedMoEQuantConfig for the given MXFP4 backend.""" + if mxfp4_backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS: + # The megakernel quantizes activations itself and owns its scales. + return FusedMoEQuantConfig.make("nvfp4", weight_dtype="mxfp4") if mxfp4_backend == Mxfp4MoeBackend.B12X_MXFP4_MXFP8: return mxfp4_mxfp8_moe_quant_config( w1_bias=w1_bias, diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index e1952c0032..68ba62782c 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -50,6 +50,7 @@ class NvFp4MoeBackend(Enum): MARLIN = "MARLIN" HUMMING = "HUMMING" EMULATION = "EMULATION" + FLASHINFER_MOE_EP_CUTEDSL = "FLASHINFER_MOE_EP_CUTEDSL" FLASHINFER_NVFP4_MOE_BACKENDS = [ @@ -148,6 +149,12 @@ def backend_to_kernel_cls( ) return [Nvfp4QuantizationEmulationTritonExperts] + elif backend == NvFp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL: + from vllm.model_executor.layers.fused_moe.experts.flashinfer_moe_ep import ( # noqa: E501 + FlashInferMoeEpExperts, + ) + + return [FlashInferMoeEpExperts] else: raise ValueError(f"Unknown NvFP4 MoE backend: {backend.value}") @@ -160,6 +167,7 @@ def map_nvfp4_backend(runner_backend: MoEBackend) -> NvFp4MoeBackend: "flashinfer_trtllm": NvFp4MoeBackend.FLASHINFER_TRTLLM, "flashinfer_cutlass": NvFp4MoeBackend.FLASHINFER_CUTLASS, "flashinfer_cutedsl": NvFp4MoeBackend.FLASHINFER_CUTEDSL, + "flashinfer_moe_ep_cutedsl": NvFp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL, "flashinfer_b12x": NvFp4MoeBackend.FLASHINFER_B12X, "marlin": NvFp4MoeBackend.MARLIN, "humming": NvFp4MoeBackend.HUMMING, @@ -207,6 +215,7 @@ def select_nvfp4_moe_backend( NvFp4MoeBackend.FLASHINFER_TRTLLM, NvFp4MoeBackend.FLASHINFER_CUTLASS, NvFp4MoeBackend.FLASHINFER_CUTEDSL, + NvFp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL, NvFp4MoeBackend.VLLM_CUTLASS, NvFp4MoeBackend.MARLIN, NvFp4MoeBackend.EMULATION, @@ -281,8 +290,9 @@ def select_nvfp4_moe_backend( f"Model sets swiglu_limit={config.swiglu_limit}, but the " f"explicitly requested moe_backend={runner_backend!r} does " f"not apply the SwiGLU clamp. Use 'flashinfer_trtllm', " - f"'flashinfer_cutlass', 'flashinfer_cutedsl', 'cutlass', " - f"'b12x', 'marlin', or 'humming' instead." + f"'flashinfer_cutlass', 'flashinfer_cutedsl', " + f"'flashinfer_moe_ep_cutedsl', 'cutlass', 'b12x', 'marlin', " + f"or 'humming' instead." ) return _return_or_raise( requested_backend, config, weight_key, activation_key, activation_format @@ -491,6 +501,9 @@ def convert_to_nvfp4_moe_kernel_format( # for other experts - other selection strategies may be used. a13_scale = 1.0 / a13_scale.max().to(torch.float32) a2_scale = 1.0 / a2_scale.max().to(torch.float32) + elif nvfp4_backend == NvFp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL: + # The megakernel transforms the checkpoint tensors itself at init. + pass else: raise ValueError(f"Unknown NvFp4 backend for MoE: {nvfp4_backend}") diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/passthrough.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/passthrough.py new file mode 100644 index 0000000000..a6c541d2b9 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/passthrough.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Prepare/finalize stages that do no expert-parallel communication. + +Selected by ``all2all_backend="passthrough"``: the experts implementation +dispatches, combines and reduces itself, so nothing is left for these stages. +""" + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig +from vllm.model_executor.layers.fused_moe.prepare_finalize.no_dp_ep import ( + MoEPrepareAndFinalizeNoDPEPModular, +) + + +class PassThroughPrepareAndFinalize(MoEPrepareAndFinalizeNoDPEPModular): + """Hands activations and routing straight to experts that own their EP + communication.""" + + def output_is_reduced(self) -> bool: + return True + + def supports_deferred_moe_finalize(self) -> bool: + # The experts apply the top-k weights and combine internally; there + # is no finalize left for a consumer to take over. + return False + + def prepare( + self, + a1: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool = False, + ) -> mk.PrepareResultType: + return a1, None, None, None, None diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py index 643fd75e82..4dd1e7aafe 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py @@ -24,10 +24,15 @@ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import ( from vllm.model_executor.layers.fused_moe.experts.xpu_moe import ( XPUExpertsMxFp4, ) +from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep import ( + is_flashinfer_moe_ep_backend, +) from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( B12X_BACKENDS, + FLASHINFER_MOE_EP_MXFP4_BACKENDS, Mxfp4MoeBackend, + convert_weight_to_mxfp4_moe_kernel_format, make_mxfp4_moe_kernel, make_mxfp4_moe_quant_config, select_mxfp4_moe_backend, @@ -38,7 +43,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( prepare_moe_fp4_layer_for_marlin, ) -from vllm.model_executor.utils import set_weight_attrs +from vllm.model_executor.utils import replace_parameter, set_weight_attrs from vllm.platforms import current_platform logger = init_logger(__name__) @@ -49,12 +54,14 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod): super().__init__(moe) self.group_size = 32 self.mxfp4_backend = Mxfp4MoeBackend.MARLIN + self.experts_cls: type[mk.FusedMoEExperts] | None = None + # Backend selection must match the weight preparation below: CUTLASS - # swizzles scales, b12x and XPU consume checkpoint packing, and Marlin - # repacks weights and scales. + # swizzles scales, b12x and XPU consume checkpoint packing, Marlin + # repacks weights and scales, and the megakernels take the checkpoint + # tensors as they are. self.use_cutlass_mxfp4 = CutlassExpertsMxfp4._supports_current_device() - self.experts_cls: type[mk.FusedMoEExperts] - if moe.moe_backend == "b12x": + if moe.moe_backend == "b12x" or is_flashinfer_moe_ep_backend(moe.moe_backend): self.mxfp4_backend, experts_cls = select_mxfp4_moe_backend(moe) assert experts_cls is not None self.experts_cls = experts_cls @@ -70,6 +77,10 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod): logger.info_once("Using MarlinExperts for MXFP4 MoE") self.experts_cls = MarlinExperts + @property + def supports_eplb(self) -> bool: + return self.mxfp4_backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS + def create_weights( self, layer: torch.nn.Module, @@ -198,6 +209,21 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod): layer.w2_weight_scale = torch.nn.Parameter( torch.stack(swizzled_w2), requires_grad=False ) + elif self.mxfp4_backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS: + w13, w2, w13_scale, w2_scale, _, _ = ( + convert_weight_to_mxfp4_moe_kernel_format( + mxfp4_backend=self.mxfp4_backend, + layer=layer, + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + w13_weight_scale=layer.w13_weight_scale, + w2_weight_scale=layer.w2_weight_scale, + ) + ) + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + replace_parameter(layer, "w13_weight_scale", w13_scale) + replace_parameter(layer, "w2_weight_scale", w2_scale) elif self.mxfp4_backend in B12X_BACKENDS or current_platform.is_xpu(): pass else: @@ -211,6 +237,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod): self.moe_quant_config = self.get_fused_moe_quant_config(layer) if self.moe_quant_config is not None: + assert self.experts_cls is not None self.moe_kernel = make_mxfp4_moe_kernel( moe_quant_config=self.moe_quant_config, moe_config=self.moe, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py index a6f2879d87..dff0c62669 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py @@ -38,6 +38,7 @@ _EPLB_SUPPORTED_NVFP4_BACKENDS = frozenset( NvFp4MoeBackend.FLASHINFER_CUTEDSL, NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED, NvFp4MoeBackend.FLASHINFER_TRTLLM, + NvFp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL, } ) @@ -61,7 +62,6 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): weight_key=kNvfp4Static, activation_key=None if use_a16 else kNvfp4Dynamic, ) - self.use_global_sf = is_global_sf_supported_for_nvfp4_backend( self.nvfp4_backend ) @@ -155,7 +155,8 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): set_weight_attrs(w13_weight_scale_2, extra_weight_attrs) w2_weight_scale_2 = torch.nn.Parameter( - torch.empty(num_experts, dtype=torch.float32), requires_grad=False + torch.empty(num_experts, dtype=torch.float32), + requires_grad=False, ) layer.register_parameter("w2_weight_global_scale", w2_weight_scale_2) extra_weight_attrs.update( @@ -175,7 +176,8 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): set_weight_attrs(w13_input_scale, extra_weight_attrs) w2_input_scale = torch.nn.Parameter( - torch.empty(num_experts, dtype=torch.float32), requires_grad=False + torch.empty(num_experts, dtype=torch.float32), + requires_grad=False, ) layer.register_parameter("w2_input_global_scale", w2_input_scale) extra_weight_attrs.update( diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py index 53e8ca2b01..9db2907a26 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py @@ -4,6 +4,7 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.config.kernel import FLASHINFER_MOE_EP_BACKENDS from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( FusedMoeWeightScaleSupported, @@ -53,6 +54,11 @@ class INCMxfp4MoEMethod(FusedMoEMethodBase): def __init__(self, moe) -> None: super().__init__(moe) + if moe.moe_backend in FLASHINFER_MOE_EP_BACKENDS: + raise ValueError( + f"moe_backend={moe.moe_backend!r} is not supported for AutoRound " + "(INC) MXFP4 MoE checkpoints" + ) self.group_size = 32 # Backend selection must stay consistent with the weight preparation in # process_weights_after_loading / get_fused_moe_quant_config, which use diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 7342b4d630..861d3cb8d6 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + import torch from vllm.logger import init_logger @@ -783,7 +784,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): ) self.moe_kernel.fused_experts.process_weights_after_loading(layer) - def process_weights_after_loading(self, layer): + def process_weights_after_loading(self, layer: RoutedExperts) -> None: if self.mxfp4_backend == Mxfp4MoeBackend.NONE: return diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py index 4c0a8a3ff7..03d4df386a 100644 --- a/vllm/model_executor/layers/quantization/utils/quant_utils.py +++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py @@ -25,6 +25,11 @@ INT4_DTYPE = scalar_types.uint4b8 INT8_DTYPE = scalar_types.uint8b128 +def require_finite_positive_scales(name: str, value: torch.Tensor) -> None: + if not torch.isfinite(value).all() or not (value > 0).all(): + raise ValueError(f"{name} must contain finite positive scales") + + def _dtype_abbr(dtype: torch.dtype | ScalarType) -> str: """Return a stable short name for torch and ScalarType dtypes.""" if isinstance(dtype, ScalarType): diff --git a/vllm/models/deepseek_v4/amd/dspark.py b/vllm/models/deepseek_v4/amd/dspark.py index c3f249fa80..e0c69c0328 100644 --- a/vllm/models/deepseek_v4/amd/dspark.py +++ b/vllm/models/deepseek_v4/amd/dspark.py @@ -11,8 +11,9 @@ ROCm port of ``nvidia/dspark.py``. Follows the same nvidia->amd recipe used for tilelang / triton / torch) instead of calling the tilelang kernels directly, and gate the trailing ``mhc_post`` on ``use_fused_mhc`` (True when AITER or TileLang fused MHC is available; False only on the torch fallback); - * drop the mega-MoE weight path (``make_deepseek_v4_expert_params_mapping`` / - ``use_mega_moe`` / ``finalize_mega_moe_weights`` do not exist in amd/model.py). + * drop the mega-MoE weight path (``make_deepseek_v4_expert_params_mapping``, + ``use_native_mega_moe`` and ``finalize_mega_moe_weights`` do not exist in + amd/model.py). Everything else — the semi-autoregressive drafting hooks, the Markov head, the sliding-window context-KV insert, and the checkpoint ``mtp.*`` weight remap — is diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index 987909702c..54250185c8 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -17,7 +17,7 @@ import torch.nn as nn import vllm.envs as envs from vllm.config import VllmConfig, get_current_vllm_config -from vllm.config.kernel import MEGA_MOE_BACKENDS +from vllm.config.kernel import NATIVE_MEGA_MOE_BACKENDS from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -99,7 +99,9 @@ class DSparkDeepseekV4Model(nn.Module): self.rms_norm_eps = config.rms_norm_eps self.num_hidden_layers = config.num_hidden_layers self.target_layer_ids = tuple(config.dspark_target_layer_ids) - self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS + self.use_native_mega_moe = ( + vllm_config.kernel_config.moe_backend in NATIVE_MEGA_MOE_BACKENDS + ) self.use_sequence_parallel = _use_sequence_parallel(vllm_config) self.num_dspark_layers = getattr(config, "n_mtp_layers", None) or 3 @@ -240,7 +242,7 @@ class DSparkDeepseekV4Model(nn.Module): inputs_embeds = sp_shard(inputs_embeds) input_ids = sp_shard(input_ids) mega_gate_metadata = None - if self.use_mega_moe: + if self.use_native_mega_moe: mega_gate_metadata = prepare_mega_gate_routing_metadata( input_ids, has_hash_routing=False, @@ -437,8 +439,8 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): are skipped here. ``embed_tokens``/``lm_head`` are aliased from the target. """ first_layer = self.model.layers[0] - use_mega_moe = first_layer.ffn.use_mega_moe - if use_mega_moe: + use_native_mega_moe = first_layer.ffn.use_native_mega_moe + if use_native_mega_moe: expert_mapping = make_deepseek_v4_expert_params_mapping( self.config.n_routed_experts ) diff --git a/vllm/models/deepseek_v4/nvidia/fi_moe.py b/vllm/models/deepseek_v4/nvidia/fi_moe.py deleted file mode 100644 index 2e94e75e7d..0000000000 --- a/vllm/models/deepseek_v4/nvidia/fi_moe.py +++ /dev/null @@ -1,397 +0,0 @@ -# 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 7edff8ae02..a004120804 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -11,7 +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.config.kernel import MEGA_MOE_BACKENDS, NATIVE_MEGA_MOE_BACKENDS from vllm.distributed import ( get_ep_group, get_pp_group, @@ -85,10 +85,6 @@ 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 @@ -807,10 +803,11 @@ class DeepseekV4MoE(nn.Module): self.use_sequence_parallel = use_sequence_parallel self.reduce_results = reduce_results 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: + self.use_native_mega_moe = moe_backend in NATIVE_MEGA_MOE_BACKENDS + if ( + self.use_native_mega_moe + and not vllm_config.parallel_config.enable_expert_parallel + ): raise NotImplementedError( "DeepSeek V4 MegaMoE currently requires expert parallel. " "Enable it with --enable-expert-parallel, or pick a different " @@ -832,15 +829,15 @@ class DeepseekV4MoE(nn.Module): self.swiglu_limit = config.swiglu_limit self.renormalize = config.norm_topk_prob self.scoring_func = getattr(config, "scoring_func", "sqrtsoftplus") - if self.use_mega_moe and self.scoring_func != "sqrtsoftplus": + if self.use_native_mega_moe and self.scoring_func != "sqrtsoftplus": raise NotImplementedError( "DeepSeek V4 MegaMoE currently supports sqrtsoftplus routing only." ) - if self.use_mega_moe and not self.renormalize: + if self.use_native_mega_moe and not self.renormalize: raise NotImplementedError( "DeepSeek V4 MegaMoE requires normalized top-k probabilities." ) - if self.use_mega_moe and getattr(config, "expert_dtype", "fp4") != "fp4": + if self.use_native_mega_moe and getattr(config, "expert_dtype", "fp4") != "fp4": raise NotImplementedError( "DeepSeek V4 MegaMoE only supports fp4 experts; got expert_dtype=" f"{config.expert_dtype!r}. Drop --kernel-config moe_backend=" @@ -864,7 +861,9 @@ class DeepseekV4MoE(nn.Module): image_sentinel_lo if getattr(config, "vision_n_layers", 0) > 0 else 0 ) is_hash_moe = extract_layer_index(prefix) < num_hash_layers - self.hash_indices_dtype = torch.int64 if self.use_mega_moe else torch.int32 + self.hash_indices_dtype = ( + torch.int64 if self.use_native_mega_moe else torch.int32 + ) if is_hash_moe: # hash MoE doesn't use e_score_correction_bias # Use randint instead of empty to avoid garbage values causing @@ -908,12 +907,12 @@ class DeepseekV4MoE(nn.Module): hidden_act=config.hidden_act, swiglu_limit=self.swiglu_limit, quant_config=quant_config, - reduce_results=self.use_mega_moe, + reduce_results=self.use_native_mega_moe, is_sequence_parallel=use_sequence_parallel, prefix=f"{prefix}.shared_experts", ) - if self.use_mega_moe: + if self.use_native_mega_moe: self._init_mega_moe_experts(vllm_config, config, prefix) else: self._init_fused_moe_experts(vllm_config, config, quant_config, prefix) @@ -950,28 +949,13 @@ 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. The FlashInfer - # megakernel has no shared-expert fusion, so it keeps the serial path. + # remains tensor-sharded, so retain 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 ) - 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, @@ -983,9 +967,7 @@ 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) + self.experts = DeepseekV4MegaMoEExperts(vllm_config, **expert_kwargs) def _init_fused_moe_experts( self, @@ -1048,7 +1030,7 @@ class DeepseekV4MoE(nn.Module): if bias_vl is not None and input_ids is None: raise ValueError("DeepSeek V4 vision MoE routing requires input_ids.") - if not self.use_mega_moe: + if not self.use_native_mega_moe: return self._forward_fused_moe(hidden_states, input_ids) org_shape = hidden_states.shape @@ -1139,7 +1121,7 @@ class DeepseekV4MoE(nn.Module): return final_hidden_states.view(org_shape) def finalize_mega_moe_weights(self) -> None: - if self.use_mega_moe: + if self.use_native_mega_moe: self.experts.finalize_weights(self.shared_experts) @@ -1179,12 +1161,15 @@ 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 in MEGA_MOE_BACKENDS + moe_backend = vllm_config.kernel_config.moe_backend + moe_needs_token_sharded_input = ( + moe_backend in MEGA_MOE_BACKENDS or parallel_config.data_parallel_size > 1 + ) return ( parallel_config.pipeline_parallel_size == 1 and parallel_config.enable_expert_parallel and parallel_config.tensor_parallel_size > 1 - and (use_mega_moe or parallel_config.data_parallel_size > 1) + and moe_needs_token_sharded_input ) @@ -1425,9 +1410,14 @@ 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 in MEGA_MOE_BACKENDS + self.use_native_mega_moe = ( + vllm_config.kernel_config.moe_backend in NATIVE_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: + if ( + self.use_native_mega_moe + and not vllm_config.parallel_config.enable_expert_parallel + ): raise NotImplementedError( "DeepSeek V4 MegaMoE currently requires expert parallel. " "Enable it with --enable-expert-parallel, or pick a different " @@ -1564,7 +1554,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): assert intermediate_tensors is not None hidden_states = intermediate_tensors["hidden_states"] - if self.use_mega_moe: + if self.use_native_mega_moe: input_ids = input_ids.to(torch.int64) full_num_tokens = positions.shape[0] @@ -1578,7 +1568,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): input_ids = sp_shard(input_ids) mega_gate_metadata = None - if self.use_mega_moe: + if self.use_native_mega_moe: mega_gate_metadata = prepare_mega_gate_routing_metadata( input_ids, has_hash_routing=self.has_local_hash_moe, @@ -1802,7 +1792,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer))) - if first_layer.ffn.use_mega_moe: + if first_layer.ffn.use_native_mega_moe: return make_deepseek_v4_expert_params_mapping(self.config.n_routed_experts) # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) diff --git a/vllm/models/deepseek_v4/nvidia/mtp.py b/vllm/models/deepseek_v4/nvidia/mtp.py index 50886b0afa..64547af038 100644 --- a/vllm/models/deepseek_v4/nvidia/mtp.py +++ b/vllm/models/deepseek_v4/nvidia/mtp.py @@ -20,7 +20,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.config.kernel import NATIVE_MEGA_MOE_BACKENDS from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -95,7 +95,9 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module): self.config = config quant_config = vllm_config.quant_config self.rms_norm_eps = config.rms_norm_eps - self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS + self.use_native_mega_moe = ( + vllm_config.kernel_config.moe_backend in NATIVE_MEGA_MOE_BACKENDS + ) self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -186,7 +188,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module): inputs_embeds ).unsqueeze(-2) mega_gate_metadata = None - if self.use_mega_moe: + if self.use_native_mega_moe: routing_input_ids = input_ids if self.mtp_block.use_sequence_parallel: routing_input_ids = sp_shard(routing_input_ids) @@ -384,7 +386,7 @@ class DeepSeekV4MTP(nn.Module): # Pre-compute expert mapping ONCE. first_layer = next(iter(self.model.layers.values())) - if first_layer.mtp_block.ffn.use_mega_moe: + if first_layer.mtp_block.ffn.use_native_mega_moe: expert_mapping = make_deepseek_v4_expert_params_mapping( self.config.n_routed_experts ) diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index a477d1b5ec..d1534d5248 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/vllm/models/deepseek_v41/nvidia/dspark.py @@ -18,7 +18,7 @@ import torch.nn as nn import vllm.envs as envs from vllm.config import VllmConfig, get_current_vllm_config -from vllm.config.kernel import MEGA_MOE_BACKENDS +from vllm.config.kernel import NATIVE_MEGA_MOE_BACKENDS from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -81,7 +81,9 @@ class DSparkDeepseekV4Model(nn.Module): self.rms_norm_eps = config.rms_norm_eps self.num_hidden_layers = config.num_hidden_layers self.target_layer_ids = tuple(config.dspark_target_layer_ids) - self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS + self.use_native_mega_moe = ( + vllm_config.kernel_config.moe_backend in NATIVE_MEGA_MOE_BACKENDS + ) self.use_sequence_parallel = _use_sequence_parallel(vllm_config) self.num_dspark_layers = ( @@ -214,7 +216,7 @@ class DSparkDeepseekV4Model(nn.Module): inputs_embeds = sp_shard(inputs_embeds) input_ids = sp_shard(input_ids) mega_gate_metadata = None - if self.use_mega_moe: + if self.use_native_mega_moe: mega_gate_metadata = prepare_mega_gate_routing_metadata( input_ids, has_hash_routing=False, @@ -371,7 +373,7 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): are skipped here. ``embed_tokens``/``lm_head`` are aliased from the target. """ first_layer = self.model.layers[0] - use_mega_moe = first_layer.ffn.use_mega_moe + use_native_mega_moe = first_layer.ffn.use_native_mega_moe # Draft MoE layers use the dspark_* expert counts, not the # backbone's (see DeepseekV4MoE and the reference # ModelArgs.get_moe_config). @@ -379,7 +381,7 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): getattr(self.config, "dspark_n_routed_experts", 0) or self.config.n_routed_experts ) - if use_mega_moe: + if use_native_mega_moe: expert_mapping = make_deepseek_v4_expert_params_mapping(n_draft_experts) else: expert_mapping = fused_moe_make_expert_params_mapping( diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index b8fa7754b7..bd91e8f12b 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -11,7 +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.config.kernel import MEGA_MOE_BACKENDS, NATIVE_MEGA_MOE_BACKENDS from vllm.distributed import ( get_engram_dp_size, get_pp_group, @@ -174,7 +174,7 @@ class DeepseekV4MoE(DeepseekV4MoEBase): def defers_finalize(self, num_tokens: int) -> bool: return ( - not self.use_mega_moe + not self.use_native_mega_moe and self.experts.moe_config.should_defer_moe_finalize(num_tokens) ) @@ -244,12 +244,14 @@ 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 in MEGA_MOE_BACKENDS + moe_needs_token_sharded_input = ( + vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS + ) return ( parallel_config.pipeline_parallel_size == 1 and parallel_config.enable_expert_parallel and parallel_config.tensor_parallel_size > 1 - and (use_mega_moe or parallel_config.data_parallel_size > 1) + and (moe_needs_token_sharded_input or parallel_config.data_parallel_size > 1) ) @@ -631,9 +633,14 @@ 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 in MEGA_MOE_BACKENDS + self.use_native_mega_moe = ( + vllm_config.kernel_config.moe_backend in NATIVE_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: + if ( + self.use_native_mega_moe + and not vllm_config.parallel_config.enable_expert_parallel + ): raise NotImplementedError( "DeepSeek V4 MegaMoE currently requires expert parallel. " "Enable it with --enable-expert-parallel, or pick a different " @@ -865,7 +872,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): assert intermediate_tensors is not None hidden_states = intermediate_tensors["hidden_states"] - if self.use_mega_moe: + if self.use_native_mega_moe: input_ids = input_ids.to(torch.int64) # Engram n-gram hashes for the whole (flattened) batch, computed once @@ -1015,7 +1022,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): def _mega_gate_metadata( self, input_ids: torch.Tensor | None ) -> MegaGateRoutingMetadata | None: - if not self.use_mega_moe: + if not self.use_native_mega_moe: return None assert input_ids is not None return prepare_mega_gate_routing_metadata( @@ -1331,7 +1338,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer))) - if first_layer.ffn.use_mega_moe: + if first_layer.ffn.use_native_mega_moe: return make_deepseek_v4_expert_params_mapping(self.config.n_routed_experts) # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) diff --git a/vllm/utils/flashinfer_moe_ep.py b/vllm/utils/flashinfer_moe_ep.py deleted file mode 100644 index 5bdd5cbf01..0000000000 --- a/vllm/utils/flashinfer_moe_ep.py +++ /dev/null @@ -1,339 +0,0 @@ -# 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", -]