From 61d834a00ef31e925f93721de8703a3741bbf8da Mon Sep 17 00:00:00 2001 From: xiao Date: Wed, 24 Jun 2026 14:06:41 +0800 Subject: [PATCH 1/3] =?UTF-8?q?Add=20=E5=BC=80=E6=BA=90=E4=BB=93=E5=BA=93?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../FlashAttention关键算子迁移与优化.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashAttention关键算子迁移与优化.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashAttention关键算子迁移与优化.md index 1aa5524..81ba709 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashAttention关键算子迁移与优化.md +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashAttention关键算子迁移与优化.md @@ -255,6 +255,11 @@ * **Tensor Core:** 现代 GPU(含沐曦 C500)的矩阵乘法专用单元,要求维度对齐为 8 或 16 的倍数。 +### 5.5 开源仓库参考 + +[GitHub - MetaX-MACA/flashattn · GitHub](https://github.com/MetaX-MACA/flashattn) + +> 链接内容可供用于学习 API、算子实现思路、benchmark 方法和优化策略。选手仍需根据 XPU-OJ 题包接口**自行实现**可提交的 `run_kernel(...)` --- From 3adda68ccc50665df5acd21ee5ee78de3c577acd Mon Sep 17 00:00:00 2001 From: Dayuxiaoshui <792179245@qq.com> Date: Wed, 24 Jun 2026 16:06:08 +0800 Subject: [PATCH 2/3] ADD file via upload --- .../赛题一mla教程.md | 361 ++++++++++++++++++ 1 file changed, 361 insertions(+) create mode 100644 基于国产软件栈大模型推理前沿算子优化/赛题一mla教程.md diff --git a/基于国产软件栈大模型推理前沿算子优化/赛题一mla教程.md b/基于国产软件栈大模型推理前沿算子优化/赛题一mla教程.md new file mode 100644 index 0000000..31dc386 --- /dev/null +++ b/基于国产软件栈大模型推理前沿算子优化/赛题一mla教程.md @@ -0,0 +1,361 @@ +# DeepSeek MLA Decode 提交说明 + +## 当前结果 + +| Status | Score | Time | Memory | Platform | +| --- | ---: | ---: | ---: | --- | +| Accepted | 49.5 | 23 ms | 22.2 G | TileLang Maca C500 / 10.1 K | + +得分说明： + +- 50 分左右基本对应和题目 baseline 的加速比约为 `1:1`。 +- 当前 `49.5` 可以理解为 baseline 档位附近的 Accepted 结果。 +- 该结果主要用于确认提交接口、TileLang kernel 调用和输出正确性已经跑通。 +- 本文档只记录提交模板、改算子位置和当前结果，不展开进一步算子优化。 + +## 提交文件 + +提交文件为： + +```text +race_tests/mla/submission.py +``` + +评测只要求 Python 文件中暴露 `run_kernel` 函数，函数名、参数顺序必须和题目一致： + +```python +def run_kernel( + q, + q_pe, + kv, + k_pe, + output, + batch, + heads, + kv_heads, + kv_ctx, + dim, + pe_dim, +): + ... +``` + +提交时使用 `submission.py` 的内容即可，不需要提交 benchmark、reference 或测试脚本。 + +## 算子来源 + +当前提交模板基于 `race_tests/mla/test_tilelang_mla.py` 中的 `flashattn` 算子封装而来。 + +题目计算语义为： + +```text +score = (q @ kv^T + q_pe @ k_pe^T) / sqrt(576) +attention = softmax(score, dim=-1) +output = attention @ kv +``` + +固定约束： + +```text +dim = 512 +pe_dim = 64 +kv_heads = 1 +heads = 16 +``` + +## 模板结构 + +`submission.py` 里主要有三部分： + +```text +flashattn(...) TileLang MLA kernel +_get_kernel(...) 按 shape 缓存编译后的 kernel +run_kernel(...) OJ 调用入口，写入 output +``` + +`run_kernel` 不做同步，不分配最终输出，只负责取得缓存 kernel 并调用： + +```python +kernel = _get_kernel(...) +kernel(q, q_pe, kv, k_pe, output) +``` + +## 改算子的位置 + +只改 MLA 算子时，主要看两个位置： + +```text +race_tests/mla/submission.py +``` + +1. `flashattn(...)` + - TileLang kernel 主体。 + - QK、QK_pe、online softmax、PV 都在这里。 + +2. `_get_kernel(...)` + - 设置 `block_n`、`block_h`、`num_split`。 + - 控制不同 shape 的 kernel 缓存 key。 + +一般不要改 `run_kernel` 的函数签名；评测器按固定签名调用。 + +## 完整算子代码 + +下面是当前 `race_tests/mla/submission.py` 的完整提交代码： + +```python +import tilelang +import tilelang.language as T + + +@tilelang.jit(pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True}) +def flashattn(batch, heads, kv_head_num, seqlen_kv, dim, pe_dim, block_N, block_H, num_split, softmax_scale): + scale = float(softmax_scale * 1.44269504) + dtype = T.float16 + accum_dtype = T.float32 + kv_group_num = heads // kv_head_num + valid_block_h = min(block_H, kv_group_num) + assert kv_head_num == 1, "kv_head_num must be 1" + + @T.prim_func + def main_split( + Q: T.Tensor([batch, heads, dim], dtype), + Q_pe: T.Tensor([batch, heads, pe_dim], dtype), + KV: T.Tensor([batch, seqlen_kv, kv_head_num, dim], dtype), + K_pe: T.Tensor([batch, seqlen_kv, kv_head_num, pe_dim], dtype), + Output: T.Tensor([batch, heads, dim], dtype), + ): + glse = T.alloc_global([batch, heads, num_split], dtype) + output_partial = T.alloc_global([batch, heads, num_split, dim], dtype) + + with T.Kernel(batch, heads // min(block_H, kv_group_num), num_split, threads=256) as (bid, hid, bz): + Q_shared = T.alloc_shared([block_H, dim], dtype) + S_shared = T.alloc_shared([block_H, block_N], dtype) + Q_pe_shared = T.alloc_shared([block_H, pe_dim], dtype) + KV_shared = T.alloc_shared([block_N, dim], dtype) + K_pe_shared = T.alloc_shared([block_N, pe_dim], dtype) + O_shared = T.alloc_shared([block_H, dim], dtype) + acc_s = T.alloc_fragment([block_H, block_N], accum_dtype) + acc_s_cast = T.alloc_fragment([block_H, block_N], dtype) + acc_o = T.alloc_fragment([block_H, dim], accum_dtype) + scores_max = T.alloc_fragment([block_H], accum_dtype) + scores_max_prev = T.alloc_fragment([block_H], accum_dtype) + scores_scale = T.alloc_fragment([block_H], accum_dtype) + scores_sum = T.alloc_fragment([block_H], accum_dtype) + logsum = T.alloc_fragment([block_H], accum_dtype) + cur_kv_head = hid // (kv_group_num // block_H) + + T.use_swizzle(10) + T.copy(Q[bid, hid * valid_block_h : (hid + 1) * valid_block_h, :], Q_shared) + T.copy(Q_pe[bid, hid * valid_block_h : (hid + 1) * valid_block_h, :], Q_pe_shared) + T.fill(acc_o, 0) + T.fill(logsum, 0) + T.fill(scores_max, -T.infinity(accum_dtype)) + + loop_range = T.ceildiv(seqlen_kv // num_split, block_N) + for k in T.Pipelined(loop_range, num_stages=2): + kv_start = (seqlen_kv // num_split) * bz + k * block_N + kv_end = (seqlen_kv // num_split) * bz + (k + 1) * block_N + T.copy(KV[bid, kv_start:kv_end, cur_kv_head, :], KV_shared) + T.copy(K_pe[bid, kv_start:kv_end, cur_kv_head, :], K_pe_shared) + + T.clear(acc_s) + T.gemm(Q_shared, KV_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullCol) + T.gemm(Q_pe_shared, K_pe_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullCol) + + T.copy(scores_max, scores_max_prev) + T.fill(scores_max, -T.infinity(accum_dtype)) + T.reduce_max(acc_s, scores_max, dim=1, clear=False) + for i in T.Parallel(block_H): + scores_max[i] = T.max(scores_max[i], scores_max_prev[i]) + for i in T.Parallel(block_H): + scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale) + for i, j in T.Parallel(block_H, block_N): + acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale) + + T.reduce_sum(acc_s, scores_sum, dim=1) + T.copy(acc_s, S_shared) + T.copy(S_shared, acc_s_cast) + for i in T.Parallel(block_H): + logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i] + for i, j in T.Parallel(block_H, dim): + acc_o[i, j] *= scores_scale[i] + T.gemm(acc_s_cast, KV_shared, acc_o, policy=T.GemmWarpPolicy.FullCol) + + for i, j in T.Parallel(block_H, dim): + acc_o[i, j] /= logsum[i] + for i in T.Parallel(block_H): + logsum[i] = T.log2(logsum[i]) + scores_max[i] * scale + + T.copy(logsum, glse[bid, hid * valid_block_h : (hid + 1) * valid_block_h, bz]) + T.copy(acc_o, O_shared) + T.copy(O_shared, output_partial[bid, hid * valid_block_h : (hid + 1) * valid_block_h, bz, :]) + + with T.Kernel(heads, batch, threads=128) as (hid, bz): + po_local = T.alloc_fragment([dim], dtype) + o_accum_local = T.alloc_fragment([dim], accum_dtype) + lse_local_split = T.alloc_var(accum_dtype) + lse_logsum_local = T.alloc_var(accum_dtype) + lse_max_local = T.alloc_var(accum_dtype) + scale_local = T.alloc_var(accum_dtype) + + T.clear(lse_logsum_local) + T.clear(o_accum_local) + lse_max_local = -T.infinity(accum_dtype) + + for k in T.serial(num_split): + lse_max_local = T.max(lse_max_local, glse[bz, hid, k]) + for k in T.Pipelined(num_split, num_stages=1): + lse_local_split = glse[bz, hid, k] + lse_logsum_local += T.exp2(lse_local_split - lse_max_local) + lse_logsum_local = T.log2(lse_logsum_local) + lse_max_local + + for k in T.serial(num_split): + for i in T.Parallel(dim): + po_local[i] = output_partial[bz, hid, k, i] + lse_local_split = glse[bz, hid, k] + scale_local = T.exp2(lse_local_split - lse_logsum_local) + for i in T.Parallel(dim): + o_accum_local[i] += po_local[i] * scale_local + + for i in T.Parallel(dim): + Output[bz, hid, i] = o_accum_local[i] + + @T.prim_func + def main_no_split( + Q: T.Tensor([batch, heads, dim], dtype), + Q_pe: T.Tensor([batch, heads, pe_dim], dtype), + KV: T.Tensor([batch, seqlen_kv, kv_head_num, dim], dtype), + K_pe: T.Tensor([batch, seqlen_kv, kv_head_num, pe_dim], dtype), + Output: T.Tensor([batch, heads, dim], dtype), + ): + with T.Kernel(heads // min(block_H, kv_group_num), batch, threads=128) as (hid, bid): + Q_shared = T.alloc_shared([block_H, dim], dtype) + S_shared = T.alloc_shared([block_H, block_N], dtype) + Q_pe_shared = T.alloc_shared([block_H, pe_dim], dtype) + KV_shared = T.alloc_shared([block_N, dim], dtype) + K_pe_shared = T.alloc_shared([block_N, pe_dim], dtype) + O_shared = T.alloc_shared([block_H, dim], dtype) + acc_s = T.alloc_fragment([block_H, block_N], accum_dtype) + acc_o = T.alloc_fragment([block_H, dim], accum_dtype) + scores_max = T.alloc_fragment([block_H], accum_dtype) + scores_max_prev = T.alloc_fragment([block_H], accum_dtype) + scores_scale = T.alloc_fragment([block_H], accum_dtype) + scores_sum = T.alloc_fragment([block_H], accum_dtype) + logsum = T.alloc_fragment([block_H], accum_dtype) + cur_kv_head = hid // (kv_group_num // block_H) + + T.copy(Q[bid, hid * valid_block_h : (hid + 1) * valid_block_h, :], Q_shared) + T.copy(Q_pe[bid, hid * valid_block_h : (hid + 1) * valid_block_h, :], Q_pe_shared) + T.fill(acc_o, 0) + T.fill(logsum, 0) + T.fill(scores_max, -T.infinity(accum_dtype)) + + loop_range = T.ceildiv(seqlen_kv, block_N) + for k in T.Pipelined(loop_range, num_stages=0): + T.copy(KV[bid, k * block_N : (k + 1) * block_N, cur_kv_head, :], KV_shared) + T.copy(K_pe[bid, k * block_N : (k + 1) * block_N, cur_kv_head, :], K_pe_shared) + T.gemm( + Q_shared, + KV_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullCol, + clear_accum=True, + ) + T.gemm(Q_pe_shared, K_pe_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullCol) + + T.copy(scores_max, scores_max_prev) + T.fill(scores_max, -T.infinity(accum_dtype)) + T.reduce_max(acc_s, scores_max, dim=1, clear=False) + for i in T.Parallel(block_H): + scores_max[i] = T.max(scores_max[i], scores_max_prev[i]) + for i in T.Parallel(block_H): + scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale) + for i, j in T.Parallel(block_H, block_N): + acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale) + + T.reduce_sum(acc_s, scores_sum, dim=1) + T.copy(acc_s, S_shared) + for i in T.Parallel(block_H): + logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i] + for i, j in T.Parallel(block_H, dim): + acc_o[i, j] *= scores_scale[i] + T.gemm(S_shared, KV_shared, acc_o, policy=T.GemmWarpPolicy.FullCol) + + for i, j in T.Parallel(block_H, dim): + acc_o[i, j] /= logsum[i] + T.copy(acc_o, O_shared) + T.copy(O_shared, Output[bid, hid * valid_block_h : (hid + 1) * valid_block_h, :]) + + if num_split > 1: + return main_split + return main_no_split + + +_KERNEL_CACHE = {} + + +def _get_kernel(batch, heads, kv_heads, kv_ctx, dim, pe_dim): + block_n = 32 + block_h = min(16, heads // kv_heads) + num_split = 1 + softmax_scale = (dim + pe_dim) ** -0.5 + key = (batch, heads, kv_heads, kv_ctx, dim, pe_dim, block_n, block_h, num_split) + kernel = _KERNEL_CACHE.get(key) + if kernel is None: + kernel = flashattn(batch, heads, kv_heads, kv_ctx, dim, pe_dim, block_n, block_h, num_split, softmax_scale) + _KERNEL_CACHE[key] = kernel + return kernel + + +def run_kernel( + q, + q_pe, + kv, + k_pe, + output, + batch, + heads, + kv_heads, + kv_ctx, + dim, + pe_dim, +): + kernel = _get_kernel(int(batch), int(heads), int(kv_heads), int(kv_ctx), int(dim), int(pe_dim)) + kernel(q, q_pe, kv, k_pe, output) +``` + +## 本地验证 + +本地验证时需要使用已经编译好的 TileLang，并把仓库根目录和 TVM Python 路径加入 `PYTHONPATH`。 + +在仓库根目录执行： + +```bash +TILELANG_CACHE_DIR=/tmp/tilelang-cache \ +MACA_PATH=${MACA_PATH:-/opt/maca} \ +LD_LIBRARY_PATH="$(pwd)/build/lib:${MACA_PATH:-/opt/maca}/lib:${MACA_PATH:-/opt/maca}/mxgpu_llvm/lib:${LD_LIBRARY_PATH}" \ +PATH="${MACA_PATH:-/opt/maca}/bin:${MACA_PATH:-/opt/maca}/mxgpu_llvm/bin:${PATH}" \ +PYTHONPATH="$(pwd):$(pwd)/3rdparty/tvm/python:$(pwd)/race_tests/mla:${PYTHONPATH}" \ +python race_tests/mla/test_tilelang_mla.py \ + --no-json \ + --batch 1 \ + --heads 16 \ + --kv_heads 1 \ + --kv_ctx 2048 \ + --dim 512 \ + --pe_dim 64 +``` + +已用 `submission.run_kernel` 验证过： + +```text +kv_ctx=2048 allclose True +kv_ctx=8192 allclose True +``` + +## 注意事项 + +- `run_kernel` 内不要调用 `torch.cuda.synchronize()`。 +- `output` 是评测器传入的缓冲区，必须原地写入。 +- 当前文档只说明提交模板和改算子入口，不涉及进一步算子优化。 From eb9815e03d89ed4cc9a640c49c0c9b7eb9f5f8fa Mon Sep 17 00:00:00 2001 From: Dayuxiaoshui <792179245@qq.com> Date: Wed, 24 Jun 2026 16:06:35 +0800 Subject: [PATCH 3/3] ADD file via upload --- .../赛题一nsa教程.md | 283 ++++++++++++++++++ 1 file changed, 283 insertions(+) create mode 100644 基于国产软件栈大模型推理前沿算子优化/赛题一nsa教程.md diff --git a/基于国产软件栈大模型推理前沿算子优化/赛题一nsa教程.md b/基于国产软件栈大模型推理前沿算子优化/赛题一nsa教程.md new file mode 100644 index 0000000..53a6733 --- /dev/null +++ b/基于国产软件栈大模型推理前沿算子优化/赛题一nsa教程.md @@ -0,0 +1,283 @@ +# Native Sparse Attention 提交说明 + +## 当前结果 + +| Status | Score | Case | Time | Memory | Platform | Submit Time | +| --- | ---: | --- | ---: | ---: | --- | --- | +| Accepted | 53.64 | muxitest001 | 619 us | 22.2 G | TileLang Maca C500 / 4.9 K | 06/18 15:53:16 | + +得分说明： + +- 50 分左右基本对应和题目 baseline 的加速比约为 `1:1`。 +- 当前 `53.64` 是 baseline 档位附近、略高于 50 分线的 Accepted 结果。 +- 该结果主要用于确认提交接口、TileLang kernel 调用和输出正确性已经跑通。 +- 本文档只记录提交模板、改算子位置和当前结果，不展开进一步算子优化。 + +## 提交文件 + +提交文件为： + +```text +race_tests/nsa/submission.py +``` + +评测只要求 Python 文件中暴露 `run_kernel` 函数，函数名、参数顺序必须和题目一致： + +```python +def run_kernel( + q, + k, + v, + block_indices, + output, + B, + seq_len, + H, + HQ, + D, + S, + block_size, + is_causal, +): + ... +``` + +提交时使用 `submission.py` 的内容即可，不需要提交 benchmark、reference 或测试脚本。 + +## 算子来源 + +当前提交模板基于 `race_tests/nsa/test_tilelang_nsa_fwd.py` 中的 `native_sparse_attention` 算子封装而来。 + +题目计算语义为： + +```text +score = q @ k_selected^T / sqrt(D) +attention = softmax(score) +output = attention @ v_selected +``` + +其中 `block_indices` 指定每个 query token 选中的 KV block，`is_causal=1` 时需要屏蔽未来 token。 + +## 模板结构 + +`submission.py` 里主要有三部分： + +```text +native_sparse_attention(...) TileLang NSA kernel +_get_kernel(...) 按 shape 缓存编译后的 kernel +run_kernel(...) OJ 调用入口，写入 output +``` + +`run_kernel` 不做同步，不分配最终输出，只负责取得缓存 kernel 并调用： + +```python +kernel = _get_kernel(...) +kernel(q, k, v, block_indices, output) +``` + +## 改算子的位置 + +只改 NSA 算子时，主要看两个位置： + +```text +race_tests/nsa/submission.py +``` + +1. `native_sparse_attention(...)` + - TileLang kernel 主体。 + - block 读取、causal mask、online softmax、PV 都在这里。 + +2. `_get_kernel(...)` + - 设置 shape cache key。 + - 控制不同 `(B, seq_len, H, HQ, D, S, block_size, is_causal)` 的 kernel 缓存。 + +一般不要改 `run_kernel` 的函数签名；评测器按固定签名调用。 + +## 完整算子代码 + +下面是当前 `race_tests/nsa/submission.py` 的完整提交代码： + +```python +import tilelang +import tilelang.language as T + + +@tilelang.jit( + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + }, +) +def native_sparse_attention(batch, heads, seq_len, dim, is_causal, block_size, groups, selected_blocks): + scale = float((dim**-0.5) * 1.44269504) + head_kv = heads // groups + q_shape = [batch, seq_len, heads, dim] + kv_shape = [batch, seq_len, head_kv, dim] + block_indices_shape = [batch, seq_len, head_kv, selected_blocks] + + dtype = T.float16 + accum_dtype = T.float32 + block_t = min(128, tilelang.math.next_power_of_2(dim)) + + assert tilelang.cdiv(dim, block_t) == 1, "The key dimension can not be larger than 128" + + S = selected_blocks + G = groups + BS = block_size + BK = BV = block_t + num_stages = 2 + threads = 64 + + @T.prim_func + def kernel( + Q: T.Tensor(q_shape, dtype), + K: T.Tensor(kv_shape, dtype), + V: T.Tensor(kv_shape, dtype), + BlockIndices: T.Tensor(block_indices_shape, T.int32), + Output: T.Tensor(q_shape, dtype), + ): + with T.Kernel(seq_len, tilelang.cdiv(dim, BV), batch * head_kv, threads=threads) as (bx, by, bz): + Q_shared = T.alloc_shared([G, BK], dtype) + K_shared = T.alloc_shared([BS, BK], dtype) + V_shared = T.alloc_shared([BS, BV], dtype) + O_shared = T.alloc_shared([G, BV], dtype) + + acc_s = T.alloc_fragment([G, BS], accum_dtype) + acc_s_cast = T.alloc_fragment([G, BS], dtype) + acc_o = T.alloc_fragment([G, BV], accum_dtype) + scores_max = T.alloc_fragment([G], accum_dtype) + scores_max_prev = T.alloc_fragment([G], accum_dtype) + scores_scale = T.alloc_fragment([G], accum_dtype) + scores_sum = T.alloc_fragment([G], accum_dtype) + logsum = T.alloc_fragment([G], accum_dtype) + + i_t = bx + i_v = by + i_bh = bz + i_b = i_bh // head_kv + i_h = i_bh % head_kv + + T.copy(Q[i_b, i_t, i_h * G : (i_h + 1) * G, :], Q_shared) + T.fill(acc_o, 0) + T.fill(logsum, 0) + T.fill(scores_max, -T.infinity(accum_dtype)) + + for s in T.Pipelined(S, num_stages=num_stages): + i_s = BlockIndices[i_b, i_t, i_h, s] * BS + if i_s <= i_t and i_s >= 0: + T.copy(K[i_b, i_s : i_s + BS, i_h, :], K_shared) + + if is_causal: + for i, j in T.Parallel(G, BS): + acc_s[i, j] = T.if_then_else(i_t >= i_s + j, 0, -T.infinity(acc_s.dtype)) + else: + T.clear(acc_s) + + T.gemm(Q_shared, K_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) + + T.copy(scores_max, scores_max_prev) + T.fill(scores_max, -T.infinity(accum_dtype)) + T.reduce_max(acc_s, scores_max, dim=1, clear=True) + for i in T.Parallel(G): + scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale) + for i, j in T.Parallel(G, BS): + acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale) + + T.reduce_sum(acc_s, scores_sum, dim=1) + for i in T.Parallel(G): + logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i] + T.copy(acc_s, acc_s_cast) + + for i, j in T.Parallel(G, BV): + acc_o[i, j] *= scores_scale[i] + + T.copy(V[i_b, i_s : i_s + BS, i_h, i_v * BV : (i_v + 1) * BV], V_shared) + T.gemm(acc_s_cast, V_shared, acc_o, policy=T.GemmWarpPolicy.FullRow) + + for i, j in T.Parallel(G, BV): + acc_o[i, j] /= logsum[i] + T.copy(acc_o, O_shared) + T.copy(O_shared, Output[i_b, i_t, i_h * G : (i_h + 1) * G, i_v * BV : (i_v + 1) * BV]) + + return kernel + + +_KERNEL_CACHE = {} + + +def _get_kernel(B, seq_len, H, HQ, D, S, block_size, is_causal): + groups = HQ // H + key = (B, seq_len, H, HQ, D, S, block_size, int(is_causal)) + kernel = _KERNEL_CACHE.get(key) + if kernel is None: + kernel = native_sparse_attention( + batch=B, + heads=HQ, + seq_len=seq_len, + dim=D, + is_causal=bool(is_causal), + block_size=block_size, + groups=groups, + selected_blocks=S, + ) + _KERNEL_CACHE[key] = kernel + return kernel + + +def run_kernel( + q, + k, + v, + block_indices, + output, + B, + seq_len, + H, + HQ, + D, + S, + block_size, + is_causal, +): + kernel = _get_kernel( + int(B), + int(seq_len), + int(H), + int(HQ), + int(D), + int(S), + int(block_size), + int(is_causal), + ) + kernel(q, k, v, block_indices, output) +``` + +## 本地验证 + +本地验证时需要使用已经编译好的 TileLang，并把仓库根目录和 TVM Python 路径加入 `PYTHONPATH`。 + +在仓库根目录执行： + +```bash +TILELANG_CACHE_DIR=/tmp/tilelang-cache \ +MACA_PATH=${MACA_PATH:-/opt/maca} \ +LD_LIBRARY_PATH="$(pwd)/build/lib:${MACA_PATH:-/opt/maca}/lib:${MACA_PATH:-/opt/maca}/mxgpu_llvm/lib:${LD_LIBRARY_PATH}" \ +PATH="${MACA_PATH:-/opt/maca}/bin:${MACA_PATH:-/opt/maca}/mxgpu_llvm/bin:${PATH}" \ +PYTHONPATH="$(pwd):$(pwd)/3rdparty/tvm/python:$(pwd)/race_tests/nsa:${PYTHONPATH}" \ +python race_tests/nsa/test_tilelang_nsa_fwd.py +``` + +已用 `submission.run_kernel` 验证过： + +```text +(B=1, seq_len=64, H=1, HQ=16, D=32, S=1, block_size=16) allclose True +(B=1, seq_len=128, H=2, HQ=32, D=64, S=4, block_size=16) allclose True +(B=1, seq_len=128, H=1, HQ=16, D=128, S=2, block_size=32) allclose True +``` + +## 注意事项 + +- `run_kernel` 内不要调用 `torch.cuda.synchronize()`。 +- `output` 是评测器传入的缓冲区，必须原地写入。 +- `block_indices` 中无效 block 使用 `seq_len` 作为哨兵值；当前 kernel 通过 `i_s <= i_t` 和 `i_s >= 0` 跳过无效或未来 block。 +- 当前文档只说明提交模板和改算子入口，不涉及进一步算子优化。