| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """FlashInfer ``moe_ep`` helpers for DeepSeek V4 vLLM integration.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | from dataclasses import dataclass |
| 8 | from typing import TYPE_CHECKING, Any |
| 9 | |
| 10 | import torch |
| 11 | import torch.nn as nn |
| 12 | |
| 13 | from vllm.config.kernel import FLASHINFER_MOE_EP_BACKENDS |
| 14 | from vllm.distributed import get_ep_group |
| 15 | from vllm.platforms import current_platform |
| 16 | |
| 17 | if TYPE_CHECKING: |
| 18 | from flashinfer.moe_ep import BootstrapConfig, MoEEpMegaLayer |
| 19 | |
| 20 | from vllm.config import VllmConfig |
| 21 | |
| 22 | |
| 23 | # Every mega kernel is Blackwell-only, so the arch is a property of the family |
| 24 | # rather than of a particular backend: validate it against the live device |
| 25 | # instead of encoding it in the backend name. An explicit allowlist so new |
| 26 | # archs (SM110/SM120, Rubin) fail loudly until flashinfer supports them. |
| 27 | FI_MOE_EP_SUPPORTED_CAPABILITIES = frozenset({(10, 0), (10, 3)}) |
| 28 | |
| 29 | |
| 30 | @dataclass(frozen=True) |
| 31 | class FiMoeEpBackendSpec: |
| 32 | """Static properties of one ``flashinfer_moe_ep_*`` backend string. |
| 33 | |
| 34 | The backend names the kernel *family*; the arch comes from the device and |
| 35 | the weight handling from the checkpoint, so this only has to carry which |
| 36 | megakernel to build and whether it needs NVSHMEM in the runtime set. |
| 37 | """ |
| 38 | |
| 39 | megakernel: str |
| 40 | needs_nvshmem: bool |
| 41 | |
| 42 | |
| 43 | FI_MOE_EP_BACKEND_SPECS: dict[str, FiMoeEpBackendSpec] = { |
| 44 | # Consumes an MXFP4 checkpoint verbatim (e2m1 weights + E8M0 per-32 |
| 45 | # scales) -- the same recipe the native deep_gemm mega path uses. |
| 46 | "flashinfer_moe_ep_mega_deep_gemm": FiMoeEpBackendSpec( |
| 47 | megakernel="deep_gemm_mega", |
| 48 | needs_nvshmem=False, |
| 49 | ), |
| 50 | # The checkpoint picks the weight path, not the kernel: an NVFP4 |
| 51 | # checkpoint is consumed prequantized (no round trip), while MXFP4 weights |
| 52 | # are dequantized to bf16 and requantized. See |
| 53 | # models/deepseek_v4/nvidia/fi_moe.py:ckpt_uses_nvfp4_experts. |
| 54 | "flashinfer_moe_ep_mega_cutedsl": FiMoeEpBackendSpec( |
| 55 | megakernel="nvfp4_cutedsl", |
| 56 | needs_nvshmem=True, |
| 57 | ), |
| 58 | } |
| 59 | |
| 60 | assert set(FI_MOE_EP_BACKEND_SPECS) == FLASHINFER_MOE_EP_BACKENDS |
| 61 | |
| 62 | _FI_RUNTIME_HANDLE: Any = None |
| 63 | |
| 64 | |
| 65 | def is_fi_moe_ep_backend(moe_backend: str) -> bool: |
| 66 | return moe_backend in FI_MOE_EP_BACKEND_SPECS |
| 67 | |
| 68 | |
| 69 | def fi_moe_ep_backend_spec(moe_backend: str) -> FiMoeEpBackendSpec: |
| 70 | try: |
| 71 | return FI_MOE_EP_BACKEND_SPECS[moe_backend] |
| 72 | except KeyError: |
| 73 | raise ValueError( |
| 74 | f"{moe_backend!r} is not a flashinfer moe_ep backend; expected " |
| 75 | f"one of {sorted(FI_MOE_EP_BACKEND_SPECS)}" |
| 76 | ) from None |
| 77 | |
| 78 | |
| 79 | def validate_fi_moe_ep_config(vllm_config: VllmConfig) -> None: |
| 80 | """Config-time checks for the mega-MoE backends, native and flashinfer.""" |
| 81 | moe_backend = vllm_config.kernel_config.moe_backend |
| 82 | if not is_fi_moe_ep_backend(moe_backend): |
| 83 | return |
| 84 | |
| 85 | # flashinfer validates the arch too, but not until the layer constructor |
| 86 | # runs during weight load; check here so the error names the flag the user |
| 87 | # actually typed. |
| 88 | capability = current_platform.get_device_capability() |
| 89 | if capability is not None: |
| 90 | cc = (capability.major, capability.minor) |
| 91 | if cc not in FI_MOE_EP_SUPPORTED_CAPABILITIES: |
| 92 | supported = ", ".join( |
| 93 | f"{m}.{n}" for m, n in sorted(FI_MOE_EP_SUPPORTED_CAPABILITIES) |
| 94 | ) |
| 95 | raise ValueError( |
| 96 | f"moe_backend={moe_backend!r} is only supported on compute " |
| 97 | f"capability {supported} (SM100/SM103), but this device is " |
| 98 | f"{cc[0]}.{cc[1]}." |
| 99 | ) |
| 100 | |
| 101 | if vllm_config.parallel_config.enable_eplb: |
| 102 | raise NotImplementedError( |
| 103 | f"EPLB is not supported with moe_backend={moe_backend!r}: the " |
| 104 | "flashinfer moe_ep experts neither apply the logical-to-physical " |
| 105 | "expert map nor report per-expert load, so rebalancing would move " |
| 106 | "weights without moving routing. Use " |
| 107 | "moe_backend=deep_gemm_mega_moe to run the mega path with EPLB." |
| 108 | ) |
| 109 | |
| 110 | |
| 111 | def make_fi_moe_ep_bootstrap() -> BootstrapConfig: |
| 112 | from flashinfer.moe_ep import BootstrapConfig |
| 113 | |
| 114 | ep = get_ep_group() |
| 115 | return BootstrapConfig( |
| 116 | world_size=ep.world_size, |
| 117 | rank=ep.rank_in_group, |
| 118 | process_group=ep.device_group, |
| 119 | auto_bootstrap=False, |
| 120 | # vLLM already bound this worker's (possibly remapped) device; |
| 121 | # without this the runtime would rebind to cuda:LOCAL_RANK|rank and |
| 122 | # launch weight transforms against another device's pointers. |
| 123 | device=torch.accelerator.current_device_index(), |
| 124 | ) |
| 125 | |
| 126 | |
| 127 | def megakernel_runtime_requirements(spec: FiMoeEpBackendSpec) -> frozenset[str]: |
| 128 | from flashinfer.moe_ep.core.runtime import NVSHMEM, TORCH_DIST |
| 129 | |
| 130 | if spec.needs_nvshmem: |
| 131 | return frozenset({TORCH_DIST, NVSHMEM}) |
| 132 | return frozenset({TORCH_DIST}) |
| 133 | |
| 134 | |
| 135 | def ensure_fi_moe_ep_runtime(vllm_config: VllmConfig) -> None: |
| 136 | """Acquire the process-wide flashinfer moe_ep runtime once per worker.""" |
| 137 | global _FI_RUNTIME_HANDLE |
| 138 | if _FI_RUNTIME_HANDLE is not None: |
| 139 | return |
| 140 | |
| 141 | from flashinfer.moe_ep import bootstrap_moe_ep_runtime |
| 142 | |
| 143 | bootstrap = make_fi_moe_ep_bootstrap() |
| 144 | spec = fi_moe_ep_backend_spec(vllm_config.kernel_config.moe_backend) |
| 145 | _FI_RUNTIME_HANDLE = bootstrap_moe_ep_runtime( |
| 146 | bootstrap, |
| 147 | megakernel_runtime_requirements(spec), |
| 148 | ) |
| 149 | |
| 150 | |
| 151 | def finalize_fi_moe_ep_runtime() -> None: |
| 152 | """Release the process-wide flashinfer moe_ep runtime.""" |
| 153 | global _FI_RUNTIME_HANDLE |
| 154 | if _FI_RUNTIME_HANDLE is None: |
| 155 | return |
| 156 | |
| 157 | from flashinfer.moe_ep import finalize_moe_ep_runtime |
| 158 | |
| 159 | finalize_moe_ep_runtime(_FI_RUNTIME_HANDLE) |
| 160 | _FI_RUNTIME_HANDLE = None |
| 161 | |
| 162 | |
| 163 | _E2M1_LUT = ( |
| 164 | 0.0, |
| 165 | 0.5, |
| 166 | 1.0, |
| 167 | 1.5, |
| 168 | 2.0, |
| 169 | 3.0, |
| 170 | 4.0, |
| 171 | 6.0, |
| 172 | -0.0, |
| 173 | -0.5, |
| 174 | -1.0, |
| 175 | -1.5, |
| 176 | -2.0, |
| 177 | -3.0, |
| 178 | -4.0, |
| 179 | -6.0, |
| 180 | ) |
| 181 | |
| 182 | |
| 183 | def _dequant_fp4_ue8m0_gran32( |
| 184 | packed: torch.Tensor, sf_ue8m0: torch.Tensor |
| 185 | ) -> torch.Tensor: |
| 186 | """[rows, K//2] packed e2m1 + [rows, K//32] ue8m0-uint8 scales -> bf16 [rows, K].""" |
| 187 | raw = packed.view(torch.uint8) |
| 188 | lut = torch.tensor(_E2M1_LUT, dtype=torch.float32, device=raw.device) |
| 189 | vals = torch.empty( |
| 190 | raw.shape[0], raw.shape[1] * 2, dtype=torch.float32, device=raw.device |
| 191 | ) |
| 192 | vals[:, ::2] = lut[(raw & 0x0F).to(torch.int64)] |
| 193 | vals[:, 1::2] = lut[(raw >> 4).to(torch.int64)] |
| 194 | sf = (sf_ue8m0.to(torch.int32) << 23).view(torch.float32) |
| 195 | return (vals * sf.repeat_interleave(32, dim=-1)).to(torch.bfloat16) |
| 196 | |
| 197 | |
| 198 | def _dequant_expert_weights_to_bf16( |
| 199 | weight: torch.Tensor, scale: torch.Tensor |
| 200 | ) -> torch.Tensor: |
| 201 | """[E, N, K//2] fp4 + [E, N, K//32] ue8m0 -> [E, N, K] bf16 (expert loop).""" |
| 202 | num_experts, n, k_half = weight.shape |
| 203 | out = torch.empty( |
| 204 | num_experts, n, k_half * 2, dtype=torch.bfloat16, device=weight.device |
| 205 | ) |
| 206 | for e in range(num_experts): |
| 207 | out[e] = _dequant_fp4_ue8m0_gran32(weight[e], scale[e]) |
| 208 | return out |
| 209 | |
| 210 | |
| 211 | def mega_moe_weight_pack_from_params( |
| 212 | w13_weight: nn.Parameter, |
| 213 | w13_weight_scale: nn.Parameter, |
| 214 | w2_weight: nn.Parameter, |
| 215 | w2_weight_scale: nn.Parameter, |
| 216 | *, |
| 217 | megakernel: str = "deep_gemm_mega", |
| 218 | ): |
| 219 | from flashinfer.moe_ep import MoEWeightPack |
| 220 | |
| 221 | if megakernel == "deep_gemm_mega": |
| 222 | # Same fp4-e2m1 + ue8m0-per-32 recipe as the native path: pass verbatim, |
| 223 | # flashinfer runs the identical deep_gemm transform. |
| 224 | return MoEWeightPack( |
| 225 | w13=w13_weight.data, |
| 226 | w2=w2_weight.data, |
| 227 | w13_scale=w13_weight_scale.data, |
| 228 | w2_scale=w2_weight_scale.data, |
| 229 | ) |
| 230 | # The cutedsl kernel quantizes with its own recipe (nvfp4 |
| 231 | # e2m1+e4m3-per-16): dequantize the checkpoint fp4 to bf16 and let the |
| 232 | # backend preprocess requantize. Double quantization: outputs are close to |
| 233 | # but not bit-identical with the native path. |
| 234 | return MoEWeightPack( |
| 235 | w13=_dequant_expert_weights_to_bf16(w13_weight.data, w13_weight_scale.data), |
| 236 | w2=_dequant_expert_weights_to_bf16(w2_weight.data, w2_weight_scale.data), |
| 237 | ) |
| 238 | |
| 239 | |
| 240 | def build_fi_mega_config( |
| 241 | *, |
| 242 | intermediate_size: int, |
| 243 | top_k: int, |
| 244 | activation_clamp: float | None, |
| 245 | megakernel: str, |
| 246 | ): |
| 247 | from flashinfer.moe_ep import ( |
| 248 | DeepGemmMegaMoeConfig, |
| 249 | MegaConfig, |
| 250 | Nvfp4CutedslMegaMoeConfig, |
| 251 | ) |
| 252 | |
| 253 | # fast_math selects approximate exp/rcp in DeepGEMM's fused SwiGLU |
| 254 | # epilogue; the cutedsl kernels accept it for API parity only. |
| 255 | if megakernel == "deep_gemm_mega": |
| 256 | mk = DeepGemmMegaMoeConfig( |
| 257 | intermediate_size=intermediate_size, |
| 258 | top_k=top_k, |
| 259 | activation_clamp=activation_clamp, |
| 260 | fast_math=True, |
| 261 | ) |
| 262 | elif megakernel == "nvfp4_cutedsl": |
| 263 | mk = Nvfp4CutedslMegaMoeConfig( |
| 264 | intermediate_size=intermediate_size, |
| 265 | top_k=top_k, |
| 266 | activation_clamp=activation_clamp, |
| 267 | fast_math=True, |
| 268 | ) |
| 269 | else: |
| 270 | raise ValueError(f"Unsupported fi_moe_ep megakernel {megakernel!r}") |
| 271 | |
| 272 | return MegaConfig( |
| 273 | megakernel=mk, |
| 274 | preprocess_weights=True, |
| 275 | quantize_input=True, |
| 276 | ) |
| 277 | |
| 278 | |
| 279 | # All MoE layers share one symmetric workspace, like the native path's |
| 280 | # class-level DeepseekV4MegaMoEExperts._symm_buffer_cache. Without this the |
| 281 | # fi path allocates one symm buffer PER LAYER (43x memory + cold working |
| 282 | # sets); the workspace is stateless across forwards (kernel tail-cleans) and |
| 283 | # layers execute sequentially on one stream, so sharing is safe. |
| 284 | def build_fi_mega_layer( |
| 285 | bootstrap: BootstrapConfig, |
| 286 | *, |
| 287 | vllm_config: VllmConfig, |
| 288 | num_experts: int, |
| 289 | max_tokens_per_rank: int, |
| 290 | hidden_size: int, |
| 291 | intermediate_size: int, |
| 292 | top_k: int, |
| 293 | activation_clamp: float | None, |
| 294 | weights, |
| 295 | ) -> MoEEpMegaLayer: |
| 296 | from flashinfer.moe_ep import FleetParams, MoEEpLayer |
| 297 | |
| 298 | megakernel = fi_moe_ep_backend_spec( |
| 299 | vllm_config.kernel_config.moe_backend |
| 300 | ).megakernel |
| 301 | mega_config = build_fi_mega_config( |
| 302 | intermediate_size=intermediate_size, |
| 303 | top_k=top_k, |
| 304 | activation_clamp=activation_clamp, |
| 305 | megakernel=megakernel, |
| 306 | ) |
| 307 | layer = MoEEpLayer( |
| 308 | bootstrap=bootstrap, |
| 309 | fleet_params=FleetParams( |
| 310 | num_experts=num_experts, |
| 311 | max_tokens_per_rank=max_tokens_per_rank, |
| 312 | token_hidden_size=hidden_size, |
| 313 | ), |
| 314 | weights=weights, |
| 315 | backend=mega_config, |
| 316 | ) |
| 317 | from flashinfer.moe_ep import MoEEpMegaLayer |
| 318 | |
| 319 | if not isinstance(layer, MoEEpMegaLayer): |
| 320 | raise TypeError( |
| 321 | f"fi_moe_ep expected MoEEpMegaLayer, got {type(layer).__name__}" |
| 322 | ) |
| 323 | return layer |
| 324 | |
| 325 | |
| 326 | __all__ = [ |
| 327 | "FI_MOE_EP_BACKEND_SPECS", |
| 328 | "FiMoeEpBackendSpec", |
| 329 | "build_fi_mega_config", |
| 330 | "build_fi_mega_layer", |
| 331 | "ensure_fi_moe_ep_runtime", |
| 332 | "fi_moe_ep_backend_spec", |
| 333 | "finalize_fi_moe_ep_runtime", |
| 334 | "is_fi_moe_ep_backend", |
| 335 | "make_fi_moe_ep_bootstrap", |
| 336 | "mega_moe_weight_pack_from_params", |
| 337 | "megakernel_runtime_requirements", |
| 338 | "validate_fi_moe_ep_config", |
| 339 | ] |