Sparse Attention 方法综述


从 8K 到 1M token 的长上下文之路:NSA → DSA → V4 → GLM IndexShare → LongCat → MiniMax MSA,以及视频 DiT 的探索

Contents

1. 背景:为什么稀疏注意力重要

随着上下文从 8K 扩展到 128K、200K 乃至 1M token,标准 full attention 的三大瓶颈急剧放大:

O(n²)Attention FLOPs
O(n)KV Cache 占用
随机访存模式
Fig 1. Full Attention FLOPs 随上下文长度的二次方增长 vs 稀疏 Attention 线性增长(示意)
8K token window
128K 长文档
1M token memory
视频时空 token

稀疏化是这条演进路线上的必然选择:只让每个 query 访问少数最相关的历史 token,把 O(n²) 的 full attention 降低到接近 O(n·k),其中 k ≪ n。

2. Native Sparse Attention (NSA)

NSA 是最早将稀疏注意力做成端到端可训练基础模块的方法,基于 GQA 架构,设计了三条并行分支:

① Compressed Branch

对历史 KV 做粗粒度压缩,block size=32,stride=16。提供全局上下文视野,成本低。

② Selected Branch

基于当前 query 从历史 block 中选 top-k,block size=64,n=16 blocks。提供细粒度精确检索

③ Sliding Window Branch

保留最近 w=512 token 的局部窗口,保证局部语言连续性

NSA 注意力掩码可视化

下图展示 NSA 三分支在注意力矩阵上的覆盖模式(token 序列长 N,当前 query 位于最后几行):

Full Attention(基准)

每个 query 看所有历史

NSA — Compressed Branch

粗粒度压缩块覆盖全程

NSA — Selected Branch

top-k block 稀疏选择

NSA — Window Branch

局部滑动窗口

NSA — 三支合并

压缩 + 选择 + 窗口并行

精确 selected
compressed 覆盖
sliding window
不访问
当前 query

NSA 关键参数

参数含义
Compression block size l32多少 token 压成一个 compressed KV
Compression stride d16压缩块滑动步长
Selected block size l'64精确检索的粒度
Selected block count n16每 query 选多少 block
Sliding window w512局部窗口 token 数
实际 attend token(selected)102416 × 64

NSA 局限:是 block 级别检索,无法精确定位单个重要 token;但 block 粒度索引次数少,硬件访存更连续,适合从训练阶段引入。

3. DeepSeek Sparse Attention (DSA)

DSA 由 DeepSeek-V3.2-Exp 引入,接在 MLA(Multi-head Latent Attention) 上,核心是两步 token 级精确检索:

Step 1: Lightning Indexer

轻量 attention 打分头,对全量历史 token 计算相关性分数,成本低但仍是 O(n)。

Step 2: Fine-grained Top-k Selection

每 query 选 top-k=2048 历史 token,主 MLA attention 只在这些 token 上计算。上下文无论多长,主 attention 成本固定。

Fig 2. DSA vs Full Attention:随上下文增长,主 Attention FLOPs 对比(k=2048 固定)

DSA — Lightning Indexer

扫描全量历史(轻量)

DSA — Main Attention

只看 top-k token(token 级精确)

DSA 的关键优势:decode 时主 attention 只看固定 k=2048 个 token,与上下文总长无关——无论 128K 还是 1M,计算量相同。

但 indexer 本身仍要扫描全量历史,top-k 操作和分散 gather 带来新的系统瓶颈,这催生了后续 GLM、LongCat、V4、MiniMax 的改进。

4. DeepSeek-V4:压缩后的稀疏 Cache (CSA/HCA)

DeepSeek-V4 引入 hybrid attention,将 CSA 和 HCA 在不同层交替使用:

CSA — Compressed Sparse Attention

4x 压缩 KV(4 token → 1 entry),再在压缩序列上做 top-k sparse selection。搜索空间缩小 4x,检索仍有较高精度。

HCA — Heavily Compressed Attention

激进的 128x 压缩,在极短压缩序列上做 dense attention。不追求精确检索,提供极低成本的全局摘要通道

CSA — 压缩空间 top-k

4x 压缩后 sparse select

HCA — 极度压缩全局

128x 压缩后 dense attend

V4 — 多层交替策略

CSA / HCA / Window 交替

graph LR
  subgraph V4["DeepSeek-V4 Hybrid Layers"]
    L1["Layer N\nCSA\n(精确检索)"] --> L2["Layer N+1\nHCA\n(全局摘要)"] --> L3["Layer N+2\nLocal Window\n(局部连续)"] --> L4["Layer N+3\nCSA..."]
  end
  KV["原始 KV Cache"] -->|"4x 压缩"| CSA_C["压缩 KV\n(1/4 长度)"]
  KV -->|"128x 压缩"| HCA_C["极压 KV\n(1/128 长度)"]
  CSA_C --> L1
  HCA_C --> L2
  KV --> L3
