karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

handle topk_ids padding in align sum kernel (#47785)

Browse Source 
Signed-off-by: gnovack <novackgm@gmail.com>
This commit is contained in:
gnovack
2026-07-10 13:33:28 -07:00
committed by GitHub
parent 735def4fcf
commit f378f79b7c
2 changed files with 62 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/libtorch_stable/moe/moe_align_sum_kernels.cu
+46 -47
View File
@@ -82,6 +82,21 @@ __global__ void batched_moe_align_block_size_kernel(
}
} // namespace batched_moe_align_block_size
template <typename scalar_t>
__device__ __forceinline__ int get_local_expert_id(
size_t idx, const scalar_t* __restrict__ topk_ids,
int32_t* __restrict__ expert_map, int32_t num_experts,
bool has_expert_map) {
int expert_id = topk_ids[idx];
if (expert_id >= num_experts || expert_id < 0) {
return -1;
}
if (has_expert_map) {
expert_id = expert_map[expert_id];
}
return expert_id;
}
template <typename scalar_t>
__device__ void _moe_align_block_size(
const scalar_t* __restrict__ topk_ids,
@@ -126,20 +141,15 @@ __device__ void _moe_align_block_size(
const size_t stride = blockDim.x;
for (size_t i = tid; i < numel; i += stride) {
int expert_id = topk_ids[i];
if (expert_id >= num_experts) {
continue;
if (int expert_id = get_local_expert_id(i, topk_ids, expert_map,
num_experts, has_expert_map);
expert_id != -1) {
int warp_idx = expert_id / experts_per_warp;
int expert_offset = expert_id % experts_per_warp;
int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset],
mask);
}
if (has_expert_map) {
expert_id = expert_map[expert_id];
// filter invalid experts
if (expert_id == -1) continue;
}
int warp_idx = expert_id / experts_per_warp;
int expert_offset = expert_id % experts_per_warp;
int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset],
mask);
}
__syncthreads();
@@ -227,14 +237,12 @@ __device__ void _moe_align_block_size_small_batch_expert(
}
for (size_t i = tid; i < numel; i += stride) {
int32_t expert_id = topk_ids[i];
if (has_expert_map) {
expert_id = expert_map[expert_id];
// filter invalid expert
if (expert_id == -1) continue;
if (int expert_id = get_local_expert_id(i, topk_ids, expert_map,
num_experts, has_expert_map);
expert_id != -1) {
int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
tokens_cnts[(tid + 1) * num_experts + expert_id] += mask;
}
int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
tokens_cnts[(tid + 1) * num_experts + expert_id] += mask;
}
__syncthreads();
@@ -276,18 +284,16 @@ __device__ void _moe_align_block_size_small_batch_expert(
}
for (size_t i = tid; i < numel; i += stride) {
int32_t expert_id = topk_ids[i];
if (has_expert_map) {
expert_id = expert_map[expert_id];
// filter invalid expert
if (expert_id == -1) continue;
}
int32_t rank_post_pad =
tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id];
if (int expert_id = get_local_expert_id(i, topk_ids, expert_map,
num_experts, has_expert_map);
expert_id != -1) {
int32_t rank_post_pad =
tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id];
if (token_mask == nullptr || token_mask[i / topk_num]) {
sorted_token_ids[sorted_token_ids_offset + rank_post_pad] = i;
++tokens_cnts[tid * num_experts + expert_id];
if (token_mask == nullptr || token_mask[i / topk_num]) {
sorted_token_ids[sorted_token_ids_offset + rank_post_pad] = i;
++tokens_cnts[tid * num_experts + expert_id];
}
}
}
}
@@ -303,22 +309,15 @@ __device__ void _count_and_sort_expert_tokens(
const size_t stride = blockDim.x * gridDim.y;
for (size_t i = tid; i < numel; i += stride) {
int32_t expert_id = topk_ids[i];
if (expert_id >= num_experts) {
continue;
}
if (has_expert_map) {
expert_id = expert_map[expert_id];
// filter invalid experts
if (expert_id == -1) continue;
}
if (token_mask == nullptr || token_mask[i / topk_num]) {
int32_t rank_post_pad = atomicAdd(
&cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1);
sorted_token_ids[max_num_tokens_padded * model_offset + rank_post_pad] =
i;
if (int expert_id = get_local_expert_id(i, topk_ids, expert_map,
num_experts, has_expert_map);
expert_id != -1) {
if (token_mask == nullptr || token_mask[i / topk_num]) {
int32_t rank_post_pad = atomicAdd(
&cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1);
sorted_token_ids[max_num_tokens_padded * model_offset + rank_post_pad] =
i;
}
}
}
}
tests/kernels/moe/test_moe_align_block_size.py
+16 -6
View File
@@ -246,20 +246,30 @@ def test_moe_align_block_size(
@pytest.mark.parametrize("topk", [2, 4])
@pytest.mark.parametrize("num_experts", [8, 64])
@pytest.mark.parametrize("block_size", [64])
@pytest.mark.parametrize("mask_inactive_experts", [False, True])
def test_moe_align_block_size_with_expert_map(
m: int, topk: int, num_experts: int, block_size: int
m: int,
topk: int,
num_experts: int,
block_size: int,
mask_inactive_experts: bool,
):
"""Test moe_align_block_size with expert mapping (EP scenario)"""
topk_ids = torch.zeros((m, topk), device="cuda", dtype=torch.int32)
for i in range(m):
experts = torch.randperm(num_experts, device="cuda")[:topk]
topk_ids[i] = experts
expert_map = torch.full((num_experts,), -1, device="cuda", dtype=torch.int32)
local_experts = list(range(0, num_experts, 2))
for i, expert_id in enumerate(local_experts):
expert_map[expert_id] = i
topk_ids = torch.empty((m, topk), device="cuda", dtype=torch.int32)
for i in range(m):
experts = torch.randperm(num_experts, device="cuda")[:topk]
for k in range(topk):
topk_ids[i, k] = (
experts[k]
if (experts[k] in local_experts) or not mask_inactive_experts
else -1
)
actual_sorted_ids, actual_expert_ids, actual_num_tokens = moe_align_block_size(
topk_ids=topk_ids,
block_size=block_size,