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

dlblas/kernels/ks_competition/Clike-AscendC/ENVIRONMENT.md

@@ -0,0 +1,89 @@
+# 环境配置
+
+该实现的构建、正确性验证和性能测试使用 HiDevLab 平台提供的 Ascend 算力
+资源完成。
+
+## 已验证环境
+
+该实现的验证环境如下。
+
+| 项目 | 版本或规格 |
+|---|---|
+| NPU | Ascend 910B1 / Atlas A2 |
+| CANN | 9.0.0 |
+| Python | 3.11.15，代码兼容 Python 3.10 及以上 |
+| PyTorch | 2.10.0+cpu |
+| torch-npu | 2.10.0 |
+| CMake | 3.27.9，最低要求 3.16 |
+| AscendC 架构参数 | `dav-2201` |
+
+表中的 `torch 2.10.0+cpu` 是 torch-npu 配套 PyTorch 包的版本标记，算子
+实际在 NPU 上执行。
+
+## CANN 环境
+
+CANN 默认安装路径为 `/usr/local/Ascend/cann-9.0.0`。使用者可在进入
+DLBlas 仓库后执行：
+
+```bash
+source /usr/local/Ascend/cann-9.0.0/set_env.sh
+export ASCEND_HOME_PATH=/usr/local/Ascend/cann-9.0.0
+```
+
+如果 CANN 安装在其他位置，`ASCEND_HOME_PATH` 需要指向对应根目录。构建
+脚本从以下位置查找 AscendC CMake 工具链：
+
+```text
+${ASCEND_HOME_PATH}/aarch64-linux/tikcpp/ascendc_kernel_cmake
+```
+
+## Python 环境检查
+
+运行环境至少需要 `torch` 和 `torch_npu`。以下命令可用于确认版本和 NPU：
+
+构建过程只使用 Python 解释器定位这两个包，不依赖 Python C API，因此无需
+安装 `python3-dev`、`python3.10-dev` 或提供 `Python.h`。
+
+```bash
+python3 - <<'PY'
+import torch
+import torch_npu
+
+print("torch:", torch.__version__)
+print("torch_npu:", torch_npu.__version__)
+print("NPU available:", torch.npu.is_available())
+PY
+```
+
+正常环境中的 `NPU available` 应为 `True`。设备状态可通过以下命令检查：
+
+```bash
+npu-smi info
+```
+
+## 可选环境变量
+
+| 变量 | 作用 |
+|---|---|
+| `ASCEND_HOME_PATH` | 指定 CANN 根目录 |
+| `BUILD_JOBS` | 指定 CMake 并行编译任务数，默认值为 64 |
+| `DLBLAS_PYTHON_EXECUTABLE` | 指定同时安装了 `torch` 和 `torch_npu` 的 Python 解释器 |
+| `DLBLAS_KS_ASCENDC_LIBRARY` | 指定已经编译好的自定义动态库绝对路径 |
+
+构建脚本默认自动查找能够导入 `torch` 和 `torch_npu` 的解释器。存在多个
+Python 环境时，可显式指定验证环境中的解释器：
+
+```bash
+DLBLAS_PYTHON_EXECUTABLE=/usr/local/python3.11.15/bin/python3 \
+    bash dlblas/kernels/ks_competition/ascend/clike_910b/build.sh
+```
+
+未指定 `DLBLAS_KS_ASCENDC_LIBRARY` 时，加载器的默认查找位置为：
+
+```text
+dlblas/kernels/ks_competition/ascend/clike_910b/build/
+libdlblas_ks_ascendc_ops.so
+```
+
+`build/` 为生成目录，已由 DLBlas 的 `.gitignore` 排除，不属于 PR 的提交内容。
+

dlblas/kernels/ks_competition/Clike-AscendC/PERFORMANCE.md

