提交提案之前，请先检索仓库内是否已有相同的提案，如已有请在同一提案中进行讨论。

💻 需求背景

_make_chunk_sort_map_kernel 根据 split_sizes 和 sorted_indices 生成 row_id_map，供后续 copy kernel 完成 MoE token dispatch。

原实现按 token block 与所有 expert 的 cumulative split 做二维比较和 reduction。随着 expert 数增大，临时向量规模随 token × expert 增长，存在较高 UB 占用。

当前现状

算子_make_chunk_sort_map_kernel 接入整网中有偶现的ub overflow ，但单算子又不复现，所以整改修改算子实现

期望实现的功能内容

希望通过修改算子算法，减少ub 占用

修改方案

将 make kernel 改为 expert-tile 实现：

• 每个 expert 的 token 区间按 BLOCK_TOKENS=128 切分；
• 使用 pid 和运行时 vector core 数进行 tile 轮转分配；
• 每个 tile 直接计算并写入对应的 row_id_map；
• 使用尾块 mask 防止最后一个 tile 越界；
• grid 根据设备实际 vector core 数动态获取，不固定使用 48 核。

该方式避免了原有 token-expert 二维临时矩阵，UB 占用主要由固定 tile 大小决定。

测试结果

精度验证

参考实现使用 torch.split + torch.cat megatron的小算子标杆，并通过：

torch.testing.assert_close(
    actual,
    expected,
    rtol=0.0,
    atol=0.0,
    equal_nan=True,
)

已覆盖以下场景：

• zero-sized expert；
• 127/128/129、255/256/257 等 tile 边界；
• 32768 tokens / 256 experts；
• 43731 tokens / 320 experts；
• 65536 tokens / 512 experts；
• forward、backward 以及带 probs 的完整 API。
  以上用例均通过，输出与参考实现完全一致，且 row_id_map 为合法 permutation。

内存安全验证

在 Ascend 910B3 上使用 mssanitizer 验证：

• 600000 tokens / 7168 experts；
• 1048577 tokens / 2048 experts。
  memcheck 未发现 target kernel 的 illegal read/write，racecheck 未发现跨 program 重叠写问题。

性能结果

测试指标为 _make_chunk_sort_map_kernel 的 device kernel time，使用Ascend 910B3 、warmup 5 次后采样。
Shape 原二维版本 Expert-tile 版本 加速比

当前性能数据仅针对 make kernel，不包含完整 sort API 的 host、cumsum、scatter 和 copy kernel 时间。完整端到端收益还需要结合实际 MoE 模型继续验证。

欢迎加入社区，感谢您对社区的贡献 🎉!