Fig 3. DeepSeek-V4 多层 hybrid attention 流程:不同层使用不同粒度的 KV 压缩

5. GLM-5.2 IndexShare:跨层共享 Indexer

GLM-5.2 在 DSA 基础上提出 IndexShare:相邻 Transformer 层的 token 重要性高度相关,可以共享同一组 selected indices,大幅减少 indexer 运行次数。

共享粒度每 4 个 sparse attention layer 共享一次 indexer
FLOPs 节省1M context 下约 2.9x per-token FLOPs 减少
关键假设相邻层 token 重要性不会剧烈变化

DSA(每层独立 indexer)

4层,每层独立运行 indexer

IndexShare(4层共享)

1次 indexer,4层复用 indices

graph TD
  Q["Query (1M context)"]
  Q --> IDX["Lightning Indexer\n(运行 1 次)"]
  IDX --> TOPK["Top-k Indices"]
  TOPK --> L1["Layer N Attention"]
  TOPK --> L2["Layer N+1 Attention"]
  TOPK --> L3["Layer N+2 Attention"]
  TOPK --> L4["Layer N+3 Attention"]
  L1 -. "各层自己算 attention weights" .-> OUT["输出"]
  L2 --> OUT
  L3 --> OUT
  L4 --> OUT
Fig 4. IndexShare:Indexer 只跑一次,4 层复用相同的 top-k indices

潜在风险:共享 indices 降低每层选择自由度。GLM-5.2 选择每 4 层共享(而非更多)正是在表达能力和效率之间取得平衡。

6. LongCat:流式感知 + 层级 Indexing

LongCat-Flash 与 LongCat-2.0 在 DSA/GLM 基础上进一步工程化,提出三项改进:

Streaming-aware Indexing (SI)

选 token/block 时考虑硬件访存,避免 selected KV 过于碎片化,尽量变成连续/可合并的内存读取。

Cross-Layer Indexing (CLI)

相邻层 attention saliency 稳定,一次 index 可给多层使用。与 GLM-5.2 IndexShare 同源思路。

Hierarchical Indexing (HI)

先粗粒度 block recall,再在候选 block 内做细粒度 token 选择。避免对全量 token 做高成本 top-k。

LongCat HI — 粗粒度 block recall

先选候选 block

LongCat HI — 细粒度 token 选择

候选 block 内 token 级 top-k

7. MiniMax Sparse Attention (MSA)

MSA 基于 GQA,结构最为简洁:仅两支,避免 NSA 三支并行的复杂度,专注于 block-level 连续访存:

Index Branch

为每个 query / GQA group 选择 top-k block,k=16 blocks(每 block 128 tokens)。

Main Branch

只对被选中的 block 做精确 softmax attention。实际 attend = 16 × 128 = 2048 tokens,与上下文总长无关。

128tokens/block
16selected blocks
2048实际 attend tokens

MSA — Block-level 选择

16 个 block × 128 tokens

MSA vs DSA 粒度对比

block sparse(左)vs token sparse(右)

MSA 的工程优势:block 粒度保证连续访存,kernel 实现简洁,在 H800 等 GPU 上实测加速效果明显。

8. 横向对比

方法 架构基础 主要粒度 核心创新 适合场景
NSA GQA compressed + selected + window (block) 端到端可训练三分支,hardware-aligned 训练推理统一的长上下文基础模块
DeepSeek DSA MLA / MQA token / fine-grained top-k Lightning Indexer + 固定 k=2048 主 attention MLA 模型降低长上下文 attention 成本
DeepSeek-V4 CSA/HCA MQA compressed block top-k / heavily compressed 4x/128x 压缩 + CSA/HCA 层间交替 1M token 级上下文和多粒度记忆
GLM-5.2 IndexShare MLA / MQA+ token top-k + cross-layer shared indices 每 4 层共享 indexer,降低 top-k 打分成本 长上下文 agent / coding,indexer 成本显著时
LongCat-Flash/2.0 MLA / MQA+ GLM++ 三项工程优化 SI + CLI + HI,流式感知 + 层级 indexing 代码 agent、仓库级输入、长上下文产品
MiniMax MSA GQA block-level top-k 简洁双支,连续 block 访存,kernel 友好 生产部署,H800 等 GPU 真实加速
Fig 5. 各方法能力维度雷达图(主观评分,1–5 分)

10.1 "可寻址记忆系统"

新方法不再把上下文视为等价 token 数组,而是分成多种记忆形态:

Local Window

最近几十到几百 token,精确访问

Sparse Retrieval

关键历史片段,top-k 精确检索

