线性注意力模型如何实现固定推理状态

大多数大语言模型（LLM）采用带有 softmax attention 的 transformer 架构。在推理过程中，模型的键值（KV）缓存会随着对话或文档上下文的增长而增大。这可能降低并发能力、增加服务成本，并迫使 Agent 应用丢弃有用的历史信息。

i14 为此设计了采用 linear attention 的替代方案 Laplace 3B。它的循环状态始终保持在每条序列 29.5 MB，不随上下文长度变化。通过与 AWS 合作，i14 利用 TorchNeuron Native 提供的原生 PyTorch 支持，将这一自定义架构迁移到了 AWS Trainium2，同时保留了 Gated Linear Attention（GLA）设计和 Muon optimizer。

在一台 trn2.48xlarge 实例上，Laplace 3B 的训练吞吐量约为每秒 120,000 tokens。推理服务在同时处理 1,280 条序列时，输出吞吐量达到每秒 11,544 tokens。本文将介绍模型设计、i14 为适配 Trainium2 所做的改动，以及训练和推理结果。

i14 是一家 AI 研究与技术公司，开发模型架构、组件和运行时。其工作借鉴了 AI、计算认知科学，以及生物系统如何随时间持续学习的相关研究。i14 将这些思路应用到模型中，使模型能够保留长距离上下文，同时避免推理服务状态随每个 token 不断增长。

为什么选择 linear attention 模型？

Transformer 的 softmax attention 会对上下文中的 token 进行相互比较。在训练阶段，这部分计算量随序列长度呈平方增长。在推理阶段，KV cache 则会随着模型处理的 token 增多而持续增长。

Linear attention 围绕一个逐 token 更新的循环状态重新组织计算。Laplace 3B 使用 GLA，将之前的上下文压缩到一个大小固定的矩阵状态中。这将训练时 attention 的复杂度从 O(T²) 降至 O(T)，并将推理状态固定在每条序列 29.5 MB。

Linear attention 会压缩历史信息，而不是保留每一个 key 和 value。因此，你需要针对自己的工作负载评估检索质量。对于长时间运行的 Agent、大型文档分析，以及需要持久记忆且希望服务资源需求可预测的应用，固定大小的状态可能带来优势。

i14 如何为 AWS Trainium2 适配 Laplace 3B

Laplace 3B 最初的实现使用了针对 GPU 优化的 Triton kernel。AWS Generative AI Innovation Center（GenAIIC）团队首先使用标准 PyTorch 算子，为 Trainium2 重写了 GLA kernel。在比较两种方案后，团队选择了 beta 版 TorchNeuron Native，也就是 AWS Trainium 的原生 PyTorch 后端，因为它的性能更好。

这次实现需要进行三项改动：

• 使用原生 PyTorch 实现 GLA attention kernel。 团队用标准 PyTorch 算子替换了 GPU 专用的 fused_chunk_simple_gla Triton kernel，采用 16 个 chunk、每个 chunk 128 tokens 的划分方式，以匹配 NeuronCore 片上缓冲区的 128 个分区。
• 重新设计 Muon 的通信方式。 MuonFSDP2 中的两次 all_to_all 集合通信将每个 rank 的通信量降至原来的 1/24，并将吞吐量提升至原来的 1.63 倍。
• 编译顺序执行的反向传播。 团队将 GLA 递推的 16 个 chunk 步骤静态展开，以满足编译器对计算图大小和片上内存的限制。

后面的工程挑战部分会详细介绍每项改动。

为什么选择 AWS Trainium 加速器？

AWS Trainium 是 AWS 推出的一系列机器学习（ML）加速器，专为模型训练和推理设计。一台 trn2.48xlarge 实例配备 16 颗 Trainium2 芯片、64 个逻辑 NeuronCore、高带宽内存，以及 NeuronLink-v3 互连。这些资源支撑了本文介绍的分布式训练和推理配置。

Trainium2 的三个特性决定了实现方式：由软件管理的片上内存（挑战 1）、高速芯片间互连（挑战 2），以及静态编译器（挑战 3）。编译器会将 GLA 递推追踪为固定的计算图。

