@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 879ddb09c3
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 11, 21:15 UTC

[Kernel][MoE] optimize `moe_align_block_size` (#29642)

Signed-off-by: Jinzhen Lin <jinzhen.ljz@antgroup.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>

Jinzhen Lin committed Dec 7, 2025 at 17:58 UTC 879ddb09c3b7b5a21f75e46ded148ce70f1c486e
10 files changed +195 -63
benchmarks/kernels/benchmark_moe_align_block_size.py
+17 -4
@@ -24,12 +24,15 @@ def get_topk_ids(num_tokens: int, num_experts: int, topk: int) -> torch.Tensor:
24 num_tokens_range = [1, 16, 256, 4096]
25 num_experts_range = [16, 64, 224, 256, 280, 512]
26 topk_range = [1, 2, 8]
27 -configs = list(itertools.product(num_tokens_range, num_experts_range, topk_range))
27 +ep_size_range = [1, 8]
28 +configs = list(
29 + itertools.product(num_tokens_range, num_experts_range, topk_range, ep_size_range)
30 +)
31
32
33 @triton.testing.perf_report(
34 triton.testing.Benchmark(
32 - x_names=["num_tokens", "num_experts", "topk"],
35 + x_names=["num_tokens", "num_experts", "topk", "ep_size"],
36 x_vals=configs,
37 line_arg="provider",
38 line_vals=["vllm"],
@@ -38,16 +41,26 @@ configs = list(itertools.product(num_tokens_range, num_experts_range, topk_range
41 args={},
42 )
43 )
41 -def benchmark(num_tokens, num_experts, topk, provider):
44 +def benchmark(num_tokens, num_experts, topk, ep_size, provider):
45 """Benchmark function for Triton."""
46 block_size = 256
47 + torch.cuda.manual_seed_all(0)
48 topk_ids = get_topk_ids(num_tokens, num_experts, topk)
49
50 + e_map = None
51 + if ep_size != 1:
52 + local_e = num_experts // ep_size
53 + e_ids = torch.randperm(num_experts, device="cuda", dtype=torch.int32)[:local_e]
54 + e_map = torch.full((num_experts,), -1, device="cuda", dtype=torch.int32)
55 + e_map[e_ids] = torch.arange(local_e, device="cuda", dtype=torch.int32)
56 +
57 quantiles = [0.5, 0.2, 0.8]
58
59 if provider == "vllm":
60 ms, min_ms, max_ms = triton.testing.do_bench(
50 - lambda: moe_align_block_size(topk_ids, block_size, num_experts),
61 + lambda: moe_align_block_size(
62 + topk_ids, block_size, num_experts, e_map, ignore_invalid_experts=True
63 + ),
64 quantiles=quantiles,
65 )
66
csrc/moe/moe_align_sum_kernels.cu
+106 -46
@@ -83,14 +83,22 @@ template <typename scalar_t>
83 __global__ void moe_align_block_size_kernel(
84 const scalar_t* __restrict__ topk_ids,
85 int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids,
86 - int32_t* __restrict__ total_tokens_post_pad, int32_t num_experts,
86 + int32_t* __restrict__ total_tokens_post_pad,
87 + int32_t* __restrict__ expert_map, int32_t num_experts,
88 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) {
89 + size_t numel, int32_t* __restrict__ cumsum, int32_t max_num_tokens_padded,
90 + bool has_expert_map) {
91 extern __shared__ int32_t shared_counts[];
92
91 - // Initialize sorted_token_ids with numel
92 - for (size_t it = threadIdx.x; it < max_num_tokens_padded; it += blockDim.x) {
93 - sorted_token_ids[it] = numel;
93 + // Use a separate threadblock to fill sorted_token_ids.
94 + // This is safe since the current kernel does not use sorted_token_ids.
95 + if (blockIdx.x == 1) {
96 + // Initialize sorted_token_ids with numel
97 + for (size_t it = threadIdx.x; it < max_num_tokens_padded;
98 + it += blockDim.x) {
99 + sorted_token_ids[it] = numel;
100 + }
101 + return;
102 }
103
104 const int warp_id = threadIdx.x / WARP_SIZE;
@@ -112,6 +120,11 @@ __global__ void moe_align_block_size_kernel(
120 if (expert_id >= num_experts) {
121 continue;
122 }
123 + if (has_expert_map) {
124 + expert_id = expert_map[expert_id];
125 + // filter invalid experts
126 + if (expert_id == -1) continue;
127 + }
128 int warp_idx = expert_id / experts_per_warp;
129 int expert_offset = expert_id % experts_per_warp;
130 atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset], 1);
@@ -163,7 +176,8 @@ 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,
166 - size_t numel, int32_t num_experts) {
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
@@ -172,6 +186,11 @@ __global__ void count_and_sort_expert_tokens_kernel(
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 }
@@ -193,50 +212,69 @@ __global__ void moe_sum_kernel(
212 }
213 }
214
196 -template <typename scalar_t>
215 +template <typename scalar_t, int32_t fill_threads>
216 __global__ void moe_align_block_size_small_batch_expert_kernel(
217 const scalar_t* __restrict__ topk_ids,
218 int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids,
200 - int32_t* __restrict__ total_tokens_post_pad, int32_t num_experts,
201 - int32_t block_size, size_t numel, int32_t max_num_tokens_padded) {
202 - // Initialize sorted_token_ids with numel
203 - for (size_t it = threadIdx.x; it < max_num_tokens_padded; it += blockDim.x) {
204 - sorted_token_ids[it] = numel;
219 + int32_t* __restrict__ total_tokens_post_pad,
220 + 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) {
222 + // Use an additional group of threads to fill sorted_token_ids.
223 + // Since the current kernel will use sorted_token_ids afterward,
224 + // we fill sorted_token_ids within the same threadblock to make
225 + // synchronization easier.
226 + if (threadIdx.x < fill_threads) {
227 + // Initialize sorted_token_ids with numel
228 + for (size_t it = threadIdx.x; it < max_num_tokens_padded;
229 + it += fill_threads) {
230 + sorted_token_ids[it] = numel;
231 + }
232 + // Three __syncthreads() corresponding to the other threads
233 + __syncthreads();
234 + __syncthreads();
235 + __syncthreads();
236 + return;
237 }
238
207 - const size_t tid = threadIdx.x;
208 - const size_t stride = blockDim.x;
239 + const size_t tid = threadIdx.x - fill_threads;
240 + const size_t stride = blockDim.x - fill_threads;
241
242 extern __shared__ int32_t shared_mem[];
243 int32_t* cumsum = shared_mem;
244 int32_t* tokens_cnts = (int32_t*)(shared_mem + num_experts + 1);
245
246 for (int i = 0; i < num_experts; ++i) {
215 - tokens_cnts[(threadIdx.x + 1) * num_experts + i] = 0;
247 + tokens_cnts[(tid + 1) * num_experts + i] = 0;
248 }
249
250 for (size_t i = tid; i < numel; i += stride) {
219 - ++tokens_cnts[(threadIdx.x + 1) * num_experts + topk_ids[i]];
251 + int32_t expert_id = topk_ids[i];
252 + if (has_expert_map) {
253 + expert_id = expert_map[expert_id];
254 + // filter invalid expert
255 + if (expert_id == -1) continue;
256 + }
257 + ++tokens_cnts[(tid + 1) * num_experts + expert_id];
258 }
259
260 __syncthreads();
261
224 - if (threadIdx.x < num_experts) {
225 - tokens_cnts[threadIdx.x] = 0;
226 - for (int i = 1; i <= blockDim.x; ++i) {
227 - tokens_cnts[i * num_experts + threadIdx.x] +=
228 - tokens_cnts[(i - 1) * num_experts + threadIdx.x];
262 + if (tid < num_experts) {
263 + tokens_cnts[tid] = 0;
264 + for (int i = 1; i <= stride; ++i) {
265 + tokens_cnts[i * num_experts + tid] +=
266 + tokens_cnts[(i - 1) * num_experts + tid];
267 }
268 }
269
270 __syncthreads();
271
234 - if (threadIdx.x == 0) {
272 + if (tid == 0) {
273 cumsum[0] = 0;
274 for (int i = 1; i <= num_experts; ++i) {
275 cumsum[i] =
276 cumsum[i - 1] +
239 - CEILDIV(tokens_cnts[blockDim.x * num_experts + i - 1], block_size) *
277 + CEILDIV(tokens_cnts[stride * num_experts + i - 1], block_size) *
278 block_size;
279 }
280 *total_tokens_post_pad = static_cast<int32_t>(cumsum[num_experts]);
@@ -244,26 +282,30 @@ __global__ void moe_align_block_size_small_batch_expert_kernel(
282
283 __syncthreads();
284
247 - if (threadIdx.x < num_experts) {
248 - for (int i = cumsum[threadIdx.x]; i < cumsum[threadIdx.x + 1];
249 - i += block_size) {
250 - expert_ids[i / block_size] = threadIdx.x;
285 + if (tid < num_experts) {
286 + for (int i = cumsum[tid]; i < cumsum[tid + 1]; i += block_size) {
287 + expert_ids[i / block_size] = tid;
288 }
289 }
290
291 // Fill remaining expert_ids with 0
255 - const size_t fill_start_idx = cumsum[num_experts] / block_size + threadIdx.x;
292 + 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);
257 - for (size_t i = fill_start_idx; i < expert_ids_size; i += blockDim.x) {
294 + for (size_t i = fill_start_idx; i < expert_ids_size; i += stride) {
295 expert_ids[i] = 0;
296 }
297
298 for (size_t i = tid; i < numel; i += stride) {
299 int32_t expert_id = topk_ids[i];
300 + if (has_expert_map) {
301 + expert_id = expert_map[expert_id];
302 + // filter invalid expert
303 + if (expert_id == -1) continue;
304 + }
305 int32_t rank_post_pad =
264 - tokens_cnts[threadIdx.x * num_experts + expert_id] + cumsum[expert_id];
306 + tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id];
307 sorted_token_ids[rank_post_pad] = i;
266 - ++tokens_cnts[threadIdx.x * num_experts + expert_id];
308 + ++tokens_cnts[tid * num_experts + expert_id];
309 }
310 }
311
@@ -275,7 +317,8 @@ __global__ void moe_align_block_size_small_batch_expert_kernel(
317 void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
318 int64_t block_size, torch::Tensor sorted_token_ids,
319 torch::Tensor experts_ids,
278 - torch::Tensor num_tokens_post_pad) {
320 + torch::Tensor num_tokens_post_pad,
321 + std::optional<torch::Tensor> maybe_expert_map) {
322 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
323
324 int64_t padded_num_experts =
@@ -287,14 +330,19 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
330 // BlockScan uses 1024 threads and assigns one thread per expert.
331 TORCH_CHECK(padded_num_experts < 1024,
332 "padded_num_experts must be less than 1024");
333 + auto options_int =
334 + torch::TensorOptions().dtype(torch::kInt).device(topk_ids.device());
335 + bool has_expert_map = maybe_expert_map.has_value();
336 + torch::Tensor expert_map;
337 + if (has_expert_map) {
338 + expert_map = maybe_expert_map.value();
339 + } else {
340 + expert_map = torch::empty({0}, options_int);
341 + }
342
343 VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(
344 topk_ids.scalar_type(), "moe_align_block_size_kernel", [&] {
345 // calc needed amount of shared mem for `cumsum` tensors
294 - auto options_int =
295 - torch::TensorOptions().dtype(torch::kInt).device(topk_ids.device());
296 - torch::Tensor cumsum_buffer =
297 - torch::empty({num_experts + 1}, options_int);
346 bool small_batch_expert_mode =
347 (topk_ids.numel() < 1024) && (num_experts <= 64);
348
@@ -304,30 +352,41 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
352 ((threads + 1) * num_experts + (num_experts + 1)) *
353 sizeof(int32_t);
354
355 + // threadIdx.x >= fill_threads: counting experts and aligning
356 + // threadIdx.x < fill_threads: filling sorted_token_ids
357 + constexpr int32_t fill_threads = 256;
358 auto small_batch_expert_kernel =
359 vllm::moe::moe_align_block_size_small_batch_expert_kernel<
309 - scalar_t>;
310 - small_batch_expert_kernel<<<1, threads, shared_mem_size, stream>>>(
360 + scalar_t, fill_threads>;
361 + small_batch_expert_kernel<<<1, fill_threads + threads,
362 + shared_mem_size, stream>>>(
363 topk_ids.data_ptr<scalar_t>(),
364 sorted_token_ids.data_ptr<int32_t>(),
365 experts_ids.data_ptr<int32_t>(),
314 - num_tokens_post_pad.data_ptr<int32_t>(), num_experts, block_size,
315 - topk_ids.numel(), sorted_token_ids.size(0));
366 + num_tokens_post_pad.data_ptr<int32_t>(),
367 + expert_map.data_ptr<int32_t>(), num_experts, block_size,
368 + topk_ids.numel(), sorted_token_ids.size(0), has_expert_map);
369 } else {
370 + torch::Tensor cumsum_buffer =
371 + torch::empty({num_experts + 1}, options_int);
372 auto align_kernel = vllm::moe::moe_align_block_size_kernel<scalar_t>;
373
374 size_t num_warps = CEILDIV(padded_num_experts, experts_per_warp);
375 size_t shared_mem_size =
376 num_warps * experts_per_warp * sizeof(int32_t);
377
323 - align_kernel<<<1, threads, shared_mem_size, stream>>>(
378 + // launch two threadblocks
379 + // blockIdx.x == 0: counting experts and aligning
380 + // blockIdx.x == 1: filling sorted_token_ids
381 + align_kernel<<<2, threads, shared_mem_size, stream>>>(
382 topk_ids.data_ptr<scalar_t>(),
383 sorted_token_ids.data_ptr<int32_t>(),
384 experts_ids.data_ptr<int32_t>(),
327 - num_tokens_post_pad.data_ptr<int32_t>(), num_experts,
328 - padded_num_experts, experts_per_warp, block_size,
329 - topk_ids.numel(), cumsum_buffer.data_ptr<int32_t>(),
330 - sorted_token_ids.size(0));
385 + num_tokens_post_pad.data_ptr<int32_t>(),
386 + expert_map.data_ptr<int32_t>(), num_experts, padded_num_experts,
387 + experts_per_warp, block_size, topk_ids.numel(),
388 + cumsum_buffer.data_ptr<int32_t>(), sorted_token_ids.size(0),
389 + has_expert_map);
390
391 const int block_threads = std::min(256, (int)threads);
392 const int num_blocks =
@@ -340,7 +399,8 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
399 sort_kernel<<<actual_blocks, block_threads, 0, stream>>>(
400 topk_ids.data_ptr<scalar_t>(),
401 sorted_token_ids.data_ptr<int32_t>(),
343 - cumsum_buffer.data_ptr<int32_t>(), topk_ids.numel(), num_experts);
402 + cumsum_buffer.data_ptr<int32_t>(), expert_map.data_ptr<int32_t>(),
403 + topk_ids.numel(), num_experts, has_expert_map);
404 }
405 });
406 }
csrc/moe/moe_ops.h
+2 -1
@@ -11,7 +11,8 @@ void moe_sum(torch::Tensor& input, torch::Tensor& output);
11 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);
14 + torch::Tensor num_tokens_post_pad,
15 + std::optional<torch::Tensor> maybe_expert_map);
16
17 void batched_moe_align_block_size(int64_t max_tokens_per_batch,
18 int64_t block_size,
csrc/moe/torch_bindings.cpp
+2 -1
@@ -19,7 +19,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
19 "moe_align_block_size(Tensor topk_ids, int num_experts,"
20 " int block_size, Tensor! sorted_token_ids,"
21 " Tensor! experts_ids,"
22 - " Tensor! num_tokens_post_pad) -> ()");
22 + " Tensor! num_tokens_post_pad,"
23 + " Tensor? maybe_expert_map) -> ()");
24 m.impl("moe_align_block_size", torch::kCUDA, &moe_align_block_size);
25
26 // Aligning the number of tokens to be processed by each expert such
tests/kernels/moe/test_moe.py
+15 -1
@@ -955,9 +955,22 @@ def test_fused_marlin_moe_with_bias(m):
955 torch.testing.assert_close(marlin_output, torch_output, atol=5e-2, rtol=0)
956
957
958 -def test_moe_align_block_size_opcheck():
958 +@pytest.mark.parametrize("ep_size", [1, 2])
959 +def test_moe_align_block_size_opcheck(ep_size):
960 num_experts = 4
961 block_size = 4
962 +
963 + expert_map = None
964 + if ep_size != 1:
965 + local_num_experts = num_experts // ep_size
966 + expert_ids = torch.randint(
967 + 0, num_experts, (local_num_experts,), device="cuda", dtype=torch.int32
968 + )
969 + expert_map = torch.full((num_experts,), -1, device="cuda", dtype=torch.int32)
970 + expert_map[expert_ids] = torch.arange(
971 + local_num_experts, device="cuda", dtype=torch.int32
972 + )
973 +
974 topk_ids = torch.randint(0, num_experts, (3, 4), dtype=torch.int32, device="cuda")
975
976 max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
@@ -980,6 +993,7 @@ def test_moe_align_block_size_opcheck():
993 sorted_ids,
994 expert_ids,
995 num_tokens_post_pad,
996 + expert_map,
997 ),
998 )
999
tests/kernels/moe/test_moe_align_block_size.py
+17 -5
@@ -106,6 +106,8 @@ def torch_moe_align_block_size(
106 max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
107 if pad_sorted_ids:
108 max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
109 + if topk_ids.numel() < num_experts:
110 + max_num_tokens_padded = topk_ids.numel() * block_size
111
112 flattened_token_indices = torch.arange(
113 topk_ids.numel(), device=topk_ids.device, dtype=torch.int32
@@ -126,6 +128,8 @@ def torch_moe_align_block_size(
128 )
129 for expert_id in range(num_experts):
130 original_count = expert_token_counts[expert_id]
131 + if expert_map is not None and expert_map[expert_id] == -1:
132 + continue
133 if original_count > 0:
134 expert_padded_counts[expert_id] = (
135 (original_count + block_size - 1) // block_size
@@ -143,6 +147,9 @@ def torch_moe_align_block_size(
147 current_pos = 0
148 current_block = 0
149 for expert_id in range(num_experts):
150 + if expert_map is not None and expert_map[expert_id] == -1:
151 + continue
152 +
153 expert_mask = sorted_expert_ids == expert_id
154 expert_tokens = sorted_token_indices[expert_mask]
155 num_expert_tokens = expert_tokens.shape[0]
@@ -153,7 +160,13 @@ def torch_moe_align_block_size(
160 )
161
162 expert_blocks_needed = expert_padded_counts[expert_id] // block_size
156 - expert_ids[current_block : current_block + expert_blocks_needed] = expert_id
163 +
164 + expert_id_new = expert_id
165 + if expert_map is not None:
166 + expert_id_new = expert_map[expert_id]
167 + expert_ids[current_block : current_block + expert_blocks_needed] = (
168 + expert_id_new
169 + )
170
171 current_pos += expert_padded_counts[expert_id]
172 current_block += expert_blocks_needed
@@ -163,8 +176,6 @@ def torch_moe_align_block_size(
176 [total_padded_tokens], dtype=torch.int32, device=topk_ids.device
177 )
178
166 - if expert_map is not None:
167 - expert_ids = expert_map[expert_ids]
179 return sorted_token_ids, expert_ids, num_tokens_post_pad
180
181
@@ -229,9 +240,9 @@ def test_moe_align_block_size(
240 )
241
242
232 -@pytest.mark.parametrize("m", [16, 32])
243 +@pytest.mark.parametrize("m", [16, 32, 2048])
244 @pytest.mark.parametrize("topk", [2, 4])
234 -@pytest.mark.parametrize("num_experts", [8])
245 +@pytest.mark.parametrize("num_experts", [8, 64])
246 @pytest.mark.parametrize("block_size", [64])
247 @pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
248 def test_moe_align_block_size_with_expert_map(
@@ -253,6 +264,7 @@ def test_moe_align_block_size_with_expert_map(
264 block_size=block_size,
265 num_experts=num_experts,
266 expert_map=expert_map,
267 + ignore_invalid_experts=True,
268 )
269 golden_sorted_ids, golden_expert_ids, golden_num_tokens = (
270 torch_moe_align_block_size(
vllm/_custom_ops.py
+2
@@ -1877,6 +1877,7 @@ def moe_align_block_size(
1877 sorted_token_ids: torch.Tensor,
1878 experts_ids: torch.Tensor,
1879 num_tokens_post_pad: torch.Tensor,
1880 + expert_map: torch.Tensor | None = None,
1881 ) -> None:
1882 torch.ops._moe_C.moe_align_block_size(
1883 topk_ids,
@@ -1885,6 +1886,7 @@ def moe_align_block_size(
1886 sorted_token_ids,
1887 experts_ids,
1888 num_tokens_post_pad,
1889 + expert_map,
1890 )
1891
1892
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+5 -1
@@ -316,7 +316,11 @@ def fused_marlin_moe(
316 if global_num_experts == -1:
317 global_num_experts = E
318 sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
319 - topk_ids, block_size_m, global_num_experts, expert_map
319 + topk_ids,
320 + block_size_m,
321 + global_num_experts,
322 + expert_map,
323 + ignore_invalid_experts=True,
324 )
325
326 assert activation is not None
vllm/model_executor/layers/fused_moe/fused_moe.py
+8 -1
@@ -1887,7 +1887,11 @@ def fused_experts_impl(
1887 )
1888
1889 sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
1890 - curr_topk_ids, config["BLOCK_SIZE_M"], global_num_experts, expert_map
1890 + curr_topk_ids,
1891 + config["BLOCK_SIZE_M"],
1892 + global_num_experts,
1893 + expert_map,
1894 + ignore_invalid_experts=True,
1895 )
1896
1897 invoke_fused_moe_kernel(
@@ -1946,6 +1950,9 @@ def fused_experts_impl(
1950 block_shape=block_shape,
1951 )
1952
1953 + if expert_map is not None:
1954 + intermediate_cache3.zero_()
1955 +
1956 invoke_fused_moe_kernel(
1957 qintermediate_cache2,
1958 w2,
vllm/model_executor/layers/fused_moe/moe_align_block_size.py
+21 -3
@@ -14,6 +14,7 @@ def moe_align_block_size(
14 num_experts: int,
15 expert_map: torch.Tensor | None = None,
16 pad_sorted_ids: bool = False,
17 + ignore_invalid_experts: bool = False,
18 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
19 """
20 Aligns the token distribution across experts to be compatible with block
@@ -35,7 +36,13 @@ def moe_align_block_size(
36 expert parallel shard. If the expert is not in the current expert
37 parallel shard, the mapping is set to -1.
38 - pad_sorted_ids: A flag indicating whether the sorted_token_ids length
38 - should be padded to a multiple of block_size,
39 + should be padded to a multiple of block_size,
40 + - ignore_invalid_experts: A flag indicating whether to ignore invalid
41 + experts. When False, all expert_ids in topk_ids will participate in
42 + counting and ranking, but invalid experts in expert_ids will be marked
43 + as -1. When True, all invalid expert_ids in topk_ids will be ignored
44 + and will not participate in counting or ranking, and there will be no
45 + -1 in expert_ids.
46
47 Returns:
48 - sorted_token_ids: A tensor containing the sorted token indices according
@@ -67,6 +74,10 @@ def moe_align_block_size(
74 max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
75 if pad_sorted_ids:
76 max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
77 + if topk_ids.numel() < num_experts:
78 + max_num_tokens_padded = min(
79 + topk_ids.numel() * block_size, max_num_tokens_padded
80 + )
81 sorted_ids = torch.empty(
82 (max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device
83 )
@@ -77,9 +88,16 @@ def moe_align_block_size(
88 num_tokens_post_pad = torch.empty((1), dtype=torch.int32, device=topk_ids.device)
89
90 ops.moe_align_block_size(
80 - topk_ids, num_experts, block_size, sorted_ids, expert_ids, num_tokens_post_pad
91 + topk_ids,
92 + num_experts,
93 + block_size,
94 + sorted_ids,
95 + expert_ids,
96 + num_tokens_post_pad,
97 + expert_map if ignore_invalid_experts else None,
98 )
82 - if expert_map is not None:
99 +
100 + if expert_map is not None and not ignore_invalid_experts:
101 expert_ids = expert_map[expert_ids]
102
103 return sorted_ids, expert_ids, num_tokens_post_pad