A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 20:00 UTC
Lora MoE Align Improvements (#29257)
Signed-off-by: gnovack <gnovack@amazon.com>
gnovack committed
Dec 8, 2025 at 18:35 UTC
ea657f2078072541f68fbc11ab83f81a585a9dc4
7 files changed
+360
-249
CMakeLists.txt
-1
@@ -944,7 +944,6 @@ target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
944
set(VLLM_MOE_EXT_SRC
945
"csrc/moe/torch_bindings.cpp"
946
"csrc/moe/moe_align_sum_kernels.cu"
947
- "csrc/moe/moe_lora_align_sum_kernels.cu"
947
"csrc/moe/topk_softmax_kernels.cu")
948
949
if(VLLM_GPU_LANG STREQUAL "CUDA")
csrc/moe/moe_align_sum_kernels.cu
+354
-71
@@ -14,7 +14,6 @@
14
15
namespace vllm {
16
namespace moe {
17
-
17
namespace batched_moe_align_block_size {
18
19
// Note num_threads needs to be 1024 for BlockScan Reduction in the kernel.
@@ -80,23 +79,30 @@ __global__ void batched_moe_align_block_size_kernel(
79
} // namespace batched_moe_align_block_size
80
81
template <typename scalar_t>
83
-__global__ void moe_align_block_size_kernel(
82
+__device__ void _moe_align_block_size(
83
const scalar_t* __restrict__ topk_ids,
84
int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids,
85
int32_t* __restrict__ total_tokens_post_pad,
86
int32_t* __restrict__ expert_map, int32_t num_experts,
87
int32_t padded_num_experts, int32_t experts_per_warp, int32_t block_size,
88
size_t numel, int32_t* __restrict__ cumsum, int32_t max_num_tokens_padded,
90
- bool has_expert_map) {
89
+ int32_t max_num_m_blocks, int32_t model_offset, int32_t inactive_expert_id,
90
+ int32_t topk_num, int32_t* token_mask, bool has_expert_map) {
91
extern __shared__ int32_t shared_counts[];
92
93
- // Use a separate threadblock to fill sorted_token_ids.
93
+ // Compute input buffer offsets. Typically these will all be 0, except when
94
+ // using Multi LoRA.
95
+ int sorted_token_ids_offset = max_num_tokens_padded * model_offset;
96
+ int expert_ids_offset = max_num_m_blocks * model_offset;
97
+ int cumsum_offset = (num_experts + 1) * model_offset;
98
+
99
+ // Use separate threadblocks to fill sorted_token_ids.
100
// This is safe since the current kernel does not use sorted_token_ids.
95
- if (blockIdx.x == 1) {
101
+ if (blockIdx.x % 2) {
102
// Initialize sorted_token_ids with numel
103
for (size_t it = threadIdx.x; it < max_num_tokens_padded;
104
it += blockDim.x) {
99
- sorted_token_ids[it] = numel;
105
+ sorted_token_ids[sorted_token_ids_offset + it] = numel;
106
}
107
return;
108
}
@@ -127,7 +133,9 @@ __global__ void moe_align_block_size_kernel(
133
}
134
int warp_idx = expert_id / experts_per_warp;
135
int expert_offset = expert_id % experts_per_warp;
130
- atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset], 1);
136
+ int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
137
+ atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset],
138
+ mask);
139
}
140
141
__syncthreads();
@@ -148,77 +156,44 @@ __global__ void moe_align_block_size_kernel(
156
int cumsum_val;
157
BlockScan(temp_storage).ExclusiveSum(expert_count, cumsum_val);
158
if (expert_id <= num_experts) {
151
- cumsum[expert_id] = cumsum_val;
159
+ cumsum[cumsum_offset + expert_id] = cumsum_val;
160
}
161
162
if (expert_id == num_experts) {
155
- *total_tokens_post_pad = cumsum_val;
163
+ total_tokens_post_pad[model_offset] = cumsum_val;
164
}
165
166
__syncthreads();
167
168
if (threadIdx.x < num_experts) {
161
- for (int i = cumsum[threadIdx.x]; i < cumsum[threadIdx.x + 1];
162
- i += block_size) {
163
- expert_ids[i / block_size] = threadIdx.x;
169
+ for (int i = cumsum[cumsum_offset + threadIdx.x];
170
+ i < cumsum[cumsum_offset + threadIdx.x + 1]; i += block_size) {
171
+ expert_ids[expert_ids_offset + i / block_size] = threadIdx.x;
172
}
173
}
174
175
// Fill remaining expert_ids with 0
168
- const size_t fill_start_idx = cumsum[num_experts] / block_size + threadIdx.x;
169
- const size_t expert_ids_size = CEILDIV(max_num_tokens_padded, block_size);
170
- for (size_t i = fill_start_idx; i < expert_ids_size; i += blockDim.x) {
171
- expert_ids[i] = 0;
172
- }
173
-}
174
-
175
-template <typename scalar_t>
176
-__global__ void count_and_sort_expert_tokens_kernel(
177
- const scalar_t* __restrict__ topk_ids,
178
- int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer,
179
- int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts,
180
- bool has_expert_map) {
181
- const size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
182
- const size_t stride = blockDim.x * gridDim.x;
183
-
184
- for (size_t i = tid; i < numel; i += stride) {
185
- int32_t expert_id = topk_ids[i];
186
- if (expert_id >= num_experts) {
187
- continue;
188
- }
189
- if (has_expert_map) {
190
- expert_id = expert_map[expert_id];
191
- // filter invalid experts
192
- if (expert_id == -1) continue;
193
- }
194
- int32_t rank_post_pad = atomicAdd(&cumsum_buffer[expert_id], 1);
195
- sorted_token_ids[rank_post_pad] = i;
196
- }
197
-}
198
-
199
-template <typename scalar_t, int TOPK>
200
-__global__ void moe_sum_kernel(
201
- scalar_t* __restrict__ out, // [..., d]
202
- const scalar_t* __restrict__ input, // [..., topk, d]
203
- const int d) {
204
- const int64_t token_idx = blockIdx.x;
205
- for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
206
- scalar_t x = 0.0;
207
-#pragma unroll
208
- for (int k = 0; k < TOPK; ++k) {
209
- x += VLLM_LDG(&input[token_idx * TOPK * d + k * d + idx]);
210
- }
211
- out[token_idx * d + idx] = x;
176
+ const size_t fill_start_idx =
177
+ cumsum[cumsum_offset + num_experts] / block_size + threadIdx.x;
178
+ for (size_t i = fill_start_idx; i < max_num_m_blocks; i += blockDim.x) {
179
+ expert_ids[expert_ids_offset + i] = inactive_expert_id;
180
}
181
}
182
183
template <typename scalar_t, int32_t fill_threads>
216
-__global__ void moe_align_block_size_small_batch_expert_kernel(
184
+__device__ void _moe_align_block_size_small_batch_expert(
185
const scalar_t* __restrict__ topk_ids,
186
int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids,
187
int32_t* __restrict__ total_tokens_post_pad,
188
int32_t* __restrict__ expert_map, int32_t num_experts, int32_t block_size,
221
- size_t numel, int32_t max_num_tokens_padded, bool has_expert_map) {
189
+ size_t numel, int32_t max_num_tokens_padded, int32_t max_num_m_blocks,
190
+ int32_t inactive_expert_id, int32_t model_offset, int32_t topk_num,
191
+ int32_t* token_mask, bool has_expert_map) {
192
+ // Compute input buffer offsets. Typically these will all be 0, except when
193
+ // using Multi LoRA.
194
+ int sorted_token_ids_offset = max_num_tokens_padded * model_offset;
195
+ int expert_ids_offset = max_num_m_blocks * model_offset;
196
+
197
// Use an additional group of threads to fill sorted_token_ids.
198
// Since the current kernel will use sorted_token_ids afterward,
199
// we fill sorted_token_ids within the same threadblock to make
@@ -227,7 +202,7 @@ __global__ void moe_align_block_size_small_batch_expert_kernel(
202
// Initialize sorted_token_ids with numel
203
for (size_t it = threadIdx.x; it < max_num_tokens_padded;
204
it += fill_threads) {
230
- sorted_token_ids[it] = numel;
205
+ sorted_token_ids[sorted_token_ids_offset + it] = numel;
206
}
207
// Three __syncthreads() corresponding to the other threads
208
__syncthreads();
@@ -254,7 +229,8 @@ __global__ void moe_align_block_size_small_batch_expert_kernel(
229
// filter invalid expert
230
if (expert_id == -1) continue;
231
}
257
- ++tokens_cnts[(tid + 1) * num_experts + expert_id];
232
+ int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
233
+ tokens_cnts[(tid + 1) * num_experts + expert_id] += mask;
234
}
235
236
__syncthreads();
@@ -277,22 +253,22 @@ __global__ void moe_align_block_size_small_batch_expert_kernel(
253
CEILDIV(tokens_cnts[stride * num_experts + i - 1], block_size) *
254
block_size;
255
}
280
- *total_tokens_post_pad = static_cast<int32_t>(cumsum[num_experts]);
256
+ total_tokens_post_pad[model_offset] =
257
+ static_cast<int32_t>(cumsum[num_experts]);
258
}
259
260
__syncthreads();
261
262
if (tid < num_experts) {
263
for (int i = cumsum[tid]; i < cumsum[tid + 1]; i += block_size) {
287
- expert_ids[i / block_size] = tid;
264
+ expert_ids[expert_ids_offset + i / block_size] = tid;
265
}
266
}
267
268
// Fill remaining expert_ids with 0
269
const size_t fill_start_idx = cumsum[num_experts] / block_size + tid;
293
- const size_t expert_ids_size = CEILDIV(max_num_tokens_padded, block_size);
294
- for (size_t i = fill_start_idx; i < expert_ids_size; i += stride) {
295
- expert_ids[i] = 0;
270
+ for (size_t i = fill_start_idx; i < max_num_m_blocks; i += stride) {
271
+ expert_ids[expert_ids_offset + i] = inactive_expert_id;
272
}
273
274
for (size_t i = tid; i < numel; i += stride) {
@@ -304,11 +280,195 @@ __global__ void moe_align_block_size_small_batch_expert_kernel(
280
}
281
int32_t rank_post_pad =
282
tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id];
307
- sorted_token_ids[rank_post_pad] = i;
308
- ++tokens_cnts[tid * num_experts + expert_id];
283
+
284
+ if (token_mask == nullptr || token_mask[i / topk_num]) {
285
+ sorted_token_ids[sorted_token_ids_offset + rank_post_pad] = i;
286
+ ++tokens_cnts[tid * num_experts + expert_id];
287
+ }
288
+ }
289
+}
290
+
291
+template <typename scalar_t>
292
+__device__ void _count_and_sort_expert_tokens(
293
+ const scalar_t* __restrict__ topk_ids,
294
+ int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer,
295
+ int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts,
296
+ int32_t max_num_tokens_padded, int32_t* __restrict__ token_mask,
297
+ int32_t model_offset, int32_t topk_num, bool has_expert_map) {
298
+ const size_t tid = blockIdx.y * blockDim.x + threadIdx.x;
299
+ const size_t stride = blockDim.x * gridDim.y;
300
+
301
+ for (size_t i = tid; i < numel; i += stride) {
302
+ int32_t expert_id = topk_ids[i];
303
+ if (expert_id >= num_experts) {
304
+ continue;
305
+ }
306
+
307
+ if (has_expert_map) {
308
+ expert_id = expert_map[expert_id];
309
+ // filter invalid experts
310
+ if (expert_id == -1) continue;
311
+ }
312
+
313
+ if (token_mask == nullptr || token_mask[i / topk_num]) {
314
+ int32_t rank_post_pad = atomicAdd(
315
+ &cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1);
316
+ sorted_token_ids[max_num_tokens_padded * model_offset + rank_post_pad] =
317
+ i;
318
+ }
319
+ }
320
+}
321
+
322
+template <typename scalar_t>
323
+__global__ void moe_align_block_size_kernel(
324
+ const scalar_t* __restrict__ topk_ids,
325
+ int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids,
326
+ int32_t* __restrict__ total_tokens_post_pad,
327
+ int32_t* __restrict__ expert_map, int32_t num_experts,
328
+ int32_t padded_num_experts, int32_t experts_per_warp, int32_t block_size,
329
+ size_t numel, int32_t* __restrict__ cumsum, int32_t max_num_tokens_padded,
330
+ int32_t topk_num, bool has_expert_map) {
331
+ _moe_align_block_size(
332
+ topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map,
333
+ num_experts, padded_num_experts, experts_per_warp, block_size, numel,
334
+ cumsum, max_num_tokens_padded, CEILDIV(max_num_tokens_padded, block_size),
335
+ 0, 0, topk_num, nullptr, has_expert_map);
336
+}
337
+
338
+template <typename scalar_t>
339
+__global__ void count_and_sort_expert_tokens_kernel(
340
+ const scalar_t* __restrict__ topk_ids,
341
+ int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer,
342
+ int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts,
343
+ int32_t max_num_tokens_padded, int32_t topk_num, bool has_expert_map) {
344
+ _count_and_sort_expert_tokens(
345
+ topk_ids, sorted_token_ids, cumsum_buffer, expert_map, numel, num_experts,
346
+ max_num_tokens_padded, nullptr, 0, topk_num, has_expert_map);
347
+}
348
+
349
+template <typename scalar_t, int TOPK>
350
+__global__ void moe_sum_kernel(
351
+ scalar_t* __restrict__ out, // [..., d]
352
+ const scalar_t* __restrict__ input, // [..., topk, d]
353
+ const int d) {
354
+ const int64_t token_idx = blockIdx.x;
355
+ for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
356
+ scalar_t x = 0.0;
357
+#pragma unroll
358
+ for (int k = 0; k < TOPK; ++k) {
359
+ x += VLLM_LDG(&input[token_idx * TOPK * d + k * d + idx]);
360
+ }
361
+ out[token_idx * d + idx] = x;
362
}
363
}
364
365
+template <typename scalar_t, int32_t fill_threads>
366
+__global__ void moe_align_block_size_small_batch_expert_kernel(
367
+ const scalar_t* __restrict__ topk_ids,
368
+ int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids,
369
+ int32_t* __restrict__ total_tokens_post_pad,
370
+ int32_t* __restrict__ expert_map, int32_t num_experts, int32_t block_size,
371
+ size_t numel, int32_t max_num_tokens_padded, int32_t topk_num,
372
+ bool has_expert_map) {
373
+ _moe_align_block_size_small_batch_expert<scalar_t, fill_threads>(
374
+ topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map,
375
+ num_experts, block_size, numel, max_num_tokens_padded,
376
+ CEILDIV(max_num_tokens_padded, block_size), 0, 0, topk_num, nullptr,
377
+ has_expert_map);
378
+}
379
+
380
+template <typename scalar_t>
381
+__global__ void moe_lora_align_block_size_kernel(
382
+ scalar_t* __restrict__ topk_ids, int32_t* __restrict__ token_lora_mapping,
383
+ int64_t block_size, int32_t* __restrict__ expert_map, int num_experts,
384
+ int max_loras, size_t numel, int max_num_tokens_padded,
385
+ int max_num_m_blocks, int32_t* __restrict__ sorted_token_ids,
386
+ int32_t* __restrict__ expert_ids, int32_t topk_num,
387
+ int32_t* total_tokens_post_pad, int32_t* adapter_enabled,
388
+ int32_t* __restrict__ cumsum, int32_t experts_per_warp,
389
+ int32_t padded_num_experts, int32_t* lora_ids,
390
+ int32_t* __restrict__ token_mask, bool has_expert_map) {
391
+ int lora_idx = blockIdx.x / 2;
392
+ int lora_id = lora_ids[lora_idx];
393
+ if (lora_id == -1 || adapter_enabled[lora_id] == 0) {
394
+ return;
395
+ }
396
+
397
+ // Populate the token_mask based on the token-LoRA mapping
398
+ int num_tokens = numel / topk_num;
399
+ if (threadIdx.x == 0) {
400
+ total_tokens_post_pad[lora_id] = 0;
401
+
402
+ for (int i = 0; i < num_tokens; i++) {
403
+ token_mask[(lora_id * num_tokens) + i] =
404
+ (int)token_lora_mapping[i] == lora_id;
405
+ }
406
+ }
407
+
408
+ __syncthreads();
409
+
410
+ _moe_align_block_size(
411
+ topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map,
412
+ num_experts, padded_num_experts, experts_per_warp, block_size, numel,
413
+ cumsum, max_num_tokens_padded, max_num_m_blocks, lora_id, -1, topk_num,
414
+ &token_mask[(lora_id * num_tokens)], has_expert_map);
415
+}
416
+
417
+template <typename scalar_t>
418
+__global__ void lora_count_and_sort_expert_tokens_kernel(
419
+ const scalar_t* __restrict__ topk_ids,
420
+ int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer,
421
+ int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts,
422
+ int32_t max_num_tokens_padded, int32_t topk_num, int32_t* token_mask,
423
+ int32_t* lora_ids, bool has_expert_map) {
424
+ int lora_idx = blockIdx.x;
425
+ int lora_id = lora_ids[lora_idx];
426
+ if (lora_id == -1) {
427
+ return;
428
+ }
429
+
430
+ int num_tokens = numel / topk_num;
431
+
432
+ _count_and_sort_expert_tokens(
433
+ topk_ids, sorted_token_ids, cumsum_buffer, expert_map, numel, num_experts,
434
+ max_num_tokens_padded, &token_mask[(lora_id * num_tokens)], lora_id,
435
+ topk_num, has_expert_map);
436
+}
437
+
438
+template <typename scalar_t, int32_t fill_threads>
439
+__global__ void moe_lora_align_block_size_small_batch_expert_kernel(
440
+ scalar_t* __restrict__ topk_ids, int32_t* token_lora_mapping,
441
+ int64_t block_size, int32_t* __restrict__ expert_map, int num_experts,
442
+ int max_loras, size_t numel, int max_num_tokens_padded,
443
+ int max_num_m_blocks, int32_t* __restrict__ sorted_token_ids,
444
+ int32_t* __restrict__ expert_ids, int topk_num,
445
+ int32_t* total_tokens_post_pad, int32_t* adapter_enabled, int32_t* lora_ids,
446
+ int32_t* token_mask, bool has_expert_map) {
447
+ int lora_idx = blockIdx.x;
448
+ int lora_id = lora_ids[lora_idx];
449
+ if (lora_id == -1 || adapter_enabled[lora_id] == 0) {
450
+ return;
451
+ }
452
+
453
+ int num_tokens = numel / topk_num;
454
+ if (threadIdx.x == 0) {
455
+ total_tokens_post_pad[lora_id] = 0;
456
+
457
+ for (int i = 0; i < num_tokens; i++) {
458
+ token_mask[(lora_id * num_tokens) + i] =
459
+ (int)token_lora_mapping[i] == lora_id;
460
+ }
461
+ }
462
+
463
+ __syncthreads();
464
+
465
+ _moe_align_block_size_small_batch_expert<scalar_t, fill_threads>(
466
+ topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map,
467
+ num_experts, block_size, numel, max_num_tokens_padded, max_num_m_blocks,
468
+ -1, lora_id, topk_num, &token_mask[(lora_id * num_tokens)],
469
+ has_expert_map);
470
+}
471
+
472
} // namespace moe
473
} // namespace vllm
474
@@ -365,7 +525,8 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
525
experts_ids.data_ptr<int32_t>(),
526
num_tokens_post_pad.data_ptr<int32_t>(),
527
expert_map.data_ptr<int32_t>(), num_experts, block_size,
368
- topk_ids.numel(), sorted_token_ids.size(0), has_expert_map);
528
+ topk_ids.numel(), sorted_token_ids.size(0), topk_ids.size(1),
529
+ has_expert_map);
530
} else {
531
torch::Tensor cumsum_buffer =
532
torch::empty({num_experts + 1}, options_int);
@@ -386,21 +547,23 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
547
expert_map.data_ptr<int32_t>(), num_experts, padded_num_experts,
548
experts_per_warp, block_size, topk_ids.numel(),
549
cumsum_buffer.data_ptr<int32_t>(), sorted_token_ids.size(0),
389
- has_expert_map);
550
+ topk_ids.size(1), has_expert_map);
551
552
const int block_threads = std::min(256, (int)threads);
553
const int num_blocks =
554
(topk_ids.numel() + block_threads - 1) / block_threads;
555
const int max_blocks = 65535;
556
const int actual_blocks = std::min(num_blocks, max_blocks);
557
+ dim3 gridDims(1, actual_blocks);
558
559
auto sort_kernel =
560
vllm::moe::count_and_sort_expert_tokens_kernel<scalar_t>;
399
- sort_kernel<<<actual_blocks, block_threads, 0, stream>>>(
561
+ sort_kernel<<<gridDims, block_threads, 0, stream>>>(
562
topk_ids.data_ptr<scalar_t>(),
563
sorted_token_ids.data_ptr<int32_t>(),
564
cumsum_buffer.data_ptr<int32_t>(), expert_map.data_ptr<int32_t>(),
403
- topk_ids.numel(), num_experts, has_expert_map);
565
+ topk_ids.numel(), num_experts, sorted_token_ids.size(0),
566
+ topk_ids.size(1), has_expert_map);
567
}
568
});
569
}
@@ -474,3 +637,123 @@ void moe_sum(torch::Tensor& input, // [num_tokens, topk, hidden_size]
637
break;
638
}
639
}
640
+
641
+void moe_lora_align_block_size(
642
+ torch::Tensor topk_ids, torch::Tensor token_lora_mapping,
643
+ int64_t num_experts, int64_t block_size, int64_t max_loras,
644
+ int64_t max_num_tokens_padded, int64_t max_num_m_blocks,
645
+ torch::Tensor sorted_token_ids, torch::Tensor expert_ids,
646
+ torch::Tensor num_tokens_post_pad, torch::Tensor adapter_enabled,
647
+ torch::Tensor lora_ids, std::optional<torch::Tensor> maybe_expert_map) {
648
+ const int topk_num = topk_ids.size(1);
649
+
650
+ TORCH_CHECK(block_size > 0, "block_size should be greater than 0. ");
651
+
652
+ int device_max_shared_mem;
653
+ auto dev = topk_ids.get_device();
654
+ cudaDeviceGetAttribute(&device_max_shared_mem,
655
+ cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
656
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
657
+
658
+ int64_t padded_num_experts =
659
+ ((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
660
+
661
+ // BlockScan uses 1024 threads and assigns one thread per expert.
662
+ TORCH_CHECK(padded_num_experts < 1024,
663
+ "padded_num_experts must be less than 1024");
664
+
665
+ auto options_int =
666
+ torch::TensorOptions().dtype(torch::kInt).device(topk_ids.device());
667
+ torch::Tensor token_mask =
668
+ torch::empty({max_loras * topk_ids.size(0)}, options_int);
669
+ bool has_expert_map = maybe_expert_map.has_value();
670
+ torch::Tensor expert_map;
671
+ if (has_expert_map) {
672
+ expert_map = maybe_expert_map.value();
673
+ } else {
674
+ expert_map = torch::empty({0}, options_int);
675
+ }
676
+
677
+ VLLM_DISPATCH_INTEGRAL_TYPES(
678
+ topk_ids.scalar_type(), "moe_lora_align_sum_kernel", [&] {
679
+ bool small_batch_expert_mode =
680
+ (topk_ids.numel() < 1024) && (num_experts <= 64);
681
+
682
+ if (small_batch_expert_mode) {
683
+ const int32_t num_thread = max((int32_t)num_experts, 128);
684
+ const int32_t shared_mem =
685
+ (num_thread + 1) * num_experts * sizeof(int32_t) +
686
+ (num_experts + 1) * sizeof(int32_t);
687
+ if (shared_mem > device_max_shared_mem) {
688
+ TORCH_CHECK(false, "Shared memory usage exceeds device limit.");
689
+ }
690
+
691
+ // threadIdx.x >= fill_threads: counting experts and aligning
692
+ // threadIdx.x < fill_threads: filling sorted_token_ids
693
+ constexpr int32_t fill_threads = 256;
694
+
695
+ dim3 blockDim(num_thread + fill_threads);
696
+ auto kernel =
697
+ vllm::moe::moe_lora_align_block_size_small_batch_expert_kernel<
698
+ scalar_t, fill_threads>;
699
+ AT_CUDA_CHECK(VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(
700
+ (void*)kernel, shared_mem));
701
+ kernel<<<max_loras, blockDim, shared_mem, stream>>>(
702
+ topk_ids.data_ptr<scalar_t>(),
703
+ token_lora_mapping.data_ptr<int32_t>(), block_size,
704
+ expert_map.data_ptr<int32_t>(), num_experts, max_loras,
705
+ topk_ids.numel(), max_num_tokens_padded, max_num_m_blocks,
706
+ sorted_token_ids.data_ptr<int32_t>(),
707
+ expert_ids.data_ptr<int32_t>(), topk_num,
708
+ num_tokens_post_pad.data_ptr<int32_t>(),
709
+ adapter_enabled.data_ptr<int32_t>(), lora_ids.data_ptr<int32_t>(),
710
+ token_mask.data_ptr<int32_t>(), has_expert_map);
711
+ } else {
712
+ int num_thread = 1024;
713
+ dim3 blockDim(num_thread);
714
+ size_t num_warps = CEILDIV(padded_num_experts, WARP_SIZE);
715
+
716
+ size_t shared_mem_size = num_warps * WARP_SIZE * sizeof(int32_t);
717
+
718
+ // cumsum buffer
719
+ torch::Tensor cumsum =
720
+ torch::zeros({max_loras * (num_experts + 1)}, options_int);
721
+
722
+ auto align_kernel =
723
+ vllm::moe::moe_lora_align_block_size_kernel<scalar_t>;
724
+
725
+ // launch two threadblocks for each lora
726
+ // blockIdx.x % 2 == 0: counting experts and aligning
727
+ // blockIdx.x % 2 == 1: filling sorted_token_ids
728
+ align_kernel<<<max_loras * 2, blockDim, shared_mem_size, stream>>>(
729
+ topk_ids.data_ptr<scalar_t>(),
730
+ token_lora_mapping.data_ptr<int32_t>(), block_size,
731
+ expert_map.data_ptr<int32_t>(), num_experts, max_loras,
732
+ topk_ids.numel(), max_num_tokens_padded, max_num_m_blocks,
733
+ sorted_token_ids.data_ptr<int32_t>(),
734
+ expert_ids.data_ptr<int32_t>(), topk_num,
735
+ num_tokens_post_pad.data_ptr<int32_t>(),
736
+ adapter_enabled.data_ptr<int32_t>(), cumsum.data_ptr<int32_t>(),
737
+ WARP_SIZE, padded_num_experts, lora_ids.data_ptr<int32_t>(),
738
+ token_mask.data_ptr<int32_t>(), has_expert_map);
739
+
740
+ const int block_threads = std::min(256, (int)num_thread);
741
+ const int num_blocks =
742
+ (topk_ids.numel() + block_threads - 1) / block_threads;
743
+
744
+ const int max_blocks = 65535;
745
+ const int actual_blocks = std::min(num_blocks, max_blocks);
746
+
747
+ dim3 gridDims(max_loras, actual_blocks);
748
+ auto sort_kernel =
749
+ vllm::moe::lora_count_and_sort_expert_tokens_kernel<scalar_t>;
750
+
751
+ sort_kernel<<<gridDims, block_threads, 0, stream>>>(
752
+ topk_ids.data_ptr<scalar_t>(),
753
+ sorted_token_ids.data_ptr<int32_t>(), cumsum.data_ptr<int32_t>(),
754
+ expert_map.data_ptr<int32_t>(), topk_ids.numel(), num_experts,
755
+ max_num_tokens_padded, topk_num, token_mask.data_ptr<int32_t>(),
756
+ lora_ids.data_ptr<int32_t>(), has_expert_map);
757
+ }
758
+ });
759
+}
\ No newline at end of file
csrc/moe/moe_lora_align_sum_kernels.cu
deleted
-174
@@ -1,174 +0,0 @@
1
-#include <stdio.h>
2
-#include <stdlib.h>
3
-#include <time.h>
4
-#include <torch/all.h>
5
-#include <ATen/cuda/CUDAContext.h>
6
-#include <c10/cuda/CUDAGuard.h>
7
-
8
-#include <ATen/ATen.h>
9
-#include <ATen/cuda/Atomic.cuh>
10
-
11
-#include "../cuda_compat.h"
12
-#include "../dispatch_utils.h"
13
-#include "core/math.hpp"
14
-
15
-namespace {
16
-
17
-__device__ __forceinline__ int32_t index(int32_t total_col, int32_t row,
18
- int32_t col) {
19
- return row * total_col + col;
20
-}
21
-
22
-} // namespace
23
-
24
-// TODO: Refactor common parts with moe_align_sum_kernels
25
-template <typename scalar_t, typename token_cnts_t>
26
-__global__ void moe_lora_align_sum_kernel(
27
- scalar_t* __restrict__ topk_ids, int32_t* token_lora_mapping,
28
- int64_t block_size, int num_experts, int max_loras, size_t numel,
29
- int max_num_tokens_padded, int max_num_m_blocks,
30
- int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids,
31
- int topk_num, int32_t* total_tokens_post_pad, int32_t* adapter_enabled,
32
- int32_t* lora_ids) {
33
- const size_t tokens_per_thread = div_ceil(numel, blockDim.x);
34
- const size_t start_idx = threadIdx.x * tokens_per_thread;
35
-
36
- int lora_idx = blockIdx.x;
37
- int lora_id = lora_ids[lora_idx];
38
- if (lora_id == -1 || adapter_enabled[lora_id] == 0) {
39
- return;
40
- }
41
- extern __shared__ int32_t shared_mem[];
42
- int32_t* cumsum = shared_mem;
43
- token_cnts_t* tokens_cnts = (token_cnts_t*)(shared_mem + num_experts + 1);
44
-
45
- // Initialize sorted_token_ids with numel
46
- for (size_t it = threadIdx.x; it < max_num_tokens_padded; it += blockDim.x) {
47
- sorted_token_ids[lora_id * max_num_tokens_padded + it] = numel;
48
- }
49
-
50
- // Initialize expert_ids with -1
51
- for (size_t it = threadIdx.x; it < max_num_m_blocks; it += blockDim.x) {
52
- expert_ids[lora_id * max_num_m_blocks + it] = -1;
53
- }
54
-
55
- // Initialize total_tokens_post_pad with 0
56
- if (threadIdx.x == 0) {
57
- total_tokens_post_pad[lora_id] = 0;
58
- }
59
-
60
- for (int i = 0; i < num_experts; ++i) {
61
- tokens_cnts[index(num_experts, threadIdx.x + 1, i)] = 0;
62
- }
63
-
64
- for (int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) {
65
- int mask = token_lora_mapping[i / topk_num] == lora_id;
66
- int idx = index(num_experts, threadIdx.x + 1, topk_ids[i]);
67
- tokens_cnts[idx] += mask;
68
- }
69
-
70
- __syncthreads();
71
-
72
- // For each expert we accumulate the token counts from the different threads.
73
- if (threadIdx.x < num_experts) {
74
- tokens_cnts[index(num_experts, 0, threadIdx.x)] = 0;
75
- for (int i = 1; i <= blockDim.x; ++i) {
76
- tokens_cnts[index(num_experts, i, threadIdx.x)] +=
77
- tokens_cnts[index(num_experts, i - 1, threadIdx.x)];
78
- }
79
- }
80
-
81
- __syncthreads();
82
-
83
- // We accumulate the token counts of all experts in thread 0.
84
- if (threadIdx.x == 0) {
85
- cumsum[0] = 0;
86
- for (int i = 1; i <= num_experts; ++i) {
87
- cumsum[i] = cumsum[i - 1] +
88
- div_ceil(tokens_cnts[index(num_experts, blockDim.x, i - 1)],
89
- block_size) *
90
- block_size;
91
- }
92
- total_tokens_post_pad[lora_id] = static_cast<int32_t>(cumsum[num_experts]);
93
- }
94
-
95
- __syncthreads();
96
-
97
- /**
98
- * For each expert, each thread processes the tokens of the corresponding
99
- * blocks and stores the corresponding expert_id for each block.
100
- */
101
- if (threadIdx.x < num_experts) {
102
- for (int i = cumsum[threadIdx.x]; i < cumsum[threadIdx.x + 1];
103
- i += block_size) {
104
- expert_ids[index(max_num_m_blocks, lora_id, i / block_size)] =
105
- threadIdx.x;
106
- }
107
- }
108
-
109
- for (int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) {
110
- int32_t expert_id = topk_ids[i];
111
- /** The cumsum[expert_id] stores the starting index of the tokens that the
112
- * expert with expert_id needs to process, and
113
- * tokens_cnts[threadIdx.x][expert_id] stores the indices of the tokens
114
- * processed by the expert with expert_id within the current thread's token
115
- * shard.
116
- */
117
- int32_t rank_post_pad =
118
- tokens_cnts[index(num_experts, threadIdx.x, expert_id)] +
119
- cumsum[expert_id];
120
-
121
- int mask = (int)token_lora_mapping[i / topk_num] == lora_id;
122
- atomicAdd(
123
- &sorted_token_ids[index(max_num_tokens_padded, lora_id, rank_post_pad)],
124
- (i - numel) * mask);
125
- tokens_cnts[index(num_experts, threadIdx.x, expert_id)] += mask;
126
- }
127
-}
128
-
129
-void moe_lora_align_block_size(
130
- torch::Tensor topk_ids, torch::Tensor token_lora_mapping,
131
- int64_t num_experts, int64_t block_size, int64_t max_loras,
132
- int64_t max_num_tokens_padded, int64_t max_num_m_blocks,
133
- torch::Tensor sorted_token_ids, torch::Tensor expert_ids,
134
- torch::Tensor num_tokens_post_pad, torch::Tensor adapter_enabled,
135
- torch::Tensor lora_ids) {
136
- const int topk_num = topk_ids.size(1);
137
-
138
- TORCH_CHECK(block_size > 0, "block_size should be greater than 0. ");
139
-
140
- int device_max_shared_mem;
141
- auto dev = topk_ids.get_device();
142
- cudaDeviceGetAttribute(&device_max_shared_mem,
143
- cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
144
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
145
-
146
- const int32_t num_thread = max((int32_t)num_experts, 128); // WARP_SIZE,
147
- TORCH_CHECK(num_thread <= 1024,
148
- "num_thread must be less than 1024, "
149
- "and fallback is not implemented yet.");
150
- const int32_t shared_mem = (num_thread + 1) * num_experts * sizeof(int32_t) +
151
- (num_experts + 1) * sizeof(int32_t);
152
-
153
- if (shared_mem > device_max_shared_mem) {
154
- TORCH_CHECK(false,
155
- "Shared memory usage exceeds device limit, and global memory "
156
- "fallback is not implemented yet.");
157
- }
158
-
159
- VLLM_DISPATCH_INTEGRAL_TYPES(
160
- topk_ids.scalar_type(), "moe_lora_align_sum_kernel", [&] {
161
- dim3 blockDim(num_thread);
162
- auto kernel = moe_lora_align_sum_kernel<scalar_t, int32_t>;
163
- AT_CUDA_CHECK(VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(
164
- (void*)kernel, shared_mem));
165
- kernel<<<max_loras, blockDim, shared_mem, stream>>>(
166
- topk_ids.data_ptr<scalar_t>(),
167
- token_lora_mapping.data_ptr<int32_t>(), block_size, num_experts,
168
- max_loras, topk_ids.numel(), max_num_tokens_padded,
169
- max_num_m_blocks, sorted_token_ids.data_ptr<int32_t>(),
170
- expert_ids.data_ptr<int32_t>(), topk_num,
171
- num_tokens_post_pad.data_ptr<int32_t>(),
172
- adapter_enabled.data_ptr<int32_t>(), lora_ids.data_ptr<int32_t>());
173
- });
174
-}
\ No newline at end of file
csrc/moe/moe_ops.h
+1
-1
@@ -27,7 +27,7 @@ void moe_lora_align_block_size(
27
int64_t max_num_tokens_padded, int64_t max_num_m_blocks,
28
torch::Tensor sorted_token_ids, torch::Tensor expert_ids,
29
torch::Tensor num_tokens_post_pad, torch::Tensor adapter_enabled,
30
- torch::Tensor lora_ids);
30
+ torch::Tensor lora_ids, std::optional<torch::Tensor> maybe_expert_map);
31
#ifndef USE_ROCM
32
torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output,
33
torch::Tensor b_qweight, torch::Tensor b_scales,
csrc/moe/torch_bindings.cpp
+2
-1
@@ -47,7 +47,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
47
" Tensor !experts_ids,"
48
" Tensor !num_tokens_post_pad,"
49
" Tensor !adapter_enabled,"
50
- " Tensor !lora_ids) -> () ");
50
+ " Tensor !lora_ids,"
51
+ " Tensor? maybe_expert_map) -> () ");
52
m.impl("moe_lora_align_block_size", torch::kCUDA, &moe_lora_align_block_size);
53
54
#ifndef USE_ROCM
tests/lora/test_moe_lora_align_sum.py
+1
-1
@@ -32,7 +32,7 @@ def sample_data(num_experts, max_loras, num_tokens, topk_num):
32
33
@pytest.mark.parametrize("num_tokens", [100, 200, 1024, 4096]) # 81920
34
@pytest.mark.parametrize("topk_num", [6])
35
-@pytest.mark.parametrize("num_experts", [64, 128])
35
+@pytest.mark.parametrize("num_experts", [64, 128, 256, 512])
36
@pytest.mark.parametrize("max_loras", [2, 32])
37
@pytest.mark.parametrize("block_size", [16])
38
def test_moe_lora_align_block_size(
vllm/_custom_ops.py
+2
@@ -1961,6 +1961,7 @@ def moe_lora_align_block_size(
1961
num_tokens_post_pad: torch.Tensor,
1962
adapter_enabled: torch.Tensor,
1963
lora_ids: torch.Tensor,
1964
+ expert_map: torch.Tensor | None = None,
1965
) -> None:
1966
torch.ops._moe_C.moe_lora_align_block_size(
1967
topk_ids,
@@ -1975,6 +1976,7 @@ def moe_lora_align_block_size(
1976
num_tokens_post_pad,
1977
adapter_enabled,
1978
lora_ids,
1979
+ expert_map,
1980
)
1981
1982