From 78ab4010471883ff2f152ef5815053bd431761cf Mon Sep 17 00:00:00 2001 From: Zhao Shijie Date: Mon, 31 Aug 2026 18:27:38 +0800 Subject: [PATCH] add AscndC kernels --- .../Clike-AscendC/ENVIRONMENT.md | 89 +++++++ .../Clike-AscendC/PERFORMANCE.md | 79 ++++++ .../ks_competition/Clike-AscendC/README.md | 62 +++++ .../ks_competition/Clike-AscendC/RUNNING.md | 93 +++++++ .../kernels/ks_competition/ascend/__init__.py | 1 + .../ascend/clike_910b/.gitattributes | 1 + .../ascend/clike_910b/CMakeLists.txt | 58 +++++ .../ascend/clike_910b/__init__.py | 3 + .../ks_competition/ascend/clike_910b/build.sh | 17 ++ .../ascend/clike_910b/csrc/indexer.asc | 208 ++++++++++++++++ .../ascend/clike_910b/csrc/sinkhorn.asc | 151 +++++++++++ .../clike_910b/csrc/sparse_attention.asc | 134 ++++++++++ .../ascend/clike_910b/loader.py | 39 +++ .../ascend/clike_910b/python_env.sh | 32 +++ .../ascend/clike_910b/run_benchmarks.sh | 29 +++ .../ks_competition/ascend/clike_indexer.py | 219 ++++++++++++++++ .../ks_competition/ascend/clike_sinkhorn.py | 32 +++ .../ascend/clike_sparse_attention.py | 65 +++++ .../kernels/ks_competition/torch/indexer.py | 235 ++++++++++++++++++ .../kernels/ks_competition/torch/sinkhorn.py | 37 +++ .../ks_competition/torch/sparse_attention.py | 126 ++++++++++ 21 files changed, 1710 insertions(+) create mode 100644 dlblas/kernels/ks_competition/Clike-AscendC/ENVIRONMENT.md create mode 100644 dlblas/kernels/ks_competition/Clike-AscendC/PERFORMANCE.md create mode 100644 dlblas/kernels/ks_competition/Clike-AscendC/README.md create mode 100644 dlblas/kernels/ks_competition/Clike-AscendC/RUNNING.md create mode 100644 dlblas/kernels/ks_competition/ascend/__init__.py create mode 100644 dlblas/kernels/ks_competition/ascend/clike_910b/.gitattributes create mode 100644 dlblas/kernels/ks_competition/ascend/clike_910b/CMakeLists.txt create mode 100644 dlblas/kernels/ks_competition/ascend/clike_910b/__init__.py create mode 100644 dlblas/kernels/ks_competition/ascend/clike_910b/build.sh create mode 100644 dlblas/kernels/ks_competition/ascend/clike_910b/csrc/indexer.asc create mode 100644 dlblas/kernels/ks_competition/ascend/clike_910b/csrc/sinkhorn.asc create mode 100644 dlblas/kernels/ks_competition/ascend/clike_910b/csrc/sparse_attention.asc create mode 100644 dlblas/kernels/ks_competition/ascend/clike_910b/loader.py create mode 100644 dlblas/kernels/ks_competition/ascend/clike_910b/python_env.sh create mode 100644 dlblas/kernels/ks_competition/ascend/clike_910b/run_benchmarks.sh create mode 100644 dlblas/kernels/ks_competition/ascend/clike_indexer.py create mode 100644 dlblas/kernels/ks_competition/ascend/clike_sinkhorn.py create mode 100644 dlblas/kernels/ks_competition/ascend/clike_sparse_attention.py create mode 100644 dlblas/kernels/ks_competition/torch/indexer.py create mode 100644 dlblas/kernels/ks_competition/torch/sinkhorn.py create mode 100644 dlblas/kernels/ks_competition/torch/sparse_attention.py diff --git a/dlblas/kernels/ks_competition/Clike-AscendC/ENVIRONMENT.md b/dlblas/kernels/ks_competition/Clike-AscendC/ENVIRONMENT.md new file mode 100644 index 000000000..76f4e4e77 --- /dev/null +++ b/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 的提交内容。 + diff --git a/dlblas/kernels/ks_competition/Clike-AscendC/PERFORMANCE.md b/dlblas/kernels/ks_competition/Clike-AscendC/PERFORMANCE.md new file mode 100644 index 000000000..bc3ac1b03 --- /dev/null +++ b/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 使用完全相同的中间精度，因此允许 +混合精度优化；最终输出仍须满足上述浮点容差或整数精确比较，并且实际执行 +路径必须调用自定义算子。 + diff --git a/dlblas/kernels/ks_competition/Clike-AscendC/README.md b/dlblas/kernels/ks_competition/Clike-AscendC/README.md new file mode 100644 index 000000000..e342460fd --- /dev/null +++ b/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) + diff --git a/dlblas/kernels/ks_competition/Clike-AscendC/RUNNING.md b/dlblas/kernels/ks_competition/Clike-AscendC/RUNNING.md new file mode 100644 index 000000000..7cc644884 --- /dev/null +++ b/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=, v1=, speedup=x +``` + +Task02 执行期间可能出现 NPU internal format 警告。该警告不影响输出正确性 +或计时完成，测试状态以 `PASS accuracy` 为准。 + +`--warmup` 和 `--repeat` 是 `auto_bench.py` 提供的公开命令行参数，本地测试 +可以按需要调整，不会修改赛事提供的 benchmark 文件。性能结果需要同时记录 +这两个参数；正式成绩以赛事评测环境实际采用的参数为准。 + diff --git a/dlblas/kernels/ks_competition/ascend/__init__.py b/dlblas/kernels/ks_competition/ascend/__init__.py new file mode 100644 index 000000000..aae7a02ea --- /dev/null +++ b/dlblas/kernels/ks_competition/ascend/__init__.py @@ -0,0 +1 @@ +"""Ascend implementations for KernelSwift competition operators.""" diff --git a/dlblas/kernels/ks_competition/ascend/clike_910b/.gitattributes b/dlblas/kernels/ks_competition/ascend/clike_910b/.gitattributes new file mode 100644 index 000000000..dfdb8b771 --- /dev/null +++ b/dlblas/kernels/ks_competition/ascend/clike_910b/.gitattributes @@ -0,0 +1 @@ +*.sh text eol=lf diff --git a/dlblas/kernels/ks_competition/ascend/clike_910b/CMakeLists.txt b/dlblas/kernels/ks_competition/ascend/clike_910b/CMakeLists.txt new file mode 100644 index 000000000..2e9712aaa --- /dev/null +++ b/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 + $<$:--npu-arch=dav-2201> + $<$:-O3> +) diff --git a/dlblas/kernels/ks_competition/ascend/clike_910b/__init__.py b/dlblas/kernels/ks_competition/ascend/clike_910b/__init__.py new file mode 100644 index 000000000..2d4c7fe41 --- /dev/null +++ b/dlblas/kernels/ks_competition/ascend/clike_910b/__init__.py @@ -0,0 +1,3 @@ +from .loader import load_library + +__all__ = ["load_library"] diff --git a/dlblas/kernels/ks_competition/ascend/clike_910b/build.sh b/dlblas/kernels/ks_competition/ascend/clike_910b/build.sh new file mode 100644 index 000000000..3b39a75ec --- /dev/null +++ b/dlblas/kernels/ks_competition/ascend/clike_910b/build.sh @@ -0,0 +1,17 @@ +#!/usr/bin/env bash +# Keep shell scripts LF-only; Bash treats a trailing CR as part of option names. +set -euo pipefail + +PROJECT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ASCEND_ROOT="${ASCEND_HOME_PATH:-/usr/local/Ascend/cann-9.0.0}" +ASCEND_CMAKE="${ASCEND_ROOT}/aarch64-linux/tikcpp/ascendc_kernel_cmake" + +source "${PROJECT_DIR}/python_env.sh" +PYTHON_BIN="$(find_dlblas_python)" +echo "Using Python: ${PYTHON_BIN}" + +cmake -S "${PROJECT_DIR}" -B "${PROJECT_DIR}/build" \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_PREFIX_PATH="${ASCEND_CMAKE}" \ + -DPython3_EXECUTABLE:FILEPATH="${PYTHON_BIN}" +cmake --build "${PROJECT_DIR}/build" --parallel "${BUILD_JOBS:-64}" diff --git a/dlblas/kernels/ks_competition/ascend/clike_910b/csrc/indexer.asc b/dlblas/kernels/ks_competition/ascend/clike_910b/csrc/indexer.asc new file mode 100644 index 000000000..74e4bcd78 --- /dev/null +++ b/dlblas/kernels/ks_competition/ascend/clike_910b/csrc/indexer.asc @@ -0,0 +1,208 @@ +#include +#include + +#include "kernel_operator.h" +#include "torch_npu/csrc/core/npu/NPUStream.h" + +namespace { + +constexpr uint32_t kVectorCores = 40; +constexpr uint32_t kIndexerHeads = 16; +constexpr uint32_t kIndexerMaxKeys = 656; + +template +__aicore__ inline void SyncPipelines() +{ + auto eventId = GetTPipePtr()->FetchEventID(Event); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); +} + +class IndexerReduceKernel { +public: + __aicore__ inline void Init(GM_ADDR scores, GM_ADDR weights, GM_ADDR output, uint32_t batch, + uint32_t seqLength, uint32_t physicalKeys, uint32_t keyLength, uint32_t ratio, bool causal, + AscendC::TPipe* pipe) + { + batch_ = batch; + seqLength_ = seqLength; + physicalKeys_ = physicalKeys; + keyLength_ = keyLength; + alignedKeys_ = (keyLength + 15U) & ~15U; + ratio_ = ratio; + causal_ = causal; + totalTasks_ = batch * seqLength; + scoreGm_.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t*>(scores), + totalTasks_ * kIndexerHeads * physicalKeys); + weightGm_.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t*>(weights), totalTasks_ * kIndexerHeads); + outputGm_.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t*>(output), totalTasks_ * keyLength); + + pipe->InitBuffer(scoreQueue_, 2, kIndexerHeads * kIndexerMaxKeys * sizeof(bfloat16_t)); + pipe->InitBuffer(scoreFloatBuf_, kIndexerHeads * kIndexerMaxKeys * sizeof(float)); + pipe->InitBuffer(weightQueue_, 2, 32 * sizeof(bfloat16_t)); + pipe->InitBuffer(weightFloatBuf_, 32 * sizeof(float)); + pipe->InitBuffer(tmpBfBuf_, kIndexerMaxKeys * sizeof(bfloat16_t)); + pipe->InitBuffer(accBuf_, kIndexerMaxKeys * sizeof(float)); + pipe->InitBuffer(tmpBuf_, kIndexerMaxKeys * sizeof(float)); + pipe->InitBuffer(outQueue_, 2, kIndexerMaxKeys * sizeof(bfloat16_t)); + } + + __aicore__ inline void Process() + { + for (uint32_t task = AscendC::GetBlockIdx(); task < totalTasks_; task += AscendC::GetBlockNum()) { + CopyIn(task); + Compute(task); + CopyOut(task); + } + } + +private: + __aicore__ inline void CopyIn(uint32_t task) + { + auto scoreBf = scoreQueue_.AllocTensor(); + auto weightBf = weightQueue_.AllocTensor(); + const uint32_t scoreBase = task * kIndexerHeads * physicalKeys_; + if (physicalKeys_ == alignedKeys_) { + AscendC::DataCopy(scoreBf, scoreGm_[scoreBase], kIndexerHeads * alignedKeys_); + } else { + AscendC::DataCopyExtParams copyParams = {static_cast(kIndexerHeads), + static_cast(keyLength_ * sizeof(bfloat16_t)), 0, 0, 0}; + const uint8_t rightPad = static_cast(alignedKeys_ - keyLength_); + AscendC::DataCopyPadExtParams padParams = {false, 0, rightPad, 0}; + AscendC::DataCopyPad(scoreBf, scoreGm_[scoreBase], copyParams, padParams); + } + AscendC::DataCopy(weightBf, weightGm_[task * kIndexerHeads], 16); + scoreQueue_.EnQue(scoreBf); + weightQueue_.EnQue(weightBf); + } + + __aicore__ inline void Compute(uint32_t task) + { + auto scoreBf = scoreQueue_.DeQue(); + auto weightBf = weightQueue_.DeQue(); + auto outBf = outQueue_.AllocTensor(); + auto scoreFloat = scoreFloatBuf_.Get(); + auto weightFloat = weightFloatBuf_.Get(); + auto tmpBf = tmpBfBuf_.Get(); + auto acc = accBuf_.Get(); + auto tmp = tmpBuf_.Get(); + + AscendC::Cast(scoreFloat, scoreBf, AscendC::RoundMode::CAST_NONE, + kIndexerHeads * alignedKeys_); + AscendC::Relu(scoreFloat, scoreFloat, kIndexerHeads * alignedKeys_); + AscendC::Cast(weightFloat, weightBf, AscendC::RoundMode::CAST_NONE, kIndexerHeads); + AscendC::Duplicate(acc, 0.0f, alignedKeys_); + SyncPipelines(); + for (uint32_t head = 0; head < kIndexerHeads; ++head) { + const float weight = weightFloat.GetValue(head); + AscendC::Muls(tmp, scoreFloat[head * alignedKeys_], weight, alignedKeys_); + AscendC::Cast(tmpBf, tmp, AscendC::RoundMode::CAST_RINT, alignedKeys_); + AscendC::Cast(tmp, tmpBf, AscendC::RoundMode::CAST_NONE, alignedKeys_); + AscendC::Add(acc, acc, tmp, alignedKeys_); + } + + SyncPipelines(); + if (causal_) { + const uint32_t seq = task % seqLength_; + uint32_t validKeys = (seq + 1) / ratio_; + validKeys = validKeys < keyLength_ ? validKeys : keyLength_; + uint32_t alignedStart = (validKeys + 7U) & ~7U; + alignedStart = alignedStart < alignedKeys_ ? alignedStart : alignedKeys_; + for (uint32_t key = validKeys; key < alignedStart; ++key) { + acc.SetValue(key, -3.402823466e+38f); + } + SyncPipelines(); + if (alignedStart < alignedKeys_) { + AscendC::Duplicate( + acc[alignedStart], -3.402823466e+38f, alignedKeys_ - alignedStart); + } + } else if (keyLength_ < alignedKeys_) { + for (uint32_t key = keyLength_; key < alignedKeys_; ++key) { + acc.SetValue(key, -3.402823466e+38f); + } + } + + SyncPipelines(); + AscendC::Cast(outBf, acc, AscendC::RoundMode::CAST_RINT, alignedKeys_); + outQueue_.EnQue(outBf); + scoreQueue_.FreeTensor(scoreBf); + weightQueue_.FreeTensor(weightBf); + } + + __aicore__ inline void CopyOut(uint32_t task) + { + auto outBf = outQueue_.DeQue(); + AscendC::DataCopyExtParams outParams = { + 1, static_cast(keyLength_ * sizeof(bfloat16_t)), 0, 0, 0}; + AscendC::DataCopyPad(outputGm_[task * keyLength_], outBf, outParams); + outQueue_.FreeTensor(outBf); + } + + AscendC::TQue scoreQueue_, weightQueue_; + AscendC::TQue outQueue_; + AscendC::TBuf scoreFloatBuf_; + AscendC::TBuf weightFloatBuf_; + AscendC::TBuf tmpBfBuf_, accBuf_, tmpBuf_; + AscendC::GlobalTensor scoreGm_, weightGm_, outputGm_; + uint32_t batch_, seqLength_, physicalKeys_, keyLength_, alignedKeys_, ratio_, totalTasks_; + bool causal_; +}; + +__global__ __vector__ void indexer_reduce_kernel(GM_ADDR scores, GM_ADDR weights, GM_ADDR output, + uint32_t batch, uint32_t seqLength, uint32_t physicalKeys, uint32_t keyLength, + uint32_t ratio, bool causal) +{ + AscendC::TPipe pipe; + IndexerReduceKernel kernel; + kernel.Init(scores, weights, output, batch, seqLength, physicalKeys, keyLength, ratio, causal, &pipe); + kernel.Process(); +} + +inline void CheckNpuContiguous(const at::Tensor& tensor, const char* name) +{ + TORCH_CHECK(tensor.device().type() == c10::DeviceType::PrivateUse1, name, " must be an NPU tensor"); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); +} + +} // namespace + +namespace dlblas_ks_ascendc { + +at::Tensor indexer_reduce(const at::Tensor& scores, const at::Tensor& weights, int64_t ratio, + bool causal, int64_t keyLength) +{ + CheckNpuContiguous(scores, "scores"); + CheckNpuContiguous(weights, "weights"); + TORCH_CHECK(scores.scalar_type() == at::kBFloat16 && weights.scalar_type() == at::kBFloat16, + "scores and weights must be bfloat16"); + TORCH_CHECK(scores.dim() == 4 && weights.dim() == 3, "invalid indexer ranks"); + TORCH_CHECK(scores.size(2) == kIndexerHeads && weights.size(2) == kIndexerHeads, + "optimized indexer reduction requires 16 heads"); + TORCH_CHECK(scores.size(3) <= kIndexerMaxKeys && keyLength > 0 && keyLength <= scores.size(3) && ratio > 0, + "unsupported key length or ratio"); + TORCH_CHECK(scores.size(0) == weights.size(0) && scores.size(1) == weights.size(1), + "incompatible score and weight shapes"); + + auto output = at::empty({scores.size(0), scores.size(1), keyLength}, scores.options()); + auto stream = c10_npu::getCurrentNPUStream().stream(true); + indexer_reduce_kernel<<>>( + reinterpret_cast(scores.mutable_data_ptr()), + reinterpret_cast(weights.mutable_data_ptr()), + reinterpret_cast(output.mutable_data_ptr()), + static_cast(scores.size(0)), static_cast(scores.size(1)), + static_cast(scores.size(3)), static_cast(keyLength), + static_cast(ratio), causal); + return output; +} + +} // namespace dlblas_ks_ascendc + +TORCH_LIBRARY_FRAGMENT(dlblas_ks_ascendc, m) +{ + m.def("indexer_reduce(Tensor scores, Tensor weights, int ratio, bool causal, int key_length) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(dlblas_ks_ascendc, PrivateUse1, m) +{ + m.impl("indexer_reduce", TORCH_FN(dlblas_ks_ascendc::indexer_reduce)); +} diff --git a/dlblas/kernels/ks_competition/ascend/clike_910b/csrc/sinkhorn.asc b/dlblas/kernels/ks_competition/ascend/clike_910b/csrc/sinkhorn.asc new file mode 100644 index 000000000..a8077e9ae --- /dev/null +++ b/dlblas/kernels/ks_competition/ascend/clike_910b/csrc/sinkhorn.asc @@ -0,0 +1,151 @@ +#include +#include + +#include "kernel_operator.h" +#include "torch_npu/csrc/core/npu/NPUStream.h" + +namespace { + +constexpr uint32_t kVectorCores = 40; + +class SinkhornKernel { +public: + __aicore__ inline void Init(GM_ADDR x, GM_ADDR output, uint32_t matrices, uint32_t repeat, + float eps, AscendC::TPipe* pipe) + { + matrices_ = matrices; + repeat_ = repeat; + eps_ = eps; + xGm_.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(x), matrices * 16); + outputGm_.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(output), matrices * 16); + pipe->InitBuffer(valueBuf_, 32 * sizeof(float)); + } + + __aicore__ inline void Process() + { + auto values = valueBuf_.Get(); + for (uint32_t matrix = AscendC::GetBlockIdx(); matrix < matrices_; matrix += AscendC::GetBlockNum()) { + AscendC::DataCopy(values, xGm_[matrix * 16], 16); + AscendC::PipeBarrier(); + + for (uint32_t row = 0; row < 4; ++row) { + float maximum = values.GetValue(row * 4); + for (uint32_t col = 1; col < 4; ++col) { + const float value = values.GetValue(row * 4 + col); + maximum = value > maximum ? value : maximum; + } + for (uint32_t col = 0; col < 4; ++col) { + const uint32_t offset = row * 4 + col; + values.SetValue(offset, values.GetValue(offset) - maximum); + } + } + AscendC::Exp(values, values, 16); + AscendC::PipeBarrier(); + + for (uint32_t row = 0; row < 4; ++row) { + float sum = 0.0f; + for (uint32_t col = 0; col < 4; ++col) { + sum += values.GetValue(row * 4 + col); + } + const float inverse = 1.0f / sum; + for (uint32_t col = 0; col < 4; ++col) { + const uint32_t offset = row * 4 + col; + values.SetValue(offset, values.GetValue(offset) * inverse + eps_); + } + } + NormalizeColumns(values); + + for (uint32_t iteration = 1; iteration < repeat_; ++iteration) { + NormalizeRows(values); + NormalizeColumns(values); + } + + AscendC::DataCopy(outputGm_[matrix * 16], values, 16); + AscendC::PipeBarrier(); + } + } + +private: + __aicore__ inline void NormalizeRows(AscendC::LocalTensor& values) + { + for (uint32_t row = 0; row < 4; ++row) { + float sum = 0.0f; + for (uint32_t col = 0; col < 4; ++col) { + sum += values.GetValue(row * 4 + col); + } + const float inverse = 1.0f / (sum + eps_); + for (uint32_t col = 0; col < 4; ++col) { + const uint32_t offset = row * 4 + col; + values.SetValue(offset, values.GetValue(offset) * inverse); + } + } + } + + __aicore__ inline void NormalizeColumns(AscendC::LocalTensor& values) + { + for (uint32_t col = 0; col < 4; ++col) { + float sum = 0.0f; + for (uint32_t row = 0; row < 4; ++row) { + sum += values.GetValue(row * 4 + col); + } + const float inverse = 1.0f / (sum + eps_); + for (uint32_t row = 0; row < 4; ++row) { + const uint32_t offset = row * 4 + col; + values.SetValue(offset, values.GetValue(offset) * inverse); + } + } + } + + AscendC::TBuf valueBuf_; + AscendC::GlobalTensor xGm_, outputGm_; + uint32_t matrices_, repeat_; + float eps_; +}; + +__global__ __vector__ void sinkhorn_kernel(GM_ADDR x, GM_ADDR output, uint32_t matrices, + uint32_t repeat, float eps) +{ + AscendC::TPipe pipe; + SinkhornKernel kernel; + kernel.Init(x, output, matrices, repeat, eps, &pipe); + kernel.Process(); +} + +inline void CheckNpuContiguous(const at::Tensor& tensor, const char* name) +{ + TORCH_CHECK(tensor.device().type() == c10::DeviceType::PrivateUse1, name, " must be an NPU tensor"); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); +} + +} // namespace + +namespace dlblas_ks_ascendc { + +at::Tensor sinkhorn(const at::Tensor& x, int64_t repeat, double eps) +{ + CheckNpuContiguous(x, "x"); + TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32"); + TORCH_CHECK(x.dim() >= 2 && x.size(-1) == 4 && x.size(-2) == 4, + "optimized sinkhorn requires 4x4 matrices"); + TORCH_CHECK(repeat >= 1, "repeat must be positive"); + auto output = at::empty_like(x); + const uint32_t matrices = static_cast(x.numel() / 16); + auto stream = c10_npu::getCurrentNPUStream().stream(true); + sinkhorn_kernel<<>>( + reinterpret_cast(x.mutable_data_ptr()), + reinterpret_cast(output.mutable_data_ptr()), matrices, + static_cast(repeat), static_cast(eps)); + return output; +} + +} // namespace dlblas_ks_ascendc + +TORCH_LIBRARY_FRAGMENT(dlblas_ks_ascendc, m) +{ + m.def("sinkhorn(Tensor x, int repeat, float eps) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(dlblas_ks_ascendc, PrivateUse1, m) +{ + m.impl("sinkhorn", TORCH_FN(dlblas_ks_ascendc::sinkhorn)); +} diff --git a/dlblas/kernels/ks_competition/ascend/clike_910b/csrc/sparse_attention.asc b/dlblas/kernels/ks_competition/ascend/clike_910b/csrc/sparse_attention.asc new file mode 100644 index 000000000..71096dd64 --- /dev/null +++ b/dlblas/kernels/ks_competition/ascend/clike_910b/csrc/sparse_attention.asc @@ -0,0 +1,134 @@ +#include +#include + +#include "kernel_operator.h" +#include "torch_npu/csrc/core/npu/NPUStream.h" + +namespace { + +constexpr uint32_t kVectorCores = 40; +constexpr uint32_t kSparseDim = 128; +constexpr uint32_t kSparseTopK = 16; +constexpr uint32_t kSparseMaxKeys = 32; + +template +__aicore__ inline void SyncPipelines() +{ + auto eventId = GetTPipePtr()->FetchEventID(Event); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); +} + +class SparseGatherKernel { +public: + __aicore__ inline void Init(GM_ADDR kv, GM_ADDR indices, GM_ADDR output, uint32_t batch, + uint32_t queryLength, uint32_t kvLength, AscendC::TPipe* pipe) + { + batch_ = batch; + queryLength_ = queryLength; + kvLength_ = kvLength; + totalTasks_ = batch * queryLength; + kvGm_.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t*>(kv), batch * kvLength * kSparseDim); + indexGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(indices), totalTasks_ * kSparseTopK); + outputGm_.SetGlobalBuffer( + reinterpret_cast<__gm__ bfloat16_t*>(output), totalTasks_ * kSparseTopK * kSparseDim); + pipe->InitBuffer(indexBuf_, 32 * sizeof(int32_t)); + pipe->InitBuffer(kvCacheBuf_, kSparseMaxKeys * kSparseDim * sizeof(bfloat16_t)); + pipe->InitBuffer(gatherQueue_, 2, kSparseTopK * kSparseDim * sizeof(bfloat16_t)); + } + + __aicore__ inline void Process() + { + auto indices = indexBuf_.Get(); + auto kvCache = kvCacheBuf_.Get(); + const uint32_t block = AscendC::GetBlockIdx(); + if (block >= queryLength_) { + return; + } + for (uint32_t batch = 0; batch < batch_; ++batch) { + AscendC::DataCopy( + kvCache, kvGm_[batch * kvLength_ * kSparseDim], kvLength_ * kSparseDim); + SyncPipelines(); + for (uint32_t query = block; query < queryLength_; query += AscendC::GetBlockNum()) { + const uint32_t task = batch * queryLength_ + query; + AscendC::DataCopy(indices, indexGm_[task * kSparseTopK], kSparseTopK); + SyncPipelines(); + auto gathered = gatherQueue_.AllocTensor(); + for (uint32_t i = 0; i < kSparseTopK; ++i) { + const int32_t index = indices.GetValue(i); + const uint32_t source = index >= 0 && static_cast(index) < kvLength_ + ? static_cast(index) * kSparseDim : 0; + AscendC::DataCopy(gathered[i * kSparseDim], kvCache[source], kSparseDim); + } + gatherQueue_.EnQue(gathered); + auto output = gatherQueue_.DeQue(); + AscendC::DataCopy( + outputGm_[task * kSparseTopK * kSparseDim], output, kSparseTopK * kSparseDim); + gatherQueue_.FreeTensor(output); + } + SyncPipelines(); + } + } + +private: + AscendC::TBuf indexBuf_, kvCacheBuf_; + AscendC::TQueBind gatherQueue_; + AscendC::GlobalTensor kvGm_, outputGm_; + AscendC::GlobalTensor indexGm_; + uint32_t batch_, queryLength_, kvLength_, totalTasks_; +}; + +__global__ __vector__ void sparse_gather_kernel(GM_ADDR kv, GM_ADDR indices, GM_ADDR output, + uint32_t batch, uint32_t queryLength, uint32_t kvLength) +{ + AscendC::TPipe pipe; + SparseGatherKernel kernel; + kernel.Init(kv, indices, output, batch, queryLength, kvLength, &pipe); + kernel.Process(); +} + +inline void CheckNpuContiguous(const at::Tensor& tensor, const char* name) +{ + TORCH_CHECK(tensor.device().type() == c10::DeviceType::PrivateUse1, name, " must be an NPU tensor"); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); +} + +} // namespace + +namespace dlblas_ks_ascendc { + +at::Tensor sparse_gather(const at::Tensor& kv, const at::Tensor& indices) +{ + CheckNpuContiguous(kv, "kv"); + CheckNpuContiguous(indices, "indices"); + TORCH_CHECK(kv.scalar_type() == at::kBFloat16, "kv must be bfloat16"); + TORCH_CHECK(indices.scalar_type() == at::kInt, "indices must be int32"); + TORCH_CHECK(kv.dim() == 3 && indices.dim() == 3 && kv.size(0) == indices.size(0), + "invalid sparse gather shapes"); + TORCH_CHECK(kv.size(2) == kSparseDim && indices.size(2) == kSparseTopK, + "optimized sparse gather requires dim=128 and topk=16"); + TORCH_CHECK(kv.size(1) <= kSparseMaxKeys, "optimized sparse gather supports at most 32 KV rows"); + auto output = at::empty( + {indices.size(0), indices.size(1), indices.size(2), kv.size(2)}, kv.options()); + auto stream = c10_npu::getCurrentNPUStream().stream(true); + sparse_gather_kernel<<>>( + reinterpret_cast(kv.mutable_data_ptr()), + reinterpret_cast(indices.mutable_data_ptr()), + reinterpret_cast(output.mutable_data_ptr()), + static_cast(kv.size(0)), static_cast(indices.size(1)), + static_cast(kv.size(1))); + return output; +} + +} // namespace dlblas_ks_ascendc + +TORCH_LIBRARY_FRAGMENT(dlblas_ks_ascendc, m) +{ + m.def("sparse_gather(Tensor kv, Tensor indices) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(dlblas_ks_ascendc, PrivateUse1, m) +{ + m.impl("sparse_gather", TORCH_FN(dlblas_ks_ascendc::sparse_gather)); +} diff --git a/dlblas/kernels/ks_competition/ascend/clike_910b/loader.py b/dlblas/kernels/ks_competition/ascend/clike_910b/loader.py new file mode 100644 index 000000000..bd0eda44f --- /dev/null +++ b/dlblas/kernels/ks_competition/ascend/clike_910b/loader.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import os +from pathlib import Path + +import torch +import torch_npu # noqa: F401 + + +_LOADED = False + + +def load_library() -> None: + global _LOADED + if _LOADED: + return + + override = os.getenv("DLBLAS_KS_ASCENDC_LIBRARY") + candidates = [] + if override: + candidates.append(Path(override)) + root = Path(__file__).resolve().parent + candidates.extend( + [ + root / "build" / "libdlblas_ks_ascendc_ops.so", + root / "lib" / "libdlblas_ks_ascendc_ops.so", + ] + ) + for candidate in candidates: + if candidate.is_file(): + torch.ops.load_library(str(candidate)) + _LOADED = True + return + + searched = ", ".join(str(path) for path in candidates) + raise RuntimeError( + "DLBlas KernelSwift AscendC library was not found. " + f"Searched: {searched}. Run ascend/clike_910b/build.sh first." + ) diff --git a/dlblas/kernels/ks_competition/ascend/clike_910b/python_env.sh b/dlblas/kernels/ks_competition/ascend/clike_910b/python_env.sh new file mode 100644 index 000000000..217f122b7 --- /dev/null +++ b/dlblas/kernels/ks_competition/ascend/clike_910b/python_env.sh @@ -0,0 +1,32 @@ +find_dlblas_python() { + local candidate resolved + local -a candidates=() + + if [[ -n "${DLBLAS_PYTHON_EXECUTABLE:-}" ]]; then + candidates+=("${DLBLAS_PYTHON_EXECUTABLE}") + else + candidates+=(python3 python) + for candidate in /usr/local/python*/bin/python3 /usr/local/bin/python3; do + candidates+=("${candidate}") + done + fi + + for candidate in "${candidates[@]}"; do + if [[ "${candidate}" == */* ]]; then + resolved="${candidate}" + else + resolved="$(command -v "${candidate}" 2>/dev/null || true)" + fi + if [[ -z "${resolved}" || ! -x "${resolved}" ]]; then + continue + fi + if "${resolved}" -c "import torch, torch_npu" >/dev/null 2>&1; then + "${resolved}" -c "import os, sys; print(os.path.realpath(sys.executable))" + return 0 + fi + done + + echo "No Python interpreter with both torch and torch_npu was found." >&2 + echo "Set DLBLAS_PYTHON_EXECUTABLE to the intended Python executable." >&2 + return 1 +} diff --git a/dlblas/kernels/ks_competition/ascend/clike_910b/run_benchmarks.sh b/dlblas/kernels/ks_competition/ascend/clike_910b/run_benchmarks.sh new file mode 100644 index 000000000..172907c7d --- /dev/null +++ b/dlblas/kernels/ks_competition/ascend/clike_910b/run_benchmarks.sh @@ -0,0 +1,29 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../../.." && pwd)" +source "${SCRIPT_DIR}/python_env.sh" +PYTHON_BIN="$(find_dlblas_python)" + +cd "${REPO_ROOT}" +export PYTHONDONTWRITEBYTECODE=1 +echo "Using Python: ${PYTHON_BIN}" + +"${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 + +"${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 + +"${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 diff --git a/dlblas/kernels/ks_competition/ascend/clike_indexer.py b/dlblas/kernels/ks_competition/ascend/clike_indexer.py new file mode 100644 index 000000000..573539666 --- /dev/null +++ b/dlblas/kernels/ks_competition/ascend/clike_indexer.py @@ -0,0 +1,219 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal, Tuple + +import torch +import torch.nn.functional as F +from torch import nn + +from clike_910b import load_library + + +world_size = 1 + + +@dataclass +class ModelArgs: + max_batch_size: int = 4 + max_seq_len: int = 4096 + dtype: Literal["bf16", "fp8"] = "fp8" + scale_fmt: Literal[None, "ue8m0"] = "ue8m0" + expert_dtype: Literal[None, "fp4"] = None + scale_dtype: Literal["fp32", "fp8"] = "fp8" + vocab_size: int = 129280 + dim: int = 4096 + moe_inter_dim: int = 4096 + n_layers: int = 7 + n_hash_layers: int = 0 + n_mtp_layers: int = 1 + n_heads: int = 64 + n_routed_experts: int = 8 + n_shared_experts: int = 1 + n_activated_experts: int = 2 + score_func: Literal["softmax", "sigmoid", "sqrtsoftplus"] = "sqrtsoftplus" + route_scale: float = 1.0 + swiglu_limit: float = 0.0 + q_lora_rank: int = 1024 + head_dim: int = 512 + rope_head_dim: int = 64 + norm_eps: float = 1e-6 + o_groups: int = 8 + o_lora_rank: int = 1024 + window_size: int = 128 + compress_ratios: Tuple[int, ...] = (0, 0, 4, 128, 4, 128, 4, 0) + compress_rope_theta: float = 40000.0 + original_seq_len: int = 0 + rope_theta: float = 10000.0 + rope_factor: float = 40 + beta_fast: int = 32 + beta_slow: int = 1 + index_n_heads: int = 64 + index_head_dim: int = 128 + index_topk: int = 512 + hc_mult: int = 4 + hc_sinkhorn_iters: int = 20 + hc_eps: float = 1e-6 + + +class ModelNew(nn.Module): + def __init__( + self, + args: ModelArgs, + freqs_cis: torch.Tensor, + kv_cache: torch.Tensor, + compress_ratio: int = 4, + ): + super().__init__() + if ( + args.index_n_heads != 16 + or args.index_head_dim != 64 + or args.rope_head_dim != 32 + ): + raise ValueError( + "The optimized Indexer expects 16 heads, head_dim=64, and rope_head_dim=32" + ) + self.dim = args.dim + self.n_heads = args.index_n_heads + self.n_local_heads = args.index_n_heads // world_size + self.head_dim = args.index_head_dim + self.rope_head_dim = args.rope_head_dim + self.index_topk = args.index_topk + self.q_lora_rank = args.q_lora_rank + self.wq_b = nn.Linear( + self.q_lora_rank, + self.n_heads * self.head_dim, + bias=False, + dtype=torch.bfloat16, + ) + self.weights_proj = nn.Linear( + self.dim, self.n_heads, bias=False, dtype=torch.bfloat16 + ) + self.softmax_scale = self.head_dim**-0.5 + self.compress_ratio = compress_ratio + self.register_buffer("freqs_cis", freqs_cis, persistent=False) + self.register_buffer("kv_cache", kv_cache, persistent=False) + + aligned_cache_len = (kv_cache.size(1) + 15) // 16 * 16 + self.register_buffer( + "_kv_cache_aligned", + F.pad(kv_cache, (0, 0, 0, aligned_cache_len - kv_cache.size(1))), + persistent=False, + ) + self.register_buffer( + "_valid_key_counts", + torch.arange( + 1, + args.max_seq_len + 1, + dtype=torch.int64, + device=kv_cache.device, + ).unsqueeze(1) + // compress_ratio, + persistent=False, + ) + + def forward(self, x: torch.Tensor, qr: torch.Tensor, start_pos: int, offset: int): + load_library() + batch_size, seq_len, _ = x.shape + end_pos = start_pos + seq_len + key_len = end_pos // self.compress_ratio + if key_len > 650: + raise ValueError( + "The optimized Indexer supports at most 650 compressed keys" + ) + + q = F.linear(qr, self.wq_b.weight) + q = q.unflatten(-1, (self.n_local_heads, self.head_dim)).contiguous() + rope = torch.view_as_complex( + q[..., -self.rope_head_dim :].float().unflatten(-1, (-1, 2)) + ) + freqs = self.freqs_cis[start_pos:end_pos].view(1, seq_len, 1, -1) + q[..., -self.rope_head_dim :].copy_( + torch.view_as_real(rope * freqs).flatten(-2) + ) + + weight_scale = self.softmax_scale * self.n_heads**-0.5 + weights = (F.linear(x, self.weights_proj.weight) * weight_scale).contiguous() + physical_key_len = (key_len + 15) // 16 * 16 + key = self._kv_cache_aligned[:batch_size, :physical_key_len] + + scores = torch.bmm( + q.view(batch_size, seq_len * self.n_local_heads, self.head_dim), + key.transpose(1, 2), + ).view(batch_size, seq_len, self.n_local_heads, physical_key_len) + reduced = torch.ops.dlblas_ks_ascendc.indexer_reduce( + scores, + weights, + self.compress_ratio, + start_pos == 0, + key_len, + ) + + count = min(self.index_topk, key_len) + topk_idxs = reduced.topk(count, dim=-1).indices + if start_pos == 0: + valid = self._valid_key_counts[:seq_len] + topk_idxs = torch.where(topk_idxs >= valid, -1, topk_idxs + offset) + else: + topk_idxs = topk_idxs + offset + return topk_idxs + + +def _make_args(): + return ModelArgs( + max_batch_size=8, + max_seq_len=2600, + dim=1024, + index_n_heads=16, + index_head_dim=64, + index_topk=128, + q_lora_rank=256, + rope_head_dim=32, + ) + + +def get_inputs(): + config = _make_args() + batch_size = 8 + seq_len = 2600 + x = torch.randn( + batch_size, + seq_len, + config.dim, + dtype=torch.bfloat16, + device="npu", + ) + qr = torch.randn( + batch_size, + seq_len, + config.q_lora_rank, + dtype=torch.bfloat16, + device="npu", + ) + return [x, qr, 0, 0] + + +def get_init_inputs(): + config = _make_args() + compress_ratio = 4 + freqs = 1.0 / ( + 10000.0 + ** ( + torch.arange(0, config.rope_head_dim, 2)[ + : config.rope_head_dim // 2 + ].float() + / config.rope_head_dim + ) + ) + positions = torch.arange(config.max_seq_len, dtype=torch.float32) + angles = torch.outer(positions, freqs).float().npu() + freqs_cis = torch.polar(torch.ones_like(angles).npu(), angles).view( + config.max_seq_len, -1 + ) + kv_cache = torch.randn( + config.max_batch_size, + config.max_seq_len // compress_ratio, + config.index_head_dim, + dtype=torch.bfloat16, + ).npu() + return [config, freqs_cis, kv_cache, compress_ratio] diff --git a/dlblas/kernels/ks_competition/ascend/clike_sinkhorn.py b/dlblas/kernels/ks_competition/ascend/clike_sinkhorn.py new file mode 100644 index 000000000..034bf41d5 --- /dev/null +++ b/dlblas/kernels/ks_competition/ascend/clike_sinkhorn.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +import torch +from torch import nn + +from clike_910b import load_library + + +class ModelNew(nn.Module): + def __init__(self, repeat: int = 10, eps: float = 1e-6): + super().__init__() + self.repeat = repeat + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + load_library() + return torch.ops.dlblas_ks_ascendc.sinkhorn( + x.contiguous(), self.repeat, self.eps + ) + + +n0 = 1 +n1 = 1024 +mhc = 4 + + +def get_inputs(): + return [torch.randn(n0, n1, mhc, mhc)] + + +def get_init_inputs(): + return [] diff --git a/dlblas/kernels/ks_competition/ascend/clike_sparse_attention.py b/dlblas/kernels/ks_competition/ascend/clike_sparse_attention.py new file mode 100644 index 000000000..b17e368e7 --- /dev/null +++ b/dlblas/kernels/ks_competition/ascend/clike_sparse_attention.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +import torch +from torch import nn + +from clike_910b import load_library + + +class ModelNew(nn.Module): + def __init__(self, n_heads: int, head_dim: int): + super().__init__() + if head_dim != 128: + raise ValueError("The optimized AscendC kernel requires head_dim=128") + self.n_heads = n_heads + self.head_dim = head_dim + self.softmax_scale = head_dim**-0.5 + self.attn_sink = nn.Parameter(torch.zeros(n_heads, dtype=torch.float32)) + + def forward( + self, + q: torch.Tensor, + kv: torch.Tensor, + topk_idxs: torch.Tensor, + ) -> torch.Tensor: + if topk_idxs.shape[-1] != 16: + raise ValueError("The optimized AscendC kernel requires topk=16") + load_library() + indices = topk_idxs.contiguous() + valid = indices >= 0 + gathered = torch.ops.dlblas_ks_ascendc.sparse_gather(kv.contiguous(), indices) + q_cube = q.to(torch.float16) + gathered_cube = gathered.to(torch.float16) + scores = ( + torch.einsum("bmhd,bmtd->bmht", q_cube, gathered_cube).float() + * self.softmax_scale + ) + scores = scores.masked_fill(~valid.unsqueeze(2), float("-inf")) + sink = self.attn_sink.float().view(1, 1, self.n_heads, 1) + maximum = torch.maximum(scores.amax(-1, keepdim=True), sink) + exponent = torch.exp(scores - maximum).masked_fill(~valid.unsqueeze(2), 0) + weights = exponent / ( + exponent.sum(-1, keepdim=True) + torch.exp(sink - maximum) + ) + return torch.einsum( + "bmht,bmtd->bmhd", weights.to(torch.float16), gathered_cube + ).to(q.dtype) + + +batch_size = 8 +seq_len = 2600 +n_kv = 32 +n_heads = 64 +head_dim = 128 +topk = 16 + + +def get_inputs(): + q = torch.randn(batch_size, seq_len, n_heads, head_dim, dtype=torch.bfloat16) + kv = torch.randn(batch_size, n_kv, head_dim, dtype=torch.bfloat16) + topk_idxs = torch.randint(0, n_kv, (batch_size, seq_len, topk), dtype=torch.int32) + return [q, kv, topk_idxs] + + +def get_init_inputs(): + return [n_heads, head_dim] diff --git a/dlblas/kernels/ks_competition/torch/indexer.py b/dlblas/kernels/ks_competition/torch/indexer.py new file mode 100644 index 000000000..92c25e692 --- /dev/null +++ b/dlblas/kernels/ks_competition/torch/indexer.py @@ -0,0 +1,235 @@ +# fmt: off +import torch +from torch import nn +import torch.nn.functional as F +from dataclasses import dataclass +from typing import Tuple, Optional, Literal +from functools import lru_cache +from contextlib import contextmanager +import math + + +world_size = 1 +rank = 0 +block_size = 128 +fp4_block_size = 32 +default_dtype = torch.bfloat16 + + +@contextmanager +def set_dtype(dtype): + """Temporarily override torch default dtype, restoring it on exit (even if an exception occurs).""" + prev = torch.get_default_dtype() + torch.set_default_dtype(dtype) + try: + yield + finally: + torch.set_default_dtype(prev) + +@dataclass +class ModelArgs: + """Model hyperparameters. Field names match the config JSON keys.""" + max_batch_size: int = 4 + max_seq_len: int = 4096 + dtype: Literal["bf16", "fp8"] = "fp8" + scale_fmt: Literal[None, "ue8m0"] = "ue8m0" + expert_dtype: Literal[None, "fp4"] = None + scale_dtype: Literal["fp32", "fp8"] = "fp8" + vocab_size: int = 129280 + dim: int = 4096 + moe_inter_dim: int = 4096 + n_layers: int = 7 + n_hash_layers: int = 0 + n_mtp_layers: int = 1 + n_heads: int = 64 + # moe + n_routed_experts: int = 8 + n_shared_experts: int = 1 + n_activated_experts: int = 2 + score_func: Literal["softmax", "sigmoid", "sqrtsoftplus"] = "sqrtsoftplus" + route_scale: float = 1. + swiglu_limit: float = 0. + # mqa + q_lora_rank: int = 1024 + head_dim: int = 512 + rope_head_dim: int = 64 + norm_eps: float = 1e-6 + o_groups: int = 8 + o_lora_rank: int = 1024 + window_size: int = 128 + compress_ratios: Tuple[int] = (0, 0, 4, 128, 4, 128, 4, 0) + # yarn + compress_rope_theta: float = 40000.0 + original_seq_len: int = 0 + rope_theta: float = 10000.0 + rope_factor: float = 40 + beta_fast: int = 32 + beta_slow: int = 1 + # index + index_n_heads: int = 64 + index_head_dim: int = 128 + index_topk: int = 512 + # hc + hc_mult: int = 4 + hc_sinkhorn_iters: int = 20 + hc_eps: float = 1e-6 + +class Linear(nn.Module): + """Linear layer supporting BF16, FP8, and FP4 weight formats with per-block scaling.""" + + def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None): + super().__init__() + self.in_features = in_features + self.out_features = out_features + dtype = dtype or torch.bfloat16 + self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype)) + nn.init.kaiming_uniform_( + self.weight, + a=math.sqrt(5), + ) + self.register_parameter("scale", None) + if bias: + self.bias = nn.Parameter(torch.empty(out_features)) + bound = 1 / math.sqrt(in_features) + nn.init.uniform_( + self.bias, + -bound, + bound, + ) + else: + self.register_parameter("bias", None) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return linear(x, self.weight, self.bias) + +def linear(x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] = None) -> torch.Tensor: + """Dispatches to fp4_gemm / fp8_gemm / F.linear based on weight dtype. + For quantized weights, x is first quantized to FP8 via act_quant.""" + assert bias is None + return F.linear(x, weight) + +class ColumnParallelLinear(Linear): + """Shards output dim across TP ranks. No all-reduce needed on output.""" + def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None): + assert out_features % world_size == 0, f"Output features must be divisible by world size (world_size={world_size})" + self.part_out_features = out_features // world_size + super().__init__(in_features, self.part_out_features, bias, dtype) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return linear(x, self.weight, self.bias) + + +class RowParallelLinear(Linear): + """Shards input dim across TP ranks. All-reduce on output to sum partial results.""" + def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None): + assert in_features % world_size == 0, f"Input features must be divisible by world size (world_size={world_size})" + self.part_in_features = in_features // world_size + super().__init__(self.part_in_features, out_features, bias, dtype) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = linear(x, self.weight, None) + if self.bias is not None: + y += self.bias + return y.type_as(x) + +def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False) -> torch.Tensor: + """Applies rotary positional embeddings in-place. Uses conjugate for inverse (de-rotation).""" + y = x + x = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) + if inverse: + freqs_cis = freqs_cis.conj() + if x.ndim == 3: + freqs_cis = freqs_cis.view(1, x.size(1), x.size(-1)) + else: + freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1)) + x = torch.view_as_real(x * freqs_cis).flatten(-2) + y.copy_(x) + return y + + +class Model(torch.nn.Module): + """Selects top-k compressed KV positions for sparse attention via learned scoring. + Has its own Compressor (with Hadamard rotation) to build compressed KV for scoring.""" + + def __init__(self, args: ModelArgs, freqs_cis: torch.Tensor, kv_cache: torch.Tensor, compress_ratio: int = 4): + super().__init__() + self.dim = args.dim + self.n_heads = args.index_n_heads + self.n_local_heads = args.index_n_heads // world_size + self.head_dim = args.index_head_dim + self.rope_head_dim = args.rope_head_dim + self.index_topk = args.index_topk + self.q_lora_rank = args.q_lora_rank + self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.head_dim) + self.weights_proj = ColumnParallelLinear(self.dim, self.n_heads, dtype=torch.bfloat16) + self.softmax_scale = self.head_dim ** -0.5 + self.compress_ratio = compress_ratio + self.kv_cache = kv_cache + self.freqs_cis = freqs_cis + + def forward(self, x: torch.Tensor, qr: torch.Tensor, start_pos: int, offset: int): + bsz, seqlen, _ = x.size() + freqs_cis = self.freqs_cis[start_pos:start_pos+seqlen] + ratio = self.compress_ratio + rd = self.rope_head_dim + end_pos = start_pos + seqlen + q = self.wq_b(qr) + q = q.unflatten(-1, (self.n_local_heads, self.head_dim)) + apply_rotary_emb(q[..., -rd:], freqs_cis) + weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads ** -0.5) + index_score = torch.einsum("bshd,btd->bsht", q, self.kv_cache[:bsz, :end_pos // ratio]) + index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2) + if start_pos == 0: + mask = torch.arange(seqlen // ratio).npu().repeat(seqlen, 1) >= torch.arange(1, seqlen + 1).npu().unsqueeze(1) // ratio + index_score += torch.where(mask, float("-inf"), 0) + topk_idxs = index_score.topk(min(self.index_topk, end_pos // ratio), dim=-1)[1] + if start_pos == 0: + mask = topk_idxs >= torch.arange(1, seqlen + 1).npu().unsqueeze(1) // ratio + topk_idxs = torch.where(mask, -1, topk_idxs + offset) + else: + topk_idxs += offset + return topk_idxs + + +def _make_args(): + return ModelArgs( + max_batch_size=8, + max_seq_len=2600, + dim=1024, + index_n_heads=16, + index_head_dim=64, + index_topk=128, + q_lora_rank=256, + rope_head_dim=32 + ) + +def get_inputs(): + args = _make_args() + batch_size = 8 + seq_len = 2600 + x = torch.randn(batch_size, seq_len, args.dim, dtype=torch.bfloat16).npu() + qr = torch.randn(batch_size, seq_len, args.q_lora_rank, dtype=torch.bfloat16).npu() + start_pos = 0 + offset = 0 + return [x, qr, start_pos, offset] + + +def get_init_inputs(): + args = _make_args() + compress_ratio=4 + max_seq_len = args.max_seq_len + rope_theta = 10000.0 + # freqs_cis = torch.zeros(max_seq_len, args.rope_head_dim).npu() + freqs = 1.0 / (rope_theta ** (torch.arange(0, args.rope_head_dim, 2)[:args.rope_head_dim//2].float() / args.rope_head_dim)) + t = torch.arange(max_seq_len, dtype=torch.float32) + freqs = torch.outer(t, freqs).float().npu() + freqs_cis = torch.polar(torch.ones_like(freqs).npu(), freqs).view(max_seq_len, -1) + kv_cache = torch.randn(args.max_batch_size, args.max_seq_len // compress_ratio, args.index_head_dim, dtype=torch.bfloat16).npu() + return [args, freqs_cis, kv_cache, compress_ratio] + + +if __name__ == "__main__": + result = Model(*get_init_inputs()).npu().forward(*get_inputs()) + print(f"Forward pass successful! Output shape: {result.shape}, dtype: {result.dtype}") + print(f"Sample output values: {result[0, :5]}") +# fmt: on diff --git a/dlblas/kernels/ks_competition/torch/sinkhorn.py b/dlblas/kernels/ks_competition/torch/sinkhorn.py new file mode 100644 index 000000000..0a8761f4d --- /dev/null +++ b/dlblas/kernels/ks_competition/torch/sinkhorn.py @@ -0,0 +1,37 @@ +# fmt: off +import torch +import torch.nn as nn +class Model(nn.Module): + """ + Pure PyTorch implementation of sinkhorn_normalize. + Iteratively normalizes a matrix to be doubly stochastic: + 1. softmax(x, dim=-1) + eps + 2. column-normalize: x / (x.sum(-2) + eps) + 3. repeat (row-normalize then column-normalize) for repeat-1 iterations + """ + def __init__(self, repeat: int = 10, eps: float = 1e-6): + super().__init__() + self.repeat = repeat + self.eps = eps + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Args: + x: [..., mhc, mhc] float32 + Returns: + [..., mhc, mhc] float32 (doubly stochastic) + """ + x = x.softmax(-1) + self.eps + x = x / (x.sum(-2, keepdim=True) + self.eps) + for _ in range(self.repeat - 1): + x = x / (x.sum(-1, keepdim=True) + self.eps) + x = x / (x.sum(-2, keepdim=True) + self.eps) + return x +n0 = 1 +n1 = 1024 +mhc = 4 +def get_inputs(): + x = torch.randn(n0, n1, mhc, mhc) + return [x] +def get_init_inputs(): + return [] +# fmt: on diff --git a/dlblas/kernels/ks_competition/torch/sparse_attention.py b/dlblas/kernels/ks_competition/torch/sparse_attention.py new file mode 100644 index 000000000..65951924e --- /dev/null +++ b/dlblas/kernels/ks_competition/torch/sparse_attention.py @@ -0,0 +1,126 @@ +# fmt: off +import torch +import torch.nn as nn + + +def sparse_attn_ref( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + """ + Pure PyTorch equivalent of sparse_attn_kernel (kernel.py). + + TileLang uses block-wise online softmax (FlashAttention style). + This implementation is mathematically equivalent using a single-pass + masked softmax with an attention sink contributing only to the denominator. + + Args: + q: [b, m, h, d] bfloat16 + kv: [b, n, d] bfloat16 (shared key-value per position) + attn_sink: [h] float32 (learnable sink; only in denominator) + topk_idxs: [b, m, topk] int32 (-1 = invalid / padding) + softmax_scale: float (typically head_dim ** -0.5) + + Returns: + o: [b, m, h, d] bfloat16 + """ + b, m, h, d = q.shape + topk = topk_idxs.shape[-1] + + valid_mask = topk_idxs >= 0 # [b, m, topk] + safe_idxs = topk_idxs.clamp(min=0).long() # replace -1 with 0 for safe gather + + # Gather KV: [b, m, topk, d] + b_idx = torch.arange(b, device=q.device)[:, None, None].expand(b, m, topk) + gathered_kv = kv[b_idx, safe_idxs] # [b, m, topk, d] + # Zero out positions that came from invalid (-1) indices + gathered_kv = gathered_kv.masked_fill(~valid_mask.unsqueeze(-1), 0.0) + + # Attention scores: [b, m, h, topk] + scores = torch.einsum("bmhd,bmtd->bmht", + q.float(), gathered_kv.float()) * softmax_scale + # Mask invalid positions to -inf so they don't affect softmax + scores = scores.masked_fill(~valid_mask.unsqueeze(2), float("-inf")) + + # Numerically stable softmax with attn_sink only in the denominator. + # Equivalent to the TileLang kernel line: + # sum_exp[i] += T.exp(attn_sink[i] - scores_max[i]) + sink = attn_sink.float().view(1, 1, h, 1) # broadcast over b, m, topk + + max_scores = torch.amax(scores, dim=-1, keepdim=True) # [b, m, h, 1] + # When all topk positions are invalid, max_scores = -inf; clamp with sink to stay finite + max_scores = torch.maximum(max_scores, sink) + + exp_scores = torch.exp(scores - max_scores) + # Re-zero invalid positions (exp(-inf - finite) = 0, but be explicit) + exp_scores = exp_scores.masked_fill(~valid_mask.unsqueeze(2), 0.0) + + exp_sink = torch.exp(sink - max_scores) # [b, m, h, 1] + sum_exp = exp_scores.sum(dim=-1, keepdim=True) + exp_sink + + attn_weights = exp_scores / sum_exp # [b, m, h, topk] + + # Weighted sum of gathered KV + output = torch.einsum("bmht,bmtd->bmhd", + attn_weights, gathered_kv.float()) # [b, m, h, d] + return output.to(q.dtype) + + +class Model(nn.Module): + """ + Pure PyTorch implementation of sparse_attn from + DeepSeek-V4-Pro/inference/kernel.py. + + Wraps sparse_attn_ref with a learnable attn_sink parameter. + """ + + def __init__(self, n_heads: int, head_dim: int): + super().__init__() + self.n_heads = n_heads + self.head_dim = head_dim + self.softmax_scale = head_dim ** -0.5 + # Learnable per-head sink bias (only affects softmax denominator) + self.attn_sink = nn.Parameter(torch.zeros(n_heads, dtype=torch.float32)) + + def forward( + self, + q: torch.Tensor, + kv: torch.Tensor, + topk_idxs: torch.Tensor, + ) -> torch.Tensor: + """ + Args: + q: [b, m, h, d] bfloat16 + kv: [b, n, d] bfloat16 + topk_idxs: [b, m, topk] int32, -1 for invalid positions + + Returns: + o: [b, m, h, d] bfloat16 + """ + return sparse_attn_ref(q, kv, self.attn_sink, topk_idxs, self.softmax_scale) + + +# --------------------------------------------------------------------------- +# Default config for get_inputs / get_init_inputs +# --------------------------------------------------------------------------- +batch_size = 8 +seq_len = 2600 +n_kv = 32 +n_heads = 64 +head_dim = 128 +topk = 16 + + +def get_inputs(): + q = torch.randn(batch_size, seq_len, n_heads, head_dim, dtype=torch.bfloat16) + kv = torch.randn(batch_size, n_kv, head_dim, dtype=torch.bfloat16) + topk_idxs = torch.randint(0, n_kv, (batch_size, seq_len, topk), dtype=torch.int32) + return [q, kv, topk_idxs] + + +def get_init_inputs(): + return [n_heads, head_dim] +# fmt: on