| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for the flashinfer moe_ep backend plumbing. |
| 4 | |
| 5 | Everything here runs without a GPU or a flashinfer install: the flashinfer |
| 6 | modules the helpers import lazily are replaced with capture fakes. |
| 7 | """ |
| 8 | |
| 9 | import sys |
| 10 | from dataclasses import dataclass, field |
| 11 | from types import ModuleType, SimpleNamespace |
| 12 | from typing import Any |
| 13 | |
| 14 | import pytest |
| 15 | import torch |
| 16 | |
| 17 | from vllm.config.kernel import ( |
| 18 | FLASHINFER_MOE_EP_BACKENDS, |
| 19 | MEGA_MOE_BACKENDS, |
| 20 | validate_flashinfer_moe_ep_model, |
| 21 | ) |
| 22 | from vllm.utils.flashinfer_moe_ep import ( |
| 23 | _E2M1_LUT, |
| 24 | FI_MOE_EP_BACKEND_SPECS, |
| 25 | _dequant_fp4_ue8m0_gran32, |
| 26 | build_fi_mega_config, |
| 27 | fi_moe_ep_backend_spec, |
| 28 | make_fi_moe_ep_bootstrap, |
| 29 | megakernel_runtime_requirements, |
| 30 | ) |
| 31 | |
| 32 | |
| 33 | @dataclass |
| 34 | class _FakeBootstrapConfig: |
| 35 | world_size: int |
| 36 | rank: int |
| 37 | process_group: Any = None |
| 38 | auto_bootstrap: bool = True |
| 39 | device: int | None = field(default=None, kw_only=True) |
| 40 | |
| 41 | |
| 42 | @dataclass |
| 43 | class _FakeDeepGemmMegaMoeConfig: |
| 44 | intermediate_size: int |
| 45 | top_k: int |
| 46 | activation_clamp: float | None |
| 47 | fast_math: bool |
| 48 | |
| 49 | |
| 50 | @dataclass |
| 51 | class _FakeNvfp4CutedslMegaMoeConfig: |
| 52 | intermediate_size: int |
| 53 | top_k: int |
| 54 | activation_clamp: float | None |
| 55 | fast_math: bool |
| 56 | |
| 57 | |
| 58 | @dataclass |
| 59 | class _FakeMegaConfig: |
| 60 | megakernel: Any |
| 61 | preprocess_weights: bool |
| 62 | quantize_input: bool |
| 63 | |
| 64 | |
| 65 | @pytest.fixture |
| 66 | def fake_flashinfer(monkeypatch): |
| 67 | """Install a minimal fake flashinfer.moe_ep for the lazy imports.""" |
| 68 | moe_ep = ModuleType("flashinfer.moe_ep") |
| 69 | core = ModuleType("flashinfer.moe_ep.core") |
| 70 | runtime = ModuleType("flashinfer.moe_ep.core.runtime") |
| 71 | flashinfer = ModuleType("flashinfer") |
| 72 | fake_attrs: dict[ModuleType, dict[str, Any]] = { |
| 73 | moe_ep: { |
| 74 | "BootstrapConfig": _FakeBootstrapConfig, |
| 75 | "DeepGemmMegaMoeConfig": _FakeDeepGemmMegaMoeConfig, |
| 76 | "Nvfp4CutedslMegaMoeConfig": _FakeNvfp4CutedslMegaMoeConfig, |
| 77 | "MegaConfig": _FakeMegaConfig, |
| 78 | "core": core, |
| 79 | }, |
| 80 | runtime: {"TORCH_DIST": "torch_dist", "NVSHMEM": "nvshmem"}, |
| 81 | flashinfer: {"moe_ep": moe_ep}, |
| 82 | core: {"runtime": runtime}, |
| 83 | } |
| 84 | for mod, attrs in fake_attrs.items(): |
| 85 | for attr, value in attrs.items(): |
| 86 | setattr(mod, attr, value) |
| 87 | |
| 88 | for name, mod in { |
| 89 | "flashinfer": flashinfer, |
| 90 | "flashinfer.moe_ep": moe_ep, |
| 91 | "flashinfer.moe_ep.core": core, |
| 92 | "flashinfer.moe_ep.core.runtime": runtime, |
| 93 | }.items(): |
| 94 | monkeypatch.setitem(sys.modules, name, mod) |
| 95 | return moe_ep |
| 96 | |
| 97 | |
| 98 | def test_fi_backend_strings_are_registered_mega_moe_backends(): |
| 99 | assert set(FI_MOE_EP_BACKEND_SPECS) == FLASHINFER_MOE_EP_BACKENDS |
| 100 | assert FLASHINFER_MOE_EP_BACKENDS < MEGA_MOE_BACKENDS |
| 101 | |
| 102 | |
| 103 | @pytest.mark.parametrize("moe_backend", sorted(FLASHINFER_MOE_EP_BACKENDS)) |
| 104 | def test_fi_moe_ep_backend_rejected_for_non_dsv4(moe_backend): |
| 105 | """An FI moe_ep backend with a non-DSv4 model must fail at config time |
| 106 | instead of silently falling through to the generic FusedMoE path.""" |
| 107 | with pytest.raises(ValueError, match="only supported for DeepSeek-V4"): |
| 108 | validate_flashinfer_moe_ep_model(moe_backend, ["MixtralForCausalLM"]) |
| 109 | |
| 110 | |
| 111 | @pytest.mark.parametrize("moe_backend", sorted(FLASHINFER_MOE_EP_BACKENDS)) |
| 112 | def test_fi_moe_ep_backend_accepted_for_dsv4(moe_backend): |
| 113 | validate_flashinfer_moe_ep_model(moe_backend, ["DeepseekV4ForCausalLM"]) |
| 114 | |
| 115 | |
| 116 | @pytest.mark.parametrize( |
| 117 | "architectures", |
| 118 | [["KimiK3ForConditionalGeneration"], ["MixtralForCausalLM"]], |
| 119 | ) |
| 120 | def test_native_deep_gemm_mega_moe_not_arch_gated(architectures): |
| 121 | """VLLM's own deep_gemm mega path is not DSv4-only (Kimi K3 uses it); |
| 122 | models validate their own constraints at construction time.""" |
| 123 | validate_flashinfer_moe_ep_model("deep_gemm_mega_moe", architectures) |
| 124 | |
| 125 | |
| 126 | def test_non_fi_backend_ignores_architectures(): |
| 127 | validate_flashinfer_moe_ep_model("auto", ["MixtralForCausalLM"]) |
| 128 | |
| 129 | |
| 130 | @pytest.mark.parametrize("moe_backend", sorted(MEGA_MOE_BACKENDS)) |
| 131 | def test_all_mega_backends_get_sequence_parallel_moe(moe_backend): |
| 132 | """Every mega backend must qualify for sequence-parallel MoE at |
| 133 | TP>1/EP: the predicate once matched only the native backend string, |
| 134 | which silently ran the fi backends full-batch with an all-reduce on |
| 135 | every rank — 0.42-0.65x native e2e at TP8.""" |
| 136 | from vllm.models.deepseek_v4.nvidia.model import _use_sequence_parallel |
| 137 | |
| 138 | vllm_config = SimpleNamespace( |
| 139 | parallel_config=SimpleNamespace( |
| 140 | pipeline_parallel_size=1, |
| 141 | enable_expert_parallel=True, |
| 142 | tensor_parallel_size=8, |
| 143 | data_parallel_size=1, |
| 144 | ), |
| 145 | kernel_config=SimpleNamespace(moe_backend=moe_backend), |
| 146 | ) |
| 147 | assert _use_sequence_parallel(vllm_config) |
| 148 | |
| 149 | |
| 150 | def test_fi_moe_ep_backend_spec_kernel_and_nvshmem_contract(): |
| 151 | dg = fi_moe_ep_backend_spec("flashinfer_moe_ep_mega_deep_gemm") |
| 152 | assert dg.megakernel == "deep_gemm_mega" |
| 153 | assert not dg.needs_nvshmem |
| 154 | |
| 155 | cd = fi_moe_ep_backend_spec("flashinfer_moe_ep_mega_cutedsl") |
| 156 | assert cd.megakernel == "nvfp4_cutedsl" |
| 157 | assert cd.needs_nvshmem |
| 158 | |
| 159 | with pytest.raises(ValueError, match="not a flashinfer moe_ep backend"): |
| 160 | fi_moe_ep_backend_spec("deep_gemm_mega_moe") |
| 161 | |
| 162 | |
| 163 | def test_megakernel_runtime_requirements(fake_flashinfer): |
| 164 | dg = megakernel_runtime_requirements( |
| 165 | fi_moe_ep_backend_spec("flashinfer_moe_ep_mega_deep_gemm") |
| 166 | ) |
| 167 | assert dg == frozenset({"torch_dist"}) |
| 168 | |
| 169 | cd = megakernel_runtime_requirements( |
| 170 | fi_moe_ep_backend_spec("flashinfer_moe_ep_mega_cutedsl") |
| 171 | ) |
| 172 | assert cd == frozenset({"torch_dist", "nvshmem"}) |
| 173 | |
| 174 | |
| 175 | def test_bootstrap_pins_the_device_vllm_bound(fake_flashinfer, monkeypatch): |
| 176 | """The runtime must not rederive the device from LOCAL_RANK/rank: under a |
| 177 | remapped CUDA_VISIBLE_DEVICES that ordinal points at the wrong GPU |
| 178 | (CUDA_ERROR_ILLEGAL_ADDRESS in the weight transforms). vLLM passes the |
| 179 | device it already bound via BootstrapConfig.device.""" |
| 180 | import vllm.utils.flashinfer_moe_ep as mod |
| 181 | |
| 182 | pg = object() |
| 183 | monkeypatch.setattr( |
| 184 | mod, |
| 185 | "get_ep_group", |
| 186 | lambda: SimpleNamespace(world_size=4, rank_in_group=2, device_group=pg), |
| 187 | ) |
| 188 | monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 3) |
| 189 | |
| 190 | bootstrap = make_fi_moe_ep_bootstrap() |
| 191 | |
| 192 | assert bootstrap.world_size == 4 |
| 193 | assert bootstrap.rank == 2 |
| 194 | assert bootstrap.process_group is pg |
| 195 | assert bootstrap.auto_bootstrap is False |
| 196 | assert bootstrap.device == 3 |
| 197 | |
| 198 | |
| 199 | def test_build_fi_mega_config_selects_kernel_config(fake_flashinfer): |
| 200 | dg = build_fi_mega_config( |
| 201 | intermediate_size=2048, |
| 202 | top_k=8, |
| 203 | activation_clamp=7.0, |
| 204 | megakernel="deep_gemm_mega", |
| 205 | ) |
| 206 | assert isinstance(dg.megakernel, _FakeDeepGemmMegaMoeConfig) |
| 207 | assert dg.megakernel.intermediate_size == 2048 |
| 208 | assert dg.megakernel.top_k == 8 |
| 209 | assert dg.megakernel.activation_clamp == 7.0 |
| 210 | assert dg.preprocess_weights and dg.quantize_input |
| 211 | |
| 212 | cd = build_fi_mega_config( |
| 213 | intermediate_size=2048, |
| 214 | top_k=8, |
| 215 | activation_clamp=None, |
| 216 | megakernel="nvfp4_cutedsl", |
| 217 | ) |
| 218 | assert isinstance(cd.megakernel, _FakeNvfp4CutedslMegaMoeConfig) |
| 219 | |
| 220 | with pytest.raises(ValueError, match="Unsupported fi_moe_ep megakernel"): |
| 221 | build_fi_mega_config( |
| 222 | intermediate_size=2048, |
| 223 | top_k=8, |
| 224 | activation_clamp=None, |
| 225 | megakernel="deep_gemm", |
| 226 | ) |
| 227 | |
| 228 | |
| 229 | def test_ckpt_uses_nvfp4_experts_reads_moe_quant_algo(): |
| 230 | from vllm.models.deepseek_v4.nvidia.fi_moe import ckpt_uses_nvfp4_experts |
| 231 | |
| 232 | nvfp4 = SimpleNamespace(quant_config=SimpleNamespace(moe_quant_algo="NVFP4")) |
| 233 | assert ckpt_uses_nvfp4_experts(nvfp4) |
| 234 | |
| 235 | mxfp4 = SimpleNamespace(quant_config=SimpleNamespace(moe_quant_algo=None)) |
| 236 | assert not ckpt_uses_nvfp4_experts(mxfp4) |
| 237 | |
| 238 | no_algo = SimpleNamespace(quant_config=SimpleNamespace()) |
| 239 | assert not ckpt_uses_nvfp4_experts(no_algo) |
| 240 | |
| 241 | |
| 242 | def test_dequant_fp4_ue8m0_gran32_decodes_lut_and_scales(): |
| 243 | """One 32-element scale group per row: low nibble is the even element, |
| 244 | high nibble the odd one, ue8m0 scale applies to the whole group.""" |
| 245 | packed = torch.arange(32, dtype=torch.uint8).reshape(2, 16) |
| 246 | sf = torch.tensor([[127], [128]], dtype=torch.uint8) # 2**0, 2**1 |
| 247 | |
| 248 | out = _dequant_fp4_ue8m0_gran32(packed, sf) |
| 249 | |
| 250 | assert out.shape == (2, 32) |
| 251 | assert out.dtype == torch.bfloat16 |
| 252 | expected = torch.empty(2, 32) |
| 253 | for row in range(2): |
| 254 | for col in range(16): |
| 255 | byte = int(packed[row, col]) |
| 256 | expected[row, 2 * col] = _E2M1_LUT[byte & 0x0F] |
| 257 | expected[row, 2 * col + 1] = _E2M1_LUT[byte >> 4] |
| 258 | expected[row] *= 2.0**row |
| 259 | assert torch.equal(out, expected.to(torch.bfloat16)) |