@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b1c4f0b265
Fork 0

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 8, 01:15 UTC

[Kernel] Optimize grouped topk kernel (#34206)

Signed-off-by: Xin Yang <xyangx@amazon.com>

Xin Yang committed Feb 20, 2026 at 01:34 UTC b1c4f0b26548d36fca304b298957e4791eafa09b
3 files changed +643 -100
csrc/moe/grouped_topk_kernels.cu
+370 -91
@@ -1,6 +1,6 @@
1 /*
2 * Adapted from
3 - * https://github.com/NVIDIA/TensorRT-LLM/blob/v0.21.0/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu
3 + * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu
4 * Copyright (c) 2025, The vLLM team.
5 * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION &
6 * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
@@ -17,8 +17,10 @@
17 * See the License for the specific language governing permissions and
18 * limitations under the License.
19 */
20 +#include "moeTopKFuncs.cuh"
21 #include <c10/cuda/CUDAStream.h>
22 #include <torch/all.h>
23 +#include <cmath>
24 #include <cuda_fp16.h>
25 #include <cuda_bf16.h>
26 #include <cuda/std/limits>
@@ -30,7 +32,17 @@ namespace vllm {
32 namespace moe {
33
34 constexpr unsigned FULL_WARP_MASK = 0xffffffff;
33 -constexpr int32_t WARP_SIZE = 32;
35 +static constexpr int WARP_SIZE = 32;
36 +static constexpr int NumNemotronExperts = 512;
37 +static constexpr int NumKimiK2Experts = 384;
38 +static constexpr int NumDeepseekExperts = 256;
39 +static constexpr int MaxSupportedExpertCount =
40 + std::max({NumNemotronExperts, NumKimiK2Experts, NumDeepseekExperts});
41 +static constexpr int MaxNumExpertsUnit = 128;
42 +static constexpr int NumTopGroupScores = 2;
43 +static constexpr int DefaultMaxNumTopExperts = 8;
44 +static constexpr int MaxSupportedTopExperts = 22;
45 +static constexpr int MaxNumTopGroups = 4;
46
47 namespace warp_topk {
48
@@ -657,76 +669,335 @@ __global__ void grouped_topk_fused_kernel(
669 #endif
670 }
671
660 -template <typename T, typename BiasT, typename IdxT>
672 +template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
673 + int MaxNumExperts, bool UseGroups,
674 + int MaxNumTopExperts = DefaultMaxNumTopExperts>
675 +__global__ void grouped_topk_fused_small_expert_count_kernel(
676 + T* scores, float* topkValues, IdxT* topkIndices, BiasT const* routingBias,
677 + int64_t const numTokens, int64_t const numGroup, int64_t const topkGroup,
678 + int64_t const topk, int64_t const numExperts,
679 + int64_t const numExpertsPerGroup, bool const renormalize,
680 + double const routedScalingFactor) {
681 +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
682 + cudaGridDependencySynchronize();
683 +#endif
684 + // declare shared memory structure
685 + // number of experts is bounded by number of threads
686 + __shared__ float __attribute((aligned(128))) smemScoreSigmoid[MaxNumExperts];
687 + __shared__ float __attribute((aligned(128))) smemScoreBias[MaxNumExperts];
688 + // number of expert groups is bounded by number of warps
689 + int constexpr NumWarps = MaxNumExperts / WARP_SIZE;
690 + __shared__ float __attribute((aligned(128))) smemGroupScores[NumWarps];
691 +
692 + // needed for warp reduce
693 + auto block = cg::this_thread_block();
694 + auto warp = cg::tiled_partition<WARP_SIZE>(block);
695 +
696 + // for the final reduction of weight norm, only some lanes need to participate
697 + int32_t laneIdx = threadIdx.x % WARP_SIZE;
698 + int32_t warpIdx = __shfl_sync(0xffffffff, threadIdx.x / WARP_SIZE, 0);
699 +
700 + if constexpr (UseGroups) {
701 + if (warpIdx >= numGroup) {
702 + return;
703 + }
704 + }
705 + // note that for invalid scores, we simply use a negative value:
706 + // they work well even with the compacted format used in topK, and
707 + // sigmoid / bias activated scores cannot be negative
708 + const float invalidScoreFloat = float{-INFINITY};
709 +
710 + // load bias already; each warp represents one expert group
711 + auto threadExpert = threadIdx.x;
712 + bool expertSelected = threadExpert < numExperts;
713 + if constexpr (UseGroups) {
714 + threadExpert = warpIdx * numExpertsPerGroup + laneIdx;
715 + expertSelected = laneIdx < numExpertsPerGroup;
716 + }
717 +
718 + auto scoreIdx = int64_t{blockIdx.x} * int64_t{numExperts} + threadExpert;
719 + auto biasVal = expertSelected ? static_cast<float>(routingBias[threadExpert])
720 + : invalidScoreFloat;
721 + topkValues += blockIdx.x * topk;
722 + topkIndices += blockIdx.x * topk;
723 +
724 + // get our assigned thread score; each warp represents one expert group
725 + float score =
726 + expertSelected ? static_cast<float>(scores[scoreIdx]) : invalidScoreFloat;
727 + auto scoreSigmoid = apply_scoring<SF>(score);
728 + // write the sigmoid score to shared for later use
729 + if (expertSelected) {
730 + smemScoreSigmoid[threadExpert] = scoreSigmoid;
731 + }
732 +
733 + // get the score with bias
734 + // note that with invalid values, because sigmoid is < 1 and bias is -1,
735 + // we must get a negative value, which is smaller than any valid value
736 + auto scoreBias = float{scoreSigmoid + float{biasVal}};
737 +
738 + if (expertSelected) {
739 + smemScoreBias[threadExpert] = scoreBias;
740 + }
741 +
742 + // registers for top group score reduction
743 + float topExpGroupScores[NumTopGroupScores];
744 + [[maybe_unused]] int32_t topExpGroupIdx[NumTopGroupScores];
745 + float topGroups[MaxNumTopGroups]; // bound of numGroup
746 + int32_t topGroupIdx[MaxNumTopGroups];
747 + float expertScoreGroup[MaxNumTopGroups];
748 + int32_t expertIdxGroup[MaxNumTopGroups];
749 + float topScores[MaxNumTopExperts]; // bound of topk
750 + int32_t topExperts[MaxNumTopExperts];
751 +
752 + if constexpr (UseGroups) {
753 + reduce_topk::reduceTopK(warp, topExpGroupScores, topExpGroupIdx, scoreBias,
754 + threadExpert,
755 + /* minValue */ invalidScoreFloat);
756 +
757 + // get the final group score and write it to shared
758 + if (warp.thread_rank() == 0) {
759 + auto groupScore = topExpGroupScores[0] + topExpGroupScores[1];
760 + smemGroupScores[warpIdx] = groupScore;
761 + }
762 + }
763 +
764 + // make group scores available to all warps
765 + __syncthreads();
766 +
767 + if constexpr (UseGroups) {
768 + if (warpIdx == 0) {
769 + // a single warp performs the selection of top groups, and goes on to
770 + // select the final experts
771 + float groupScore =
772 + laneIdx < numGroup ? smemGroupScores[laneIdx] : invalidScoreFloat;
773 +
774 + reduce_topk::reduceTopK(warp, topGroups, topGroupIdx, groupScore, laneIdx,
775 + /* minValue */ invalidScoreFloat);
776 + // final expert selection: get relevant indexes and scores from shared
777 +#pragma unroll
778 + for (int ii = 0; ii < MaxNumTopGroups; ++ii) { // bound of numGroup
779 + auto groupIdx = topGroupIdx[ii];
780 + expertIdxGroup[ii] = groupIdx * numExpertsPerGroup + laneIdx;
781 +
782 + expertScoreGroup[ii] = (ii < topkGroup) && expertSelected
783 + ? smemScoreBias[expertIdxGroup[ii]]
784 + : invalidScoreFloat;
785 + }
786 +
787 + reduce_topk::reduceTopK(warp, topScores, topExperts, expertScoreGroup,
788 + expertIdxGroup, /* minValue */ invalidScoreFloat,
789 + topk);
790 + }
791 + } else if constexpr (MaxNumExperts > MaxNumExpertsUnit) {
792 + // without groups, and the expert number is larger than MaxNumExpertsUnit,
793 + // we need to use multiple warps to calculate the intermediate topk results
794 +
795 + int constexpr NumExpertWarps = (MaxNumExperts - 1) / MaxNumExpertsUnit + 1;
796 + int constexpr NumInterTopK = NumExpertWarps * MaxNumTopExperts;
797 + __shared__ float
798 + __attribute((aligned(128))) smemInterTopScores[NumInterTopK];
799 + __shared__ int32_t
800 + __attribute((aligned(128))) smemInterTopExperts[NumInterTopK];
801 + if (warpIdx < NumExpertWarps) {
802 + int offset = warpIdx * WARP_SIZE * MaxNumTopGroups;
803 +#pragma unroll
804 + for (int ii = 0; ii < MaxNumTopGroups; ++ii) {
805 + auto expertIdx = ii * WARP_SIZE + laneIdx;
806 + expertIdxGroup[ii] = offset + expertIdx;
807 + expertScoreGroup[ii] = offset + expertIdx < numExperts
808 + ? smemScoreBias[offset + expertIdx]
809 + : invalidScoreFloat;
810 + }
811 + reduce_topk::reduceTopK(warp, topScores, topExperts, expertScoreGroup,
812 + expertIdxGroup,
813 + /* minValue */ invalidScoreFloat, topk);
814 +
815 + if (laneIdx < topk) {
816 + smemInterTopScores[warpIdx * MaxNumTopExperts + laneIdx] =
817 + topScores[laneIdx];
818 + smemInterTopExperts[warpIdx * MaxNumTopExperts + laneIdx] =
819 + topExperts[laneIdx];
820 + } else if (laneIdx >= topk && laneIdx < MaxNumTopExperts) {
821 + smemInterTopScores[warpIdx * MaxNumTopExperts + laneIdx] =
822 + invalidScoreFloat;
823 + smemInterTopExperts[warpIdx * MaxNumTopExperts + laneIdx] =
824 + MaxNumExperts - 1;
825 + }
826 + }
827 + __syncthreads();
828 + if (warpIdx == 0) {
829 + int constexpr NumInterTopKPerThread = (NumInterTopK - 1) / WARP_SIZE + 1;
830 + float intermediateScore[NumInterTopKPerThread];
831 + int32_t intermediateExpert[NumInterTopKPerThread];
832 + for (int i = laneIdx; i < NumInterTopKPerThread * WARP_SIZE;
833 + i += WARP_SIZE) {
834 + int ii = i / WARP_SIZE;
835 + if (i < NumInterTopK) {
836 + intermediateScore[ii] = smemInterTopScores[i];
837 + intermediateExpert[ii] = smemInterTopExperts[i];
838 + } else {
839 + intermediateScore[ii] = invalidScoreFloat;
840 + intermediateExpert[ii] = MaxNumExperts - 1;
841 + }
842 + }
843 + reduce_topk::reduceTopK(warp, topScores, topExperts, intermediateScore,
844 + intermediateExpert,
845 + /* minValue */ invalidScoreFloat, topk);
846 + }
847 + } else {
848 + // without groups, and the expert number is smaller than MaxNumExpertsUnit
849 + // each thread just takes `MaxNumTopGroups` experts
850 + if (warpIdx == 0) {
851 +#pragma unroll
852 + for (int ii = 0; ii < MaxNumTopGroups; ++ii) {
853 + auto expertIdx = ii * WARP_SIZE + laneIdx;
854 + expertIdxGroup[ii] = expertIdx;
855 + expertScoreGroup[ii] = expertIdx < numExperts ? smemScoreBias[expertIdx]
856 + : invalidScoreFloat;
857 + }
858 + reduce_topk::reduceTopK(warp, topScores, topExperts, expertScoreGroup,
859 + expertIdxGroup,
860 + /* minValue */ invalidScoreFloat, topk);
861 + }
862 + }
863 +
864 + if (warpIdx == 0) {
865 + // determine our lane's expert index and write to output
866 + int32_t expertIdx =
867 + laneIdx < topk ? topExperts[laneIdx] : MaxNumExperts - 1;
868 + float scoreNorm = laneIdx < topk ? smemScoreSigmoid[expertIdx] : 0.F;
869 + float finalScore = static_cast<float>(scoreNorm * routedScalingFactor);
870 + // norm the value
871 + if (renormalize) {
872 + auto redNorm = cg::reduce(warp, scoreNorm, cg::plus<float>{});
873 + finalScore /= (redNorm + 1e-20);
874 + }
875 + // store the topk scores and experts to output
876 + if (laneIdx < topk) {
877 + topkValues[laneIdx] = finalScore;
878 + topkIndices[laneIdx] = expertIdx;
879 + }
880 + }
881 +
882 +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
883 + cudaTriggerProgrammaticLaunchCompletion();
884 +#endif
885 +}
886 +
887 +template <typename T, typename BiasT, typename IdxT, ScoringFunc SF>
888 void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,
889 BiasT const* bias, int64_t const num_tokens,
890 int64_t const num_experts, int64_t const n_group,
891 int64_t const topk_group, int64_t const topk,
892 bool const renormalize, double const routed_scaling_factor,
666 - int const scoring_func, bool enable_pdl = false,
667 - cudaStream_t const stream = 0) {
893 + bool enable_pdl = false, cudaStream_t const stream = 0) {
894 cudaLaunchConfig_t config;
669 - // One block per token; one warp per group.
670 - config.gridDim = static_cast<uint32_t>(num_tokens);
671 - config.blockDim = static_cast<uint32_t>(n_group) * WARP_SIZE;
672 - // Dynamic shared memory: WarpSelect staging + per-group topk buffers.
673 - int32_t const num_warps = static_cast<int32_t>(n_group);
674 - size_t const val_bytes =
675 - static_cast<size_t>(num_warps) * WARP_SIZE * sizeof(T);
676 - size_t const val_bytes_aligned =
677 - warp_topk::round_up_to_multiple_of<256>(val_bytes);
678 - size_t const idx_bytes =
679 - static_cast<size_t>(num_warps) * WARP_SIZE * sizeof(int32_t);
680 - size_t const internal_bytes = val_bytes_aligned + idx_bytes;
681 - size_t const extra_bytes = 16 + static_cast<size_t>(n_group) * sizeof(T);
682 - config.dynamicSmemBytes = internal_bytes + extra_bytes;
895 config.stream = stream;
896 cudaLaunchAttribute attrs[1];
897 attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
898 attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl;
899 config.numAttrs = 1;
900 config.attrs = attrs;
689 - auto const sf = static_cast<ScoringFunc>(scoring_func);
690 - switch (sf) {
691 - case SCORING_NONE: {
692 - auto* kernel_instance =
693 - &grouped_topk_fused_kernel<T, BiasT, IdxT, SCORING_NONE>;
694 - cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values,
695 - topk_indices, bias, num_tokens, num_experts, n_group,
696 - topk_group, topk, renormalize, routed_scaling_factor);
697 - return;
698 - }
699 - case SCORING_SIGMOID: {
700 - auto* kernel_instance =
701 - &grouped_topk_fused_kernel<T, BiasT, IdxT, SCORING_SIGMOID>;
702 - cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values,
703 - topk_indices, bias, num_tokens, num_experts, n_group,
704 - topk_group, topk, renormalize, routed_scaling_factor);
705 - return;
901 +
902 + // Check if we can use the optimized
903 + // grouped_topk_fused_small_expert_count_kernel
904 + bool const is_single_group =
905 + (n_group == 1) && (topk_group == 1) &&
906 + (num_experts <= MaxSupportedExpertCount) &&
907 + (topk <= DefaultMaxNumTopExperts || topk == MaxSupportedTopExperts);
908 +
909 + int64_t const experts_per_group = num_experts / n_group;
910 + bool const is_multi_group =
911 + (n_group > 1) && (num_experts <= NumDeepseekExperts) &&
912 + (experts_per_group <= WARP_SIZE) &&
913 + (experts_per_group * topk_group <= MaxNumExpertsUnit) &&
914 + (topk <= DefaultMaxNumTopExperts) && (topk_group <= MaxNumTopGroups);
915 +
916 + if (is_single_group || is_multi_group) {
917 + auto* kernel_instance =
918 + &grouped_topk_fused_small_expert_count_kernel<T, BiasT, IdxT, SF,
919 + NumDeepseekExperts, true>;
920 + int num_threads = NumDeepseekExperts;
921 + if (is_single_group) {
922 + // Special case for Nemotron, which selects top 22 from 512 experts, and 1
923 + // group only.
924 + if (num_experts == NumNemotronExperts && n_group == 1 &&
925 + topk == MaxSupportedTopExperts) {
926 + kernel_instance = &grouped_topk_fused_small_expert_count_kernel<
927 + T, BiasT, IdxT, SF, NumNemotronExperts, false,
928 + MaxSupportedTopExperts>;
929 + num_threads = NumNemotronExperts;
930 + } else if (num_experts > NumKimiK2Experts &&
931 + num_experts <= MaxSupportedExpertCount) {
932 + kernel_instance = &grouped_topk_fused_small_expert_count_kernel<
933 + T, BiasT, IdxT, SF, MaxSupportedExpertCount, false>;
934 + num_threads = MaxSupportedExpertCount;
935 + } else if (num_experts > MaxNumExpertsUnit &&
936 + num_experts <= NumKimiK2Experts) {
937 + kernel_instance = &grouped_topk_fused_small_expert_count_kernel<
938 + T, BiasT, IdxT, SF, NumKimiK2Experts, false>;
939 + num_threads = NumKimiK2Experts;
940 + } else {
941 + kernel_instance = &grouped_topk_fused_small_expert_count_kernel<
942 + T, BiasT, IdxT, SF, MaxNumExpertsUnit, false>;
943 + num_threads = MaxNumExpertsUnit;
944 + }
945 }
707 - default:
708 - // should be guarded by higher level checks.
709 - TORCH_CHECK(false, "Unsupported scoring_func in invokeNoAuxTc");
946 + config.gridDim = num_tokens;
947 + config.blockDim = num_threads;
948 + config.dynamicSmemBytes = 0;
949 + cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values,
950 + topk_indices, bias, num_tokens, n_group, topk_group,
951 + topk, num_experts, num_experts / n_group, renormalize,
952 + routed_scaling_factor);
953 + } else {
954 + auto* kernel_instance = &grouped_topk_fused_kernel<T, BiasT, IdxT, SF>;
955 + // One block per token; one warp per group.
956 + config.gridDim = static_cast<uint32_t>(num_tokens);
957 + config.blockDim = static_cast<uint32_t>(n_group) * WARP_SIZE;
958 + // Dynamic shared memory: WarpSelect staging + per-group topk buffers.
959 + int32_t const num_warps = static_cast<int32_t>(n_group);
960 + size_t const val_bytes =
961 + static_cast<size_t>(num_warps) * WARP_SIZE * sizeof(T);
962 + size_t const val_bytes_aligned =
963 + warp_topk::round_up_to_multiple_of<256>(val_bytes);
964 + size_t const idx_bytes =
965 + static_cast<size_t>(num_warps) * WARP_SIZE * sizeof(int32_t);
966 + size_t const internal_bytes = val_bytes_aligned + idx_bytes;
967 + size_t const extra_bytes = 16 + static_cast<size_t>(n_group) * sizeof(T);
968 + config.dynamicSmemBytes = internal_bytes + extra_bytes;
969 + cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values,
970 + topk_indices, bias, num_tokens, num_experts, n_group,
971 + topk_group, topk, renormalize, routed_scaling_factor);
972 }
973 }
974
713 -#define INSTANTIATE_NOAUX_TC(T, BiasT, IdxT) \
714 - template void invokeNoAuxTc<T, BiasT, IdxT>( \
975 +#define INSTANTIATE_NOAUX_TC(T, BiasT, IdxT, SF) \
976 + template void invokeNoAuxTc<T, BiasT, IdxT, SF>( \
977 T * scores, float* topk_values, IdxT* topk_indices, BiasT const* bias, \
978 int64_t const num_tokens, int64_t const num_experts, \
979 int64_t const n_group, int64_t const topk_group, int64_t const topk, \
980 bool const renormalize, double const routed_scaling_factor, \
719 - int const scoring_func, bool enable_pdl, cudaStream_t const stream);
720 -
721 -INSTANTIATE_NOAUX_TC(float, float, int32_t);
722 -INSTANTIATE_NOAUX_TC(float, half, int32_t);
723 -INSTANTIATE_NOAUX_TC(float, __nv_bfloat16, int32_t);
724 -INSTANTIATE_NOAUX_TC(half, float, int32_t);
725 -INSTANTIATE_NOAUX_TC(half, half, int32_t);
726 -INSTANTIATE_NOAUX_TC(half, __nv_bfloat16, int32_t);
727 -INSTANTIATE_NOAUX_TC(__nv_bfloat16, float, int32_t);
728 -INSTANTIATE_NOAUX_TC(__nv_bfloat16, half, int32_t);
729 -INSTANTIATE_NOAUX_TC(__nv_bfloat16, __nv_bfloat16, int32_t);
981 + bool enable_pdl, cudaStream_t const stream);
982 +
983 +INSTANTIATE_NOAUX_TC(float, float, int32_t, SCORING_SIGMOID);
984 +INSTANTIATE_NOAUX_TC(float, half, int32_t, SCORING_SIGMOID);
985 +INSTANTIATE_NOAUX_TC(float, __nv_bfloat16, int32_t, SCORING_SIGMOID);
986 +INSTANTIATE_NOAUX_TC(half, float, int32_t, SCORING_SIGMOID);
987 +INSTANTIATE_NOAUX_TC(half, half, int32_t, SCORING_SIGMOID);
988 +INSTANTIATE_NOAUX_TC(half, __nv_bfloat16, int32_t, SCORING_SIGMOID);
989 +INSTANTIATE_NOAUX_TC(__nv_bfloat16, float, int32_t, SCORING_SIGMOID);
990 +INSTANTIATE_NOAUX_TC(__nv_bfloat16, half, int32_t, SCORING_SIGMOID);
991 +INSTANTIATE_NOAUX_TC(__nv_bfloat16, __nv_bfloat16, int32_t, SCORING_SIGMOID);
992 +INSTANTIATE_NOAUX_TC(float, float, int32_t, SCORING_NONE);
993 +INSTANTIATE_NOAUX_TC(float, half, int32_t, SCORING_NONE);
994 +INSTANTIATE_NOAUX_TC(float, __nv_bfloat16, int32_t, SCORING_NONE);
995 +INSTANTIATE_NOAUX_TC(half, float, int32_t, SCORING_NONE);
996 +INSTANTIATE_NOAUX_TC(half, half, int32_t, SCORING_NONE);
997 +INSTANTIATE_NOAUX_TC(half, __nv_bfloat16, int32_t, SCORING_NONE);
998 +INSTANTIATE_NOAUX_TC(__nv_bfloat16, float, int32_t, SCORING_NONE);
999 +INSTANTIATE_NOAUX_TC(__nv_bfloat16, half, int32_t, SCORING_NONE);
1000 +INSTANTIATE_NOAUX_TC(__nv_bfloat16, __nv_bfloat16, int32_t, SCORING_NONE);
1001 } // end namespace moe
1002 } // namespace vllm
1003
@@ -762,46 +1033,53 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
1033 {num_tokens, topk}, torch::dtype(torch::kInt32).device(torch::kCUDA));
1034
1035 auto stream = c10::cuda::getCurrentCUDAStream(scores.get_device());
1036 + auto const sf = static_cast<vllm::moe::ScoringFunc>(scoring_func);
1037 +
1038 +#define LAUNCH_KERNEL_SF(T, BiasT, IdxT) \
1039 + do { \
1040 + switch (sf) { \
1041 + case vllm::moe::SCORING_NONE: \
1042 + vllm::moe::invokeNoAuxTc<T, BiasT, IdxT, vllm::moe::SCORING_NONE>( \
1043 + reinterpret_cast<T*>(scores.mutable_data_ptr()), \
1044 + reinterpret_cast<float*>(topk_values.mutable_data_ptr()), \
1045 + reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()), \
1046 + reinterpret_cast<BiasT const*>(bias.data_ptr()), num_tokens, \
1047 + num_experts, n_group, topk_group, topk, renormalize, \
1048 + routed_scaling_factor, false, stream); \
1049 + break; \
1050 + case vllm::moe::SCORING_SIGMOID: \
1051 + vllm::moe::invokeNoAuxTc<T, BiasT, IdxT, vllm::moe::SCORING_SIGMOID>( \
1052 + reinterpret_cast<T*>(scores.mutable_data_ptr()), \
1053 + reinterpret_cast<float*>(topk_values.mutable_data_ptr()), \
1054 + reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()), \
1055 + reinterpret_cast<BiasT const*>(bias.data_ptr()), num_tokens, \
1056 + num_experts, n_group, topk_group, topk, renormalize, \
1057 + routed_scaling_factor, false, stream); \
1058 + break; \
1059 + default: \
1060 + throw std::invalid_argument("Unsupported scoring_func"); \
1061 + break; \
1062 + } \
1063 + } while (0)
1064
766 -#define LAUNCH_KERNEL(T, IdxT) \
767 - do { \
768 - switch (bias_type) { \
769 - case torch::kFloat16: \
770 - vllm::moe::invokeNoAuxTc<T, half, IdxT>( \
771 - reinterpret_cast<T*>(scores.mutable_data_ptr()), \
772 - reinterpret_cast<float*>(topk_values.mutable_data_ptr()), \
773 - reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()), \
774 - reinterpret_cast<half const*>(bias.data_ptr()), num_tokens, \
775 - num_experts, n_group, topk_group, topk, renormalize, \
776 - routed_scaling_factor, static_cast<int>(scoring_func), false, \
777 - stream); \
778 - break; \
779 - case torch::kFloat32: \
780 - vllm::moe::invokeNoAuxTc<T, float, IdxT>( \
781 - reinterpret_cast<T*>(scores.mutable_data_ptr()), \
782 - reinterpret_cast<float*>(topk_values.mutable_data_ptr()), \
783 - reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()), \
784 - reinterpret_cast<float const*>(bias.data_ptr()), num_tokens, \
785 - num_experts, n_group, topk_group, topk, renormalize, \
786 - routed_scaling_factor, static_cast<int>(scoring_func), false, \
787 - stream); \
788 - break; \
789 - case torch::kBFloat16: \
790 - vllm::moe::invokeNoAuxTc<T, __nv_bfloat16, IdxT>( \
791 - reinterpret_cast<T*>(scores.mutable_data_ptr()), \
792 - reinterpret_cast<float*>(topk_values.mutable_data_ptr()), \
793 - reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()), \
794 - reinterpret_cast<__nv_bfloat16 const*>(bias.data_ptr()), \
795 - num_tokens, num_experts, n_group, topk_group, topk, renormalize, \
796 - routed_scaling_factor, static_cast<int>(scoring_func), false, \
797 - stream); \
798 - break; \
799 - default: \
800 - throw std::invalid_argument( \
801 - "Invalid bias dtype, only supports float16, float32, and " \
802 - "bfloat16"); \
803 - break; \
804 - } \
1065 +#define LAUNCH_KERNEL(T, IdxT) \
1066 + do { \
1067 + switch (bias_type) { \
1068 + case torch::kFloat16: \
1069 + LAUNCH_KERNEL_SF(T, half, IdxT); \
1070 + break; \
1071 + case torch::kFloat32: \
1072 + LAUNCH_KERNEL_SF(T, float, IdxT); \
1073 + break; \
1074 + case torch::kBFloat16: \
1075 + LAUNCH_KERNEL_SF(T, __nv_bfloat16, IdxT); \
1076 + break; \
1077 + default: \
1078 + throw std::invalid_argument( \
1079 + "Invalid bias dtype, only supports float16, float32, and " \
1080 + "bfloat16"); \
1081 + break; \
1082 + } \
1083 } while (0)
1084
1085 switch (data_type) {
@@ -824,5 +1102,6 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
1102 break;
1103 }
1104 #undef LAUNCH_KERNEL
1105 +#undef LAUNCH_KERNEL_SF
1106 return {topk_values, topk_indices};
1107 }
csrc/moe/moeTopKFuncs.cuh new
+257
@@ -0,0 +1,257 @@
1 +/*
2 + * Adapted from
3 + * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh
4 + * Copyright (c) 2026, The vLLM team.
5 + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION. All rights
6 + * reserved. SPDX-License-Identifier: Apache-2.0
7 + *
8 + * Licensed under the Apache License, Version 2.0 (the "License");
9 + * you may not use this file except in compliance with the License.
10 + * You may obtain a copy of the License at
11 + *
12 + * http://www.apache.org/licenses/LICENSE-2.0
13 + *
14 + * Unless required by applicable law or agreed to in writing, software
15 + * distributed under the License is distributed on an "AS IS" BASIS,
16 + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17 + * See the License for the specific language governing permissions and
18 + * limitations under the License.
19 + */
20 +#pragma once
21 +
22 +#include <cooperative_groups.h>
23 +#include <cooperative_groups/reduce.h>
24 +#include <cub/cub.cuh>
25 +
26 +namespace vllm {
27 +namespace moe {
28 +namespace reduce_topk {
29 +namespace cg = cooperative_groups;
30 +static constexpr int kWARP_SIZE = 32;
31 +
32 +template <typename T_>
33 +struct TopKRedType {
34 + using T = T_;
35 + static_assert(
36 + std::is_same_v<T, float> || std::is_same_v<T, half> ||
37 + std::is_same_v<T, __nv_bfloat16> || std::is_same_v<T, int>,
38 + "Top K reduction only implemented for int, float, float16 and bfloat16");
39 +
40 + using TypeCmp = std::conditional_t<sizeof(T) == 4, uint64_t, uint32_t>;
41 + using IdxT = std::conditional_t<sizeof(T) == 4, int32_t, int16_t>;
42 +
43 + static constexpr int kMoveBits = (sizeof(T) == 4) ? 32 : 16;
44 + static constexpr int kMaxIdx = 65535;
45 + TypeCmp compValIdx;
46 +
47 + static __host__ __device__ inline TypeCmp makeCmpVal(T val, int32_t idx = 0) {
48 + auto valueBits = cub::Traits<T>::TwiddleIn(
49 + reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(val));
50 + TypeCmp compactTmp = valueBits;
51 + compactTmp = (compactTmp << kMoveBits) | (0xFFFF & (kMaxIdx - idx));
52 + // Use 65535 minus idx to give higher priority to elements with smaller
53 + // indices.
54 + return compactTmp;
55 + }
56 +
57 + static __host__ __device__ void unpack(T& value, int32_t& index,
58 + TypeCmp cmp) {
59 + // Since “65535-idx” is always smaller than 65536 and positive, we can
60 + // directly use it as the lower 16 bits
61 + index = kMaxIdx - static_cast<int32_t>((cmp & 0xFFFF));
62 +
63 + auto compactTmp = cmp >> kMoveBits;
64 + auto valueBits = cub::Traits<T>::TwiddleOut(
65 + reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(compactTmp));
66 + value = reinterpret_cast<T&>(valueBits);
67 + }
68 +
69 + __host__ __device__ TopKRedType() = default;
70 +
71 + __host__ __device__ TopKRedType(T val, int32_t idx)
72 + : compValIdx(makeCmpVal(val, idx)) {}
73 +
74 + __host__ __device__ operator TypeCmp() const noexcept { return compValIdx; }
75 +
76 + __device__ inline TypeCmp reduce(
77 + cg::thread_block_tile<kWARP_SIZE> const& warp) {
78 + return cg::reduce(warp, compValIdx, cg::greater<TypeCmp>{});
79 + }
80 +};
81 +
82 +////////////////////////////////////////////////////////////////////////////////////////////////////
83 +
84 +template <int K_, bool Enable_>
85 +struct TopKIdx {
86 + // by default, empty
87 +};
88 +
89 +template <int K_>
90 +struct TopKIdx<K_, true> {
91 + static constexpr int K = K_;
92 + int32_t val[K];
93 +};
94 +
95 +////////////////////////////////////////////////////////////////////////////////////////////////////
96 +
97 +#define TOPK_SWAP(I, J) \
98 + { \
99 + auto pairMin = min(topK[I].compValIdx, topK[J].compValIdx); \
100 + auto pairMax = max(topK[I].compValIdx, topK[J].compValIdx); \
101 + topK[I].compValIdx = pairMax; \
102 + topK[J].compValIdx = pairMin; \
103 + }
104 +
105 +template <int N, typename RedType>
106 +struct Sort;
107 +
108 +template <typename RedType>
109 +struct Sort<1, RedType> {
110 + static __device__ void run(RedType* topK) {}
111 +};
112 +
113 +template <typename RedType>
114 +struct Sort<2, RedType> {
115 + static __device__ void run(RedType* topK) { TOPK_SWAP(0, 1); }
116 +};
117 +
118 +template <typename RedType>
119 +struct Sort<3, RedType> {
120 + static __device__ void run(RedType* topK) {
121 + TOPK_SWAP(0, 1);
122 + TOPK_SWAP(1, 2);
123 + TOPK_SWAP(0, 1);
124 + }
125 +};
126 +
127 +template <typename RedType>
128 +struct Sort<4, RedType> {
129 + static __device__ void run(RedType* topK) {
130 + TOPK_SWAP(0, 2);
131 + TOPK_SWAP(1, 3);
132 + TOPK_SWAP(0, 1);
133 + TOPK_SWAP(2, 3);
134 + TOPK_SWAP(1, 2);
135 + }
136 +};
137 +
138 +template <int K, typename Type>
139 +__forceinline__ __device__ void reduceTopK(
140 + cg::thread_block_tile<kWARP_SIZE> const& warp, Type (&out)[K],
141 + int32_t (&outIdx)[K], Type value, int32_t idx, Type const minValue,
142 + int actualK = K) {
143 + static_assert(K > 0, "Top K must have K > 0");
144 + static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
145 + using RedType = TopKRedType<Type>;
146 + RedType topK{value, idx};
147 + typename RedType::TypeCmp packedMax{};
148 +#pragma unroll
149 + for (int kk = 0; kk < actualK; ++kk) {
150 + topK =
151 + kk > 0 && packedMax == topK.compValIdx ? RedType{minValue, idx} : topK;
152 + // get the next largest value
153 + packedMax = topK.reduce(warp);
154 + RedType::unpack(out[kk], outIdx[kk], packedMax);
155 + }
156 +};
157 +
158 +template <int K, typename Type, int N, bool IsSorted = false>
159 +__device__ void reduceTopKFunc(cg::thread_block_tile<kWARP_SIZE> const& warp,
160 + Type (&out)[K], int32_t (&outIdx)[K],
161 + Type (&value)[N], int32_t (&idx)[N],
162 + Type minValue, int actualK = K) {
163 + static_assert(K > 0, "Top K must have K > 0");
164 + static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
165 + static_assert(N > 0, "Top K must have N > 0");
166 + static_assert(N < 5,
167 + "Only support candidates number less than or equal to 128");
168 + using RedType = TopKRedType<Type>;
169 + RedType topK[N];
170 +#pragma unroll
171 + for (int nn = 0; nn < N; ++nn) {
172 + topK[nn] = RedType{value[nn], idx[nn]};
173 + }
174 +
175 + if constexpr (!IsSorted) {
176 + Sort<N, RedType>::run(topK);
177 + }
178 + typename RedType::TypeCmp packedMax{};
179 +#pragma unroll
180 + for (int kk = 0; kk < actualK; ++kk) {
181 + bool update = kk > 0 && packedMax == topK[0].compValIdx;
182 +#pragma unroll
183 + for (int nn = 0; nn < N; ++nn) {
184 + topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]}
185 + : update ? topK[nn + 1]
186 + : topK[nn];
187 + }
188 + // get the next largest value
189 + packedMax = topK[0].reduce(warp);
190 + RedType::unpack(out[kk], outIdx[kk], packedMax);
191 + }
192 +};
193 +
194 +template <int K, typename Type, int N>
195 +__forceinline__ __device__ void reduceTopK(
196 + cg::thread_block_tile<kWARP_SIZE> const& warp, Type (&out)[K],
197 + int32_t (&outIdx)[K], Type (&value)[N], int32_t (&idx)[N],
198 + Type const minValue, int actualK = K) {
199 + static_assert(K > 0, "Top K must have K > 0");
200 + static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
201 + static_assert(N > 0, "Top K must have N > 0");
202 + static_assert(
203 + N <= 16,
204 + "Only support candidates number less than or equal to 16*32=512");
205 + static_assert(N <= 4 || N % 4 == 0,
206 + "Only support candidates number is a multiple of 4*32=128 or "
207 + "less than or equal to 4");
208 + using RedType = TopKRedType<Type>;
209 +
210 + if constexpr (N <= 4) {
211 + reduceTopKFunc<K, Type, N>(warp, out, outIdx, value, idx, minValue,
212 + actualK);
213 + } else {
214 + constexpr int numLoops = N / 4;
215 + constexpr int numResults = (numLoops * K - 1) / kWARP_SIZE + 1;
216 +
217 + Type topKBufferValue[numResults];
218 + int32_t topKBufferIdx[numResults];
219 + int32_t laneIdx = threadIdx.x % kWARP_SIZE;
220 +
221 + for (int ii = 0; ii < numResults; ++ii) {
222 + topKBufferValue[ii] = minValue;
223 + topKBufferIdx[ii] = ii * kWARP_SIZE - 1;
224 + }
225 + for (int loop = 0; loop < numLoops; ++loop) {
226 + int start = loop * 4;
227 + Type topKValue[K];
228 + int32_t topKIdx[K];
229 + Type inValue[4];
230 + int32_t inIdx[4];
231 + for (int i = 0; i < 4; ++i) {
232 + inValue[i] = value[start + i];
233 + inIdx[i] = idx[start + i];
234 + }
235 + reduceTopKFunc<K, Type, 4>(warp, topKValue, topKIdx, inValue, inIdx,
236 + minValue, actualK);
237 + int inOffset = laneIdx % K;
238 + if (laneIdx >= loop * K && laneIdx < (loop + 1) * K) {
239 + topKBufferValue[0] = topKValue[inOffset];
240 + topKBufferIdx[0] = topKIdx[inOffset];
241 + }
242 + if (loop == numLoops - 1 && (laneIdx < (numLoops * K - kWARP_SIZE))) {
243 + topKBufferValue[1] = topKValue[inOffset];
244 + topKBufferIdx[1] = topKIdx[inOffset];
245 + }
246 + }
247 +
248 + reduceTopKFunc<K, Type, numResults>(warp, out, outIdx, topKBufferValue,
249 + topKBufferIdx, minValue, actualK);
250 + }
251 +};
252 +
253 +#undef TOPK_SWAP
254 +
255 +} // namespace reduce_topk
256 +} // namespace moe
257 +} // namespace vllm
tests/kernels/moe/test_grouped_topk.py
+16 -9
@@ -8,6 +8,7 @@ Run `pytest tests/kernels/moe/test_grouped_topk.py`.
8 import pytest
9 import torch
10
11 +import vllm.model_executor.layers.batch_invariant as batch_invariant
12 from vllm.config import (
13 CompilationConfig,
14 VllmConfig,
@@ -27,11 +28,17 @@ from vllm.utils.torch_utils import set_random_seed
28 )
29 @pytest.mark.parametrize("n_token", [1, 33, 64])
30 @pytest.mark.parametrize("n_hidden", [1024, 2048])
30 -@pytest.mark.parametrize("n_expert", [16])
31 -@pytest.mark.parametrize("topk", [2])
31 +@pytest.mark.parametrize(
32 + "n_expert,topk,num_expert_group,topk_group",
33 + [
34 + (16, 2, 8, 2),
35 + (128, 2, 8, 2),
36 + (256, 8, 8, 4),
37 + (384, 8, 1, 1),
38 + (512, 22, 1, 1),
39 + ],
40 +)
41 @pytest.mark.parametrize("renormalize", [True, False])
33 -@pytest.mark.parametrize("num_expert_group", [8])
34 -@pytest.mark.parametrize("topk_group", [2])
42 @pytest.mark.parametrize("scoring_func", ["softmax", "sigmoid"])
43 @pytest.mark.parametrize("routed_scaling_factor", [1.0, 2.5])
44 @pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float32])
@@ -42,9 +49,9 @@ def test_grouped_topk(
49 n_hidden: int,
50 n_expert: int,
51 topk: int,
45 - renormalize: bool,
52 num_expert_group: int,
53 topk_group: int,
54 + renormalize: bool,
55 scoring_func: str,
56 routed_scaling_factor: float,
57 input_dtype: torch.dtype,
@@ -62,6 +69,7 @@ def test_grouped_topk(
69
70 with set_current_vllm_config(vllm_config), monkeypatch.context() as m:
71 m.setenv("VLLM_USE_FUSED_MOE_GROUPED_TOPK", "0")
72 + m.setattr(batch_invariant, "VLLM_BATCH_INVARIANT", True)
73 grouped_topk = GroupedTopk(
74 topk=topk,
75 renormalize=renormalize,
@@ -89,8 +97,7 @@ def test_grouped_topk(
97 e_score_correction_bias=e_score_correction_bias,
98 )
99
92 - if renormalize:
93 - torch.testing.assert_close(
94 - baseline_topk_weights, test_topk_weights, atol=2e-2, rtol=0
95 - )
100 + torch.testing.assert_close(
101 + baseline_topk_weights, test_topk_weights, atol=2e-2, rtol=0
102 + )
103 torch.testing.assert_close(baseline_topk_ids, test_topk_ids, atol=0, rtol=0)