handle topk_ids padding in align sum kernel (#47785)
Signed-off-by: gnovack <novackgm@gmail.com>
gnovack committed
Jul 10, 2026 at 13:33 UTC
f378f79b7c34d7ca94d53db34837929da2db03ed
2 files changed
+63
-54
csrc/libtorch_stable/moe/moe_align_sum_kernels.cu
+47
-48
@@ -82,6 +82,21 @@ __global__ void batched_moe_align_block_size_kernel(
82
}
83
} // namespace batched_moe_align_block_size
84
85
+template <typename scalar_t>
86
+__device__ __forceinline__ int get_local_expert_id(
87
+ size_t idx, const scalar_t* __restrict__ topk_ids,
88
+ int32_t* __restrict__ expert_map, int32_t num_experts,
89
+ bool has_expert_map) {
90
+ int expert_id = topk_ids[idx];
91
+ if (expert_id >= num_experts || expert_id < 0) {
92
+ return -1;
93
+ }
94
+ if (has_expert_map) {
95
+ expert_id = expert_map[expert_id];
96
+ }
97
+ return expert_id;
98
+}
99
+
100
template <typename scalar_t>
101
__device__ void _moe_align_block_size(
102
const scalar_t* __restrict__ topk_ids,
@@ -126,20 +141,15 @@ __device__ void _moe_align_block_size(
141
const size_t stride = blockDim.x;
142
143
for (size_t i = tid; i < numel; i += stride) {
129
- int expert_id = topk_ids[i];
130
- if (expert_id >= num_experts) {
131
- continue;
144
+ if (int expert_id = get_local_expert_id(i, topk_ids, expert_map,
145
+ num_experts, has_expert_map);
146
+ expert_id != -1) {
147
+ int warp_idx = expert_id / experts_per_warp;
148
+ int expert_offset = expert_id % experts_per_warp;
149
+ int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
150
+ atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset],
151
+ mask);
152
}
133
- if (has_expert_map) {
134
- expert_id = expert_map[expert_id];
135
- // filter invalid experts
136
- if (expert_id == -1) continue;
137
- }
138
- int warp_idx = expert_id / experts_per_warp;
139
- int expert_offset = expert_id % experts_per_warp;
140
- int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
141
- atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset],
142
- mask);
153
}
154
155
__syncthreads();
@@ -227,14 +237,12 @@ __device__ void _moe_align_block_size_small_batch_expert(
237
}
238
239
for (size_t i = tid; i < numel; i += stride) {
230
- int32_t expert_id = topk_ids[i];
231
- if (has_expert_map) {
232
- expert_id = expert_map[expert_id];
233
- // filter invalid expert
234
- if (expert_id == -1) continue;
240
+ if (int expert_id = get_local_expert_id(i, topk_ids, expert_map,
241
+ num_experts, has_expert_map);
242
+ expert_id != -1) {
243
+ int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
244
+ tokens_cnts[(tid + 1) * num_experts + expert_id] += mask;
245
}
236
- int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
237
- tokens_cnts[(tid + 1) * num_experts + expert_id] += mask;
246
}
247
248
__syncthreads();
@@ -276,18 +284,16 @@ __device__ void _moe_align_block_size_small_batch_expert(
284
}
285
286
for (size_t i = tid; i < numel; i += stride) {
279
- int32_t expert_id = topk_ids[i];
280
- if (has_expert_map) {
281
- expert_id = expert_map[expert_id];
282
- // filter invalid expert
283
- if (expert_id == -1) continue;
284
- }
285
- int32_t rank_post_pad =
286
- tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id];
287
-
288
- if (token_mask == nullptr || token_mask[i / topk_num]) {
289
- sorted_token_ids[sorted_token_ids_offset + rank_post_pad] = i;
290
- ++tokens_cnts[tid * num_experts + expert_id];
287
+ if (int expert_id = get_local_expert_id(i, topk_ids, expert_map,
288
+ num_experts, has_expert_map);
289
+ expert_id != -1) {
290
+ int32_t rank_post_pad =
291
+ tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id];
292
+
293
+ if (token_mask == nullptr || token_mask[i / topk_num]) {
294
+ sorted_token_ids[sorted_token_ids_offset + rank_post_pad] = i;
295
+ ++tokens_cnts[tid * num_experts + expert_id];
296
+ }
297
}
298
}
299
}
@@ -303,22 +309,15 @@ __device__ void _count_and_sort_expert_tokens(
309
const size_t stride = blockDim.x * gridDim.y;
310
311
for (size_t i = tid; i < numel; i += stride) {
306
- int32_t expert_id = topk_ids[i];
307
- if (expert_id >= num_experts) {
308
- continue;
309
- }
310
-
311
- if (has_expert_map) {
312
- expert_id = expert_map[expert_id];
313
- // filter invalid experts
314
- if (expert_id == -1) continue;
315
- }
316
-
317
- if (token_mask == nullptr || token_mask[i / topk_num]) {
318
- int32_t rank_post_pad = atomicAdd(
319
- &cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1);
320
- sorted_token_ids[max_num_tokens_padded * model_offset + rank_post_pad] =
321
- i;
312
+ if (int expert_id = get_local_expert_id(i, topk_ids, expert_map,
313
+ num_experts, has_expert_map);
314
+ expert_id != -1) {
315
+ if (token_mask == nullptr || token_mask[i / topk_num]) {
316
+ int32_t rank_post_pad = atomicAdd(
317
+ &cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1);
318
+ sorted_token_ids[max_num_tokens_padded * model_offset + rank_post_pad] =
319
+ i;
320
+ }
321
}
322
}
323
}
tests/kernels/moe/test_moe_align_block_size.py
+16
-6
@@ -246,20 +246,30 @@ def test_moe_align_block_size(
246
@pytest.mark.parametrize("topk", [2, 4])
247
@pytest.mark.parametrize("num_experts", [8, 64])
248
@pytest.mark.parametrize("block_size", [64])
249
+@pytest.mark.parametrize("mask_inactive_experts", [False, True])
250
def test_moe_align_block_size_with_expert_map(
250
- m: int, topk: int, num_experts: int, block_size: int
251
+ m: int,
252
+ topk: int,
253
+ num_experts: int,
254
+ block_size: int,
255
+ mask_inactive_experts: bool,
256
):
257
"""Test moe_align_block_size with expert mapping (EP scenario)"""
253
- topk_ids = torch.zeros((m, topk), device="cuda", dtype=torch.int32)
254
- for i in range(m):
255
- experts = torch.randperm(num_experts, device="cuda")[:topk]
256
- topk_ids[i] = experts
257
-
258
expert_map = torch.full((num_experts,), -1, device="cuda", dtype=torch.int32)
259
local_experts = list(range(0, num_experts, 2))
260
for i, expert_id in enumerate(local_experts):
261
expert_map[expert_id] = i
262
263
+ topk_ids = torch.empty((m, topk), device="cuda", dtype=torch.int32)
264
+ for i in range(m):
265
+ experts = torch.randperm(num_experts, device="cuda")[:topk]
266
+ for k in range(topk):
267
+ topk_ids[i, k] = (
268
+ experts[k]
269
+ if (experts[k] in local_experts) or not mask_inactive_experts
270
+ else -1
271
+ )
272
+
273
actual_sorted_ids, actual_expert_ids, actual_num_tokens = moe_align_block_size(
274
topk_ids=topk_ids,
275
block_size=block_size,