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 12, 22:15 UTC
[Refactor] Optimize `select_experts` (#28069)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Nov 19, 2025 at 18:53 UTC
5031cd5d55ad99e8f9b31dd0020a06b346f6e493
5 files changed
+7
-15
vllm/model_executor/layers/fused_moe/fused_moe.py
-5
@@ -1246,7 +1246,6 @@ def eplb_map_to_physical_and_record(
1246
expert_load_view: torch.Tensor,
1247
logical_to_physical_map: torch.Tensor,
1248
logical_replica_count: torch.Tensor,
1249
- indices_type: torch.dtype | None = None,
1249
) -> torch.Tensor:
1250
"""
1251
Map the logical expert ids to physical expert ids
@@ -1260,7 +1259,6 @@ def eplb_map_to_physical_and_record(
1259
expert_load_view: The expert load view.
1260
logical_to_physical_map: The logical to physical map.
1261
logical_replica_count: The logical replica count.
1263
- indices_type: The indices type.
1262
1263
Returns:
1264
The physical expert ids.
@@ -1310,9 +1308,6 @@ def eplb_map_to_physical_and_record(
1308
index=topk_ids_flatten.long(),
1309
src=torch.ones_like(topk_ids_flatten).to(expert_load_view),
1310
)
1313
-
1314
- if indices_type is not None:
1315
- topk_ids = topk_ids.to(dtype=indices_type)
1311
return topk_ids
1312
1313
vllm/model_executor/layers/fused_moe/layer.py
+4
-7
@@ -68,7 +68,6 @@ else:
68
expert_load_view: torch.Tensor,
69
logical_to_physical_map: torch.Tensor,
70
logical_replica_count: torch.Tensor,
71
- indices_type: torch.dtype | None,
71
) -> torch.Tensor:
72
# CPU fallback: no EPLB so just return as is
73
return topk_ids
@@ -1509,8 +1508,6 @@ class FusedMoE(CustomOp):
1508
routed_scaling_factor=routed_scaling_factor,
1509
e_score_correction_bias=e_score_correction_bias,
1510
)
1512
- if indices_type is not None:
1513
- topk_ids = topk_ids.to(dtype=indices_type)
1511
elif e_score_correction_bias is not None:
1512
topk_weights, topk_ids = fused_topk_bias(
1513
hidden_states=hidden_states,
@@ -1519,7 +1516,7 @@ class FusedMoE(CustomOp):
1516
topk=top_k,
1517
renormalize=renormalize,
1518
)
1522
- if routed_scaling_factor is not None:
1519
+ if routed_scaling_factor != 1.0:
1520
topk_weights *= routed_scaling_factor
1521
elif custom_routing_function is None:
1522
topk_weights, topk_ids, token_expert_indices = fused_topk(
@@ -1536,8 +1533,6 @@ class FusedMoE(CustomOp):
1533
topk=top_k,
1534
renormalize=renormalize,
1535
)
1539
- if indices_type is not None:
1540
- topk_ids = topk_ids.to(dtype=indices_type)
1536
1537
if enable_eplb:
1538
assert expert_load_view is not None
@@ -1549,9 +1544,11 @@ class FusedMoE(CustomOp):
1544
expert_load_view=expert_load_view,
1545
logical_to_physical_map=logical_to_physical_map,
1546
logical_replica_count=logical_replica_count,
1552
- indices_type=indices_type,
1547
)
1548
1549
+ if (indices_type is not None) and topk_ids.dtype != indices_type:
1550
+ topk_ids = topk_ids.to(dtype=indices_type)
1551
+
1552
assert topk_ids.dtype == indices_type or indices_type is None
1553
1554
# Compute zero expert result if needed
vllm/model_executor/layers/quantization/modelopt.py
+1
-1
@@ -1706,7 +1706,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1706
intermediate_size=layer.intermediate_size_per_partition,
1707
local_expert_offset=layer.ep_rank * layer.local_num_experts,
1708
local_num_experts=layer.local_num_experts,
1709
- routed_scaling_factor=None,
1709
+ routed_scaling_factor=1.0,
1710
tile_tokens_dim=None,
1711
routing_method_type=routing_method_type,
1712
do_finalize=True,
vllm/model_executor/models/longcat_flash.py
+1
-1
@@ -118,7 +118,7 @@ class FlashConfig(PretrainedConfig):
118
router_dtype="float32",
119
router_bias=False,
120
topk_method=None,
121
- routed_scaling_factor=None,
121
+ routed_scaling_factor=1.0,
122
zero_expert_num=0,
123
zero_expert_type=None,
124
nextn_use_scmoe=False,
vllm/model_executor/models/openpangu.py
+1
-1
@@ -625,7 +625,7 @@ class OpenPanguDecoderLayer(nn.Module):
625
bias=getattr(config, "mlp_bias", False),
626
prefix=f"{prefix}.mlp",
627
)
628
- self.routed_scaling_factor = getattr(config, "routed_scaling_factor", None)
628
+ self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
629
self.num_hidden_layers = config.num_hidden_layers
630
self.first_k_dense_replace = getattr(
631
config, "first_k_dense_replace", self.num_hidden_layers