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 15, 16:00 UTC
[V1] TPU - Fix fused MOE (#15834)
Signed-off-by: Alexander Matveev <amatveev@redhat.com>
Alexander Matveev committed
Apr 1, 2025 at 01:58 UTC
7e4e709b43b6b35011b78247f8b0afde7e41961b
1 file changed
+1
-1
vllm/model_executor/layers/fused_moe/layer.py
+1
-1
@@ -309,7 +309,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
309
expert_map=expert_map,
310
renormalize=renormalize)
311
312
- forward_native = forward_cuda
312
+ forward_native = forward_tpu if current_platform.is_tpu else forward_cuda
313
314
315
def determine_expert_map(