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 10, 22:45 UTC
CustomOp: Unify aiter impl into GroupedTopk (#31221)
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com>
Xinyu Chen committed
Dec 27, 2025 at 01:44 UTC
87f1b8ca2c94d173ac8b2f149ffbdd55b9adc800
2 files changed
+40
-26
vllm/model_executor/layers/fused_moe/fused_moe.py
+31
@@ -35,6 +35,9 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
35
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
36
MoEPrepareAndFinalizeNoEP,
37
)
38
+from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa: E501
39
+ rocm_aiter_grouped_topk,
40
+)
41
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
42
TopKWeightAndReduceNoOP,
43
)
@@ -1295,6 +1298,7 @@ class GroupedTopk(CustomOp):
1298
topk_group: int = 0,
1299
scoring_func: str = "softmax",
1300
routed_scaling_factor: float = 1.0,
1301
+ num_fused_shared_experts: int = 0,
1302
) -> None:
1303
super().__init__()
1304
self.native_impl = grouped_topk
@@ -1304,6 +1308,7 @@ class GroupedTopk(CustomOp):
1308
self.topk_group = topk_group
1309
self.scoring_func = scoring_func
1310
self.routed_scaling_factor = routed_scaling_factor
1311
+ self.num_fused_shared_experts = num_fused_shared_experts
1312
1313
def forward_native(
1314
self,
@@ -1333,6 +1338,32 @@ class GroupedTopk(CustomOp):
1338
hidden_states, gating_output, e_score_correction_bias
1339
)
1340
1341
+ def forward_hip(
1342
+ self,
1343
+ hidden_states: torch.Tensor,
1344
+ gating_output: torch.Tensor,
1345
+ e_score_correction_bias: torch.Tensor | None = None,
1346
+ ) -> tuple[torch.Tensor, torch.Tensor]:
1347
+ if rocm_aiter_ops.is_fused_moe_enabled():
1348
+ if not rocm_aiter_ops.is_fusion_moe_shared_experts_enabled():
1349
+ assert self.num_fused_shared_experts == 0
1350
+ return rocm_aiter_grouped_topk(
1351
+ hidden_states,
1352
+ gating_output,
1353
+ self.topk,
1354
+ self.renormalize,
1355
+ self.num_expert_group,
1356
+ self.topk_group,
1357
+ self.scoring_func,
1358
+ self.routed_scaling_factor,
1359
+ e_score_correction_bias,
1360
+ self.num_fused_shared_experts,
1361
+ )
1362
+ else:
1363
+ return self.forward_native(
1364
+ hidden_states, gating_output, e_score_correction_bias
1365
+ )
1366
+
1367
1368
@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
1369
def eplb_map_to_physical_and_record(
vllm/model_executor/layers/fused_moe/layer.py
+9
-26
@@ -4,7 +4,6 @@
4
from collections.abc import Callable, Iterable
5
from contextlib import nullcontext
6
from enum import Enum
7
-from functools import partial
7
from typing import Literal, cast, get_args, overload
8
9
import torch
@@ -67,9 +66,6 @@ else:
66
67
eplb_map_to_physical_and_record = _eplb_map_to_physical_and_record
68
from vllm.model_executor.layers.fused_moe.fused_moe import GroupedTopk
70
-from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa: E501
71
- rocm_aiter_grouped_topk,
72
-)
69
70
if current_platform.is_tpu():
71
from .moe_pallas import fused_moe as fused_moe_pallas
@@ -1583,28 +1579,15 @@ class FusedMoE(CustomOp):
1579
elif self.use_grouped_topk and valid_grouping():
1580
assert self.topk_group is not None
1581
assert self.num_expert_group is not None
1586
- if rocm_aiter_ops.is_fused_moe_enabled():
1587
- if not rocm_aiter_ops.is_fusion_moe_shared_experts_enabled():
1588
- assert self.num_fused_shared_experts == 0
1589
- grouped_topk_impl = partial(
1590
- rocm_aiter_grouped_topk,
1591
- num_fused_shared_experts=self.num_fused_shared_experts,
1592
- topk=self.top_k,
1593
- renormalize=self.renormalize,
1594
- num_expert_group=self.num_expert_group,
1595
- topk_group=self.topk_group,
1596
- scoring_func=self.scoring_func,
1597
- routed_scaling_factor=self.routed_scaling_factor,
1598
- )
1599
- else:
1600
- grouped_topk_impl = GroupedTopk(
1601
- topk=self.top_k,
1602
- renormalize=self.renormalize,
1603
- num_expert_group=self.num_expert_group,
1604
- topk_group=self.topk_group,
1605
- scoring_func=self.scoring_func,
1606
- routed_scaling_factor=self.routed_scaling_factor,
1607
- )
1582
+ grouped_topk_impl = GroupedTopk(
1583
+ topk=self.top_k,
1584
+ renormalize=self.renormalize,
1585
+ num_expert_group=self.num_expert_group,
1586
+ topk_group=self.topk_group,
1587
+ scoring_func=self.scoring_func,
1588
+ routed_scaling_factor=self.routed_scaling_factor,
1589
+ num_fused_shared_experts=self.num_fused_shared_experts,
1590
+ )
1591
1592
topk_weights, topk_ids = grouped_topk_impl(
1593
hidden_states=hidden_states,