| 1 | #include <array> |
| 2 | #include <cub/cub.cuh> |
| 3 | |
| 4 | #include <cuda_runtime.h> |
| 5 | #include <torch/csrc/stable/macros.h> |
| 6 | #include <torch/csrc/stable/accelerator.h> |
| 7 | #include <torch/csrc/stable/ops.h> |
| 8 | #include <torch/csrc/stable/tensor.h> |
| 9 | #include <torch/headeronly/core/ScalarType.h> |
| 10 | |
| 11 | #include "../../cuda_compat.h" |
| 12 | #include "libtorch_stable/core/math.hpp" |
| 13 | #include "libtorch_stable/dispatch_utils.h" |
| 14 | #include "libtorch_stable/quantization/vectorization.cuh" |
| 15 | #include "libtorch_stable/torch_utils.h" |
| 16 | |
| 17 | #define CEILDIV(x, y) (((x) + (y) - 1) / (y)) |
| 18 | |
| 19 | namespace vllm { |
| 20 | namespace moe { |
| 21 | namespace batched_moe_align_block_size { |
| 22 | |
| 23 | // Note num_threads needs to be 1024 for BlockScan Reduction in the kernel. |
| 24 | static constexpr int32_t num_threads = 1024; |
| 25 | static constexpr int32_t num_blocks = 1; |
| 26 | template <bool cooperative_writes> |
| 27 | __global__ void batched_moe_align_block_size_kernel( |
| 28 | int32_t const num_batches, int32_t const max_tokens_per_batch, |
| 29 | int32_t const block_size, int32_t const* __restrict__ batch_num_tokens, |
| 30 | int32_t* __restrict__ sorted_ids, int32_t* __restrict__ block_ids, |
| 31 | int32_t* __restrict__ num_tokens_post_pad) { |
| 32 | size_t const stride = blockDim.x * gridDim.x; |
| 33 | int32_t const num_blocks_per_batch = |
| 34 | CEILDIV(max_tokens_per_batch, block_size); |
| 35 | size_t const sorted_ids_size = |
| 36 | static_cast<size_t>(num_blocks_per_batch) * num_batches * block_size; |
| 37 | size_t const block_ids_size = sorted_ids_size / block_size; |
| 38 | int32_t const SENTINEL = |
| 39 | num_batches * max_tokens_per_batch; // To denote invalid entries. |
| 40 | // Initialize sorted_ids |
| 41 | for (size_t i = threadIdx.x; i < sorted_ids_size; i += stride) { |
| 42 | sorted_ids[i] = SENTINEL; |
| 43 | } |
| 44 | // Initialize expert_ids with -1 |
| 45 | for (size_t i = threadIdx.x; i < block_ids_size; i += stride) { |
| 46 | block_ids[i] = -1; |
| 47 | } |
| 48 | |
| 49 | int32_t b_num_tokens = 0; |
| 50 | if (threadIdx.x < num_batches) { |
| 51 | b_num_tokens = batch_num_tokens[threadIdx.x]; |
| 52 | } |
| 53 | int32_t const ceil_b_num_tokens = |
| 54 | CEILDIV(b_num_tokens, block_size) * block_size; |
| 55 | |
| 56 | // Compute prefix sum over token counts per expert |
| 57 | using BlockScan = cub::BlockScan<int32_t, 1024>; |
| 58 | __shared__ typename BlockScan::TempStorage temp_storage; |
| 59 | int cumsum_val; |
| 60 | BlockScan(temp_storage).ExclusiveSum(ceil_b_num_tokens, cumsum_val); |
| 61 | __syncthreads(); |
| 62 | |
| 63 | if (threadIdx.x == num_batches - 1) { |
| 64 | *num_tokens_post_pad = cumsum_val + ceil_b_num_tokens; |
| 65 | } |
| 66 | |
| 67 | if constexpr (cooperative_writes) { |
| 68 | __shared__ int32_t batch_cumsum[num_threads]; |
| 69 | __shared__ int32_t valid_tokens[num_threads]; |
| 70 | __shared__ int32_t batch_num_blocks[num_threads]; |
| 71 | if (threadIdx.x < num_batches) { |
| 72 | batch_cumsum[threadIdx.x] = cumsum_val; |
| 73 | valid_tokens[threadIdx.x] = b_num_tokens; |
| 74 | batch_num_blocks[threadIdx.x] = ceil_b_num_tokens / block_size; |
| 75 | } |
| 76 | __syncthreads(); |
| 77 | |
| 78 | int32_t const max_num_groups = blockDim.x / WARP_SIZE; |
| 79 | int32_t num_groups = 1; |
| 80 | while (num_groups < num_batches && num_groups < max_num_groups) { |
| 81 | num_groups *= 2; |
| 82 | } |
| 83 | int32_t const threads_per_batch = blockDim.x / num_groups; |
| 84 | int32_t const group_id = threadIdx.x / threads_per_batch; |
| 85 | int32_t const group_offset = threadIdx.x % threads_per_batch; |
| 86 | |
| 87 | // Assign at least one warp to each batch when possible. |
| 88 | for (int32_t batch_id = group_id; batch_id < num_batches; |
| 89 | batch_id += num_groups) { |
| 90 | size_t const batch_offset = |
| 91 | static_cast<size_t>(batch_id) * max_tokens_per_batch; |
| 92 | size_t const cumsum = batch_cumsum[batch_id]; |
| 93 | for (size_t i = group_offset; i < valid_tokens[batch_id]; |
| 94 | i += threads_per_batch) { |
| 95 | sorted_ids[cumsum + i] = static_cast<int32_t>(batch_offset + i); |
| 96 | } |
| 97 | |
| 98 | size_t const block_start = cumsum / block_size; |
| 99 | for (size_t i = group_offset; i < batch_num_blocks[batch_id]; |
| 100 | i += threads_per_batch) { |
| 101 | block_ids[block_start + i] = batch_id; |
| 102 | } |
| 103 | } |
| 104 | } else if (threadIdx.x < num_batches) { |
| 105 | size_t const batch_id = threadIdx.x; |
| 106 | size_t const batch_offset = batch_id * max_tokens_per_batch; |
| 107 | for (size_t i = 0; i < b_num_tokens; ++i) { |
| 108 | sorted_ids[cumsum_val + i] = static_cast<int32_t>(batch_offset + i); |
| 109 | } |
| 110 | |
| 111 | size_t const block_start = cumsum_val / block_size; |
| 112 | for (size_t i = 0; i < ceil_b_num_tokens / block_size; ++i) { |
| 113 | block_ids[block_start + i] = batch_id; |
| 114 | } |
| 115 | } |
| 116 | } |
| 117 | } // namespace batched_moe_align_block_size |
| 118 | |
| 119 | template <typename scalar_t> |
| 120 | __device__ __forceinline__ int get_local_expert_id( |
| 121 | size_t idx, const scalar_t* __restrict__ topk_ids, |
| 122 | int32_t* __restrict__ expert_map, int32_t num_experts, |
| 123 | bool has_expert_map) { |
| 124 | int expert_id = topk_ids[idx]; |
| 125 | if (expert_id >= num_experts || expert_id < 0) { |
| 126 | return -1; |
| 127 | } |
| 128 | if (has_expert_map) { |
| 129 | expert_id = expert_map[expert_id]; |
| 130 | } |
| 131 | return expert_id; |
| 132 | } |
| 133 | |
| 134 | template <typename scalar_t> |
| 135 | __device__ void _moe_align_block_size( |
| 136 | const scalar_t* __restrict__ topk_ids, |
| 137 | int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids, |
| 138 | int32_t* __restrict__ total_tokens_post_pad, |
| 139 | int32_t* __restrict__ expert_map, int32_t num_experts, |
| 140 | int32_t padded_num_experts, int32_t experts_per_warp, int32_t block_size, |
| 141 | size_t numel, int32_t* __restrict__ cumsum, int32_t max_num_tokens_padded, |
| 142 | int32_t max_num_m_blocks, int32_t model_offset, int32_t inactive_expert_id, |
| 143 | int32_t topk_num, int32_t* token_mask, bool has_expert_map) { |
| 144 | extern __shared__ int32_t shared_counts[]; |
| 145 | |
| 146 | // Compute input buffer offsets. Typically these will all be 0, except when |
| 147 | // using Multi LoRA. |
| 148 | int sorted_token_ids_offset = max_num_tokens_padded * model_offset; |
| 149 | int expert_ids_offset = max_num_m_blocks * model_offset; |
| 150 | int cumsum_offset = (num_experts + 1) * model_offset; |
| 151 | |
| 152 | // Use separate threadblocks to fill sorted_token_ids. |
| 153 | // This is safe since the current kernel does not use sorted_token_ids. |
| 154 | if (blockIdx.x % 2) { |
| 155 | // Initialize sorted_token_ids with numel |
| 156 | for (size_t it = threadIdx.x; it < max_num_tokens_padded; |
| 157 | it += blockDim.x) { |
| 158 | sorted_token_ids[sorted_token_ids_offset + it] = numel; |
| 159 | } |
| 160 | return; |
| 161 | } |
| 162 | |
| 163 | const int warp_id = threadIdx.x / WARP_SIZE; |
| 164 | const int my_expert_start = warp_id * experts_per_warp; |
| 165 | |
| 166 | for (int i = 0; i < experts_per_warp; ++i) { |
| 167 | if (my_expert_start + i < padded_num_experts) { |
| 168 | shared_counts[warp_id * experts_per_warp + i] = 0; |
| 169 | } |
| 170 | } |
| 171 | |
| 172 | __syncthreads(); |
| 173 | |
| 174 | const size_t tid = threadIdx.x; |
| 175 | const size_t stride = blockDim.x; |
| 176 | |
| 177 | for (size_t i = tid; i < numel; i += stride) { |
| 178 | if (int expert_id = get_local_expert_id(i, topk_ids, expert_map, |
| 179 | num_experts, has_expert_map); |
| 180 | expert_id != -1) { |
| 181 | int warp_idx = expert_id / experts_per_warp; |
| 182 | int expert_offset = expert_id % experts_per_warp; |
| 183 | int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num]; |
| 184 | atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset], |
| 185 | mask); |
| 186 | } |
| 187 | } |
| 188 | |
| 189 | __syncthreads(); |
| 190 | |
| 191 | // Compute prefix sum over token counts per expert |
| 192 | using BlockScan = cub::BlockScan<int32_t, 1024>; |
| 193 | __shared__ typename BlockScan::TempStorage temp_storage; |
| 194 | |
| 195 | int expert_count = 0; |
| 196 | int expert_id = threadIdx.x; |
| 197 | if (expert_id < num_experts) { |
| 198 | int warp_idx = expert_id / experts_per_warp; |
| 199 | int expert_offset = expert_id % experts_per_warp; |
| 200 | expert_count = shared_counts[warp_idx * experts_per_warp + expert_offset]; |
| 201 | expert_count = CEILDIV(expert_count, block_size) * block_size; |
| 202 | } |
| 203 | |
| 204 | int cumsum_val; |
| 205 | BlockScan(temp_storage).ExclusiveSum(expert_count, cumsum_val); |
| 206 | if (expert_id <= num_experts) { |
| 207 | cumsum[cumsum_offset + expert_id] = cumsum_val; |
| 208 | } |
| 209 | |
| 210 | if (expert_id == num_experts) { |
| 211 | total_tokens_post_pad[model_offset] = cumsum_val; |
| 212 | } |
| 213 | |
| 214 | __syncthreads(); |
| 215 | |
| 216 | if (threadIdx.x < num_experts) { |
| 217 | for (int i = cumsum[cumsum_offset + threadIdx.x]; |
| 218 | i < cumsum[cumsum_offset + threadIdx.x + 1]; i += block_size) { |
| 219 | expert_ids[expert_ids_offset + i / block_size] = threadIdx.x; |
| 220 | } |
| 221 | } |
| 222 | |
| 223 | // Fill remaining expert_ids with -1 |
| 224 | const size_t fill_start_idx = |
| 225 | cumsum[cumsum_offset + num_experts] / block_size + threadIdx.x; |
| 226 | for (size_t i = fill_start_idx; i < max_num_m_blocks; i += blockDim.x) { |
| 227 | expert_ids[expert_ids_offset + i] = inactive_expert_id; |
| 228 | } |
| 229 | } |
| 230 | |
| 231 | template <typename scalar_t, int32_t fill_threads> |
| 232 | __device__ void _moe_align_block_size_small_batch_expert( |
| 233 | const scalar_t* __restrict__ topk_ids, |
| 234 | int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids, |
| 235 | int32_t* __restrict__ total_tokens_post_pad, |
| 236 | int32_t* __restrict__ expert_map, int32_t num_experts, int32_t block_size, |
| 237 | size_t numel, int32_t max_num_tokens_padded, int32_t max_num_m_blocks, |
| 238 | int32_t inactive_expert_id, int32_t model_offset, int32_t topk_num, |
| 239 | int32_t* token_mask, bool has_expert_map, int32_t* scatter_idx = nullptr) { |
| 240 | // Compute input buffer offsets. Typically these will all be 0, except when |
| 241 | // using Multi LoRA. |
| 242 | int sorted_token_ids_offset = max_num_tokens_padded * model_offset; |
| 243 | int expert_ids_offset = max_num_m_blocks * model_offset; |
| 244 | |
| 245 | // Use an additional group of threads to fill sorted_token_ids. |
| 246 | // Since the current kernel will use sorted_token_ids afterward, |
| 247 | // we fill sorted_token_ids within the same threadblock to make |
| 248 | // synchronization easier. |
| 249 | if (threadIdx.x < fill_threads) { |
| 250 | // Initialize sorted_token_ids with numel |
| 251 | for (size_t it = threadIdx.x; it < max_num_tokens_padded; |
| 252 | it += fill_threads) { |
| 253 | sorted_token_ids[sorted_token_ids_offset + it] = numel; |
| 254 | } |
| 255 | // Three __syncthreads() corresponding to the other threads |
| 256 | __syncthreads(); |
| 257 | __syncthreads(); |
| 258 | __syncthreads(); |
| 259 | return; |
| 260 | } |
| 261 | |
| 262 | const size_t tid = threadIdx.x - fill_threads; |
| 263 | const size_t stride = blockDim.x - fill_threads; |
| 264 | |
| 265 | extern __shared__ int32_t shared_mem[]; |
| 266 | int32_t* cumsum = shared_mem; |
| 267 | int32_t* tokens_cnts = (int32_t*)(shared_mem + num_experts + 1); |
| 268 | |
| 269 | for (int i = 0; i < num_experts; ++i) { |
| 270 | tokens_cnts[(tid + 1) * num_experts + i] = 0; |
| 271 | } |
| 272 | |
| 273 | for (size_t i = tid; i < numel; i += stride) { |
| 274 | if (int expert_id = get_local_expert_id(i, topk_ids, expert_map, |
| 275 | num_experts, has_expert_map); |
| 276 | expert_id != -1) { |
| 277 | int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num]; |
| 278 | tokens_cnts[(tid + 1) * num_experts + expert_id] += mask; |
| 279 | } |
| 280 | } |
| 281 | |
| 282 | __syncthreads(); |
| 283 | |
| 284 | if (tid < num_experts) { |
| 285 | tokens_cnts[tid] = 0; |
| 286 | for (int i = 1; i <= stride; ++i) { |
| 287 | tokens_cnts[i * num_experts + tid] += |
| 288 | tokens_cnts[(i - 1) * num_experts + tid]; |
| 289 | } |
| 290 | } |
| 291 | |
| 292 | __syncthreads(); |
| 293 | |
| 294 | if (tid == 0) { |
| 295 | cumsum[0] = 0; |
| 296 | for (int i = 1; i <= num_experts; ++i) { |
| 297 | cumsum[i] = |
| 298 | cumsum[i - 1] + |
| 299 | CEILDIV(tokens_cnts[stride * num_experts + i - 1], block_size) * |
| 300 | block_size; |
| 301 | } |
| 302 | total_tokens_post_pad[model_offset] = |
| 303 | static_cast<int32_t>(cumsum[num_experts]); |
| 304 | } |
| 305 | |
| 306 | __syncthreads(); |
| 307 | |
| 308 | if (tid < num_experts) { |
| 309 | for (int i = cumsum[tid]; i < cumsum[tid + 1]; i += block_size) { |
| 310 | expert_ids[expert_ids_offset + i / block_size] = tid; |
| 311 | } |
| 312 | } |
| 313 | |
| 314 | // Fill remaining expert_ids with -1 |
| 315 | const size_t fill_start_idx = cumsum[num_experts] / block_size + tid; |
| 316 | for (size_t i = fill_start_idx; i < max_num_m_blocks; i += stride) { |
| 317 | expert_ids[expert_ids_offset + i] = inactive_expert_id; |
| 318 | } |
| 319 | |
| 320 | for (size_t i = tid; i < numel; i += stride) { |
| 321 | int expert_id = get_local_expert_id(i, topk_ids, expert_map, num_experts, |
| 322 | has_expert_map); |
| 323 | bool valid = |
| 324 | expert_id != -1 && (token_mask == nullptr || token_mask[i / topk_num]); |
| 325 | if (scatter_idx != nullptr) { |
| 326 | scatter_idx[i] = valid ? static_cast<int32_t>(i) : -1; |
| 327 | } |
| 328 | if (expert_id != -1) { |
| 329 | int32_t rank_post_pad = |
| 330 | tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id]; |
| 331 | |
| 332 | if (token_mask == nullptr || token_mask[i / topk_num]) { |
| 333 | sorted_token_ids[sorted_token_ids_offset + rank_post_pad] = i; |
| 334 | ++tokens_cnts[tid * num_experts + expert_id]; |
| 335 | } |
| 336 | } |
| 337 | } |
| 338 | } |
| 339 | |
| 340 | template <typename scalar_t> |
| 341 | __device__ void _count_and_sort_expert_tokens( |
| 342 | const scalar_t* __restrict__ topk_ids, |
| 343 | int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer, |
| 344 | int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts, |
| 345 | int32_t max_num_tokens_padded, int32_t* __restrict__ token_mask, |
| 346 | int32_t model_offset, int32_t topk_num, bool has_expert_map, |
| 347 | int32_t* scatter_idx = nullptr) { |
| 348 | const size_t tid = blockIdx.y * blockDim.x + threadIdx.x; |
| 349 | const size_t stride = blockDim.x * gridDim.y; |
| 350 | |
| 351 | for (size_t i = tid; i < numel; i += stride) { |
| 352 | int expert_id = get_local_expert_id(i, topk_ids, expert_map, num_experts, |
| 353 | has_expert_map); |
| 354 | bool valid = |
| 355 | expert_id != -1 && (token_mask == nullptr || token_mask[i / topk_num]); |
| 356 | if (scatter_idx != nullptr) { |
| 357 | scatter_idx[i] = valid ? static_cast<int32_t>(i) : -1; |
| 358 | } |
| 359 | if (expert_id != -1) { |
| 360 | if (token_mask == nullptr || token_mask[i / topk_num]) { |
| 361 | int32_t rank_post_pad = atomicAdd( |
| 362 | &cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1); |
| 363 | sorted_token_ids[max_num_tokens_padded * model_offset + rank_post_pad] = |
| 364 | i; |
| 365 | } |
| 366 | } |
| 367 | } |
| 368 | } |
| 369 | |
| 370 | template <typename scalar_t> |
| 371 | __global__ void moe_align_block_size_kernel( |
| 372 | const scalar_t* __restrict__ topk_ids, |
| 373 | int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids, |
| 374 | int32_t* __restrict__ total_tokens_post_pad, |
| 375 | int32_t* __restrict__ expert_map, int32_t num_experts, |
| 376 | int32_t padded_num_experts, int32_t experts_per_warp, int32_t block_size, |
| 377 | size_t numel, int32_t* __restrict__ cumsum, int32_t max_num_tokens_padded, |
| 378 | int32_t topk_num, bool has_expert_map) { |
| 379 | _moe_align_block_size( |
| 380 | topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map, |
| 381 | num_experts, padded_num_experts, experts_per_warp, block_size, numel, |
| 382 | cumsum, max_num_tokens_padded, CEILDIV(max_num_tokens_padded, block_size), |
| 383 | 0, -1, topk_num, nullptr, has_expert_map); |
| 384 | } |
| 385 | |
| 386 | template <typename scalar_t> |
| 387 | __global__ void count_and_sort_expert_tokens_kernel( |
| 388 | const scalar_t* __restrict__ topk_ids, |
| 389 | int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer, |
| 390 | int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts, |
| 391 | int32_t max_num_tokens_padded, int32_t topk_num, bool has_expert_map, |
| 392 | int32_t* scatter_idx) { |
| 393 | _count_and_sort_expert_tokens( |
| 394 | topk_ids, sorted_token_ids, cumsum_buffer, expert_map, numel, num_experts, |
| 395 | max_num_tokens_padded, nullptr, 0, topk_num, has_expert_map, scatter_idx); |
| 396 | } |
| 397 | |
| 398 | // Reduce the topk expert outputs per token (summed in fp32). The output is |
| 399 | // dense [num_tokens, d]; the input is addressed by its strides so non- |
| 400 | // contiguous inputs work without a copy. A 16B-vectorized path is used when |
| 401 | // the hidden dim is contiguous (innermost stride 1) and aligned; otherwise a |
| 402 | // scalar kernel reads via arbitrary strides. topk is a compile-time constant |
| 403 | // for common values and runtime otherwise. |
| 404 | |
| 405 | // Elements per 16-byte vector (8 for bf16/fp16, 4 for fp32). |
| 406 | template <typename scalar_t> |
| 407 | constexpr int MOE_SUM_VEC = 16 / sizeof(scalar_t); |
| 408 | |
| 409 | template <typename idx_t> |
| 410 | __device__ __forceinline__ bool moe_sum_pad_aware_skip( |
| 411 | const idx_t* __restrict__ topk_ids, const int32_t* __restrict__ expert_map, |
| 412 | int64_t idx) { |
| 413 | int64_t expert_id = static_cast<int64_t>(topk_ids[idx]); |
| 414 | if (expert_id < 0) return true; |
| 415 | if (expert_map != nullptr && expert_map[expert_id] < 0) return true; |
| 416 | return false; |
| 417 | } |
| 418 | |
| 419 | template <typename scalar_t, typename idx_t, int TOPK, bool PAD_AWARE> |
| 420 | __global__ void moe_sum_vec_kernel( |
| 421 | scalar_t* __restrict__ out, // [num_tokens, d], contiguous |
| 422 | const scalar_t* __restrict__ input, // [num_tokens, topk, d], d contiguous |
| 423 | const int64_t num_tokens, const int d, const int64_t stride_token, |
| 424 | const int64_t stride_topk, const idx_t* __restrict__ topk_ids, |
| 425 | const int32_t* __restrict__ expert_map, const int64_t stride_tk_token, |
| 426 | const int64_t stride_tk_k) { |
| 427 | using vec_t = vllm::vec_n_t<scalar_t, MOE_SUM_VEC<scalar_t>>; // 16-byte pack |
| 428 | constexpr int VEC = MOE_SUM_VEC<scalar_t>; |
| 429 | const int64_t n_vec = d / VEC; |
| 430 | const int64_t total = num_tokens * n_vec; |
| 431 | for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < total; |
| 432 | i += (int64_t)gridDim.x * blockDim.x) { |
| 433 | const int64_t token = i / n_vec; |
| 434 | const int64_t v = i % n_vec; |
| 435 | const scalar_t* in_tok = input + token * stride_token + v * VEC; |
| 436 | const idx_t* tk_tok = nullptr; |
| 437 | if constexpr (PAD_AWARE) { |
| 438 | tk_tok = topk_ids + token * stride_tk_token; |
| 439 | } |
| 440 | |
| 441 | float acc[VEC]; |
| 442 | #pragma unroll |
| 443 | for (int j = 0; j < VEC; ++j) acc[j] = 0.f; |
| 444 | |
| 445 | #pragma unroll |
| 446 | for (int k = 0; k < TOPK; ++k) { |
| 447 | if constexpr (PAD_AWARE) { |
| 448 | if (moe_sum_pad_aware_skip(tk_tok, expert_map, k * stride_tk_k)) { |
| 449 | continue; |
| 450 | } |
| 451 | } |
| 452 | vec_t packed = *reinterpret_cast<const vec_t*>(in_tok + k * stride_topk); |
| 453 | #pragma unroll |
| 454 | for (int j = 0; j < VEC; ++j) acc[j] += static_cast<float>(packed.val[j]); |
| 455 | } |
| 456 | |
| 457 | vec_t outp; |
| 458 | #pragma unroll |
| 459 | for (int j = 0; j < VEC; ++j) outp.val[j] = static_cast<scalar_t>(acc[j]); |
| 460 | *reinterpret_cast<vec_t*>(out + token * d + v * VEC) = outp; |
| 461 | } |
| 462 | } |
| 463 | |
| 464 | // Runtime-topk variant of the above, for topk values outside the templated |
| 465 | // set. |
| 466 | template <typename scalar_t, typename idx_t, bool PAD_AWARE> |
| 467 | __global__ void moe_sum_vec_dynamic_kernel( |
| 468 | scalar_t* __restrict__ out, // [num_tokens, d], contiguous |
| 469 | const scalar_t* __restrict__ input, // [num_tokens, topk, d], d contiguous |
| 470 | const int64_t num_tokens, const int d, const int topk, |
| 471 | const int64_t stride_token, const int64_t stride_topk, |
| 472 | const idx_t* __restrict__ topk_ids, const int32_t* __restrict__ expert_map, |
| 473 | const int64_t stride_tk_token, const int64_t stride_tk_k) { |
| 474 | using vec_t = vllm::vec_n_t<scalar_t, MOE_SUM_VEC<scalar_t>>; |
| 475 | constexpr int VEC = MOE_SUM_VEC<scalar_t>; |
| 476 | const int64_t n_vec = d / VEC; |
| 477 | const int64_t total = num_tokens * n_vec; |
| 478 | for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < total; |
| 479 | i += (int64_t)gridDim.x * blockDim.x) { |
| 480 | const int64_t token = i / n_vec; |
| 481 | const int64_t v = i % n_vec; |
| 482 | const scalar_t* in_tok = input + token * stride_token + v * VEC; |
| 483 | const idx_t* tk_tok = nullptr; |
| 484 | if constexpr (PAD_AWARE) { |
| 485 | tk_tok = topk_ids + token * stride_tk_token; |
| 486 | } |
| 487 | |
| 488 | float acc[VEC]; |
| 489 | #pragma unroll |
| 490 | for (int j = 0; j < VEC; ++j) acc[j] = 0.f; |
| 491 | |
| 492 | for (int k = 0; k < topk; ++k) { |
| 493 | if constexpr (PAD_AWARE) { |
| 494 | if (moe_sum_pad_aware_skip(tk_tok, expert_map, k * stride_tk_k)) { |
| 495 | continue; |
| 496 | } |
| 497 | } |
| 498 | vec_t packed = *reinterpret_cast<const vec_t*>(in_tok + k * stride_topk); |
| 499 | #pragma unroll |
| 500 | for (int j = 0; j < VEC; ++j) acc[j] += static_cast<float>(packed.val[j]); |
| 501 | } |
| 502 | |
| 503 | vec_t outp; |
| 504 | #pragma unroll |
| 505 | for (int j = 0; j < VEC; ++j) outp.val[j] = static_cast<scalar_t>(acc[j]); |
| 506 | *reinterpret_cast<vec_t*>(out + token * d + v * VEC) = outp; |
| 507 | } |
| 508 | } |
| 509 | |
| 510 | // Stride-aware scalar fallback: handles unaligned/non-vectorizable hidden dims |
| 511 | // (including a non-contiguous hidden stride) via per-element strided reads. |
| 512 | template <typename scalar_t, typename idx_t, bool PAD_AWARE> |
| 513 | __global__ void moe_sum_scalar_kernel( |
| 514 | scalar_t* __restrict__ out, // [num_tokens, d], contiguous |
| 515 | const scalar_t* __restrict__ input, // [num_tokens, topk, d] |
| 516 | const int d, const int topk, const int64_t stride_token, |
| 517 | const int64_t stride_topk, const int64_t stride_hidden, |
| 518 | const idx_t* __restrict__ topk_ids, const int32_t* __restrict__ expert_map, |
| 519 | const int64_t stride_tk_token, const int64_t stride_tk_k) { |
| 520 | const int64_t token_idx = blockIdx.x; |
| 521 | const scalar_t* in_tok = input + token_idx * stride_token; |
| 522 | const idx_t* tk_tok = nullptr; |
| 523 | if constexpr (PAD_AWARE) { |
| 524 | tk_tok = topk_ids + token_idx * stride_tk_token; |
| 525 | } |
| 526 | for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) { |
| 527 | float x = 0.f; |
| 528 | for (int k = 0; k < topk; ++k) { |
| 529 | if constexpr (PAD_AWARE) { |
| 530 | if (moe_sum_pad_aware_skip(tk_tok, expert_map, k * stride_tk_k)) { |
| 531 | continue; |
| 532 | } |
| 533 | } |
| 534 | x += static_cast<float>( |
| 535 | VLLM_LDG(&in_tok[k * stride_topk + idx * stride_hidden])); |
| 536 | } |
| 537 | out[token_idx * d + idx] = static_cast<scalar_t>(x); |
| 538 | } |
| 539 | } |
| 540 | |
| 541 | template <typename scalar_t, int32_t fill_threads> |
| 542 | __global__ void moe_align_block_size_small_batch_expert_kernel( |
| 543 | const scalar_t* __restrict__ topk_ids, |
| 544 | int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids, |
| 545 | int32_t* __restrict__ total_tokens_post_pad, |
| 546 | int32_t* __restrict__ expert_map, int32_t num_experts, int32_t block_size, |
| 547 | size_t numel, int32_t max_num_tokens_padded, int32_t topk_num, |
| 548 | bool has_expert_map, int32_t* scatter_idx) { |
| 549 | _moe_align_block_size_small_batch_expert<scalar_t, fill_threads>( |
| 550 | topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map, |
| 551 | num_experts, block_size, numel, max_num_tokens_padded, |
| 552 | CEILDIV(max_num_tokens_padded, block_size), -1, 0, topk_num, nullptr, |
| 553 | has_expert_map, scatter_idx); |
| 554 | } |
| 555 | |
| 556 | template <typename scalar_t> |
| 557 | __global__ void moe_lora_align_block_size_kernel( |
| 558 | scalar_t* __restrict__ topk_ids, int32_t* __restrict__ token_lora_mapping, |
| 559 | int64_t block_size, int32_t* __restrict__ expert_map, int num_experts, |
| 560 | int max_loras, size_t numel, int max_num_tokens_padded, |
| 561 | int max_num_m_blocks, int32_t* __restrict__ sorted_token_ids, |
| 562 | int32_t* __restrict__ expert_ids, int32_t topk_num, |
| 563 | int32_t* total_tokens_post_pad, int32_t* adapter_enabled, |
| 564 | int32_t* __restrict__ cumsum, int32_t experts_per_warp, |
| 565 | int32_t padded_num_experts, int32_t* lora_ids, |
| 566 | int32_t* __restrict__ token_mask, bool has_expert_map) { |
| 567 | int lora_idx = blockIdx.x / 2; |
| 568 | int lora_id = lora_ids[lora_idx]; |
| 569 | // Output buffers are indexed by lora_id (in [0, max_loras)). The grid |
| 570 | // iterates one extra slot to accommodate the "-1" entry that |
| 571 | // active_lora_ids may hold in position 0 for mixed base + LoRA batches; |
| 572 | // guard against any other unexpected lora_id >= max_loras to avoid |
| 573 | // out-of-bounds writes. This mirrors the `lora_id >= max_loras` guard in |
| 574 | // the Triton _fused_moe_lora_kernel. |
| 575 | if (lora_id == -1 || lora_id >= max_loras || adapter_enabled[lora_id] == 0) { |
| 576 | return; |
| 577 | } |
| 578 | |
| 579 | // Populate the token_mask based on the token-LoRA mapping |
| 580 | int num_tokens = numel / topk_num; |
| 581 | // Only the even counting block owns per-LoRA metadata. The odd block only |
| 582 | // initializes sorted_token_ids and must not race the final count write. |
| 583 | if (blockIdx.x % 2 == 0 && threadIdx.x == 0) { |
| 584 | total_tokens_post_pad[lora_id] = 0; |
| 585 | |
| 586 | for (int i = 0; i < num_tokens; i++) { |
| 587 | token_mask[(lora_id * num_tokens) + i] = |
| 588 | (int)token_lora_mapping[i] == lora_id; |
| 589 | } |
| 590 | } |
| 591 | |
| 592 | __syncthreads(); |
| 593 | |
| 594 | _moe_align_block_size( |
| 595 | topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map, |
| 596 | num_experts, padded_num_experts, experts_per_warp, block_size, numel, |
| 597 | cumsum, max_num_tokens_padded, max_num_m_blocks, lora_id, -1, topk_num, |
| 598 | &token_mask[(lora_id * num_tokens)], has_expert_map); |
| 599 | } |
| 600 | |
| 601 | template <typename scalar_t> |
| 602 | __global__ void lora_count_and_sort_expert_tokens_kernel( |
| 603 | const scalar_t* __restrict__ topk_ids, |
| 604 | int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer, |
| 605 | int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts, |
| 606 | int32_t max_num_tokens_padded, int32_t topk_num, int32_t* token_mask, |
| 607 | int32_t max_loras, int32_t* lora_ids, int32_t* adapter_enabled, |
| 608 | bool has_expert_map) { |
| 609 | int lora_idx = blockIdx.x; |
| 610 | int lora_id = lora_ids[lora_idx]; |
| 611 | // Same guard rationale as moe_lora_align_block_size_kernel. Additionally |
| 612 | // skip disabled adapter slots: moe_lora_align_block_size_kernel early-returns |
| 613 | // for them and leaves token_mask[lora_id, :] uninitialized (token_mask is |
| 614 | // allocated with torch::empty), so running the sort loop here would traverse |
| 615 | // garbage mask bits and pollute this slot's rows of sorted_token_ids and |
| 616 | // cumsum_buffer. Downstream consumers already skip disabled slots, so the |
| 617 | // pollution is dormant today, but the check keeps behavior symmetric with |
| 618 | // the other two align kernels and avoids O(numel) wasted work per disabled |
| 619 | // slot. Short-circuit evaluation ensures adapter_enabled is only indexed |
| 620 | // after lora_id is confirmed to be in [0, max_loras). |
| 621 | if (lora_id == -1 || lora_id >= max_loras || adapter_enabled[lora_id] == 0) { |
| 622 | return; |
| 623 | } |
| 624 | |
| 625 | int num_tokens = numel / topk_num; |
| 626 | |
| 627 | _count_and_sort_expert_tokens( |
| 628 | topk_ids, sorted_token_ids, cumsum_buffer, expert_map, numel, num_experts, |
| 629 | max_num_tokens_padded, &token_mask[(lora_id * num_tokens)], lora_id, |
| 630 | topk_num, has_expert_map); |
| 631 | } |
| 632 | |
| 633 | template <typename scalar_t, int32_t fill_threads> |
| 634 | __global__ void moe_lora_align_block_size_small_batch_expert_kernel( |
| 635 | scalar_t* __restrict__ topk_ids, int32_t* token_lora_mapping, |
| 636 | int64_t block_size, int32_t* __restrict__ expert_map, int num_experts, |
| 637 | int max_loras, size_t numel, int max_num_tokens_padded, |
| 638 | int max_num_m_blocks, int32_t* __restrict__ sorted_token_ids, |
| 639 | int32_t* __restrict__ expert_ids, int topk_num, |
| 640 | int32_t* total_tokens_post_pad, int32_t* adapter_enabled, int32_t* lora_ids, |
| 641 | int32_t* token_mask, bool has_expert_map) { |
| 642 | int lora_idx = blockIdx.x; |
| 643 | int lora_id = lora_ids[lora_idx]; |
| 644 | // Same guard rationale as moe_lora_align_block_size_kernel. |
| 645 | if (lora_id == -1 || lora_id >= max_loras || adapter_enabled[lora_id] == 0) { |
| 646 | return; |
| 647 | } |
| 648 | |
| 649 | int num_tokens = numel / topk_num; |
| 650 | if (threadIdx.x == 0) { |
| 651 | total_tokens_post_pad[lora_id] = 0; |
| 652 | |
| 653 | for (int i = 0; i < num_tokens; i++) { |
| 654 | token_mask[(lora_id * num_tokens) + i] = |
| 655 | (int)token_lora_mapping[i] == lora_id; |
| 656 | } |
| 657 | } |
| 658 | |
| 659 | __syncthreads(); |
| 660 | |
| 661 | _moe_align_block_size_small_batch_expert<scalar_t, fill_threads>( |
| 662 | topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map, |
| 663 | num_experts, block_size, numel, max_num_tokens_padded, max_num_m_blocks, |
| 664 | -1, lora_id, topk_num, &token_mask[(lora_id * num_tokens)], |
| 665 | has_expert_map); |
| 666 | } |
| 667 | |
| 668 | } // namespace moe |
| 669 | } // namespace vllm |
| 670 | |
| 671 | // taken from |
| 672 | // https://github.com/sgl-project/sglang/blob/8b5f83ed3b7d2a49ad5c5cd5aa61c5d502f47dbc |
| 673 | void moe_align_block_size(torch::stable::Tensor topk_ids, int64_t num_experts, |
| 674 | int64_t block_size, |
| 675 | torch::stable::Tensor sorted_token_ids, |
| 676 | torch::stable::Tensor experts_ids, |
| 677 | torch::stable::Tensor num_tokens_post_pad, |
| 678 | std::optional<torch::stable::Tensor> maybe_expert_map, |
| 679 | std::optional<torch::stable::Tensor> scatter_idx) { |
| 680 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 681 | topk_ids.get_device_index()); |
| 682 | const cudaStream_t stream = |
| 683 | get_current_cuda_stream(topk_ids.get_device_index()); |
| 684 | |
| 685 | int32_t* scatter_ptr = nullptr; |
| 686 | if (scatter_idx.has_value()) { |
| 687 | STD_TORCH_CHECK( |
| 688 | scatter_idx->scalar_type() == torch::headeronly::ScalarType::Int, |
| 689 | "scatter_idx must be int32"); |
| 690 | STD_TORCH_CHECK(scatter_idx->device() == topk_ids.device() && |
| 691 | scatter_idx->is_contiguous() && |
| 692 | scatter_idx->numel() == topk_ids.numel(), |
| 693 | "scatter_idx must be contiguous, on the same device, and " |
| 694 | "match topk_ids.numel()"); |
| 695 | scatter_ptr = reinterpret_cast<int32_t*>(scatter_idx->mutable_data_ptr()); |
| 696 | } |
| 697 | |
| 698 | int64_t padded_num_experts = |
| 699 | ((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE; |
| 700 | int experts_per_warp = WARP_SIZE; |
| 701 | int threads = 1024; |
| 702 | threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE; |
| 703 | |
| 704 | // BlockScan uses 1024 threads and assigns one thread per expert. |
| 705 | STD_TORCH_CHECK(padded_num_experts < 1024, |
| 706 | "padded_num_experts must be less than 1024"); |
| 707 | bool has_expert_map = maybe_expert_map.has_value(); |
| 708 | torch::stable::Tensor expert_map; |
| 709 | if (has_expert_map) { |
| 710 | expert_map = maybe_expert_map.value(); |
| 711 | } else { |
| 712 | expert_map = torch::stable::new_empty(topk_ids, {0}, |
| 713 | torch::headeronly::ScalarType::Int); |
| 714 | } |
| 715 | |
| 716 | VLLM_STABLE_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES( |
| 717 | topk_ids.scalar_type(), "moe_align_block_size_kernel", [&] { |
| 718 | // calc needed amount of shared mem for `cumsum` tensors |
| 719 | bool small_batch_expert_mode = |
| 720 | (topk_ids.numel() < 1024) && (num_experts <= 64); |
| 721 | |
| 722 | if (small_batch_expert_mode) { |
| 723 | const int32_t threads = max((int32_t)num_experts, WARP_SIZE); |
| 724 | const int32_t shared_mem_size = |
| 725 | ((threads + 1) * num_experts + (num_experts + 1)) * |
| 726 | sizeof(int32_t); |
| 727 | |
| 728 | // threadIdx.x >= fill_threads: counting experts and aligning |
| 729 | // threadIdx.x < fill_threads: filling sorted_token_ids |
| 730 | constexpr int32_t fill_threads = 256; |
| 731 | auto small_batch_expert_kernel = |
| 732 | vllm::moe::moe_align_block_size_small_batch_expert_kernel< |
| 733 | scalar_t, fill_threads>; |
| 734 | small_batch_expert_kernel<<<1, fill_threads + threads, |
| 735 | shared_mem_size, stream>>>( |
| 736 | reinterpret_cast<const scalar_t*>(topk_ids.const_data_ptr()), |
| 737 | reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()), |
| 738 | reinterpret_cast<int32_t*>(experts_ids.mutable_data_ptr()), |
| 739 | reinterpret_cast<int32_t*>( |
| 740 | num_tokens_post_pad.mutable_data_ptr()), |
| 741 | reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()), |
| 742 | num_experts, block_size, topk_ids.numel(), |
| 743 | sorted_token_ids.size(0), topk_ids.size(1), has_expert_map, |
| 744 | scatter_ptr); |
| 745 | } else { |
| 746 | torch::stable::Tensor cumsum_buffer = torch::stable::new_empty( |
| 747 | topk_ids, {num_experts + 1}, torch::headeronly::ScalarType::Int); |
| 748 | auto align_kernel = vllm::moe::moe_align_block_size_kernel<scalar_t>; |
| 749 | |
| 750 | size_t num_warps = CEILDIV(padded_num_experts, experts_per_warp); |
| 751 | size_t shared_mem_size = |
| 752 | num_warps * experts_per_warp * sizeof(int32_t); |
| 753 | |
| 754 | // launch two threadblocks |
| 755 | // blockIdx.x == 0: counting experts and aligning |
| 756 | // blockIdx.x == 1: filling sorted_token_ids |
| 757 | align_kernel<<<2, threads, shared_mem_size, stream>>>( |
| 758 | reinterpret_cast<const scalar_t*>(topk_ids.const_data_ptr()), |
| 759 | reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()), |
| 760 | reinterpret_cast<int32_t*>(experts_ids.mutable_data_ptr()), |
| 761 | reinterpret_cast<int32_t*>( |
| 762 | num_tokens_post_pad.mutable_data_ptr()), |
| 763 | reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()), |
| 764 | num_experts, padded_num_experts, experts_per_warp, block_size, |
| 765 | topk_ids.numel(), |
| 766 | reinterpret_cast<int32_t*>(cumsum_buffer.mutable_data_ptr()), |
| 767 | sorted_token_ids.size(0), topk_ids.size(1), has_expert_map); |
| 768 | |
| 769 | const int block_threads = std::min(256, (int)threads); |
| 770 | const int num_blocks = |
| 771 | (topk_ids.numel() + block_threads - 1) / block_threads; |
| 772 | const int max_blocks = 65535; |
| 773 | const int actual_blocks = std::min(num_blocks, max_blocks); |
| 774 | dim3 gridDims(1, actual_blocks); |
| 775 | |
| 776 | auto sort_kernel = |
| 777 | vllm::moe::count_and_sort_expert_tokens_kernel<scalar_t>; |
| 778 | sort_kernel<<<gridDims, block_threads, 0, stream>>>( |
| 779 | reinterpret_cast<const scalar_t*>(topk_ids.const_data_ptr()), |
| 780 | reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()), |
| 781 | reinterpret_cast<int32_t*>(cumsum_buffer.mutable_data_ptr()), |
| 782 | reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()), |
| 783 | topk_ids.numel(), num_experts, sorted_token_ids.size(0), |
| 784 | topk_ids.size(1), has_expert_map, scatter_ptr); |
| 785 | } |
| 786 | }); |
| 787 | } |
| 788 | |
| 789 | void batched_moe_align_block_size(int64_t max_tokens_per_batch, |
| 790 | int64_t block_size, |
| 791 | const torch::stable::Tensor& batch_num_tokens, |
| 792 | torch::stable::Tensor sorted_ids, |
| 793 | torch::stable::Tensor batch_ids, |
| 794 | torch::stable::Tensor num_tokens_post_pad) { |
| 795 | namespace batched_kernel = vllm::moe::batched_moe_align_block_size; |
| 796 | |
| 797 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 798 | batch_num_tokens.get_device_index()); |
| 799 | const cudaStream_t stream = |
| 800 | get_current_cuda_stream(batch_num_tokens.get_device_index()); |
| 801 | int32_t const B = batch_num_tokens.size(0); |
| 802 | int32_t const num_blocks_per_batch = |
| 803 | round_to_next_multiple_of(max_tokens_per_batch, block_size) / block_size; |
| 804 | int32_t const num_blocks = num_blocks_per_batch * B; |
| 805 | int64_t const sorted_ids_size = num_blocks * block_size; |
| 806 | |
| 807 | STD_TORCH_CHECK(sorted_ids.size(0) == sorted_ids_size); |
| 808 | STD_TORCH_CHECK(batch_ids.size(0) == sorted_ids_size / block_size); |
| 809 | STD_TORCH_CHECK(num_tokens_post_pad.size(0) == 1); |
| 810 | STD_TORCH_CHECK(B <= batched_kernel::num_threads); |
| 811 | |
| 812 | // Avoid coordination overhead for small capacities or many batches. |
| 813 | int64_t const cooperative_threshold = std::max<int64_t>(256, 8 * B); |
| 814 | bool const use_cooperative_writes = |
| 815 | max_tokens_per_batch >= cooperative_threshold; |
| 816 | auto kernel = |
| 817 | use_cooperative_writes |
| 818 | ? batched_kernel::batched_moe_align_block_size_kernel<true> |
| 819 | : batched_kernel::batched_moe_align_block_size_kernel<false>; |
| 820 | kernel<<<batched_kernel::num_blocks, batched_kernel::num_threads, 0, |
| 821 | stream>>>( |
| 822 | B, max_tokens_per_batch, block_size, |
| 823 | reinterpret_cast<const int32_t*>(batch_num_tokens.const_data_ptr()), |
| 824 | reinterpret_cast<int32_t*>(sorted_ids.mutable_data_ptr()), |
| 825 | reinterpret_cast<int32_t*>(batch_ids.mutable_data_ptr()), |
| 826 | reinterpret_cast<int32_t*>(num_tokens_post_pad.mutable_data_ptr())); |
| 827 | } |
| 828 | |
| 829 | void moe_sum(torch::stable::Tensor& input, // [num_tokens, topk, hidden_size] |
| 830 | torch::stable::Tensor& output, // [num_tokens, hidden_size] |
| 831 | std::optional<torch::stable::Tensor> topk_ids, |
| 832 | std::optional<torch::stable::Tensor> expert_map) { |
| 833 | // Output is dense and written in place, so it must be contiguous. The input |
| 834 | // is read by its strides (no copy); only the hidden dim needs to be |
| 835 | // contiguous to take the vectorized path. |
| 836 | STD_TORCH_CHECK(output.is_contiguous(), |
| 837 | "moe_sum expects a contiguous output"); |
| 838 | |
| 839 | const int hidden_size = input.size(-1); |
| 840 | const int64_t num_tokens = output.numel() / hidden_size; |
| 841 | const int topk = input.size(1); |
| 842 | const int64_t stride_token = input.stride(0); |
| 843 | const int64_t stride_topk = input.stride(1); |
| 844 | const int64_t stride_hidden = input.stride(2); |
| 845 | |
| 846 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 847 | output.get_device_index()); |
| 848 | const cudaStream_t stream = |
| 849 | get_current_cuda_stream(output.get_device_index()); |
| 850 | |
| 851 | if (topk_ids.has_value()) { |
| 852 | // Pad-aware reduce path |
| 853 | const torch::stable::Tensor& tk = topk_ids.value(); |
| 854 | STD_TORCH_CHECK(tk.size(0) == num_tokens && tk.size(1) == topk, |
| 855 | "moe_sum: topk_ids must have shape [num_tokens, topk]"); |
| 856 | const int64_t stride_tk_token = tk.stride(0); |
| 857 | const int64_t stride_tk_k = tk.stride(1); |
| 858 | |
| 859 | const int32_t* expert_map_ptr = nullptr; |
| 860 | if (expert_map.has_value()) { |
| 861 | STD_TORCH_CHECK( |
| 862 | expert_map->scalar_type() == torch::headeronly::ScalarType::Int, |
| 863 | "moe_sum: expert_map must be int32"); |
| 864 | expert_map_ptr = |
| 865 | reinterpret_cast<const int32_t*>(expert_map->const_data_ptr()); |
| 866 | } |
| 867 | |
| 868 | #define LAUNCH_MOE_SUM_PAD_AWARE_VEC(TOPK) \ |
| 869 | vllm::moe::moe_sum_vec_kernel<scalar_t, idx_t, TOPK, true> \ |
| 870 | <<<grid, dim3(block), 0, stream>>>( \ |
| 871 | out_ptr, in_ptr, num_tokens, hidden_size, stride_token, stride_topk, \ |
| 872 | topk_ids_ptr, expert_map_ptr, stride_tk_token, stride_tk_k) |
| 873 | |
| 874 | VLLM_STABLE_DISPATCH_FLOATING_TYPES( |
| 875 | input.scalar_type(), "moe_sum_pad_aware", [&] { |
| 876 | constexpr int VEC = vllm::moe::MOE_SUM_VEC<scalar_t>; |
| 877 | constexpr int WIDTH = VEC * sizeof(scalar_t); |
| 878 | auto* out_ptr = |
| 879 | reinterpret_cast<scalar_t*>(output.mutable_data_ptr()); |
| 880 | auto* in_ptr = |
| 881 | reinterpret_cast<const scalar_t*>(input.const_data_ptr()); |
| 882 | |
| 883 | const bool can_vec = |
| 884 | (stride_hidden == 1) && (hidden_size % VEC == 0) && |
| 885 | (stride_token % VEC == 0) && (stride_topk % VEC == 0) && |
| 886 | (reinterpret_cast<uintptr_t>(in_ptr) % WIDTH == 0) && |
| 887 | (reinterpret_cast<uintptr_t>(out_ptr) % WIDTH == 0); |
| 888 | |
| 889 | VLLM_STABLE_DISPATCH_IDX_TYPES( |
| 890 | tk.scalar_type(), "moe_sum_pad_aware_idx", [&] { |
| 891 | auto* topk_ids_ptr = |
| 892 | reinterpret_cast<const idx_t*>(tk.const_data_ptr()); |
| 893 | if (can_vec) { |
| 894 | const int64_t n_vec = hidden_size / VEC; |
| 895 | const int64_t total = num_tokens * n_vec; |
| 896 | const int block = 256; |
| 897 | const dim3 grid( |
| 898 | std::min<int64_t>((total + block - 1) / block, 65535)); |
| 899 | switch (topk) { |
| 900 | case 1: |
| 901 | LAUNCH_MOE_SUM_PAD_AWARE_VEC(1); |
| 902 | break; |
| 903 | case 2: |
| 904 | LAUNCH_MOE_SUM_PAD_AWARE_VEC(2); |
| 905 | break; |
| 906 | case 4: |
| 907 | LAUNCH_MOE_SUM_PAD_AWARE_VEC(4); |
| 908 | break; |
| 909 | case 6: |
| 910 | LAUNCH_MOE_SUM_PAD_AWARE_VEC(6); |
| 911 | break; |
| 912 | case 8: |
| 913 | LAUNCH_MOE_SUM_PAD_AWARE_VEC(8); |
| 914 | break; |
| 915 | case 9: |
| 916 | LAUNCH_MOE_SUM_PAD_AWARE_VEC(9); |
| 917 | break; |
| 918 | default: |
| 919 | vllm::moe::moe_sum_vec_dynamic_kernel<scalar_t, idx_t, |
| 920 | true> |
| 921 | <<<grid, dim3(block), 0, stream>>>( |
| 922 | out_ptr, in_ptr, num_tokens, hidden_size, topk, |
| 923 | stride_token, stride_topk, topk_ids_ptr, |
| 924 | expert_map_ptr, stride_tk_token, stride_tk_k); |
| 925 | break; |
| 926 | } |
| 927 | } else { |
| 928 | dim3 grid(num_tokens); |
| 929 | dim3 block(std::min(hidden_size, 1024)); |
| 930 | vllm::moe::moe_sum_scalar_kernel<scalar_t, idx_t, true> |
| 931 | <<<grid, block, 0, stream>>>( |
| 932 | out_ptr, in_ptr, hidden_size, topk, stride_token, |
| 933 | stride_topk, stride_hidden, topk_ids_ptr, |
| 934 | expert_map_ptr, stride_tk_token, stride_tk_k); |
| 935 | } |
| 936 | }); |
| 937 | }); |
| 938 | #undef LAUNCH_MOE_SUM_PAD_AWARE_VEC |
| 939 | return; |
| 940 | } |
| 941 | |
| 942 | #define LAUNCH_MOE_SUM_VEC(TOPK) \ |
| 943 | vllm::moe::moe_sum_vec_kernel<scalar_t, int32_t, TOPK, false> \ |
| 944 | <<<grid, dim3(block), 0, stream>>>(out_ptr, in_ptr, num_tokens, \ |
| 945 | hidden_size, stride_token, \ |
| 946 | stride_topk, nullptr, nullptr, 0, 0) |
| 947 | |
| 948 | VLLM_STABLE_DISPATCH_FLOATING_TYPES(input.scalar_type(), "moe_sum", [&] { |
| 949 | constexpr int VEC = vllm::moe::MOE_SUM_VEC<scalar_t>; |
| 950 | constexpr int WIDTH = VEC * sizeof(scalar_t); // 16 bytes |
| 951 | auto* out_ptr = reinterpret_cast<scalar_t*>(output.mutable_data_ptr()); |
| 952 | auto* in_ptr = reinterpret_cast<const scalar_t*>(input.const_data_ptr()); |
| 953 | |
| 954 | // Vectorize along hidden only when it is contiguous (innermost stride 1), |
| 955 | // a whole number of vectors, and every row offset stays 16B-aligned. |
| 956 | const bool can_vec = (stride_hidden == 1) && (hidden_size % VEC == 0) && |
| 957 | (stride_token % VEC == 0) && |
| 958 | (stride_topk % VEC == 0) && |
| 959 | (reinterpret_cast<uintptr_t>(in_ptr) % WIDTH == 0) && |
| 960 | (reinterpret_cast<uintptr_t>(out_ptr) % WIDTH == 0); |
| 961 | if (can_vec) { |
| 962 | const int64_t n_vec = hidden_size / VEC; |
| 963 | const int64_t total = num_tokens * n_vec; |
| 964 | const int block = 256; |
| 965 | const dim3 grid(std::min<int64_t>((total + block - 1) / block, 65535)); |
| 966 | switch (topk) { |
| 967 | case 1: |
| 968 | LAUNCH_MOE_SUM_VEC(1); |
| 969 | break; |
| 970 | case 2: |
| 971 | LAUNCH_MOE_SUM_VEC(2); |
| 972 | break; |
| 973 | case 4: |
| 974 | LAUNCH_MOE_SUM_VEC(4); |
| 975 | break; |
| 976 | case 6: |
| 977 | LAUNCH_MOE_SUM_VEC(6); |
| 978 | break; |
| 979 | case 8: |
| 980 | LAUNCH_MOE_SUM_VEC(8); |
| 981 | break; |
| 982 | case 9: |
| 983 | LAUNCH_MOE_SUM_VEC(9); |
| 984 | break; |
| 985 | default: |
| 986 | vllm::moe::moe_sum_vec_dynamic_kernel<scalar_t, int32_t, false> |
| 987 | <<<grid, dim3(block), 0, stream>>>( |
| 988 | out_ptr, in_ptr, num_tokens, hidden_size, topk, stride_token, |
| 989 | stride_topk, nullptr, nullptr, 0, 0); |
| 990 | break; |
| 991 | } |
| 992 | } else { |
| 993 | dim3 grid(num_tokens); |
| 994 | dim3 block(std::min(hidden_size, 1024)); |
| 995 | vllm::moe::moe_sum_scalar_kernel<scalar_t, int32_t, false> |
| 996 | <<<grid, block, 0, stream>>>(out_ptr, in_ptr, hidden_size, topk, |
| 997 | stride_token, stride_topk, stride_hidden, |
| 998 | nullptr, nullptr, 0, 0); |
| 999 | } |
| 1000 | }); |
| 1001 | #undef LAUNCH_MOE_SUM_VEC |
| 1002 | } |
| 1003 | |
| 1004 | void moe_lora_align_block_size( |
| 1005 | torch::stable::Tensor topk_ids, torch::stable::Tensor token_lora_mapping, |
| 1006 | int64_t num_experts, int64_t block_size, int64_t max_loras, |
| 1007 | int64_t max_num_tokens_padded, int64_t max_num_m_blocks, |
| 1008 | torch::stable::Tensor sorted_token_ids, torch::stable::Tensor expert_ids, |
| 1009 | torch::stable::Tensor num_tokens_post_pad, |
| 1010 | torch::stable::Tensor adapter_enabled, torch::stable::Tensor lora_ids, |
| 1011 | std::optional<torch::stable::Tensor> maybe_expert_map) { |
| 1012 | const int topk_num = topk_ids.size(1); |
| 1013 | |
| 1014 | STD_TORCH_CHECK(block_size > 0, "block_size should be greater than 0. "); |
| 1015 | |
| 1016 | int device_max_shared_mem; |
| 1017 | int dev = topk_ids.get_device_index(); |
| 1018 | const torch::stable::accelerator::DeviceGuard device_guard(dev); |
| 1019 | cudaDeviceGetAttribute(&device_max_shared_mem, |
| 1020 | cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); |
| 1021 | const cudaStream_t stream = get_current_cuda_stream(dev); |
| 1022 | |
| 1023 | int64_t padded_num_experts = |
| 1024 | ((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE; |
| 1025 | |
| 1026 | // BlockScan uses 1024 threads and assigns one thread per expert. |
| 1027 | STD_TORCH_CHECK(padded_num_experts < 1024, |
| 1028 | "padded_num_experts must be less than 1024"); |
| 1029 | |
| 1030 | torch::stable::Tensor token_mask = |
| 1031 | torch::stable::new_empty(topk_ids, {max_loras * topk_ids.size(0)}, |
| 1032 | torch::headeronly::ScalarType::Int); |
| 1033 | bool has_expert_map = maybe_expert_map.has_value(); |
| 1034 | torch::stable::Tensor expert_map; |
| 1035 | if (has_expert_map) { |
| 1036 | expert_map = maybe_expert_map.value(); |
| 1037 | } else { |
| 1038 | expert_map = torch::stable::new_empty(topk_ids, {0}, |
| 1039 | torch::headeronly::ScalarType::Int); |
| 1040 | } |
| 1041 | |
| 1042 | VLLM_STABLE_DISPATCH_INTEGRAL_TYPES( |
| 1043 | topk_ids.scalar_type(), "moe_lora_align_sum_kernel", [&] { |
| 1044 | bool small_batch_expert_mode = |
| 1045 | (topk_ids.numel() < 1024) && (num_experts <= 64); |
| 1046 | |
| 1047 | if (small_batch_expert_mode) { |
| 1048 | const int32_t num_thread = max((int32_t)num_experts, 128); |
| 1049 | const int32_t shared_mem = |
| 1050 | (num_thread + 1) * num_experts * sizeof(int32_t) + |
| 1051 | (num_experts + 1) * sizeof(int32_t); |
| 1052 | if (shared_mem > device_max_shared_mem) { |
| 1053 | STD_TORCH_CHECK(false, "Shared memory usage exceeds device limit."); |
| 1054 | } |
| 1055 | |
| 1056 | // threadIdx.x >= fill_threads: counting experts and aligning |
| 1057 | // threadIdx.x < fill_threads: filling sorted_token_ids |
| 1058 | constexpr int32_t fill_threads = 256; |
| 1059 | |
| 1060 | dim3 blockDim(num_thread + fill_threads); |
| 1061 | auto kernel = |
| 1062 | vllm::moe::moe_lora_align_block_size_small_batch_expert_kernel< |
| 1063 | scalar_t, fill_threads>; |
| 1064 | STD_CUDA_CHECK(VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize( |
| 1065 | (void*)kernel, shared_mem)); |
| 1066 | // Grid size is (max_loras + 1) because active_lora_ids has length |
| 1067 | // max_loras + 1: sorted-unique values of token_lora_mapping, which |
| 1068 | // can include -1 (base-model tokens) in addition to up to max_loras |
| 1069 | // real LoRA slots. Using max_loras would drop the real LoRA slot |
| 1070 | // when -1 is present at position 0 and leave output buffers |
| 1071 | // uninitialized, causing illegal memory accesses in downstream |
| 1072 | // MoE-LoRA kernels. This mirrors the fix made for the Triton |
| 1073 | // _fused_moe_lora_kernel grid in vllm-project/vllm#32277. |
| 1074 | kernel<<<max_loras + 1, blockDim, shared_mem, stream>>>( |
| 1075 | reinterpret_cast<scalar_t*>(topk_ids.mutable_data_ptr()), |
| 1076 | reinterpret_cast<int32_t*>(token_lora_mapping.mutable_data_ptr()), |
| 1077 | block_size, |
| 1078 | reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()), |
| 1079 | num_experts, max_loras, topk_ids.numel(), max_num_tokens_padded, |
| 1080 | max_num_m_blocks, |
| 1081 | reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()), |
| 1082 | reinterpret_cast<int32_t*>(expert_ids.mutable_data_ptr()), |
| 1083 | topk_num, |
| 1084 | reinterpret_cast<int32_t*>( |
| 1085 | num_tokens_post_pad.mutable_data_ptr()), |
| 1086 | reinterpret_cast<int32_t*>(adapter_enabled.mutable_data_ptr()), |
| 1087 | reinterpret_cast<int32_t*>(lora_ids.mutable_data_ptr()), |
| 1088 | reinterpret_cast<int32_t*>(token_mask.mutable_data_ptr()), |
| 1089 | has_expert_map); |
| 1090 | } else { |
| 1091 | int num_thread = 1024; |
| 1092 | dim3 blockDim(num_thread); |
| 1093 | size_t num_warps = CEILDIV(padded_num_experts, WARP_SIZE); |
| 1094 | |
| 1095 | size_t shared_mem_size = num_warps * WARP_SIZE * sizeof(int32_t); |
| 1096 | |
| 1097 | // cumsum buffer |
| 1098 | torch::stable::Tensor cumsum = torch::stable::new_zeros( |
| 1099 | topk_ids, {max_loras * (num_experts + 1)}, |
| 1100 | torch::headeronly::ScalarType::Int); |
| 1101 | |
| 1102 | auto align_kernel = |
| 1103 | vllm::moe::moe_lora_align_block_size_kernel<scalar_t>; |
| 1104 | |
| 1105 | // Launch two threadblocks per LoRA slot, across max_loras + 1 slots |
| 1106 | // to cover the extra "-1" (base-model tokens) entry that |
| 1107 | // active_lora_ids may contain in addition to up to max_loras real |
| 1108 | // LoRA slots. Using max_loras would drop the real LoRA slot when -1 |
| 1109 | // occupies position 0 and leave the output buffers uninitialized, |
| 1110 | // causing illegal memory accesses downstream. Mirrors the grid fix |
| 1111 | // applied to _fused_moe_lora_kernel in vllm-project/vllm#32277. |
| 1112 | // blockIdx.x % 2 == 0: counting experts and aligning |
| 1113 | // blockIdx.x % 2 == 1: filling sorted_token_ids |
| 1114 | align_kernel<<<(max_loras + 1) * 2, blockDim, shared_mem_size, |
| 1115 | stream>>>( |
| 1116 | reinterpret_cast<scalar_t*>(topk_ids.mutable_data_ptr()), |
| 1117 | reinterpret_cast<int32_t*>(token_lora_mapping.mutable_data_ptr()), |
| 1118 | block_size, |
| 1119 | reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()), |
| 1120 | num_experts, max_loras, topk_ids.numel(), max_num_tokens_padded, |
| 1121 | max_num_m_blocks, |
| 1122 | reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()), |
| 1123 | reinterpret_cast<int32_t*>(expert_ids.mutable_data_ptr()), |
| 1124 | topk_num, |
| 1125 | reinterpret_cast<int32_t*>( |
| 1126 | num_tokens_post_pad.mutable_data_ptr()), |
| 1127 | reinterpret_cast<int32_t*>(adapter_enabled.mutable_data_ptr()), |
| 1128 | reinterpret_cast<int32_t*>(cumsum.mutable_data_ptr()), WARP_SIZE, |
| 1129 | padded_num_experts, |
| 1130 | reinterpret_cast<int32_t*>(lora_ids.mutable_data_ptr()), |
| 1131 | reinterpret_cast<int32_t*>(token_mask.mutable_data_ptr()), |
| 1132 | has_expert_map); |
| 1133 | |
| 1134 | const int block_threads = std::min(256, (int)num_thread); |
| 1135 | const int num_blocks = |
| 1136 | (topk_ids.numel() + block_threads - 1) / block_threads; |
| 1137 | |
| 1138 | const int max_blocks = 65535; |
| 1139 | const int actual_blocks = std::min(num_blocks, max_blocks); |
| 1140 | |
| 1141 | // Same rationale as align_kernel above: iterate over max_loras + 1 |
| 1142 | // slots so the sort kernel processes the real LoRA slot even when |
| 1143 | // active_lora_ids has -1 at position 0. |
| 1144 | dim3 gridDims(max_loras + 1, actual_blocks); |
| 1145 | auto sort_kernel = |
| 1146 | vllm::moe::lora_count_and_sort_expert_tokens_kernel<scalar_t>; |
| 1147 | |
| 1148 | sort_kernel<<<gridDims, block_threads, 0, stream>>>( |
| 1149 | reinterpret_cast<const scalar_t*>(topk_ids.const_data_ptr()), |
| 1150 | reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()), |
| 1151 | reinterpret_cast<int32_t*>(cumsum.mutable_data_ptr()), |
| 1152 | reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()), |
| 1153 | topk_ids.numel(), num_experts, max_num_tokens_padded, topk_num, |
| 1154 | reinterpret_cast<int32_t*>(token_mask.mutable_data_ptr()), |
| 1155 | max_loras, |
| 1156 | reinterpret_cast<int32_t*>(lora_ids.mutable_data_ptr()), |
| 1157 | reinterpret_cast<int32_t*>(adapter_enabled.mutable_data_ptr()), |
| 1158 | has_expert_map); |
| 1159 | } |
| 1160 | }); |
| 1161 | } |