所有对 expert_map 的 All-to-All 后端（如 DeepEP 弹性调度逻辑）可能将非本地专家编码为 -1，但现有内核在访问 topk_ids[i] 后未检查该值，直接触发 expert_map[-1] 导致错误。此 PR 通过在校验前剔除 -1 及越界值来解决该问题。
此 PR 是对关键路径的鲁棒性修复，建议尽快合并。设计决策（将合法性检查封装为独立函数）值得借鉴，便于后续维护和扩展。
review 中无实质性讨论；WoosukKwon 直接批准了变更。
moe_align_sum_kernels.cu 中新增 get_local_expert_id 核函数，封装检查逻辑；_moe_align_block_size 和 _moe_align_block_size_small_batch_expert 中直接读取 topk_ids 的位置替换为调用 get_local_expert_id，仅在返回值不为 -1 时处理；test_moe_align_block_size.py 中为 test_moe_align_block_size_with_expert_map 增加 mask_inactive_experts 参数化测试，验证含 -1 填充的输入能正确对齐。| 文件 | 模块 | 状态 | 重要度 |
|---|---|---|---|
tests/kernels/moe/test_moe_align_block_size.py |
MoE 测试 | modified | 4.7 |
csrc/libtorch_stable/moe/moe_align_sum_kernels.cu |
MoE 内核 | modified | 5.57 |
tests/kernels/moe/test_moe_align_block_size.py
test-coverage
测试覆盖增强，新增 mask_inactive_experts 参数验证含 -1 填充的输入。
@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,
mask_inactive_experts: bool,
):
"""Test moe_align_block_size with expert mapping (EP scenario)"""
# 构造 expert_map：一半是本地专家
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：当 mask_inactive_experts 时，非本地专家用 -1
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 = moe_align_block_size(topk_ids, block_size, num_experts,
expert_map=expert_map,
ignore_invalid_experts=True)
golden = torch_moe_align_block_size(topk_ids, block_size, num_experts,
expert_map=expert_map)
# 验证对齐结果
torch.testing.assert_close(actual.num_tokens, golden.num_tokens)
torch.testing.assert_close(actual.expert_ids, golden.expert_ids)
_verify_expert_level_sorting(...)
csrc/libtorch_stable/moe/moe_align_sum_kernels.cu
core-logic
核心逻辑修改，添加 get_local_expert_id 函数并替换所有直接索引。
// 辅助函数：安全获取本地专家 ID
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];
// 有效性检查：拒绝越界或 -1
if (expert_id >= num_experts || expert_id < 0) {
return -1;
}
if (has_expert_map) {
expert_id = expert_map[expert_id];
}
return expert_id;
}
// 在 _moe_align_block_size 中的使用示例
for (size_t i = tid; i < numel; i += stride) {
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);
}
}
当前评论区没有形成足够清晰的争议点或结论，后续有更多讨论时会体现在这里。
仅影响启用了 expert_map 的 MoE 对齐路径。核心逻辑转为先校验再访问，消除越界风险。新增的测试覆盖了含 -1 和不含 -1 两种情况，回归风险低。性能影响仅在于每个 token 增加了分支判断，但由于 -1 使用频率低且仅增加少量条件，影响可忽略。
对使用 All-to-All 后端（如 DeepEP）进行专家并行的用户影响最大：修复了可能出现的内核崩溃。对其他使用默认 MoE 对齐的用户无影响（因为默认 expert_map 为 None，不会调用新路径）。变更完全向后兼容。
当前没有检测到明确关联的 Issue 链接，后续同步到相关引用后会出现在这里。
参与讨论