[V1] [ROCm] Enable EP with AITER Fused MoE (#20270)

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
2025-07-01 09:48:30 -07:00
parent 3d19d47d91
commit 02cabff207
4 changed files with 15 additions and 5 deletions
--- a/vllm/model_executor/layers/fused_moe/layer.py
+++ b/vllm/model_executor/layers/fused_moe/layer.py
@@ -646,13 +646,13 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
            indices_type=self.topk_indices_dtype)

        if self.rocm_aiter_moe_enabled:
-            assert expert_map is None
            return self.rocm_aiter_fused_experts(
                hidden_states=x,
                w1=layer.w13_weight,
                w2=layer.w2_weight,
                topk_weights=topk_weights,
                topk_ids=topk_ids,
+                expert_map=expert_map,
                activation=activation,
                apply_router_weight_on_input=apply_router_weight_on_input)
        else: