NOSA-1B 源码解析（三）：Triton 内核如何将注意力计算加速到极致？

NOSA-1B 源码解析三Triton 内核如何将注意力计算加速到极致【免费下载链接】NOSA-1B项目地址: https://ai.gitcode.com/OpenBMB/NOSA-1BOpenBMB 开源的 NOSA-1B 是一个主打 KV-cache 卸载与原生稀疏注意力Native and Offloadable Sparse Attention的大模型它能把长上下文生成时的解码吞吐提升最高 5 倍以上。本文是 NOSA-1B 源码解析系列第三篇聚焦藏在cis_pooling.py里的Triton 内核一步步拆解它是如何用 GPU 并行计算把注意力打分pooling这一步加速到极致的非常适合想学习 Triton 内核优化技巧的新手开发者。为什么注意力计算会成为卡脖子的瓶颈大模型每生成一个 token都要拿着当前的 Query 去和历史上所有的 Key/Value 做一次点积注意力。序列越长参与计算的 KV 越多计算量呈平方级增长这就是长文本生成慢的根本原因。NOSA-1B 的思路很聪明先粗筛、再精算。它用一块很小的压缩 KV快速算出每个块的重要性分数CIS即压缩重要性分数选出 top-k 个关键块然后只对这少数几个块做完整的稀疏注意力。这样一来绝大部分计算都被跳过了。但粗筛这一步本身也得够快否则省下的时间又会被填回去——这正是 Triton 内核登场的地方。Triton 内核藏在哪里两条源码调用链我们先定位源码。整个加速路径涉及两个关键文件cis_pooling.pyTriton 内核的唯一实现文件包含nosa_mean_pool_kernel内核和封装函数nosa_mean_poolingmodeling_llama_long_infllmv2.py在compressed_attention函数里真正调用它第 796 行score_cis nosa_mean_pooling(...)在modeling_minicpm.py中同样实现了compressed_attention稀疏注意力流程两套实现共用同一套稀疏思想。调用链大致是MiniCPMInfLLMv2Attention/ Llama 稀疏注意力层进入compressed_attention先用infllmv2_attn_stage1得到初步的块级注意力分数再用nosa_mean_poolingTriton 内核对 CIS 分数做滑窗均值池化得到每个块的综合重要性最后topk选出关键块交给infllmv2_attn_varlen_func做真正的稀疏注意力Triton 内核加速的三大核心机制机制一三维网格并行一个线程块负责一个窗口Triton 内核的并行模型是网格 程序你定义一个逻辑网格gridGPU 会同时派出成千上万个程序program去执行同一个内核。nosa_mean_pool_kernel的网格设计非常巧妙grid (H, B, M) # (注意力头数, batch 数, 滑窗个数)内核开头用tl.program_id(0/1/2)取出当前程序负责的头部、样本和窗口编号然后只处理属于自己的那一段数据。也就是说每个头、每个样本、每个滑窗之间的池化计算完全互不干扰天然并行。序列越长、头数越多并行度越高加速比越明显。机制二块式加载block load 向量化归约传统 PyTorch 写法比如avg_pool1d会经过多次内核启动和中间张量落盘开销很大。而 Triton 内核把数据一次性地加载进寄存器/SRAMblock_idx tl.arange(0, kernel_size) # 0~31 的窗口内偏移 block_cis_ptrs beg_pos block_idx * cis_stride_n block_scores tl.load(block_cis_ptrs, maskmask, other0.0)tl.load一次性读入 32 个分数kernel_size32随后tl.sum(block_scores, axis0) / val_cnt在寄存器里完成求和与平均。整个过程零中间张量、零 Python 循环内存访问完全贴合 GPU 的带宽特性。机制三mask 处理变长序列避免 padding 浪费真实推理时每个样本长度不一样cu_seqlens记录每个样本的起止位置。内核用mask (block_idx tidx_m * stride) (batch_end - batch_start)判断哪些位置有效并用tl.sum(mask.to(tl.int32), axis0)统计有效元素个数作为分母。这样短序列不会因为补齐而浪费算力长序列也不会算错属于工程上很实用的细节。Triton 与 PyTorch 的期末考一致性校验作者在cis_pooling.py的main()里写了一段自检代码用torch.nn.functional.avg_pool1d做标准答案与 Triton 内核的输出逐元素对比打印最大误差max_diff (result - baseline).abs().max() print(Triton vs PyTorch max diff:, max_diff.item())这段代码非常值得新手学习每写一个自定义内核都要先和 PyTorch 参考实现做数值对齐确保加速不会牺牲正确性。读者可以直接运行python cis_pooling.py复现验证。内核之外滑动窗口与 top-k 如何协同Triton 内核只是粗筛环节的一环它产出的score_cis还要和max_pooling_1d_varlen得到的块分数融合kernel_size32、kernel_stride16每 32 个 token 一个压缩块滑动步长 16相邻块有 50% 重叠保证边界信息不丢失融合后的分数先保留init_blocks local_blocks select_blocks个候选块对应开头固定块、局部窗口块等其余位置填-inf排除最终topk选出真正参与计算的块索引按序排序后交给稀疏注意力内核正是这套Triton 粗筛 稀疏精算的组合拳让 NOSA-1B 在保持长文本质量的同时把解码吞吐做到远超全量注意力的水平。小结三行代码看懂 Triton 加速的本质回顾整个内核加速本质可以浓缩为三点优化手段代码体现收益三维网格并行tl.program_id(0/1/2)头×样本×窗口全并行块式加载与归约tl.loadtl.sum零中间张量、减少访存编译期常量折叠tl.constexpr参数循环展开、极致指令优化如果你也想给自己的模型写高性能算子强烈建议从cis_pooling.py这个 100 行不到的小文件入手它麻雀虽小、五脏俱全triton.jit装饰器、网格设计、mask 处理、数值校验一应俱全是学习 Triton 内核最理想的入门范本。下一篇源码解析我们将继续深入compressed_attention看看 top-k 稀疏注意力是怎么和 Flash Attention 无缝衔接的。【免费下载链接】NOSA-1B项目地址: https://ai.gitcode.com/OpenBMB/NOSA-1B创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考