A high-throughput and memory-efficient inference and serving engine for LLMs
[ROCm][Perf] Use fused softplus-sqrt-topk router under AITER fused-MoE (#44945)
Co-authored-by: vLLM Contributor <contributor@vllm.ai>
Fangzhou Ai committed
Jun 9, 2026 at 13:53 UTC
01d8cd92dd963140bddc1e82df7df55d0efa920c
1 file changed
+26
-2
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
+26
-2
@@ -262,13 +262,37 @@ def fused_topk_bias(
262
topk_weights *= routed_scaling_factor
263
return topk_weights, topk_ids
264
265
+ if scoring_func == "sqrtsoftplus":
266
+ M = hidden_states.size(0)
267
+ topk_weights = torch.empty(
268
+ M, topk, dtype=torch.float32, device=hidden_states.device
269
+ )
270
+ topk_ids = torch.empty(
271
+ M,
272
+ topk,
273
+ dtype=torch.int32 if indices_type is None else indices_type,
274
+ device=hidden_states.device,
275
+ )
276
+ token_expert_indices = torch.empty(
277
+ M, topk, dtype=torch.int32, device=hidden_states.device
278
+ )
279
+ return vllm_topk_softplus_sqrt(
280
+ topk_weights,
281
+ topk_ids,
282
+ token_expert_indices,
283
+ gating_output,
284
+ renormalize,
285
+ e_score_correction_bias,
286
+ input_tokens,
287
+ hash_indices_table,
288
+ routed_scaling_factor,
289
+ )
290
+
291
n_routed_experts = gating_output.shape[-1]
292
if scoring_func == "softmax":
293
scores = gating_output.softmax(dim=-1)
294
elif scoring_func == "sigmoid":
295
scores = gating_output.sigmoid()
270
- elif scoring_func == "sqrtsoftplus":
271
- scores = F.softplus(gating_output).sqrt()
296
else:
297
raise ValueError(f"Unsupported scoring function: {scoring_func}")
298
if e_score_correction_bias is not None: