[Perf] Optimize dsv4 routing using specialized kernel, 2.94% E2E TPOT improvement (#48660)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jul 17, 2026 at 16:35 UTC
b5433b6f5079feb32f9f278cf4ae23bd87375148
4 files changed
+264
csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu
+78
@@ -71,6 +71,73 @@ __device__ __forceinline__ float toFloat(T value) {
71
}
72
}
73
74
+#ifndef USE_ROCM
75
+// Adapted from:
76
+// https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/csrc/deepseek_v4/hash_topk.cuh
77
+template <typename OutIndType, typename HashIndType>
78
+__launch_bounds__(128) __global__
79
+ void dsv4HashTopkSoftplusSqrt(const float* input, float* output,
80
+ OutIndType* indices, int num_rows,
81
+ int num_experts, float routed_scaling_factor,
82
+ const HashIndType* input_ids,
83
+ const HashIndType* tid2eid) {
84
+ const int warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
85
+ const int lane = threadIdx.x % 32;
86
+ if (warp >= num_rows) return;
87
+ const int64_t token_id = load_index_as_int64(input_ids, warp);
88
+
89
+ #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
90
+ cudaGridDependencySynchronize();
91
+ #endif
92
+ int expert = 0;
93
+ float weight = 0.f;
94
+ if (lane < 6) {
95
+ // only load and calculate for 6 experts
96
+ expert = static_cast<int>(tid2eid[token_id * 6 + lane]);
97
+ const float x = input[warp * num_experts + expert];
98
+ weight = sqrtf(fmaxf(x, 0.f) + __logf(1.f + __expf(-fabsf(x))));
99
+ }
100
+ float weight_sum = weight;
101
+ #pragma unroll
102
+ for (int mask = 16; mask > 0; mask >>= 1) {
103
+ // sum in warp
104
+ weight_sum += VLLM_SHFL_XOR_SYNC(weight_sum, mask);
105
+ }
106
+
107
+ #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
108
+ cudaTriggerProgrammaticLaunchCompletion();
109
+ #endif
110
+ if (lane < 6) {
111
+ const int offset = warp * 6 + lane;
112
+ output[offset] =
113
+ weight * routed_scaling_factor / (weight_sum > 0.f ? weight_sum : 1.f);
114
+ indices[offset] = static_cast<OutIndType>(expert);
115
+ }
116
+}
117
+
118
+template <typename OutIndType, typename HashIndType>
119
+void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,
120
+ int num_rows, int num_experts,
121
+ double routed_scaling_factor,
122
+ const HashIndType* input_ids,
123
+ const HashIndType* tid2eid, cudaStream_t stream) {
124
+ if (num_rows == 0) return;
125
+ auto* kernel = &dsv4HashTopkSoftplusSqrt<OutIndType, HashIndType>;
126
+ cudaLaunchConfig_t config = {};
127
+ config.gridDim = (num_rows + 3) / 4;
128
+ config.blockDim = 128;
129
+ config.stream = stream;
130
+ cudaLaunchAttribute attr;
131
+ attr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
132
+ attr.val.programmaticStreamSerializationAllowed = 1;
133
+ config.attrs = &attr;
134
+ config.numAttrs = 1;
135
+ const float scale = static_cast<float>(routed_scaling_factor);
136
+ cudaLaunchKernelEx(&config, kernel, input, output, indices, num_rows,
137
+ num_experts, scale, input_ids, tid2eid);
138
+}
139
+#endif
140
+
141
// ====================== TopK softplus_sqrt things
142
// ===============================
143
@@ -556,6 +623,17 @@ void topkGatingSoftplusSqrtKernelLauncher(
623
const float* correction_bias, const bool use_hash,
624
const HashIndType* input_ids, const HashIndType* tid2eid,
625
cudaStream_t stream) {
626
+#ifndef USE_ROCM
627
+ if constexpr (std::is_same_v<InputType, float>) {
628
+ if (use_hash && topk == 6 && renormalize &&
629
+ (num_experts == 256 || num_experts == 384)) {
630
+ launchDsv4HashTopk<IndType, HashIndType>(
631
+ gating_output, topk_weights, topk_indices, num_tokens, num_experts,
632
+ routed_scaling_factor, input_ids, tid2eid, stream);
633
+ return;
634
+ }
635
+ }
636
+#endif
637
static constexpr int WARPS_PER_TB = 4;
638
static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
639
// for bfloat16 dtype, we need 4 bytes loading to make sure num_experts
tests/kernels/moe/test_topk_softplus_sqrt.py
+45
@@ -10,6 +10,7 @@ from vllm.model_executor.layers.fused_moe.config import (
10
RoutingMethodType,
11
get_routing_method_type,
12
)
13
+from vllm.model_executor.layers.fused_moe.router.dsv4_topk import dsv4_topk
14
from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
15
fused_topk_bias,
16
)
@@ -186,3 +187,47 @@ def test_fused_topk_softplus_sqrt_hash(
187
sorted_w_ref = topk_weights_ref.gather(1, idx_ref)
188
sorted_w = topk_weights.gather(1, idx_ops)
189
torch.testing.assert_close(sorted_w_ref, sorted_w, atol=2e-2, rtol=1e-2)
190
+
191
+
192
+@pytest.mark.skipif(
193
+ not current_platform.is_cuda(),
194
+ reason="The DeepSeek V4 fast path is CUDA-only.",
195
+)
196
+@pytest.mark.parametrize(
197
+ ("num_tokens", "num_experts", "indices_type"),
198
+ [
199
+ (0, 256, torch.uint32),
200
+ (17, 256, torch.uint32),
201
+ (17, 384, torch.int64),
202
+ ],
203
+)
204
+def test_dsv4_fast_topk(
205
+ num_tokens: int,
206
+ num_experts: int,
207
+ indices_type: torch.dtype,
208
+):
209
+ torch.manual_seed(0)
210
+ gating_output = torch.randn(
211
+ (num_tokens, num_experts), dtype=torch.float32, device="cuda"
212
+ )
213
+ correction_bias = torch.randn(num_experts, dtype=torch.float32, device="cuda")
214
+
215
+ topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt(
216
+ gating_output=gating_output,
217
+ topk=6,
218
+ renormalize=True,
219
+ routed_scaling_factor=1.5,
220
+ e_score_correction_bias=correction_bias,
221
+ )
222
+ topk_weights, topk_ids = dsv4_topk(
223
+ gating_output, correction_bias, indices_type, 1.5
224
+ )
225
+
226
+ assert topk_ids.dtype == indices_type
227
+ torch.testing.assert_close(topk_ids_ref.to(indices_type), topk_ids, atol=0, rtol=0)
228
+ torch.testing.assert_close(
229
+ topk_weights_ref,
230
+ topk_weights,
231
+ atol=2e-5,
232
+ rtol=2e-5,
233
+ )
vllm/model_executor/layers/fused_moe/router/dsv4_topk.py
new
+121
@@ -0,0 +1,121 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+from vllm.platforms import current_platform
7
+from vllm.triton_utils import tl, triton
8
+
9
+_TOPK = 6
10
+
11
+# Adapted from:
12
+# https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/moe_fused_gate.py
13
+
14
+
15
+def can_use_dsv4_topk(
16
+ gating_output: torch.Tensor,
17
+ correction_bias: torch.Tensor | None,
18
+ topk: int,
19
+ renormalize: bool,
20
+ indices_dtype: torch.dtype,
21
+) -> bool:
22
+ return (
23
+ current_platform.is_cuda()
24
+ and gating_output.dtype == torch.float32
25
+ and gating_output.ndim == 2
26
+ and gating_output.shape[1] in (256, 384)
27
+ and gating_output.is_contiguous()
28
+ and correction_bias is not None
29
+ and correction_bias.dtype == torch.float32
30
+ and correction_bias.shape == (gating_output.shape[1],)
31
+ and correction_bias.is_contiguous()
32
+ and topk == _TOPK
33
+ and renormalize
34
+ and indices_dtype in (torch.int32, torch.uint32, torch.int64)
35
+ )
36
+
37
+
38
+if current_platform.is_cuda():
39
+
40
+ @triton.jit
41
+ def _dsv4_topk_kernel(
42
+ gating_output_ptr,
43
+ correction_bias_ptr,
44
+ topk_weights_ptr,
45
+ topk_ids_ptr,
46
+ routed_scaling_factor,
47
+ NUM_EXPERTS: tl.constexpr,
48
+ BLOCK_N: tl.constexpr,
49
+ launch_pdl: tl.constexpr,
50
+ ):
51
+ row = tl.program_id(0)
52
+ expert_offsets = tl.arange(0, BLOCK_N)
53
+ expert_mask = expert_offsets < NUM_EXPERTS
54
+ bias = tl.load(
55
+ correction_bias_ptr + expert_offsets, mask=expert_mask, other=0.0
56
+ ).to(tl.float32)
57
+
58
+ if launch_pdl:
59
+ tl.extra.cuda.gdc_wait()
60
+
61
+ logits = tl.load(
62
+ gating_output_ptr + row * NUM_EXPERTS + expert_offsets,
63
+ mask=expert_mask,
64
+ other=0.0,
65
+ ).to(tl.float32)
66
+ weights = tl.sqrt(tl.where(logits > 20.0, logits, tl.log(1.0 + tl.exp(logits))))
67
+ current = tl.where(expert_mask, weights + bias, -float("inf"))
68
+ current = tl.where(current == current, current, -1e30)
69
+
70
+ topk_offsets = tl.arange(0, 8)
71
+ selected_weights = tl.zeros([8], dtype=tl.float32)
72
+ selected_ids = tl.zeros([8], dtype=tl.int32)
73
+ for slot in tl.static_range(6):
74
+ max_value = tl.max(current, axis=0)
75
+ candidate = tl.where(current == max_value, expert_offsets, NUM_EXPERTS)
76
+ expert_id = tl.min(candidate, axis=0).to(tl.int32)
77
+ selected_weight = tl.sum(
78
+ tl.where(expert_offsets == expert_id, weights, 0.0), axis=0
79
+ )
80
+ is_slot = topk_offsets == slot
81
+ selected_weights = tl.where(is_slot, selected_weight, selected_weights)
82
+ selected_ids = tl.where(is_slot, expert_id, selected_ids)
83
+ current = tl.where(expert_offsets == expert_id, -float("inf"), current)
84
+
85
+ weight_sum = tl.sum(selected_weights, axis=0)
86
+ selected_weights *= routed_scaling_factor / tl.where(
87
+ weight_sum > 0.0, weight_sum, 1.0
88
+ )
89
+ output_mask = topk_offsets < 6
90
+ output_offsets = row * 6 + topk_offsets
91
+
92
+ if launch_pdl:
93
+ tl.extra.cuda.gdc_launch_dependents()
94
+
95
+ tl.store(topk_weights_ptr + output_offsets, selected_weights, mask=output_mask)
96
+ tl.store(topk_ids_ptr + output_offsets, selected_ids, mask=output_mask)
97
+
98
+
99
+def dsv4_topk(
100
+ gating_output: torch.Tensor,
101
+ correction_bias: torch.Tensor,
102
+ indices_dtype: torch.dtype,
103
+ routed_scaling_factor: float,
104
+) -> tuple[torch.Tensor, torch.Tensor]:
105
+ num_tokens, num_experts = gating_output.shape
106
+ shape = (num_tokens, _TOPK)
107
+ topk_weights = gating_output.new_empty(shape, dtype=torch.float32)
108
+ topk_ids = gating_output.new_empty(shape, dtype=indices_dtype)
109
+ if num_tokens > 0:
110
+ _dsv4_topk_kernel[(num_tokens,)](
111
+ gating_output,
112
+ correction_bias,
113
+ topk_weights,
114
+ topk_ids,
115
+ routed_scaling_factor,
116
+ NUM_EXPERTS=num_experts,
117
+ BLOCK_N=triton.next_power_of_2(num_experts),
118
+ num_warps=1,
119
+ launch_pdl=current_platform.is_arch_support_pdl(),
120
+ )
121
+ return topk_weights, topk_ids
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
+20
@@ -14,6 +14,10 @@ from vllm.model_executor.layers.fused_moe.config import (
14
get_routing_method_type,
15
)
16
from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
17
+from vllm.model_executor.layers.fused_moe.router.dsv4_topk import (
18
+ can_use_dsv4_topk,
19
+ dsv4_topk,
20
+)
21
22
23
def vllm_topk_softmax(
@@ -182,6 +186,22 @@ def fused_topk_bias(
186
"Number of tokens mismatch"
187
)
188
189
+ output_indices_dtype = torch.int32 if indices_type is None else indices_type
190
+ if scoring_func == "sqrtsoftplus" and can_use_dsv4_topk(
191
+ gating_output,
192
+ e_score_correction_bias,
193
+ topk,
194
+ renormalize,
195
+ output_indices_dtype,
196
+ ):
197
+ assert e_score_correction_bias is not None
198
+ return dsv4_topk(
199
+ gating_output,
200
+ e_score_correction_bias,
201
+ output_indices_dtype,
202
+ routed_scaling_factor,
203
+ )
204
+
205
M, _ = hidden_states.size()
206
207
topk_weights = torch.empty(