[Perf][DSV4.1] Pad shared experts for native MegaMoE fusion (#56568)
Signed-off-by: Canlin <canlinguosdu@gmail.com> Co-authored-by: Codex <noreply@openai.com>
Canlin Guo committed
Sep 16, 2026 at 17:08 UTC
03f67b3ad1e61d12c21d42e20e69cb5f99bdd82e
2 files changed
+128
-23
tests/models/test_deepseek_v4_mega_moe.py
+80
-18
@@ -367,17 +367,21 @@ def test_deepseek_v4_mega_moe_padding_preserves_weights(monkeypatch, intermediat
367
[
368
pytest.param(128, 512, 128, True, False, id="aligned-block128"),
369
pytest.param(128, 512, 128, True, True, id="aligned-block128-mxfp8"),
370
- pytest.param(5120, 2304, 32, False, False, id="deepseek-v41-flash"),
370
+ pytest.param(5120, 2304, 32, True, False, id="deepseek-v41-flash"),
371
+ pytest.param(128, 2304, 128, True, False, id="padded-block128"),
372
+ pytest.param(5120, 2304, 32, True, True, id="deepseek-v41-flash-mxfp8"),
373
+ pytest.param(128, 2304, 128, False, False, id="padded-packed-scale-fallback"),
374
],
375
)
376
def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(
377
monkeypatch, hidden_size, intermediate_size, block_size, fused, mxfp8
378
):
376
- """V4.1-Flash's routed padding must preserve the serial shared-expert fallback."""
379
+ """Shared fusion preserves checkpoint channels and zeros padded channels."""
380
381
class FakeDeepGemm:
382
transformed_dims: list[tuple[int, int]] = []
383
scale_inputs: list[tuple[int, ...]] = []
384
+ scale_values: list[torch.Tensor] = []
385
386
@staticmethod
387
def get_symm_buffer_for_mega_moe(*args, num_shared_experts=0, **kwargs):
@@ -402,6 +406,7 @@ def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(
406
@classmethod
407
def transform_sf_into_required_layout(cls, sf, mn, k, *args, **kwargs):
408
cls.scale_inputs.append(tuple(sf.shape))
409
+ cls.scale_values.append(sf)
410
return torch.empty((sf.shape[0], mn, k // 128), dtype=torch.int32)
411
412
@classmethod
@@ -429,7 +434,10 @@ def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(
434
435
def fp8_parameter(*shape):
436
return torch.nn.Parameter(
432
- torch.empty(*shape, dtype=torch.float8_e4m3fn), requires_grad=False
437
+ torch.empty(*shape, dtype=torch.uint8)
438
+ .random_(1, 119)
439
+ .view(torch.float8_e4m3fn),
440
+ requires_grad=False,
441
)
442
443
def scale_parameter(*shape, dtype=torch.int32):
@@ -460,21 +468,34 @@ def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(
468
linear.weight_block_size = (1, 32)
469
linear.weight_scale = torch.nn.Parameter(
470
linear.weight_scale_inv.view(torch.uint8)
463
- .repeat_interleave(128, dim=0)
464
- .repeat_interleave(4, dim=1),
471
+ .repeat_interleave(block_size, dim=0)
472
+ .repeat_interleave(block_size // 32, dim=1),
473
requires_grad=False,
474
)
475
del linear.weight_scale_inv
476
+ originals = []
477
+ for linear in (shared_experts.gate_up_proj, shared_experts.down_proj):
478
+ scale = linear.weight_scale if mxfp8 else linear.weight_scale_inv
479
+ if not fused:
480
+ # Already packed scales cannot be padded as checkpoint UE8M0 bytes.
481
+ linear.weight_scale_inv = scale_parameter(
482
+ linear.weight.shape[0], linear.weight.shape[1] // 128
483
+ )
484
+ continue
485
+ scale.data.view(torch.uint8).random_(124, 130)
486
+ block_m, block_k = linear.weight_block_size
487
+ scale_1x32 = (
488
+ experts._ue8m0_uint8_to_float(scale.view(torch.uint8))
489
+ .repeat_interleave(block_m, dim=0)
490
+ .repeat_interleave(block_k // 32, dim=1)
491
+ )
492
+ originals.append((linear.weight.view(torch.uint8).clone(), scale_1x32))
493
monkeypatch.setattr("vllm.utils.deep_gemm._import_deep_gemm", lambda: FakeDeepGemm)
494
495
original_gate_up_ptr = shared_experts.gate_up_proj.weight.data_ptr()
496
original_down_ptr = shared_experts.down_proj.weight.data_ptr()
497
experts.finalize_weights(shared_experts)
498
474
- assert FakeDeepGemm.scale_inputs[-2:] == [
475
- (1, 2 * intermediate_size, hidden_size // 32),
476
- (1, hidden_size, intermediate_size // 32),
477
- ]
499
assert experts.has_fused_shared_experts is fused
500
if not fused:
501
assert experts.intermediate_size == 2560
@@ -490,15 +511,56 @@ def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(
511
return
512
513
assert FakeDeepGemm.transformed_dims == [(3, 3), (2, 2)]
493
- assert shared_experts.gate_up_proj.weight.data_ptr() != original_gate_up_ptr
494
- assert (
495
- experts._transformed_shared_l1_weights[0].data_ptr()
496
- == shared_experts.gate_up_proj.weight.data_ptr()
497
- )
498
- assert (
499
- experts._transformed_shared_l2_weights[0].data_ptr()
500
- == shared_experts.down_proj.weight.data_ptr()
501
- )
514
+ padded_size = experts.intermediate_size
515
+ assert FakeDeepGemm.scale_inputs[-2:] == [
516
+ (1, 2 * padded_size, hidden_size // 32),
517
+ (1, hidden_size, padded_size // 32),
518
+ ]
519
+ for actual, original, fill in (
520
+ (
521
+ experts._transformed_shared_l1_weights[0].view(torch.uint8),
522
+ originals[0][0],
523
+ 0,
524
+ ),
525
+ (FakeDeepGemm.scale_values[-2][0], originals[0][1], 1),
526
+ ):
527
+ actual = actual.unflatten(0, (2, padded_size))
528
+ assert torch.equal(
529
+ actual[:, :intermediate_size], original.unflatten(0, (2, intermediate_size))
530
+ )
531
+ assert torch.all(actual[:, intermediate_size:] == fill)
532
+ for actual, original, fill in (
533
+ (
534
+ experts._transformed_shared_l2_weights[0].view(torch.uint8),
535
+ originals[1][0],
536
+ 0,
537
+ ),
538
+ (FakeDeepGemm.scale_values[-1][0], originals[1][1], 1),
539
+ ):
540
+ assert torch.equal(actual[:, : original.shape[1]], original)
541
+ assert torch.all(actual[:, original.shape[1] :] == fill)
542
+ if padded_size != intermediate_size:
543
+ # Generic linear post-load hooks still receive checkpoint-shaped weights.
544
+ for linear, (weight, _) in zip(
545
+ (shared_experts.gate_up_proj, shared_experts.down_proj), originals
546
+ ):
547
+ assert torch.equal(linear.weight.view(torch.uint8), weight)
548
+ assert shared_experts.gate_up_proj.weight.data_ptr() == original_gate_up_ptr
549
+ assert shared_experts.down_proj.weight.data_ptr() == original_down_ptr
550
+ else:
551
+ assert shared_experts.gate_up_proj.weight.data_ptr() != original_gate_up_ptr
552
+ assert (
553
+ experts._transformed_shared_l1_weights[0].data_ptr()
554
+ == shared_experts.gate_up_proj.weight.data_ptr()
555
+ )
556
+ assert (
557
+ experts._transformed_shared_l2_weights[0].data_ptr()
558
+ == shared_experts.down_proj.weight.data_ptr()
559
+ )
560
+
561
+ transformed_l1 = experts._transformed_shared_l1_weights
562
+ experts.finalize_weights(shared_experts)
563
+ assert experts._transformed_shared_l1_weights is transformed_l1
564
565
566
@pytest.mark.parametrize("fused", [False, True])
vllm/models/deepseek_v4/nvidia/model.py
+48
-5
@@ -204,6 +204,7 @@ class DeepseekV4MegaMoEExperts(nn.Module):
204
self.top_k = top_k
205
self.hidden_size = hidden_size
206
self.intermediate_size = intermediate_size
207
+ self.unpadded_intermediate_size = intermediate_size
208
self.num_shared_experts = num_shared_experts
209
self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
210
@@ -386,6 +387,17 @@ class DeepseekV4MegaMoEExperts(nn.Module):
387
# the generic linear post-load hook replaces the raw checkpoint scales
388
# with its 128x128 DeepGEMM layout.
389
checkpoint_scale_dtypes = (torch.float8_e8m0fnu, torch.uint8)
390
+ unpadded_size = self.unpadded_intermediate_size * self.num_shared_experts
391
+ padding = self.intermediate_size * self.num_shared_experts - unpadded_size
392
+ pad_weights = (
393
+ padding > 0
394
+ and gate_up_weight.dtype == torch.float8_e4m3fn
395
+ and down_weight.dtype == torch.float8_e4m3fn
396
+ and gate_up_weight.shape == (2 * unpadded_size, self.hidden_size)
397
+ and down_weight.shape == (self.hidden_size, unpadded_size)
398
+ and gate_up_scale.dtype in checkpoint_scale_dtypes
399
+ and down_scale.dtype in checkpoint_scale_dtypes
400
+ )
401
if (
402
gate_up_scale.dtype in checkpoint_scale_dtypes
403
and down_scale.dtype in checkpoint_scale_dtypes
@@ -396,6 +408,7 @@ class DeepseekV4MegaMoEExperts(nn.Module):
408
gate_up_scale,
409
gate_up_weight.shape[0],
410
gate_up_weight.shape[1],
411
+ padding=(padding, 0) if pad_weights else (0, 0),
412
)
413
down_scale = self._prepare_shared_expert_scale(
414
deep_gemm,
@@ -403,12 +416,27 @@ class DeepseekV4MegaMoEExperts(nn.Module):
416
down_scale,
417
down_weight.shape[0],
418
down_weight.shape[1],
419
+ padding=(0, padding) if pad_weights else (0, 0),
420
)
421
422
if gate_up_scale is None or down_scale is None:
423
self.num_shared_experts = 0
424
return
425
426
+ if pad_weights:
427
+ # Pad gate/up separately so the SwiGLU split stays at the midpoint.
428
+ gate_up_weight = (
429
+ torch.nn.functional.pad(
430
+ gate_up_weight.view(torch.uint8).unflatten(0, (2, unpadded_size)),
431
+ (0, 0, 0, padding),
432
+ )
433
+ .flatten(0, 1)
434
+ .view(gate_up_weight.dtype)
435
+ )
436
+ down_weight = torch.nn.functional.pad(
437
+ down_weight.view(torch.uint8), (0, padding)
438
+ ).view(down_weight.dtype)
439
+
440
shared_intermediate_size = self.intermediate_size * self.num_shared_experts
441
expected_gate_up_shape = (
442
2 * shared_intermediate_size,
@@ -449,11 +477,11 @@ class DeepseekV4MegaMoEExperts(nn.Module):
477
# released instead of adding roughly 0.7 GiB per rank on DSV4-Flash.
478
# The generic linear post-load hook may still repack the serial scales,
479
# but this shared MLP is never called after native fusion is enabled.
452
- gate_up.weight.data = transformed_l1[0]
453
- self._transformed_shared_l1_weights = (
454
- gate_up.weight.data,
455
- transformed_l1[1],
456
- )
480
+ # Padded weights need separate storage: the generic linear post-load
481
+ # hooks still expect the original checkpoint weight/scale shapes.
482
+ if not pad_weights:
483
+ gate_up.weight.data = transformed_l1[0]
484
+ self._transformed_shared_l1_weights = transformed_l1
485
self._transformed_shared_l2_weights = transformed_l2
486
487
def _prepare_shared_expert_scale(
@@ -463,6 +491,8 @@ class DeepseekV4MegaMoEExperts(nn.Module):
491
scale: torch.Tensor,
492
mn: int,
493
k: int,
494
+ *,
495
+ padding: tuple[int, int] = (0, 0),
496
) -> torch.Tensor | None:
497
block_size = getattr(linear, "weight_block_size", None)
498
if block_size is None or len(block_size) != 2:
@@ -497,6 +527,19 @@ class DeepseekV4MegaMoEExperts(nn.Module):
527
.repeat_interleave(block_k // 32, dim=1)[:mn, : k // 32]
528
.contiguous()
529
)
530
+ pad_m, pad_k = padding
531
+ if pad_m:
532
+ scale_1x32 = torch.nn.functional.pad(
533
+ scale_1x32.unflatten(0, (2, mn // 2)),
534
+ (0, 0, 0, pad_m),
535
+ value=1.0,
536
+ ).flatten(0, 1)
537
+ mn += 2 * pad_m
538
+ if pad_k:
539
+ scale_1x32 = torch.nn.functional.pad(
540
+ scale_1x32, (0, pad_k // 32), value=1.0
541
+ )
542
+ k += pad_k
543
# The grouped API is used with a singleton dimension to request the
544
# MN-major, TMA-aligned packed-UE8M0 strides, then squeezed back to the
545
# 2D layout required for a shared expert.