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 7, 00:15 UTC
[ROCM] Optimize the fused_topk_bias to use aiter instead of fallback torch ops. (#36253)
Signed-off-by: zhutaoyu <zhutaoyu97@gmail.com>
Taoyu Zhu committed
Mar 10, 2026 at 00:30 UTC
70485a11bd83afa50e6ecc8e9619d9bdd0ff2039
1 file changed
+38
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
+38
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import functools
4
from collections.abc import Callable
5
6
import torch
@@ -57,6 +58,19 @@ def vllm_topk_sigmoid(
58
return topk_weights, topk_indices
59
60
61
+@functools.lru_cache(maxsize=8)
62
+def _aiter_get_num_expert_group(num_experts: int) -> int:
63
+ _AITER_MAX_EXPERTS_PER_GROUP = 32
64
+ g = max(1, -(-num_experts // _AITER_MAX_EXPERTS_PER_GROUP))
65
+ while num_experts % g != 0:
66
+ g += 1
67
+ assert num_experts % g == 0, f"{num_experts=} not divisible by {g=}"
68
+ assert num_experts // g <= _AITER_MAX_EXPERTS_PER_GROUP, (
69
+ f"group size {num_experts // g} exceeds limit {_AITER_MAX_EXPERTS_PER_GROUP}"
70
+ )
71
+ return g
72
+
73
+
74
def fused_topk_bias(
75
hidden_states: torch.Tensor,
76
gating_output: torch.Tensor,
@@ -108,6 +122,30 @@ def fused_topk_bias(
122
return topk_weights, topk_ids
123
else:
124
raise ValueError(f"Unsupported scoring function: {scoring_func}")
125
+ elif rocm_aiter_ops.is_fused_moe_enabled() and scoring_func == "sigmoid":
126
+ M = hidden_states.size(0)
127
+ num_experts = gating_output.shape[-1]
128
+ num_expert_group = _aiter_get_num_expert_group(num_experts)
129
+ if topk >= num_expert_group:
130
+ topk_weights = torch.empty(
131
+ M, topk, dtype=torch.float32, device=hidden_states.device
132
+ )
133
+ topk_ids = torch.empty(
134
+ M,
135
+ topk,
136
+ dtype=torch.int32 if indices_type is None else indices_type,
137
+ device=hidden_states.device,
138
+ )
139
+ rocm_aiter_ops.biased_grouped_topk(
140
+ gating_output,
141
+ e_score_correction_bias.to(gating_output.dtype),
142
+ topk_weights,
143
+ topk_ids,
144
+ num_expert_group=num_expert_group,
145
+ topk_group=num_expert_group,
146
+ need_renorm=renormalize,
147
+ )
148
+ return topk_weights, topk_ids
149
150
n_routed_experts = gating_output.shape[-1]
151
if scoring_func == "softmax":