| 1 | /* |
| 2 | * Adapted from |
| 3 | * https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu |
| 4 | * Copyright (c) 2024, The vLLM team. |
| 5 | * SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & |
| 6 | * AFFILIATES. All rights 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 | #include <type_traits> |
| 21 | |
| 22 | #include <cuda_runtime.h> |
| 23 | #include <torch/csrc/stable/accelerator.h> |
| 24 | #include <torch/csrc/stable/tensor.h> |
| 25 | #include <torch/headeronly/core/ScalarType.h> |
| 26 | #include <torch/headeronly/util/Exception.h> |
| 27 | |
| 28 | #include "../../cuda_compat.h" |
| 29 | #include "../cub_helpers.h" |
| 30 | #include "libtorch_stable/torch_utils.h" |
| 31 | #ifndef USE_ROCM |
| 32 | #include <cuda_bf16.h> |
| 33 | #include <cuda_fp16.h> |
| 34 | #else |
| 35 | #include <hip/hip_bf16.h> |
| 36 | #include <hip/hip_fp16.h> |
| 37 | typedef __hip_bfloat16 __nv_bfloat16; |
| 38 | typedef __hip_bfloat162 __nv_bfloat162; |
| 39 | #endif |
| 40 | |
| 41 | #define MAX(a, b) ((a) > (b) ? (a) : (b)) |
| 42 | #define MIN(a, b) ((a) < (b) ? (a) : (b)) |
| 43 | |
| 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 |
| 56 | int N, |
| 57 | /// Alignment requirement in bytes |
| 58 | int Alignment = sizeof(T) * N> |
| 59 | struct alignas(Alignment) AlignedArray { |
| 60 | T data[N]; |
| 61 | }; |
| 62 | |
| 63 | template <typename T> |
| 64 | __device__ __forceinline__ float toFloat(T value) { |
| 65 | if constexpr (std::is_same_v<T, float>) { |
| 66 | return value; |
| 67 | } else if constexpr (std::is_same_v<T, __nv_bfloat16>) { |
| 68 | return __bfloat162float(value); |
| 69 | } else if constexpr (std::is_same_v<T, __half>) { |
| 70 | return __half2float(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 bool* is_padding, const float* bias_vl, |
| 85 | int64_t image_sentinel_lo) { |
| 86 | const int warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32; |
| 87 | const int lane = threadIdx.x % 32; |
| 88 | if (warp >= num_rows) return; |
| 89 | const int64_t token_id = load_index_as_int64(input_ids, warp); |
| 90 | const bool is_pad_row = is_padding != nullptr && is_padding[warp]; |
| 91 | // Image tokens carry five consecutive in-vocab sentinel ids starting at |
| 92 | // image_sentinel_lo (0 = disabled); they skip the tid2eid lookup and |
| 93 | // select experts by score + bias_vl instead. Ids above the sentinel block |
| 94 | // are regular special tokens and must not match. |
| 95 | const bool is_image = bias_vl != nullptr && image_sentinel_lo > 0 && |
| 96 | token_id >= image_sentinel_lo && |
| 97 | token_id < image_sentinel_lo + 5; |
| 98 | |
| 99 | #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) |
| 100 | cudaGridDependencySynchronize(); |
| 101 | #endif |
| 102 | int expert = 0; |
| 103 | float weight = 0.f; |
| 104 | if (is_image && !is_pad_row) { |
| 105 | // Top-6 by (sqrtsoftplus(x) + bias_vl) over the full row; the routing |
| 106 | // weights are the raw sqrtsoftplus scores of the selected experts. |
| 107 | constexpr int MAX_PER_LANE = 12; // supports num_experts <= 384 |
| 108 | const int per_lane = num_experts / 32; |
| 109 | float vals[MAX_PER_LANE]; |
| 110 | float sels[MAX_PER_LANE]; |
| 111 | #pragma unroll |
| 112 | for (int i = 0; i < MAX_PER_LANE; ++i) { |
| 113 | if (i < per_lane) { |
| 114 | const int e = lane + 32 * i; |
| 115 | const float x = input[warp * num_experts + e]; |
| 116 | float v = sqrtf(fmaxf(x, 0.f) + __logf(1.f + __expf(-fabsf(x)))); |
| 117 | if (isnan(v)) { |
| 118 | v = 0.f; |
| 119 | } |
| 120 | vals[i] = v; |
| 121 | sels[i] = v + bias_vl[e]; |
| 122 | } else { |
| 123 | vals[i] = 0.f; |
| 124 | sels[i] = -INFINITY; |
| 125 | } |
| 126 | } |
| 127 | int slot_expert[6]; |
| 128 | float slot_weight[6]; |
| 129 | #pragma unroll |
| 130 | for (int j = 0; j < 6; ++j) { |
| 131 | float best = -INFINITY; |
| 132 | int best_e = num_experts; |
| 133 | #pragma unroll |
| 134 | for (int i = 0; i < MAX_PER_LANE; ++i) { |
| 135 | if (sels[i] > best) { |
| 136 | best = sels[i]; |
| 137 | best_e = lane + 32 * i; |
| 138 | } |
| 139 | } |
| 140 | // Warp argmax reduce; ties go to the lower expert id. |
| 141 | #pragma unroll |
| 142 | for (int mask = 16; mask > 0; mask >>= 1) { |
| 143 | const float other_best = VLLM_SHFL_XOR_SYNC(best, mask); |
| 144 | const int other_e = VLLM_SHFL_XOR_SYNC(best_e, mask); |
| 145 | if (other_best > best || (other_best == best && other_e < best_e)) { |
| 146 | best = other_best; |
| 147 | best_e = other_e; |
| 148 | } |
| 149 | } |
| 150 | float w = 0.f; |
| 151 | #pragma unroll |
| 152 | for (int i = 0; i < MAX_PER_LANE; ++i) { |
| 153 | if (lane + 32 * i == best_e) { |
| 154 | w = vals[i]; |
| 155 | sels[i] = -INFINITY; |
| 156 | } |
| 157 | } |
| 158 | #pragma unroll |
| 159 | for (int mask = 16; mask > 0; mask >>= 1) { |
| 160 | w += VLLM_SHFL_XOR_SYNC(w, mask); |
| 161 | } |
| 162 | slot_expert[j] = best_e; |
| 163 | slot_weight[j] = w; |
| 164 | } |
| 165 | if (lane < 6) { |
| 166 | expert = slot_expert[lane]; |
| 167 | weight = slot_weight[lane]; |
| 168 | } |
| 169 | } else if (lane < 6 && !is_pad_row) { |
| 170 | // only load and calculate for 6 experts |
| 171 | expert = static_cast<int>(tid2eid[token_id * 6 + lane]); |
| 172 | const float x = input[warp * num_experts + expert]; |
| 173 | weight = sqrtf(fmaxf(x, 0.f) + __logf(1.f + __expf(-fabsf(x)))); |
| 174 | if (isnan(weight)) { |
| 175 | weight = 0.f; |
| 176 | } |
| 177 | } |
| 178 | float weight_sum = weight; |
| 179 | #pragma unroll |
| 180 | for (int mask = 16; mask > 0; mask >>= 1) { |
| 181 | // sum in warp |
| 182 | weight_sum += VLLM_SHFL_XOR_SYNC(weight_sum, mask); |
| 183 | } |
| 184 | |
| 185 | #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) |
| 186 | cudaTriggerProgrammaticLaunchCompletion(); |
| 187 | #endif |
| 188 | if (lane < 6) { |
| 189 | const int offset = warp * 6 + lane; |
| 190 | output[offset] = |
| 191 | weight * routed_scaling_factor / (weight_sum > 0.f ? weight_sum : 1.f); |
| 192 | indices[offset] = !is_pad_row ? static_cast<OutIndType>(expert) |
| 193 | : static_cast<OutIndType>(-1); |
| 194 | } |
| 195 | } |
| 196 | |
| 197 | template <typename OutIndType, typename HashIndType> |
| 198 | void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices, |
| 199 | int num_rows, int num_experts, |
| 200 | double routed_scaling_factor, |
| 201 | const HashIndType* input_ids, |
| 202 | const HashIndType* tid2eid, cudaStream_t stream, |
| 203 | const bool* is_padding, const float* bias_vl, |
| 204 | int64_t image_sentinel_lo) { |
| 205 | if (num_rows == 0) return; |
| 206 | auto* kernel = &dsv4HashTopkSoftplusSqrt<OutIndType, HashIndType>; |
| 207 | cudaLaunchConfig_t config = {}; |
| 208 | config.gridDim = (num_rows + 3) / 4; |
| 209 | config.blockDim = 128; |
| 210 | config.stream = stream; |
| 211 | cudaLaunchAttribute attr; |
| 212 | attr.id = cudaLaunchAttributeProgrammaticStreamSerialization; |
| 213 | attr.val.programmaticStreamSerializationAllowed = 1; |
| 214 | config.attrs = &attr; |
| 215 | config.numAttrs = 1; |
| 216 | const float scale = static_cast<float>(routed_scaling_factor); |
| 217 | cudaLaunchKernelEx(&config, kernel, input, output, indices, num_rows, |
| 218 | num_experts, scale, input_ids, tid2eid, is_padding, |
| 219 | bias_vl, image_sentinel_lo); |
| 220 | } |
| 221 | #endif |
| 222 | |
| 223 | // ====================== TopK softplus_sqrt things |
| 224 | // =============================== |
| 225 | |
| 226 | /* |
| 227 | A Top-K gating softplus_sqrt written to exploit when the number of experts in |
| 228 | the MoE layers are a small power of 2. This allows us to cleanly share the |
| 229 | rows among the threads in a single warp and eliminate communication between |
| 230 | warps (so no need to use shared mem). |
| 231 | |
| 232 | It fuses the sigmoid, max and argmax into a single kernel. |
| 233 | |
| 234 | Limitations: |
| 235 | 1) This implementation is optimized for when the number of experts is a small |
| 236 | power of 2. Additionally it also supports when number of experts is multiple |
| 237 | of 64 which is still faster than the computing sigmoid and topK separately |
| 238 | (only tested on CUDA yet). 2) This implementation assumes k is small, but will |
| 239 | work for any k. |
| 240 | */ |
| 241 | |
| 242 | template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, |
| 243 | int WARP_SIZE_PARAM, bool USE_HASH, typename IndType, |
| 244 | typename HashIndType, typename InputType = float> |
| 245 | __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ |
| 246 | void topkGatingSoftplusSqrt( |
| 247 | const InputType* input, const bool* finished, float* output, |
| 248 | const int num_rows, IndType* indices, int* source_rows, const int k, |
| 249 | const int start_expert, const int end_expert, const bool renormalize, |
| 250 | double routed_scaling_factor, const float* correction_bias, |
| 251 | const HashIndType* input_ids, const HashIndType* tid2eid, |
| 252 | const bool* is_padding, const float* bias_vl, |
| 253 | int64_t image_sentinel_lo) { |
| 254 | static_assert(std::is_same_v<InputType, float> || |
| 255 | std::is_same_v<InputType, __nv_bfloat16> || |
| 256 | std::is_same_v<InputType, __half>, |
| 257 | "InputType must be float, __nv_bfloat16, or __half"); |
| 258 | |
| 259 | // We begin by enforcing compile time assertions and setting up compile time |
| 260 | // constants. |
| 261 | static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG), |
| 262 | "BYTES_PER_LDG must be power of 2"); |
| 263 | static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16"); |
| 264 | |
| 265 | // Number of bytes each thread pulls in per load |
| 266 | static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType); |
| 267 | static constexpr int ELTS_PER_ROW = NUM_EXPERTS; |
| 268 | static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT; |
| 269 | static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG; |
| 270 | |
| 271 | if constexpr (std::is_same_v<InputType, __nv_bfloat16> || |
| 272 | std::is_same_v<InputType, __half>) { |
| 273 | static_assert(ELTS_PER_LDG == 1 || ELTS_PER_LDG % 2 == 0, |
| 274 | "ELTS_PER_LDG must be 1 or even for 16-bit conversion"); |
| 275 | } |
| 276 | |
| 277 | // Restrictions based on previous section. |
| 278 | static_assert( |
| 279 | VPT % ELTS_PER_LDG == 0, |
| 280 | "The elements per thread must be a multiple of the elements per ldg"); |
| 281 | static_assert(WARP_SIZE_PARAM % THREADS_PER_ROW == 0, |
| 282 | "The threads per row must cleanly divide the threads per warp"); |
| 283 | static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW), |
| 284 | "THREADS_PER_ROW must be power of 2"); |
| 285 | static_assert(THREADS_PER_ROW <= WARP_SIZE_PARAM, |
| 286 | "THREADS_PER_ROW can be at most warp size"); |
| 287 | |
| 288 | // We have NUM_EXPERTS elements per row. We specialize for small #experts |
| 289 | static constexpr int ELTS_PER_WARP = WARP_SIZE_PARAM * VPT; |
| 290 | static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW; |
| 291 | static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP; |
| 292 | |
| 293 | // Restrictions for previous section. |
| 294 | static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0, |
| 295 | "The elts per row must cleanly divide the total elt per warp"); |
| 296 | |
| 297 | // ===================== From this point, we finally start computing run-time |
| 298 | // variables. ======================== |
| 299 | |
| 300 | // Compute CTA and warp rows. We pack multiple rows into a single warp, and a |
| 301 | // block contains WARPS_PER_CTA warps. This, each block processes a chunk of |
| 302 | // rows. We start by computing the start row for each block. |
| 303 | const int cta_base_row = blockIdx.x * ROWS_PER_CTA; |
| 304 | |
| 305 | // Now, using the base row per thread block, we compute the base row per warp. |
| 306 | const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP; |
| 307 | |
| 308 | // The threads in a warp are split into sub-groups that will work on a row. |
| 309 | // We compute row offset for each thread sub-group |
| 310 | const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW; |
| 311 | const int thread_row = warp_base_row + thread_row_in_warp; |
| 312 | |
| 313 | // Threads with indices out of bounds should early exit here. |
| 314 | if (thread_row >= num_rows) { |
| 315 | return; |
| 316 | } |
| 317 | const bool row_is_active = finished ? !finished[thread_row] : true; |
| 318 | const bool is_pad_row = is_padding != nullptr && is_padding[thread_row]; |
| 319 | |
| 320 | // Image tokens carry five consecutive in-vocab sentinel ids starting at |
| 321 | // image_sentinel_lo (0 = disabled); they select experts with bias_vl |
| 322 | // instead of correction_bias / the tid2eid hash table. Ids above the |
| 323 | // sentinel block are regular special tokens and must not match. |
| 324 | int64_t token_id = 0; |
| 325 | if (input_ids != nullptr) { |
| 326 | token_id = load_index_as_int64(input_ids, thread_row); |
| 327 | } |
| 328 | const bool use_vl_bias = bias_vl != nullptr && image_sentinel_lo > 0 && |
| 329 | token_id >= image_sentinel_lo && |
| 330 | token_id < image_sentinel_lo + 5; |
| 331 | const float* row_bias = use_vl_bias ? bias_vl : correction_bias; |
| 332 | |
| 333 | // We finally start setting up the read pointers for each thread. First, each |
| 334 | // thread jumps to the start of the row it will read. |
| 335 | const InputType* thread_row_ptr = input + thread_row * ELTS_PER_ROW; |
| 336 | |
| 337 | // Now, we compute the group each thread belong to in order to determine the |
| 338 | // first column to start loads. |
| 339 | const int thread_group_idx = threadIdx.x % THREADS_PER_ROW; |
| 340 | const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG; |
| 341 | const InputType* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread; |
| 342 | |
| 343 | // Finally, we pull in the data from global mem |
| 344 | float row_chunk[VPT]; |
| 345 | |
| 346 | #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) |
| 347 | cudaGridDependencySynchronize(); |
| 348 | #endif |
| 349 | |
| 350 | if (is_pad_row) { |
| 351 | #pragma unroll |
| 352 | for (int ii = 0; ii < VPT; ++ii) { |
| 353 | row_chunk[ii] = 0.f; |
| 354 | } |
| 355 | } else if constexpr (std::is_same_v<InputType, float>) { |
| 356 | using VecType = AlignedArray<float, ELTS_PER_LDG>; |
| 357 | VecType* row_chunk_vec_ptr = reinterpret_cast<VecType*>(&row_chunk); |
| 358 | const VecType* vec_thread_read_ptr = |
| 359 | reinterpret_cast<const VecType*>(thread_read_ptr); |
| 360 | #pragma unroll |
| 361 | for (int ii = 0; ii < LDG_PER_THREAD; ++ii) { |
| 362 | row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW]; |
| 363 | } |
| 364 | } else if constexpr (std::is_same_v<InputType, __nv_bfloat16>) { |
| 365 | if constexpr (ELTS_PER_LDG >= 2) { |
| 366 | using VecType = AlignedArray<__nv_bfloat16, ELTS_PER_LDG>; |
| 367 | float2* row_chunk_f2 = reinterpret_cast<float2*>(row_chunk); |
| 368 | const VecType* vec_thread_read_ptr = |
| 369 | reinterpret_cast<const VecType*>(thread_read_ptr); |
| 370 | #pragma unroll |
| 371 | for (int ii = 0; ii < LDG_PER_THREAD; ++ii) { |
| 372 | VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW]; |
| 373 | int base_idx_f2 = ii * ELTS_PER_LDG / 2; |
| 374 | #pragma unroll |
| 375 | for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) { |
| 376 | row_chunk_f2[base_idx_f2 + jj] = __bfloat1622float2( |
| 377 | *reinterpret_cast<const __nv_bfloat162*>(vec.data + jj * 2)); |
| 378 | } |
| 379 | } |
| 380 | } else { // ELTS_PER_LDG == 1 |
| 381 | #pragma unroll |
| 382 | for (int ii = 0; ii < LDG_PER_THREAD; ++ii) { |
| 383 | const __nv_bfloat16* scalar_ptr = |
| 384 | thread_read_ptr + ii * THREADS_PER_ROW; |
| 385 | row_chunk[ii] = __bfloat162float(*scalar_ptr); |
| 386 | } |
| 387 | } |
| 388 | } else if constexpr (std::is_same_v<InputType, __half>) { |
| 389 | if constexpr (ELTS_PER_LDG >= 2) { |
| 390 | using VecType = AlignedArray<__half, ELTS_PER_LDG>; |
| 391 | float2* row_chunk_f2 = reinterpret_cast<float2*>(row_chunk); |
| 392 | const VecType* vec_thread_read_ptr = |
| 393 | reinterpret_cast<const VecType*>(thread_read_ptr); |
| 394 | #pragma unroll |
| 395 | for (int ii = 0; ii < LDG_PER_THREAD; ++ii) { |
| 396 | VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW]; |
| 397 | int base_idx_f2 = ii * ELTS_PER_LDG / 2; |
| 398 | #pragma unroll |
| 399 | for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) { |
| 400 | row_chunk_f2[base_idx_f2 + jj] = __half22float2( |
| 401 | *reinterpret_cast<const __half2*>(vec.data + jj * 2)); |
| 402 | } |
| 403 | } |
| 404 | } else { // ELTS_PER_LDG == 1 |
| 405 | #pragma unroll |
| 406 | for (int ii = 0; ii < LDG_PER_THREAD; ++ii) { |
| 407 | const __half* scalar_ptr = thread_read_ptr + ii * THREADS_PER_ROW; |
| 408 | row_chunk[ii] = __half2float(*scalar_ptr); |
| 409 | } |
| 410 | } |
| 411 | } |
| 412 | constexpr float threshold = 20.0f; |
| 413 | constexpr float beta = 1.0f; |
| 414 | |
| 415 | // Hash MoE path: indices are predetermined from lookup table. Image rows |
| 416 | // (sentinel ids) fall through to the score-based top-k path with bias_vl. |
| 417 | const bool hash_row = USE_HASH && !use_vl_bias; |
| 418 | if (hash_row) { |
| 419 | const int64_t token_expert_offset = token_id * static_cast<int64_t>(k); |
| 420 | if (!is_pad_row) { |
| 421 | #pragma unroll |
| 422 | for (int ii = 0; ii < VPT; ++ii) { |
| 423 | float val = row_chunk[ii]; |
| 424 | float val_b = val * beta; |
| 425 | val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta; |
| 426 | val = sqrtf(val); |
| 427 | |
| 428 | // Dummy/padding tokens can result in NaN values, so |
| 429 | // clamp them to 0.0. Note: this clamp could likely be removed if |
| 430 | // 'is_padding' is made mandatory |
| 431 | if (isnan(val)) { |
| 432 | val = 0.f; |
| 433 | } |
| 434 | row_chunk[ii] = val; |
| 435 | } |
| 436 | } |
| 437 | float selected_sum = 0.f; |
| 438 | #pragma unroll |
| 439 | for (int k_idx = 0; k_idx < k; ++k_idx) { |
| 440 | const int expert = static_cast<int>( |
| 441 | load_index_as_int64(tid2eid, token_expert_offset + k_idx)); |
| 442 | const int idx = k * thread_row + k_idx; |
| 443 | for (int ii = 0; ii < VPT; ++ii) { |
| 444 | const int group_id = ii / ELTS_PER_LDG; |
| 445 | const int local_id = ii % ELTS_PER_LDG; |
| 446 | const int expert_idx = first_elt_read_by_thread + |
| 447 | group_id * THREADS_PER_ROW * ELTS_PER_LDG + |
| 448 | local_id; |
| 449 | if (expert == expert_idx) { |
| 450 | indices[idx] = !is_pad_row ? static_cast<IndType>(expert) |
| 451 | : static_cast<IndType>(-1); |
| 452 | selected_sum += row_chunk[ii]; |
| 453 | break; |
| 454 | } |
| 455 | } |
| 456 | } |
| 457 | // Compute per-thread scale (using warp reduction when renormalizing). |
| 458 | // THREADS_PER_ROW-parameterized butterfly works for both warp sizes (32 |
| 459 | // on CUDA, 64 on ROCm CDNA) and any THREADS_PER_ROW the dispatch picks. |
| 460 | if (renormalize) { |
| 461 | #pragma unroll |
| 462 | for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) { |
| 463 | selected_sum += |
| 464 | VLLM_SHFL_XOR_SYNC_WIDTH(selected_sum, mask, THREADS_PER_ROW); |
| 465 | } |
| 466 | } |
| 467 | float scale = static_cast<float>(routed_scaling_factor); |
| 468 | if (renormalize) { |
| 469 | const float denom = selected_sum > 0.f ? selected_sum : 1.f; |
| 470 | scale /= denom; |
| 471 | } |
| 472 | |
| 473 | #pragma unroll |
| 474 | for (int k_idx = 0; k_idx < k; ++k_idx) { |
| 475 | const int expert = static_cast<int>( |
| 476 | load_index_as_int64(tid2eid, token_expert_offset + k_idx)); |
| 477 | const int idx = k * thread_row + k_idx; |
| 478 | for (int ii = 0; ii < VPT; ++ii) { |
| 479 | const int group_id = ii / ELTS_PER_LDG; |
| 480 | const int local_id = ii % ELTS_PER_LDG; |
| 481 | const int expert_idx = first_elt_read_by_thread + |
| 482 | group_id * THREADS_PER_ROW * ELTS_PER_LDG + |
| 483 | local_id; |
| 484 | if (expert == expert_idx) { |
| 485 | output[idx] = row_chunk[ii] * scale; |
| 486 | break; |
| 487 | } |
| 488 | } |
| 489 | } |
| 490 | #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) |
| 491 | cudaTriggerProgrammaticLaunchCompletion(); |
| 492 | #endif |
| 493 | return; |
| 494 | } else { |
| 495 | if (!is_pad_row) { |
| 496 | #pragma unroll |
| 497 | for (int ii = 0; ii < VPT; ++ii) { |
| 498 | float val = row_chunk[ii]; |
| 499 | float val_b = val * beta; |
| 500 | // Compute softplus: log(1 + exp(val)) with numerical stability |
| 501 | // When val > threshold, softplus(x) ≈ x to avoid exp overflow |
| 502 | val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta; |
| 503 | val = sqrtf(val); |
| 504 | // Dummy/padding tokens can result in NaN values, so |
| 505 | // clamp them to 0.0. Note: this clamp could likely be removed if |
| 506 | // 'is_padding' is made mandatory |
| 507 | if (isnan(val)) { |
| 508 | val = 0.f; |
| 509 | } |
| 510 | if (row_bias) { |
| 511 | const int group_id = ii / ELTS_PER_LDG; |
| 512 | const int local_id = ii % ELTS_PER_LDG; |
| 513 | const int expert_idx = first_elt_read_by_thread + |
| 514 | group_id * THREADS_PER_ROW * ELTS_PER_LDG + |
| 515 | local_id; |
| 516 | val = val + row_bias[expert_idx]; |
| 517 | } |
| 518 | row_chunk[ii] = val; |
| 519 | } |
| 520 | } |
| 521 | |
| 522 | // Original TopK path: find top-k experts by score |
| 523 | // Now, sigmoid_res contains the sigmoid of the row chunk. Now, I want to |
| 524 | // find the topk elements in each row, along with the max index. |
| 525 | int start_col = first_elt_read_by_thread; |
| 526 | static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW; |
| 527 | |
| 528 | float selected_sum = 0.f; |
| 529 | for (int k_idx = 0; k_idx < k; ++k_idx) { |
| 530 | // First, each thread does the local argmax |
| 531 | float max_val = row_chunk[0]; |
| 532 | int expert = start_col; |
| 533 | #pragma unroll |
| 534 | for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD; |
| 535 | ++ldg, col += COLS_PER_GROUP_LDG) { |
| 536 | #pragma unroll |
| 537 | for (int ii = 0; ii < ELTS_PER_LDG; ++ii) { |
| 538 | float val = row_chunk[ldg * ELTS_PER_LDG + ii]; |
| 539 | |
| 540 | // No check on the experts here since columns with the smallest index |
| 541 | // are processed first and only updated if > (not >=) |
| 542 | if (val > max_val) { |
| 543 | max_val = val; |
| 544 | expert = col + ii; |
| 545 | } |
| 546 | } |
| 547 | } |
| 548 | |
| 549 | // Now, we perform the argmax reduce. We use the butterfly pattern so threads |
| 550 | // reach consensus about the max. This will be useful for K > 1 so that the |
| 551 | // threads can agree on "who" had the max value. That thread can then blank out |
| 552 | // their max with -inf and the warp can run more iterations... |
| 553 | #pragma unroll |
| 554 | for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) { |
| 555 | float other_max = |
| 556 | VLLM_SHFL_XOR_SYNC_WIDTH(max_val, mask, THREADS_PER_ROW); |
| 557 | int other_expert = |
| 558 | VLLM_SHFL_XOR_SYNC_WIDTH(expert, mask, THREADS_PER_ROW); |
| 559 | |
| 560 | // We want lower indices to "win" in every thread so we break ties this |
| 561 | // way |
| 562 | if (other_max > max_val || |
| 563 | (other_max == max_val && other_expert < expert)) { |
| 564 | max_val = other_max; |
| 565 | expert = other_expert; |
| 566 | } |
| 567 | } |
| 568 | |
| 569 | // Write the max for this k iteration to global memory. |
| 570 | if (thread_group_idx == 0) { |
| 571 | // Add a guard to ignore experts not included by this node |
| 572 | const bool node_uses_expert = |
| 573 | expert >= start_expert && expert < end_expert; |
| 574 | const bool should_process_row = |
| 575 | row_is_active && node_uses_expert && !is_pad_row; |
| 576 | |
| 577 | // The lead thread from each sub-group will write out the final results |
| 578 | // to global memory. (This will be a single) thread per row of the |
| 579 | // input/output matrices. |
| 580 | const int idx = k * thread_row + k_idx; |
| 581 | if (row_bias != nullptr && should_process_row) { |
| 582 | max_val -= row_bias[expert]; |
| 583 | } |
| 584 | output[idx] = max_val; |
| 585 | indices[idx] = |
| 586 | !is_pad_row ? expert - start_expert : static_cast<IndType>(-1); |
| 587 | source_rows[idx] = k_idx * num_rows + thread_row; |
| 588 | if (renormalize) { |
| 589 | selected_sum += max_val; |
| 590 | } |
| 591 | } |
| 592 | |
| 593 | // Finally, we clear the value in the thread with the current max if there |
| 594 | // is another iteration to run. |
| 595 | if (k_idx + 1 < k) { |
| 596 | const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG; |
| 597 | const int thread_to_clear_in_group = |
| 598 | (expert / ELTS_PER_LDG) % THREADS_PER_ROW; |
| 599 | |
| 600 | // Only the thread in the group which produced the max will reset the |
| 601 | // "winning" value to -inf. |
| 602 | if (thread_group_idx == thread_to_clear_in_group) { |
| 603 | const int offset_for_expert = expert % ELTS_PER_LDG; |
| 604 | // Safe to set to any negative value since row_chunk values must be |
| 605 | // between 0 and 1. |
| 606 | row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] = |
| 607 | -10000.f; |
| 608 | } |
| 609 | } |
| 610 | } |
| 611 | |
| 612 | // Apply renormalization and routed scaling factor to final weights. |
| 613 | if (thread_group_idx == 0) { |
| 614 | float scale = static_cast<float>(routed_scaling_factor); |
| 615 | if (renormalize) { |
| 616 | const float denom = selected_sum > 0.f ? selected_sum : 1.f; |
| 617 | scale /= denom; |
| 618 | } |
| 619 | for (int k_idx = 0; k_idx < k; ++k_idx) { |
| 620 | const int idx = k * thread_row + k_idx; |
| 621 | output[idx] = output[idx] * scale; |
| 622 | } |
| 623 | } |
| 624 | #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) |
| 625 | cudaTriggerProgrammaticLaunchCompletion(); |
| 626 | #endif |
| 627 | } |
| 628 | } |
| 629 | |
| 630 | namespace detail { |
| 631 | // Constructs some constants needed to partition the work across threads at |
| 632 | // compile time. |
| 633 | template <int EXPERTS, int BYTES_PER_LDG, int WARP_SIZE_PARAM, |
| 634 | typename InputType> |
| 635 | struct TopkConstants { |
| 636 | static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType); |
| 637 | static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0 || |
| 638 | EXPERTS % (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0, |
| 639 | ""); |
| 640 | static constexpr int VECs_PER_THREAD = |
| 641 | MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM)); |
| 642 | static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG; |
| 643 | static constexpr int THREADS_PER_ROW = EXPERTS / VPT; |
| 644 | static const int ROWS_PER_WARP = WARP_SIZE_PARAM / THREADS_PER_ROW; |
| 645 | }; |
| 646 | } // namespace detail |
| 647 | |
| 648 | #define DISPATCH_HASH(use_hash, USE_HASH, ...) \ |
| 649 | if (use_hash) { \ |
| 650 | const bool USE_HASH = true; \ |
| 651 | static_assert(USE_HASH == true, "USE_HASH must be compile-time constant"); \ |
| 652 | __VA_ARGS__ \ |
| 653 | } else { \ |
| 654 | const bool USE_HASH = false; \ |
| 655 | static_assert(USE_HASH == false, \ |
| 656 | "USE_HASH must be compile-time constant"); \ |
| 657 | __VA_ARGS__ \ |
| 658 | } |
| 659 | |
| 660 | template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, |
| 661 | int MAX_BYTES_PER_LDG, typename IndType, typename HashIndType, |
| 662 | typename InputType> |
| 663 | void topkGatingSoftplusSqrtLauncherHelper( |
| 664 | const InputType* input, const bool* finished, float* output, |
| 665 | IndType* indices, int* source_row, const int num_rows, const int k, |
| 666 | const int start_expert, const int end_expert, const bool renormalize, |
| 667 | double routed_scaling_factor, const float* correction_bias, |
| 668 | const bool use_hash, const HashIndType* input_ids, |
| 669 | const HashIndType* tid2eid, cudaStream_t stream, const bool* is_padding, |
| 670 | const float* bias_vl, int64_t image_sentinel_lo) { |
| 671 | static constexpr int BYTES_PER_LDG = |
| 672 | MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS); |
| 673 | using Constants = |
| 674 | detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM, InputType>; |
| 675 | static constexpr int VPT = Constants::VPT; |
| 676 | static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP; |
| 677 | const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP; |
| 678 | const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB; |
| 679 | dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB); |
| 680 | DISPATCH_HASH(use_hash, USE_HASH, { |
| 681 | auto* kernel = |
| 682 | &topkGatingSoftplusSqrt<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, |
| 683 | WARP_SIZE_PARAM, USE_HASH, IndType, HashIndType, |
| 684 | InputType>; |
| 685 | #ifndef USE_ROCM |
| 686 | cudaLaunchConfig_t config = {}; |
| 687 | config.gridDim = num_blocks; |
| 688 | config.blockDim = block_dim; |
| 689 | config.dynamicSmemBytes = 0; |
| 690 | config.stream = stream; |
| 691 | cudaLaunchAttribute attrs[1]; |
| 692 | attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; |
| 693 | attrs[0].val.programmaticStreamSerializationAllowed = 1; |
| 694 | config.numAttrs = 1; |
| 695 | config.attrs = attrs; |
| 696 | cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows, |
| 697 | indices, source_row, k, start_expert, end_expert, |
| 698 | renormalize, routed_scaling_factor, correction_bias, |
| 699 | input_ids, tid2eid, is_padding, bias_vl, |
| 700 | image_sentinel_lo); |
| 701 | #else |
| 702 | kernel<<<num_blocks, block_dim, 0, stream>>>( |
| 703 | input, finished, output, num_rows, indices, source_row, k, start_expert, |
| 704 | end_expert, renormalize, routed_scaling_factor, correction_bias, |
| 705 | input_ids, tid2eid, is_padding, bias_vl, image_sentinel_lo); |
| 706 | #endif |
| 707 | }) |
| 708 | } |
| 709 | |
| 710 | #ifndef USE_ROCM |
| 711 | #define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \ |
| 712 | static_assert(WARP_SIZE == 32, \ |
| 713 | "Unsupported warp size. Only 32 is supported for CUDA"); \ |
| 714 | topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, WARP_SIZE, \ |
| 715 | MAX_BYTES>( \ |
| 716 | gating_output, nullptr, topk_weights, topk_indices, \ |
| 717 | token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ |
| 718 | routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid, \ |
| 719 | stream, is_padding, bias_vl, image_sentinel_lo); |
| 720 | #else |
| 721 | #define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \ |
| 722 | if (WARP_SIZE == 64) { \ |
| 723 | topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 64, \ |
| 724 | MAX_BYTES>( \ |
| 725 | gating_output, nullptr, topk_weights, topk_indices, \ |
| 726 | token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ |
| 727 | routed_scaling_factor, correction_bias, use_hash, input_ids, \ |
| 728 | tid2eid, stream, is_padding, bias_vl, image_sentinel_lo); \ |
| 729 | } else if (WARP_SIZE == 32) { \ |
| 730 | topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, \ |
| 731 | MAX_BYTES>( \ |
| 732 | gating_output, nullptr, topk_weights, topk_indices, \ |
| 733 | token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ |
| 734 | routed_scaling_factor, correction_bias, use_hash, input_ids, \ |
| 735 | tid2eid, stream, is_padding, bias_vl, image_sentinel_lo); \ |
| 736 | } else { \ |
| 737 | assert(false && \ |
| 738 | "Unsupported warp size. Only 32 and 64 are supported for ROCm"); \ |
| 739 | } |
| 740 | #endif |
| 741 | |
| 742 | template <typename IndType, typename InputType, typename HashIndType = IndType> |
| 743 | void topkGatingSoftplusSqrtKernelLauncher( |
| 744 | const InputType* gating_output, float* topk_weights, IndType* topk_indices, |
| 745 | int* token_expert_indices, const int num_tokens, const int num_experts, |
| 746 | const int topk, const bool renormalize, double routed_scaling_factor, |
| 747 | const float* correction_bias, const bool use_hash, |
| 748 | const HashIndType* input_ids, const HashIndType* tid2eid, |
| 749 | cudaStream_t stream, const bool* is_padding, const float* bias_vl, |
| 750 | int64_t image_sentinel_lo) { |
| 751 | #ifndef USE_ROCM |
| 752 | if constexpr (std::is_same_v<InputType, float>) { |
| 753 | if (use_hash && topk == 6 && renormalize && |
| 754 | (num_experts == 256 || num_experts == 384)) { |
| 755 | launchDsv4HashTopk<IndType, HashIndType>( |
| 756 | gating_output, topk_weights, topk_indices, num_tokens, num_experts, |
| 757 | routed_scaling_factor, input_ids, tid2eid, stream, is_padding, |
| 758 | bias_vl, image_sentinel_lo); |
| 759 | return; |
| 760 | } |
| 761 | } |
| 762 | #endif |
| 763 | static constexpr int WARPS_PER_TB = 4; |
| 764 | static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16; |
| 765 | // for bfloat16 dtype, we need 4 bytes loading to make sure num_experts |
| 766 | // elements can be loaded by a warp |
| 767 | static constexpr int BYTES_PER_LDG_MULTIPLE_64 = |
| 768 | (std::is_same_v<InputType, __nv_bfloat16> || |
| 769 | std::is_same_v<InputType, __half>) |
| 770 | ? 4 |
| 771 | : 8; |
| 772 | // Narrower LDG (ELTS_PER_LDG=1) used by 192/320/448/576 on ROCm WARP_SIZE=64 |
| 773 | // where ELTS_PER_LDG=2 fails the EXPERTS%(ELTS_PER_LDG*WARP_SIZE)==0 check. |
| 774 | // On CUDA WARP_SIZE=32 the wider LDG already aligns, so the alias collapses |
| 775 | // back to BYTES_PER_LDG_MULTIPLE_64 — no behavioral change for CUDA. |
| 776 | #ifdef USE_ROCM |
| 777 | static constexpr int BYTES_PER_LDG_MULTIPLE_64_NARROW = |
| 778 | (std::is_same_v<InputType, __nv_bfloat16> || |
| 779 | std::is_same_v<InputType, __half>) |
| 780 | ? 2 |
| 781 | : 4; |
| 782 | #else |
| 783 | static constexpr int BYTES_PER_LDG_MULTIPLE_64_NARROW = |
| 784 | BYTES_PER_LDG_MULTIPLE_64; |
| 785 | #endif |
| 786 | switch (num_experts) { |
| 787 | case 1: |
| 788 | LAUNCH_SOFTPLUS_SQRT(1, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); |
| 789 | break; |
| 790 | case 2: |
| 791 | LAUNCH_SOFTPLUS_SQRT(2, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); |
| 792 | break; |
| 793 | case 4: |
| 794 | LAUNCH_SOFTPLUS_SQRT(4, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); |
| 795 | break; |
| 796 | case 8: |
| 797 | LAUNCH_SOFTPLUS_SQRT(8, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); |
| 798 | break; |
| 799 | case 16: |
| 800 | LAUNCH_SOFTPLUS_SQRT(16, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); |
| 801 | break; |
| 802 | case 32: |
| 803 | LAUNCH_SOFTPLUS_SQRT(32, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); |
| 804 | break; |
| 805 | case 64: |
| 806 | LAUNCH_SOFTPLUS_SQRT(64, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); |
| 807 | break; |
| 808 | case 128: |
| 809 | LAUNCH_SOFTPLUS_SQRT(128, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); |
| 810 | break; |
| 811 | case 256: |
| 812 | LAUNCH_SOFTPLUS_SQRT(256, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); |
| 813 | break; |
| 814 | case 512: |
| 815 | LAUNCH_SOFTPLUS_SQRT(512, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); |
| 816 | break; |
| 817 | // Multiples of 64 that are not powers of 2. The kernel requires |
| 818 | // EXPERTS % (ELTS_PER_LDG * WARP_SIZE) == 0. With ELTS_PER_LDG=2 |
| 819 | // (BYTES_PER_LDG_MULTIPLE_64), this holds for all five values on CUDA |
| 820 | // WARP_SIZE=32 but only for 384 on ROCm WARP_SIZE=64. The other four |
| 821 | // use BYTES_PER_LDG_MULTIPLE_64_NARROW (ELTS_PER_LDG=1), which |
| 822 | // satisfies the assertion for any multiple of 64 on either backend; |
| 823 | // on CUDA the narrow alias collapses back to the wider load, so CUDA |
| 824 | // behavior is unchanged. |
| 825 | case 192: |
| 826 | LAUNCH_SOFTPLUS_SQRT(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW); |
| 827 | break; |
| 828 | case 320: |
| 829 | LAUNCH_SOFTPLUS_SQRT(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW); |
| 830 | break; |
| 831 | case 384: |
| 832 | LAUNCH_SOFTPLUS_SQRT(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64); |
| 833 | break; |
| 834 | case 448: |
| 835 | LAUNCH_SOFTPLUS_SQRT(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW); |
| 836 | break; |
| 837 | case 576: |
| 838 | LAUNCH_SOFTPLUS_SQRT(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW); |
| 839 | break; |
| 840 | default: { |
| 841 | STD_TORCH_CHECK(false, "Unsupported expert number: ", num_experts); |
| 842 | } |
| 843 | } |
| 844 | } |
| 845 | |
| 846 | } // namespace moe |
| 847 | } // namespace vllm |
| 848 | |
| 849 | template <typename ComputeType> |
| 850 | void dispatch_topk_softplus_sqrt_launch( |
| 851 | const ComputeType* gating_output, torch::stable::Tensor& topk_weights, |
| 852 | torch::stable::Tensor& topk_indices, |
| 853 | torch::stable::Tensor& token_expert_indices, int num_tokens, |
| 854 | int num_experts, int topk, bool renormalize, double routed_scaling_factor, |
| 855 | const std::optional<torch::stable::Tensor>& correction_bias, |
| 856 | const std::optional<torch::stable::Tensor>& input_ids, |
| 857 | const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream, |
| 858 | const std::optional<torch::stable::Tensor>& is_padding, |
| 859 | const std::optional<torch::stable::Tensor>& bias_vl, |
| 860 | int64_t image_sentinel_lo) { |
| 861 | const float* bias_ptr = nullptr; |
| 862 | if (correction_bias.has_value()) { |
| 863 | bias_ptr = correction_bias.value().const_data_ptr<float>(); |
| 864 | } |
| 865 | |
| 866 | const float* bias_vl_ptr = nullptr; |
| 867 | if (bias_vl.has_value()) { |
| 868 | const torch::stable::Tensor& bias_vl_tensor = bias_vl.value(); |
| 869 | STD_TORCH_CHECK(input_ids.has_value(), |
| 870 | "input_ids is required when bias_vl is set"); |
| 871 | STD_TORCH_CHECK( |
| 872 | bias_vl_tensor.scalar_type() == torch::headeronly::ScalarType::Float, |
| 873 | "bias_vl tensor must be float32"); |
| 874 | STD_TORCH_CHECK(bias_vl_tensor.dim() == 1, "bias_vl tensor must be 1D"); |
| 875 | STD_TORCH_CHECK(bias_vl_tensor.size(0) == num_experts, |
| 876 | "bias_vl size mismatch, expected: ", num_experts); |
| 877 | STD_TORCH_CHECK(bias_vl_tensor.is_contiguous(), |
| 878 | "bias_vl tensor must be contiguous"); |
| 879 | bias_vl_ptr = bias_vl_tensor.const_data_ptr<float>(); |
| 880 | } |
| 881 | |
| 882 | auto launch = [&](auto* topk_indices_ptr) { |
| 883 | using OutIndType = |
| 884 | typename std::remove_pointer<decltype(topk_indices_ptr)>::type; |
| 885 | |
| 886 | const bool* is_padding_ptr = nullptr; |
| 887 | if (is_padding.has_value()) { |
| 888 | const torch::stable::Tensor& is_padding_tensor = is_padding.value(); |
| 889 | STD_TORCH_CHECK(is_padding_tensor.scalar_type() == |
| 890 | torch::headeronly::ScalarType::Bool, |
| 891 | "is_padding tensor must be bool"); |
| 892 | STD_TORCH_CHECK(is_padding_tensor.dim() == 1, |
| 893 | "is_padding tensor must be 1D"); |
| 894 | STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens, |
| 895 | "is_padding size mismatch, expected: ", num_tokens); |
| 896 | STD_TORCH_CHECK(is_padding_tensor.is_contiguous(), |
| 897 | "is_padding tensor must be contiguous"); |
| 898 | is_padding_ptr = is_padding_tensor.const_data_ptr<bool>(); |
| 899 | } |
| 900 | |
| 901 | if (tid2eid.has_value()) { |
| 902 | STD_TORCH_CHECK(input_ids.has_value(), |
| 903 | "input_ids is required for hash MoE"); |
| 904 | STD_TORCH_CHECK( |
| 905 | input_ids.value().scalar_type() == tid2eid.value().scalar_type(), |
| 906 | "input_ids and tid2eid must have the same dtype"); |
| 907 | if (tid2eid.value().scalar_type() == |
| 908 | torch::headeronly::ScalarType::Long) { |
| 909 | vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType, |
| 910 | int64_t>( |
| 911 | gating_output, topk_weights.mutable_data_ptr<float>(), |
| 912 | topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(), |
| 913 | num_tokens, num_experts, topk, renormalize, routed_scaling_factor, |
| 914 | bias_ptr, true, input_ids.value().const_data_ptr<int64_t>(), |
| 915 | tid2eid.value().const_data_ptr<int64_t>(), stream, is_padding_ptr, |
| 916 | bias_vl_ptr, image_sentinel_lo); |
| 917 | } else { |
| 918 | STD_TORCH_CHECK(tid2eid.value().scalar_type() == |
| 919 | torch::headeronly::ScalarType::Int); |
| 920 | vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType, |
| 921 | int>( |
| 922 | gating_output, topk_weights.mutable_data_ptr<float>(), |
| 923 | topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(), |
| 924 | num_tokens, num_experts, topk, renormalize, routed_scaling_factor, |
| 925 | bias_ptr, true, input_ids.value().const_data_ptr<int>(), |
| 926 | tid2eid.value().const_data_ptr<int>(), stream, is_padding_ptr, |
| 927 | bias_vl_ptr, image_sentinel_lo); |
| 928 | } |
| 929 | } else if (bias_vl_ptr != nullptr) { |
| 930 | // Non-hash rows with bias_vl: image tokens (five sentinel ids starting |
| 931 | // at image_sentinel_lo) select experts by score + bias_vl. input_ids |
| 932 | // drives the per-row branch, so |
| 933 | // dispatch the hash-index type on its dtype. |
| 934 | if (input_ids.value().scalar_type() == |
| 935 | torch::headeronly::ScalarType::Long) { |
| 936 | vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType, |
| 937 | int64_t>( |
| 938 | gating_output, topk_weights.mutable_data_ptr<float>(), |
| 939 | topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(), |
| 940 | num_tokens, num_experts, topk, renormalize, routed_scaling_factor, |
| 941 | bias_ptr, false, input_ids.value().const_data_ptr<int64_t>(), |
| 942 | static_cast<const int64_t*>(nullptr), stream, is_padding_ptr, |
| 943 | bias_vl_ptr, image_sentinel_lo); |
| 944 | } else { |
| 945 | STD_TORCH_CHECK(input_ids.value().scalar_type() == |
| 946 | torch::headeronly::ScalarType::Int); |
| 947 | vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType, |
| 948 | int>( |
| 949 | gating_output, topk_weights.mutable_data_ptr<float>(), |
| 950 | topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(), |
| 951 | num_tokens, num_experts, topk, renormalize, routed_scaling_factor, |
| 952 | bias_ptr, false, input_ids.value().const_data_ptr<int>(), |
| 953 | static_cast<const int*>(nullptr), stream, is_padding_ptr, |
| 954 | bias_vl_ptr, image_sentinel_lo); |
| 955 | } |
| 956 | } else { |
| 957 | vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType>( |
| 958 | gating_output, topk_weights.mutable_data_ptr<float>(), |
| 959 | topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(), |
| 960 | num_tokens, num_experts, topk, renormalize, routed_scaling_factor, |
| 961 | bias_ptr, false, static_cast<const OutIndType*>(nullptr), |
| 962 | static_cast<const OutIndType*>(nullptr), stream, is_padding_ptr, |
| 963 | nullptr, 0); |
| 964 | } |
| 965 | }; |
| 966 | |
| 967 | if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) { |
| 968 | launch(topk_indices.mutable_data_ptr<int>()); |
| 969 | } else if (topk_indices.scalar_type() == |
| 970 | torch::headeronly::ScalarType::UInt32) { |
| 971 | launch(topk_indices.mutable_data_ptr<uint32_t>()); |
| 972 | } else { |
| 973 | STD_TORCH_CHECK(topk_indices.scalar_type() == |
| 974 | torch::headeronly::ScalarType::Long); |
| 975 | launch(topk_indices.mutable_data_ptr<int64_t>()); |
| 976 | } |
| 977 | } |
| 978 | |
| 979 | void topk_softplus_sqrt( |
| 980 | torch::stable::Tensor& topk_weights, // [num_tokens, topk] |
| 981 | torch::stable::Tensor& topk_indices, // [num_tokens, topk] |
| 982 | torch::stable::Tensor& token_expert_indices, // [num_tokens, topk] |
| 983 | torch::stable::Tensor& gating_output, // [num_tokens, num_experts] |
| 984 | bool renormalize, double routed_scaling_factor, |
| 985 | const std::optional<torch::stable::Tensor>& correction_bias, |
| 986 | const std::optional<torch::stable::Tensor>& input_ids, |
| 987 | const std::optional<torch::stable::Tensor>& tid2eid, |
| 988 | const std::optional<torch::stable::Tensor>& is_padding, |
| 989 | const std::optional<torch::stable::Tensor>& bias_vl, |
| 990 | int64_t image_sentinel_lo) { |
| 991 | const int num_experts = gating_output.size(-1); |
| 992 | const auto num_tokens = gating_output.numel() / num_experts; |
| 993 | const int topk = topk_weights.size(-1); |
| 994 | const torch::stable::accelerator::DeviceGuard guard( |
| 995 | gating_output.get_device_index()); |
| 996 | const cudaStream_t stream = |
| 997 | get_current_cuda_stream(gating_output.get_device_index()); |
| 998 | |
| 999 | if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) { |
| 1000 | dispatch_topk_softplus_sqrt_launch<float>( |
| 1001 | gating_output.const_data_ptr<float>(), topk_weights, topk_indices, |
| 1002 | token_expert_indices, num_tokens, num_experts, topk, renormalize, |
| 1003 | routed_scaling_factor, correction_bias, input_ids, tid2eid, stream, |
| 1004 | is_padding, bias_vl, image_sentinel_lo); |
| 1005 | } else if (gating_output.scalar_type() == |
| 1006 | torch::headeronly::ScalarType::Half) { |
| 1007 | dispatch_topk_softplus_sqrt_launch<__half>( |
| 1008 | reinterpret_cast<const __half*>(gating_output.const_data_ptr()), |
| 1009 | topk_weights, topk_indices, token_expert_indices, num_tokens, |
| 1010 | num_experts, topk, renormalize, routed_scaling_factor, correction_bias, |
| 1011 | input_ids, tid2eid, stream, is_padding, bias_vl, image_sentinel_lo); |
| 1012 | } else if (gating_output.scalar_type() == |
| 1013 | torch::headeronly::ScalarType::BFloat16) { |
| 1014 | dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>( |
| 1015 | reinterpret_cast<const __nv_bfloat16*>(gating_output.const_data_ptr()), |
| 1016 | topk_weights, topk_indices, token_expert_indices, num_tokens, |
| 1017 | num_experts, topk, renormalize, routed_scaling_factor, correction_bias, |
| 1018 | input_ids, tid2eid, stream, is_padding, bias_vl, image_sentinel_lo); |
| 1019 | } else { |
| 1020 | STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", |
| 1021 | gating_output.scalar_type()); |
| 1022 | } |
| 1023 | } |