NOSA-1B 性能基准全解析:5.04× 解码吞吐提升是怎么测出来的?

📅 2026/8/17 17:09:23
NOSA-1B 性能基准全解析:5.04× 解码吞吐提升是怎么测出来的?
NOSA-1B 性能基准全解析5.04× 解码吞吐提升是怎么测出来的【免费下载链接】NOSA-1B项目地址: https://ai.gitcode.com/OpenBMB/NOSA-1BNOSA-1B 是 OpenBMB 开源社区发布的可训练稀疏注意力模型专为 KV-Cache 卸载与解码吞吐优化而生。官方基准显示NOSA-1B 相比 FullAttn 基线可实现最高5.04× 解码吞吐提升相比 InfLLMv2 提升 1.92×相比 ShadowKV 提升 1.83×。本文为你逐层拆解这组性能基准的测试方法、加速原理与复现要点让新手也能看懂 NOSA-1B 的性能从何而来。一、NOSA-1B 是什么一句话看懂核心功能NOSANative and Offloadable Sparse Attention是一种**可训练 可卸载的稀疏注意力机制**配套推理系统NOSI将其落地为实际加速效果。拆开来看就两个关键词Native原生模型在训练阶段就学会稀疏注意力模式而不是推理时临时剪枝质量更有保障Offloadable可卸载KV-Cache 可以卸到显存之外如 CPU 内存大幅降低 GPU 显存压力从而支持更长上下文。NOSA-1B 是这一机制在 1B 参数规模上的落地模型模型实现集中在modeling_llama_long_infllmv2.py中的SparseLlamaForCausalLM配合cis_pooling.py中的 Triton 池化算子完成加速。二、5.04× 解码吞吐提升的基准是怎么测出来的对比基线有哪些性能数字只有放在同一把尺子下才可信NOSA 团队选择了三组主流基线| 基线方案 | 类型 | 相对 NOSA-1B 的吞吐提升 | | :- | :- | :- | | FullAttn | 全量注意力传统方案 |5.04×| | InfLLMv2 | 长上下文稀疏注意力 | 1.92× | | ShadowKV | KV-Cache 卸载方案 | 1.83× |测试条件与指标模型规模1B / 3B / 8B 三档 LLM 均参与评测评测场景长上下文 长文本生成这正是 KV-Cache 最容易爆显存的场景核心指标解码吞吐每秒生成 token 数这是衡量生成速度的最直观指标。结果如何解读相对FullAttn同样显存预算下吞吐提升约5 倍意味着生成长文本时等待时间大幅缩短相对同为稀疏方案的InfLLMv2提升近 2 倍说明可训练的稀疏注意力比纯推理期稀疏更高效相对同为卸载方案的ShadowKV提升 1.83×说明显式局部性约束带来的收益是实打实的。三、解码吞吐为什么能提升 5 倍四大机制拆解1. KV-Cache 卸载把记忆搬出显存 生成每个 token 都要读取全部历史 KV-Cache序列越长显存占用越大。NOSA 的思路是把大部分 KV-Cache 卸载到显存之外GPU 只保留当前正在用的部分从而用更少显存跑更长序列。2. 显式局部性约束只算该算的块 这是 NOSA 区别于传统稀疏注意力的关键。模型在训练时被施加了显式局部性约束——长距离 token 的注意力只需通过压缩表示来传递不必每个 token 都做全量计算。见modeling_llama_long_infllmv2.py中的稀疏注意力参数窗口大小 1024、块大小 64即每个 query 固定关注附近约 16 个 KV 块TopK 选择 64 个高分块其余历史只需读压缩表示。3. 两阶段注意力先压缩、再精挑 解码时先对 KV 做压缩注意力kernel_size32、stride16 的池化得到粗粒度得分再通过 TopK 选出高分块做精确注意力。compressed_attention函数完整实现了这套先筛块、后精算的流程见modeling_llama_long_infllmv2.py。4. Triton 算子把池化做到接近零开销 ⚡均值池化这一步若用 PyTorch 原生实现会有不小开销NOSA 用 Triton 手写了nosa_mean_pool_kernel内核见cis_pooling.py把池化成本压到极低让压缩这一步几乎不拖后腿。四、NOSA-1B 模型配置速览想深入了解的朋友可以先看config.json核心配置如下| 配置项 | 数值 | | :- | :- | | 架构 | SparseLlamaForCausalLMLlama 系 | | 层数 / 隐藏维度 | 28 层 / 2048 | | 注意力头 / KV 头 | 16 / 2 | | 精度 | bfloat16 | | 词表大小 | 73448 | | 稀疏参数 | TopK64、窗口1024、块64、kernel32/stride16 |可以看到NOSA-1B 的 KV 头只有 2 个配合稀疏注意力KV-Cache 的体量本身就被压得很小这也是高吞吐的重要前提。五、如何自己复现与体验 NOSA-1B想亲手跑一遍只需克隆仓库并按需加载git clone https://gitcode.com/OpenBMB/NOSA-1B仓库内已包含完整权重pytorch_model.bin、分词器tokenizer.model与配置config.json模型入口为SparseLlamaForCausalLM可直接通过 Transformers 的AutoModelForCausalLM加载体验。需要对比测试时可以把同样的长文本生成任务分别跑在 NOSA-1B 与 FullAttn 基线上观察吞吐与显存占用的差异。六、除了吞吐长上下文质量同样关键速度再快如果生成质量崩了也没有意义。NOSA 团队强调相比此前的卸载基线方案NOSA-1B 在长上下文 / 长生成场景下的质量同样更优——这正是训练阶段学稀疏带来的红利模型知道自己该看哪里而不是靠启发式规则硬选。小结总结一下 NOSA-1B 性能基准的关键信息✅ 5.04× 解码吞吐提升是对比 FullAttn 的实测结果对比 InfLLMv2、ShadowKV 同样领先✅ 提升来源是「KV-Cache 卸载 显式局部性约束 两阶段稀疏注意力 Triton 加速」的组合拳✅ 1B / 3B / 8B 三档模型均受益长上下文场景收益最明显。对想降低长文本生成成本、又不想牺牲质量的开发者来说NOSA-1B 是一个值得放进候选清单的开源方案。下一篇文章我们继续聊聊如何把 NOSA 部署到自己的推理服务里。【免费下载链接】NOSA-1B项目地址: https://ai.gitcode.com/OpenBMB/NOSA-1B创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考