由软件管理的片上内存。 每个 NeuronCore 都有一个 28 MiB、包含 128 个分区的 State Buffer（SBUF），以及一个 2 MiB 的部分和缓冲区（PSUM）。PSUM 直接连接到一个 128×128 的脉动阵列 Tensor Engine，旁边还有独立的向量和标量计算单元。SBUF 并不是 cache：编译器会提前安排每一次从 HBM 到 SBUF 的数据传输，确保数据及时就位，不会发生缓存驱逐或未命中（图 1）。

图 1：NeuronCore-v3 内部架构。包括由软件管理的 SBUF（28 MiB，128 个分区）、Tensor Engine（128×128 脉动阵列）、Vector Engine、Scalar Engine 和 GPSIMD 核心。

i14 与 Laplace 3B 架构

Laplace 3B 是 i14 的一个 24 层 decoder，每一层都用 Simple GLA 替代了 softmax attention。图 2 展示了单层的数据流：从顶部的输入 xₜ，到底部的输出 xₜ₊₁。本节将逐一介绍归一化与卷积、QKVG 投影、GLA mixer，以及前馈 channel mixer。

归一化与卷积。 每个 token 的残差流向量 xₜ ∈ ℝᴰ（D = 3,072）先经过 LayerNorm，再经过一个宽度 k = 4、带有 SiLU 的逐通道因果 Conv1d。这会在进入 mixer 前，为每个位置提供一个 4-gram 摘要。随后，GLA mixer 通过带衰减的状态，对整个历史执行基于内容的检索。

QKVG 投影。 卷积后的输出经过一个融合矩阵 Wqkvg ∈ ℝ^(9,216 × 3,072)，其中 9,216 = H × (2dₖ + 2dᵥ)。结果被拆分为每个 token 的四路数据流：query qₜ、key kₜ、value vₜ 和 gate gₜ（H = 16，dₖ = 96，dᵥ = 192）。gate 在 mixer 之后使用。

图 2：Laplace decoder 层架构。绿色表示固定模块（无参数），灰色表示可学习模块，黄色表示循环状态，蓝色表示投影。

GLA mixer。 每个 head h 都维护一个大小为 dₖ × dᵥ（96 × 192）的矩阵状态 Sₜ⁽ʰ⁾，以压缩且逐渐衰减的形式记录过去的 key-value 对。每处理一个 token，状态都会按以下方式更新：

GLA 循环状态更新公式

GLA 状态读出公式

GLA 自增强公式

衰减方案。 每个 head 的衰减率 λₕ 决定了它遗忘历史 token 的速度。该衰减率在初始化时固定，不参与学习，并按照几何规律设置。这一方案为 16 个 head 分配了不同的记忆长度，并定义了每个 head 的对数衰减率 γₕ：

GLA 衰减方案公式

Channel mixer（FFN）。 在 attention 模块及其残差连接之后，第二个 LayerNorm 将数据送入内部维度为 8,192 的 GLU（SiLU）FFN。它包含三个矩阵：两个升维投影 W₁、W₂ ∈ ℝ^(8,192 × D)，以及一个降维投影 W₃ ∈ ℝ^(D × 8,192)。之后再接第二个残差连接。

无论序列多长，每个 head 的状态 S⁽ʰ⁾ 都固定为 dₖ × dᵥ = 96 × 192。因此，推理时每个 token 的复杂度为 O(1)，训练复杂度为 O(T)，而不是 O(T²)。

工程挑战

这三项改动各自带来了不同的工程挑战。

挑战 1：在 Trainium 上使用原生 PyTorch 实现 GLA

在 GPU 上，GLA 通过 flash-linear-attention 库中的融合 Triton kernel fused_chunk_simple_gla 运行。GenAIIC 团队用标准 PyTorch 算子重新表达了这套计算，通过 TorchNeuron Native 执行，并显式循环处理大小固定的 chunk。chunk 数量在编译时固定，同时决定了反向传播的展开深度。团队采用 fla_num_chunks = 16，因此每个 chunk 包含 C = 128 tokens。

在每个 chunk c 内，chunk 内部的输出由带掩码的 linear attention 计算，持续维护的状态 S 则补充更早的上下文：

分块 GLA 的 chunk 内输出

随后，状态推进到下一个 chunk：

分块 GLA 的状态更新

这里的 λᶜ 表示每个 head 的衰减率在整个 chunk 上的累积效果。16 个 chunk 的内部计算并行执行，只有状态更新需要顺序执行（图 3）。