@@ -0,0 +1,79 @@
+# 性能结果与测试精度
+
+## 测试口径
+
+以下性能数据采用统一测试口径记录。
+
+- 硬件：Ascend 910B1 / Atlas A2。
+- 软件：CANN 9.0.0、PyTorch 2.10.0、torch-npu 2.10.0。
+- 测试工具：赛事提供的 `benchmarks/ks/auto_bench.py`。
+- 统计方式：每个正式样本执行一次 `forward` 后同步 NPU，报告中位数。
+- 预热次数：每个实现 100 次。
+- 正式次数：SparseAttention 和 Indexer 为 1000 次，Sinkhorn 为 10000 次。
+- Sinkhorn 单次执行时间较短，增加测量次数可尽量降低不同测试运行之间的
+  波动，使中位数结果更稳定。
+- 浮点正确性阈值：`atol=1e-2`、`rtol=1e-2`、`equal_nan=True`。
+- 整数输出：使用 `torch.equal` 逐元素精确比较。
+
+## 性能结果
+
+| 算子 | 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 |
+
+记录结果显示，三个算子均通过赛事 `auto_bench.py` 的 accuracy 检查。
+表中数据为按当前次数配置完成的手动复测结果。
+
+## 精度说明
+
+### SparseAttention
+
+| 阶段 | 精度 |
+|---|---|
+| `q`、`kv` 输入 | BF16 |
+| `topk_idxs` 输入 | INT32 |
+| `attn_sink` | FP32 |
+| Reference 点积、softmax、加权求和 | FP32 |
+| 优化实现 Cube contraction | FP16 |
+| 优化实现 score 与 softmax | FP32 |
+| 输出 | BF16 |
+
+赛事文档明确标注了输入输出类型。该实现使用混合精度，最终 BF16 输出通过
+`atol=rtol=1e-2` 检查。
+
+### Indexer
+
+| 阶段 | 精度 |
+|---|---|
+| `x`、`qr`、KV cache | BF16 |
+| Linear 权重及主路径 | BF16 |
+| RoPE 临时计算 | FP32 / complex64 |
+| 优化实现 QK Cube `bmm` | BF16 |
+| 融合 reduction 累加 | FP32，并保留 BF16 乘法舍入 |
+| TopK 输出 | INT64 |
+
+虽然 `ModelArgs.dtype` 的默认字符串是 `"fp8"`，该组测试实际使用
+`default_dtype=torch.bfloat16`，输入、KV cache 和 Linear 权重均为 BF16，
+没有执行 FP8 路径。TopK 输出是整数，`auto_bench.py` 要求逐元素完全一致。
+
+### Sinkhorn
+
+| 阶段 | 精度 |
+|---|---|
+| 输入 | FP32 |
+| softmax、行归一化、列归一化 | FP32 |
+| 输出 | FP32 |
+
+赛事文档在 Task03 的 `forward` 注释中明确要求 FP32 输入和输出。该实现将
+全部迭代融合到一个 AscendC kernel，但不降低计算精度。
+
+## 文档要求的边界
+
+根据赛事要求，优化实现需要保持与 reference 相同的 `Model` 初始化参数和
+`forward` 参数，并通过 reference 正确性校验。赛事文档没有要求自定义
+kernel 的每个内部阶段必须与 reference 使用完全相同的中间精度，因此允许
+混合精度优化；最终输出仍须满足上述浮点容差或整数精确比较，并且实际执行
+路径必须调用自定义算子。
+

dlblas/kernels/ks_competition/Clike-AscendC/README.md

@@ -0,0 +1,62 @@
+# Clike-AscendC
+
+Clike-AscendC 是 KernelSwift Clike 华为昇腾赛道三个算子的 AscendC
+实现，面向 Ascend 910B（Atlas A2），覆盖 SparseAttention、Indexer 和
+Sinkhorn。本文档汇总其代码位置、优化方法及配套说明。
+
+## 代码位置
+
+下表中的路径均以 DLBlas 仓库根目录为基准。
+
+| 算子 | 赛事参考实现 | 优化后的 Python 入口 |
+|---|---|---|
+| SparseAttention | `dlblas/kernels/ks_competition/torch/sparse_attention.py` | `dlblas/kernels/ks_competition/ascend/clike_sparse_attention.py` |
+| Indexer | `dlblas/kernels/ks_competition/torch/indexer.py` | `dlblas/kernels/ks_competition/ascend/clike_indexer.py` |
+| Sinkhorn | `dlblas/kernels/ks_competition/torch/sinkhorn.py` | `dlblas/kernels/ks_competition/ascend/clike_sinkhorn.py` |
+
+三个优化入口均提供赛事评测工具要求的 `ModelNew`、`get_inputs` 和
+`get_init_inputs`。公共 AscendC 实现及构建支持的位置如下：
+
+| 内容 | 路径 |
+|---|---|
+| SparseAttention AscendC kernel | `dlblas/kernels/ks_competition/ascend/clike_910b/csrc/sparse_attention.asc` |
+| Indexer AscendC kernel | `dlblas/kernels/ks_competition/ascend/clike_910b/csrc/indexer.asc` |
+| Sinkhorn AscendC kernel | `dlblas/kernels/ks_competition/ascend/clike_910b/csrc/sinkhorn.asc` |
+| PyTorch 动态库加载器 | `dlblas/kernels/ks_competition/ascend/clike_910b/loader.py` |
+| CMake 配置 | `dlblas/kernels/ks_competition/ascend/clike_910b/CMakeLists.txt` |
+| 构建脚本 | `dlblas/kernels/ks_competition/ascend/clike_910b/build.sh` |
+| 性能测试脚本 | `dlblas/kernels/ks_competition/ascend/clike_910b/run_benchmarks.sh` |
+| 赛事性能工具 | `benchmarks/ks/auto_bench.py` |
+
+仓库中的 Task02 参考实现相比赛题文档将 `.cuda()` 改为 `.npu()`，将全局
+`ModelArgs` 实例改为函数内按需构造，并直接使用 `torch.bfloat16` 表示默认
+类型，以兼容赛事工具的安全 AST 加载。其 `Model`、`forward`、配置值、
+数据类型、随机数调用、计算顺序、mask 和 TopK 写法均保持赛事 reference
+的定义。
+
+## 优化说明
+
+### SparseAttention
+
+- 该实现由每个 Vector core 将当前 batch 的 `32 x 128` KV 表驻留在 UB。
+- 两级 `TQueBind` 流水用于重叠 gather 搬入与搬出。
+- contraction 采用 FP16 Cube 路径，score 和 softmax 保留 FP32。
+
+### Indexer
+
+- 该实现将 650 个逻辑 key 补齐到 656，使 score task 保持 512B 对齐。
+- 序列和 head 展平后通过一次 `bmm` 完成计算，以减少 broadcast matmul 开销。
+- AscendC kernel 融合 ReLU、BF16 权重乘法、16-head reduction 和 causal
+  mask，并使用双缓冲。
+
+### Sinkhorn
+
+- 该实现将 softmax 与 10 轮行列归一化融合为一次 kernel launch。
+- 每个 `4 x 4` 矩阵在整个迭代期间驻留 UB。
+
+## 文档索引
+
+- [环境配置](ENVIRONMENT.md)
+- [构建与运行脚本](RUNNING.md)
+- [性能结果与测试精度](PERFORMANCE.md)
+

