原创 一支烟一朵花 2025-02-19 07:30 上海
长上下文建模对于下一代语言模型至关重要，但标准注意力机制的高计算成本带来了巨大的计算挑战。
研究界越来越认识到长上下文建模是下一代大型语言模型的关键能力，由各种现实世界应用驱动，包括深入推理、存储库级代码生成和多轮自主代理系
长上下文之熵
长上下文建模对于下一代语言模型至关重要，但标准注意力机制的高计算成本带来了巨大的计算挑战。
研究界越来越认识到长上下文建模是下一代大型语言模型的关键能力，由各种现实世界应用驱动，包括深入推理、存储库级代码生成和多轮自主代理系统。
最近的突破，包括 OpenAI 的o系列模型 DeepSeek-R1和 Gemini 1.5 Pro，使模型能够处理整个代码库、冗长的文档、在数千个令牌上保持连贯的多轮对话，并跨长期依赖关系执行复杂的推理。
随着序列长度的增加，原版注意力机制的高复杂性成为关键的延迟瓶颈。
理论估计表明，在解码 64k 长度的上下文时，使用 softmax 架构的注意力计算占总延迟的 70-80%，这凸显了对更高效注意力机制的迫切需求。
NSA
长上下文建模对于下一代语言模型至关重要，但标准注意力机制的高计算成本带来了巨大的计算挑战。
稀疏注意力为提高效率同时保持模型功能提供了一个有前途的方向。
稀疏注意力的核心思想就是减少不必要的计算，不像传统的“全注意力”机制那样每一对词之间都要计算一次注意力，而是挑选出最有用的部分进行计算。
这样，计算量大大减少，效率自然提高。
就像你做饭时，不需要把整个厨房的食材都搬出来，只取最重要的几样，就能做出美味佳肴，节省了时间和精力。
为了实现更有效和高效的稀疏注意力，DeepSeek研发团队提出了 NSA，这是一种原生可训练的稀疏注意力架构，集成了分层令牌建模。
它将算法创新与硬件对齐的优化集成在一起，以实现高效的长上下文建模。
NSA采用了动态分层稀疏策略，通过将粗粒度的令牌压缩和细粒度的令牌选择相结合，来保证模型既能捕捉到全局上下文的感知，又能保持局部信息的精确度。
这就像你在看一场足球比赛时，既能看到整个场地的布局，又不至于错过关键球员的动作，保持对比赛的全面把握。它通过两项关键创新推进了稀疏注意力设计：
（1）通过算术强度平衡算法设计实现了大幅加速，并针对现代硬件进行了实现优化。
（2）支持端到端训练，在不牺牲模型性能的情况下减少训练前计算。
效果惊人
如图 1 所示，实验表明，使用 NSA 预训练的模型在一般基准、长上下文任务和基于指令的推理中保持或超过全注意力模型。 就像是你在比赛中，虽然没有用上最贵的装备，但依然能跑出好成绩。
同时，NSA在处理64k长度序列时，在解码、前向传播和后向传播方面的速度比传统的全注意力机制快得多，证明了它在整个模型生命周期中的高效性。
NSA 引入了与上述关键要求相对应的两项核心创新：
硬件对齐系统：优化 Tensor Core 利用率和内存访问的逐块稀疏注意力，确保平衡的算术强度。
训练感知设计：通过高效的算法和反向算子实现稳定的端到端训练。
这种优化使NSA能够支持高效部署和端到端培训。
框架通过三个并行的注意力分支处理输入序列：
对于给定的查询，前面的键和值被处理为
粗粒度模式的压缩注意力、
重要标记块的选定注意力和
本地上下文的滑动注意力。
右图：每个分支产生的不同注意力模式的可视化。绿色区域表示需要计算注意力分数的区域，
而白色区域表示可以跳过的区域。
NSA的内核设计
内核通过 GQA 组加载查询（网格循环），获取相应的稀疏 KV 块（Inner Loop），并在 SRAM 上执行注意力计算。
绿色块表示 SRAM 上的数据，而蓝色块表示 HBM 上的数据。
详细时序
假设我们用一种叫“稀疏注意力内核”的方法来优化计算处理速度，就像优化一条生产线一样。
想象有一条生产线，分成三个主要部分：
外层循环、内层循环（组内的处理）和处理单元
外层循环就像生产调度员，负责把控整个生产流程。
它会告诉内层循环（组内的处理）*按顺序处理每个查询块的位置，就像一位指挥员，指挥工人依次完成每一段工作任务。
在每个查询块的位置处理时，内层循环（组内的处理）会先找到该位置属于某个组的所有任务（查询头）以及这些任务需要的特定工具（KV块索引），然后要求处理单元把这些任务和工具都搬到一块高效的工作区域（SRAM）里。
处理单元会根据索引，按顺序把需要的材料（K、V块）搬到工作区域，就像是工人按照清单搬运物料到工作台，这样可以减少搬运的次数，提高效率。 处理单元开始用这些材料进行核心的工作（注意力计算），也就是把不同的材料组合起来，制造出产品。
当这个位置的工作完成后，内层循环（组内的处理）会告诉调度员（外层循环）已经完成了这一步，调度员会继续指挥下一个位置的处理。
在整个过程中，外层循环（网格上的外环）还负责合理安排工人的工作，确保各个部分的流水线（GPU流式多处理器）都能高效运转，不会因为任务分配不均而出现闲置或者过载的情况。
最后，这个生产流程通过合理组织和优化搬运（减少冗余的工具和材料运输）以及合理安排工作（平衡各部分的工作负载），达到了高效的生产速度。
就像我们在做一项复杂的工作，但通过把任务分组、合理搬运工具材料、优化工作流程，就能提高效率。
关键创新
- 查询分组策略：不再按时间连续块加载查询和KV对，而是将属于同一GQA组的查询头及其共享的KV块加载到SRAM中，减少内存访问的不连续性。共享KV获取：在内层循环中按序加载连续的KV块，减少内存加载开销。外层网格调度：利用Triton的网格调度器处理查询/输出循环，简化内核逻辑并提升效率。
性能优势
- 消除了冗余的KV传输，通过组内共享KV块实现数据复用。均衡了GPU流式多处理器上的计算工作负载，提高整体算术强度。
在硬件层面优化了内存访问模式和计算调度，使得稀疏注意力机制在大规模LLM训练和解码过程中具备更高的效率和可扩展性。
快速总结
论文链接
https://arxiv.org/pdf/2502.11089v1
研究背景
- 长序列建模的重要性：长序列建模对于下一代语言模型至关重要，应用场景包括深度推理、代码生成和多轮对话系统。标准注意力机制的挑战：标准注意力机制（如 Transformer 中的 Full Attention）在处理长序列时计算成本高，导致显著的延迟瓶颈。
核心创新
动态层次稀疏策略：
- 结合粗粒度的 token 压缩和细粒度的 token 选择，既保留全局上下文感知能力，又保持局部精度。通过压缩块和选择块的动态组合，实现高效的长序列建模。
算法创新与硬件优化：
- 算术强度平衡的算法设计：通过优化现代硬件的实现，显著提高计算速度。端到端训练：减少预训练计算，同时保持模型性能。
关键数值结果：
- 在 64k 序列长度上，NSA 在解码、前向传播和反向传播阶段均实现了显著的速度提升，分别达到 11.6×、9.0× 和 6.0× 的加速比。在长序列任务中，NSA 的性能优于 Full Attention 和其他稀疏注意力方法。
核心设计
稀疏注意力机制设计：
- 压缩块：通过聚合连续的 token 块，生成压缩的键和值，减少计算量。选择块：选择重要的 token 块进行注意力计算，保留关键信息。滑动窗口：处理局部上下文，防止模型过度依赖局部模式。
硬件对齐的内核设计：
- 块级数据加载：针对 GQA 和 MQA 架构，优化数据加载和计算。共享 KV 获取：减少内存访问，提高计算效率。
训练和推理优化：
- 端到端训练：通过可训练的稀疏注意力模块，减少预训练计算。推理加速：在解码阶段，通过减少 KV 缓存加载，显著降低延迟。
实验结果
性能评估：
- 通用基准测试：NSA 在多个基准测试中表现优于 Full Attention 和其他稀疏注意力方法。长序列任务：在 64k 序列长度的针-in-a-haystack 测试中，NSA 实现了完美的检索精度。推理任务：在 AIME 数学推理基准测试中，NSA 在 8k 和 16k 序列长度上均优于 Full Attention。
效率分析：
- 训练速度：在 64k 序列长度上，NSA 的前向传播速度提升 9.0×，反向传播速度提升 6.0×。解码速度：在 64k 序列长度上，NSA 的解码速度提升 11.6×。
NSA 通过算法创新和硬件优化，实现了高效的长序列建模，显著提高了计算效率和模型性能。
这一研究为长序列建模提供了新的思路和方法，具有重要的理论和实践意义。
DeepSeek最近似乎像是开了挂一样，一个月一篇硬核论文，他们在大模型训练，推理效率和成本优化的路上越走越顺。
Long live DeepSeek,Long live 开源！
