OpenBMB 开源 NOSA-1BKV Cache 卸载加持长文本推理提速 5.04 倍【免费下载链接】NOSA-1B项目地址: https://ai.gitcode.com/OpenBMB/NOSA-1BOpenBMB 开源社区正式发布NOSA-1B一款基于 NOSANative and Offloadable Sparse Attention原生可卸载稀疏注意力机制的 10 亿参数长文本推理模型。它通过 KV Cache 卸载技术大幅压缩显存占用让长文本生成吞吐量相比全量注意力FullAttn最高提升5.04 倍相比 InfLLMv2 提升 1.92 倍。本文将从零开始用通俗的语言拆解 NOSA-1B 的原理、性能亮点与快速上手方法。 太长不看版长文本推理的瓶颈是 KV Cache 吃显存NOSA-1B 把不重要的KV Cache 卸载到 CPU/磁盘只把关键部分留在 GPU从而又快又省。长文本推理为什么又慢又贵先认识 KV Cache大模型生成每个 token 时都要重新计算之前所有 token 的 KeyK和 ValueV于是系统会把它们缓存下来这就是KV Cache。问题在于KV Cache 会随上下文长度线性膨胀。处理 32K、128K 甚至更长的上下文时KV Cache 动辄占几个 GB 到几十 GB 显存直接把显卡塞满导致显存不够长文本根本跑不起来 被迫用更大的 batch速度骤降推理成本成倍上升所以怎么处理 KV Cache 就成了长文本推理加速的核心战场。NOSA-1B 是什么一文看懂原生可卸载稀疏注意力NOSA-1B是 OpenBMB 开源社区基于清华 THUNLP 提出的NOSA 稀疏注意力机制训练的 1B 模型配套推理系统名为NOSI。它的核心思路很巧妙给注意力加本地性约束Locality Constraint训练时就告诉模型——注意力要尽量集中在附近 token 和少量关键远处 token 上让模型学会哪些位置重要。训练时就考虑卸载不是推理时才临时想办法而是把KV Cache 可卸载直接写进训练目标所以卸载后依然能保持高质量输出。推理时只保留关键 KVGPU 上只留当前窗口和 Top-K 个最重要的块其余 KV Cache 全部卸载到 CPU 内存显存压力骤降。和同类方案相比NOSA 是原生可训练 可卸载的稀疏注意力因此长上下文/长生成质量明显优于此前的卸载方案。KV Cache 卸载如何带来 5.04 倍加速根据官方在 1B / 3B / 8B 三档模型上的测试NOSA 配合 NOSI 推理系统的解码吞吐量提升非常直观对比基线吞吐提升说明FullAttn全量注意力最高 5.04×不牺牲上下文长度的前提下大幅提速InfLLMv2最高 1.92×同属 KV 卸载路线NOSA 更优ShadowKV最高 1.83×稀疏化方案对比也就是说同样一张显卡、同样 128K 长文本NOSA-1B 生成速度可以快出好几倍长文本场景下的性价比非常突出。⚡工程实现亮点Triton 算子 两阶段注意力NOSA-1B 的仓库实现也很有看点适合想深入研究或二次开发的工程师稀疏注意力核心模型架构为SparseLlamaForCausalLM定义在modeling_llama_long_infllmv2.py采用两阶段注意力——先用压缩后的 KV 快速选出 Top-K 关键块再做精确注意力块大小与 Top-K 均设为 64。Triton 加速算子cis_pooling.py中用 Triton 手写了nosa_mean_pool_kernel均值池化核用于计算压缩重要性分数CIS性能开销极低。MiniCPM 适配modeling_minicpm.py与configuration_minicpm.py提供了基于 MiniCPM 架构的 NOSA 实现说明该方法可迁移到不同底座模型上。配置细节config.json显示模型为 28 层、16 头、bf16 精度并采用 LongRoPE 位置编码将原始 32K 上下文进一步外推。 小知识CISCompressed Informative Score压缩重要性分数就是 NOSA 判断哪个历史块值得保留在 GPU的依据这也是 KV Cache 卸载能否卸载得准的关键。快速上手体验 NOSA-1B 只需三步NOSA-1B 以 Hugging Face Transformers 格式发布仓库内包含config.json模型配置pytorch_model.bin权重文件tokenizer.json/tokenizer.model/special_tokens_map.json分词器generation_config.json生成参数默认 temperature 0.8想获取全部文件可以克隆仓库git clone https://gitcode.com/OpenBMB/NOSA-1B克隆后配合 Transformers 加载即可使用模型类会自动通过auto_map映射到modeling_llama_long_infllmv2.py中的SparseLlamaForCausalLM。需要注意的是KV 卸载相关的高性能算子依赖flash-attn与infllm_v2扩展建议在安装好这些依赖的 CUDA 环境中运行。NOSA-1B 适合哪些场景超长文档问答几十万字的论文、财报、代码库直接喂给模型长对话 Agent多轮会话历史不丢、不爆显存️单卡部署长文本显存有限的开发者和研究者用一张消费级显卡跑长文本研究稀疏注意力想学习 KV Cache 卸载、可训练稀疏注意力实现的工程师小结NOSA-1B 用原生可卸载稀疏注意力回答了长文本推理的两大难题显存不够和速度太慢。5.04 倍的吞吐提升、对长上下文质量的保持加上 Apache-2.0 开源协议让它成为长文本推理领域非常值得一试的模型。如果你正被长文本显存问题困扰不妨从克隆 NOSA-1B 开始亲手感受 KV Cache 卸载带来的速度飞跃。【免费下载链接】NOSA-1B项目地址: https://ai.gitcode.com/OpenBMB/NOSA-1B创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考