Remove the XPU branch of topk_softplus_sqrt (#51242)
Signed-off-by: xiaolong <xiaolong.guo@intel.com>
Liangqiusong committed
Aug 6, 2026 at 20:18 UTC
1e05b21d61e6126e4811313f39c961bf8b314470
1 file changed
-62
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
-62
@@ -3,7 +3,6 @@
3
import functools
4
5
import torch
6
-import torch.nn.functional as F
6
7
import vllm._custom_ops as ops
8
import vllm.envs as envs
@@ -72,52 +71,6 @@ def vllm_topk_sigmoid(
71
return topk_weights, topk_indices
72
73
75
-def _topk_softplus_sqrt_torch(
76
- topk_weights: torch.Tensor,
77
- topk_indices: torch.Tensor,
78
- token_expert_indices: torch.Tensor,
79
- gating_output: torch.Tensor,
80
- renormalize: bool = False,
81
- e_score_correction_bias: torch.Tensor | None = None,
82
- input_tokens: torch.Tensor | None = None,
83
- hash_indices_table: torch.Tensor | None = None,
84
- routed_scaling_factor: float = 1.0,
85
-) -> tuple[torch.Tensor, ...]:
86
- """Pure PyTorch fallback for topk_softplus_sqrt (XPU/CPU)."""
87
- # scores = sqrt(softplus(gating_output))
88
- scores = torch.sqrt(F.softplus(gating_output.float()))
89
-
90
- # Bias is used for expert SELECTION only, not for weight computation.
91
- # Using biased scores as weights flattens the distribution when the bias
92
- # is near-uniform (e.g., DSv4-Flash where all biases ≈ 8.08).
93
- if e_score_correction_bias is not None:
94
- scores_for_choice = scores + e_score_correction_bias.float()
95
- else:
96
- scores_for_choice = scores
97
-
98
- topk = topk_weights.shape[-1]
99
-
100
- if hash_indices_table is not None and input_tokens is not None:
101
- # Hash MoE: expert indices predetermined by lookup table
102
- # hash_indices_table: [vocab_size, topk] mapping token_id -> expert_ids
103
- expert_ids = hash_indices_table[input_tokens.long()] # [M, topk]
104
- topk_indices.copy_(expert_ids)
105
- # Gather weights from unbiased scores
106
- weights = scores.gather(1, expert_ids.long())
107
- else:
108
- # Standard topk selection using biased scores
109
- _, indices = torch.topk(scores_for_choice, k=topk, dim=-1)
110
- topk_indices.copy_(indices)
111
- # Gather weights from unbiased scores
112
- weights = scores.gather(1, indices)
113
-
114
- if renormalize:
115
- weights = weights / (weights.sum(dim=-1, keepdim=True).clamp(min=1e-20))
116
-
117
- topk_weights.copy_(weights * routed_scaling_factor)
118
- return topk_weights, topk_indices
119
-
120
-
74
def vllm_topk_softplus_sqrt(
75
topk_weights: torch.Tensor,
76
topk_indices: torch.Tensor,
@@ -129,21 +82,6 @@ def vllm_topk_softplus_sqrt(
82
hash_indices_table: torch.Tensor | None = None,
83
routed_scaling_factor: float = 1.0,
84
) -> tuple[torch.Tensor, ...]:
132
- from vllm.platforms import current_platform
133
-
134
- if current_platform.is_xpu():
135
- return _topk_softplus_sqrt_torch(
136
- topk_weights,
137
- topk_indices,
138
- token_expert_indices,
139
- gating_output,
140
- renormalize,
141
- e_score_correction_bias,
142
- input_tokens,
143
- hash_indices_table,
144
- routed_scaling_factor,
145
- )
146
-
85
ops.topk_hash_softplus_sqrt(
86
topk_weights,
87
topk_indices,