Compressed Memory

全局背景,压缩 dense memory

以及 Cross-layer cached indices:高频复用的 top-k 候选集,跨层/跨步共享。

10.2 Indexer 成为新瓶颈

DSA 之后,主 attention 成本大幅下降,但 indexer 的全历史打分、top-k 操作、候选 gather 变成新的系统瓶颈。

  • GLM IndexShare → 跨 4 层共享,减少 indexer 调用次数
  • MiniMax block index branch → block 粒度,降低候选数量
  • LongCat HI → 层级 indexing,先粗后细
  • MoBA block-level routing → 也是降低候选选择成本

本质都在回答同一个问题:如何让 indexer 更便宜、候选更连续、KV 更少、跨层更可复用?

10.3 压缩和稀疏正在融合

两条路线:

压缩后稀疏(CSA/HCA)

先压缩 KV,再稀疏选择;或压缩到足够短后 dense attend。代表:DeepSeek-V4。

直接 Block Sparse(MSA/MoBA)

不先压缩,让模型或 index branch 选择少数 block。代表:MiniMax MSA。

未来长上下文架构很可能不是 pure sparse 或 pure compression,而是多种记忆并存

10.4 训练阶段原生支持比后验稀疏更重要。NSA、DeepSeek、GLM、LongCat、MiniMax 都倾向在 continued training / mid-training 中引入稀疏机制。硬改成 sparse 往往损失长程能力。

10. 视频模型中的 Sparse Attention

视频和 NLP 都是长序列建模,sparse attention + SWA 在视频里是合理方向;但视频 token 数量和 diffusion 推理方式让 LLM 式 token-level sparse 很难直接搬过来。

共同点 vs 关键差异

  • 都是序列数据:文本是 token 序列,视频是时空 patch token 序列
  • 都需要长程记忆:长文本需上下文记忆,长视频需历史帧/场景/角色状态
  • 都强依赖局部邻域:当前 token 依赖附近 token / 帧 / 空间邻域
  • 都需要少量远程信息:文本需前文设定,视频需第一帧、场景布局、动作历史

基本结构可以是:local window(短程)+ sparse memory(远程一致性)

  • 时空结构:视频 token 不是一维,一帧有几百到几千 spatial tokens
  • Diffusion 推理:每个 block 不是 forward 一次,而是多个 denoising steps(T 步)
  • 成本爆炸:token-level indexer 成本 ≈ T × K × S × history_tokens,比文本高两个数量级以上
  • 架构限制:现有视频 DiT 大多是 MHA/full attention 预训练,不是 MLA/MQA,直接改动风险高

视频 Token-level Sparse 成本估算

视频 token-level indexer 成本: T × K × S × history_tokens T = diffusion steps (e.g., 20-50) K = block 内帧数 (e.g., 16) S = 每帧 spatial tokens (e.g., 256-4096) history_tokens = 历史 token 总数 vs 文本 decode: 1 × 1 × 1 × history_tokens 视频比文本高 T×K×S ≈ 10,000x 以上

更可行的第一步:替换 Attention Mask

保留原视频模型主体

不动 DiT 架构,不引入新 indexer

替换 Attention Mask

full attention → block sparse + SWA

Minecraft 实验结果

实验设置:wan1.3b(128×128, 前5s推5s)、wan5b(128×128, 前10s推20s),对比不同 KV Cache 策略。

方法KV 压缩比描述细节质量
MHA full hist1x(无压缩)完整历史 KV Cache,基准最好
MQA~num_heads xMulti-Query Attention下降
kv3d c88x对 KV Cache 做 8 倍压缩明显下降
kv3d 122 2664x对 KV Cache 做 4 倍压缩下降

MQA 或压缩都对回访的细节影响很大——数值指标难以反映细节好坏的差异。

关键发现:动态 Memory 任务

静态回访任务处于 retrieval 表达能力范围内,但 retrieval 不能解决动态 memory 能力。需要设计能体现访问历史 KV Cache 优越性的场景。

2D 炸毁地图障碍物小游戏(动态 memory 任务):无法靠静态记忆完成的任务(需安放炸弹、引爆炸弹,追踪动态状态)。

训练流程

1. wan2.1 初始化
训 128×128×10s AR Diffusion Full Hist
2. 改训 Sparse
sparse + first frame + SWA,512×512×10s
3. Consistency Distill
4. DMD

结论:记录所有 KV Cache 是学会动态 memory 的必要条件。直接训练 sparse 会爆炸(训练不稳定)。

消融结果对比(定性)

配置顺序动态任务表现
GT(Ground Truth)从左完整动态记忆
sparse + first frame + SWA部分动态能力
first frame + SWA(无 sparse)从右缺失动态记忆

To Do