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 7, 19:00 UTC
[Bugfix] Fix expert_ids padding values in moe_align_block_size kernel (#35161)
Signed-off-by: Xin Yang <xyangx@amazon.com>
Xin Yang committed
Feb 24, 2026 at 17:14 UTC
3bbb2046ff320395c80c139e55e7c1947c3fb5e1
2 files changed
+12
-9
csrc/moe/moe_align_sum_kernels.cu
+4
-4
@@ -172,7 +172,7 @@ __device__ void _moe_align_block_size(
172
}
173
}
174
175
- // Fill remaining expert_ids with 0
175
+ // Fill remaining expert_ids with -1
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) {
@@ -265,7 +265,7 @@ __device__ void _moe_align_block_size_small_batch_expert(
265
}
266
}
267
268
- // Fill remaining expert_ids with 0
268
+ // Fill remaining expert_ids with -1
269
const size_t fill_start_idx = cumsum[num_experts] / block_size + tid;
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;
@@ -332,7 +332,7 @@ __global__ void moe_align_block_size_kernel(
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);
335
+ 0, -1, topk_num, nullptr, has_expert_map);
336
}
337
338
template <typename scalar_t>
@@ -373,7 +373,7 @@ __global__ void moe_align_block_size_small_batch_expert_kernel(
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,
376
+ CEILDIV(max_num_tokens_padded, block_size), -1, 0, topk_num, nullptr,
377
has_expert_map);
378
}
379
tests/kernels/moe/test_moe_align_block_size.py
+8
-5
@@ -12,7 +12,7 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
12
batched_moe_align_block_size,
13
moe_align_block_size,
14
)
15
-from vllm.utils.math_utils import round_up
15
+from vllm.utils.math_utils import cdiv, round_up
16
from vllm.utils.torch_utils import set_random_seed
17
18
NUM_TOKENS = [1, 3, 256, 2256, 4096]
@@ -142,7 +142,9 @@ def torch_moe_align_block_size(
142
device=topk_ids.device,
143
)
144
max_num_blocks = (max_num_tokens_padded + block_size - 1) // block_size
145
- expert_ids = torch.zeros(max_num_blocks, dtype=torch.int32, device=topk_ids.device)
145
+ expert_ids = torch.full(
146
+ (max_num_blocks,), -1, dtype=torch.int32, device=topk_ids.device
147
+ )
148
149
current_pos = 0
150
current_block = 0
@@ -234,9 +236,10 @@ def test_moe_align_block_size(
236
assert len(valid_tokens) == total_tokens, (
237
f"Should have exactly {total_tokens} valid tokens, got {len(valid_tokens)}"
238
)
237
- assert (actual_expert_ids >= 0).all() and (actual_expert_ids < num_experts).all(), (
238
- "expert_ids should contain valid expert indices"
239
- )
239
+ actual_num_blocks = cdiv(int(actual_num_tokens.item()), block_size)
240
+ assert (actual_expert_ids[:actual_num_blocks] >= 0).all() and (
241
+ actual_expert_ids[:actual_num_blocks] < num_experts
242
+ ).all(), "expert_ids should contain valid expert indices"
243
244
245
@pytest.mark.parametrize("m", [16, 32, 2048])