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 10, 14:00 UTC
[Kernel] Support bias type in grouped_topk kernel (#31781)
Signed-off-by: Xin Yang <xyangx@amazon.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Xin Yang committed
Jan 7, 2026 at 12:16 UTC
0ada960a20766fa1de9063576979f9d807583785
3 files changed
+104
-72
csrc/moe/grouped_topk_kernels.cu
+95
-63
@@ -457,8 +457,8 @@ __device__ inline T apply_scoring(T val) {
457
}
458
}
459
460
-template <typename T, ScoringFunc SF>
461
-__device__ void topk_with_k2(T* output, T const* input, T const* bias,
460
+template <typename T, typename BiasT, ScoringFunc SF>
461
+__device__ void topk_with_k2(T* output, T const* input, BiasT const* bias,
462
cg::thread_block_tile<32> const& tile,
463
int32_t const lane_id,
464
int const num_experts_per_group) {
@@ -469,7 +469,7 @@ __device__ void topk_with_k2(T* output, T const* input, T const* bias,
469
if (num_experts_per_group > WARP_SIZE) {
470
for (int i = lane_id; i < num_experts_per_group; i += WARP_SIZE) {
471
T value = apply_scoring<SF>(input[i]);
472
- value = value + bias[i];
472
+ value = value + static_cast<T>(bias[i]);
473
474
if (value > largest) {
475
second_largest = largest;
@@ -481,7 +481,7 @@ __device__ void topk_with_k2(T* output, T const* input, T const* bias,
481
} else {
482
for (int i = lane_id; i < num_experts_per_group; i += WARP_SIZE) {
483
T value = apply_scoring<SF>(input[i]);
484
- value = value + bias[i];
484
+ value = value + static_cast<T>(bias[i]);
485
largest = value;
486
}
487
}
@@ -503,8 +503,8 @@ __device__ void topk_with_k2(T* output, T const* input, T const* bias,
503
}
504
}
505
506
-template <typename T, ScoringFunc SF>
507
-__global__ void topk_with_k2_kernel(T* output, T* input, T const* bias,
506
+template <typename T, typename BiasT, ScoringFunc SF>
507
+__global__ void topk_with_k2_kernel(T* output, T* input, BiasT const* bias,
508
int64_t const num_tokens,
509
int64_t const num_cases,
510
int64_t const n_group,
@@ -517,7 +517,7 @@ __global__ void topk_with_k2_kernel(T* output, T* input, T const* bias,
517
input += case_id * num_experts_per_group;
518
// bias is per expert group, offset to current group
519
int32_t group_id = case_id % n_group;
520
- T const* group_bias = bias + group_id * num_experts_per_group;
520
+ BiasT const* group_bias = bias + group_id * num_experts_per_group;
521
output += case_id;
522
523
cg::thread_block block = cg::this_thread_block();
@@ -526,18 +526,19 @@ __global__ void topk_with_k2_kernel(T* output, T* input, T const* bias,
526
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
527
asm volatile("griddepcontrol.wait;");
528
#endif
529
- topk_with_k2<T, SF>(output, input, group_bias, tile, lane_id,
530
- num_experts_per_group);
529
+ topk_with_k2<T, BiasT, SF>(output, input, group_bias, tile, lane_id,
530
+ num_experts_per_group);
531
}
532
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
533
asm volatile("griddepcontrol.launch_dependents;");
534
#endif
535
}
536
537
-template <typename T, typename IdxT, ScoringFunc SF, int NGroup = -1>
537
+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
538
+ int NGroup = -1>
539
__global__ void group_idx_and_topk_idx_kernel(
540
T* scores, T const* group_scores, float* topk_values, IdxT* topk_indices,
540
- T const* bias, int64_t const num_tokens, int64_t const n_group,
541
+ BiasT const* bias, int64_t const num_tokens, int64_t const n_group,
542
int64_t const topk_group, int64_t const topk, int64_t const num_experts,
543
int64_t const num_experts_per_group, bool renormalize,
544
double routed_scaling_factor) {
@@ -623,7 +624,7 @@ __global__ void group_idx_and_topk_idx_kernel(
624
T input = scores[offset + i];
625
if (is_finite(input)) {
626
T score = apply_scoring<SF>(input);
626
- candidates = score + bias[offset + i];
627
+ candidates = score + static_cast<T>(bias[offset + i]);
628
}
629
}
630
queue.add(candidates, offset + i);
@@ -698,10 +699,10 @@ __global__ void group_idx_and_topk_idx_kernel(
699
#endif
700
}
701
701
-template <typename T, typename IdxT, ScoringFunc SF>
702
+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF>
703
inline void launch_group_idx_and_topk_kernel(
704
cudaLaunchConfig_t const& config, T* scores, T* group_scores,
704
- float* topk_values, IdxT* topk_indices, T const* bias,
705
+ float* topk_values, IdxT* topk_indices, BiasT const* bias,
706
int64_t const num_tokens, int64_t const n_group, int64_t const topk_group,
707
int64_t const topk, int64_t const num_experts,
708
int64_t const num_experts_per_group, bool const renormalize,
@@ -715,36 +716,36 @@ inline void launch_group_idx_and_topk_kernel(
716
717
switch (n_group) {
718
case 4: {
718
- launch(&group_idx_and_topk_idx_kernel<T, IdxT, SF, 4>);
719
+ launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF, 4>);
720
break;
721
}
722
case 8: {
722
- launch(&group_idx_and_topk_idx_kernel<T, IdxT, SF, 8>);
723
+ launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF, 8>);
724
break;
725
}
726
case 16: {
726
- launch(&group_idx_and_topk_idx_kernel<T, IdxT, SF, 16>);
727
+ launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF, 16>);
728
break;
729
}
730
case 32: {
730
- launch(&group_idx_and_topk_idx_kernel<T, IdxT, SF, 32>);
731
+ launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF, 32>);
732
break;
733
}
734
default: {
734
- launch(&group_idx_and_topk_idx_kernel<T, IdxT, SF>);
735
+ launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF>);
736
break;
737
}
738
}
739
}
740
740
-template <typename T, typename IdxT>
741
+template <typename T, typename BiasT, typename IdxT>
742
void invokeNoAuxTc(T* scores, T* group_scores, float* topk_values,
742
- IdxT* topk_indices, T const* bias, int64_t const num_tokens,
743
- int64_t const num_experts, int64_t const n_group,
744
- int64_t const topk_group, int64_t const topk,
745
- bool const renormalize, double const routed_scaling_factor,
746
- int const scoring_func, bool enable_pdl = false,
747
- cudaStream_t const stream = 0) {
743
+ IdxT* topk_indices, BiasT const* bias,
744
+ int64_t const num_tokens, int64_t const num_experts,
745
+ int64_t const n_group, int64_t const topk_group,
746
+ int64_t const topk, bool const renormalize,
747
+ double const routed_scaling_factor, int const scoring_func,
748
+ bool enable_pdl = false, cudaStream_t const stream = 0) {
749
int64_t num_cases = num_tokens * n_group;
750
int64_t topk_with_k2_num_blocks = (num_cases - 1) / NUM_WARPS_PER_BLOCK + 1;
751
cudaLaunchConfig_t config;
@@ -765,12 +766,12 @@ void invokeNoAuxTc(T* scores, T* group_scores, float* topk_values,
766
};
767
switch (sf) {
768
case SCORING_NONE: {
768
- auto* kernel_instance1 = &topk_with_k2_kernel<T, SCORING_NONE>;
769
+ auto* kernel_instance1 = &topk_with_k2_kernel<T, BiasT, SCORING_NONE>;
770
launch_topk_with_k2(kernel_instance1);
771
break;
772
}
773
case SCORING_SIGMOID: {
773
- auto* kernel_instance1 = &topk_with_k2_kernel<T, SCORING_SIGMOID>;
774
+ auto* kernel_instance1 = &topk_with_k2_kernel<T, BiasT, SCORING_SIGMOID>;
775
launch_topk_with_k2(kernel_instance1);
776
break;
777
}
@@ -794,14 +795,14 @@ void invokeNoAuxTc(T* scores, T* group_scores, float* topk_values,
795
config.attrs = attrs;
796
switch (sf) {
797
case SCORING_NONE: {
797
- launch_group_idx_and_topk_kernel<T, IdxT, SCORING_NONE>(
798
+ launch_group_idx_and_topk_kernel<T, BiasT, IdxT, SCORING_NONE>(
799
config, scores, group_scores, topk_values, topk_indices, bias,
800
num_tokens, n_group, topk_group, topk, num_experts,
801
num_experts_per_group, renormalize, routed_scaling_factor);
802
break;
803
}
804
case SCORING_SIGMOID: {
804
- launch_group_idx_and_topk_kernel<T, IdxT, SCORING_SIGMOID>(
805
+ launch_group_idx_and_topk_kernel<T, BiasT, IdxT, SCORING_SIGMOID>(
806
config, scores, group_scores, topk_values, topk_indices, bias,
807
num_tokens, n_group, topk_group, topk, num_experts,
808
num_experts_per_group, renormalize, routed_scaling_factor);
@@ -812,17 +813,23 @@ void invokeNoAuxTc(T* scores, T* group_scores, float* topk_values,
813
}
814
}
815
815
-#define INSTANTIATE_NOAUX_TC(T, IdxT) \
816
- template void invokeNoAuxTc<T, IdxT>( \
817
- T * scores, T * group_scores, float* topk_values, IdxT* topk_indices, \
818
- T const* bias, int64_t const num_tokens, int64_t const num_experts, \
819
- int64_t const n_group, int64_t const topk_group, int64_t const topk, \
820
- bool const renormalize, double const routed_scaling_factor, \
816
+#define INSTANTIATE_NOAUX_TC(T, BiasT, IdxT) \
817
+ template void invokeNoAuxTc<T, BiasT, IdxT>( \
818
+ T * scores, T * group_scores, float* topk_values, IdxT* topk_indices, \
819
+ BiasT const* bias, int64_t const num_tokens, int64_t const num_experts, \
820
+ int64_t const n_group, int64_t const topk_group, int64_t const topk, \
821
+ bool const renormalize, double const routed_scaling_factor, \
822
int const scoring_func, bool enable_pdl, cudaStream_t const stream);
823
823
-INSTANTIATE_NOAUX_TC(float, int32_t);
824
-INSTANTIATE_NOAUX_TC(half, int32_t);
825
-INSTANTIATE_NOAUX_TC(__nv_bfloat16, int32_t);
824
+INSTANTIATE_NOAUX_TC(float, float, int32_t);
825
+INSTANTIATE_NOAUX_TC(float, half, int32_t);
826
+INSTANTIATE_NOAUX_TC(float, __nv_bfloat16, int32_t);
827
+INSTANTIATE_NOAUX_TC(half, float, int32_t);
828
+INSTANTIATE_NOAUX_TC(half, half, int32_t);
829
+INSTANTIATE_NOAUX_TC(half, __nv_bfloat16, int32_t);
830
+INSTANTIATE_NOAUX_TC(__nv_bfloat16, float, int32_t);
831
+INSTANTIATE_NOAUX_TC(__nv_bfloat16, half, int32_t);
832
+INSTANTIATE_NOAUX_TC(__nv_bfloat16, __nv_bfloat16, int32_t);
833
} // end namespace moe
834
} // namespace vllm
835
@@ -831,6 +838,7 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
838
int64_t topk, bool renormalize, double routed_scaling_factor,
839
torch::Tensor const& bias, int64_t scoring_func = 0) {
840
auto data_type = scores.scalar_type();
841
+ auto bias_type = bias.scalar_type();
842
auto input_size = scores.sizes();
843
int64_t num_tokens = input_size[0];
844
int64_t num_experts = input_size[1];
@@ -854,39 +862,62 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
862
863
auto stream = c10::cuda::getCurrentCUDAStream(scores.get_device());
864
865
+#define LAUNCH_KERNEL(T, IdxT) \
866
+ do { \
867
+ switch (bias_type) { \
868
+ case torch::kFloat16: \
869
+ vllm::moe::invokeNoAuxTc<T, half, IdxT>( \
870
+ reinterpret_cast<T*>(scores.mutable_data_ptr()), \
871
+ reinterpret_cast<T*>(group_scores.mutable_data_ptr()), \
872
+ reinterpret_cast<float*>(topk_values.mutable_data_ptr()), \
873
+ reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()), \
874
+ reinterpret_cast<half const*>(bias.data_ptr()), num_tokens, \
875
+ num_experts, n_group, topk_group, topk, renormalize, \
876
+ routed_scaling_factor, static_cast<int>(scoring_func), false, \
877
+ stream); \
878
+ break; \
879
+ case torch::kFloat32: \
880
+ vllm::moe::invokeNoAuxTc<T, float, IdxT>( \
881
+ reinterpret_cast<T*>(scores.mutable_data_ptr()), \
882
+ reinterpret_cast<T*>(group_scores.mutable_data_ptr()), \
883
+ reinterpret_cast<float*>(topk_values.mutable_data_ptr()), \
884
+ reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()), \
885
+ reinterpret_cast<float const*>(bias.data_ptr()), num_tokens, \
886
+ num_experts, n_group, topk_group, topk, renormalize, \
887
+ routed_scaling_factor, static_cast<int>(scoring_func), false, \
888
+ stream); \
889
+ break; \
890
+ case torch::kBFloat16: \
891
+ vllm::moe::invokeNoAuxTc<T, __nv_bfloat16, IdxT>( \
892
+ reinterpret_cast<T*>(scores.mutable_data_ptr()), \
893
+ reinterpret_cast<T*>(group_scores.mutable_data_ptr()), \
894
+ reinterpret_cast<float*>(topk_values.mutable_data_ptr()), \
895
+ reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()), \
896
+ reinterpret_cast<__nv_bfloat16 const*>(bias.data_ptr()), \
897
+ num_tokens, num_experts, n_group, topk_group, topk, renormalize, \
898
+ routed_scaling_factor, static_cast<int>(scoring_func), false, \
899
+ stream); \
900
+ break; \
901
+ default: \
902
+ throw std::invalid_argument( \
903
+ "Invalid bias dtype, only supports float16, float32, and " \
904
+ "bfloat16"); \
905
+ break; \
906
+ } \
907
+ } while (0)
908
+
909
switch (data_type) {
910
case torch::kFloat16:
911
// Handle Float16
860
- vllm::moe::invokeNoAuxTc<half, int32_t>(
861
- reinterpret_cast<half*>(scores.mutable_data_ptr()),
862
- reinterpret_cast<half*>(group_scores.mutable_data_ptr()),
863
- reinterpret_cast<float*>(topk_values.mutable_data_ptr()),
864
- reinterpret_cast<int32_t*>(topk_indices.mutable_data_ptr()),
865
- reinterpret_cast<half const*>(bias.data_ptr()), num_tokens,
866
- num_experts, n_group, topk_group, topk, renormalize,
867
- routed_scaling_factor, static_cast<int>(scoring_func), false, stream);
912
+ LAUNCH_KERNEL(half, int32_t);
913
break;
914
case torch::kFloat32:
915
// Handle Float32
871
- vllm::moe::invokeNoAuxTc<float, int32_t>(
872
- reinterpret_cast<float*>(scores.mutable_data_ptr()),
873
- reinterpret_cast<float*>(group_scores.mutable_data_ptr()),
874
- reinterpret_cast<float*>(topk_values.mutable_data_ptr()),
875
- reinterpret_cast<int32_t*>(topk_indices.mutable_data_ptr()),
876
- reinterpret_cast<float const*>(bias.data_ptr()), num_tokens,
877
- num_experts, n_group, topk_group, topk, renormalize,
878
- routed_scaling_factor, static_cast<int>(scoring_func), false, stream);
916
+ LAUNCH_KERNEL(float, int32_t);
917
break;
918
case torch::kBFloat16:
919
// Handle BFloat16
882
- vllm::moe::invokeNoAuxTc<__nv_bfloat16, int32_t>(
883
- reinterpret_cast<__nv_bfloat16*>(scores.mutable_data_ptr()),
884
- reinterpret_cast<__nv_bfloat16*>(group_scores.mutable_data_ptr()),
885
- reinterpret_cast<float*>(topk_values.mutable_data_ptr()),
886
- reinterpret_cast<int32_t*>(topk_indices.mutable_data_ptr()),
887
- reinterpret_cast<__nv_bfloat16 const*>(bias.data_ptr()), num_tokens,
888
- num_experts, n_group, topk_group, topk, renormalize,
889
- routed_scaling_factor, static_cast<int>(scoring_func), false, stream);
920
+ LAUNCH_KERNEL(__nv_bfloat16, int32_t);
921
break;
922
default:
923
// Handle other data types
@@ -894,5 +925,6 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
925
"Invalid dtype, only supports float16, float32, and bfloat16");
926
break;
927
}
928
+#undef LAUNCH_KERNEL
929
return {topk_values, topk_indices};
930
}
tests/kernels/moe/test_grouped_topk.py
+7
-7
@@ -34,7 +34,8 @@ from vllm.utils.torch_utils import set_random_seed
34
@pytest.mark.parametrize("topk_group", [2])
35
@pytest.mark.parametrize("scoring_func", ["softmax", "sigmoid"])
36
@pytest.mark.parametrize("routed_scaling_factor", [1.0, 2.5])
37
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
37
+@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float32])
38
+@pytest.mark.parametrize("bias_dtype", [torch.float32])
39
def test_grouped_topk(
40
monkeypatch: pytest.MonkeyPatch,
41
n_token: int,
@@ -46,7 +47,8 @@ def test_grouped_topk(
47
topk_group: int,
48
scoring_func: str,
49
routed_scaling_factor: float,
49
- dtype: torch.dtype,
50
+ input_dtype: torch.dtype,
51
+ bias_dtype: torch.dtype,
52
):
53
vllm_config = VllmConfig(
54
compilation_config=CompilationConfig(custom_ops=["all", "+grouped_topk"])
@@ -54,11 +56,9 @@ def test_grouped_topk(
56
get_cached_compilation_config.cache_clear()
57
58
set_random_seed(0)
57
- hidden_states = torch.randn((n_token, n_hidden), dtype=dtype, device="cuda")
58
- gating_output = torch.randn((n_token, n_expert), dtype=dtype, device="cuda")
59
- e_score_correction_bias = torch.randn(
60
- (n_expert,), dtype=torch.float32, device="cuda"
61
- )
59
+ hidden_states = torch.randn((n_token, n_hidden), dtype=input_dtype, device="cuda")
60
+ gating_output = torch.randn((n_token, n_expert), dtype=input_dtype, device="cuda")
61
+ e_score_correction_bias = torch.randn((n_expert,), dtype=bias_dtype, device="cuda")
62
63
with set_current_vllm_config(vllm_config), monkeypatch.context() as m:
64
m.setenv("VLLM_USE_FUSED_MOE_GROUPED_TOPK", "0")
vllm/model_executor/layers/fused_moe/fused_moe.py
+2
-2
@@ -1627,7 +1627,7 @@ def fused_grouped_topk(
1627
topk,
1628
renormalize,
1629
routed_scaling_factor,
1630
- e_score_correction_bias.to(gating_output.dtype),
1630
+ e_score_correction_bias,
1631
1, # scoring_func=1 for sigmoid
1632
)
1633
elif scoring_func == "softmax":
@@ -1641,7 +1641,7 @@ def fused_grouped_topk(
1641
topk,
1642
renormalize,
1643
routed_scaling_factor,
1644
- e_score_correction_bias.to(gating_output.dtype),
1644
+ e_score_correction_bias,
1645
0, # scoring_func=0 (no activation, scores already computed)
1646
)
1647
else: