| 1 | /* |
| 2 | * Adapted from |
| 3 | * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh |
| 4 | * https://github.com/flashinfer-ai/flashinfer/blob/06400d062a2d51564bbe781f6f811d0b75ca593e/include/flashinfer/trtllm/fused_moe/RoutingKernelTopK.cuh |
| 5 | * Copyright (c) 2026, The vLLM team. |
| 6 | * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION. All rights |
| 7 | * reserved. SPDX-License-Identifier: Apache-2.0 |
| 8 | * |
| 9 | * Licensed under the Apache License, Version 2.0 (the "License"); |
| 10 | * you may not use this file except in compliance with the License. |
| 11 | * You may obtain a copy of the License at |
| 12 | * |
| 13 | * http://www.apache.org/licenses/LICENSE-2.0 |
| 14 | * |
| 15 | * Unless required by applicable law or agreed to in writing, software |
| 16 | * distributed under the License is distributed on an "AS IS" BASIS, |
| 17 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 18 | * See the License for the specific language governing permissions and |
| 19 | * limitations under the License. |
| 20 | */ |
| 21 | #pragma once |
| 22 | |
| 23 | #include <cooperative_groups.h> |
| 24 | #include <cooperative_groups/reduce.h> |
| 25 | #include <cub/cub.cuh> |
| 26 | |
| 27 | #include <cstdint> |
| 28 | #include <type_traits> |
| 29 | |
| 30 | namespace vllm { |
| 31 | namespace moe { |
| 32 | namespace reduce_topk { |
| 33 | namespace cg = cooperative_groups; |
| 34 | static constexpr int kWARP_SIZE = 32; |
| 35 | |
| 36 | template <typename T_> |
| 37 | struct TopKRedType { |
| 38 | using T = T_; |
| 39 | static_assert( |
| 40 | std::is_same_v<T, float> || std::is_same_v<T, half> || |
| 41 | std::is_same_v<T, __nv_bfloat16> || std::is_same_v<T, int>, |
| 42 | "Top K reduction only implemented for int, float, float16 and bfloat16"); |
| 43 | |
| 44 | using TypeCmp = std::conditional_t<sizeof(T) == 4, uint64_t, uint32_t>; |
| 45 | |
| 46 | static constexpr int kMoveBits = (sizeof(T) == 4) ? 32 : 16; |
| 47 | static constexpr int kMaxIdx = 65535; |
| 48 | TypeCmp compVal; |
| 49 | |
| 50 | static __host__ __device__ inline TypeCmp makeCmpVal(T val, int32_t idx = 0) { |
| 51 | auto valueBits = cub::Traits<T>::TwiddleIn( |
| 52 | reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(val)); |
| 53 | TypeCmp compactTmp = valueBits; |
| 54 | compactTmp = (compactTmp << kMoveBits) | (0xFFFF & (kMaxIdx - idx)); |
| 55 | // Use 65535 minus idx to give higher priority to elements with smaller |
| 56 | // indices. |
| 57 | return compactTmp; |
| 58 | } |
| 59 | |
| 60 | static __host__ __device__ void unpack(T& value, int32_t& index, |
| 61 | TypeCmp cmp) { |
| 62 | // Since “65535-idx” is always smaller than 65536 and positive, we can |
| 63 | // directly use it as the lower 16 bits |
| 64 | index = kMaxIdx - static_cast<int32_t>((cmp & 0xFFFF)); |
| 65 | |
| 66 | auto compactTmp = cmp >> kMoveBits; |
| 67 | auto valueBits = cub::Traits<T>::TwiddleOut( |
| 68 | reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(compactTmp)); |
| 69 | value = reinterpret_cast<T&>(valueBits); |
| 70 | } |
| 71 | |
| 72 | __host__ __device__ TopKRedType() = default; |
| 73 | |
| 74 | __host__ __device__ TopKRedType(T val, int32_t idx) |
| 75 | : compVal(makeCmpVal(val, idx)) {} |
| 76 | |
| 77 | __host__ __device__ operator TypeCmp() const noexcept { return compVal; } |
| 78 | |
| 79 | __device__ inline TypeCmp reduce( |
| 80 | cg::thread_block_tile<kWARP_SIZE> const& warp) { |
| 81 | #ifdef __CUDA_ARCH__ |
| 82 | static constexpr bool kHAS_FAST_REDUX = (__CUDA_ARCH__ / 100) >= 10; |
| 83 | #else |
| 84 | static constexpr bool kHAS_FAST_REDUX = false; |
| 85 | #endif |
| 86 | if constexpr (!kHAS_FAST_REDUX) { |
| 87 | return cg::reduce(warp, compVal, cg::greater<TypeCmp>{}); |
| 88 | } else if constexpr (sizeof(TypeCmp) == 8) { |
| 89 | uint32_t hi = static_cast<uint32_t>(compVal >> 32); |
| 90 | uint32_t lo = static_cast<uint32_t>(compVal & 0xffffffffu); |
| 91 | uint32_t maxHi; |
| 92 | asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n" |
| 93 | : "=r"(maxHi) |
| 94 | : "r"(hi)); |
| 95 | uint32_t loContrib = hi == maxHi ? lo : 0u; |
| 96 | uint32_t maxLo; |
| 97 | asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n" |
| 98 | : "=r"(maxLo) |
| 99 | : "r"(loContrib)); |
| 100 | return (static_cast<TypeCmp>(maxHi) << 32) | static_cast<TypeCmp>(maxLo); |
| 101 | } else { |
| 102 | TypeCmp result; |
| 103 | asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n" |
| 104 | : "=r"(result) |
| 105 | : "r"(compVal)); |
| 106 | return result; |
| 107 | } |
| 108 | } |
| 109 | }; |
| 110 | |
| 111 | template <int N> |
| 112 | struct IsPowerOf2 { |
| 113 | static constexpr bool value = N > 0 && (N & (N - 1)) == 0; |
| 114 | }; |
| 115 | |
| 116 | template <int N> |
| 117 | struct NextPow2 { |
| 118 | private: |
| 119 | static constexpr unsigned u = static_cast<unsigned>(N - 1); |
| 120 | static constexpr unsigned s1 = u | (u >> 1); |
| 121 | static constexpr unsigned s2 = s1 | (s1 >> 2); |
| 122 | static constexpr unsigned s3 = s2 | (s2 >> 4); |
| 123 | static constexpr unsigned s4 = s3 | (s3 >> 8); |
| 124 | static constexpr unsigned s5 = s4 | (s4 >> 16); |
| 125 | |
| 126 | public: |
| 127 | static constexpr int value = N <= 1 ? 1 : static_cast<int>(s5 + 1); |
| 128 | }; |
| 129 | |
| 130 | template <int A, int B, int Size, typename T> |
| 131 | __device__ __forceinline__ void topkCompareSwap(T* a) { |
| 132 | if constexpr (A < Size && B < Size) { |
| 133 | if (a[A] < a[B]) { |
| 134 | T tmp = a[A]; |
| 135 | a[A] = a[B]; |
| 136 | a[B] = tmp; |
| 137 | } |
| 138 | } else { |
| 139 | (void)a; |
| 140 | } |
| 141 | } |
| 142 | |
| 143 | template <int I, int End, int Step, int PairStride, int Size, typename T> |
| 144 | __device__ __forceinline__ void topkMergePairs(T* a) { |
| 145 | if constexpr (I + Step < End) { |
| 146 | topkCompareSwap<I, I + Step, Size, T>(a); |
| 147 | topkMergePairs<I + PairStride, End, Step, PairStride, Size, T>(a); |
| 148 | } else { |
| 149 | (void)a; |
| 150 | } |
| 151 | } |
| 152 | |
| 153 | template <int Lo, int N, int R, int Size, typename T> |
| 154 | __device__ __forceinline__ void topkOEM(T* a) { |
| 155 | constexpr int M = R * 2; |
| 156 | if constexpr (M < N) { |
| 157 | topkOEM<Lo, N, M, Size, T>(a); |
| 158 | topkOEM<Lo + R, N - R, M, Size, T>(a); |
| 159 | topkMergePairs<Lo + R, Lo + N, R, M, Size, T>(a); |
| 160 | } else if constexpr (R < N) { |
| 161 | topkCompareSwap<Lo, Lo + R, Size, T>(a); |
| 162 | } else { |
| 163 | (void)a; |
| 164 | } |
| 165 | } |
| 166 | |
| 167 | template <int Lo, int N, int Size, typename T> |
| 168 | __device__ __forceinline__ void topkSortBatcher(T* a) { |
| 169 | if constexpr (N > 1) { |
| 170 | constexpr int Half = N / 2; |
| 171 | topkSortBatcher<Lo, Half, Size, T>(a); |
| 172 | topkSortBatcher<Lo + Half, N - Half, Size, T>(a); |
| 173 | topkOEM<Lo, N, 1, Size, T>(a); |
| 174 | } else { |
| 175 | (void)a; |
| 176 | } |
| 177 | } |
| 178 | |
| 179 | template <int N, typename RedType> |
| 180 | struct Sort { |
| 181 | static_assert(N > 0 && N <= 64, "Sort only supports N in range [1, 64]"); |
| 182 | |
| 183 | static __device__ void run(RedType* topK) { |
| 184 | if constexpr (IsPowerOf2<N>::value) { |
| 185 | #pragma unroll |
| 186 | for (int k = 2; k <= N; k *= 2) { |
| 187 | #pragma unroll |
| 188 | for (int j = k / 2; j > 0; j /= 2) { |
| 189 | #pragma unroll |
| 190 | for (int i = 0; i < N; ++i) { |
| 191 | int ixj = i ^ j; |
| 192 | if (ixj > i) { |
| 193 | if ((i & k) == 0) { |
| 194 | if (topK[i].compVal < topK[ixj].compVal) { |
| 195 | auto tmp = topK[i].compVal; |
| 196 | topK[i].compVal = topK[ixj].compVal; |
| 197 | topK[ixj].compVal = tmp; |
| 198 | } |
| 199 | } else { |
| 200 | if (topK[i].compVal > topK[ixj].compVal) { |
| 201 | auto tmp = topK[i].compVal; |
| 202 | topK[i].compVal = topK[ixj].compVal; |
| 203 | topK[ixj].compVal = tmp; |
| 204 | } |
| 205 | } |
| 206 | } |
| 207 | } |
| 208 | } |
| 209 | } |
| 210 | } else { |
| 211 | constexpr int P = NextPow2<N>::value; |
| 212 | topkSortBatcher<0, P, N, RedType>(topK); |
| 213 | } |
| 214 | } |
| 215 | }; |
| 216 | |
| 217 | template <typename RedType> |
| 218 | struct Sort<1, RedType> { |
| 219 | static __device__ void run(RedType*) {} |
| 220 | }; |
| 221 | |
| 222 | template <typename RedType> |
| 223 | struct Sort<2, RedType> { |
| 224 | static __device__ void run(RedType* topK) { topkCompareSwap<0, 1, 2>(topK); } |
| 225 | }; |
| 226 | |
| 227 | template <typename RedType> |
| 228 | struct Sort<3, RedType> { |
| 229 | static __device__ void run(RedType* topK) { |
| 230 | topkCompareSwap<0, 1, 3>(topK); |
| 231 | topkCompareSwap<1, 2, 3>(topK); |
| 232 | topkCompareSwap<0, 1, 3>(topK); |
| 233 | } |
| 234 | }; |
| 235 | |
| 236 | template <typename RedType> |
| 237 | struct Sort<4, RedType> { |
| 238 | static __device__ void run(RedType* topK) { |
| 239 | topkCompareSwap<0, 2, 4>(topK); |
| 240 | topkCompareSwap<1, 3, 4>(topK); |
| 241 | topkCompareSwap<0, 1, 4>(topK); |
| 242 | topkCompareSwap<2, 3, 4>(topK); |
| 243 | topkCompareSwap<1, 2, 4>(topK); |
| 244 | } |
| 245 | }; |
| 246 | |
| 247 | template <int K, typename Type> |
| 248 | __forceinline__ __device__ void reduceTopK( |
| 249 | cg::thread_block_tile<kWARP_SIZE> const& warp, Type (&out)[K], |
| 250 | int32_t (&outIdx)[K], Type value, int32_t idx, Type const minValue, |
| 251 | int actualK = K) { |
| 252 | static_assert(K > 0, "Top K must have K > 0"); |
| 253 | static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE"); |
| 254 | using RedType = TopKRedType<Type>; |
| 255 | RedType topK{value, idx}; |
| 256 | typename RedType::TypeCmp packedMax{}; |
| 257 | #pragma unroll |
| 258 | for (int kk = 0; kk < actualK; ++kk) { |
| 259 | topK = kk > 0 && packedMax == topK.compVal ? RedType{minValue, idx} : topK; |
| 260 | packedMax = topK.reduce(warp); |
| 261 | RedType::unpack(out[kk], outIdx[kk], packedMax); |
| 262 | } |
| 263 | }; |
| 264 | |
| 265 | template <int K, typename Type, int N> |
| 266 | __forceinline__ __device__ void reduceTopK( |
| 267 | cg::thread_block_tile<kWARP_SIZE> const& warp, Type (&out)[K], |
| 268 | int32_t (&outIdx)[K], Type (&value)[N], int32_t (&idx)[N], |
| 269 | Type const minValue, int actualK = K) { |
| 270 | static_assert(K > 0, "Top K must have K > 0"); |
| 271 | static_assert(K <= kWARP_SIZE, "Top K must have K <= kWARP_SIZE"); |
| 272 | static_assert(N > 0, "Top K must have N > 0"); |
| 273 | static_assert(N <= 64, |
| 274 | "Only support candidates number less than or equal to " |
| 275 | "64*32=2048"); |
| 276 | using RedType = TopKRedType<Type>; |
| 277 | RedType topK[N]; |
| 278 | #pragma unroll |
| 279 | for (int nn = 0; nn < N; ++nn) { |
| 280 | topK[nn] = RedType{value[nn], idx[nn]}; |
| 281 | } |
| 282 | |
| 283 | Sort<N, RedType>::run(topK); |
| 284 | |
| 285 | typename RedType::TypeCmp packedMax{}; |
| 286 | for (int kk = 0; kk < actualK; ++kk) { |
| 287 | bool update = kk > 0 && packedMax == topK[0].compVal; |
| 288 | #pragma unroll |
| 289 | for (int nn = 0; nn < N; ++nn) { |
| 290 | topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]} |
| 291 | : update ? topK[nn + 1] |
| 292 | : topK[nn]; |
| 293 | } |
| 294 | packedMax = topK[0].reduce(warp); |
| 295 | RedType::unpack(out[kk], outIdx[kk], packedMax); |
| 296 | } |
| 297 | }; |
| 298 | |
| 299 | template <int NumExperts, int NumTopExperts, int MinExperts, int MaxExperts, |
| 300 | int MinTopExperts, int MaxTopExperts> |
| 301 | struct LaneOwnedTopKRange { |
| 302 | static_assert(MinExperts > 0 && MinExperts <= MaxExperts); |
| 303 | static_assert(MinTopExperts > 0 && MinTopExperts <= MaxTopExperts); |
| 304 | static constexpr bool kEnabled = |
| 305 | NumExperts >= MinExperts && NumExperts <= MaxExperts && |
| 306 | NumTopExperts >= MinTopExperts && NumTopExperts <= MaxTopExperts; |
| 307 | }; |
| 308 | |
| 309 | static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MIN_EXPERTS = 512; |
| 310 | static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MAX_EXPERTS = 1024; |
| 311 | static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MIN_TOP_EXPERTS = 9; |
| 312 | static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MAX_TOP_EXPERTS = 16; |
| 313 | |
| 314 | template <int NumExperts, int NumTopExperts> |
| 315 | using HighExpertLaneOwnedTopKRange = |
| 316 | LaneOwnedTopKRange<NumExperts, NumTopExperts, |
| 317 | kHIGH_EXPERT_LANE_OWNED_TOPK_MIN_EXPERTS, |
| 318 | kHIGH_EXPERT_LANE_OWNED_TOPK_MAX_EXPERTS, |
| 319 | kHIGH_EXPERT_LANE_OWNED_TOPK_MIN_TOP_EXPERTS, |
| 320 | kHIGH_EXPERT_LANE_OWNED_TOPK_MAX_TOP_EXPERTS>; |
| 321 | |
| 322 | template <int K, typename Type, int N> |
| 323 | __forceinline__ __device__ void reduceTopKForLane( |
| 324 | cg::thread_block_tile<kWARP_SIZE> const& warp, Type& out, int32_t& outIdx, |
| 325 | Type (&value)[N], int32_t (&idx)[N], Type const minValue, int32_t laneIdx) { |
| 326 | static_assert(K > 0, "Top K must have K > 0"); |
| 327 | static_assert(K <= kWARP_SIZE, "Top K must have K <= kWARP_SIZE"); |
| 328 | static_assert(N > 0, "Top K must have N > 0"); |
| 329 | static_assert(N <= 64, |
| 330 | "Only support candidates number less than or equal to " |
| 331 | "64*32=2048"); |
| 332 | using RedType = TopKRedType<Type>; |
| 333 | RedType topK[N]; |
| 334 | #pragma unroll |
| 335 | for (int nn = 0; nn < N; ++nn) { |
| 336 | topK[nn] = RedType{value[nn], idx[nn]}; |
| 337 | } |
| 338 | |
| 339 | Sort<N, RedType>::run(topK); |
| 340 | |
| 341 | typename RedType::TypeCmp packedMax{}; |
| 342 | typename RedType::TypeCmp lanePacked{}; |
| 343 | #pragma unroll |
| 344 | for (int kk = 0; kk < K; ++kk) { |
| 345 | bool update = kk > 0 && packedMax == topK[0].compVal; |
| 346 | #pragma unroll |
| 347 | for (int nn = 0; nn < N; ++nn) { |
| 348 | topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]} |
| 349 | : update ? topK[nn + 1] |
| 350 | : topK[nn]; |
| 351 | } |
| 352 | packedMax = topK[0].reduce(warp); |
| 353 | if (laneIdx == kk) { |
| 354 | lanePacked = packedMax; |
| 355 | } |
| 356 | } |
| 357 | |
| 358 | if (laneIdx < K) { |
| 359 | RedType::unpack(out, outIdx, lanePacked); |
| 360 | } else { |
| 361 | out = minValue; |
| 362 | outIdx = -1; |
| 363 | } |
| 364 | } |
| 365 | |
| 366 | } // namespace reduce_topk |
| 367 | } // namespace moe |
| 368 | } // namespace vllm |