dlblas/kernels/ks_competition/Clike-AscendC/RUNNING.md

@@ -0,0 +1,93 @@
+# 构建与运行脚本
+
+以下复现命令均以 DLBlas 仓库根目录为工作目录。
+
+## 编译 AscendC 动态库
+
+```bash
+source /usr/local/Ascend/cann-9.0.0/set_env.sh
+bash dlblas/kernels/ks_competition/ascend/clike_910b/build.sh
+```
+
+编译完成后生成以下动态库：
+
+```text
+dlblas/kernels/ks_competition/ascend/clike_910b/build/
+libdlblas_ks_ascendc_ops.so
+```
+
+该构建脚本使用 CMake Release 模式、`-O3` 和
+`--npu-arch=dav-2201`。
+
+## 一键性能测试
+
+```bash
+bash dlblas/kernels/ks_competition/ascend/clike_910b/run_benchmarks.sh
+```
+
+性能脚本依次调用赛事提供的 `benchmarks/ks/auto_bench.py`：
+
+- SparseAttention：预热 100 次，正式测量 1000 次。
+- Indexer：预热 100 次，正式测量 1000 次。
+- Sinkhorn：预热 100 次，正式测量 10000 次。
+
+Sinkhorn 的单次执行时间较短，因此增加正式测量次数，以尽量降低不同测试
+运行之间的波动，使中位数结果更稳定。
+
+测试过程中，参考实现和优化实现分别连续测量，不交替执行。每个正式样本
+执行一次 `forward`，随后同步 NPU，最终报告所有样本的中位数。
+构建脚本和性能脚本使用相同的 Python 探测逻辑，只选择能够同时导入
+`torch` 和 `torch_npu` 的解释器。
+
+单独执行下列命令前，可通过公共 helper 解析同一个 Python：
+
+```bash
+source dlblas/kernels/ks_competition/ascend/clike_910b/python_env.sh
+PYTHON_BIN="$(find_dlblas_python)"
+```
+
+## 单独测试 SparseAttention
+
+```bash
+"${PYTHON_BIN}" -u benchmarks/ks/auto_bench.py \
+    --v0_file dlblas/kernels/ks_competition/torch/sparse_attention.py \
+    --v1_file dlblas/kernels/ks_competition/ascend/clike_sparse_attention.py \
+    --warmup 100 \
+    --repeat 1000
+```
+
+## 单独测试 Indexer
+
+```bash
+"${PYTHON_BIN}" -u benchmarks/ks/auto_bench.py \
+    --v0_file dlblas/kernels/ks_competition/torch/indexer.py \
+    --v1_file dlblas/kernels/ks_competition/ascend/clike_indexer.py \
+    --warmup 100 \
+    --repeat 1000
+```
+
+## 单独测试 Sinkhorn
+
+```bash
+"${PYTHON_BIN}" -u benchmarks/ks/auto_bench.py \
+    --v0_file dlblas/kernels/ks_competition/torch/sinkhorn.py \
+    --v1_file dlblas/kernels/ks_competition/ascend/clike_sinkhorn.py \
+    --warmup 100 \
+    --repeat 10000
+```
+
+`auto_bench.py` 先加载由相同随机种子生成的模型和输入，将 reference 的
+state dict 加载到 `ModelNew`，检查输出正确性，再分别计时。测试通过时的
+输出格式如下：
+
+```text
+PASS accuracy; v0=<reference ms>, v1=<optimized ms>, speedup=<ratio>x
+```
+
+Task02 执行期间可能出现 NPU internal format 警告。该警告不影响输出正确性
+或计时完成，测试状态以 `PASS accuracy` 为准。
+
+`--warmup` 和 `--repeat` 是 `auto_bench.py` 提供的公开命令行参数，本地测试
+可以按需要调整，不会修改赛事提供的 benchmark 文件。性能结果需要同时记录
+这两个参数；正式成绩以赛事评测环境实际采用的参数为准。
+