图 3：GLA 分块递推，包含并行的 chunk 内 attention，以及顺序执行的状态链。

由于 SBUF 有 128 个分区，每个包含 128 tokens 的 chunk 可以让每个 token 对应一个分区，无需填充，也无需跨分区移动数据。这种对齐方式使实测吞吐量提高了 14%（图 4）。

图 4：chunk 大小 C = 128，与 SBUF 的 128 个分区一一对应。

挑战 2：FSDP2 下的分布式 Muon 正交化

Muon 会在更新前对每个梯度进行正交化。对于 96 个二维权重矩阵中的每一个，总计 27.18 亿参数，它都会执行五次 Newton–Schulz 迭代（算法 1）。

Newton–Schulz 算法

分布式场景的问题。 Gram 矩阵 A = XX⊤ 计算权重矩阵中每一行与其他所有行之间的相关性。在 Fully Sharded Data Parallel（FSDP2）下，64 个 rank 中的每一个都只持有 1/64 的行。因此，仅在本地分片上计算 XX⊤，只能得到一个对角块，所有跨分片项都会丢失（图 5）。

图 5：FSDP2 下的 Newton–Schulz：分片本地的 Gram 矩阵块，以及 gather/broadcast 与 all_to_all 的对比。

朴素方案：Gather 后再 broadcast。 每个 rank 都收集完整梯度（5.44 GB），由一个 rank 执行 Newton–Schulz，然后广播结果。这种方式每步需要传输 10.87 GB 数据，而另外 63 个 rank 都在等待。

生产方案：通过 all_to_all 转置分布。 在 MuonFSDP2 中，第一次 all_to_all 让每个 rank 获得少量完整矩阵，随后全部 64 个 rank 并行执行 Newton–Schulz，再通过第二次 all_to_all 将结果送回。每个 rank 的通信量降至 0.45 GB，只有原来的 1/24。吞吐量从每秒 73,378 tokens 提升至 119,934 tokens，达到原来的 1.63 倍，且 loss 曲线逐位一致（图 6）。

图 6：每步通信量：使用 all_gather 时为 10.87 GB，使用 all_to_all 时为 0.45 GB。

通信成本。 Newton–Schulz 只占单个训练步骤 FLOPs 的约 2%（总计 9.53 × 10¹⁵，其中它占 1.95 × 10¹⁴），因此瓶颈在通信。在带宽为 1.28 TB/s 的 NeuronLink-v3 上，交换 0.45 GB 数据耗时不到 0.4 ms，而整个训练步骤约为 4.4 秒。

挑战 3：为顺序递推展开反向传播

反向传播需要对全部 16 次顺序状态更新求导。在 GPU 上，autograd 会在运行时构建这张计算图。在 Trainium2 上，编译器必须提前看到完整展开的 16 步计算图，这带来了三项约束。

计算图深度会随 chunk 数量增长：当使用 32 个 chunk（C = 64）时，计算图超出了调度器的容量。每个 chunk 的内存占用会随 chunk 大小增长：当使用 8 个 chunk（C = 256）时，中间结果溢出到 HBM。所有形状都在追踪时固定，因此循环中的每个中间结果都需要单独的静态内存分配。

16 个 chunk（C = 128）同时满足了这三项约束：展开后的计算图没有超出调度器容量，中间结果可以留在片上，计算分块也能与 SBUF 的 128 个分区对齐。

GenAIIC 团队在一台 trn2.48xlarge 实例上训练 Laplace 3B，使用全部 64 个逻辑 NeuronCore，以及 Muon 的 all_to_all 实现路径。完整训练方案包含 190,735 步，对应 1,000 亿 FineWeb-Edu tokens。以下结果截至第 84,210 步，约处理了 440 亿 tokens。

吞吐量。 编译预热结束后，吞吐量稳定在每秒约 120,000 tokens，并一直保持到第 84,210 步。

Loss 与梯度稳定性。 Loss（图 7a）从约 11.4 平滑下降，在第 84,210 步达到 2.3729。梯度范数（图 7b）在训练早期出现预期中的峰值，随后稳步降至约 0.05；预热结束后，没有出现尖峰、发散或 NaN/Inf 值。

