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 11, 08:45 UTC
CustomOp: grouped topk (#29575)
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com>
Xinyu Chen committed
Dec 17, 2025 at 17:43 UTC
3b1d440ede42855f031ba72af4817583e5dddba0
4 files changed
+75
-14
tests/kernels/moe/test_grouped_topk.py
+6
-4
@@ -9,8 +9,8 @@ import pytest
9
import torch
10
11
from vllm.model_executor.layers.fused_moe.fused_moe import (
12
+ GroupedTopk,
13
fused_grouped_topk,
13
- grouped_topk,
14
)
15
from vllm.platforms import current_platform
16
@@ -50,15 +50,17 @@ def test_grouped_topk(
50
51
with monkeypatch.context() as m:
52
m.setenv("VLLM_USE_FUSED_MOE_GROUPED_TOPK", "0")
53
- baseline_topk_weights, baseline_topk_ids = grouped_topk(
54
- hidden_states=hidden_states,
55
- gating_output=gating_output,
53
+ grouped_topk = GroupedTopk(
54
topk=topk,
55
renormalize=renormalize,
56
num_expert_group=num_expert_group,
57
topk_group=topk_group,
58
scoring_func=scoring_func,
59
routed_scaling_factor=routed_scaling_factor,
60
+ )
61
+ baseline_topk_weights, baseline_topk_ids = grouped_topk(
62
+ hidden_states=hidden_states,
63
+ gating_output=gating_output,
64
e_score_correction_bias=e_score_correction_bias,
65
)
66
vllm/model_executor/layers/fused_moe/__init__.py
+2
-2
@@ -77,11 +77,11 @@ if HAS_TRITON:
77
BatchedTritonExperts,
78
)
79
from vllm.model_executor.layers.fused_moe.fused_moe import (
80
+ GroupedTopk,
81
TritonExperts,
82
fused_experts,
83
fused_topk,
84
get_config_file_name,
84
- grouped_topk,
85
)
86
from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
87
TritonOrDeepGemmExperts,
@@ -91,7 +91,7 @@ if HAS_TRITON:
91
"fused_topk",
92
"fused_experts",
93
"get_config_file_name",
94
- "grouped_topk",
94
+ "GroupedTopk",
95
"cutlass_moe_fp8",
96
"cutlass_moe_fp4",
97
"cutlass_moe_w4a8_fp8",
vllm/model_executor/layers/fused_moe/fused_moe.py
+52
@@ -16,6 +16,7 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
16
from vllm import _custom_ops as ops
17
from vllm._aiter_ops import rocm_aiter_ops
18
from vllm.logger import init_logger
19
+from vllm.model_executor.custom_op import CustomOp
20
from vllm.model_executor.layers.batch_invariant import (
21
vllm_is_batch_invariant,
22
)
@@ -1286,6 +1287,57 @@ def grouped_topk(
1287
return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
1288
1289
1290
+@CustomOp.register("grouped_topk")
1291
+class GroupedTopk(CustomOp):
1292
+ """GroupedTopk used by the Deepseek-V2 and Deepseek-V3 model."""
1293
+
1294
+ def __init__(
1295
+ self,
1296
+ topk: int,
1297
+ renormalize: bool,
1298
+ num_expert_group: int = 0,
1299
+ topk_group: int = 0,
1300
+ scoring_func: str = "softmax",
1301
+ routed_scaling_factor: float = 1.0,
1302
+ ) -> None:
1303
+ super().__init__()
1304
+ self.native_impl = grouped_topk
1305
+ self.topk = topk
1306
+ self.renormalize = renormalize
1307
+ self.num_expert_group = num_expert_group
1308
+ self.topk_group = topk_group
1309
+ self.scoring_func = scoring_func
1310
+ self.routed_scaling_factor = routed_scaling_factor
1311
+
1312
+ def forward_native(
1313
+ self,
1314
+ hidden_states: torch.Tensor,
1315
+ gating_output: torch.Tensor,
1316
+ e_score_correction_bias: torch.Tensor | None = None,
1317
+ ) -> tuple[torch.Tensor, torch.Tensor]:
1318
+ return self.native_impl(
1319
+ hidden_states,
1320
+ gating_output,
1321
+ self.topk,
1322
+ self.renormalize,
1323
+ self.num_expert_group,
1324
+ self.topk_group,
1325
+ self.scoring_func,
1326
+ self.routed_scaling_factor,
1327
+ e_score_correction_bias,
1328
+ )
1329
+
1330
+ def forward_cuda(
1331
+ self,
1332
+ hidden_states: torch.Tensor,
1333
+ gating_output: torch.Tensor,
1334
+ e_score_correction_bias: torch.Tensor | None = None,
1335
+ ) -> tuple[torch.Tensor, torch.Tensor]:
1336
+ return self.forward_native(
1337
+ hidden_states, gating_output, e_score_correction_bias
1338
+ )
1339
+
1340
+
1341
@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
1342
def eplb_map_to_physical_and_record(
1343
topk_ids: torch.Tensor,
vllm/model_executor/layers/fused_moe/layer.py
+15
-8
@@ -67,7 +67,7 @@ else:
67
return topk_ids
68
69
eplb_map_to_physical_and_record = _eplb_map_to_physical_and_record
70
-from vllm.model_executor.layers.fused_moe.fused_moe import grouped_topk
70
+from vllm.model_executor.layers.fused_moe.fused_moe import GroupedTopk
71
from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa: E501
72
rocm_aiter_grouped_topk,
73
)
@@ -1594,19 +1594,26 @@ class FusedMoE(CustomOp):
1594
grouped_topk_impl = partial(
1595
rocm_aiter_grouped_topk,
1596
num_fused_shared_experts=self.num_fused_shared_experts,
1597
+ topk=self.top_k,
1598
+ renormalize=self.renormalize,
1599
+ num_expert_group=self.num_expert_group,
1600
+ topk_group=self.topk_group,
1601
+ scoring_func=self.scoring_func,
1602
+ routed_scaling_factor=self.routed_scaling_factor,
1603
)
1604
else:
1599
- grouped_topk_impl = grouped_topk
1605
+ grouped_topk_impl = GroupedTopk(
1606
+ topk=self.top_k,
1607
+ renormalize=self.renormalize,
1608
+ num_expert_group=self.num_expert_group,
1609
+ topk_group=self.topk_group,
1610
+ scoring_func=self.scoring_func,
1611
+ routed_scaling_factor=self.routed_scaling_factor,
1612
+ )
1613
1614
topk_weights, topk_ids = grouped_topk_impl(
1615
hidden_states=hidden_states,
1616
gating_output=router_logits,
1604
- topk=self.top_k,
1605
- renormalize=self.renormalize,
1606
- num_expert_group=self.num_expert_group,
1607
- topk_group=self.topk_group,
1608
- scoring_func=self.scoring_func,
1609
- routed_scaling_factor=self.routed_scaling_factor,
1617
e_score_correction_bias=self.e_score_correction_bias,
1618
)
1619
elif self.e_score_correction_bias is not None: