@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ffb2cd6b54
Fork 0

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, 06:00 UTC

[Perf] Optimize `moe_align_block_size` CUDA kernel (#19572)

Signed-off-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: mgoin <mgoin64@gmail.com>

Wentao Ye committed Jun 17, 2025 at 14:49 UTC ffb2cd6b5441af28da22694a19838a19507a15e1
7 files changed +386 -365
benchmarks/kernels/benchmark_moe_align_block_size.py new
+159
@@ -0,0 +1,159 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import argparse
4 +import itertools
5 +
6 +import torch
7 +import triton
8 +
9 +from vllm import _custom_ops as ops
10 +from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
11 + moe_align_block_size_triton,
12 +)
13 +
14 +
15 +def get_topk_ids(num_tokens: int, num_experts: int, topk: int) -> torch.Tensor:
16 + return torch.stack(
17 + [
18 + torch.randperm(num_experts, dtype=torch.int32, device="cuda")[:topk]
19 + for _ in range(num_tokens)
20 + ]
21 + )
22 +
23 +
24 +def check_correctness(num_tokens, num_experts=256, block_size=256, topk=8):
25 + """
26 + Verifies vllm vs. Triton
27 + """
28 + topk_ids = get_topk_ids(num_tokens, num_experts, topk)
29 +
30 + # 1. malloc space for triton and vllm
31 + # malloc enough space (max_num_tokens_padded) for the sorted ids
32 + max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
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(
38 + (max_num_tokens_padded // block_size,), dtype=torch.int32, device="cuda"
39 + )
40 + num_tokens_post_pad_triton = torch.empty((1,), dtype=torch.int32, device="cuda")
41 +
42 + 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)
45 + num_tokens_post_pad_vllm = torch.empty_like(num_tokens_post_pad_triton)
46 +
47 + # 2. run implementations
48 + moe_align_block_size_triton(
49 + topk_ids,
50 + num_experts,
51 + block_size,
52 + sorted_ids_triton,
53 + expert_ids_triton,
54 + num_tokens_post_pad_triton,
55 + )
56 +
57 + ops.moe_align_block_size(
58 + topk_ids,
59 + num_experts,
60 + block_size,
61 + sorted_ids_vllm,
62 + expert_ids_vllm,
63 + num_tokens_post_pad_vllm,
64 + )
65 + print(f"✅ VLLM implementation works with {num_experts} experts!")
66 +
67 + # 3. compare results
68 + if torch.allclose(expert_ids_triton, expert_ids_vllm) and torch.allclose(
69 + num_tokens_post_pad_triton, num_tokens_post_pad_vllm
70 + ):
71 + print("✅ Triton and VLLM implementations match.")
72 + else:
73 + print("❌ Triton and VLLM implementations DO NOT match.")
74 + print("Triton expert_ids:", expert_ids_triton)
75 + print("VLLM expert_ids:", expert_ids_vllm)
76 + print("Triton num_tokens_post_pad:", num_tokens_post_pad_triton)
77 + print("VLLM num_tokens_post_pad:", num_tokens_post_pad_vllm)
78 +
79 +
80 +# test configurations
81 +num_tokens_range = [1, 16, 256, 4096]
82 +num_experts_range = [16, 64, 224, 256, 280, 512]
83 +topk_range = [1, 2, 8]
84 +configs = list(itertools.product(num_tokens_range, num_experts_range, topk_range))
85 +
86 +
87 +@triton.testing.perf_report(
88 + triton.testing.Benchmark(
89 + x_names=["num_tokens", "num_experts", "topk"],
90 + x_vals=configs,
91 + line_arg="provider",
92 + line_vals=["vllm", "triton"], # "triton"
93 + line_names=["VLLM", "Triton"], # "Triton"
94 + plot_name="moe-align-block-size-performance",
95 + args={},
96 + )
97 +)
98 +def benchmark(num_tokens, num_experts, topk, provider):
99 + """Benchmark function for Triton."""
100 + block_size = 256
101 + topk_ids = get_topk_ids(num_tokens, num_experts, topk)
102 +
103 + max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
104 + sorted_ids = torch.empty((max_num_tokens_padded,), dtype=torch.int32, device="cuda")
105 + sorted_ids.fill_(topk_ids.numel())
106 + max_num_m_blocks = max_num_tokens_padded // block_size
107 + expert_ids = torch.empty((max_num_m_blocks,), dtype=torch.int32, device="cuda")
108 + num_tokens_post_pad = torch.empty((1,), dtype=torch.int32, device="cuda")
109 +
110 + quantiles = [0.5, 0.2, 0.8]
111 +
112 + if provider == "vllm":
113 + ms, min_ms, max_ms = triton.testing.do_bench(
114 + lambda: ops.moe_align_block_size(
115 + topk_ids,
116 + num_experts,
117 + block_size,
118 + sorted_ids.clone(),
119 + expert_ids.clone(),
120 + num_tokens_post_pad.clone(),
121 + ),
122 + quantiles=quantiles,
123 + )
124 + elif provider == "triton":
125 + ms, min_ms, max_ms = triton.testing.do_bench(
126 + lambda: moe_align_block_size_triton(
127 + topk_ids,
128 + num_experts,
129 + block_size,
130 + sorted_ids.clone(),
131 + expert_ids.clone(),
132 + num_tokens_post_pad.clone(),
133 + ),
134 + quantiles=quantiles,
135 + )
136 +
137 + return 1000 * ms, 1000 * max_ms, 1000 * min_ms
138 +
139 +
140 +if __name__ == "__main__":
141 + parser = argparse.ArgumentParser()
142 + parser.add_argument(
143 + "--num_experts",
144 + type=int,
145 + default=64,
146 + choices=[8, 16, 32, 64, 128, 256],
147 + )
148 + parser.add_argument(
149 + "--topk",
150 + type=int,
151 + default=8,
152 + choices=[2, 4, 8],
153 + help="Top-k value for correctness check.",
154 + )
155 + args = parser.parse_args()
156 +
157 + print("Running correctness check...")
158 + check_correctness(num_tokens=1024, num_experts=args.num_experts, topk=args.topk)
159 + benchmark.run(print_data=True, show_plots=True)
csrc/moe/moe_align_sum_kernels.cu
+133 -317
@@ -13,232 +13,45 @@
13 namespace vllm {
14 namespace moe {
15
16 -namespace {
17 -__device__ __forceinline__ int32_t index(int32_t total_col, int32_t row,
18 - int32_t col) {
19 - // don't worry about overflow because num_experts is relatively small
20 - return row * total_col + col;
21 -}
22 -} // namespace
23 -
24 -template <typename scalar_t, typename token_cnts_t>
25 -__global__ void moe_align_block_size_kernel(scalar_t* __restrict__ topk_ids,
26 - int32_t* sorted_token_ids,
27 - int32_t* expert_ids,
28 - int32_t* total_tokens_post_pad,
29 - int32_t num_experts,
30 - int32_t block_size, size_t numel) {
31 - const size_t tokens_per_thread = CEILDIV(numel, blockDim.x);
32 - const size_t start_idx = threadIdx.x * tokens_per_thread;
33 -
34 - extern __shared__ int32_t shared_mem[];
35 - int32_t* cumsum = shared_mem; // 1d tensor with shape (num_experts + 1)
36 - token_cnts_t* tokens_cnts =
37 - (token_cnts_t*)(shared_mem + num_experts +
38 - 1); // 2d tensor with shape (blockDim.x + 1, num_experts)
39 -
40 - for (int i = 0; i < num_experts; ++i) {
41 - tokens_cnts[index(num_experts, threadIdx.x + 1, i)] = 0;
42 - }
43 -
44 - /**
45 - * In the first step we compute token_cnts[thread_index + 1][expert_index],
46 - * which counts how many tokens in the token shard of thread_index are
47 - * assigned to expert expert_index.
48 - */
49 - for (int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) {
50 - ++tokens_cnts[index(num_experts, threadIdx.x + 1, topk_ids[i])];
51 - }
52 -
53 - __syncthreads();
54 -
55 - // For each expert we accumulate the token counts from the different threads.
56 - if (threadIdx.x < num_experts) {
57 - tokens_cnts[index(num_experts, 0, threadIdx.x)] = 0;
58 - for (int i = 1; i <= blockDim.x; ++i) {
59 - tokens_cnts[index(num_experts, i, threadIdx.x)] +=
60 - tokens_cnts[index(num_experts, i - 1, threadIdx.x)];
61 - }
62 - }
63 -
64 - __syncthreads();
65 -
66 - // We accumulate the token counts of all experts in thread 0.
67 - if (threadIdx.x == 0) {
68 - cumsum[0] = 0;
69 - for (int i = 1; i <= num_experts; ++i) {
70 - cumsum[i] = cumsum[i - 1] +
71 - CEILDIV(tokens_cnts[index(num_experts, blockDim.x, i - 1)],
72 - block_size) *
73 - block_size;
74 - }
75 - *total_tokens_post_pad = static_cast<int32_t>(cumsum[num_experts]);
76 - }
77 -
78 - __syncthreads();
79 -
80 - /**
81 - * For each expert, each thread processes the tokens of the corresponding
82 - * blocks and stores the corresponding expert_id for each block.
83 - */
84 - if (threadIdx.x < num_experts) {
85 - for (int i = cumsum[threadIdx.x]; i < cumsum[threadIdx.x + 1];
86 - i += block_size) {
87 - expert_ids[i / block_size] = threadIdx.x;
88 - }
89 - }
90 -
91 - /**
92 - * Each thread processes a token shard, calculating the index of each token
93 - * after sorting by expert number. Given the example topk_ids =
94 - * [0,1,2,1,2,3,0,3,4] and block_size = 4, then the output would be [0, 6, *,
95 - * *, 1, 3, *, *, 2, 4, *, *, 5, 7, *, *, 8, *, *, *], where * represents a
96 - * padding value(preset in python).
97 - */
98 - for (int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) {
99 - int32_t expert_id = topk_ids[i];
100 - /** The cumsum[expert_id] stores the starting index of the tokens that the
101 - * expert with expert_id needs to process, and
102 - * tokens_cnts[threadIdx.x][expert_id] stores the indices of the tokens
103 - * processed by the expert with expert_id within the current thread's token
104 - * shard.
105 - */
106 - int32_t rank_post_pad =
107 - tokens_cnts[index(num_experts, threadIdx.x, expert_id)] +
108 - cumsum[expert_id];
109 - sorted_token_ids[rank_post_pad] = i;
110 - ++tokens_cnts[index(num_experts, threadIdx.x, expert_id)];
111 - }
112 -}
113 -
114 -// TODO(simon): this is temporarily adapted from
115 -// https://github.com/sgl-project/sglang/commit/31548116a8dc8c6df7e146e0587335a59fc5b9d7
116 -// we did this to unblock Deepseek V3 but there should be a better
117 -// implementation to manage shared memory.
118 -template <typename scalar_t>
119 -__global__ void moe_align_block_size_global_mem_kernel(
120 - scalar_t* __restrict__ topk_ids, int32_t* sorted_token_ids,
121 - int32_t* expert_ids, int32_t* total_tokens_post_pad, int32_t num_experts,
122 - int32_t block_size, size_t numel, int32_t* tokens_cnts, int32_t* cumsum) {
123 - const size_t tokens_per_thread = CEILDIV(numel, blockDim.x);
124 - const size_t start_idx = threadIdx.x * tokens_per_thread;
125 -
126 - for (int i = 0; i < num_experts; ++i) {
127 - tokens_cnts[index(num_experts, threadIdx.x + 1, i)] = 0;
128 - }
129 -
130 - /**
131 - * In the first step we compute token_cnts[thread_index + 1][expert_index],
132 - * which counts how many tokens in the token shard of thread_index are
133 - * assigned to expert expert_index.
134 - */
135 - for (int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) {
136 - ++tokens_cnts[index(num_experts, threadIdx.x + 1, topk_ids[i])];
137 - }
138 -
139 - __syncthreads();
140 -
141 - // For each expert we accumulate the token counts from the different threads.
142 - if (threadIdx.x < num_experts) {
143 - tokens_cnts[index(num_experts, 0, threadIdx.x)] = 0;
144 - for (int i = 1; i <= blockDim.x; ++i) {
145 - tokens_cnts[index(num_experts, i, threadIdx.x)] +=
146 - tokens_cnts[index(num_experts, i - 1, threadIdx.x)];
147 - }
148 - }
149 -
150 - __syncthreads();
151 -
152 - // We accumulate the token counts of all experts in thread 0.
153 - if (threadIdx.x == 0) {
154 - cumsum[0] = 0;
155 - for (int i = 1; i <= num_experts; ++i) {
156 - cumsum[i] = cumsum[i - 1] +
157 - CEILDIV(tokens_cnts[index(num_experts, blockDim.x, i - 1)],
158 - block_size) *
159 - block_size;
160 - }
161 - *total_tokens_post_pad = cumsum[num_experts];
162 - }
163 -
164 - __syncthreads();
165 -
166 - /**
167 - * For each expert, each thread processes the tokens of the corresponding
168 - * blocks and stores the corresponding expert_id for each block.
169 - */
170 - if (threadIdx.x < num_experts) {
171 - for (int i = cumsum[threadIdx.x]; i < cumsum[threadIdx.x + 1];
172 - i += block_size) {
173 - expert_ids[i / block_size] = threadIdx.x;
174 - }
175 - }
176 -
177 - /**
178 - * Each thread processes a token shard, calculating the index of each token
179 - * after sorting by expert number. Given the example topk_ids =
180 - * [0,1,2,1,2,3,0,3,4] and block_size = 4, then the output would be [0, 6, *,
181 - * *, 1, 3, *, *, 2, 4, *, *, 5, 7, *, *, 8, *, *, *], where * represents a
182 - * padding value(preset in python).
183 - */
184 - for (int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) {
185 - int32_t expert_id = topk_ids[i];
186 - /** The cumsum[expert_id] stores the starting index of the tokens that the
187 - * expert with expert_id needs to process, and
188 - * tokens_cnts[threadIdx.x][expert_id] stores the indices of the tokens
189 - * processed by the expert with expert_id within the current thread's token
190 - * shard.
191 - */
192 - int32_t rank_post_pad =
193 - tokens_cnts[index(num_experts, threadIdx.x, expert_id)] +
194 - cumsum[expert_id];
195 - sorted_token_ids[rank_post_pad] = i;
196 - ++tokens_cnts[index(num_experts, threadIdx.x, expert_id)];
197 - }
198 -}
199 -
200 -// taken from
201 -// https://github.com/sgl-project/sglang/commit/cdae77b03dfc6fec3863630550b45bbfc789f957
16 template <typename scalar_t>
203 -__global__ void sgl_moe_align_block_size_kernel(
204 - scalar_t* __restrict__ topk_ids, int32_t* sorted_token_ids,
205 - int32_t* expert_ids, int32_t* total_tokens_post_pad, int32_t num_experts,
206 - int32_t block_size, size_t numel, int32_t* cumsum) {
207 - __shared__ int32_t shared_counts[32][8];
208 -
209 - const int warp_id = threadIdx.x / 32;
210 - const int experts_per_warp = 8;
17 +__global__ void moe_align_block_size_kernel(
18 + const scalar_t* __restrict__ topk_ids,
19 + int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids,
20 + int32_t* __restrict__ total_tokens_post_pad, int32_t num_experts,
21 + int32_t padded_num_experts, int32_t experts_per_warp, int32_t block_size,
22 + size_t numel, int32_t* __restrict__ cumsum) {
23 + extern __shared__ int32_t shared_counts[];
24 +
25 + const int warp_id = threadIdx.x / WARP_SIZE;
26 const int my_expert_start = warp_id * experts_per_warp;
27
213 - // Initialize shared_counts for this warp's experts
28 for (int i = 0; i < experts_per_warp; ++i) {
215 - if (my_expert_start + i < num_experts) {
216 - shared_counts[warp_id][i] = 0;
29 + if (my_expert_start + i < padded_num_experts) {
30 + shared_counts[warp_id * experts_per_warp + i] = 0;
31 }
32 }
33
34 __syncthreads();
35
222 - const size_t tokens_per_thread = CEILDIV(numel, blockDim.x);
223 - const size_t start_idx = threadIdx.x * tokens_per_thread;
36 + const size_t tid = threadIdx.x;
37 + const size_t stride = blockDim.x;
38
225 - for (int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) {
39 + for (size_t i = tid; i < numel; i += stride) {
40 int expert_id = topk_ids[i];
41 int warp_idx = expert_id / experts_per_warp;
42 int expert_offset = expert_id % experts_per_warp;
229 - atomicAdd(&shared_counts[warp_idx][expert_offset], 1);
43 + atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset], 1);
44 }
45
46 __syncthreads();
47
234 - // Single thread computes cumulative sum and total tokens
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;
241 - expert_count = shared_counts[warp_idx][expert_offset];
54 + expert_count = shared_counts[warp_idx * experts_per_warp + expert_offset];
55
56 cumsum[i] =
57 cumsum[i - 1] + CEILDIV(expert_count, block_size) * block_size;
@@ -248,7 +61,6 @@ __global__ void sgl_moe_align_block_size_kernel(
61
62 __syncthreads();
63
251 - // Assign expert IDs to blocks
64 if (threadIdx.x < num_experts) {
65 for (int i = cumsum[threadIdx.x]; i < cumsum[threadIdx.x + 1];
66 i += block_size) {
@@ -257,13 +69,11 @@ __global__ void sgl_moe_align_block_size_kernel(
69 }
70 }
71
260 -// taken from
261 -// https://github.com/sgl-project/sglang/commit/cdae77b03dfc6fec3863630550b45bbfc789f957
72 template <typename scalar_t>
263 -__global__ void sgl_moe_token_sort_kernel(scalar_t* __restrict__ topk_ids,
264 - int32_t* sorted_token_ids,
265 - int32_t* cumsum_buffer,
266 - size_t numel) {
73 +__global__ void count_and_sort_expert_tokens_kernel(
74 + const scalar_t* __restrict__ topk_ids,
75 + int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer,
76 + size_t numel) {
77 const size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
78 const size_t stride = blockDim.x * gridDim.x;
79
@@ -290,132 +100,138 @@ __global__ void moe_sum_kernel(
100 }
101 }
102
103 +template <typename scalar_t>
104 +__global__ void moe_align_block_size_small_batch_expert_kernel(
105 + const scalar_t* __restrict__ topk_ids,
106 + int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids,
107 + int32_t* __restrict__ total_tokens_post_pad, int32_t num_experts,
108 + int32_t block_size, size_t numel) {
109 + const size_t tid = threadIdx.x;
110 + const size_t stride = blockDim.x;
111 +
112 + extern __shared__ int32_t shared_mem[];
113 + int32_t* cumsum = shared_mem;
114 + int32_t* tokens_cnts = (int32_t*)(shared_mem + num_experts + 1);
115 +
116 + for (int i = 0; i < num_experts; ++i) {
117 + tokens_cnts[(threadIdx.x + 1) * num_experts + i] = 0;
118 + }
119 +
120 + for (size_t i = tid; i < numel; i += stride) {
121 + ++tokens_cnts[(threadIdx.x + 1) * num_experts + topk_ids[i]];
122 + }
123 +
124 + __syncthreads();
125 +
126 + if (threadIdx.x < num_experts) {
127 + tokens_cnts[threadIdx.x] = 0;
128 + for (int i = 1; i <= blockDim.x; ++i) {
129 + tokens_cnts[i * num_experts + threadIdx.x] +=
130 + tokens_cnts[(i - 1) * num_experts + threadIdx.x];
131 + }
132 + }
133 +
134 + __syncthreads();
135 +
136 + if (threadIdx.x == 0) {
137 + cumsum[0] = 0;
138 + for (int i = 1; i <= num_experts; ++i) {
139 + cumsum[i] =
140 + cumsum[i - 1] +
141 + CEILDIV(tokens_cnts[blockDim.x * num_experts + i - 1], block_size) *
142 + block_size;
143 + }
144 + *total_tokens_post_pad = static_cast<int32_t>(cumsum[num_experts]);
145 + }
146 +
147 + __syncthreads();
148 +
149 + if (threadIdx.x < num_experts) {
150 + for (int i = cumsum[threadIdx.x]; i < cumsum[threadIdx.x + 1];
151 + i += block_size) {
152 + expert_ids[i / block_size] = threadIdx.x;
153 + }
154 + }
155 +
156 + for (size_t i = tid; i < numel; i += stride) {
157 + int32_t expert_id = topk_ids[i];
158 + int32_t rank_post_pad =
159 + tokens_cnts[threadIdx.x * num_experts + expert_id] + cumsum[expert_id];
160 + sorted_token_ids[rank_post_pad] = i;
161 + ++tokens_cnts[threadIdx.x * num_experts + expert_id];
162 + }
163 +}
164 +
165 } // namespace moe
166 } // namespace vllm
167
168 +// taken from
169 +// https://github.com/sgl-project/sglang/blob/8b5f83ed3b7d2a49ad5c5cd5aa61c5d502f47dbc
170 void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
171 int64_t block_size, torch::Tensor sorted_token_ids,
172 torch::Tensor experts_ids,
173 torch::Tensor num_tokens_post_pad) {
174 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
175
302 - int device_max_shared_mem;
303 - auto dev = topk_ids.get_device();
304 - cudaDeviceGetAttribute(&device_max_shared_mem,
305 - cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
306 -
307 - const int32_t num_thread = max((int32_t)num_experts, WARP_SIZE);
308 - const int32_t shared_mem_i32 =
309 - ((num_thread + 1) * num_experts + (num_experts + 1)) * sizeof(int32_t);
310 - const int32_t shared_mem_i16 =
311 - ((num_thread + 1) * num_experts) * sizeof(uint16_t) +
312 - (num_experts + 1) * sizeof(int32_t);
313 -
314 - bool use_global_memory = false;
315 - bool use_i16 = false; // Use uint16_t for shared memory token counts
316 - if (shared_mem_i32 < device_max_shared_mem) {
317 - // Do nothing in this case. We're all set to use int32_t token counts
318 - } else if (shared_mem_i16 < device_max_shared_mem &&
319 - topk_ids.numel() <= 65535) {
320 - // when nelements of topk_ids is smaller than 65535 (max value of uint16),
321 - // element value of token_cnts would also smaller than 65535,
322 - // so we can use uint16 as dtype of token_cnts
323 - use_i16 = true;
324 - } else {
325 - use_global_memory = true;
326 - }
176 + int64_t padded_num_experts =
177 + ((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
178 + int experts_per_warp = WARP_SIZE;
179 + int threads = 1024;
180 + threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
181
328 - if (use_global_memory) {
329 - VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(
330 - topk_ids.scalar_type(), "moe_align_block_size_global_mem_kernel", [&] {
331 - // calc needed amount of shared mem for `tokens_cnts` and `cumsum`
332 - // tensors
333 - const int32_t num_thread = max((int32_t)num_experts, WARP_SIZE);
334 -
335 - auto options_int = torch::TensorOptions()
336 - .dtype(torch::kInt)
337 - .device(topk_ids.device());
338 - torch::Tensor token_cnts_buffer =
339 - torch::empty({(num_experts + 1) * num_experts}, options_int);
340 - torch::Tensor cumsum_buffer =
341 - torch::empty({num_experts + 1}, options_int);
342 -
343 - auto kernel =
344 - vllm::moe::moe_align_block_size_global_mem_kernel<scalar_t>;
345 - kernel<<<1, num_thread, 0, stream>>>(
182 + VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(
183 + topk_ids.scalar_type(), "moe_align_block_size_kernel", [&] {
184 + // calc needed amount of shared mem for `cumsum` tensors
185 + auto options_int =
186 + torch::TensorOptions().dtype(torch::kInt).device(topk_ids.device());
187 + torch::Tensor cumsum_buffer =
188 + torch::zeros({num_experts + 1}, options_int);
189 + bool small_batch_expert_mode =
190 + (topk_ids.numel() < 1024) && (num_experts <= 64);
191 +
192 + if (small_batch_expert_mode) {
193 + const int32_t threads = max((int32_t)num_experts, WARP_SIZE);
194 + const int32_t shared_mem_size =
195 + ((threads + 1) * num_experts + (num_experts + 1)) *
196 + sizeof(int32_t);
197 +
198 + auto small_batch_expert_kernel =
199 + vllm::moe::moe_align_block_size_small_batch_expert_kernel<
200 + scalar_t>;
201 + small_batch_expert_kernel<<<1, threads, shared_mem_size, stream>>>(
202 topk_ids.data_ptr<scalar_t>(),
203 sorted_token_ids.data_ptr<int32_t>(),
204 experts_ids.data_ptr<int32_t>(),
205 num_tokens_post_pad.data_ptr<int32_t>(), num_experts, block_size,
350 - topk_ids.numel(), token_cnts_buffer.data_ptr<int32_t>(),
351 - cumsum_buffer.data_ptr<int32_t>());
352 - });
353 - } else if (use_i16) {
354 - VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(
355 - topk_ids.scalar_type(), "moe_align_block_size_kernel", [&] {
356 - // set dynamic shared mem
357 - auto kernel =
358 - vllm::moe::moe_align_block_size_kernel<scalar_t, uint16_t>;
359 - AT_CUDA_CHECK(VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(
360 - (void*)kernel, shared_mem_i16));
361 - kernel<<<1, num_thread, shared_mem_i16, stream>>>(
206 + topk_ids.numel());
207 + } else {
208 + auto align_kernel = vllm::moe::moe_align_block_size_kernel<scalar_t>;
209 +
210 + size_t num_warps = CEILDIV(padded_num_experts, experts_per_warp);
211 + size_t shared_mem_size =
212 + num_warps * experts_per_warp * sizeof(int32_t);
213 +
214 + align_kernel<<<1, threads, shared_mem_size, stream>>>(
215 topk_ids.data_ptr<scalar_t>(),
216 sorted_token_ids.data_ptr<int32_t>(),
217 experts_ids.data_ptr<int32_t>(),
365 - num_tokens_post_pad.data_ptr<int32_t>(), num_experts, block_size,
366 - topk_ids.numel());
367 - });
368 - } else {
369 - VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(
370 - topk_ids.scalar_type(), "moe_align_block_size_kernel", [&] {
371 - auto kernel =
372 - vllm::moe::moe_align_block_size_kernel<scalar_t, int32_t>;
373 - AT_CUDA_CHECK(VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(
374 - (void*)kernel, shared_mem_i32));
375 - kernel<<<1, num_thread, shared_mem_i32, stream>>>(
218 + num_tokens_post_pad.data_ptr<int32_t>(), num_experts,
219 + padded_num_experts, experts_per_warp, block_size,
220 + topk_ids.numel(), cumsum_buffer.data_ptr<int32_t>());
221 +
222 + const int block_threads = std::min(256, (int)threads);
223 + const int num_blocks =
224 + (topk_ids.numel() + block_threads - 1) / block_threads;
225 + const int max_blocks = 65535;
226 + const int actual_blocks = std::min(num_blocks, max_blocks);
227 +
228 + auto sort_kernel =
229 + vllm::moe::count_and_sort_expert_tokens_kernel<scalar_t>;
230 + sort_kernel<<<actual_blocks, block_threads, 0, stream>>>(
231 topk_ids.data_ptr<scalar_t>(),
232 sorted_token_ids.data_ptr<int32_t>(),
378 - experts_ids.data_ptr<int32_t>(),
379 - num_tokens_post_pad.data_ptr<int32_t>(), num_experts, block_size,
380 - topk_ids.numel());
381 - });
382 - }
383 -}
384 -
385 -void sgl_moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
386 - int64_t block_size,
387 - torch::Tensor sorted_token_ids,
388 - torch::Tensor experts_ids,
389 - torch::Tensor num_tokens_post_pad) {
390 - const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
391 - TORCH_CHECK(num_experts == 256,
392 - "sgl_moe_align_block_size kernel only supports deepseek v3.");
393 -
394 - VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(
395 - topk_ids.scalar_type(), "sgl_moe_align_block_size_kernel", [&] {
396 - // calc needed amount of shared mem for `cumsum` tensors
397 - auto options_int =
398 - torch::TensorOptions().dtype(torch::kInt).device(topk_ids.device());
399 - torch::Tensor cumsum_buffer =
400 - torch::zeros({num_experts + 1}, options_int);
401 -
402 - auto align_kernel =
403 - vllm::moe::sgl_moe_align_block_size_kernel<scalar_t>;
404 - align_kernel<<<1, 1024, 0, stream>>>(
405 - topk_ids.data_ptr<scalar_t>(), sorted_token_ids.data_ptr<int32_t>(),
406 - experts_ids.data_ptr<int32_t>(),
407 - num_tokens_post_pad.data_ptr<int32_t>(), num_experts, block_size,
408 - topk_ids.numel(), cumsum_buffer.data_ptr<int32_t>());
409 -
410 - const int block_threads = 256;
411 - const int num_blocks =
412 - (topk_ids.numel() + block_threads - 1) / block_threads;
413 - const int max_blocks = 65535;
414 - const int actual_blocks = std::min(num_blocks, max_blocks);
415 - auto sort_kernel = vllm::moe::sgl_moe_token_sort_kernel<scalar_t>;
416 - sort_kernel<<<actual_blocks, block_threads, 0, stream>>>(
417 - topk_ids.data_ptr<scalar_t>(), sorted_token_ids.data_ptr<int32_t>(),
418 - cumsum_buffer.data_ptr<int32_t>(), topk_ids.numel());
233 + cumsum_buffer.data_ptr<int32_t>(), topk_ids.numel());
234 + }
235 });
236 }
237
csrc/moe/moe_ops.h
-6
@@ -12,12 +12,6 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
12 int64_t block_size, torch::Tensor sorted_token_ids,
13 torch::Tensor experts_ids,
14 torch::Tensor num_tokens_post_pad);
15 -
16 -void sgl_moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
17 - int64_t block_size,
18 - torch::Tensor sorted_token_ids,
19 - torch::Tensor experts_ids,
20 - torch::Tensor num_tokens_post_pad);
15 #ifndef USE_ROCM
16 torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output,
17 torch::Tensor b_qweight, torch::Tensor b_scales,
csrc/moe/torch_bindings.cpp
-9
@@ -22,15 +22,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
22 " Tensor! num_tokens_post_pad) -> ()");
23 m.impl("moe_align_block_size", torch::kCUDA, &moe_align_block_size);
24
25 - // temporarily adapted from
26 - // https://github.com/sgl-project/sglang/commit/ded9fcd09a43d5e7d5bb31a2bc3e9fc21bf65d2a
27 - m.def(
28 - "sgl_moe_align_block_size(Tensor topk_ids, int num_experts,"
29 - " int block_size, Tensor! sorted_token_ids,"
30 - " Tensor! experts_ids,"
31 - " Tensor! num_tokens_post_pad) -> ()");
32 - m.impl("sgl_moe_align_block_size", torch::kCUDA, &sgl_moe_align_block_size);
33 -
25 #ifndef USE_ROCM
26 m.def(
27 "moe_wna16_gemm(Tensor input, Tensor! output, Tensor b_qweight, "
tests/kernels/moe/test_moe_align_block_size.py new
+90
@@ -0,0 +1,90 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import itertools
4 +
5 +import pytest
6 +import torch
7 +
8 +from vllm import _custom_ops as ops
9 +from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
10 + moe_align_block_size_triton)
11 +
12 +
13 +@pytest.mark.parametrize(
14 + "block_size,num_tokens,topk,num_experts",
15 + list(
16 + itertools.product(
17 + [32, 64, 128, 256], # block_size
18 + [
19 + 1,
20 + 3,
21 + 7,
22 + 16,
23 + 256,
24 + 2256,
25 + 4096,
26 + ], # num_tokens
27 + [1, 4, 16, 64], # topk
28 + [64, 160, 256, 257, 260, 264], # num_experts
29 + )),
30 +)
31 +def test_moe_align_block_size_compare_implementations(block_size, num_tokens,
32 + topk, num_experts):
33 + topk_ids = torch.stack([
34 + torch.randperm(num_experts, dtype=torch.int32, device="cuda")[:topk]
35 + for _ in range(num_tokens)
36 + ])
37 +
38 + max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
39 +
40 + sorted_ids_cuda = torch.empty((max_num_tokens_padded, ),
41 + dtype=torch.int32,
42 + device=topk_ids.device)
43 + sorted_ids_cuda.fill_(topk_ids.numel())
44 + max_num_m_blocks = max_num_tokens_padded // block_size
45 + expert_ids_cuda = torch.zeros((max_num_m_blocks, ),
46 + dtype=torch.int32,
47 + device=topk_ids.device)
48 + num_tokens_post_pad_cuda = torch.empty((1),
49 + dtype=torch.int32,
50 + device=topk_ids.device)
51 +
52 + sorted_ids_triton = torch.empty_like(sorted_ids_cuda)
53 + sorted_ids_triton.fill_(topk_ids.numel())
54 + expert_ids_triton = torch.zeros_like(expert_ids_cuda)
55 + num_tokens_post_pad_triton = torch.empty_like(num_tokens_post_pad_cuda)
56 +
57 + ops.moe_align_block_size(
58 + topk_ids,
59 + num_experts,
60 + block_size,
61 + sorted_ids_cuda,
62 + expert_ids_cuda,
63 + num_tokens_post_pad_cuda,
64 + )
65 +
66 + moe_align_block_size_triton(
67 + topk_ids,
68 + num_experts,
69 + block_size,
70 + sorted_ids_triton,
71 + expert_ids_triton,
72 + num_tokens_post_pad_triton,
73 + )
74 +
75 + assert torch.allclose(expert_ids_cuda, expert_ids_triton), (
76 + f"Expert IDs mismatch for block_size={block_size}, "
77 + f"num_tokens={num_tokens}, topk={topk}\n"
78 + f"CUDA expert_ids: {expert_ids_cuda}\n"
79 + f"Triton expert_ids: {expert_ids_triton}")
80 +
81 + assert torch.allclose(
82 + num_tokens_post_pad_cuda, num_tokens_post_pad_triton), (
83 + f"Num tokens post pad mismatch for block_size={block_size}, "
84 + f"num_tokens={num_tokens}, topk={topk}\n"
85 + f"CUDA num_tokens_post_pad: {num_tokens_post_pad_cuda}\n"
86 + f"Triton num_tokens_post_pad: {num_tokens_post_pad_triton}")
87 +
88 +
89 +if __name__ == "__main__":
90 + pytest.main([__file__])
vllm/_custom_ops.py
-9
@@ -1524,15 +1524,6 @@ def moe_align_block_size(topk_ids: torch.Tensor, num_experts: int,
1524 num_tokens_post_pad)
1525
1526
1527 -def sgl_moe_align_block_size(topk_ids: torch.Tensor, num_experts: int,
1528 - block_size: int, sorted_token_ids: torch.Tensor,
1529 - experts_ids: torch.Tensor,
1530 - num_tokens_post_pad: torch.Tensor) -> None:
1531 - torch.ops._moe_C.sgl_moe_align_block_size(topk_ids, num_experts,
1532 - block_size, sorted_token_ids,
1533 - experts_ids, num_tokens_post_pad)
1534 -
1535 -
1527 def moe_wna16_gemm(input: torch.Tensor, output: torch.Tensor,
1528 b_qweight: torch.Tensor, b_scales: torch.Tensor,
1529 b_qzeros: Optional[torch.Tensor],
vllm/model_executor/layers/fused_moe/moe_align_block_size.py
+4 -24
@@ -4,7 +4,6 @@ from typing import Optional
4
5 import torch
6
7 -import vllm.envs as envs
7 from vllm import _custom_ops as ops
8 from vllm.triton_utils import tl, triton
9 from vllm.utils import round_up
@@ -99,6 +98,7 @@ def moe_align_block_size_stage4(
98
99 # Triton implementation based on:
100 # https://github.com/sgl-project/sglang/commit/ba5112ff691d791a9e38c6c71f59324a5fcb49d0
101 +# TODO(wentao): Deprecated this function in the future.
102 def moe_align_block_size_triton(
103 topk_ids: torch.Tensor,
104 num_experts: int,
@@ -220,29 +220,9 @@ def moe_align_block_size(
220 num_tokens_post_pad = torch.empty((1),
221 dtype=torch.int32,
222 device=topk_ids.device)
223 - if num_experts >= 224:
224 - if envs.VLLM_ENABLE_MOE_ALIGN_BLOCK_SIZE_TRITON or num_experts != 256:
225 - moe_align_block_size_triton(
226 - topk_ids,
227 - num_experts,
228 - block_size,
229 - sorted_ids,
230 - expert_ids,
231 - num_tokens_post_pad,
232 - )
233 - else:
234 - # Currently requires num_experts=256
235 - ops.sgl_moe_align_block_size(
236 - topk_ids,
237 - num_experts,
238 - block_size,
239 - sorted_ids,
240 - expert_ids,
241 - num_tokens_post_pad,
242 - )
243 - else:
244 - ops.moe_align_block_size(topk_ids, num_experts, block_size, sorted_ids,
245 - expert_ids, num_tokens_post_pad)
223 +
224 + ops.moe_align_block_size(topk_ids, num_experts, block_size, sorted_ids,
225 + expert_ids, num_tokens_post_pad)
226 if expert_map is not None:
227 expert_ids = expert_map[expert_ids]
228