[Perf][DSV4.1] Remove MegaMoE padding and shared padding workaround (#57204)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Sep 16, 2026 at 12:09 UTC
30b847c1fa5461a590606cbc9b8dae4a490c6cdd
2 files changed
+42
-157
tests/models/test_deepseek_v4_mega_moe.py
+37
-92
@@ -300,8 +300,10 @@ def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership():
300
301
302
@pytest.mark.parametrize("intermediate_size", [512, 2304])
303
-def test_deepseek_v4_mega_moe_padding_preserves_weights(monkeypatch, intermediate_size):
304
- """Pad gate/up separately and zero every added down-projection column."""
303
+def test_deepseek_v4_mega_moe_preserves_checkpoint_dimensions(
304
+ monkeypatch, intermediate_size
305
+):
306
+ """Keep native widths so V4.1 routed and shared experts can fuse."""
307
experts = DeepseekV4MegaMoEExperts(
308
SimpleNamespace(
309
scheduler_config=SimpleNamespace(max_num_batched_tokens=4),
@@ -333,29 +335,15 @@ def test_deepseek_v4_mega_moe_padding_preserves_weights(monkeypatch, intermediat
335
)
336
337
experts.finalize_weights()
336
- padded_size = experts.intermediate_size
337
- assert padded_size == (intermediate_size + 511) // 512 * 512
338
- for actual, original in zip(experts._transformed_l1_weights, originals[:2]):
339
- expected = (
340
- original
341
- if actual.dtype == torch.int8
342
- else (experts._ue8m0_uint8_to_float(original))
343
- )
344
- actual = actual.view(expected.dtype).unflatten(1, (2, padded_size))
345
- assert torch.equal(
346
- actual[:, :, :intermediate_size],
347
- expected.unflatten(1, (2, intermediate_size)),
348
- )
349
- assert torch.count_nonzero(actual[:, :, intermediate_size:]) == 0
350
- for actual, original in zip(experts._transformed_l2_weights, originals[2:]):
338
+ assert experts.intermediate_size == intermediate_size
339
+ transformed = (*experts._transformed_l1_weights, *experts._transformed_l2_weights)
340
+ for actual, original in zip(transformed, originals):
341
expected = (
342
original
343
if actual.dtype == torch.int8
354
- else (experts._ue8m0_uint8_to_float(original))
344
+ else experts._ue8m0_uint8_to_float(original)
345
)
356
- actual = actual.view(expected.dtype)
357
- assert torch.equal(actual[..., : original.shape[-1]], expected)
358
- assert torch.count_nonzero(actual[..., original.shape[-1] :]) == 0
346
+ assert torch.equal(actual.view(expected.dtype), expected)
347
348
transformed_l1 = experts._transformed_l1_weights
349
experts.finalize_weights()
@@ -363,20 +351,19 @@ def test_deepseek_v4_mega_moe_padding_preserves_weights(monkeypatch, intermediat
351
352
353
@pytest.mark.parametrize(
366
- "hidden_size,intermediate_size,block_size,fused,mxfp8",
354
+ "hidden_size,intermediate_size,block_size,mxfp8",
355
[
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, 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"),
356
+ pytest.param(128, 512, 128, False, id="aligned-block128"),
357
+ pytest.param(128, 512, 128, True, id="aligned-block128-mxfp8"),
358
+ pytest.param(5120, 2304, 32, False, id="deepseek-v41-flash"),
359
+ pytest.param(128, 2304, 128, False, id="native-block128"),
360
+ pytest.param(5120, 2304, 32, True, id="deepseek-v41-flash-mxfp8"),
361
],
362
)
363
def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(
377
- monkeypatch, hidden_size, intermediate_size, block_size, fused, mxfp8
364
+ monkeypatch, hidden_size, intermediate_size, block_size, mxfp8
365
):
379
- """Shared fusion preserves checkpoint channels and zeros padded channels."""
366
+ """Shared fusion preserves checkpoint weights and scales at native widths."""
367
368
class FakeDeepGemm:
369
transformed_dims: list[tuple[int, int]] = []
@@ -476,12 +463,6 @@ def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(
463
originals = []
464
for linear in (shared_experts.gate_up_proj, shared_experts.down_proj):
465
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
466
scale.data.view(torch.uint8).random_(124, 130)
467
block_m, block_k = linear.weight_block_size
468
scale_1x32 = (
@@ -493,70 +474,34 @@ def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(
474
monkeypatch.setattr("vllm.utils.deep_gemm._import_deep_gemm", lambda: FakeDeepGemm)
475
476
original_gate_up_ptr = shared_experts.gate_up_proj.weight.data_ptr()
496
- original_down_ptr = shared_experts.down_proj.weight.data_ptr()
477
experts.finalize_weights(shared_experts)
478
499
- assert experts.has_fused_shared_experts is fused
500
- if not fused:
501
- assert experts.intermediate_size == 2560
502
- assert FakeDeepGemm.transformed_dims == [(3, 3)]
503
- assert experts.num_shared_experts == 0
504
- assert shared_experts.gate_up_proj.weight.data_ptr() == original_gate_up_ptr
505
- assert shared_experts.down_proj.weight.data_ptr() == original_down_ptr
506
- assert shared_experts.gate_up_proj.weight.shape == (
507
- 2 * intermediate_size,
508
- hidden_size,
509
- )
510
- assert shared_experts.down_proj.weight.shape == (hidden_size, intermediate_size)
511
- return
512
-
479
+ assert experts.has_fused_shared_experts
480
+ assert experts.intermediate_size == intermediate_size
481
assert FakeDeepGemm.transformed_dims == [(3, 3), (2, 2)]
514
- padded_size = experts.intermediate_size
482
assert FakeDeepGemm.scale_inputs[-2:] == [
516
- (1, 2 * padded_size, hidden_size // 32),
517
- (1, hidden_size, padded_size // 32),
483
+ (1, 2 * intermediate_size, hidden_size // 32),
484
+ (1, hidden_size, intermediate_size // 32),
485
]
519
- for actual, original, fill in (
486
+ for transformed, actual_scale, (weight, scale) in zip(
487
(
521
- experts._transformed_shared_l1_weights[0].view(torch.uint8),
522
- originals[0][0],
523
- 0,
488
+ experts._transformed_shared_l1_weights,
489
+ experts._transformed_shared_l2_weights,
490
),
525
- (FakeDeepGemm.scale_values[-2][0], originals[0][1], 1),
491
+ FakeDeepGemm.scale_values[-2:],
492
+ originals,
493
):
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
- )
494
+ assert torch.equal(transformed[0].view(torch.uint8), weight)
495
+ assert torch.equal(actual_scale[0], scale)
496
+ assert shared_experts.gate_up_proj.weight.data_ptr() != original_gate_up_ptr
497
+ assert (
498
+ experts._transformed_shared_l1_weights[0].data_ptr()
499
+ == shared_experts.gate_up_proj.weight.data_ptr()
500
+ )
501
+ assert (
502
+ experts._transformed_shared_l2_weights[0].data_ptr()
503
+ == shared_experts.down_proj.weight.data_ptr()
504
+ )
505
506
transformed_l1 = experts._transformed_shared_l1_weights
507
experts.finalize_weights(shared_experts)
vllm/models/deepseek_v4/nvidia/model.py
+5
-65
@@ -204,7 +204,6 @@ 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
207
self.num_shared_experts = num_shared_experts
208
self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
209
@@ -387,17 +386,6 @@ class DeepseekV4MegaMoEExperts(nn.Module):
386
# the generic linear post-load hook replaces the raw checkpoint scales
387
# with its 128x128 DeepGEMM layout.
388
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
- )
389
if (
390
gate_up_scale.dtype in checkpoint_scale_dtypes
391
and down_scale.dtype in checkpoint_scale_dtypes
@@ -408,7 +396,6 @@ class DeepseekV4MegaMoEExperts(nn.Module):
396
gate_up_scale,
397
gate_up_weight.shape[0],
398
gate_up_weight.shape[1],
411
- padding=(padding, 0) if pad_weights else (0, 0),
399
)
400
down_scale = self._prepare_shared_expert_scale(
401
deep_gemm,
@@ -416,27 +403,12 @@ class DeepseekV4MegaMoEExperts(nn.Module):
403
down_scale,
404
down_weight.shape[0],
405
down_weight.shape[1],
419
- padding=(0, padding) if pad_weights else (0, 0),
406
)
407
408
if gate_up_scale is None or down_scale is None:
409
self.num_shared_experts = 0
410
return
411
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
-
412
shared_intermediate_size = self.intermediate_size * self.num_shared_experts
413
expected_gate_up_shape = (
414
2 * shared_intermediate_size,
@@ -477,11 +449,11 @@ class DeepseekV4MegaMoEExperts(nn.Module):
449
# released instead of adding roughly 0.7 GiB per rank on DSV4-Flash.
450
# The generic linear post-load hook may still repack the serial scales,
451
# but this shared MLP is never called after native fusion is enabled.
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
452
+ gate_up.weight.data = transformed_l1[0]
453
+ self._transformed_shared_l1_weights = (
454
+ gate_up.weight.data,
455
+ transformed_l1[1],
456
+ )
457
self._transformed_shared_l2_weights = transformed_l2
458
459
def _prepare_shared_expert_scale(
@@ -491,8 +463,6 @@ class DeepseekV4MegaMoEExperts(nn.Module):
463
scale: torch.Tensor,
464
mn: int,
465
k: int,
494
- *,
495
- padding: tuple[int, int] = (0, 0),
466
) -> torch.Tensor | None:
467
block_size = getattr(linear, "weight_block_size", None)
468
if block_size is None or len(block_size) != 2:
@@ -527,19 +497,6 @@ class DeepseekV4MegaMoEExperts(nn.Module):
497
.repeat_interleave(block_k // 32, dim=1)[:mn, : k // 32]
498
.contiguous()
499
)
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
500
# The grouped API is used with a singleton dimension to request the
501
# MN-major, TMA-aligned packed-UE8M0 strides, then squeezed back to the
502
# 2D layout required for a shared expert.
@@ -558,23 +515,6 @@ class DeepseekV4MegaMoEExperts(nn.Module):
515
516
if self._transformed_l1_weights is None:
517
self._check_runtime_supported()
561
- # MegaMoE's 1x32 activation scales need 16-byte TMA rows, so
562
- # pad each gate/up half and the down projection to a multiple of 512.
563
- padded_size = (self.intermediate_size + 511) // 512 * 512
564
- padding = padded_size - self.intermediate_size
565
- if padding:
566
- for param in (self.w13_weight, self.w13_weight_scale):
567
- gate_up = param.data.unflatten(1, (2, self.intermediate_size))
568
- param.data = torch.nn.functional.pad(
569
- gate_up, (0, 0, 0, padding)
570
- ).flatten(1, 2)
571
- self.w2_weight.data = torch.nn.functional.pad(
572
- self.w2_weight.data, (0, padding // 2)
573
- )
574
- self.w2_weight_scale.data = torch.nn.functional.pad(
575
- self.w2_weight_scale.data, (0, padding // 32)
576
- )
577
- self.intermediate_size = padded_size
518
w13_scale = deep_gemm.transform_sf_into_required_layout(
519
self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(),
520
2 * self.intermediate_size,