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 5, 21:00 UTC
[ROCm] Enable MORI EP for unquantized MoE with AITER backend (#37529)
Signed-off-by: Tan Pin Siang <pinsiang.tan@amd.com>
Tan Pin Siang committed
Mar 30, 2026 at 15:19 UTC
85c0950b1f647e0b0654fbf3e91a9757b8233752
2 files changed
+24
-9
vllm/model_executor/layers/fused_moe/all2all_utils.py
+12
-5
@@ -186,16 +186,23 @@ def maybe_make_prepare_finalize(
186
use_fp8_dispatch = (
187
quant_config.is_per_act_token or quant_config.is_block_quantized
188
)
189
- # For PTPC (per token per channel) quant, the scale dim for each token is 1
190
- # For 1x128 quant, the scale dim for each token is hidden_dim // 128
191
- scale_dim = 1 if quant_config.is_per_act_token else moe.hidden_dim // 128
189
+ if use_fp8_dispatch:
190
+ # For PTPC (per token per channel) quant, scale dim is 1
191
+ # For 1x128 quant, scale dim is hidden_dim // 128
192
+ quant_dtype = quant_config.quant_dtype
193
+ scale_dim = 1 if quant_config.is_per_act_token else moe.hidden_dim // 128
194
+ else:
195
+ # Unquantized dispatch (e.g. AITER with defer_input_quant):
196
+ # dispatch raw BF16/FP16 data, no scales needed.
197
+ quant_dtype = moe.in_dtype
198
+ scale_dim = 0
199
all_to_all_args = dict(
200
rank=all2all_manager.rank,
201
num_ep_ranks=all2all_manager.world_size,
195
- quant_dtype=quant_config.quant_dtype,
202
+ quant_dtype=quant_dtype,
203
token_hidden_size=moe.hidden_dim,
204
scale_dim=scale_dim,
198
- scale_type_size=torch.float32.itemsize,
205
+ scale_type_size=0 if scale_dim == 0 else torch.float32.itemsize,
206
max_num_tokens_per_dp_rank=moe.max_num_tokens,
207
input_dtype=moe.in_dtype,
208
num_local_experts=moe.num_experts // all2all_manager.world_size,
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+12
-4
@@ -108,10 +108,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
108
self,
109
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
110
) -> FusedMoEPrepareAndFinalizeModular | None:
111
- if self.unquantized_backend == UnquantizedMoeBackend.AITER:
112
- return None
113
- else:
114
- return super().maybe_make_prepare_finalize(routing_tables)
111
+ return super().maybe_make_prepare_finalize(routing_tables)
112
113
def select_gemm_impl(
114
self,
@@ -130,6 +127,17 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
127
max_num_tokens=self.moe.max_num_tokens,
128
num_dispatchers=prepare_finalize.num_dispatchers(),
129
)
130
+ elif (
131
+ self.unquantized_backend == UnquantizedMoeBackend.AITER
132
+ and rocm_aiter_ops.is_fused_moe_enabled()
133
+ ):
134
+ from .rocm_aiter_fused_moe import AiterExperts
135
+
136
+ logger.debug("AiterExperts %s", self.moe)
137
+ return AiterExperts(
138
+ moe_config=self.moe,
139
+ quant_config=self.moe_quant_config,
140
+ )
141
else:
142
logger.debug("TritonExperts %s", self.moe)
143
return TritonExperts(