图 7：训练动态。（a）Loss 从 11.4 降至 2.3729。（b）根据 TensorBoard 数据绘制的梯度范数。

Trainium2 上的推理优化

GenAIIC 团队通过未经修改的 vLLM 引擎，在一台 trn2.48xlarge 实例上同时处理 1,280 条序列，实现了每秒 11,544 个输出 tokens。

固定大小的服务状态。 每条序列的循环状态为 29.5 MB，约为 8,192 tokens 上下文下同类 KV cache 大小的 1/82。将 prompt 从 2,048 tokens 增加到 8,192 tokens，每个输出 token 的成本仅增加 0.5%。

未经修改的 vLLM。 一个 wrapper 会替换每个 decoder 层中的一个属性，也就是 token mixer，以便在 vLLM 池中读写循环状态。由于模型将自身声明为 attention-free，vLLM 会为每条序列分配一个固定 page。输出结果与 i14 参考实现逐 token 一致：在 BF16 下，3/3 个 prompt 的 48/48 个步骤全部匹配。

定位瓶颈。 将 decode 时间相对于 decoder 层数 L 进行线性拟合，可得：

扩展到整台实例。 在不使用 tensor parallelism 的情况下，通过增加副本扩展吞吐量，每个逻辑 NeuronCore 运行一个引擎，处理 20 条序列。Eager 模式的吞吐量从单副本的每秒 133 个输出 tokens，增长到 64 个副本时的 6,772，此时主机 CPU 使用率达到 82%。编译 decode 步骤消除了逐算子分发开销，将吞吐量提升至原来的 1.7 倍，达到 11,544（图 8）。

图 8：整台实例的推理吞吐量。（a）Eager 与编译后 decode 随副本数量增加的扩展表现。（b）编译 decode 将吞吐量从每秒 6,772 个输出 tokens 提升到 11,544。

固定 batch 基准测试掩盖了什么。 在真实流量下，每一种新的 decode batch 宽度都是需要单独编译的形状，首次使用的成本约为 11 秒。这些步骤仅占总步骤数的 1.1%，却消耗了 47.5% 的 decode 时间。通过填充到固定宽度，可以消除这项成本，并保持结果逐位一致。在请求交错到达的场景下，编译 decode 将每个副本的吞吐量提升至原来的 3.52 倍，而固定 batch 下仅为 1.46 倍。由于 BF16 运算重新结合会使编译后的输出与 eager 输出略有差异，当严格可复现性很重要时，应使用 eager decode。

这个项目在训练和推理两方面都带来了一些经验：

• 原生 PyTorch 可以在 Trainium2 上支持自定义架构。 i14 保留了 GLA 架构，并使用 TorchNeuron Native 完成最终实现。
• 对于自定义 optimizer，通信设计很关键。 Muon 的 all_to_all 设计将每个 rank 的通信量降至原来的 1/24，并将训练吞吐量提升至原来的 1.63 倍。
• 固定状态让推理内存需求可预测。 无论上下文多长，每条序列都使用 29.5 MB。
• 先分析性能，再编译。 占据大部分 decode 时间的是算子分发，而不是算术计算或 HBM 带宽。
• 用真实的请求到达模式做基准测试。 固定 batch 掩盖了 batch 宽度变化时编译新形状的成本。

Laplace 3B 表明，AWS Trainium2 能够支持标准 transformer 之外的模型架构。借助 TorchNeuron Native 提供的原生 PyTorch 支持和 AWS Neuron SDK，i14 完成了基于 GLA 的模型训练与推理服务，同时保留了固定大小的循环状态。

对于 Agent 系统，这改变了一项基本的服务约束。更长的历史不再要求更大的单序列状态，有助于团队保留更多上下文，同时让内存使用保持可预测。

如果不断增长的上下文正在推高服务成本、降低并发能力，或限制 Agent 记忆，可以试用 i14，或联系你的 AWS 客户团队，评估在 AWS Trainium2 上运行 Laplace 3B。

AWS Trainium 与 Trainium2

AWS Neuron SDK 文档

通过 TorchNeuron Native 在 AWS Trainium 上使用原生 PyTorch

采用硬件高效训练的 Gated Linear Attention Transformers

Muon optimizer 参考实现

AWS Generative AI Innovation Center