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 14, 23:45 UTC
[TPU][Bugfix] fix moe layer (#21340)
Signed-off-by: Chengji Yao <chengjiyao@google.com> Co-authored-by: Simon Mo <simon.mo@hey.com>
Chengji Yao committed
Jul 24, 2025 at 00:38 UTC
e74bfc70e4105017f6a73d1797be64021d268cdf
2 files changed
+19
-1
tests/v1/tpu/test_basic.py
+1
@@ -18,6 +18,7 @@ if TYPE_CHECKING:
18
19
MODELS = [
20
"Qwen/Qwen2.5-1.5B-Instruct",
21
+ "Qwen/Qwen1.5-MoE-A2.7B",
22
# TODO: Enable this models with v6e
23
# "Qwen/Qwen2-7B-Instruct",
24
# "meta-llama/Llama-3.1-8B",
vllm/model_executor/layers/fused_moe/layer.py
+18
-1
@@ -481,8 +481,16 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
481
e_score_correction_bias: Optional[torch.Tensor] = None,
482
apply_router_weight_on_input: bool = False,
483
activation: str = "silu",
484
- **kwargs,
484
+ enable_eplb: bool = False,
485
+ expert_load_view: Optional[torch.Tensor] = None,
486
+ logical_to_physical_map: Optional[torch.Tensor] = None,
487
+ logical_replica_count: Optional[torch.Tensor] = None,
488
):
489
+ if enable_eplb is not False or expert_load_view is not None or \
490
+ logical_to_physical_map is not None or \
491
+ logical_replica_count is not None:
492
+ raise NotImplementedError("Expert load balancing is not supported "
493
+ "for CPU.")
494
return layer.cpu_fused_moe(
495
layer,
496
x,
@@ -518,6 +526,10 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
526
e_score_correction_bias: Optional[torch.Tensor] = None,
527
apply_router_weight_on_input: bool = False,
528
activation: str = "silu",
529
+ enable_eplb: bool = False,
530
+ expert_load_view: Optional[torch.Tensor] = None,
531
+ logical_to_physical_map: Optional[torch.Tensor] = None,
532
+ logical_replica_count: Optional[torch.Tensor] = None,
533
) -> torch.Tensor:
534
assert not use_grouped_topk
535
assert num_expert_group is None
@@ -531,6 +543,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
543
raise NotImplementedError(
544
"Expert score correction bias is not supported for TPU.")
545
assert activation == "silu", f"{activation} is not supported for TPU."
546
+ if enable_eplb is not False or expert_load_view is not None or \
547
+ logical_to_physical_map is not None or \
548
+ logical_replica_count is not None:
549
+ raise NotImplementedError("Expert load balancing is not supported "
550
+ "for TPU.")
551
return fused_moe_pallas(hidden_states=x,
552
w1=layer.w13_weight,
553
w2=layer.w2_weight,