[ROCm] [Performance] Optimize aiter moe for DeepSeekV4 (#46122)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
TJian committed
Jun 26, 2026 at 21:43 UTC
8921c4be88effbd295dd7c2410fd21411256f819
1 file changed
+28
-15
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+28
-15
@@ -1429,38 +1429,51 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1429
)
1430
1431
elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
1432
- from vllm._aiter_ops import rocm_aiter_ops # noqa: F401
1432
+ # Initially introduced for DeepSeekV4
1433
1434
if w13_bias is not None:
1435
w13_bias = w13_bias.data.to(torch.float32)
1436
if w2_bias is not None:
1437
w2_bias = w2_bias.data.to(torch.float32)
1438
1439
- e, n, k = w13_weight.shape
1439
+ import os
1440
1441
- # No de-interleave: standard _load_w13 already produces
1442
- # [gate_all, up_all] layout. Use aiter-native shuffle functions
1443
- # (matching aiter/ops/flydsl/test_flydsl_moe_a4w4.py pattern).
1441
+ from aiter.ops.shuffle import shuffle_scale as _shuf_s
1442
from aiter.ops.shuffle import shuffle_weight as _shuf_w
1445
- from aiter.utility.fp4_utils import e8m0_shuffle as _e8m0_shuf
1443
1447
- # w13 (gate+up, stage1): shuffle_weight with layout (16,16)
1444
+ # TODO: Remove this once AITER is fixed
1445
+ # Necessary for AITER side from crashing
1446
+ os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0"
1447
+
1448
w13_weight = torch.nn.Parameter(
1449
- _shuf_w(w13_weight.data.view(torch.float4_e2m1fn_x2), (16, 16)),
1449
+ _shuf_w(
1450
+ w13_weight.data.view(torch.float4_e2m1fn_x2),
1451
+ is_guinterleave=True,
1452
+ gate_up=True,
1453
+ ),
1454
requires_grad=False,
1455
)
1452
- shuffled_w13_scale = _e8m0_shuf(
1453
- w13_weight_scale.view(-1, w13_weight_scale.shape[-1])
1456
+ shuffled_w13_scale = _shuf_s(
1457
+ w13_weight_scale.reshape(-1, w13_weight_scale.shape[-1]),
1458
+ num_experts,
1459
+ True,
1460
+ True,
1461
)
1462
1456
- # w2 (down-proj, stage2): same shuffle as w13 for a4w4 fp4x2
1457
- # (tuning script uses shuffle_weight((16,16)) + e8m0_shuffle for both)
1463
w2_weight = torch.nn.Parameter(
1459
- _shuf_w(w2_weight.data.view(torch.float4_e2m1fn_x2), (16, 16)),
1464
+ _shuf_w(
1465
+ w2_weight.data.view(torch.float4_e2m1fn_x2),
1466
+ is_guinterleave=True,
1467
+ gate_up=False,
1468
+ ),
1469
requires_grad=False,
1470
)
1462
- shuffled_w2_scale = _e8m0_shuf(
1463
- w2_weight_scale.view(-1, w2_weight_scale.shape[-1])
1471
+ # use_gu_interleave
1472
+ shuffled_w2_scale = _shuf_s(
1473
+ w2_weight_scale.reshape(-1, w2_weight_scale.shape[-1]),
1474
+ num_experts,
1475
+ True,
1476
+ False,
1477
)
1478
1479
return (