From ff86daa9a4b9b932a0e636c5e8894b32ad3ada0d Mon Sep 17 00:00:00 2001 From: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com> Date: Fri, 23 Oct 2026 17:37:56 +0800 Subject: [PATCH] [Experiment][MoE] Stage SM90 NVFP4 MegaMoE integration Archive the legacy flashinfer.megamoe experiment for H200. The integration is held in favor of FlashInfer #6999. Co-authored-by: Codex Signed-off-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com> --- .buildkite/test_areas/distributed.yaml | 4 + .buildkite/test_areas/kernels.yaml | 2 + docs/serving/expert_parallel_deployment.md | 21 ++- tests/kernels/moe/test_flashinfer_moe_ep.py | 149 +++++++++++++++- tests/models/test_deepseek_v4_fi_moe_ep.py | 134 +++++++++++++- vllm/config/kernel.py | 5 +- .../device_communicators/all2all.py | 24 ++- .../layers/fused_moe/all2all_utils.py | 9 + .../fused_moe/experts/flashinfer_moe_ep.py | 14 +- .../layers/fused_moe/flashinfer_moe_ep.py | 28 ++- .../fused_moe/flashinfer_moe_ep_sm90_nvfp4.py | 163 ++++++++++++++++++ .../layers/fused_moe/modular_kernel.py | 16 +- .../layers/fused_moe/oracle/unquantized.py | 9 +- 13 files changed, 540 insertions(+), 38 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/flashinfer_moe_ep_sm90_nvfp4.py diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 303798f25fb2..6065e785465e 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -428,12 +428,16 @@ steps: - vllm/v1/executor/uniproc_executor.py - vllm/v1/worker/gpu_worker.py - tests/distributed/test_mnnvl_alltoall.py + - tests/kernels/moe/test_flashinfer_moe_ep.py + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/layers/quantization/ commands: # https://github.com/NVIDIA/nccl/issues/1838 - export NCCL_CUMEM_HOST_ENABLE=0 # test with torchrun tp=2 and dp=4 with ep - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep + - pytest -v -s kernels/moe/test_flashinfer_moe_ep.py -k sm90_nvfp4 mirror: amd: label: ":amd: (MI355) Distributed DP + EP" diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 44811a4085c1..0f15ce92cf38 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -906,12 +906,14 @@ steps: - csrc/quantization/w8a8/cutlass/moe/ - csrc/moe/ - tests/kernels/moe + - tests/models/test_deepseek_v4_fi_moe_ep.py - vllm/model_executor/layers/fused_moe/ - vllm/model_executor/layers/quantization/ - vllm/distributed/device_communicators/ - vllm/config commands: - pytest -v -s kernels/moe/test_moe_layer.py + - pytest -v -s models/test_deepseek_v4_fi_moe_ep.py -k sm90_nvfp4 mirror: amd: label: ":amd: (MI355) FusedMoE Layer Kernels" diff --git a/docs/serving/expert_parallel_deployment.md b/docs/serving/expert_parallel_deployment.md index f3b329610216..606223b9b3fd 100644 --- a/docs/serving/expert_parallel_deployment.md +++ b/docs/serving/expert_parallel_deployment.md @@ -27,7 +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) | +| `passthrough` | Any | Internal, not accepted on the CLI: bound automatically for MoE backends that dispatch and combine themselves; excluded unquantized layers retain framework collectives | `--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 @@ -80,6 +80,25 @@ vllm serve deepseek-ai/DeepSeek-V3-0324 \ --enable-expert-parallel # Enable expert parallelism ``` +### Hopper NVFP4 MegaMoE + +On H20/H200 with four or eight NVLink-connected GPUs, ModelOpt NVFP4 checkpoints +with a SwiGLU clamp (`swiglu_limit`) can run FlashInfer's SM90 MegaMoE (NVFP4 +weights, FP8 activations) with `--moe-backend flashinfer_moe_ep_cutedsl`. It needs +a FlashInfer build that provides [`flashinfer.megamoe`](https://github.com/flashinfer-ai/flashinfer/tree/main/flashinfer/megamoe) and does not support EPLB, +microbatching/dual batch overlap or runtime weight transfer. Unquantized MoE layers, such as +BF16 MTP experts, keep the default backend. +Gate/up must share each expert's global weight scale; local scales must remain +exactly representable during weight preparation. +This backend is opt-in: `auto` keeps its existing selection, and unsupported +explicit configurations fail at startup. + +```bash +vllm serve nvidia/GLM-5.3-Flash-NVFP4 \ + --data-parallel-size 8 --enable-expert-parallel \ + --moe-backend flashinfer_moe_ep_cutedsl +``` + ## Multi-Node Deployment For multi-node deployment, use the DeepEP communication kernel with one of two modes (see [Backend Selection Guide](#backend-selection-guide) above). diff --git a/tests/kernels/moe/test_flashinfer_moe_ep.py b/tests/kernels/moe/test_flashinfer_moe_ep.py index 7ce3124ae144..2f03ef20252e 100644 --- a/tests/kernels/moe/test_flashinfer_moe_ep.py +++ b/tests/kernels/moe/test_flashinfer_moe_ep.py @@ -1,6 +1,6 @@ # 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. +"""End-to-end tests for the FlashInfer MoE-EP adapters on Hopper and Blackwell. Each test spawns one process per GPU, bootstraps vLLM's expert-parallel groups, builds a ``FlashInferMoeEp`` adapter from bf16 expert weights at the @@ -367,3 +367,150 @@ def destroy(self) -> None: else moe_ep.DeepGemmMegaMoeConfig ) assert isinstance(mega_config.megakernel, expected) + + +@torch.inference_mode() +def _sm90_nvfp4_mixed_moe(env: dict[str, str], world: int, tp_size: int) -> None: + from vllm.forward_context import get_forward_context, set_forward_context + from vllm.model_executor.layers.fused_moe import FusedMoEFactory + from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4Config + from vllm.v1.worker.workspace import init_workspace_manager + + set_env_vars_and_device(env) + rank = torch.distributed.get_rank() + device = torch.device("cuda", rank) + dp_size = world // tp_size + dp_rank = rank // tp_size + cfg = VllmConfig() + cfg.parallel_config.tensor_parallel_size = tp_size + cfg.parallel_config.data_parallel_size = dp_size + cfg.parallel_config.data_parallel_rank = dp_rank + cfg.parallel_config.enable_expert_parallel = True + cfg.parallel_config.all2all_backend = "passthrough" + cfg.kernel_config.moe_backend = "flashinfer_moe_ep_cutedsl" + cfg.scheduler_config.max_num_batched_tokens = 512 * tp_size + init_workspace_manager(device) + hidden, intermediate, experts, topk = 4096, 2048, 288, 8 + with set_current_vllm_config(cfg): + ensure_model_parallel_initialized(tp_size, 1) + for quantized in (True, False): + with torch.device(device): + layer = FusedMoEFactory( + num_experts=experts, + top_k=topk, + hidden_size=hidden, + intermediate_size=intermediate, + params_dtype=torch.bfloat16, + quant_config=( + ModelOptNvFp4Config(is_checkpoint_nvfp4_serialized=True) + if quantized + else None + ), + swiglu_limit=10.0, + is_sequence_parallel=tp_size > 1, + prefix=f"operator_{quantized}", + ) + routed = layer.routed_experts + local = experts // world + routed.w13_weight.zero_() + routed.w2_weight.zero_() + for i in range(local): + e = rank * local + i + w1 = 2.0 / (1 + e % 2) + w2 = 1.75 * (1 + e % 3) + routed.w13_weight[i, 0, 0] = 4 if quantized else w1 + routed.w13_weight[i, intermediate, 0] = 4 if quantized else w1 + routed.w2_weight[i, e, 0] = 4 if quantized else w2 + if quantized: + routed.w13_weight_scale.fill_(128) + routed.w2_weight_scale.fill_(224) + for i in range(local): + e = rank * local + i + routed.w13_weight_scale_2[i].fill_(1 / (128 * (1 + e % 2))) + routed.w2_weight_scale_2[i].fill_((1 + e % 3) / 256) + routed.w13_input_scale.fill_(1) + routed.w2_input_scale.fill_(1) + routed.quant_method.process_weights_after_loading(routed) + + def case(routed, lengths, graph=False, quantized=quantized): + n = (lengths[dp_rank] + tp_size - 1) // tp_size + x = torch.zeros(n, hidden, dtype=torch.bfloat16, device=device) + x[:, 0] = torch.tensor([2.0, 4.0, 8.0], device=device)[ + torch.arange(n, device=device) % 3 + ] + x[:, 1] = 7.0 + ids = ( + ( + (torch.arange(n, device=device) + rank * 7)[:, None] + + torch.arange(topk, device=device)[None, :] * (experts // topk) + ) + .remainder(experts) + .to(torch.int32) + ) + weights = ( + torch.tensor([2.0 ** -(i + 1) for i in range(topk)], device=device) + .expand(n, -1) + .contiguous() + ) + expected = torch.zeros(n, hidden, device=device) + for j in range(topk): + e = ids[:, j].long() + g = (x[:, 0].float() * 2 / (1 + e % 2)).clamp(max=10) + value = torch.nn.functional.silu(g) * g + value *= weights[:, j] + if quantized: + # FI's intermediate E4M3 quantizer uses UE8M0 scales. + scale = torch.exp2(torch.ceil(torch.log2(value.abs() / 448))) + value = (value / scale).to(torch.float8_e4m3fn).float() * scale + value *= 1.75 * (1 + e % 3) + expected.scatter_add_(1, e[:, None], value[:, None]) + with set_forward_context( + None, + cfg, + num_tokens=lengths[dp_rank], + num_tokens_across_dp=torch.tensor(lengths, dtype=torch.int32), + ): + metadata = get_forward_context().dp_metadata + assert metadata is not None + with metadata.sp_local_sizes(tp_size): + y = routed.forward_modular(x, weights, ids) + torch.accelerator.synchronize() + if graph: + torch.distributed.barrier() + captured = torch.cuda.CUDAGraph() + with torch.cuda.graph(captured): + y = routed.forward_modular(x, weights, ids) + captured.replay() + captured.replay() + torch.accelerator.synchronize() + assert torch.isfinite(y).all() + torch.testing.assert_close( + y, expected.to(torch.bfloat16), rtol=0.008, atol=0.001 + ) + torch.distributed.barrier() + if rank == 0: + print( + f"PASS NVFP4={quantized} lengths={lengths} graph={graph}", + flush=True, + ) + + for n in (1, 16, 64, 257): + case(routed, [n * tp_size] * dp_size) + case(routed, [0] + [3 * i * tp_size for i in range(1, dp_size)]) + case(routed, [16 * tp_size] * dp_size, graph=True) + case(routed, [257 * tp_size] * dp_size, graph=True) + del routed, layer + torch.accelerator.memory.empty_cache() + + +@pytest.mark.skipif( + not current_platform.is_device_capability(90), reason="Requires SM90" +) +@pytest.mark.parametrize("world_size,tp_size", [(4, 1), (8, 1), (4, 2)]) +def test_sm90_nvfp4_and_bf16_mtp_experts(world_size: int, tp_size: int): + if torch.cuda.get_device_properties().multi_processor_count not in (78, 132): + pytest.skip("FlashInfer SM90 NVFP4 requires 78 or 132 SMs") + if torch.accelerator.device_count() < world_size: + pytest.skip(f"Requires {world_size} NVLink-connected Hopper GPUs") + pytest.importorskip("flashinfer.megamoe") + distributed_run(_sm90_nvfp4_mixed_moe, world_size, tp_size) diff --git a/tests/models/test_deepseek_v4_fi_moe_ep.py b/tests/models/test_deepseek_v4_fi_moe_ep.py index d98bb26a2a3b..e5f5d1c9f2a6 100644 --- a/tests/models/test_deepseek_v4_fi_moe_ep.py +++ b/tests/models/test_deepseek_v4_fi_moe_ep.py @@ -66,6 +66,91 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires a Hopper GPU") +def test_sm90_nvfp4_supports_the_glm_geometry(): + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.platforms import current_platform + + if not current_platform.is_device_capability(90): + pytest.skip("Requires a Hopper GPU") + moe = make_dummy_moe_config( + num_experts=288, + num_local_experts=36, + experts_per_token=8, + hidden_dim=4096, + intermediate_size=2048, + ) + moe.moe_backend = FLASHINFER_MOE_EP_CUTEDSL + moe.swiglu_limit = 10.0 + moe.moe_parallel_config = replace( + moe.moe_parallel_config, + use_ep=True, + ep_size=8, + dp_size=8, + all2all_backend=PASSTHROUGH_ALL2ALL_BACKEND, + ) + cfg = VllmConfig() + with set_current_vllm_config(cfg): + if torch.cuda.get_device_properties().multi_processor_count not in (78, 132): + assert not _is_supported(moe, kNvfp4Static, kNvfp4Dynamic)[0] + return + assert _is_supported(moe, kNvfp4Static, kNvfp4Dynamic) == (True, "") + cfg.parallel_config.ubatch_size = 2 + assert not cfg.parallel_config.enable_dbo + assert not _is_supported(moe, kNvfp4Static, kNvfp4Dynamic)[0] + cfg.parallel_config.ubatch_size = 1 + assert not _is_supported(moe, kMxfp4Static, None)[0] + moe.moe_parallel_config.ep_size = 2 + assert not _is_supported(moe, kNvfp4Static, kNvfp4Dynamic)[0] + moe.moe_parallel_config.ep_size = 8 + moe.hidden_dim = 4352 + assert not _is_supported(moe, kNvfp4Static, kNvfp4Dynamic)[0] + + +def test_sm90_nvfp4_marlin_byte_order(): + from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep_sm90_nvfp4 import ( + nvfp4_to_marlin, + ) + + packed = torch.tensor( + [0x10, 0x32, 0x54, 0x76, 0x98, 0xBA, 0xDC, 0xFE], dtype=torch.uint8 + ).view(1, 1, 8) + expected = torch.tensor( + [0x04, 0x15, 0x26, 0x37, 0x8C, 0x9D, 0xAE, 0xBF], dtype=torch.uint8 + ).view(1, 1, 8) + assert torch.equal(nvfp4_to_marlin(packed), expected) + + +@pytest.mark.parametrize("as_bytes", [False, True]) +def test_sm90_nvfp4_scale_rebalance_preserves_dequantization(as_bytes): + from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep_sm90_nvfp4 import ( + rebalance_nvfp4_scales, + ) + + scales = torch.tensor([0.5, 32, 224, 448], dtype=torch.float8_e4m3fn) + scales = scales.view(1, 2, 2).repeat(2, 1, 1) + global_scales = torch.tensor([0.5, 2.0]) + encoded, adjusted = rebalance_nvfp4_scales( + scales.view(torch.uint8) if as_bytes else scales, global_scales + ) + assert encoded.dtype == torch.uint8 + assert encoded.view(torch.float8_e4m3fn).float().amax() <= 448 / 6 + assert torch.equal( + encoded.view(torch.float8_e4m3fn).float() * adjusted[:, None, None], + scales.float() * global_scales[:, None, None], + ) + + +def test_sm90_nvfp4_rejects_lossy_scale_rebalance(): + from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep_sm90_nvfp4 import ( + rebalance_nvfp4_scales, + ) + + scales = torch.tensor([2**-9, 448], dtype=torch.float8_e4m3fn).view(1, 1, 2) + with pytest.raises(ValueError, match="exactly"): + rebalance_nvfp4_scales(scales, torch.ones(1)) + + 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) @@ -169,15 +254,17 @@ def _is_supported(moe, weight_key, activation_key) -> tuple[bool, str]: @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), - ) +def megakernel_host(monkeypatch: pytest.MonkeyPatch): + """A Blackwell host without DBO/EPLB/weight transfer.""" + from vllm.config import VllmConfig + from vllm.platforms.interface import DeviceCapability + + config = VllmConfig() monkeypatch.setattr(fi_ep, "get_current_vllm_config", lambda: config) monkeypatch.setattr( - FlashInferMoeEpExperts, "_supports_current_device", staticmethod(lambda: True) + fi_ep.current_platform, + "get_device_capability", + lambda *args, **kwargs: DeviceCapability(10, 0), ) return config @@ -347,6 +434,39 @@ def test_passthrough_all2all_backend_selects_the_pass_through_stages(): assert isinstance(prepare_finalize, PassThroughPrepareAndFinalize) +def test_passthrough_keeps_dispatch_for_unquantized_mtp_experts(): + from vllm.model_executor.layers.fused_moe.prepare_finalize.naive_dp_ep import ( + MoEPrepareAndFinalizeNaiveDPEPModular, + ) + + moe = make_dummy_moe_config() + moe.moe_parallel_config = replace( + moe.moe_parallel_config, + use_ep=True, + ep_size=8, + dp_size=8, + all2all_backend=PASSTHROUGH_ALL2ALL_BACKEND, + ) + stages = maybe_make_prepare_finalize(moe, FusedMoEQuantConfig.make()) + assert isinstance(stages, MoEPrepareAndFinalizeNaiveDPEPModular) + + +def test_passthrough_tp_only_bf16_does_not_dispatch_across_dp(): + from vllm.model_executor.layers.fused_moe.prepare_finalize.no_dp_ep import ( + MoEPrepareAndFinalizeNoDPEPModular, + ) + + moe = make_dummy_moe_config() + moe.moe_parallel_config = replace( + moe.moe_parallel_config, + use_ep=True, + ep_size=4, + all2all_backend=PASSTHROUGH_ALL2ALL_BACKEND, + ) + stages = maybe_make_prepare_finalize(moe, FusedMoEQuantConfig.make()) + assert isinstance(stages, MoEPrepareAndFinalizeNoDPEPModular) + + def test_passthrough_all2all_backend_is_not_a_cli_choice(): """Users cannot pick the passthrough all2all backend explicitly.""" parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 8b5fc5661844..9ad53735beee 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -312,9 +312,10 @@ 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 + - "flashinfer_moe_ep_cutedsl": Use FlashInfer's MoE-EP mega-kernel with NVFP4 weights (MXFP4 checkpoints are requantized at - load); requires Blackwell, expert parallelism, and NVSHMEM + load on Blackwell). Requires expert parallelism; Blackwell uses NVSHMEM, + Hopper NVFP4 uses Torch symmetric memory on four or eight NVLink GPUs - "flashinfer_b12x": Use FlashInfer CuteDSL fused MoE for SM12x (RTX Pro 6000 / DGX Spark) - "b12x": Use b12x FP4 MoE kernels on SM12x diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 74ee703a095b..53df1430dcad 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -38,20 +38,6 @@ 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). @@ -200,6 +186,16 @@ def destroy(self): pass +class PassThroughAll2AllManager(AgRsAll2AllManager): + """MegaMoE owns communication; BF16 excluded layers retain AG/RS fallback.""" + + def get_handle(self, kwargs): + raise RuntimeError( + "passthrough has no all2all handle: dispatch and combine run " + "inside the MoE backend." + ) + + class DeepEPAll2AllManagerBase(All2AllManagerBase): """All2All communication based on DeepEP High-Throughput kernels.""" diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 3fbeebaaa88d..ed328e75fe17 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -197,6 +197,15 @@ def maybe_make_prepare_finalize( input_dtype: torch.dtype | None = None, ) -> FusedMoEPrepareAndFinalize | None: if moe.moe_parallel_config.use_passthrough_all2all: + if quant_config is not None and quant_config.quant_dtype is None: + # Excluded/BF16 MTP experts still need the framework's collectives. + if not moe.moe_parallel_config.use_all2all_kernels and moe.dp_size == 1: + return make_moe_prepare_and_finalize_no_dp_ep(use_monolithic) + return make_moe_prepare_and_finalize_naive_dp_ep( + is_sequence_parallel=moe.moe_parallel_config.is_sequence_parallel, + num_dispatchers=moe.moe_parallel_config.ep_size, + use_monolithic=use_monolithic, + ) # The experts dispatch and combine themselves, in every EP topology. return PassThroughPrepareAndFinalize() 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 index 4bfb45bfad25..add72547fd47 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_moe_ep.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_moe_ep.py @@ -28,6 +28,9 @@ supports_current_device, validate_flashinfer_moe_ep_layer, ) +from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep_sm90_nvfp4 import ( + FlashInferSm90Nvfp4MoeEp, +) from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceNoOP, ) @@ -37,6 +40,7 @@ kNvfp4Static, require_finite_positive_scales, ) +from vllm.platforms import current_platform if TYPE_CHECKING: pass @@ -60,7 +64,7 @@ def __init__( quant_config: FusedMoEQuantConfig, ) -> None: super().__init__(moe_config, quant_config) - self._adapter: FlashInferMoeEp | None = None + self._adapter: FlashInferMoeEp | FlashInferSm90Nvfp4MoeEp | None = None @staticmethod def is_supported_config( @@ -140,7 +144,13 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: epilogue alphas. """ validate_flashinfer_moe_ep_layer(layer) # type: ignore[arg-type] - adapter = FlashInferMoeEp( + adapter_cls = ( + FlashInferSm90Nvfp4MoeEp + if current_platform.is_device_capability(90) + and self.quant_config.use_nvfp4_w4a4 + else FlashInferMoeEp + ) + adapter = adapter_cls( self.moe_config, FlashInferMoeEpWeights( w13=layer.w13_weight, diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py b/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py index 63345a417438..d6179c868433 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py @@ -37,7 +37,7 @@ kMxfp4Static: "mxfp4", } -_SUPPORTED_CAPABILITIES = frozenset({(10, 0), (10, 3)}) +_SUPPORTED_CAPABILITIES = frozenset({(9, 0), (10, 0), (10, 3)}) def supports_current_device() -> bool: @@ -155,10 +155,32 @@ def flashinfer_moe_ep_unsupported_reasons( 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 vllm_config.parallel_config.use_ubatching: + unsupported.append("microbatching/dual batch overlap") if spec.kernel == "deep_gemm" and vllm_config.parallel_config.enable_eplb: unsupported.append("EPLB") + is_sm90 = current_platform.get_device_capability() == (9, 0) + if is_sm90 and weight_key == kMxfp4Static: + unsupported.append("SM90 MXFP4 weights") + if is_sm90 and weight_key == kNvfp4Static: + if torch.cuda.get_device_properties().multi_processor_count not in (78, 132): + unsupported.append("SM90 NVFP4 requires 78 or 132 SMs") + if spec.kernel != "cutedsl": + unsupported.append("SM90 requires NVFP4 CuTeDSL") + if moe.moe_parallel_config.ep_size not in (4, 8): + unsupported.append("SM90 NVFP4 requires EP4 or EP8") + if moe.swiglu_limit is None: + unsupported.append("SM90 NVFP4 requires swiglu_limit") + if ( + moe.hidden_dim % (512 if moe.hidden_dim > 4096 else 256) + or moe.intermediate_size_per_partition % 128 + or 2 * moe.intermediate_size_per_partition > moe.hidden_dim + ): + unsupported.append("SM90 NVFP4 expert dimensions") + if vllm_config.parallel_config.enable_eplb: + unsupported.append("SM90 NVFP4 EPLB") + if moe.routing_method is RoutingMethodType.DeepseekV4: + unsupported.append("SM90 NVFP4 input router-weight placement") return tuple(unsupported) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep_sm90_nvfp4.py b/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep_sm90_nvfp4.py new file mode 100644 index 000000000000..b01bc840e27c --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep_sm90_nvfp4.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Hopper NVFP4 weight conversion and execution through FlashInfer MegaMoE.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +from weakref import WeakValueDictionary + +import torch + +from vllm.distributed import get_ep_group +from vllm.forward_context import get_forward_context, is_forward_context_available + +if TYPE_CHECKING: + from flashinfer.megamoe import MegaMoESymmBuffer + + from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig + from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep import ( + FlashInferMoeEpEpilogue, + FlashInferMoeEpWeights, + ) + +_BUFFERS: WeakValueDictionary[tuple, MegaMoESymmBuffer] = WeakValueDictionary() +# Use fused kernels while split-family serving hangs are investigated. +_KERNEL_FAMILY = "fused" + + +def nvfp4_to_marlin(packed: torch.Tensor) -> torch.Tensor: + # Low-nibble-first checkpoint bytes -> K[b] (high), K[b + 4] (low). + words = packed.view(torch.uint8).reshape(*packed.shape[:-1], -1, 4) + codes = torch.stack((words & 15, words >> 4), dim=-1).flatten(-2) + return ((codes[..., :4] << 4) | codes[..., 4:]).reshape(packed.shape) + + +def rebalance_nvfp4_scales( + scale: torch.Tensor, global_scale: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Move powers of two into the globals without changing decoded weights.""" + local = scale.view(torch.float8_e4m3fn).float() + # The decoder represents FP4 * local scale in one E4M3 byte. + shift = torch.exp2( + torch.ceil(torch.log2(local.amax(dim=(1, 2)) / (448.0 / 6))).clamp_min(0) + ) + local = local / shift[:, None, None] + encoded = local.to(torch.float8_e4m3fn) + if not torch.equal(encoded.float(), local): + raise ValueError( + "NVFP4 scales cannot be rebalanced into the SM90 decoder range exactly" + ) + return encoded.view(torch.uint8), (global_scale.float() * shift).contiguous() + + +def _get_buffer(group, num_experts, capacity, top_k, hidden, intermediate, clamp): + key = (group, num_experts, capacity, top_k, hidden, intermediate, clamp) + if (buffer := _BUFFERS.get(key)) is not None: + return buffer + try: + from flashinfer.megamoe import MegaMoESymmBuffer + except ImportError as err: + raise RuntimeError( + "SM90 NVFP4 requires a FlashInfer source build with flashinfer.megamoe" + ) from err + buffer = MegaMoESymmBuffer( + group, num_experts, capacity, top_k, hidden, intermediate, clamp + ) + buffer.prewarm(kernel_family=_KERNEL_FAMILY) + _BUFFERS[key] = buffer + return buffer + + +class FlashInferSm90Nvfp4MoeEp: + def __init__( + self, + moe: FusedMoEConfig, + weights: FlashInferMoeEpWeights, + epilogue: FlashInferMoeEpEpilogue, + *, + apply_topk_in_fc1: bool, + ) -> None: + from flashinfer.megamoe import transform_nvfp4_weights_for_mega_moe + + if apply_topk_in_fc1: + raise ValueError("SM90 NVFP4 MegaMoE does not weight expert inputs") + assert weights.w13_scale is not None and weights.w2_scale is not None + assert epilogue.fc1_alpha is not None and epilogue.fc2_alpha is not None + assert moe.swiglu_limit is not None + s1, self.g1 = rebalance_nvfp4_scales(weights.w13_scale, epilogue.fc1_alpha) + s2, self.g2 = rebalance_nvfp4_scales(weights.w2_scale, epilogue.fc2_alpha) + self.w1, self.w2 = transform_nvfp4_weights_for_mega_moe( + (nvfp4_to_marlin(weights.w13), s1), + (nvfp4_to_marlin(weights.w2), s2), + ) + self._empty_scale = torch.empty(0, dtype=torch.uint8, device=self.w1.device) + self.moe = moe + self.buffer = _get_buffer( + get_ep_group().device_group, + moe.num_experts, + moe.max_num_tokens, + moe.experts_per_token, + moe.hidden_dim, + moe.intermediate_size_per_partition, + float(moe.swiglu_limit), + ) + + def __call__(self, hidden_states, topk_ids, topk_weights): + from flashinfer.megamoe import nvfp4_mega_moe + + n = hidden_states.shape[0] + m = n + if is_forward_context_available(): + metadata = get_forward_context().dp_metadata + if metadata is not None: + if metadata.local_sizes is not None: + m = max(m, max(metadata.local_sizes)) + else: + total = int(metadata.num_tokens_across_dp_cpu.max()) + m = max(m, (total + self.moe.sp_size - 1) // self.moe.sp_size) + if m == 0: + return hidden_states.new_empty((0, self.moe.hidden_dim)) + if m > self.moe.max_num_tokens: + raise ValueError(f"SM90 MegaMoE M={m} exceeds {self.moe.max_num_tokens}") + b = self.buffer + if n: + torch.ops._C.per_token_group_fp8_quant( + hidden_states.contiguous(), + b.x[:n], + b.x_sf[:n], + 128, + 1e-10, + -448.0, + 448.0, + False, + False, + False, + ) + b.topk_idx[:n].copy_(topk_ids) + b.topk_weights[:n].copy_(topk_weights) + if n < m: + b.x[n:m].zero_() + b.x_sf[n:m].fill_(1) + b.topk_idx[n:m].fill_(-1) + b.topk_weights[n:m].zero_() + b.commit_staged(m) + output = hidden_states.new_empty((m, self.moe.hidden_dim)) + nvfp4_mega_moe( + output, + self.w1, + self.w2, + b, + num_tokens=m, + l1_global_scales=self.g1, + l2_global_scales=self.g2, + kernel_family=_KERNEL_FAMILY, + ) + return output[:n] + + def warmup(self) -> None: + pass # The shared buffer prewarms all selectable plans once at creation. + + def kernel_weights(self): + # Scales are embedded in the packed weight tensors. + return self.w1, self._empty_scale, self.w2, self._empty_scale diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 7f9cb0616735..2624bd89f28b 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1325,7 +1325,14 @@ def _fused_experts( # kernels. CUDAGraph compatible all2all kernels like the DeepEP # low-latency kernels are always batched and can never run into # the tensor.numel() == 0 case. - if M_full == 0: + from vllm.model_executor.layers.fused_moe.prepare_finalize.passthrough import ( + PassThroughPrepareAndFinalize, + ) + + # Internally communicating experts must participate with no local tokens. + if M_full == 0 and not isinstance( + self.prepare_finalize, PassThroughPrepareAndFinalize + ): out_dim = self.fused_experts.output_hidden_dim(a1q.size(-1)) return a1q.new_empty((*a1q.shape[:-1], out_dim), dtype=in_dtype) @@ -1333,8 +1340,8 @@ def _fused_experts( self._allocate_buffers( in_dtype, a1q.device, - M_full, - M_full, + max(1, M_full), + max(1, M_full), N, K, top_k, @@ -1345,6 +1352,9 @@ def _fused_experts( ) ) + if M_full == 0: + fused_out = fused_out[:0] + use_output_alias = ( not fused_out_is_specialized and output_alias is not None diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 2f05074c643c..f7bc575aed96 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -9,7 +9,7 @@ import vllm.envs as envs import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm._aiter_ops import rocm_aiter_ops -from vllm.config.kernel import MoEBackend +from vllm.config.kernel import FLASHINFER_MOE_EP_BACKENDS, MoEBackend from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.all2all_utils import ( @@ -303,10 +303,9 @@ def _return_or_raise( ) runner_backend = moe_config.moe_backend - # 'humming' and 'marlin' are quantization-only; an unquantized layer (e.g. - # excluded via modules_to_not_convert, or an unquantized MTP draft of a - # quantized model) falls through to auto instead of erroring. - if runner_backend not in ["auto", "humming", "marlin"]: + # Quantization-only backends fall through to auto for excluded layers + # and unquantized MTP experts. + if runner_backend not in {"auto", "humming", "marlin", *FLASHINFER_MOE_EP_BACKENDS}: requested_backend = map_unquantized_backend(runner_backend) if ( activation_format == mk.FusedMoEActivationFormat.BatchedExperts