@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 4f6885fffc

[DSV4][Kernel] Fuse shared experts into MegaMoE (#53040)

Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>

Canlin Guo committed Aug 21, 2026 at 03:07 UTC 4f6885fffc931ce030d55a97ffb03bbeb0934d10
5 files changed +569 -44
tests/models/test_deepseek_v4_mega_moe.py
+232
@@ -13,6 +13,7 @@ from vllm.models.deepseek_v4.nvidia.dspark import DSparkDeepseekV4ForCausalLM
13 from vllm.models.deepseek_v4.nvidia.model import (
14 DeepseekV4ForCausalLM,
15 DeepseekV4MegaMoEExperts,
16 + DeepseekV4MoE,
17 make_deepseek_v4_expert_params_mapping,
18 )
19 from vllm.models.deepseek_v4.nvidia.mtp import DeepSeekV4MTP
@@ -168,6 +169,149 @@ def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership():
169 assert torch.count_nonzero(experts.w13_weight[1]) == 0
170
171
172 +def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(monkeypatch):
173 + class FakeDeepGemm:
174 + transformed_dims: list[tuple[int, int]] = []
175 + scale_inputs: list[tuple[int, ...]] = []
176 +
177 + @staticmethod
178 + def get_symm_buffer_for_mega_moe(*args, num_shared_experts=0, **kwargs):
179 + return None
180 +
181 + @staticmethod
182 + def get_block_m_for_mega_moe(*args, **kwargs):
183 + return 128
184 +
185 + @staticmethod
186 + def fp8_fp4_mega_moe(
187 + y,
188 + l1_weights,
189 + l2_weights,
190 + sym_buffer,
191 + shared_l1_weights=None,
192 + shared_l2_weights=None,
193 + **kwargs,
194 + ):
195 + return None
196 +
197 + @classmethod
198 + def transform_sf_into_required_layout(cls, sf, mn, k, *args, **kwargs):
199 + cls.scale_inputs.append(tuple(sf.shape))
200 + return torch.empty((sf.shape[0], mn, k // 128), dtype=torch.int32)
201 +
202 + @classmethod
203 + def transform_weights_for_mega_moe(cls, l1_weights, l2_weights):
204 + cls.transformed_dims.append((l1_weights[0].dim(), l2_weights[0].dim()))
205 + if l1_weights[0].dim() == 2:
206 + return (l1_weights[0].clone(), l1_weights[1]), l2_weights
207 + return l1_weights, l2_weights
208 +
209 + vllm_config = SimpleNamespace(
210 + scheduler_config=SimpleNamespace(max_num_batched_tokens=4),
211 + compilation_config=SimpleNamespace(static_forward_context={}),
212 + )
213 + experts = DeepseekV4MegaMoEExperts(
214 + vllm_config,
215 + num_experts=2,
216 + num_local_experts=1,
217 + experts_start_idx=0,
218 + top_k=1,
219 + hidden_size=128,
220 + intermediate_size=128,
221 + num_shared_experts=1,
222 + )
223 + experts._check_runtime_supported = lambda: None
224 +
225 + def fp8_parameter(*shape):
226 + return torch.nn.Parameter(
227 + torch.empty(*shape, dtype=torch.float8_e4m3fn), requires_grad=False
228 + )
229 +
230 + def scale_parameter(*shape, dtype=torch.int32):
231 + return torch.nn.Parameter(torch.ones(*shape, dtype=dtype), requires_grad=False)
232 +
233 + shared_experts = SimpleNamespace(
234 + gate_up_proj=SimpleNamespace(
235 + weight=fp8_parameter(256, 128),
236 + weight_block_size=(128, 128),
237 + weight_scale_inv=scale_parameter(2, 1, dtype=torch.float8_e8m0fnu),
238 + ),
239 + down_proj=SimpleNamespace(
240 + weight=fp8_parameter(128, 128),
241 + weight_block_size=(128, 128),
242 + weight_scale_inv=scale_parameter(1, 1, dtype=torch.float8_e8m0fnu),
243 + ),
244 + )
245 + monkeypatch.setattr("vllm.utils.deep_gemm._import_deep_gemm", lambda: FakeDeepGemm)
246 +
247 + original_gate_up_ptr = shared_experts.gate_up_proj.weight.data_ptr()
248 + experts.finalize_weights(shared_experts)
249 +
250 + assert FakeDeepGemm.transformed_dims == [(3, 3), (2, 2)]
251 + assert FakeDeepGemm.scale_inputs[-2:] == [(1, 256, 4), (1, 128, 4)]
252 + assert experts.has_fused_shared_experts
253 + assert shared_experts.gate_up_proj.weight.data_ptr() != original_gate_up_ptr
254 + assert (
255 + experts._transformed_shared_l1_weights[0].data_ptr()
256 + == shared_experts.gate_up_proj.weight.data_ptr()
257 + )
258 + assert (
259 + experts._transformed_shared_l2_weights[0].data_ptr()
260 + == shared_experts.down_proj.weight.data_ptr()
261 + )
262 +
263 +
264 +@pytest.mark.parametrize("fused", [False, True])
265 +def test_deepseek_v4_mega_moe_does_not_double_add_fused_shared_expert(
266 + monkeypatch, fused
267 +):
268 + class FakeGate(torch.nn.Module):
269 + tid2eid = None
270 + e_score_correction_bias = None
271 +
272 + def forward(self, hidden_states):
273 + return torch.empty(hidden_states.shape[0], 2), None
274 +
275 + class FakeExperts(torch.nn.Module):
276 + has_fused_shared_experts = fused
277 +
278 + def forward(self, hidden_states, *args, **kwargs):
279 + return torch.ones_like(hidden_states)
280 +
281 + class FakeSharedExperts(torch.nn.Module):
282 + calls = 0
283 +
284 + def forward(self, hidden_states):
285 + self.calls += 1
286 + return torch.full_like(hidden_states, 2)
287 +
288 + moe = DeepseekV4MoE.__new__(DeepseekV4MoE)
289 + torch.nn.Module.__init__(moe)
290 + moe.use_mega_moe = True
291 + moe.gate = FakeGate()
292 + moe.experts = FakeExperts()
293 + moe.shared_experts = FakeSharedExperts()
294 + moe.scoring_func = "sqrtsoftplus"
295 + moe.n_activated_experts = 1
296 + moe.renormalize = True
297 + moe.hash_indices_dtype = torch.int64
298 + moe.routed_scaling_factor = 1.0
299 + moe.swiglu_limit = 10.0
300 + monkeypatch.setattr(
301 + "vllm.models.deepseek_v4.nvidia.model.fused_topk_bias",
302 + lambda **kwargs: (
303 + torch.ones(kwargs["hidden_states"].shape[0], 1),
304 + torch.zeros(kwargs["hidden_states"].shape[0], 1, dtype=torch.int64),
305 + ),
306 + )
307 +
308 + output = moe(torch.zeros(2, 128))
309 +
310 + expected = 1 if fused else 3
311 + assert torch.all(output == expected)
312 + assert moe.shared_experts.calls == (0 if fused else 1)
313 +
314 +
315 @pytest.mark.skipif(
316 not torch.cuda.is_available(),
317 reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.",
@@ -246,6 +390,94 @@ def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact():
390 )
391
392
393 +@pytest.mark.parametrize("shared_block_m", [8, 32, 96, 128, 192])
394 +def test_deepseek_v4_mega_moe_stages_shared_scale_tma_layout(shared_block_m):
395 + from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8
396 +
397 + device = torch.device("cuda")
398 + num_tokens = shared_block_m + 7
399 + hidden_size = 256
400 + top_k = 8
401 + generator = torch.Generator(device=device)
402 + generator.manual_seed(shared_block_m)
403 + hidden_states = torch.randn(
404 + num_tokens,
405 + hidden_size,
406 + device=device,
407 + dtype=torch.bfloat16,
408 + generator=generator,
409 + )
410 + topk_ids = torch.randint(
411 + 0,
412 + 256,
413 + (num_tokens, top_k),
414 + device=device,
415 + dtype=torch.int32,
416 + generator=generator,
417 + )
418 + topk_weights = torch.randn(
419 + num_tokens,
420 + top_k,
421 + device=device,
422 + dtype=torch.float32,
423 + generator=generator,
424 + )
425 +
426 + ref_x, ref_x_sf = per_token_cast_to_fp8(
427 + hidden_states,
428 + use_ue8m0=True,
429 + gran_k=32,
430 + use_packed_ue8m0=True,
431 + )
432 + aligned_block_m = ((shared_block_m + 127) // 128) * 128
433 + num_shared_rows = ((num_tokens + shared_block_m - 1) // shared_block_m) * (
434 + aligned_block_m
435 + )
436 + ref_shared_x_sf = torch.zeros(
437 + num_shared_rows,
438 + hidden_size // 128,
439 + dtype=torch.int32,
440 + device=device,
441 + )
442 + for token_id in range(num_tokens):
443 + m_in_block = token_id % shared_block_m
444 + transposed_m = (
445 + (m_in_block // 128) * 128 + (m_in_block % 32) * 4 + (m_in_block % 128) // 32
446 + )
447 + shared_row = token_id // shared_block_m * aligned_block_m + transposed_m
448 + ref_shared_x_sf[shared_row].copy_(ref_x_sf[token_id])
449 +
450 + fused_x = torch.empty_like(ref_x)
451 + fused_x_sf = torch.empty_like(ref_x_sf)
452 + fused_shared_storage = torch.full(
453 + (hidden_size // 128, num_shared_rows),
454 + -1,
455 + dtype=torch.int32,
456 + device=device,
457 + )
458 + fused_shared_x_sf = fused_shared_storage.t()
459 + fused_topk_idx = torch.empty_like(topk_ids, dtype=torch.int64)
460 + fused_topk_weights = torch.empty_like(topk_weights)
461 +
462 + prepare_megamoe_inputs(
463 + hidden_states,
464 + topk_weights,
465 + topk_ids,
466 + fused_x,
467 + fused_x_sf,
468 + fused_topk_idx,
469 + fused_topk_weights,
470 + shared_x_sf=fused_shared_x_sf,
471 + shared_block_m=shared_block_m,
472 + )
473 + torch.accelerator.synchronize()
474 +
475 + populated = ref_shared_x_sf != 0
476 + assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8))
477 + assert torch.equal(fused_x_sf, ref_x_sf)
478 + assert torch.equal(fused_shared_x_sf[populated], ref_shared_x_sf[populated])
479 +
480 +
481 def test_deepseek_v4_pwal_hook_finalizes_mega_moe_and_mhc_broadcast():
482 """The loader invokes the model-level PWAL hook for every load format,
483 so it must finalize megamoe + mhc broadcast weights to cover dummy
vllm/envs.py
+7
@@ -288,6 +288,7 @@ if TYPE_CHECKING:
288 VLLM_GC_DEBUG: str = ""
289 VLLM_DEBUG_WORKSPACE: bool = False
290 VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False
291 + VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION: bool = False
292 VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256
293 VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 1024
294 VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary"
@@ -1993,6 +1994,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
1994 "VLLM_DISABLE_SHARED_EXPERTS_STREAM": lambda: bool(
1995 int(os.getenv("VLLM_DISABLE_SHARED_EXPERTS_STREAM", "0"))
1996 ),
1997 + # Emergency rollback for the DeepSeek-V4 NVIDIA MegaMoE path. By default,
1998 + # DeepGEMM computes replicated FP8 shared experts in the same persistent
1999 + # SM100 kernel as the routed FP4 experts.
2000 + "VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION": lambda: bool(
2001 + int(os.getenv("VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION", "0"))
2002 + ),
2003 # Limits when we run shared_experts in a separate stream.
2004 # We found out that for large batch sizes, the separate stream
2005 # execution is not beneficial (most likely because of the input clone)
vllm/models/deepseek_v4/nvidia/model.py
+274 -42
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import typing
4 from collections.abc import Callable, Iterable
5 +from inspect import signature
6 from itertools import islice
7
8 import regex as re
@@ -18,6 +19,7 @@ from vllm.distributed import (
19 )
20 from vllm.distributed.eplb.eplb_state import EplbLayerState
21 from vllm.forward_context import get_forward_context, is_forward_context_available
22 +from vllm.logger import init_logger
23 from vllm.model_executor.kernels.mhc.tilelang import (
24 hc_head_fused_kernel_tilelang,
25 mhc_fused_post_pre_tilelang,
@@ -84,6 +86,8 @@ from vllm.utils.math_utils import cdiv
86 from vllm.v1.attention.backends.registry import AttentionBackendEnum
87 from vllm.v1.worker.ubatching import dbo_current_ubatch_id
88
89 +logger = init_logger(__name__)
90 +
91
92 class DeepseekV4MLP(nn.Module):
93 def __init__(
@@ -165,7 +169,7 @@ def make_deepseek_v4_expert_params_mapping(
169
170
171 class DeepseekV4MegaMoEExperts(nn.Module):
168 - _symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int], object] = {}
172 + _symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int, int], object] = {}
173
174 def __init__(
175 self,
@@ -177,6 +181,7 @@ class DeepseekV4MegaMoEExperts(nn.Module):
181 top_k: int,
182 hidden_size: int,
183 intermediate_size: int,
184 + num_shared_experts: int = 0,
185 prefix: str = "",
186 num_logical_experts: int | None = None,
187 ):
@@ -190,6 +195,7 @@ class DeepseekV4MegaMoEExperts(nn.Module):
195 self.top_k = top_k
196 self.hidden_size = hidden_size
197 self.intermediate_size = intermediate_size
198 + self.num_shared_experts = num_shared_experts
199 self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
200
201 self.num_logical_experts = (
@@ -247,6 +253,12 @@ class DeepseekV4MegaMoEExperts(nn.Module):
253
254 self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None
255 self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None
256 + self._transformed_shared_l1_weights: (
257 + tuple[torch.Tensor, torch.Tensor] | None
258 + ) = None
259 + self._transformed_shared_l2_weights: (
260 + tuple[torch.Tensor, torch.Tensor] | None
261 + ) = None
262
263 # Register in the static forward context so the custom-op wrapper
264 # can look up this module by name from within a torch.compile graph.
@@ -321,45 +333,220 @@ class DeepseekV4MegaMoEExperts(nn.Module):
333 "to be multiples of 128."
334 )
335
324 - def finalize_weights(self) -> None:
325 - if self._transformed_l1_weights is not None:
326 - return
336 + @staticmethod
337 + def _deep_gemm_supports_shared_experts(deep_gemm) -> bool:
338 + """Check the Python API before touching a symmetric-memory group.
339
328 - self._check_runtime_supported()
329 - from vllm.utils.deep_gemm import _import_deep_gemm
340 + This also gives users of an older precompiled vLLM wheel a safe serial
341 + fallback instead of failing halfway through multi-rank buffer setup.
342 + """
343 + try:
344 + buffer_params = signature(deep_gemm.get_symm_buffer_for_mega_moe).parameters
345 + kernel_params = signature(deep_gemm.fp8_fp4_mega_moe).parameters
346 + except (TypeError, ValueError):
347 + return False
348 + return (
349 + hasattr(deep_gemm, "get_block_m_for_mega_moe")
350 + and hasattr(deep_gemm, "transform_weights_for_mega_moe")
351 + and "num_shared_experts" in buffer_params
352 + and "shared_l1_weights" in kernel_params
353 + and "shared_l2_weights" in kernel_params
354 + )
355
331 - deep_gemm = _import_deep_gemm()
356 + def _finalize_shared_expert_weights(
357 + self, deep_gemm, shared_experts: DeepseekV4MLP
358 + ) -> None:
359 + gate_up = shared_experts.gate_up_proj
360 + down = shared_experts.down_proj
361 + gate_up_weight = gate_up.weight.data
362 + gate_up_scale = gate_up.weight_scale_inv.data
363 + down_weight = down.weight.data
364 + down_scale = down.weight_scale_inv.data
365 +
366 + # MegaMoE's shared FP8 MMA consumes a 1x32 scale for every weight row,
367 + # while the checkpoint uses coarser block-FP8 scales (usually
368 + # 128x128). Build a dedicated, numerically equivalent scale view before
369 + # the generic linear post-load hook replaces the raw checkpoint scales
370 + # with its 128x128 DeepGEMM layout.
371 + checkpoint_scale_dtypes = (torch.float8_e8m0fnu, torch.uint8)
372 + if (
373 + gate_up_scale.dtype in checkpoint_scale_dtypes
374 + and down_scale.dtype in checkpoint_scale_dtypes
375 + ):
376 + gate_up_scale = self._prepare_shared_expert_scale(
377 + deep_gemm,
378 + gate_up,
379 + gate_up_scale,
380 + gate_up_weight.shape[0],
381 + gate_up_weight.shape[1],
382 + )
383 + down_scale = self._prepare_shared_expert_scale(
384 + deep_gemm,
385 + down,
386 + down_scale,
387 + down_weight.shape[0],
388 + down_weight.shape[1],
389 + )
390 +
391 + if gate_up_scale is None or down_scale is None:
392 + self.num_shared_experts = 0
393 + return
394
333 - w13_scale = deep_gemm.transform_sf_into_required_layout(
334 - self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(),
335 - 2 * self.intermediate_size,
395 + shared_intermediate_size = self.intermediate_size * self.num_shared_experts
396 + expected_gate_up_shape = (
397 + 2 * shared_intermediate_size,
398 self.hidden_size,
337 - (1, 32),
338 - self.num_local_experts,
399 )
340 - w2_scale = deep_gemm.transform_sf_into_required_layout(
341 - self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(),
342 - self.hidden_size,
343 - self.intermediate_size,
344 - (1, 32),
345 - self.num_local_experts,
400 + expected_down_shape = (self.hidden_size, shared_intermediate_size)
401 + if (
402 + gate_up_weight.dtype != torch.float8_e4m3fn
403 + or down_weight.dtype != torch.float8_e4m3fn
404 + or gate_up_scale.dtype != torch.int32
405 + or down_scale.dtype != torch.int32
406 + or tuple(gate_up_weight.shape) != expected_gate_up_shape
407 + or tuple(down_weight.shape) != expected_down_shape
408 + ):
409 + logger.warning(
410 + "Disabling native MegaMoE shared-expert fusion for %s: expected "
411 + "replicated block-FP8 weights with gate_up=%s, down=%s, and "
412 + "DeepGEMM int32 scales; got gate_up=%s/%s/%s and down=%s/%s/%s.",
413 + self.prefix,
414 + expected_gate_up_shape,
415 + expected_down_shape,
416 + tuple(gate_up_weight.shape),
417 + gate_up_weight.dtype,
418 + gate_up_scale.dtype,
419 + tuple(down_weight.shape),
420 + down_weight.dtype,
421 + down_scale.dtype,
422 + )
423 + self.num_shared_experts = 0
424 + return
425 +
426 + transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe(
427 + (gate_up_weight, gate_up_scale),
428 + (down_weight, down_scale),
429 + )
430 + # L1 interleaving allocates a full copy. Re-home the loader Parameter on
431 + # that storage so the original 2*intermediate*hidden FP8 tensor can be
432 + # released instead of adding roughly 0.7 GiB per rank on DSV4-Flash.
433 + # The generic linear post-load hook may still repack the serial scales,
434 + # but this shared MLP is never called after native fusion is enabled.
435 + gate_up.weight.data = transformed_l1[0]
436 + self._transformed_shared_l1_weights = (
437 + gate_up.weight.data,
438 + transformed_l1[1],
439 )
347 - self._transformed_l1_weights, self._transformed_l2_weights = (
348 - deep_gemm.transform_weights_for_mega_moe(
349 - (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale),
350 - (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale),
440 + self._transformed_shared_l2_weights = transformed_l2
441 +
442 + def _prepare_shared_expert_scale(
443 + self,
444 + deep_gemm,
445 + linear: nn.Module,
446 + scale: torch.Tensor,
447 + mn: int,
448 + k: int,
449 + ) -> torch.Tensor | None:
450 + block_size = getattr(linear, "weight_block_size", None)
451 + if block_size is None or len(block_size) != 2:
452 + logger.warning(
453 + "Disabling native MegaMoE shared-expert fusion for %s: "
454 + "shared FP8 weight block size is unavailable.",
455 + self.prefix,
456 + )
457 + return None
458 +
459 + block_m, block_k = block_size
460 + expected_shape = (
461 + (mn + block_m - 1) // block_m,
462 + (k + block_k - 1) // block_k,
463 + )
464 + if block_k % 32 != 0 or tuple(scale.shape) != expected_shape:
465 + logger.warning(
466 + "Disabling native MegaMoE shared-expert fusion for %s: "
467 + "cannot convert shared scale shape %s with block size %s "
468 + "to MegaMoE's 1x32 layout for weight (%d, %d).",
469 + self.prefix,
470 + tuple(scale.shape),
471 + tuple(block_size),
472 + mn,
473 + k,
474 )
475 + return None
476 +
477 + scale_fp32 = self._ue8m0_uint8_to_float(scale.view(torch.uint8))
478 + scale_1x32 = (
479 + scale_fp32.repeat_interleave(block_m, dim=0)
480 + .repeat_interleave(block_k // 32, dim=1)[:mn, : k // 32]
481 + .contiguous()
482 )
353 - # Drop the original loader-side parameters: the MegaMoE kernels only
354 - # consume the transformed views above. transform_weights_for_mega_moe
355 - # allocates a fresh tensor for the L1 weight (see _interleave_l1_weights)
356 - # and fresh SF tensors for L1/L2; the L2 weight is the only tensor that
357 - # aliases the original storage, and _transformed_l2_weights still holds
358 - # it, so the storage stays live after we drop the Parameter.
359 - self.w13_weight = None
360 - self.w13_weight_scale = None
361 - self.w2_weight = None
362 - self.w2_weight_scale = None
483 + # The grouped API is used with a singleton dimension to request the
484 + # MN-major, TMA-aligned packed-UE8M0 strides, then squeezed back to the
485 + # 2D layout required for a shared expert.
486 + return deep_gemm.transform_sf_into_required_layout(
487 + scale_1x32.unsqueeze(0),
488 + mn,
489 + k,
490 + (1, 32),
491 + 1,
492 + ).squeeze(0)
493 +
494 + def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None:
495 + from vllm.utils.deep_gemm import _import_deep_gemm
496 +
497 + deep_gemm = _import_deep_gemm()
498 +
499 + if self._transformed_l1_weights is None:
500 + self._check_runtime_supported()
501 + w13_scale = deep_gemm.transform_sf_into_required_layout(
502 + self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(),
503 + 2 * self.intermediate_size,
504 + self.hidden_size,
505 + (1, 32),
506 + self.num_local_experts,
507 + )
508 + w2_scale = deep_gemm.transform_sf_into_required_layout(
509 + self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(),
510 + self.hidden_size,
511 + self.intermediate_size,
512 + (1, 32),
513 + self.num_local_experts,
514 + )
515 + self._transformed_l1_weights, self._transformed_l2_weights = (
516 + deep_gemm.transform_weights_for_mega_moe(
517 + (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale),
518 + (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale),
519 + )
520 + )
521 + # Drop the original loader-side parameters: the MegaMoE kernels only
522 + # consume the transformed views above. transform_weights_for_mega_moe
523 + # allocates a fresh tensor for the L1 weight (see
524 + # _interleave_l1_weights) and fresh SF tensors for L1/L2; the L2
525 + # weight is the only tensor that aliases the original storage, and
526 + # _transformed_l2_weights still holds it, so the storage stays live
527 + # after we drop the Parameter.
528 + self.w13_weight = None
529 + self.w13_weight_scale = None
530 + self.w2_weight = None
531 + self.w2_weight_scale = None
532 +
533 + if shared_experts is None or self.num_shared_experts == 0:
534 + return
535 + if self._transformed_shared_l1_weights is not None:
536 + return
537 + if not self._deep_gemm_supports_shared_experts(deep_gemm):
538 + logger.warning_once(
539 + "Disabling native MegaMoE shared-expert fusion because the "
540 + "installed DeepGEMM Python API is older than the vLLM "
541 + "source. Rebuild the vendored _deep_gemm_C extension to enable it.",
542 + )
543 + self.num_shared_experts = 0
544 + return
545 + self._finalize_shared_expert_weights(deep_gemm, shared_experts)
546 +
547 + @property
548 + def has_fused_shared_experts(self) -> bool:
549 + return self._transformed_shared_l1_weights is not None
550
551 def get_symm_buffer(self):
552 from vllm.utils.deep_gemm import _import_deep_gemm
@@ -376,6 +563,7 @@ class DeepseekV4MegaMoEExperts(nn.Module):
563 self.top_k,
564 self.hidden_size,
565 self.intermediate_size,
566 + self.num_shared_experts if self.has_fused_shared_experts else 0,
567 )
568 symm_buffer = self._symm_buffer_cache.get(key)
569 if symm_buffer is None:
@@ -386,6 +574,9 @@ class DeepseekV4MegaMoEExperts(nn.Module):
574 self.top_k,
575 self.hidden_size,
576 self.intermediate_size,
577 + num_shared_experts=(
578 + self.num_shared_experts if self.has_fused_shared_experts else 0
579 + ),
580 )
581 self._symm_buffer_cache[key] = symm_buffer
582 return symm_buffer
@@ -492,6 +683,19 @@ class DeepseekV4MegaMoEExperts(nn.Module):
683 else None,
684 )
685
686 + shared_x_sf = None
687 + shared_block_m = None
688 + if self.has_fused_shared_experts:
689 + shared_x_sf = symm_buffer.shared_l1_acts_sf
690 + shared_block_m = deep_gemm.get_block_m_for_mega_moe(
691 + get_ep_group().world_size,
692 + self.num_experts,
693 + symm_buffer.num_max_tokens_per_rank,
694 + num_tokens,
695 + self.top_k,
696 + "fp8xfp4",
697 + )
698 +
699 prepare_megamoe_inputs(
700 hidden_states,
701 topk_weights,
@@ -501,18 +705,32 @@ class DeepseekV4MegaMoEExperts(nn.Module):
705 symm_buffer.topk_idx[:num_tokens],
706 symm_buffer.topk_weights[:num_tokens],
707 is_padding=is_padding,
708 + shared_x_sf=shared_x_sf,
709 + shared_block_m=shared_block_m,
710 )
711
712 assert self._transformed_l1_weights is not None
713 assert self._transformed_l2_weights is not None
508 - deep_gemm.fp8_fp4_mega_moe(
509 - y,
510 - self._transformed_l1_weights,
511 - self._transformed_l2_weights,
512 - symm_buffer,
513 - activation_clamp=activation_clamp,
514 - fast_math=fast_math,
515 - )
714 + if self.has_fused_shared_experts:
715 + deep_gemm.fp8_fp4_mega_moe(
716 + y,
717 + self._transformed_l1_weights,
718 + self._transformed_l2_weights,
719 + symm_buffer,
720 + shared_l1_weights=self._transformed_shared_l1_weights,
721 + shared_l2_weights=self._transformed_shared_l2_weights,
722 + activation_clamp=activation_clamp,
723 + fast_math=fast_math,
724 + )
725 + else:
726 + deep_gemm.fp8_fp4_mega_moe(
727 + y,
728 + self._transformed_l1_weights,
729 + self._transformed_l2_weights,
730 + symm_buffer,
731 + activation_clamp=activation_clamp,
732 + fast_math=fast_math,
733 + )
734 return y
735
736
@@ -645,6 +863,16 @@ class DeepseekV4MoE(nn.Module):
863 self.experts_start_idx = self.physical_expert_start
864 self.experts_end_idx = self.physical_expert_end
865
866 + # Native DeepGEMM fusion requires each EP rank to own the complete
867 + # shared MLP. Sequence parallel replicates those weights while sharding
868 + # tokens. TP=1 is also naturally replicated. With PP+TP the shared MLP
869 + # remains tensor-sharded, so retain the serial path.
870 + fuse_shared_experts = bool(
871 + self.shared_experts is not None
872 + and not envs.VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION
873 + and (self.use_sequence_parallel or self.tp_size == 1)
874 + )
875 +
876 self.experts = DeepseekV4MegaMoEExperts(
877 vllm_config,
878 num_experts=self.n_physical_experts,
@@ -654,6 +882,7 @@ class DeepseekV4MoE(nn.Module):
882 top_k=config.num_experts_per_tok,
883 hidden_size=config.hidden_size,
884 intermediate_size=config.moe_intermediate_size,
885 + num_shared_experts=(self.n_shared_experts if fuse_shared_experts else 0),
886 prefix=f"{prefix}.experts",
887 )
888
@@ -739,7 +968,10 @@ class DeepseekV4MoE(nn.Module):
968 activation_clamp=activation_clamp,
969 )
970
742 - if self.shared_experts is not None:
971 + if (
972 + self.shared_experts is not None
973 + and not self.experts.has_fused_shared_experts
974 + ):
975 shared_output = self.shared_experts(hidden_states)
976 final_hidden_states += shared_output
977
@@ -759,7 +991,7 @@ class DeepseekV4MoE(nn.Module):
991
992 def finalize_mega_moe_weights(self) -> None:
993 if self.use_mega_moe:
762 - self.experts.finalize_weights()
994 + self.experts.finalize_weights(self.shared_experts)
995
996
997 def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]:
vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py
+51
@@ -17,6 +17,7 @@ def _prepare_megamoe_inputs_kernel(
17 hidden_states,
18 x_fp8,
19 x_sf,
20 + shared_x_sf,
21 topk_ids,
22 topk_weights,
23 is_padding,
@@ -28,6 +29,8 @@ def _prepare_megamoe_inputs_kernel(
29 x_stride_k: tl.constexpr,
30 x_sf_stride_m: tl.constexpr,
31 x_sf_stride_k: tl.constexpr,
32 + shared_x_sf_stride_m: tl.constexpr,
33 + shared_x_sf_stride_k: tl.constexpr,
34 topk_ids_stride_m: tl.constexpr,
35 topk_ids_stride_k: tl.constexpr,
36 topk_weights_stride_m: tl.constexpr,
@@ -42,6 +45,7 @@ def _prepare_megamoe_inputs_kernel(
45 BLOCK_K: tl.constexpr,
46 GROUP_K: tl.constexpr,
47 BLOCK_TOPK: tl.constexpr,
48 + SHARED_BLOCK_M: tl.constexpr,
49 ) -> None:
50 token_id = tl.program_id(0)
51 k_block_id = tl.program_id(1)
@@ -84,6 +88,25 @@ def _prepare_megamoe_inputs_kernel(
88 packed_scale,
89 )
90
91 + # DeepGEMM's SM100 shared-expert TMA loads require the activation scales
92 + # in an MN-major layout whose row permutation depends on the MegaMoE
93 + # scheduler's runtime BLOCK_M. Write that view while the packed UE8M0 scale
94 + # is already resident, avoiding another kernel and temporary tensor.
95 + if shared_x_sf is not None:
96 + m_block_id = token_id // SHARED_BLOCK_M
97 + m_in_block = token_id % SHARED_BLOCK_M
98 + aligned_block_m: tl.constexpr = triton.cdiv(SHARED_BLOCK_M, 128) * 128
99 + transposed_m = (
100 + (m_in_block // 128) * 128 + (m_in_block % 32) * 4 + (m_in_block % 128) // 32
101 + )
102 + shared_row = m_block_id * aligned_block_m + transposed_m
103 + tl.store(
104 + shared_x_sf
105 + + shared_row * shared_x_sf_stride_m
106 + + k_block_id * shared_x_sf_stride_k,
107 + packed_scale,
108 + )
109 +
110 if k_block_id == 0:
111 topk_offsets = tl.arange(0, BLOCK_TOPK)
112 topk_mask = topk_offsets < top_k
@@ -131,6 +154,8 @@ def prepare_megamoe_inputs(
154 topk_idx_out: torch.Tensor,
155 topk_weights_out: torch.Tensor,
156 is_padding: torch.Tensor | None = None,
157 + shared_x_sf: torch.Tensor | None = None,
158 + shared_block_m: int | None = None,
159 ) -> None:
160 num_tokens, hidden_size = hidden_states.shape
161 if num_tokens == 0:
@@ -146,6 +171,28 @@ def prepare_megamoe_inputs(
171 "DeepSeek V4 MegaMoE input staging requires topk_weights and "
172 "topk_ids to have the same shape."
173 )
174 + if (shared_x_sf is None) != (shared_block_m is None):
175 + raise ValueError(
176 + "DeepSeek V4 MegaMoE shared input staging requires both "
177 + "shared_x_sf and shared_block_m."
178 + )
179 + if shared_x_sf is not None:
180 + assert shared_block_m is not None
181 + if shared_block_m <= 0:
182 + raise ValueError("MegaMoE shared_block_m must be positive.")
183 + expected_sf_k = hidden_size // 128
184 + if shared_x_sf.ndim != 2 or shared_x_sf.shape[1] != expected_sf_k:
185 + raise ValueError(
186 + "MegaMoE shared_x_sf must have shape "
187 + f"(*, {expected_sf_k}), got {tuple(shared_x_sf.shape)}."
188 + )
189 + aligned_block_m = triton.cdiv(shared_block_m, 128) * 128
190 + required_rows = triton.cdiv(num_tokens, shared_block_m) * aligned_block_m
191 + if shared_x_sf.shape[0] < required_rows:
192 + raise ValueError(
193 + "MegaMoE shared_x_sf has insufficient rows: requires "
194 + f"{required_rows}, got {shared_x_sf.shape[0]}."
195 + )
196
197 block_k = 128
198 grid = (num_tokens, triton.cdiv(hidden_size, block_k))
@@ -155,6 +202,7 @@ def prepare_megamoe_inputs(
202 hidden_states,
203 x_fp8,
204 x_sf,
205 + shared_x_sf,
206 topk_ids,
207 topk_weights,
208 is_padding,
@@ -166,6 +214,8 @@ def prepare_megamoe_inputs(
214 x_fp8.stride(1),
215 x_sf.stride(0),
216 x_sf.stride(1),
217 + shared_x_sf.stride(0) if shared_x_sf is not None else 0,
218 + shared_x_sf.stride(1) if shared_x_sf is not None else 0,
219 topk_ids.stride(0),
220 topk_ids.stride(1),
221 topk_weights.stride(0),
@@ -180,5 +230,6 @@ def prepare_megamoe_inputs(
230 BLOCK_K=block_k,
231 GROUP_K=32,
232 BLOCK_TOPK=block_topk,
233 + SHARED_BLOCK_M=shared_block_m or 1,
234 num_warps=4,
235 )
vllm/models/kimi_k3/nvidia/model.py
+5 -2
@@ -93,7 +93,10 @@ from vllm.models.common.ops.sequence_parallel import (
93 sp_reduce_scatter,
94 sp_shard,
95 )
96 -from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MegaMoEExperts
96 +from vllm.models.deepseek_v4.nvidia.model import (
97 + DeepseekV4MegaMoEExperts,
98 + DeepseekV4MLP,
99 +)
100 from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
101 from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention
102 from vllm.models.kimi_k3.nvidia.latent_moe_runner import (
@@ -354,7 +357,7 @@ class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts):
357 torch.distributed.barrier(group=ep_group.cpu_group)
358 self._synchronized_ep_groups.add(key)
359
357 - def finalize_weights(self) -> None:
360 + def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None:
361 if self._transformed_l1_weights is not None:
362 return
363