| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """FlashInfer ``moe_ep`` expert module for DeepSeek V4.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | from typing import TYPE_CHECKING, Any |
| 8 | |
| 9 | import torch |
| 10 | import torch.nn as nn |
| 11 | |
| 12 | from vllm.model_executor.utils import set_weight_attrs |
| 13 | from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MegaMoEExperts |
| 14 | from vllm.utils.flashinfer_moe_ep import ( |
| 15 | build_fi_mega_layer, |
| 16 | ensure_fi_moe_ep_runtime, |
| 17 | fi_moe_ep_backend_spec, |
| 18 | make_fi_moe_ep_bootstrap, |
| 19 | mega_moe_weight_pack_from_params, |
| 20 | ) |
| 21 | |
| 22 | if TYPE_CHECKING: |
| 23 | from flashinfer.moe_ep import MoEEpMegaLayer |
| 24 | |
| 25 | from vllm.config import VllmConfig |
| 26 | from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MLP |
| 27 | |
| 28 | _MOE_SKIP_PADDING: bool | None = None |
| 29 | |
| 30 | |
| 31 | def ckpt_uses_nvfp4_experts(vllm_config: VllmConfig) -> bool: |
| 32 | """True when the loaded checkpoint quantizes experts with modelopt NVFP4 |
| 33 | (e2m1 + fp8-e4m3 per-16 block scales + per-tensor weight_scale_2), i.e. |
| 34 | the recipe the nvfp4_cutedsl prequantized-weights path consumes verbatim. |
| 35 | |
| 36 | DeepSeek-V4 NVFP4-expert checkpoints declare ``moe_quant_algo: NVFP4`` in |
| 37 | their ``quantization_config``, which DeepseekV4FP8Config surfaces as |
| 38 | ``moe_quant_algo``. |
| 39 | """ |
| 40 | return getattr(vllm_config.quant_config, "moe_quant_algo", "") == "NVFP4" |
| 41 | |
| 42 | |
| 43 | def resolve_mega_moe_is_padding(num_tokens: int) -> torch.Tensor | None: |
| 44 | from vllm.forward_context import get_forward_context, is_forward_context_available |
| 45 | |
| 46 | global _MOE_SKIP_PADDING |
| 47 | if _MOE_SKIP_PADDING is None: |
| 48 | import vllm.envs as envs |
| 49 | |
| 50 | _MOE_SKIP_PADDING = bool(envs.VLLM_MOE_SKIP_PADDING) |
| 51 | if not _MOE_SKIP_PADDING or not is_forward_context_available(): |
| 52 | return None |
| 53 | is_padding = get_forward_context().is_padding |
| 54 | if is_padding is None: |
| 55 | return None |
| 56 | return is_padding[:num_tokens] |
| 57 | |
| 58 | |
| 59 | def apply_mega_moe_routing_preprocess( |
| 60 | topk_ids: torch.Tensor, |
| 61 | *, |
| 62 | is_padding: torch.Tensor | None = None, |
| 63 | ) -> torch.Tensor: |
| 64 | """Padding-only routing preprocess (EPLB hooks go here later).""" |
| 65 | if is_padding is not None: |
| 66 | topk_ids = torch.where(is_padding.unsqueeze(1), -1, topk_ids) |
| 67 | return topk_ids |
| 68 | |
| 69 | |
| 70 | def nvfp4_prequant_pack_and_alphas( |
| 71 | w13_weight: torch.Tensor, |
| 72 | w13_weight_scale: torch.Tensor, |
| 73 | w13_weight_scale_2: torch.Tensor, # (E_local, 2) fp32: [:,0]=gate(w1), [:,1]=up(w3) |
| 74 | w2_weight: torch.Tensor, |
| 75 | w2_weight_scale: torch.Tensor, |
| 76 | w2_weight_scale_2: torch.Tensor, # (E_local,) fp32 |
| 77 | *, |
| 78 | intermediate_size: int, |
| 79 | ): |
| 80 | """NVFP4 checkpoint params -> (MoEWeightPack, fc1_alpha, fc2_alpha). |
| 81 | |
| 82 | Activation quantization is fully dynamic, so the checkpoint's static |
| 83 | ``input_scale`` drops out and each GEMM's epilogue alpha reduces to the |
| 84 | weight's per-tensor ``weight_scale_2``. |
| 85 | |
| 86 | fc1_alpha is one scalar per expert, but gate (w1) and up (w3) carry their |
| 87 | own ``weight_scale_2``. When they differ, the ratio is folded into the up |
| 88 | half's e4m3 block scales — only if the fold round-trips exactly, since a |
| 89 | lossy rescale would silently change the model. |
| 90 | """ |
| 91 | from flashinfer.moe_ep import MoEWeightPack |
| 92 | |
| 93 | inter = intermediate_size |
| 94 | gate_s2 = w13_weight_scale_2[:, 0].float() |
| 95 | up_s2 = w13_weight_scale_2[:, 1].float() |
| 96 | if (gate_s2 <= 0).any() or (up_s2 <= 0).any() or (w2_weight_scale_2 <= 0).any(): |
| 97 | raise ValueError( |
| 98 | "nvfp4 prequant: non-positive weight_scale_2 loaded — checkpoint " |
| 99 | "scale tensors missing or loader routed them wrong." |
| 100 | ) |
| 101 | |
| 102 | w13_scale = w13_weight_scale |
| 103 | if not torch.equal(gate_s2, up_s2): |
| 104 | ratio = up_s2 / gate_s2 # (E_local,) |
| 105 | up_sf = w13_scale[:, inter:, :].float() |
| 106 | folded = up_sf * ratio[:, None, None] |
| 107 | folded_e4m3 = folded.to(torch.float8_e4m3fn) |
| 108 | if not torch.equal(folded_e4m3.float(), folded): |
| 109 | raise ValueError( |
| 110 | "NVFP4 MoE cannot merge the gate and up weight_scale_2 values " |
| 111 | "because their ratio is not exactly representable in " |
| 112 | "float8_e4m3fn. Use a checkpoint quantized with a shared " |
| 113 | "gate/up weight_scale_2." |
| 114 | ) |
| 115 | w13_scale = w13_scale.clone() |
| 116 | w13_scale[:, inter:, :] = folded_e4m3 |
| 117 | |
| 118 | fc1_alpha = gate_s2.clone().contiguous() |
| 119 | fc2_alpha = w2_weight_scale_2.float().clone().contiguous() |
| 120 | pack = MoEWeightPack( |
| 121 | w13=w13_weight, |
| 122 | w2=w2_weight, |
| 123 | w13_scale=w13_scale, |
| 124 | w2_scale=w2_weight_scale, |
| 125 | ) |
| 126 | return pack, fc1_alpha, fc2_alpha |
| 127 | |
| 128 | |
| 129 | class DeepseekV4MegaMoEExpertsFI(DeepseekV4MegaMoEExperts): |
| 130 | """Same weight layout/loader as the native mega experts, FI compute path.""" |
| 131 | |
| 132 | def __init__( |
| 133 | self, |
| 134 | vllm_config: VllmConfig, |
| 135 | *, |
| 136 | activation_clamp: float | None = None, |
| 137 | **kwargs: Any, |
| 138 | ) -> None: |
| 139 | super().__init__(vllm_config, **kwargs) |
| 140 | self._vllm_config = vllm_config |
| 141 | self._activation_clamp = activation_clamp |
| 142 | self._mega_layer: MoEEpMegaLayer | None = None |
| 143 | self._fast_ctx: tuple[Any, Any, Any, int, bool] | None = None |
| 144 | self._epilogue_alphas: tuple[torch.Tensor, torch.Tensor] | None = None |
| 145 | self._nvfp4_prequant = ckpt_uses_nvfp4_experts(vllm_config) |
| 146 | if self._nvfp4_prequant: |
| 147 | megakernel = fi_moe_ep_backend_spec( |
| 148 | vllm_config.kernel_config.moe_backend |
| 149 | ).megakernel |
| 150 | if megakernel != "nvfp4_cutedsl": |
| 151 | raise ValueError( |
| 152 | "NVFP4-quantized expert checkpoint requires " |
| 153 | "moe_backend=flashinfer_moe_ep_mega_cutedsl, got a " |
| 154 | f"backend using megakernel {megakernel!r} " |
| 155 | "(deep_gemm consumes the MXFP4 checkpoint instead)." |
| 156 | ) |
| 157 | self._realloc_nvfp4_params() |
| 158 | |
| 159 | def _realloc_nvfp4_params(self) -> None: |
| 160 | """Swap the mx-recipe scale params for the NVFP4 checkpoint's: |
| 161 | fp8-e4m3 per-16 block scales plus the per-tensor second-level |
| 162 | scales (weight_scale_2, input_scale) modelopt exports.""" |
| 163 | n_e = self.num_local_experts |
| 164 | inter = self.intermediate_size |
| 165 | hidden = self.hidden_size |
| 166 | attrs = {"weight_loader": self.weight_loader} |
| 167 | |
| 168 | def _param(shape: tuple, dtype: torch.dtype) -> nn.Parameter: |
| 169 | p = nn.Parameter(torch.zeros(*shape, dtype=dtype), requires_grad=False) |
| 170 | set_weight_attrs(p, attrs) |
| 171 | return p |
| 172 | |
| 173 | self.w13_weight_scale = _param( |
| 174 | (n_e, 2 * inter, hidden // 16), torch.float8_e4m3fn |
| 175 | ) |
| 176 | self.w13_weight_scale.quant_method = "block" |
| 177 | self.w2_weight_scale = _param((n_e, hidden, inter // 16), torch.float8_e4m3fn) |
| 178 | self.w2_weight_scale.quant_method = "block" |
| 179 | # (E, 2): column 0 = gate (w1), column 1 = up (w3). |
| 180 | self.w13_weight_scale_2 = _param((n_e, 2), torch.float32) |
| 181 | self.w2_weight_scale_2 = _param((n_e,), torch.float32) |
| 182 | # Loaded for completeness; unused — activation quant is dynamic |
| 183 | # (see nvfp4_prequant_pack_and_alphas). |
| 184 | self.w13_input_scale = _param((n_e, 2), torch.float32) |
| 185 | self.w2_input_scale = _param((n_e,), torch.float32) |
| 186 | |
| 187 | def weight_loader( |
| 188 | self, |
| 189 | param: nn.Parameter, |
| 190 | loaded_weight: torch.Tensor, |
| 191 | weight_name: str, |
| 192 | shard_id: str, |
| 193 | expert_id: int, |
| 194 | return_success: bool = False, |
| 195 | ) -> bool | None: |
| 196 | # NVFP4 checkpoint second-level scalars route here; everything |
| 197 | # else (packed weights, block scales) matches the base layout. |
| 198 | if "weight_scale_2" in weight_name or "input_scale" in weight_name: |
| 199 | local_expert_ids = self._map_global_expert_id(expert_id) |
| 200 | if not local_expert_ids: |
| 201 | return False if return_success else None |
| 202 | value = loaded_weight.reshape(()).to(torch.float32) |
| 203 | for local_expert_id in local_expert_ids: |
| 204 | if shard_id in ("w1", "w3"): |
| 205 | if "w13_" not in weight_name: |
| 206 | return False if return_success else None |
| 207 | param.data[local_expert_id, 0 if shard_id == "w1" else 1] = value |
| 208 | elif shard_id == "w2": |
| 209 | if "w2_" not in weight_name: |
| 210 | return False if return_success else None |
| 211 | param.data[local_expert_id] = value |
| 212 | else: |
| 213 | raise ValueError(f"Unsupported expert shard id: {shard_id}") |
| 214 | return True if return_success else None |
| 215 | return super().weight_loader( |
| 216 | param, |
| 217 | loaded_weight, |
| 218 | weight_name, |
| 219 | shard_id, |
| 220 | expert_id, |
| 221 | return_success, |
| 222 | ) |
| 223 | |
| 224 | def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None: |
| 225 | # The FlashInfer megakernel has no shared-expert fusion; the caller's |
| 226 | # serial shared MLP path handles shared_experts. |
| 227 | if self._mega_layer is not None: |
| 228 | return |
| 229 | if self.w13_weight is None: |
| 230 | return |
| 231 | |
| 232 | self._check_runtime_supported() |
| 233 | ensure_fi_moe_ep_runtime(self._vllm_config) |
| 234 | |
| 235 | if self._nvfp4_prequant: |
| 236 | # NVFP4 checkpoint: hand the packed weights + both scale |
| 237 | # planes straight to the backend (no dequant->requant); |
| 238 | # per-expert globals become fc1/fc2 epilogue alphas staged |
| 239 | # at every forward via MoEEpTensors. |
| 240 | weights, fc1_alpha, fc2_alpha = nvfp4_prequant_pack_and_alphas( |
| 241 | self.w13_weight.data, |
| 242 | self.w13_weight_scale.data, |
| 243 | self.w13_weight_scale_2.data, |
| 244 | self.w2_weight.data, |
| 245 | self.w2_weight_scale.data, |
| 246 | self.w2_weight_scale_2.data, |
| 247 | intermediate_size=self.intermediate_size, |
| 248 | ) |
| 249 | self._epilogue_alphas = (fc1_alpha, fc2_alpha) |
| 250 | else: |
| 251 | weights = mega_moe_weight_pack_from_params( |
| 252 | self.w13_weight, |
| 253 | self.w13_weight_scale, |
| 254 | self.w2_weight, |
| 255 | self.w2_weight_scale, |
| 256 | megakernel=fi_moe_ep_backend_spec( |
| 257 | self._vllm_config.kernel_config.moe_backend |
| 258 | ).megakernel, |
| 259 | ) |
| 260 | self._mega_layer = build_fi_mega_layer( |
| 261 | make_fi_moe_ep_bootstrap(), |
| 262 | vllm_config=self._vllm_config, |
| 263 | num_experts=self.num_experts, |
| 264 | max_tokens_per_rank=self.max_num_tokens, |
| 265 | hidden_size=self.hidden_size, |
| 266 | intermediate_size=self.intermediate_size, |
| 267 | top_k=self.top_k, |
| 268 | activation_clamp=self._activation_clamp, |
| 269 | weights=weights, |
| 270 | ) |
| 271 | del weights |
| 272 | # Allocate (or attach to) the pooled workspace before first |
| 273 | # forward so warmup/capture never hits the lazy path. |
| 274 | self._mega_layer._ensure_workspace() |
| 275 | self.w13_weight = None |
| 276 | self.w13_weight_scale = None |
| 277 | self.w2_weight = None |
| 278 | self.w2_weight_scale = None |
| 279 | if self._nvfp4_prequant: |
| 280 | self.w13_weight_scale_2 = None |
| 281 | self.w2_weight_scale_2 = None |
| 282 | self.w13_input_scale = None |
| 283 | self.w2_input_scale = None |
| 284 | |
| 285 | # EPLB is rejected for these backends in validate_fi_moe_ep_config, so |
| 286 | # nothing constructs an EplbState and none of these run. They are kept |
| 287 | # as explicit errors rather than deleted because inheriting the native |
| 288 | # implementations would be worse: set_eplb_state would record a map |
| 289 | # this forward path never consults, silently splitting routing from |
| 290 | # weights, and get_expert_weights reads transformed-weight attributes |
| 291 | # the flashinfer path releases after preprocess. |
| 292 | _EPLB_UNSUPPORTED = ( |
| 293 | "EPLB is not supported with the flashinfer moe_ep backends; " |
| 294 | "validate_fi_moe_ep_config should have rejected this " |
| 295 | "configuration at startup." |
| 296 | ) |
| 297 | |
| 298 | def set_eplb_state( |
| 299 | self, |
| 300 | moe_layer_idx: int, |
| 301 | expert_load_view: torch.Tensor, |
| 302 | logical_to_physical_map: torch.Tensor, |
| 303 | logical_replica_count: torch.Tensor, |
| 304 | ) -> None: |
| 305 | raise NotImplementedError(self._EPLB_UNSUPPORTED) |
| 306 | |
| 307 | def get_expert_weights(self) -> list[torch.Tensor]: |
| 308 | raise NotImplementedError(self._EPLB_UNSUPPORTED) |
| 309 | |
| 310 | def update_expert_map(self) -> None: |
| 311 | raise NotImplementedError(self._EPLB_UNSUPPORTED) |
| 312 | |
| 313 | def forward( |
| 314 | self, |
| 315 | hidden_states: torch.Tensor, |
| 316 | topk_weights: torch.Tensor, |
| 317 | topk_ids: torch.Tensor, |
| 318 | *, |
| 319 | activation_clamp: float | None, |
| 320 | fast_math: bool = True, |
| 321 | ) -> torch.Tensor: |
| 322 | # fast_math is a native deep_gemm knob; the FI kernels have no |
| 323 | # equivalent toggle, so it is accepted for signature parity only. |
| 324 | if hidden_states.shape[0] > self.max_num_tokens: |
| 325 | raise ValueError( |
| 326 | f"DeepSeek V4 MegaMoE got {hidden_states.shape[0]} tokens, " |
| 327 | f"but the symmetric buffer was sized for {self.max_num_tokens}." |
| 328 | ) |
| 329 | |
| 330 | from flashinfer.moe_ep import MoEEpTensors |
| 331 | |
| 332 | num_tokens = hidden_states.shape[0] |
| 333 | is_padding = resolve_mega_moe_is_padding(num_tokens) |
| 334 | topk_ids = apply_mega_moe_routing_preprocess( |
| 335 | topk_ids, |
| 336 | is_padding=is_padding, |
| 337 | ) |
| 338 | |
| 339 | # Fast path: after the first successful full forward the layer is |
| 340 | # immutable, so skip MoEEpMegaLayer.forward()'s per-call validation |
| 341 | # and go straight to the kernel backend's stage_inputs + compute. |
| 342 | alphas = self._epilogue_alphas |
| 343 | fc1_alpha = alphas[0] if alphas is not None else None |
| 344 | fc2_alpha = alphas[1] if alphas is not None else None |
| 345 | |
| 346 | fast = self._fast_ctx |
| 347 | if fast is not None: |
| 348 | kernel, workspace, transformed, hidden_size, zero_copy = fast |
| 349 | t = MoEEpTensors( |
| 350 | hidden_states=hidden_states, |
| 351 | topk_ids=topk_ids, |
| 352 | topk_weights=topk_weights, |
| 353 | fc1_alpha=fc1_alpha, |
| 354 | fc2_alpha=fc2_alpha, |
| 355 | ) |
| 356 | kernel.stage_inputs(t, workspace, quantize_input=True) |
| 357 | if zero_copy: |
| 358 | # Zero-copy (cutedsl backends): consume the workspace [:n] |
| 359 | # view directly — valid under stream ordering until the next |
| 360 | # MoE layer's launch on the shared workspace. |
| 361 | return kernel.compute(workspace, transformed, output=None) |
| 362 | # deep_gemm_mega's compute() requires a real output tensor. |
| 363 | out = torch.empty( |
| 364 | num_tokens, |
| 365 | hidden_size, |
| 366 | dtype=torch.bfloat16, |
| 367 | device=hidden_states.device, |
| 368 | ) |
| 369 | return kernel.compute(workspace, transformed, output=out) |
| 370 | |
| 371 | ensure_fi_moe_ep_runtime(self._vllm_config) |
| 372 | self.finalize_weights() |
| 373 | assert self._mega_layer is not None |
| 374 | |
| 375 | y = self._mega_layer.forward( |
| 376 | MoEEpTensors( |
| 377 | hidden_states=hidden_states, |
| 378 | topk_ids=topk_ids, |
| 379 | topk_weights=topk_weights, |
| 380 | fc1_alpha=fc1_alpha, |
| 381 | fc2_alpha=fc2_alpha, |
| 382 | ) |
| 383 | ) |
| 384 | layer = self._mega_layer |
| 385 | if hidden_states.dtype == torch.bfloat16: |
| 386 | self._fast_ctx = ( |
| 387 | layer._kernel, |
| 388 | layer._ensure_workspace(), |
| 389 | layer._transformed, |
| 390 | layer._fleet_params.token_hidden_size, |
| 391 | # zero-copy output views are a cutedsl-backend contract |
| 392 | layer._kernel.kernel_name() != "deep_gemm_mega", |
| 393 | ) |
| 394 | return y |
| 395 | |
| 396 | |
| 397 | DeepseekV4MegaMoEExpertsFI.weight_loader.supports_moe_loading = True # type: ignore[attr-defined] |