A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 11:00 UTC
[Bugfix] [ROCm] [DSV4] Fix AITER MXFP4 MoE weight loading and shuffle… (#42595)
Co-authored-by: MHYangAMD <MHYangAMD@users.noreply.github.com>
MHYangAMD committed
May 29, 2026 at 19:08 UTC
ab12aab12711ca91316c58da2fc77bae32a9e6ce
2 files changed
+37
-52
vllm/model_executor/layers/fused_moe/layer.py
+14
-14
@@ -727,16 +727,19 @@ class FusedMoE(PluggableLayer):
727
# Only narrow if the loaded_weight is not a scalar (0-dim tensor)
728
# and we're not loading the full weight
729
if not load_full and loaded_weight.ndim > 0:
730
- # Handle padding: loaded_weight might be smaller than shard_size on last
731
- # TP rank
732
- start_offset = shard_size * tp_rank
730
+ # When the parameter has been padded (e.g. MXFP4 rounding up
731
+ # intermediate_size_per_partition), shard_size is the padded
732
+ # size. Compute the offset into the checkpoint weight using
733
+ # the *unpadded* per-rank size so that every TP rank lands at
734
+ # the correct slice.
735
+ tp_size = self.moe_config.moe_parallel_config.tp_size
736
+ loaded_per_rank = loaded_weight.shape[shard_dim] // tp_size
737
+ start_offset = loaded_per_rank * tp_rank
738
available = loaded_weight.shape[shard_dim] - start_offset
739
if available <= 0:
740
# If there is no available weight to load for this TP rank
736
- # (can happen on last TP rank with padding), we can skip
737
- # loading and return early
741
return
739
- narrow_size = min(shard_size, available)
742
+ narrow_size = min(loaded_per_rank, available)
743
loaded_weight = loaded_weight.narrow(shard_dim, start_offset, narrow_size)
744
# Narrow parameter and load.
745
# w1, gate_proj: Load into first logical weight of w13.
@@ -765,21 +768,18 @@ class FusedMoE(PluggableLayer):
768
):
769
# Index the loaded weight for tp sharding.
770
# down_proj: "RowParallel" so tp sharding on input_dim
768
- # Narrow parameter and load.
769
- shard_size = expert_data.shape[shard_dim]
771
# Only narrow if the loaded_weight is not a scalar (0-dim tensor)
772
# and we're not loading the full weight
773
if not load_full and loaded_weight.ndim > 0:
773
- # Handle padding: loaded_weight might be smaller than shard_size on last
774
- # TP rank
775
- start_offset = shard_size * tp_rank
774
+ # Same padding fix as _load_w13: use unpadded per-rank size.
775
+ tp_size = self.moe_config.moe_parallel_config.tp_size
776
+ loaded_per_rank = loaded_weight.shape[shard_dim] // tp_size
777
+ start_offset = loaded_per_rank * tp_rank
778
available = loaded_weight.shape[shard_dim] - start_offset
779
if available <= 0:
780
# If there is no available weight to load for this TP rank
779
- # (can happen on last TP rank with padding), we can skip
780
- # loading and return early
781
return
782
- narrow_size = min(shard_size, available)
782
+ narrow_size = min(loaded_per_rank, available)
783
loaded_weight = loaded_weight.narrow(shard_dim, start_offset, narrow_size)
784
# w2, down_proj: Load into only logical weight of w2.
785
hidden_dim = self._get_hidden_dim(shard_dim, expert_data.ndim)
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+23
-38
@@ -622,16 +622,15 @@ def select_deepseek_v4_mxfp4_moe_backend(
622
assert last_error is not None
623
raise last_error
624
625
- # DeepSeek-V4 on ROCm is more accurate with the unfused Triton MXFP4 path
626
- # than the default AITER path. Prefer Triton-unfused for this routing mode,
627
- # while keeping AITER as a fallback if Triton-unfused rejects the config.
625
+ # DeepSeek-V4 on ROCm: prefer AITER FlyDSL MoE (better perf + accuracy
626
+ # after shuffle/TP-offset fixes), with Triton-unfused as fallback.
627
if (
628
current_platform.is_rocm()
629
and config.routing_method == RoutingMethodType.DeepseekV4
630
):
631
priority_backends = [
633
- Mxfp4MoeBackend.TRITON_UNFUSED,
632
Mxfp4MoeBackend.AITER_MXFP4_BF16,
633
+ Mxfp4MoeBackend.TRITON_UNFUSED,
634
]
635
else:
636
priority_backends = _get_priority_backends()
@@ -1421,7 +1420,7 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1420
)
1421
1422
elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
1424
- from vllm._aiter_ops import rocm_aiter_ops
1423
+ from vllm._aiter_ops import rocm_aiter_ops # noqa: F401
1424
1425
if w13_bias is not None:
1426
w13_bias = w13_bias.data.to(torch.float32)
@@ -1430,45 +1429,31 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1429
1430
e, n, k = w13_weight.shape
1431
1433
- w13_weight.view(torch.uint8).copy_(
1434
- w13_weight.data.view(torch.uint8)
1435
- .view(e, n // 2, 2, k)
1436
- .permute(0, 2, 1, 3)
1437
- .contiguous()
1438
- .view(e, n, k)
1432
+ # No de-interleave: standard _load_w13 already produces
1433
+ # [gate_all, up_all] layout. Use aiter-native shuffle functions
1434
+ # (matching aiter/ops/flydsl/test_flydsl_moe_a4w4.py pattern).
1435
+ from aiter.ops.shuffle import shuffle_weight as _shuf_w
1436
+ from aiter.utility.fp4_utils import e8m0_shuffle as _e8m0_shuf
1437
+
1438
+ # w13 (gate+up, stage1): shuffle_weight with layout (16,16)
1439
+ w13_weight = torch.nn.Parameter(
1440
+ _shuf_w(w13_weight.data.view(torch.float4_e2m1fn_x2), (16, 16)),
1441
+ requires_grad=False,
1442
)
1440
- w13_weight_scale.data = (
1441
- w13_weight_scale.data.view(e, n // 2, 2, -1)
1442
- .permute(0, 2, 1, 3)
1443
- .contiguous()
1444
- .view(e, n, -1)
1443
+ shuffled_w13_scale = _e8m0_shuf(
1444
+ w13_weight_scale.view(-1, w13_weight_scale.shape[-1])
1445
)
1446
1447
- w13_weight.data = w13_weight.data.view(torch.float4_e2m1fn_x2)
1448
- w2_weight.data = w2_weight.data.view(torch.float4_e2m1fn_x2)
1449
-
1450
- w13_weight.data = rocm_aiter_ops.shuffle_weight_a16w4(w13_weight, 16, True)
1451
- shuffled_w13_scale = rocm_aiter_ops.shuffle_scale_a16w4(
1452
- w13_weight_scale.view(-1, w13_weight_scale.shape[-1]),
1453
- num_experts,
1454
- True,
1447
+ # w2 (down-proj, stage2): same shuffle as w13 for a4w4 fp4x2
1448
+ # (tuning script uses shuffle_weight((16,16)) + e8m0_shuffle for both)
1449
+ w2_weight = torch.nn.Parameter(
1450
+ _shuf_w(w2_weight.data.view(torch.float4_e2m1fn_x2), (16, 16)),
1451
+ requires_grad=False,
1452
)
1456
-
1457
- w2_weight.data = rocm_aiter_ops.shuffle_weight_a16w4(w2_weight, 16, False)
1458
- shuffled_w2_scale = rocm_aiter_ops.shuffle_scale_a16w4(
1459
- w2_weight_scale.view(-1, w2_weight_scale.shape[-1]),
1460
- num_experts,
1461
- False,
1453
+ shuffled_w2_scale = _e8m0_shuf(
1454
+ w2_weight_scale.view(-1, w2_weight_scale.shape[-1])
1455
)
1456
1464
- if w13_bias is not None:
1465
- w13_bias = (
1466
- w13_bias.data.view(-1, n // 2, 2)
1467
- .permute(0, 2, 1)
1468
- .contiguous()
1469
- .view(-1, n)
1470
- )
1471
-
1457
return (
1458
w13_weight,
1459
w2_weight,