深入浅出 NOSA-1B稀疏注意力机制完全原理解析【免费下载链接】NOSA-1B项目地址: https://ai.gitcode.com/OpenBMB/NOSA-1BNOSA-1B 是 OpenBMB 开源社区发布的可训练稀疏注意力模型它把「稀疏注意力」从论文带进了真实可用的代码在保持长文本生成质量的同时通过 KV 缓存卸载让解码吞吐最高提升 5 倍。本文将从零开始带你彻底看懂稀疏注意力机制原理理解 NOSA-1B 凭什么又快又准并附上源码导读供你亲自验证。为什么大模型需要稀疏注意力KV 缓存长上下文的第一道坎标准的 Transformer 注意力中每个新 token 都要回看所有历史 token计算量随序列长度呈平方级增长。为了不重复计算大模型会把历史 token 的 Key 和 Value 缓存下来这就是KV 缓存KV Cache。问题在于上下文越长KV 缓存越大。处理超长文本时KV 缓存会迅速吃光显存推理速度也随之下降。这成了长上下文大模型落地的最大瓶颈之一。注意力其实很稀疏研究人员很早就发现一个反直觉的事实真正的注意力分布非常稀疏。大多数 token 其实只依赖少数几个关键的历史 token——最近的局部内容以及少数重量级的远处记忆。既然大多数注意力分数趋近于零那有没有可能只计算那部分重要的呢这正是稀疏注意力机制的出发点。NOSA-1B 是什么一句话概括NOSA-1B 是一个原生且可卸载的稀疏注意力模型。它的完整名称是 NOSANative and Offloadable Sparse Attention1B 指的是约 10 亿参数的模型规模。它基于 Llama 架构改造而来代码入口在 modeling_llama_long_infllmv2.py 中的SparseLlamaForCausalLM模型共 28 层、16 个注意力头并借助 LongRoPE 将上下文窗口扩展到 32K token相关配置见 config.json。它的亮点不是参数量而是把 KV 缓存从显存卸载到 CPU/磁盘同时不牺牲生成质量。稀疏注意力机制的核心思想记住块而不是每个字NOSA-1B 的思路可以概括成三步压缩把长长的历史 KV 按块压缩成少量记忆块打分给每个记忆块算一个重要性分数选择每次生成只挑最重要的少数几个块参与注意力计算。这样显存里永远只放精选的 KV其余全部卸载到内存或磁盘长上下文推理自然又快又省。NOSA-1B 的四大核心设计详解1. KV 压缩把记忆浓缩成块 在标准模型中每个 token 都会保留一份 Key。NOSA-1B 引入CompressK模块用kernel_size32、stride16的滑动窗口对 Key 做均值压缩每 32 个 token 合并成一个压缩 Key。原始 Keyt0 t1 t2 ... t31 t32 ... 压缩 Key K0(代表 t0~t31) K1(代表 t16~t47) ...效果立竿见影参与计算的 Key 数量大幅减少KV 缓存体积随之骤降。这部分逻辑就写在 modeling_llama_long_infllmv2.py 的CompressK类中。2. CIS 压缩重要性分数给每个记忆块打分 ⭐压缩之后如何知道哪个记忆块值得关注NOSA-1B 引入了一个可学习的CISCompressed Importance Score压缩重要性分数机制模型通过一个小型线性层delta结合可学习参数A为每个位置动态生成重要性分数再经 softplus 激活后作为 Value 的缩放权重。为了高效计算块级分数项目用 Triton 手写了一个均值池化 kernel见 cis_pooling.py 中的nosa_mean_pooling把逐 token 的 CIS 平滑成每个块的分数。分数越高说明这个记忆块越值钱。3. 两级块选择先粗选、再精排 拿到块分数后NOSA-1B 采用两级选择策略兼顾效率与精度第一级粗选在压缩后的 KV 上先算一遍注意力快速挑出select_blocks个候选块第二级精排在候选中按 CIS 分数排序最终选出topk64个块真正参与注意力计算。同时模型始终固定保留两类块序列开头的 1 个初始块init_blocks1相当于文章标题/摘要以及当前位置附近的 1024 个局部 tokenwindow_size1024相当于最近几页。这套组合拳正是compressed_attention函数的核心逻辑。4. 可卸载性把 KV 缓存搬出显存 这是 NOSA 名字里 Offloadable 的关键既然每步只用到选中的少量块那么没被选中的 KV 完全没必要留在显存里可以卸载到 CPU 内存甚至磁盘。配合配套的推理系统 NOSI显存占用被压到极低长序列生成也能流畅运行。NOSA-1B 的实际效果如何官方在 1B/3B/8B 三档模型上做了对比测试解码吞吐提升非常直观对比基线吞吐提升FullAttn全注意力最高 5.04×InfLLMv2最高 1.92×ShadowKV最高 1.83×也就是说在几乎不掉点的情况下NOSA-1B 能跑出比全量注意力快数倍的解码速度这正是稀疏注意力机制的价值所在。源码导读从哪里看起想动手验证的朋友建议按下面顺序阅读仓库文件modeling_llama_long_infllmv2.py模型主体。先看CompressKKV 压缩再看compressed_attention两级块选择最后看SparseLlamaForCausalLM模型入口cis_pooling.pyCIS 均值池化的 Triton 实现理解打分如何落地config.json模型结构参数与 LongRoPE 扩展配置可以看到 32K 上下文的 rope_scaling 细节。总结NOSA-1B 用一套优雅的组合拳——KV 压缩 CIS 重要性打分 两级块选择 可卸载设计——解决了长上下文推理的显存与速度痛点。对新手来说理解稀疏注意力机制的关键就一句话不是所有的历史都值得看学会挑重点才是长上下文推理的出路。希望这篇文章能帮你迈出理解稀疏注意力的第一步。如需动手复现可通过git clone https://gitcode.com/OpenBMB/NOSA-1B获取完整代码和模型权重结合本文的源码导读逐一验证。【免费下载链接】NOSA-1B项目地址: https://ai.gitcode.com/OpenBMB/NOSA-1B创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 深入浅出 NOSA-1B：稀疏注意力机制完全原理解析
深入浅出 NOSA-1B稀疏注意力机制完全原理解析【免费下载链接】NOSA-1B项目地址: https://ai.gitcode.com/OpenBMB/NOSA-1BNOSA-1B 是 OpenBMB 开源社区发布的可训练稀疏注意力模型它把「稀疏注意力」从论文带进了真实可用的代码在保持长文本生成质量的同时通过 KV 缓存卸载让解码吞吐最高提升 5 倍。本文将从零开始带你彻底看懂稀疏注意力机制原理理解 NOSA-1B 凭什么又快又准并附上源码导读供你亲自验证。为什么大模型需要稀疏注意力KV 缓存长上下文的第一道坎标准的 Transformer 注意力中每个新 token 都要回看所有历史 token计算量随序列长度呈平方级增长。为了不重复计算大模型会把历史 token 的 Key 和 Value 缓存下来这就是KV 缓存KV Cache。问题在于上下文越长KV 缓存越大。处理超长文本时KV 缓存会迅速吃光显存推理速度也随之下降。这成了长上下文大模型落地的最大瓶颈之一。注意力其实很稀疏研究人员很早就发现一个反直觉的事实真正的注意力分布非常稀疏。大多数 token 其实只依赖少数几个关键的历史 token——最近的局部内容以及少数重量级的远处记忆。既然大多数注意力分数趋近于零那有没有可能只计算那部分重要的呢这正是稀疏注意力机制的出发点。NOSA-1B 是什么一句话概括NOSA-1B 是一个原生且可卸载的稀疏注意力模型。它的完整名称是 NOSANative and Offloadable Sparse Attention1B 指的是约 10 亿参数的模型规模。它基于 Llama 架构改造而来代码入口在 modeling_llama_long_infllmv2.py 中的SparseLlamaForCausalLM模型共 28 层、16 个注意力头并借助 LongRoPE 将上下文窗口扩展到 32K token相关配置见 config.json。它的亮点不是参数量而是把 KV 缓存从显存卸载到 CPU/磁盘同时不牺牲生成质量。稀疏注意力机制的核心思想记住块而不是每个字NOSA-1B 的思路可以概括成三步压缩把长长的历史 KV 按块压缩成少量记忆块打分给每个记忆块算一个重要性分数选择每次生成只挑最重要的少数几个块参与注意力计算。这样显存里永远只放精选的 KV其余全部卸载到内存或磁盘长上下文推理自然又快又省。NOSA-1B 的四大核心设计详解1. KV 压缩把记忆浓缩成块 在标准模型中每个 token 都会保留一份 Key。NOSA-1B 引入CompressK模块用kernel_size32、stride16的滑动窗口对 Key 做均值压缩每 32 个 token 合并成一个压缩 Key。原始 Keyt0 t1 t2 ... t31 t32 ... 压缩 Key K0(代表 t0~t31) K1(代表 t16~t47) ...效果立竿见影参与计算的 Key 数量大幅减少KV 缓存体积随之骤降。这部分逻辑就写在 modeling_llama_long_infllmv2.py 的CompressK类中。2. CIS 压缩重要性分数给每个记忆块打分 ⭐压缩之后如何知道哪个记忆块值得关注NOSA-1B 引入了一个可学习的CISCompressed Importance Score压缩重要性分数机制模型通过一个小型线性层delta结合可学习参数A为每个位置动态生成重要性分数再经 softplus 激活后作为 Value 的缩放权重。为了高效计算块级分数项目用 Triton 手写了一个均值池化 kernel见 cis_pooling.py 中的nosa_mean_pooling把逐 token 的 CIS 平滑成每个块的分数。分数越高说明这个记忆块越值钱。3. 两级块选择先粗选、再精排 拿到块分数后NOSA-1B 采用两级选择策略兼顾效率与精度第一级粗选在压缩后的 KV 上先算一遍注意力快速挑出select_blocks个候选块第二级精排在候选中按 CIS 分数排序最终选出topk64个块真正参与注意力计算。同时模型始终固定保留两类块序列开头的 1 个初始块init_blocks1相当于文章标题/摘要以及当前位置附近的 1024 个局部 tokenwindow_size1024相当于最近几页。这套组合拳正是compressed_attention函数的核心逻辑。4. 可卸载性把 KV 缓存搬出显存 这是 NOSA 名字里 Offloadable 的关键既然每步只用到选中的少量块那么没被选中的 KV 完全没必要留在显存里可以卸载到 CPU 内存甚至磁盘。配合配套的推理系统 NOSI显存占用被压到极低长序列生成也能流畅运行。NOSA-1B 的实际效果如何官方在 1B/3B/8B 三档模型上做了对比测试解码吞吐提升非常直观对比基线吞吐提升FullAttn全注意力最高 5.04×InfLLMv2最高 1.92×ShadowKV最高 1.83×也就是说在几乎不掉点的情况下NOSA-1B 能跑出比全量注意力快数倍的解码速度这正是稀疏注意力机制的价值所在。源码导读从哪里看起想动手验证的朋友建议按下面顺序阅读仓库文件modeling_llama_long_infllmv2.py模型主体。先看CompressKKV 压缩再看compressed_attention两级块选择最后看SparseLlamaForCausalLM模型入口cis_pooling.pyCIS 均值池化的 Triton 实现理解打分如何落地config.json模型结构参数与 LongRoPE 扩展配置可以看到 32K 上下文的 rope_scaling 细节。总结NOSA-1B 用一套优雅的组合拳——KV 压缩 CIS 重要性打分 两级块选择 可卸载设计——解决了长上下文推理的显存与速度痛点。对新手来说理解稀疏注意力机制的关键就一句话不是所有的历史都值得看学会挑重点才是长上下文推理的出路。希望这篇文章能帮你迈出理解稀疏注意力的第一步。如需动手复现可通过git clone https://gitcode.com/OpenBMB/NOSA-1B获取完整代码和模型权重结合本文的源码导读逐一验证。【免费下载链接】NOSA-1B项目地址: https://ai.gitcode.com/OpenBMB/NOSA-1B创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考