[KernelSwift算子优化] 程柯雷-Task02 Indexer-571900729

[Lfan-ke]

作品说明：沐曦 C500 上用 Clike 重写 KernelSwift 赛道二 Task02 Indexer，文件内同时含与参考一致的 Model 与优化实现 ModelNew，forward 内通过 python 调用自实现算子。

优化方案：把 einsum 之后的四段（relu×weights→sum、掩码、topk(128)、二次掩码，合计占参考实现的 91%）折成一个 kernel——一个 block 一行、只读因果掩码放行的那半数据、行内先按 128 一段各自排再树形归并只留 top-128；rotary 另写一个原地 kernel，一线程搬一个 uint4。

★本次更新：把 reduce+topk 的 11 次 kernel 启动合并成 1 次。原先按 valid 的 2 的幂把行分段、shared 按段申请，代价是 S=2600 时 P 取遍 1..1024 要启动 11 次，而这张卡一次 kernel 启动就要 15 µs。但 kernel 内部本来就自己按 valid 重算 P，分段只决定 shared 大小与线程数：shared 给大了无害，线程数也不改变结果（所有循环都按 threadIdx 跨步，同步次数与 blockDim 无关），所以合成一次启动是逐位等价的。整题 2.151 → 1.968 ms，只测 reduce+topk 这一个 kernel 是 1.166 → 0.981 ms，交替 A/B 六对全胜。省下的不止 10 次启动开销：2 万个块的代价相差极大（s 小的行几乎不干活、s 大的要排 1024），合成一个 grid 后调度器能把轻重块混着填满，而分段启动每一段都是一道屏障。

性能结果（MetaX C500 / MACA 3.7.0.38 / torch 2.8.0+metax3.7.0.7，官方 benchmarks/ks/auto_bench.py 默认参数，同一次独占会话连跑五次，均 PASS accuracy）：

赛题	v0（torch 参考）	v1（本提交）	Speedup
Task02 Indexer	6.3997 / 6.4047 / 6.4117 / 6.4067 / 6.4101 ms	2.2042 / 2.2109 / 2.1978 / 2.2223 / 2.2248 ms	2.903× / 2.897× / 2.917× / 2.883× / 2.881×

正确性：输出是 int64 索引，harness 用 torch.equal 逐位比对、无容差，五次均 PASS。

本作品为原创。

[CLAassistant]

All committers have signed the CLA.