[Perf] Optimize `fused_topk_bias` for DSv4, 1.5~2x kernel performance improvement (#47463)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jul 15, 2026 at 19:40 UTC
2dab187f756f2c4549fd5e57019413f3415365a4
3 files changed
+76
-67
csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu
+67
-58
@@ -44,6 +44,12 @@ typedef __hip_bfloat162 __nv_bfloat162;
44
namespace vllm {
45
namespace moe {
46
47
+template <typename HashIndType>
48
+__device__ __forceinline__ int64_t load_index_as_int64(const HashIndType* ptr,
49
+ int64_t offset) {
50
+ return static_cast<int64_t>(ptr[offset]);
51
+}
52
+
53
/// Aligned array type
54
template <typename T,
55
/// Number of elements in the array
@@ -86,14 +92,14 @@ __device__ __forceinline__ float toFloat(T value) {
92
93
template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG,
94
int WARP_SIZE_PARAM, bool USE_HASH, typename IndType,
89
- typename InputType = float>
95
+ typename HashIndType, typename InputType = float>
96
__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
97
void topkGatingSoftplusSqrt(
98
const InputType* input, const bool* finished, float* output,
99
const int num_rows, IndType* indices, int* source_rows, const int k,
100
const int start_expert, const int end_expert, const bool renormalize,
101
double routed_scaling_factor, const float* correction_bias,
96
- const IndType* input_ids, const IndType* tid2eid) {
102
+ const HashIndType* input_ids, const HashIndType* tid2eid) {
103
static_assert(std::is_same_v<InputType, float> ||
104
std::is_same_v<InputType, __nv_bfloat16> ||
105
std::is_same_v<InputType, __half>,
@@ -240,8 +246,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
246
247
// Hash MoE path: indices are predetermined from lookup table
248
if constexpr (USE_HASH) {
243
- const IndType token_id = input_ids[thread_row];
244
- const IndType* expert_indices_for_token = tid2eid + token_id * k;
249
+ const int64_t token_id = load_index_as_int64(input_ids, thread_row);
250
+ const int64_t token_expert_offset = token_id * static_cast<int64_t>(k);
251
#pragma unroll
252
for (int ii = 0; ii < VPT; ++ii) {
253
float val = row_chunk[ii];
@@ -252,7 +258,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
258
float selected_sum = 0.f;
259
#pragma unroll
260
for (int k_idx = 0; k_idx < k; ++k_idx) {
255
- const int expert = expert_indices_for_token[k_idx];
261
+ const int expert = static_cast<int>(
262
+ load_index_as_int64(tid2eid, token_expert_offset + k_idx));
263
const int idx = k * thread_row + k_idx;
264
for (int ii = 0; ii < VPT; ++ii) {
265
const int group_id = ii / ELTS_PER_LDG;
@@ -261,7 +268,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
268
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
269
local_id;
270
if (expert == expert_idx) {
264
- indices[idx] = expert;
271
+ indices[idx] = static_cast<IndType>(expert);
272
selected_sum += row_chunk[ii];
273
break;
274
}
@@ -285,7 +292,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
292
293
#pragma unroll
294
for (int k_idx = 0; k_idx < k; ++k_idx) {
288
- const int expert = expert_indices_for_token[k_idx];
295
+ const int expert = static_cast<int>(
296
+ load_index_as_int64(tid2eid, token_expert_offset + k_idx));
297
const int idx = k * thread_row + k_idx;
298
for (int ii = 0; ii < VPT; ++ii) {
299
const int group_id = ii / ELTS_PER_LDG;
@@ -461,14 +469,15 @@ struct TopkConstants {
469
}
470
471
template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM,
464
- int MAX_BYTES_PER_LDG, typename IndType, typename InputType>
472
+ int MAX_BYTES_PER_LDG, typename IndType, typename HashIndType,
473
+ typename InputType>
474
void topkGatingSoftplusSqrtLauncherHelper(
475
const InputType* input, const bool* finished, float* output,
476
IndType* indices, int* source_row, const int num_rows, const int k,
477
const int start_expert, const int end_expert, const bool renormalize,
478
double routed_scaling_factor, const float* correction_bias,
470
- const bool use_hash, const IndType* input_ids, const IndType* tid2eid,
471
- cudaStream_t stream) {
479
+ const bool use_hash, const HashIndType* input_ids,
480
+ const HashIndType* tid2eid, cudaStream_t stream) {
481
static constexpr int BYTES_PER_LDG =
482
MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
483
using Constants =
@@ -481,7 +490,8 @@ void topkGatingSoftplusSqrtLauncherHelper(
490
DISPATCH_HASH(use_hash, USE_HASH, {
491
auto* kernel =
492
&topkGatingSoftplusSqrt<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG,
484
- WARP_SIZE_PARAM, USE_HASH, IndType, InputType>;
493
+ WARP_SIZE_PARAM, USE_HASH, IndType, HashIndType,
494
+ InputType>;
495
#ifndef USE_ROCM
496
cudaLaunchConfig_t config = {};
497
config.gridDim = num_blocks;
@@ -538,13 +548,14 @@ void topkGatingSoftplusSqrtLauncherHelper(
548
}
549
#endif
550
541
-template <typename IndType, typename InputType>
551
+template <typename IndType, typename InputType, typename HashIndType = IndType>
552
void topkGatingSoftplusSqrtKernelLauncher(
553
const InputType* gating_output, float* topk_weights, IndType* topk_indices,
554
int* token_expert_indices, const int num_tokens, const int num_experts,
555
const int topk, const bool renormalize, double routed_scaling_factor,
546
- const float* correction_bias, const bool use_hash, const IndType* input_ids,
547
- const IndType* tid2eid, cudaStream_t stream) {
556
+ const float* correction_bias, const bool use_hash,
557
+ const HashIndType* input_ids, const HashIndType* tid2eid,
558
+ cudaStream_t stream) {
559
static constexpr int WARPS_PER_TB = 4;
560
static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
561
// for bfloat16 dtype, we need 4 bytes loading to make sure num_experts
@@ -644,57 +655,55 @@ void dispatch_topk_softplus_sqrt_launch(
655
if (correction_bias.has_value()) {
656
bias_ptr = correction_bias.value().const_data_ptr<float>();
657
}
647
- bool use_hash = false;
648
- if (tid2eid.has_value()) {
649
- STD_TORCH_CHECK(input_ids.has_value(),
650
- "input_ids is required for hash MoE");
651
- use_hash = true;
652
- }
653
- if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) {
654
- const int* input_ids_ptr = nullptr;
655
- const int* tid2eid_ptr = nullptr;
658
+
659
+ auto launch = [&](auto* topk_indices_ptr) {
660
+ using OutIndType =
661
+ typename std::remove_pointer<decltype(topk_indices_ptr)>::type;
662
if (tid2eid.has_value()) {
657
- input_ids_ptr = input_ids.value().const_data_ptr<int>();
658
- tid2eid_ptr = tid2eid.value().const_data_ptr<int>();
663
+ STD_TORCH_CHECK(input_ids.has_value(),
664
+ "input_ids is required for hash MoE");
665
+ STD_TORCH_CHECK(
666
+ input_ids.value().scalar_type() == tid2eid.value().scalar_type(),
667
+ "input_ids and tid2eid must have the same dtype");
668
+ if (tid2eid.value().scalar_type() ==
669
+ torch::headeronly::ScalarType::Long) {
670
+ vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType,
671
+ int64_t>(
672
+ gating_output, topk_weights.mutable_data_ptr<float>(),
673
+ topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
674
+ num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
675
+ bias_ptr, true, input_ids.value().const_data_ptr<int64_t>(),
676
+ tid2eid.value().const_data_ptr<int64_t>(), stream);
677
+ } else {
678
+ STD_TORCH_CHECK(tid2eid.value().scalar_type() ==
679
+ torch::headeronly::ScalarType::Int);
680
+ vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType,
681
+ int>(
682
+ gating_output, topk_weights.mutable_data_ptr<float>(),
683
+ topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
684
+ num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
685
+ bias_ptr, true, input_ids.value().const_data_ptr<int>(),
686
+ tid2eid.value().const_data_ptr<int>(), stream);
687
+ }
688
+ } else {
689
+ vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType>(
690
+ gating_output, topk_weights.mutable_data_ptr<float>(),
691
+ topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
692
+ num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
693
+ bias_ptr, false, static_cast<const OutIndType*>(nullptr),
694
+ static_cast<const OutIndType*>(nullptr), stream);
695
}
696
+ };
697
661
- vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int, ComputeType>(
662
- gating_output, topk_weights.mutable_data_ptr<float>(),
663
- topk_indices.mutable_data_ptr<int>(),
664
- token_expert_indices.mutable_data_ptr<int>(), num_tokens, num_experts,
665
- topk, renormalize, routed_scaling_factor, bias_ptr, use_hash,
666
- input_ids_ptr, tid2eid_ptr, stream);
698
+ if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) {
699
+ launch(topk_indices.mutable_data_ptr<int>());
700
} else if (topk_indices.scalar_type() ==
701
torch::headeronly::ScalarType::UInt32) {
669
- const uint32_t* input_ids_ptr = nullptr;
670
- const uint32_t* tid2eid_ptr = nullptr;
671
- if (tid2eid.has_value()) {
672
- input_ids_ptr = input_ids.value().const_data_ptr<uint32_t>();
673
- tid2eid_ptr = tid2eid.value().const_data_ptr<uint32_t>();
674
- }
675
- vllm::moe::topkGatingSoftplusSqrtKernelLauncher<uint32_t, ComputeType>(
676
- gating_output, topk_weights.mutable_data_ptr<float>(),
677
- topk_indices.mutable_data_ptr<uint32_t>(),
678
- token_expert_indices.mutable_data_ptr<int>(), num_tokens, num_experts,
679
- topk, renormalize, routed_scaling_factor, bias_ptr, use_hash,
680
- input_ids_ptr, tid2eid_ptr, stream);
702
+ launch(topk_indices.mutable_data_ptr<uint32_t>());
703
} else {
704
STD_TORCH_CHECK(topk_indices.scalar_type() ==
705
torch::headeronly::ScalarType::Long);
684
-
685
- const int64_t* input_ids_ptr = nullptr;
686
- const int64_t* tid2eid_ptr = nullptr;
687
- if (tid2eid.has_value()) {
688
- input_ids_ptr = input_ids.value().const_data_ptr<int64_t>();
689
- tid2eid_ptr = tid2eid.value().const_data_ptr<int64_t>();
690
- }
691
-
692
- vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int64_t, ComputeType>(
693
- gating_output, topk_weights.mutable_data_ptr<float>(),
694
- topk_indices.mutable_data_ptr<int64_t>(),
695
- token_expert_indices.mutable_data_ptr<int>(), num_tokens, num_experts,
696
- topk, renormalize, routed_scaling_factor, bias_ptr, use_hash,
697
- input_ids_ptr, tid2eid_ptr, stream);
706
+ launch(topk_indices.mutable_data_ptr<int64_t>());
707
}
708
}
709
@@ -738,4 +747,4 @@ void topk_softplus_sqrt(
747
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ",
748
gating_output.scalar_type());
749
}
741
-}
\ No newline at end of file
750
+}
tests/kernels/moe/test_topk_softplus_sqrt.py
+2
-2
@@ -153,9 +153,9 @@ def test_fused_topk_softplus_sqrt_hash(
153
# experts.
154
hash_indices_table = torch.stack(
155
[torch.randperm(num_experts)[:topk] for _ in range(vocab_size)]
156
- ).to(device="cuda", dtype=torch.int32)
156
+ ).to(device="cuda", dtype=torch.long)
157
input_ids = torch.randint(
158
- 0, vocab_size, (num_tokens,), dtype=torch.int32, device="cuda"
158
+ 0, vocab_size, (num_tokens,), dtype=torch.long, device="cuda"
159
)
160
161
topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt(
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
+7
-7
@@ -170,13 +170,12 @@ def fused_topk_bias(
170
hash_indices_table: torch.Tensor | None = None,
171
routed_scaling_factor: float = 1.0,
172
):
173
- # The topk kernel dispatches dtype based on topk_ids (set by
174
- # indices_type) and assumes input_tokens/hash_indices_table match.
175
- if indices_type is not None:
176
- if input_tokens is not None and input_tokens.dtype != indices_type:
177
- input_tokens = input_tokens.to(dtype=indices_type)
178
- if hash_indices_table is not None and hash_indices_table.dtype != indices_type:
179
- hash_indices_table = hash_indices_table.to(dtype=indices_type)
173
+ if (
174
+ input_tokens is not None
175
+ and hash_indices_table is not None
176
+ and input_tokens.dtype != hash_indices_table.dtype
177
+ ):
178
+ input_tokens = input_tokens.to(dtype=hash_indices_table.dtype)
179
180
if not rocm_aiter_ops.is_fused_moe_enabled():
181
assert hidden_states.size(0) == gating_output.size(0), (
@@ -304,6 +303,7 @@ def fused_topk_bias(
303
scores_for_choice = scores.view(-1, n_routed_experts)
304
# For batch invariance, use sorted=True to ensure deterministic expert selection
305
if hash_indices_table is not None:
306
+ assert input_tokens is not None
307
topk_indices = hash_indices_table[input_tokens]
308
else:
309
use_sorted = envs.VLLM_BATCH_INVARIANT