[KernelSwift算子优化] 赵世杰-Clike算子优化赛道华为赛题-251300387

[freebeat-boop]

作品说明

本次提交面向 Ascend 910B（Atlas A2），提供 Task01 SparseAttention、Task02 Indexer 和 Task03 Sinkhorn 三个算子的 AscendC 实现，并补充构建脚本、性能测试脚本及相关文档。三个算子均通过赛事 auto_bench.py 正确性检查。

优化方案

• SparseAttention：将 KV 数据驻留 UB，使用 TQueBind 流水优化 gather 搬运，并采用 FP16 Cube 完成 contraction，score 和 softmax 保持 FP32。
• Indexer：对 key 维度进行对齐，通过 bmm 降低 broadcast matmul 开销，并在 AscendC kernel 中融合 ReLU、权重乘法、head reduction 和 causal mask。
• Sinkhorn：将 softmax 与多轮行列归一化融合为单次 kernel launch，使 4 x 4 矩阵在迭代期间驻留 UB。

性能测试结果

测试环境为 HiDevLab 平台 Ascend 910B1 A2，三个算子均通过 accuracy 检查。

算子	PyTorch Reference	AscendC	加速比
SparseAttention	12.805535 ms	7.297500 ms	1.755x
Indexer	9.259980 ms	5.142080 ms	1.801x
Sinkhorn	1.623945 ms	0.329970 ms	4.921x

SparseAttention 和 Indexer 预热 100 次、测量 1000 次；Sinkhorn 因单次执行时间较短，预热 100 次、测量 10000 次，以降低测试波动。

原创声明

赛题参考实现和评测工具来源于 KernelSwift/DLBlas。本次提交中的 AscendC 优化实现、Torch 接入、构建与测试脚本由提交者基于赛题要求独立完成，不包含未经授权复制的第三方参赛代码。

[CLAassistant]

All committers have signed the CLA.