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 13, 17:00 UTC
[flashinfer] fix FI all2all with FI cutlass moe (#28166)
Signed-off-by: Xiaozhu <mxz297@gmail.com>
Xiaozhu Meng committed
Nov 5, 2025 at 21:52 UTC
e31946f86eb7975bb1483cff04ea52ba9b5aa79c
1 file changed
+3
-1
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
+3
-1
@@ -233,12 +233,13 @@ def flashinfer_alltoall_dispatch(
233
max_num_token = (
234
max(global_num_tokens_cpu) if global_num_tokens_cpu is not None else x.shape[0]
235
)
236
+ orig_topk_weights_dtype = topk_weights.dtype
237
alltoall_info, topk_ids, topk_weights, _ = (
238
MnnvlMoe.mnnvl_moe_alltoallv_prepare_without_allgather(
239
topk_ids,
240
topk_weights,
241
None,
241
- all2all_manager.prepare_workspace,
242
+ all2all_manager.prepare_workspace_tensor,
243
max_num_token,
244
ep_rank,
245
ep_size,
@@ -247,6 +248,7 @@ def flashinfer_alltoall_dispatch(
248
top_k,
249
)
250
)
251
+ topk_weights = topk_weights.view(dtype=orig_topk_weights_dtype)
252
253
x, x_sf = moe_kernel_quantize_input(
254
x,