复杂度到线性注意力与长文本优化)
这次聊一个很直接的问题大模型处理 1 万字可能很快但上下文来到 100 万 token 时会慢到不现实。原因不在显卡显存够不够也不在代码写得好不好而是 Transformer 的“全注意力”机制在设计上就要付出 O(n²) 级别的代价。n 每翻 10 倍单次注意力的计算量大约要翻 100 倍。所以你可以这样理解模型每生成一个新词都要把前面所有历史重新核对一遍相似度——你每写新的一页它就把之前那上百万页记录重新翻一遍。这篇文章是 Kimi Linear 核心原理系列的第一篇先解决认知问题全注意力为什么贵。我会先给出标准注意力的复杂度推导解释自回归生成里的预填充和增量生成为什么成本不同再介绍线性注意力这一类方法的通用思路最后用一个可运行的模拟实验把 O(n²) 和 O(n) 的差距直观拉出来。这篇以原理为主不涉及模型 API 调用和批量任务队列那些更适合放到工程实践篇展开。建议的阅读人群包括关心长文本大模型怎么省算力的人、想理解线性注意力到底是什么的人、以及想搞懂 Kimi 为什么要在超长上下文场景做 Linear 优化的人。1. 核心结论速览全注意力为什么贵先给结论后面再展开推导。主题说明核心问题标准自注意力对序列里每个 token都要计算它与全部历史 token 的相关性计算复杂度单层单序列注意力约为 O(n²·d)n 是序列长度d 是特征维度长文本影响上下文从 1 万 token 到 100 万 token单次注意力计算量大约放大 10000 倍自回归生成无缓存时总成本可达 O(n³)有 KV Cache 后预填充 O(n²)增量生成每步仍要和全部历史做点积Kimi Linear面向超长上下文场景的线性复杂度注意力方向目标是减少“重翻历史”的重复成本读者收益看懂注意力复杂度、理解线性注意力思路、用代码验证复杂度增长趋势这里的“全注意力”也叫全局注意力或 Full Attention指的是 Transformer 标配的自注意力每一层的每个位置都能直接看到整段序列的所有位置。这个机制表达能力很强但代价也同步放大。2. 全注意力机制每个新词都要重翻历史2.1 注意力到底在做什么标准注意力的输入是三个矩阵Q查询、K键、V值。含义可以简化成当前这个 token 提出一个“我要找什么”的查询 Q然后和序列里所有历史 token 的 K 做相似度计算得到一组权重最后用权重去加权汇总 V 里的信息。数学上写作Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V这里QK^T就是核心。如果序列长度是 n特征维度是 d那么QK^T会形成一个 n × n 的矩阵。这个矩阵里的每一个元素都表示“第 i 个 token 对第 j 个 token 的关注程度”。n 越大这个矩阵越大计算量和内存压力都呈平方级上升。2.2 “重翻百万页记录”到底指什么用一个更贴近生活的比喻假设你在写一本手册手册里的每一页是一条历史记录。每次写下新的一页之前你都要把前面所有旧页翻一遍看哪些内容和当前内容相关然后把相关的内容划重点、做摘要。等你写到第 100 万页时每写新的一页就要把前面那 100 万页重新翻一遍。Transformer 生成文本的过程与此高度类似。模型在生成第 n 个 token 时会把前面 n-1 个 token 作为上下文输入。标准全注意力中新 token 的 query 要和之前每一个 token 的 key 做点积才能知道“应该参考谁”。所以历史越长新 token 需要核对的对象就越多。这不是工程上可以简单优化的“代码问题”而是注意力机制本身的算法特性。2.3 为什么不能只算一次然后复用结果这是初学者最容易问的问题既然前面历史的 K、V 都算过了能不能直接缓存最终结果后面不用再算实际上工程里已经有 KV Cache 这个优化了而且效果非常明显。它的做法是在生成阶段把已经算好的历史 K、V 缓存下来新 token 过来时不需要重新计算历史 token 的 K、V。但是新 token 的 Q 仍然要和所有历史 key 做点积这个操作没法省。原因在于新 query 和每个历史 key 的匹配结果取决于当前这一对组合任何历史 key 都不能被提前“打包成一个固定结论”。而且 Transformer 是很多层堆叠的每一层的中间表示都会随着新 token 加入而更新严格来说模型并不是把“一段固定摘要”拿来用而是每一层都要对完整序列重新建立两两关系。KV Cache 只是在工程上缓解了重复计算并没有改变“新词要接触全部历史”这个本质。3. 复杂度推导从 O(n²) 到 O(n³) 的成本增长3.1 单次前向的计算量只看一次前向过程标准注意力主要有两笔大计算QK^T形状是 (n, d) 乘 (d, n)乘法量约 n²·d。注意力矩阵再乘 V形状是 (n, n) 乘 (n, d)乘法量约 n²·d。所以单层单头注意力的计算量量级在 O(n²·d)。如果 d 是常数那核心就是 O(n²)。模型层数越多、注意力头越多这个成本还会继续翻倍。举个例子n 10k 时单层单头注意力矩阵的元素是 10 的 8 次方量级n 100k 时是 10 的 10 次方量级n 1M 时是 10 的 12 次方量级。这还只是注意力矩阵的元素个数不是最终需要的浮点运算次数。真实模型动辄几十层、几十个头再乘上 batch size压力会进一步放大。3.2 自回归生成里的重复成本大模型生成文本是自回归的也就是一个 token 一个 token 地往外蹦。朴素做法里每生成一个 token都要把之前所有 token 重新前向计算一遍。这样生成 n 个 token 的总成本会达到 O(n³)这是完全不可接受的。工程上 KV Cache 解决了其中一部分。具体来说预填充阶段对完整输入上下文做一次前向复杂度约 O(n²)。但这次计算把历史的 K、V 都缓存了下来。增量生成阶段每生成一个新 token只需要让新 query 和缓存里的历史 key 做点积。每步复杂度约 O(n)因为要和 n 个历史 key 分别计算。如果一共生成 n 个 token增量阶段总复杂度约 O(n²)。所以有 KV Cache 后整体生成成本从 O(n³) 降到了 O(n²)。这已经是很重要的优化。但“每个新词都要重翻百万页记录”的问题依然存在上下文越长增量生成的单步延迟和计算量就越高。这也解释了一个常见现象同一个模型输入从 1 万字涨到 10 万字后不仅首 token 延迟变高后面每生成一个 token 的速度也会肉眼可见变慢。因为每个新 token 都要面对更长的历史 key 列表。3.3 显存压力不只是注意力矩阵标准注意力还有一个问题是显存。QK^T形成的 n × n 矩阵如果完整放到显存里长度大时很快会爆。很多优化的核心思路之一就是不在显存里展开完整矩阵比如 FlashAttention 通过分块计算和 IO 优化减少了显存占用。但它优化的重点是读写带宽并没有改变 O(n²) 的计算量。也就是说FlashAttention 让注意力更快但不会让注意力从平方级变成线性级。另外KV Cache 的显存占用会随历史长度线性增长而且还要乘层数和头数。长上下文场景下Cache 可能在计算还没跑完之前就先吃掉大量显存。这也是长文本推理“贵”的重要组成部分。4. Kimi Linear 是什么把“重翻记录”改成“持续记账”4.1 长文本场景为什么是痛点Kimi 系列模型给很多人留下的印象是上下文很长能一次性处理几十万字甚至百万字量级的内容。这个能力很实用但背后压着一个残酷的成本问题标准全注意力的复杂度是 O(n²)上下文越长每次生成新词的“翻页”成本越高。如果模型要连续处理百万 token 级别的内容全注意力带来的延迟和显存压力会非常明显。所以长文本场景需要一种更省成本的注意力方案。Kimi Linear 的名字里直接带“Linear”从命名和公开信息来看它应当属于“把注意力成本从平方级往线性或近线性方向压”的技术路线。这里我做一个边界说明Kimi Linear 的具体内部实现我没有拿到官方技术文档因此本文不把它当成一个公式去拆而是介绍它所属的技术方向——线性注意力。这也是理解 Kimi Linear 的最底层基础。4.2 线性注意力的通用思路传统全注意力的麻烦在于softmax(QK^T)这一步里必须先构造 n × n 的矩阵再和 V 相乘。所以无论如何优化硬件的读写算法层面的平方复杂度都躲不开。线性注意力这类方法的通用思路是想办法改变运算顺序不构造 n × n 矩阵。核心做法是把 softmax 换成合适的核函数比如用 φ 表示某种非线性映射把注意力公式改写成φ(Q) (φ(K)^T V)关键在于括号里的φ(K)^T V可以先算得到一个与序列长度无关的状态矩阵。这样每来一个新的序列批次只需要在当前状态上继续累加而不是把整段历史翻出来重新两两计算。复杂度从 O(n²) 变成 O(n·d²)其中 n 是序列长度d 是特征维度。当 d 远小于 n 时这个差距会非常可观。这也对应标题里的比喻标准全注意力是“每翻一页重翻全部旧记录”线性注意力更像“边读边记账新内容来了只更新账本摘要”。新的 token 不再需要和百万条旧记录逐一比对而是直接使用已经维护好的压缩状态。4.3 对 Kimi Linear 的合理预期真正放到大规模模型里可能不会直接用第一节里的标准线性注意力公式。工业界的实现往往更复杂比如会结合分块计算、混合注意力、状态缓存、硬件算子优化等手段。我个人的判断是Kimi Linear 更应该理解为一套面向超长上下文推理的算法与工程组合方案目标是把“随长度平方增长”的成本压到接近线性而不是某一个孤立的数学技巧。因此阅读任何介绍 Kimi Linear 的文章时如果看到“Kimi 把 softmax 替换成某个函数”“Kimi 使用固定大小的状态压缩”这类描述要意识到这些大概率是对方向性原理的简述而不是对完整实现的精确刻画。更稳妥的理解是先掌握线性注意力的通用设计再把它作为背景去理解 Kimi 长文本场景里每一步工程取舍的来源。5. 全注意力与线性注意力核心对比维度标准全注意力线性注意力通用思路计算复杂度O(n²·d)O(n·d²)注意力矩阵需要构造或分块处理 n × n 矩阵通常可以避免构造 n × n 矩阵长上下文扩展成本随长度平方增长更适合超长序列精确检索能力softmax 可以做很强的“定点关注”常见实现的表达能力可能受限长距离聚焦弱一些实现成熟度高度成熟硬件算子、分布式方案完善需要处理数值稳定、训练稳定性、质量损失典型适合场景大多数通用大模型的预训练和推理长文本、流式生成、低资源场景5.1 为什么不能无条件替换看到线性注意力有复杂度优势很容易产生一个冲动全部换成线性注意力不就完了现实没有这么顺利原因是质量很难保。第一softmax 注意力擅长“尖锐选择”。它可以让模型在某一两个历史位置上给出非常高的权重这对指代消解、逻辑推理、代码定位等任务很关键。线性注意力如果只是简单做状态累积会把信息“抹平”长距离内真正关键的少数内容可能被稀释。第二模型训练方式会放大差异。现有大模型很多是在全注意力结构下完成预训练的如果推理阶段直接换成线性注意力训练和推理的行为不一致效果会明显退化。所以线性注意力方案通常要在设计阶段就进入模型架构或者采用足够复杂的混合策略而不是把已训好的模型拿来简单替换。第三工程上的线性注意力并不简单。归一化怎么做、数值稳定性怎么处理、状态怎么和分布式并行配合都是实际问题。这也是很多长文本优化方案最终会走“混合路线”的原因比如局部窗口注意力加全局稀疏注意力或局部窗口加压缩状态而不是用纯粹线性注意力包打天下。6. 可运行实验用代码观察 O(n²) 与 O(n) 的差距6.1 环境准备这一节只需要 Python 3.8 以上版本和 PyTorchCPU 就能跑不需要 GPU也不需要下载模型。建议先建立一个干净的虚拟环境python -m venv attn-demo source attn-demo/bin/activate # Windows 下是 attn-demo\Scripts\activate pip install torch --index-url https://download.pytorch.org/whl/cpu如果本机已经装好带 CUDA 的 PyTorch也可以直接用不影响实验原理。6.2 标准注意力的实现示意先实现一个最简单的标准注意力前向计算。这个版本只用来演示复杂度没有做 FlashAttention 之类的优化import torch import torch.nn.functional as F import time def standard_attention(q, k, v): # q, k, v 形状都是 (seq_len, dim) d q.shape[-1] scores torch.matmul(q, k.transpose(0, 1)) / (d ** 0.5) weights torch.softmax(scores, dim-1) return torch.matmul(weights, v)重点观察scores的形状它是(seq_len, seq_len)。序列长度翻倍这个矩阵的元素数量会变成原来的 4 倍。6.3 线性注意力的简化示意下面是线性注意力里“先累积状态再逐 token 查询”的通用简化版本。真实实现会有更严格的核函数设计和数值处理这里只是为了展示复杂度结构def phi(x): # 简单核函数ReLU return torch.relu(x) def linear_attention(q, k, v): # q, k, v 形状都是 (seq_len, dim) q phi(q) k phi(k) # 状态矩阵大小只和 dim 有关和 seq_len 无关 kv_state k.transpose(0, 1) v # (dim, dim) k_sum k.transpose(0, 1).sum(dim-1, keepdimTrue) # (dim, 1) numerator q kv_state # (seq_len, dim) denominator q k_sum 1e-6 # (seq_len, 1) return numerator / denominator这里的kv_state是核心它在累积整个序列的键值信息但维度固定。每当新 token 到来只需把它的 k、v 累加进kv_state不需要重新构造 n × n 注意力矩阵。6.4 测量不同长度下的耗时写一个简单的计时函数对比两种注意力在序列长度变化时的耗时def bench(seq_len, dim64, modestandard, repeats2): torch.manual_seed(0) q torch.randn(seq_len, dim) k torch.randn(seq_len, dim) v torch.randn(seq_len, dim) fn standard_attention if mode standard else linear_attention fn(q, k, v) # 预热 t0 time.time() for _ in range(repeats): fn(q, k, v) return (time.time() - t0) / repeats for n in [256, 512, 1024, 2048]: t_std bench(n, modestandard) t_lin bench(n, modelinear) print(fseq_len{n:5} standard{t_std*1000:8.2f} ms linear{t_lin*1000:8.2f} ms)这里我把repeats设为 2是担心 CPU 上执行较大序列时已经比较慢。如果你的机器性能很好可以继续把 n 加到 4096、8192趋势会更明显。需要提前说明这个实验在小序列下不保证线性注意力一定更快。因为线性注意力每步也要做 d × d 的矩阵运算常数不小。真正的差距要等 n 继续增长才会拉开标准注意力的耗时会更早进入明显加速上涨而线性注意力的增长曲线相对平滑。所以重点不是看某个点谁快而是看随 n 变大两条曲线的增长趋势差异。这里也要再次强调代码里的线性注意力只是帮助理解复杂度结构的示意实现不是 Kimi Linear 在真实模型里的实现也不代表实际模型的最终质量表现。7. 长文本推理的资源观察方法与优化方向如果你不是在跑简化实验而是在真实大模型上观察长文本成本可以从四个指标入手预填充耗时给定完整输入上下文后模型处理这段上下文要多久。增量生成延迟从生成第 1 个 token 到第 N 个 token每个 token 的平均延迟。GPU 显存峰值在本地部署场景下观察显存占用。吞吐量每秒能生成多少个 token。观察方法可以这样设计固定同一段 prompt 前缀把上下文从 1k、4k、16k、64k 逐级递增记录首 token 延迟和后续生成速度。如果本地部署 GPU可以边跑边用nvidia-smi -l 1观察显存变化nvidia-smi -l 1如果是 API 服务可以记录每次请求的输入 token 数、输出 token 数和总耗时对比不同上下文长度下的处理速度分布。长文本优化是目前很活跃的方向常见思路大致有几类KV Cache 压缩与淘汰不保留全部历史 key/value而是保留重要部分。稀疏注意力与滑动窗口每个 token 只关注局部窗口和少量全局锚点。线性注意力与状态空间模型用固定大小的压缩状态替代全部历史这就是 Kimi Linear 所属方向。FlashAttention 这类 IO 优化减少显存读写压力但不改变 O(n²) 计算量。上下文缓存让不同请求共享相同前缀的 KV 缓存减少重复预填充成本。这些方案可以组合使用实际产品中很少只依赖单一招式。理解这个背景后再看 Kimi Linear 的定位就会更清楚它到底想解决哪一环的瓶颈。8. 常见误区与排查思路关于“全注意力为什么贵”和“线性注意力是不是万能”我这里整理了几个常见问题供对照理解。误区 / 现象更准确的理解怎么验证长文本慢只是因为显存不够慢主要来自 O(n²) 计算和 KV 缓存读取显存不足是另一个限制用上文 bench 增加 n 观察耗时换 GPU 显存更大但计算复杂度不变每生成一个词都会把整段历史重新前向跑一遍有 KV Cache 后历史 K/V 不用重算但新 query 仍要和全部历史 key 做点积观察增量生成延迟是否随上下文长度上升Linear 等于零成本线性注意力需要处理质量下降、训练稳定性、归一化等问题用真实任务对比标准注意力和线性注意力的效果有了 FlashAttention 就不用关心平方复杂度FlashAttention 优化 IO不改变计算量本身固定 batch 和序列长度对比普通 attention 和 flash attention 在不同长度下的耗时曲线线性注意力一定在小序列也更快线性注意力常数大序列短时不一定有优势用第一节实验代码从 256 到 2048 甚至更大范围扫描如果你发现自己的复杂度推导结果和实际观察对不上从三个地方排查是不是漏掉了特征维度 d注意力复杂度是 O(n²·d)不是只看序列长度。是不是没有区分预填充和增量生成前者与输入长度强相关后者与已生成历史长度强相关。是不是把“单层单头”和“整模型”搞混了完整模型还要乘层数、头数和 batch size。这类排查思路对后续阅读线性注意力论文也很有用因为论文里经常用不同粒度汇报复杂度。9. 工程实践建议与下一步9.1 适合直接落地的几条建议先量化成本再决定优化方向。不要一上来就改架构。用一个真实长文本场景分别记录预填充耗时、增量生成延迟和显存峰值确认瓶颈到底在计算、显存还是带宽。长文本场景优先考虑已有工程优化。比如 FlashAttention、KV Cache、上下文缓存这些方案成熟度更高影响面更小。如果这些还不够再考虑线性注意力或稀疏注意力方向。线性注意力类方案不要贪心一次全替换。可以先从局部窗口加全局压缩状态这类混合结构开始逐步验证质量损失。纯线性注意力虽然复杂度漂亮但在复杂推理类任务上可能掉点。对于文本数据本身要遵守授权和隐私要求。实验数据不要随意使用未授权的长文档、聊天记录或私人信息需要商用时要确认数据来源和版权边界。9.2 下一步可以做什么第一步把上文 bench 脚本中的 n 继续拉大最好画一张横轴为序列长度、纵轴为耗时的 log-log 图。标准注意力曲线会更接近二次方上升线性注意力曲线会更接近一次方。这样能建立对复杂度最直观的体感。第二步找一篇线性注意力方向的经典论文或开源实现在小规模模型上做真实对比。重点看三项延迟、显存、质量。你会比只看文章更容易理解“线性注意力和全注意力不是零和替换而是工程权衡”。这篇是先收藏起来后面需要推导复杂度、对比线性注意力方案或者研究长文本推理优化时可以直接翻出来对照。Kimi Linear 后续如果讲到更具体的设计细节我会在这个系列里
这次聊一个很直接的问题大模型处理 1 万字可能很快但上下文来到 100 万 token 时会慢到不现实。原因不在显卡显存够不够也不在代码写得好不好而是 Transformer 的“全注意力”机制在设计上就要付出 O(n²) 级别的代价。n 每翻 10 倍单次注意力的计算量大约要翻 100 倍。所以你可以这样理解模型每生成一个新词都要把前面所有历史重新核对一遍相似度——你每写新的一页它就把之前那上百万页记录重新翻一遍。这篇文章是 Kimi Linear 核心原理系列的第一篇先解决认知问题全注意力为什么贵。我会先给出标准注意力的复杂度推导解释自回归生成里的预填充和增量生成为什么成本不同再介绍线性注意力这一类方法的通用思路最后用一个可运行的模拟实验把 O(n²) 和 O(n) 的差距直观拉出来。这篇以原理为主不涉及模型 API 调用和批量任务队列那些更适合放到工程实践篇展开。建议的阅读人群包括关心长文本大模型怎么省算力的人、想理解线性注意力到底是什么的人、以及想搞懂 Kimi 为什么要在超长上下文场景做 Linear 优化的人。1. 核心结论速览全注意力为什么贵先给结论后面再展开推导。主题说明核心问题标准自注意力对序列里每个 token都要计算它与全部历史 token 的相关性计算复杂度单层单序列注意力约为 O(n²·d)n 是序列长度d 是特征维度长文本影响上下文从 1 万 token 到 100 万 token单次注意力计算量大约放大 10000 倍自回归生成无缓存时总成本可达 O(n³)有 KV Cache 后预填充 O(n²)增量生成每步仍要和全部历史做点积Kimi Linear面向超长上下文场景的线性复杂度注意力方向目标是减少“重翻历史”的重复成本读者收益看懂注意力复杂度、理解线性注意力思路、用代码验证复杂度增长趋势这里的“全注意力”也叫全局注意力或 Full Attention指的是 Transformer 标配的自注意力每一层的每个位置都能直接看到整段序列的所有位置。这个机制表达能力很强但代价也同步放大。2. 全注意力机制每个新词都要重翻历史2.1 注意力到底在做什么标准注意力的输入是三个矩阵Q查询、K键、V值。含义可以简化成当前这个 token 提出一个“我要找什么”的查询 Q然后和序列里所有历史 token 的 K 做相似度计算得到一组权重最后用权重去加权汇总 V 里的信息。数学上写作Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V这里QK^T就是核心。如果序列长度是 n特征维度是 d那么QK^T会形成一个 n × n 的矩阵。这个矩阵里的每一个元素都表示“第 i 个 token 对第 j 个 token 的关注程度”。n 越大这个矩阵越大计算量和内存压力都呈平方级上升。2.2 “重翻百万页记录”到底指什么用一个更贴近生活的比喻假设你在写一本手册手册里的每一页是一条历史记录。每次写下新的一页之前你都要把前面所有旧页翻一遍看哪些内容和当前内容相关然后把相关的内容划重点、做摘要。等你写到第 100 万页时每写新的一页就要把前面那 100 万页重新翻一遍。Transformer 生成文本的过程与此高度类似。模型在生成第 n 个 token 时会把前面 n-1 个 token 作为上下文输入。标准全注意力中新 token 的 query 要和之前每一个 token 的 key 做点积才能知道“应该参考谁”。所以历史越长新 token 需要核对的对象就越多。这不是工程上可以简单优化的“代码问题”而是注意力机制本身的算法特性。2.3 为什么不能只算一次然后复用结果这是初学者最容易问的问题既然前面历史的 K、V 都算过了能不能直接缓存最终结果后面不用再算实际上工程里已经有 KV Cache 这个优化了而且效果非常明显。它的做法是在生成阶段把已经算好的历史 K、V 缓存下来新 token 过来时不需要重新计算历史 token 的 K、V。但是新 token 的 Q 仍然要和所有历史 key 做点积这个操作没法省。原因在于新 query 和每个历史 key 的匹配结果取决于当前这一对组合任何历史 key 都不能被提前“打包成一个固定结论”。而且 Transformer 是很多层堆叠的每一层的中间表示都会随着新 token 加入而更新严格来说模型并不是把“一段固定摘要”拿来用而是每一层都要对完整序列重新建立两两关系。KV Cache 只是在工程上缓解了重复计算并没有改变“新词要接触全部历史”这个本质。3. 复杂度推导从 O(n²) 到 O(n³) 的成本增长3.1 单次前向的计算量只看一次前向过程标准注意力主要有两笔大计算QK^T形状是 (n, d) 乘 (d, n)乘法量约 n²·d。注意力矩阵再乘 V形状是 (n, n) 乘 (n, d)乘法量约 n²·d。所以单层单头注意力的计算量量级在 O(n²·d)。如果 d 是常数那核心就是 O(n²)。模型层数越多、注意力头越多这个成本还会继续翻倍。举个例子n 10k 时单层单头注意力矩阵的元素是 10 的 8 次方量级n 100k 时是 10 的 10 次方量级n 1M 时是 10 的 12 次方量级。这还只是注意力矩阵的元素个数不是最终需要的浮点运算次数。真实模型动辄几十层、几十个头再乘上 batch size压力会进一步放大。3.2 自回归生成里的重复成本大模型生成文本是自回归的也就是一个 token 一个 token 地往外蹦。朴素做法里每生成一个 token都要把之前所有 token 重新前向计算一遍。这样生成 n 个 token 的总成本会达到 O(n³)这是完全不可接受的。工程上 KV Cache 解决了其中一部分。具体来说预填充阶段对完整输入上下文做一次前向复杂度约 O(n²)。但这次计算把历史的 K、V 都缓存了下来。增量生成阶段每生成一个新 token只需要让新 query 和缓存里的历史 key 做点积。每步复杂度约 O(n)因为要和 n 个历史 key 分别计算。如果一共生成 n 个 token增量阶段总复杂度约 O(n²)。所以有 KV Cache 后整体生成成本从 O(n³) 降到了 O(n²)。这已经是很重要的优化。但“每个新词都要重翻百万页记录”的问题依然存在上下文越长增量生成的单步延迟和计算量就越高。这也解释了一个常见现象同一个模型输入从 1 万字涨到 10 万字后不仅首 token 延迟变高后面每生成一个 token 的速度也会肉眼可见变慢。因为每个新 token 都要面对更长的历史 key 列表。3.3 显存压力不只是注意力矩阵标准注意力还有一个问题是显存。QK^T形成的 n × n 矩阵如果完整放到显存里长度大时很快会爆。很多优化的核心思路之一就是不在显存里展开完整矩阵比如 FlashAttention 通过分块计算和 IO 优化减少了显存占用。但它优化的重点是读写带宽并没有改变 O(n²) 的计算量。也就是说FlashAttention 让注意力更快但不会让注意力从平方级变成线性级。另外KV Cache 的显存占用会随历史长度线性增长而且还要乘层数和头数。长上下文场景下Cache 可能在计算还没跑完之前就先吃掉大量显存。这也是长文本推理“贵”的重要组成部分。4. Kimi Linear 是什么把“重翻记录”改成“持续记账”4.1 长文本场景为什么是痛点Kimi 系列模型给很多人留下的印象是上下文很长能一次性处理几十万字甚至百万字量级的内容。这个能力很实用但背后压着一个残酷的成本问题标准全注意力的复杂度是 O(n²)上下文越长每次生成新词的“翻页”成本越高。如果模型要连续处理百万 token 级别的内容全注意力带来的延迟和显存压力会非常明显。所以长文本场景需要一种更省成本的注意力方案。Kimi Linear 的名字里直接带“Linear”从命名和公开信息来看它应当属于“把注意力成本从平方级往线性或近线性方向压”的技术路线。这里我做一个边界说明Kimi Linear 的具体内部实现我没有拿到官方技术文档因此本文不把它当成一个公式去拆而是介绍它所属的技术方向——线性注意力。这也是理解 Kimi Linear 的最底层基础。4.2 线性注意力的通用思路传统全注意力的麻烦在于softmax(QK^T)这一步里必须先构造 n × n 的矩阵再和 V 相乘。所以无论如何优化硬件的读写算法层面的平方复杂度都躲不开。线性注意力这类方法的通用思路是想办法改变运算顺序不构造 n × n 矩阵。核心做法是把 softmax 换成合适的核函数比如用 φ 表示某种非线性映射把注意力公式改写成φ(Q) (φ(K)^T V)关键在于括号里的φ(K)^T V可以先算得到一个与序列长度无关的状态矩阵。这样每来一个新的序列批次只需要在当前状态上继续累加而不是把整段历史翻出来重新两两计算。复杂度从 O(n²) 变成 O(n·d²)其中 n 是序列长度d 是特征维度。当 d 远小于 n 时这个差距会非常可观。这也对应标题里的比喻标准全注意力是“每翻一页重翻全部旧记录”线性注意力更像“边读边记账新内容来了只更新账本摘要”。新的 token 不再需要和百万条旧记录逐一比对而是直接使用已经维护好的压缩状态。4.3 对 Kimi Linear 的合理预期真正放到大规模模型里可能不会直接用第一节里的标准线性注意力公式。工业界的实现往往更复杂比如会结合分块计算、混合注意力、状态缓存、硬件算子优化等手段。我个人的判断是Kimi Linear 更应该理解为一套面向超长上下文推理的算法与工程组合方案目标是把“随长度平方增长”的成本压到接近线性而不是某一个孤立的数学技巧。因此阅读任何介绍 Kimi Linear 的文章时如果看到“Kimi 把 softmax 替换成某个函数”“Kimi 使用固定大小的状态压缩”这类描述要意识到这些大概率是对方向性原理的简述而不是对完整实现的精确刻画。更稳妥的理解是先掌握线性注意力的通用设计再把它作为背景去理解 Kimi 长文本场景里每一步工程取舍的来源。5. 全注意力与线性注意力核心对比维度标准全注意力线性注意力通用思路计算复杂度O(n²·d)O(n·d²)注意力矩阵需要构造或分块处理 n × n 矩阵通常可以避免构造 n × n 矩阵长上下文扩展成本随长度平方增长更适合超长序列精确检索能力softmax 可以做很强的“定点关注”常见实现的表达能力可能受限长距离聚焦弱一些实现成熟度高度成熟硬件算子、分布式方案完善需要处理数值稳定、训练稳定性、质量损失典型适合场景大多数通用大模型的预训练和推理长文本、流式生成、低资源场景5.1 为什么不能无条件替换看到线性注意力有复杂度优势很容易产生一个冲动全部换成线性注意力不就完了现实没有这么顺利原因是质量很难保。第一softmax 注意力擅长“尖锐选择”。它可以让模型在某一两个历史位置上给出非常高的权重这对指代消解、逻辑推理、代码定位等任务很关键。线性注意力如果只是简单做状态累积会把信息“抹平”长距离内真正关键的少数内容可能被稀释。第二模型训练方式会放大差异。现有大模型很多是在全注意力结构下完成预训练的如果推理阶段直接换成线性注意力训练和推理的行为不一致效果会明显退化。所以线性注意力方案通常要在设计阶段就进入模型架构或者采用足够复杂的混合策略而不是把已训好的模型拿来简单替换。第三工程上的线性注意力并不简单。归一化怎么做、数值稳定性怎么处理、状态怎么和分布式并行配合都是实际问题。这也是很多长文本优化方案最终会走“混合路线”的原因比如局部窗口注意力加全局稀疏注意力或局部窗口加压缩状态而不是用纯粹线性注意力包打天下。6. 可运行实验用代码观察 O(n²) 与 O(n) 的差距6.1 环境准备这一节只需要 Python 3.8 以上版本和 PyTorchCPU 就能跑不需要 GPU也不需要下载模型。建议先建立一个干净的虚拟环境python -m venv attn-demo source attn-demo/bin/activate # Windows 下是 attn-demo\Scripts\activate pip install torch --index-url https://download.pytorch.org/whl/cpu如果本机已经装好带 CUDA 的 PyTorch也可以直接用不影响实验原理。6.2 标准注意力的实现示意先实现一个最简单的标准注意力前向计算。这个版本只用来演示复杂度没有做 FlashAttention 之类的优化import torch import torch.nn.functional as F import time def standard_attention(q, k, v): # q, k, v 形状都是 (seq_len, dim) d q.shape[-1] scores torch.matmul(q, k.transpose(0, 1)) / (d ** 0.5) weights torch.softmax(scores, dim-1) return torch.matmul(weights, v)重点观察scores的形状它是(seq_len, seq_len)。序列长度翻倍这个矩阵的元素数量会变成原来的 4 倍。6.3 线性注意力的简化示意下面是线性注意力里“先累积状态再逐 token 查询”的通用简化版本。真实实现会有更严格的核函数设计和数值处理这里只是为了展示复杂度结构def phi(x): # 简单核函数ReLU return torch.relu(x) def linear_attention(q, k, v): # q, k, v 形状都是 (seq_len, dim) q phi(q) k phi(k) # 状态矩阵大小只和 dim 有关和 seq_len 无关 kv_state k.transpose(0, 1) v # (dim, dim) k_sum k.transpose(0, 1).sum(dim-1, keepdimTrue) # (dim, 1) numerator q kv_state # (seq_len, dim) denominator q k_sum 1e-6 # (seq_len, 1) return numerator / denominator这里的kv_state是核心它在累积整个序列的键值信息但维度固定。每当新 token 到来只需把它的 k、v 累加进kv_state不需要重新构造 n × n 注意力矩阵。6.4 测量不同长度下的耗时写一个简单的计时函数对比两种注意力在序列长度变化时的耗时def bench(seq_len, dim64, modestandard, repeats2): torch.manual_seed(0) q torch.randn(seq_len, dim) k torch.randn(seq_len, dim) v torch.randn(seq_len, dim) fn standard_attention if mode standard else linear_attention fn(q, k, v) # 预热 t0 time.time() for _ in range(repeats): fn(q, k, v) return (time.time() - t0) / repeats for n in [256, 512, 1024, 2048]: t_std bench(n, modestandard) t_lin bench(n, modelinear) print(fseq_len{n:5} standard{t_std*1000:8.2f} ms linear{t_lin*1000:8.2f} ms)这里我把repeats设为 2是担心 CPU 上执行较大序列时已经比较慢。如果你的机器性能很好可以继续把 n 加到 4096、8192趋势会更明显。需要提前说明这个实验在小序列下不保证线性注意力一定更快。因为线性注意力每步也要做 d × d 的矩阵运算常数不小。真正的差距要等 n 继续增长才会拉开标准注意力的耗时会更早进入明显加速上涨而线性注意力的增长曲线相对平滑。所以重点不是看某个点谁快而是看随 n 变大两条曲线的增长趋势差异。这里也要再次强调代码里的线性注意力只是帮助理解复杂度结构的示意实现不是 Kimi Linear 在真实模型里的实现也不代表实际模型的最终质量表现。7. 长文本推理的资源观察方法与优化方向如果你不是在跑简化实验而是在真实大模型上观察长文本成本可以从四个指标入手预填充耗时给定完整输入上下文后模型处理这段上下文要多久。增量生成延迟从生成第 1 个 token 到第 N 个 token每个 token 的平均延迟。GPU 显存峰值在本地部署场景下观察显存占用。吞吐量每秒能生成多少个 token。观察方法可以这样设计固定同一段 prompt 前缀把上下文从 1k、4k、16k、64k 逐级递增记录首 token 延迟和后续生成速度。如果本地部署 GPU可以边跑边用nvidia-smi -l 1观察显存变化nvidia-smi -l 1如果是 API 服务可以记录每次请求的输入 token 数、输出 token 数和总耗时对比不同上下文长度下的处理速度分布。长文本优化是目前很活跃的方向常见思路大致有几类KV Cache 压缩与淘汰不保留全部历史 key/value而是保留重要部分。稀疏注意力与滑动窗口每个 token 只关注局部窗口和少量全局锚点。线性注意力与状态空间模型用固定大小的压缩状态替代全部历史这就是 Kimi Linear 所属方向。FlashAttention 这类 IO 优化减少显存读写压力但不改变 O(n²) 计算量。上下文缓存让不同请求共享相同前缀的 KV 缓存减少重复预填充成本。这些方案可以组合使用实际产品中很少只依赖单一招式。理解这个背景后再看 Kimi Linear 的定位就会更清楚它到底想解决哪一环的瓶颈。8. 常见误区与排查思路关于“全注意力为什么贵”和“线性注意力是不是万能”我这里整理了几个常见问题供对照理解。误区 / 现象更准确的理解怎么验证长文本慢只是因为显存不够慢主要来自 O(n²) 计算和 KV 缓存读取显存不足是另一个限制用上文 bench 增加 n 观察耗时换 GPU 显存更大但计算复杂度不变每生成一个词都会把整段历史重新前向跑一遍有 KV Cache 后历史 K/V 不用重算但新 query 仍要和全部历史 key 做点积观察增量生成延迟是否随上下文长度上升Linear 等于零成本线性注意力需要处理质量下降、训练稳定性、归一化等问题用真实任务对比标准注意力和线性注意力的效果有了 FlashAttention 就不用关心平方复杂度FlashAttention 优化 IO不改变计算量本身固定 batch 和序列长度对比普通 attention 和 flash attention 在不同长度下的耗时曲线线性注意力一定在小序列也更快线性注意力常数大序列短时不一定有优势用第一节实验代码从 256 到 2048 甚至更大范围扫描如果你发现自己的复杂度推导结果和实际观察对不上从三个地方排查是不是漏掉了特征维度 d注意力复杂度是 O(n²·d)不是只看序列长度。是不是没有区分预填充和增量生成前者与输入长度强相关后者与已生成历史长度强相关。是不是把“单层单头”和“整模型”搞混了完整模型还要乘层数、头数和 batch size。这类排查思路对后续阅读线性注意力论文也很有用因为论文里经常用不同粒度汇报复杂度。9. 工程实践建议与下一步9.1 适合直接落地的几条建议先量化成本再决定优化方向。不要一上来就改架构。用一个真实长文本场景分别记录预填充耗时、增量生成延迟和显存峰值确认瓶颈到底在计算、显存还是带宽。长文本场景优先考虑已有工程优化。比如 FlashAttention、KV Cache、上下文缓存这些方案成熟度更高影响面更小。如果这些还不够再考虑线性注意力或稀疏注意力方向。线性注意力类方案不要贪心一次全替换。可以先从局部窗口加全局压缩状态这类混合结构开始逐步验证质量损失。纯线性注意力虽然复杂度漂亮但在复杂推理类任务上可能掉点。对于文本数据本身要遵守授权和隐私要求。实验数据不要随意使用未授权的长文档、聊天记录或私人信息需要商用时要确认数据来源和版权边界。9.2 下一步可以做什么第一步把上文 bench 脚本中的 n 继续拉大最好画一张横轴为序列长度、纵轴为耗时的 log-log 图。标准注意力曲线会更接近二次方上升线性注意力曲线会更接近一次方。这样能建立对复杂度最直观的体感。第二步找一篇线性注意力方向的经典论文或开源实现在小规模模型上做真实对比。重点看三项延迟、显存、质量。你会比只看文章更容易理解“线性注意力和全注意力不是零和替换而是工程权衡”。这篇是先收藏起来后面需要推导复杂度、对比线性注意力方案或者研究长文本推理优化时可以直接翻出来对照。Kimi Linear 后续如果讲到更具体的设计细节我会在这个系列里
