[ROCm] [Bugfix] Fix Triton fused shared expert alignment (#51632)
Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
akii96 committed
Aug 19, 2026 at 19:14 UTC
583a00257d4c5d1a54063d956057df1df6822b06
2 files changed
+50
-2
tests/kernels/moe/test_moe.py
+44
@@ -459,6 +459,50 @@ def test_fused_moe(
459
)
460
461
462
+def test_fused_shared_expert_alignment(workspace_init):
463
+ set_random_seed(7)
464
+ m, n, k = 4, 64, 128
465
+ routed_experts = 8
466
+ physical_experts = routed_experts + 1
467
+ dtype = torch.bfloat16
468
+
469
+ a = torch.randn((m, k), device=DEVICE_TYPE, dtype=dtype) / 10
470
+ w1 = torch.randn((physical_experts, 2 * n, k), device=DEVICE_TYPE, dtype=dtype) / 10
471
+ w2 = torch.randn((physical_experts, k, n), device=DEVICE_TYPE, dtype=dtype) / 10
472
+ topk_ids = torch.tensor(
473
+ [[0, 8], [1, 8], [2, 8], [3, 8]], device=DEVICE_TYPE, dtype=torch.int32
474
+ )
475
+ topk_weights = torch.tensor(
476
+ [[0.5, 1.0]] * m, device=DEVICE_TYPE, dtype=torch.float32
477
+ )
478
+
479
+ moe_config = make_dummy_moe_config(
480
+ num_experts=physical_experts,
481
+ experts_per_token=2,
482
+ hidden_dim=k,
483
+ intermediate_size=n,
484
+ in_dtype=dtype,
485
+ max_num_tokens=m,
486
+ )
487
+ modular_moe = modular_triton_fused_moe(moe_config, FUSED_MOE_UNQUANTIZED_CONFIG)
488
+
489
+ with set_current_vllm_config(vllm_config):
490
+ expected = torch_experts(a, w1, w2, topk_weights, topk_ids)
491
+ actual = modular_moe.apply(
492
+ hidden_states=a,
493
+ w1=w1,
494
+ w2=w2,
495
+ topk_weights=topk_weights,
496
+ topk_ids=topk_ids,
497
+ activation=MoEActivation.SILU,
498
+ global_num_experts=routed_experts,
499
+ expert_map=None,
500
+ apply_router_weight_on_input=False,
501
+ )
502
+
503
+ torch.testing.assert_close(actual, expected, atol=2e-2, rtol=0)
504
+
505
+
506
def test_fused_moe_int64_overflow(workspace_init):
507
"""Regression test for int32 overflow in stride*offset products.
508
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+6
-2
@@ -315,13 +315,15 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
315
)
316
intermediate_cache3 = _resize_cache(workspace2, (num_tokens, top_k_num, K))
317
318
+ # Include fused shared-expert rows while preserving EP remapping.
319
+ num_align_experts = w1.shape[0] if expert_map is None else global_num_experts
320
sorted_token_ids, expert_ids, num_tokens_post_padded = (
321
_prepare_expert_assignment(
322
topk_ids,
323
config,
324
num_tokens,
325
top_k_num,
324
- global_num_experts,
326
+ num_align_experts,
327
expert_map,
328
use_int8_w8a16=self.quant_config.use_int8_w8a16,
329
use_int4_w4a16=self.quant_config.use_int4_w4a16,
@@ -688,8 +690,10 @@ class TritonWNA16Experts(TritonExperts):
690
)
691
intermediate_cache3 = _resize_cache(workspace2, (num_tokens, top_k_num, K))
692
693
+ # Include fused shared-expert rows while preserving EP remapping.
694
+ num_align_experts = w1.shape[0] if expert_map is None else global_num_experts
695
sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
692
- topk_ids, config["BLOCK_SIZE_M"], global_num_experts, expert_map
696
+ topk_ids, config["BLOCK_SIZE_M"], num_align_experts, expert_map
697
)
698
699
invoke_fused_moe_wna16_triton_kernel(