dlblas/kernels/ks_competition/ascend/__init__.py

@@ -0,0 +1 @@
+"""Ascend implementations for KernelSwift competition operators."""

dlblas/kernels/ks_competition/ascend/clike_910b/.gitattributes

@@ -0,0 +1 @@
+*.sh text eol=lf

dlblas/kernels/ks_competition/ascend/clike_910b/CMakeLists.txt

@@ -0,0 +1,58 @@
+cmake_minimum_required(VERSION 3.16)
+
+find_package(ASC REQUIRED)
+project(dlblas_ks_ascendc_kernels LANGUAGES ASC CXX)
+
+find_package(Python3 COMPONENTS Interpreter REQUIRED)
+
+execute_process(
+    COMMAND ${Python3_EXECUTABLE} -c "import torch; print(torch.utils.cmake_prefix_path)"
+    OUTPUT_VARIABLE TORCH_CMAKE_PATH
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+    ERROR_VARIABLE TORCH_QUERY_ERROR
+    ERROR_STRIP_TRAILING_WHITESPACE
+    RESULT_VARIABLE TORCH_QUERY_STATUS
+)
+if(NOT TORCH_QUERY_STATUS EQUAL 0)
+    message(FATAL_ERROR
+        "${Python3_EXECUTABLE} cannot import torch: ${TORCH_QUERY_ERROR}"
+    )
+endif()
+find_package(Torch REQUIRED HINTS ${TORCH_CMAKE_PATH}/Torch)
+
+execute_process(
+    COMMAND ${Python3_EXECUTABLE} -c "import os, torch_npu; print(os.path.dirname(torch_npu.__file__))"
+    OUTPUT_VARIABLE TORCH_NPU_PATH
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+    ERROR_VARIABLE TORCH_NPU_QUERY_ERROR
+    ERROR_STRIP_TRAILING_WHITESPACE
+    RESULT_VARIABLE TORCH_NPU_QUERY_STATUS
+)
+if(NOT TORCH_NPU_QUERY_STATUS EQUAL 0)
+    message(FATAL_ERROR
+        "${Python3_EXECUTABLE} cannot import torch_npu: ${TORCH_NPU_QUERY_ERROR}"
+    )
+endif()
+
+add_library(dlblas_ks_ascendc_ops SHARED
+    csrc/sparse_attention.asc
+    csrc/indexer.asc
+    csrc/sinkhorn.asc
+)
+
+target_include_directories(dlblas_ks_ascendc_ops PRIVATE
+    ${TORCH_INCLUDE_DIRS}
+    ${TORCH_NPU_PATH}/include
+)
+
+target_link_directories(dlblas_ks_ascendc_ops PRIVATE ${TORCH_NPU_PATH}/lib)
+target_link_libraries(dlblas_ks_ascendc_ops PRIVATE torch_npu)
+
+target_compile_definitions(dlblas_ks_ascendc_ops PRIVATE
+    __GLIBCXX_USE_CXX11_ABI=0
+)
+
+target_compile_options(dlblas_ks_ascendc_ops PRIVATE
+    $<$<COMPILE_LANGUAGE:ASC>:--npu-arch=dav-2201>
+    $<$<COMPILE_LANGUAGE:ASC>:-O3>
+)

dlblas/kernels/ks_competition/ascend/clike_910b/__init__.py

@@ -0,0 +1,3 @@
+from .loader import load_library
+
+__all__ = ["load_library"]