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 15, 01:00 UTC
[perf] Speed up align sum kernels (#21079)
Signed-off-by: Himanshu Jaju <hj@mistral.ai>
Himanshu Jaju committed
Jul 21, 2025 at 19:19 UTC
0ec82edda59aaf5cf3b07aadf4ecce1aa1131add
3 files changed
+60
-25
benchmarks/kernels/benchmark_moe_align_block_size.py
+2
-5
@@ -33,15 +33,13 @@ def check_correctness(num_tokens, num_experts=256, block_size=256, topk=8):
33
sorted_ids_triton = torch.empty(
34
(max_num_tokens_padded,), dtype=torch.int32, device="cuda"
35
)
36
- sorted_ids_triton.fill_(topk_ids.numel()) # fill with sentinel value
37
- expert_ids_triton = torch.zeros(
36
+ expert_ids_triton = torch.empty(
37
(max_num_tokens_padded // block_size,), dtype=torch.int32, device="cuda"
38
)
39
num_tokens_post_pad_triton = torch.empty((1,), dtype=torch.int32, device="cuda")
40
41
sorted_ids_vllm = torch.empty_like(sorted_ids_triton)
43
- sorted_ids_vllm.fill_(topk_ids.numel())
44
- expert_ids_vllm = torch.zeros_like(expert_ids_triton)
42
+ expert_ids_vllm = torch.empty_like(expert_ids_triton)
43
num_tokens_post_pad_vllm = torch.empty_like(num_tokens_post_pad_triton)
44
45
# 2. run implementations
@@ -102,7 +100,6 @@ def benchmark(num_tokens, num_experts, topk, provider):
100
101
max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
102
sorted_ids = torch.empty((max_num_tokens_padded,), dtype=torch.int32, device="cuda")
105
- sorted_ids.fill_(topk_ids.numel())
103
max_num_m_blocks = max_num_tokens_padded // block_size
104
expert_ids = torch.empty((max_num_m_blocks,), dtype=torch.int32, device="cuda")
105
num_tokens_post_pad = torch.empty((1,), dtype=torch.int32, device="cuda")
csrc/moe/moe_align_sum_kernels.cu
+55
-16
@@ -1,6 +1,7 @@
1
#include <torch/all.h>
2
#include <ATen/cuda/CUDAContext.h>
3
#include <c10/cuda/CUDAGuard.h>
4
+#include <cub/cub.cuh>
5
6
#include <ATen/ATen.h>
7
#include <ATen/cuda/Atomic.cuh>
@@ -19,9 +20,14 @@ __global__ void moe_align_block_size_kernel(
20
int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids,
21
int32_t* __restrict__ total_tokens_post_pad, int32_t num_experts,
22
int32_t padded_num_experts, int32_t experts_per_warp, int32_t block_size,
22
- size_t numel, int32_t* __restrict__ cumsum) {
23
+ size_t numel, int32_t* __restrict__ cumsum, int32_t max_num_tokens_padded) {
24
extern __shared__ int32_t shared_counts[];
25
26
+ // Initialize sorted_token_ids with numel
27
+ for (size_t it = threadIdx.x; it < max_num_tokens_padded; it += blockDim.x) {
28
+ sorted_token_ids[it] = numel;
29
+ }
30
+
31
const int warp_id = threadIdx.x / WARP_SIZE;
32
const int my_expert_start = warp_id * experts_per_warp;
33
@@ -45,18 +51,27 @@ __global__ void moe_align_block_size_kernel(
51
52
__syncthreads();
53
48
- if (threadIdx.x == 0) {
49
- cumsum[0] = 0;
50
- for (int i = 1; i <= num_experts; ++i) {
51
- int expert_count = 0;
52
- int warp_idx = (i - 1) / experts_per_warp;
53
- int expert_offset = (i - 1) % experts_per_warp;
54
- expert_count = shared_counts[warp_idx * experts_per_warp + expert_offset];
54
+ // Compute prefix sum over token counts per expert
55
+ using BlockScan = cub::BlockScan<int32_t, 1024>;
56
+ __shared__ typename BlockScan::TempStorage temp_storage;
57
56
- cumsum[i] =
57
- cumsum[i - 1] + CEILDIV(expert_count, block_size) * block_size;
58
- }
59
- *total_tokens_post_pad = cumsum[num_experts];
58
+ int expert_count = 0;
59
+ int expert_id = threadIdx.x;
60
+ if (expert_id < num_experts) {
61
+ int warp_idx = expert_id / experts_per_warp;
62
+ int expert_offset = expert_id % experts_per_warp;
63
+ expert_count = shared_counts[warp_idx * experts_per_warp + expert_offset];
64
+ expert_count = CEILDIV(expert_count, block_size) * block_size;
65
+ }
66
+
67
+ int cumsum_val;
68
+ BlockScan(temp_storage).ExclusiveSum(expert_count, cumsum_val);
69
+ if (expert_id <= num_experts) {
70
+ cumsum[expert_id] = cumsum_val;
71
+ }
72
+
73
+ if (expert_id == num_experts) {
74
+ *total_tokens_post_pad = cumsum_val;
75
}
76
77
__syncthreads();
@@ -67,6 +82,13 @@ __global__ void moe_align_block_size_kernel(
82
expert_ids[i / block_size] = threadIdx.x;
83
}
84
}
85
+
86
+ // Fill remaining expert_ids with 0
87
+ const size_t fill_start_idx = cumsum[num_experts] / block_size + threadIdx.x;
88
+ const size_t expert_ids_size = CEILDIV(max_num_tokens_padded, block_size);
89
+ for (size_t i = fill_start_idx; i < expert_ids_size; i += blockDim.x) {
90
+ expert_ids[i] = 0;
91
+ }
92
}
93
94
template <typename scalar_t>
@@ -105,7 +127,12 @@ __global__ void moe_align_block_size_small_batch_expert_kernel(
127
const scalar_t* __restrict__ topk_ids,
128
int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids,
129
int32_t* __restrict__ total_tokens_post_pad, int32_t num_experts,
108
- int32_t block_size, size_t numel) {
130
+ int32_t block_size, size_t numel, int32_t max_num_tokens_padded) {
131
+ // Initialize sorted_token_ids with numel
132
+ for (size_t it = threadIdx.x; it < max_num_tokens_padded; it += blockDim.x) {
133
+ sorted_token_ids[it] = numel;
134
+ }
135
+
136
const size_t tid = threadIdx.x;
137
const size_t stride = blockDim.x;
138
@@ -153,6 +180,13 @@ __global__ void moe_align_block_size_small_batch_expert_kernel(
180
}
181
}
182
183
+ // Fill remaining expert_ids with 0
184
+ const size_t fill_start_idx = cumsum[num_experts] / block_size + threadIdx.x;
185
+ const size_t expert_ids_size = CEILDIV(max_num_tokens_padded, block_size);
186
+ for (size_t i = fill_start_idx; i < expert_ids_size; i += blockDim.x) {
187
+ expert_ids[i] = 0;
188
+ }
189
+
190
for (size_t i = tid; i < numel; i += stride) {
191
int32_t expert_id = topk_ids[i];
192
int32_t rank_post_pad =
@@ -179,13 +213,17 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
213
int threads = 1024;
214
threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
215
216
+ // BlockScan uses 1024 threads and assigns one thread per expert.
217
+ TORCH_CHECK(padded_num_experts < 1024,
218
+ "padded_num_experts must be less than 1024");
219
+
220
VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(
221
topk_ids.scalar_type(), "moe_align_block_size_kernel", [&] {
222
// calc needed amount of shared mem for `cumsum` tensors
223
auto options_int =
224
torch::TensorOptions().dtype(torch::kInt).device(topk_ids.device());
225
torch::Tensor cumsum_buffer =
188
- torch::zeros({num_experts + 1}, options_int);
226
+ torch::empty({num_experts + 1}, options_int);
227
bool small_batch_expert_mode =
228
(topk_ids.numel() < 1024) && (num_experts <= 64);
229
@@ -203,7 +241,7 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
241
sorted_token_ids.data_ptr<int32_t>(),
242
experts_ids.data_ptr<int32_t>(),
243
num_tokens_post_pad.data_ptr<int32_t>(), num_experts, block_size,
206
- topk_ids.numel());
244
+ topk_ids.numel(), sorted_token_ids.size(0));
245
} else {
246
auto align_kernel = vllm::moe::moe_align_block_size_kernel<scalar_t>;
247
@@ -217,7 +255,8 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
255
experts_ids.data_ptr<int32_t>(),
256
num_tokens_post_pad.data_ptr<int32_t>(), num_experts,
257
padded_num_experts, experts_per_warp, block_size,
220
- topk_ids.numel(), cumsum_buffer.data_ptr<int32_t>());
258
+ topk_ids.numel(), cumsum_buffer.data_ptr<int32_t>(),
259
+ sorted_token_ids.size(0));
260
261
const int block_threads = std::min(256, (int)threads);
262
const int num_blocks =
vllm/model_executor/layers/fused_moe/moe_align_block_size.py
+3
-4
@@ -111,6 +111,8 @@ def moe_align_block_size_triton(
111
dtype=torch.int32,
112
device=topk_ids.device)
113
tokens_per_thread = cdiv(numel, num_experts)
114
+ sorted_token_ids.fill_(numel)
115
+ expert_ids.zero_()
116
117
moe_align_block_size_stage1[grid](
118
topk_ids,
@@ -205,11 +207,8 @@ def moe_align_block_size(
207
sorted_ids = torch.empty((max_num_tokens_padded, ),
208
dtype=torch.int32,
209
device=topk_ids.device)
208
- sorted_ids.fill_(topk_ids.numel())
210
max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
210
- # Expert ids must be zeroed out to prevent index out of bounds error while
211
- # mapping global expert ids to local expert ids in expert parallelism.
212
- expert_ids = torch.zeros((max_num_m_blocks, ),
211
+ expert_ids = torch.empty((max_num_m_blocks, ),
212
dtype=torch.int32,
213
device=topk_ids.device)
214
num_tokens_post_pad = torch.empty((1),