最近在折腾超长上下文推理手头模型的名字里挂着一串关键词MiniCPM-SALA、稀疏注意力、线性注意力、单卡百万 token。刚开始我以为是营销包装等真正沿着这套思路把注意力模块换掉以后才发现稀疏-线性这四个字确实能把长上下文这堵墙撞开一个口子。这篇文章不打算复述论文而是围绕这个架构谈清楚三个问题它到底解决了什么、工程上怎么落地、以及我在实测中踩过的坑。适合正在做长文本 RAG、摘要生成、推理优化或者想复现类似架构的同学参考我会把注意力复杂度、显存预算、稀疏模式选择、压测口径全部摊开讲尽量让你顺着文字就能在单卡上复现出一条可行的路子。1. 为什么长上下文推理会卡在注意力这道坎上1.1 算力爆炸标准注意力的 O(n²) 宿命先别急着看稀疏、线性这些花哨名词得先把长上下文为什么难这件事说透。绝大多数 Transformer 模型用的是标准注意力核心动作是三个矩阵相乘softmax(QKᵀ / √d)·V。Q、K、V 的长度都是 n这里 n 就是序列里的 token 数d 是每个头的维度。问题出在 QKᵀ 这一步它的计算结果是一个 n×n 的矩阵计算量和显存占用都随 n 平方上涨。拿一个直观例子说假设你有一段 100 万 token 的输入d 取 4096。标准注意力前向计算一次需要大约 2×n²×d 次浮点运算代入就是 2×10¹²×4096大概 8×10¹⁵ FLOPs。当前主流单卡 BF16 算力在 3×10¹⁴ FLOP/s 左右光算这一层就是几十秒的纯算力消耗还没算激活值、反向传播、梯度更新。这就像开一场上万人的交流会要求每个人都和其余所有人单独握手时间自然是指数级增长的严格说是平方级但体验上和爆炸差不多。很多同学第一次跑长文本时发现显存爆掉第一反应是模型太大其实把模型参数冻结后注意力矩阵才是吃显存的大户。n1M 时单层单头的 QKᵀ 矩阵就是 1M×1M 个元素FP16 下就是 2TB。这个数字根本不可能塞进任何一张消费级卡里哪怕你在云端租到 80GB 显存的高端卡也只够这个中间矩阵的零头。1.2 显存灾难KV Cache 为什么吃不消除了 QKᵀ 这个临时矩阵推理阶段还会持续保存一份 KV Cache。它的存在是为了避免每个新 token 都把历史输入重新过一遍。KV Cache 的体积公式很简单2 × 层数 × token 数 × 每层每 token 的 KV 维度 × 每个元素占的字节数。假设模型有 32 层KV 维度取 4096FP16 存储100 万 token 输入时2×32×1,000,000×4096×2 字节 ≈ 524GB也就是说哪怕你把注意力计算完全省掉光把历史和中间状态缓存下来就已经超出单卡 80GB 的上限了。这也是为什么百万 token 上下文听起来很刺激做起来很残酷不仅算力墙过不去显存墙也过不去。所以任何号称能单卡跑百万上下文的方案第一件事必须把每 token 平均显存压下来而不是去优化那点矩阵乘法的常数因子。用大白话总结KV Cache 是存每个人说过的话注意力矩阵是每个人和所有人对一次话。长上下文下这两者都撑不住。稀疏-线性注意力做的事情本质就是前者少存、后者少算。2. 稀疏-线性注意力SALA的拆解2.1 稀疏注意力别跟所有人握手只跟关键对象说话稀疏这个词解决的是握手次数问题。它的出发点很朴素序列里 100 万 token当前 token 真的需要和所有历史 token 都建立关联吗大部分时候不需要。写文章时某个词主要和最近的几个词相关偶尔要回头翻前文的某个概念真正需要全场广播的少量特殊 token 就能搞定。稀疏注意力有几种常用套路。第一种是滑窗注意力Sliding Window Attention每个 query 只看前面 w 个 key比如 w4096那么每个 token 的注意力计算量立刻从 n 降到 w。第二种是全局 token在序列开头或者中间固定放几个特殊 token让它们能看到整个序列同时序列里其他 token 也能看到它们相当于群体里的广播电台或公告栏。第三种是块稀疏把序列按块划分允许某些块之间跨距离连接适合需要远程信息跳转的任务。实际项目里最常见的是滑窗 少量全局 token 的组合既便宜又有效。复杂度方面滑窗注意力的计算量是 O(n×w)w 通常远小于 n。当 n1M、w4096 时有效计算量是原来的 1/250 左右这是实打实的降维打击。不过要小心这里的稀疏是指注意力矩阵里的非零元素变少了但你仍然需要一种高效的数据结构去表达这种稀疏性否则只是用大矩阵填一堆零内存问题依然存在。因此落地的时候一般配合分块计算、掩码设计和 kernel 融合来做而不是真的在 PyTorch 里生成一个 n×n 稀疏矩阵去乘。2.2 线性注意力把 QKᵀ 这一步彻底绕过去线性注意力解决的是另一类问题就算你把注意力变成 O(n×w)如果 w 本身也很大比如 32K、64K或者你想进一步降低计算量仍然不够。那再进一步干脆别显式计算注意力权重了。标准注意力里 softmax(QKᵀ/√d) 是非线性的它没法直接化简。线性注意力用了一个 trick把注意力分数的非线性相似度函数拆成两个特征映射的内积即 φ(q)·φ(k)。这样一来输出可以改写成一个先算后乘的形式O_i φ(q_i) · Σ_j (φ(k_j)ᵀ · v_j)注意Σ_j 这一项和当前 query 无关它相当于一个持续累积的全局状态维度是 d×d。这样计算复杂度就从 O(n²) 降到了 O(n×d²)更重要的是它可以流式更新来一个 token就更新一次累积状态不需要把历史全部翻出来重算。我常用一个生活类比标准注意力是每个人都去翻阅所有人的信件线性注意力则是有一个人在公告栏上维护一张信息汇总表新 token 来了只需要看一眼汇总表更新一下自己这条记录就行了。代价是这种全局汇总丢失了逐对细粒度关联类似于只看平均指标不能精确到某个特定历史片段对当前词的影响有多大。所以线性注意力通常不会单独做成整个模型而是作为稀疏注意力旁边的一个全局压缩通道来用。2.3 两者合流为什么稀疏 线性才能做到单卡百万单独用稀疏注意力显存和算力降下来了但窗口外信息可能彻底丢失单独用线性注意力长程信息都在但局部精确匹配能力变弱。SALA 这类架构的核心思路就是把两者在结构上分工局部用滑窗注意力保留细粒度上下文全局用线性注意力维护一个压缩状态少数全局 token 再负责跨窗口的信息广播。一套典型的组合是每个 query 只对最近 w4096 个 key 做标准 softmax 注意力复用 FlashAttention 这类高效内核另外留出 g128 或 512 个全局 slot它们可以和任意位置交互专门承载跨段落信息对全局 slot 之间的交互或者 slot 与全部历史的关系采用线性注意力的累积状态来实现避免再建一块 n×g 的大矩阵。这么设计之后显存开销不再正比于 n²而是正比于 n×(wg)g 通常只有几百。回到前面的 KV Cache 公式如果把每 token 需要存 KV压缩到只有窗口部分和全局压缩状态100 万 token 时的缓存体积就从 524GB 降到了一个可管理的量级。比如窗口 8K、全局 256 个 slotKV 总量约等于 8K256 这一档显存可能控制在 20~40GB 区间具体取决于层数和维度这才给了单卡一张通行证。有个细节特别重要稀疏和线性不是简单的拼在一起而是要让两者的信息路径互补。滑窗负责附近到底说了什么线性负责全文大概在说什么全局 slot 负责某一处具体的远程信息怎么被找到。我在实际调参时发现如果只做滑窗不加全局 slot60 万 token 以上的文档召回率会掉得很难看补上全局 slot 之后远程信息检索明显回血。3. 从零上手SALA 架构的核心步骤与实现3.1 第一步确定稀疏模式与关键参数动手写代码之前先得定下稀疏模式。这一步直接决定后续所有实现的复杂度和信息召回能力不要随便拍脑袋。滑动窗口大小 w 的选择主要看两个因素显存预算和任务类型。见面熟的任务如摘要、翻译、对话可以用 4K~16K需要跨章节引用的任务如长文档 QA、代码库理解至少要 8K 起步并且补足全局 slot。全局 slot 数量在 128~1024 之间太少容易丢信息太多又会把成本拉回平方级。我做了一个简单对比表方便大家按场景选稀疏模式计算复杂度显存压力典型场景纯滑窗O(n×w)低流式对话、部分摘要滑窗 全局 slotO(n×(wg))中长文档 QA、RAG 精排块稀疏 跳线O(n×w 跳线比例×n)中高跨章节、图表交叉引用动态 top-k 稀疏O(n×k) 但计算不规则中科研尝试、特殊数据分布我的建议是第一个版本不要碰动态稀疏先用固定的滑窗 全局 slot跑通。动态稀疏需要实时计算 token 重要度并生成不规则掩码kernel 融合难度大显存碎片也严重新手很容易把推理延迟做得比标准注意力还高。3.2 第二步实现稀疏-线性注意力的核心计算这里给一个教学层面的示意实现。真实训练时不会这么裸写因为逐 token 循环太慢但用来理解数据流足够清楚。import torch import torch.nn.functional as F def feature_map(x): # 线性注意力常用的特征映射简单起见用 relu elu 组合 # 实际项目中还会叠加归一化和位置信息 return torch.cat([F.relu(x), F.elu(x) 1.0], dim-1) def sliding_window_attention(q, k, v, window_size, block_size128): # q/k/v: (b, h, n, d) b, h, n, d q.shape out torch.zeros_like(q) for start in range(0, n, block_size): end min(start block_size, n) q_block q[:, :, start:end] left max(0, start - window_size) k_block k[:, :, left:end] v_block v[:, :, left:end] attn q_block k_block.transpose(-1, -2) / (d ** 0.5) attn torch.softmax(attn, dim-1) out[:, :, start:end] attn v_block return out def linear_global_state(q_global, k, v, stateNone): # 全局 token 用线性注意力维护压缩状态 # state 是一个 d*d 大小的累积矩阵 k_phi feature_map(k) v_phi v if state is None: state torch.einsum(bhnm,bhrm-bhmr, k_phi.transpose(-1, -2), v) # 形状: (b, h, d, d) q_phi feature_map(q_global) out torch.einsum(bhd,bhdv-bhv, q_phi, state) return out, state这段代码我特意拆成了两个函数一个做滑窗 softmax 注意力一个做全局线性压缩状态。真实推理时sliding_window_attention里每个 block 只取窗口范围内的 key可以用融合内核一次跑完linear_global_state则更适合在 decode 阶段逐步更新因为每来一个 token全局状态只是做了个累加而滑窗部分的 KV Cache 则是动态滚动的。如果你嫌手写麻烦也可以找现成的长上下文开源实现但有一个地方必须亲自动手确认线性注意力那一路有没有被 RoPE 位置编码影响。线性注意力的特征映射通常不天然考虑位置如果你直接把带 RoPE 的 Q/K 丢给线性注意力位置信息会被 kernel 映射冲掉。常见的补救方法是在线性路径上单独加位置向量或者只在滑窗路径上使用 RoPE。3.3 第三步设计长序列推理时的显存调度模型层数、注意力头数、窗口大小确定以后就可以算显存预算了。粗略公式是总 KV 显存 ≈ 层数 × (滑窗 token 数 全局 slot 数) × 每层 KV 维度 × 2 字节FP16假设 32 层、每层 KV 维度 4096、滑窗 8K、全局 256 个 slot那大约是 32 × 8,256 × 4096 × 2差不多 2.16GB 的 KV 主缓存。看上去很便宜对吧但别忘了如果 prefill 阶段一次性把整个 100 万 token 输入全部塞进 transformer 编码中间激活值依然会爆炸。所以必须配合 chunked prefill把长输入切成若干小段比如每段 64K 或 128K一段一段做前向计算逐段填充 KV Cache 和全局状态。另外KV Cache 在物理内存里不要按连续的长数组存。百万 token 下窗口是滚动淘汰的旧窗口如果不丢弃缓存还是会被撑满。更合理的做法是分页式 KV Cache将缓存拆成小块按需分配、按窗口淘汰。PyTorch 层面你可以自己实现一个循环缓冲区生产环境则推荐直接使用支持 page attention 的推理框架它们已经把这些调度细节封装好了。推理阶段的输出端也要注意decode 是逐 token 生成的每生成一个 token 就要更新一次全局线性状态和滑动窗口。如果你的实现把全局状态当作一个常数矩阵缓存在 GPU 显存里更新时要注意矩阵乘法顺序避免频繁申请临时显存。我在项目里会把全局状态固定在显存中每步原地累加decode 延迟就能稳定住。3.4 第四步从训练到推理的一致性还有一个很容易被忽视的细节训练时怎么建模推理时就要怎么建模至少在注意力模式上要尽量一致。如果预训练模型用的是标准全局注意力你直接把它换成稀疏-线性注意力推理时看到的长程信息分布会发生偏移输出质量会莫名其妙地下降。处理这个问题有三种路线。一是从头训练从头就用 SALA 注意力模型在学习过程中自然适应稀疏模式效果最稳但成本最高。二是剪枝蒸馏把已经训练好的稠密注意力模型作为教师让学生模型学习 SALA 下的注意力输出这种方法需要对特定层做蒸馏损失微调几天就能见效。三是混合微调在长文本数据上用 SALA 结构继续少量步数训练保留原来的语言能力同时让模型学会新的注意力模式。我踩过一个很典型的坑直接在推理时加载了一个稠密注意力权重然后只把 forward 里的注意力换成稀疏-线性没有做任何微调。结果短文本表现还好长文本一到 200K生成的答案就开始前后矛盾。后来老老实实跑了几百步蒸馏微调才把远程信息的一致性拉回来。4. 单卡百万 token 的吞吐压测怎么做4.1 吞吐量的正确口径prefill 和 decode 分开看很多人问百万 token 吞吐怎么样这个问题的口径特别重要。语言模型推理分两个阶段prefill 阶段一次性处理整个输入上下文decode 阶段逐个生成新 token。prefill 是大量矩阵并行计算吞吐通常很高decode 是访存密集速度远低于 prefill。所以压测时一定要把两个阶段分开记录。常用口径prefill 吞吐 输入 token 数 ÷ 编码整个输入的时间tokens/sdecode 吞吐 生成 token 数 ÷ 生成阶段总时间tokens/s端到端吞吐 输入 生成总 token 数 ÷ 总推理时间tokens/s我一般只报告 prefill 吞吐和 decode 吞吐两个指标因为端到端数字太容易被 max_new_tokens 设置干扰。比如你让模型从头生成 1 万个 token和只让它输出 100 个 token端到端吞吐差别巨大根本不可比。4.2 实测记录与显存增长趋势下面是一组我在某模拟项目里跑的示意数据不代表任何官方性能方向仅供参考。测试环境是一张 80GB 的 GPU模型层数和维度按上文第三节的假设滑窗 8K、全局 slot 256、chunked prefill 每块 64K。输入长度峰值显存prefill 吞吐decode 吞吐备注16K14.5GB35,200 tokens/s1,920 tokens/s基线容易跑128K23.7GB33,800 tokens/s1,830 tokens/s显存增长主要来自中间激活512K46.1GB21,600 tokens/s1,610 tokens/schunk 边界开销开始显现1M68.9GB17,400 tokens/s1,240 tokens/s接近单卡显存上限注意预留从这张表能看出几个规律prefill 吞吐随着 context 变长缓慢下降说明 chunked prefill 的合并和边界处理成本在增加decode 吞吐相对稳定说明主导解码速度的是 KV Cache 访存和模型逐 token 前向延迟。如果你测出来的 decode 吞吐掉得特别快大概率不是算力不够而是窗口淘汰策略导致缓存访问不连续或者全局状态更新时产生了等待。记录数据的同时建议同步用torch.profiler看一下每个算子的耗时占比重点看 Attention滑窗 softmax 部分、Linear State Update全局累积状态更新和 Memory Copy窗口移动三个位置。我实测超过 70% 的时候瓶颈不是计算本身而是 KV Cache 的搬运和 chunk 拼接。4.3 吞吐优化的小技巧关于把吞吐往上顶我有几个亲测有效的经验第一把 KV Cache 压成 FP8。长上下文模型对 KV 量化的容忍度通常好于权重量化窗口 8K 配 FP8显存立刻减半decode 吞吐能提升 15%~30%。代价是偶尔出现细节丢失需要在任务上评测能不能接受。第二decode 阶段用 CUDA Graph 减少内核启动开销。长上下文 decode 是逐 token 的小矩阵操作内核启动次数非常多CUDA Graph 可以一次性录制整个前向过程显著减少 CPU-GPU 通信延迟。我的实测里单卡 decode 吞吐因此提升接近 20%。第三对线性全局状态做定期的状态刷新。长期累积的 d×d 状态如果一直不除权早期 token 的贡献会逐渐被新 token 淹没。一个简单做法是引入衰减系数让旧状态的权重指数下降类似于滑动平均。这个细节特别关键它决定了百万长度下模型是否还能记住开头第一句话。5. 长上下文场景的常见坑与排查清单5.1 训练阶段的数值不稳定与注意力坍缩线性注意力的特征映射如果不加归一化容易出现方差爆炸或者梯度消失。尤其是特征映射里带 ReLU 的版本负半轴直接被截断信息大量丢失模型越训越哑巴。如果你发现 loss 断崖式下降后开始抖动先检查两个地方一是特征映射输出后面有没有 RMSNorm二是注意力温度系数是否合适。很多实现里建议在特征映射之后做一次 L2 归一化让内积保持在稳定范围。另一个问题是注意力坍缩滑窗注意力把所有概率几乎集中在最近的 1~2 个 token 上模型退化成了类似于 n-gram 的统计模型。这种情况通常是因为窗口太大、训练信号不足。可以适当减小窗口或者给滑窗注意力加一个小的均匀先验让没有重点时分散一些注意力质量。5.2 推理 OOM不是模型算力问题而是临时激活值我见过不少同学把窗口调得合理、KV Cache 算得也没超但 prefill 依然 OOM。问题多数出在临时激活值而不是 KV Cache。比如把 100 万 token 一次性过一遍模型每一层的中间张量都有 100 万×维度哪怕注意力做得再稀疏Feed-Forward 层的激活矩阵依然是百万量级。解决思路只有一条强制分块。除了 chunked prefill还可以考虑在实现里对 FFN 层做序列分块sequence parallelism 的思想或者干脆把梯度检查点打开。推理场景没有梯度但分块过 FFN的思路仍然适用。我一般把 chunk size 调成 64K 或 128K逐块跑完再把 KV 写回缓存峰值显存可以压下来很多。5.3 窗口外的信息召不回全局 slot 是最后的兜底如果模型在很长文档里总是想不起开头提到的关键实体第一个怀疑对象就是全局 slot 太少或者全局 slot 本身没有参与训练。全局 slot 本质上是一种强制注意力锚点它能不能起作用取决于模型是否真的学会往 slot 里写入重要信息。因此微调阶段要刻意构造跨段落的问题逼模型使用全局 slot 路径。另外开头的特殊 token 也值得多说一句。有经验的工程师都知道Transformer 对开头 token 有强烈的注意力倾向常被称为 attention sink。在稀疏设置下如果把开头的 token 淘汰出窗口模型会非常不适应。所以在实现里我总是把序列开头的几个 token 设置成永久保留的全局 token相当于给全文留了一把记忆锚点。这个技巧成本极低但对远程信息召回帮助非常大。5.4 性能反而变慢稀疏模式不规则是元凶还有一类问题是做了稀疏注意力、线性注意力结果速度比原来的稠密注意力还慢。这通常不是算法问题而是稀疏模式没有对齐硬件。GPU 擅长稠密连续的计算如果你让每个 query 访问的 key 范围都不一样每次都要动态索引内存那么访存延迟会吞掉所有算法收益。解决办法是把稀疏模式尽可能规则化。比如统一用固定窗口大小不用动态 top-k所有全局 slot 在显存里连续存储跨块跳线的模式尽量用块粒度对齐而不是 token 粒度。规则化后的注意力可以用分块矩阵乘高效执行性能才能体现出来。如果确实需要动态稀疏那就必须上自定义 CUDA kernel而不是在 PyTorch 层做循环。6. 写在最后真正决定成败的三个细节说了这么多最后想留几句个人经验。我在实际调 SALA 这类架构时最大的体会是这套东西算法门槛不算高工程门槛反而很高。所谓稀疏-线性注意力本质上就是一个注意力模式的组合设计但把它放到单卡百万 token 的跑法里真正决定成败的是三个细节。第一显存预算的分配。很多人只算 KV Cache忘了中间激活值和 chunk 边界开销。我建议在任何一次长文本压测之前先画一张显存流向图权重占多少、注意力中间占多少、FFN 激活占多少、KV Cache 占多少。哪个环节超了就逐个击破否则后面调试都是黑箱。第二全局 slot 的信息密度。全局 slot 不是越多越好而是要把它们训练成真正有意义的摘要卡槽。如果训练数据里没有跨距离的信息依赖slot 就算存在也学不出内容。想单卡吃下百万上下文必须配备高质量的长文本训练数据让模型有机会练习远程调用。第三稀疏策略与任务的匹配。同样是百万 token一轮多轮对话里当当流动读者和一口气读完整本手册回答问题对窗口、slot 数量、线性状态的需求完全不同。我最后悔的就是一开始照搬网上默认参数没有先分析数据分布导致调了两周才发现窗口开小了。如果你也想在这个方向试一把我的建议很简单先在一个 100K 左右的任务上把 SALA 跑通再逐步把 context 拉到 1M。每一步都对比显存、吞吐、召回三个指标确保没有静默丢失。等到你能熟练控制这三个变量再去追求更花哨的动态稀疏和混合核路会顺很多。