大多数大型语言模型（LLM）使用带softmax注意力的transformer架构。在推理过程中，其键值（KV）缓存会随着对话或文档上下文而增长。这会降低并发能力，增加部署成本，并迫使智能体（agentic）应用丢弃有用的历史记录。
i14 将Laplace 3B设计为线性注意力替代方案。其循环状态保持在每个序列29.5 MB，与上下文长度无关。i14与AWS合作，使用了 通过TorchNeuron Native提供的原生PyTorch支持 将这一自定义架构带到 AWS Trainium2上。这保留了其门控线性注意力（GLA）设计和Muon优化器。
在一台trn2.48xlarge实例上，Laplace 3B以约每秒120,000个token的速度进行训练。该模型在1,280个并发序列上实现了每秒11,544个输出token。本文介绍了模型设计、i14为Trainium2所做的更改，以及训练和推理结果。
关于i14
i14是一家开发模型架构、组件和运行时的AI研究与科技公司。其工作借鉴了AI、计算认知科学以及关于生物系统如何随时间学习的研究。i14将这些理念应用于能够在保留长程上下文的同时，不让部署状态随每个token增长的模型。
为什么选择线性注意力模型？
Transformer softmax注意力在上下文中比较各个token。在训练过程中，这一计算量随序列长度呈二次方增长。在推理过程中，KV缓存随着模型处理更多token而增长。
线性注意力围绕一个为每个token更新的循环状态来组织计算。Laplace 3B使用GLA将先前的上下文压缩为固定大小的矩阵状态。这将训练注意力的复杂度从O(T2)降低到O(T)，并使推理状态保持在每个序列固定的29.5 MB。
| 考量 | Transformer softmax注意力 | 使用GLA的Laplace 3B |
| 训练注意力 | 以O(T2)的扩展方式构建成对token分数 | 以O(T)的扩展方式更新分块循环状态 |
| 推理内存 | KV缓存随上下文增长 | 循环状态保持固定大小 |
| 更长的上下文 | 每个序列消耗更多加速器内存 | 上下文长度不会增加循环状态的分配 |
| 实现方式 | 使用广泛优化的注意力内核 | 需要针对特定架构的内核和编译器调优 |
线性注意力压缩历史信息而不是保留每一个键和值。因此你必须针对你的工作负载评估检索质量。固定状态对于长期运行的智能体、大型文档分析以及需要具有可预测服务需求的持久内存的应用是有利的。
i14 如何为 AWS Trainium2 适配 Laplace 3B
Laplace 3B 的原始实现使用了针对 GPU 优化的 Triton 内核。AWS 生成式 AI 创新中心（GenAIIC）团队首先使用标准 PyTorch 算子为 Trainium2 重写了 GLA 内核。在比较了两种方案后，团队选择了 TorchNeuron Native，这是 AWS Trainium 的原生 PyTorch 后端的 beta 版本，因为它的性能更好。
该实现需要进行三项更改：
- 为 GLA 注意力内核使用原生 PyTorch。 团队用标准 PyTorch 算子替换了针对 GPU 的 fused_chunk_simple_gla Triton 内核，使用 16 个各含 128 个 token 的块，以匹配 NeuronCore 片上缓冲区的 128 个分区。
- 重新设计 Muon 的通信方式。 MuonFSDP2 中的两个 all_to_all 集合通信将每个 rank 的流量减少了 24 倍，并将吞吐量提高了 1.63 倍。
- 编译顺序反向传播。 团队对 GLA 递归的 16 个块步骤进行了静态展开，以适应编译器的图大小和片上内存限制。
工程挑战部分详细描述了每一项改动。
为什么选择 AWS Trainium 加速器？
AWS Trainium 是 AWS 面向模型训练和推理设计的机器学习（ML）加速器系列。trn2.48xlarge 实例结合了 16 个 Trainium2 芯片、64 个逻辑 NeuronCore、高带宽内存以及 NeuronLink-v3 互联。这些资源支持本文所述的分布式训练和推理配置。
Trainium2 的三个特性决定了实现方式：软件管理的片上内存（挑战 1）、快速的芯片间互联（挑战 2）以及静态编译器（挑战 3）。编译器将 GLA 递归跟踪为固定图。
软件管理的片上内存。 每个 NeuronCore 拥有一个 28 MiB 状态缓冲区（SBUF，含 128 个分区）和一个 2 MiB 部分和缓冲区（PSUM）。PSUM 连接到一个 128×128 脉动张量引擎，旁边还有独立的向量单元和标量单元。SBUF 并不是缓存：编译器提前调度每一次 HBM-to-SBUF 传输，因此数据在需要时就位于正确位置，不会发生逐出或缺失（图 1）。
图 1：NeuronCore-v3 内部架构。软件管理的 SBUF（28 MiB，128 个分区）、张量引擎（128×128 脉动阵列）、向量引擎、标量引擎以及 GPSIMD 核心。
i14 与 Laplace 3B 架构
Laplace 3B 是 24 层的 i14 解码器，其中每一层都用 Simple GLA 取代了 softmax 注意力。图 2 从顶部的输入 xt 一直跟踪到该层底部的输出 xt + 1 。本节解释每个模块：归一化与卷积、QKVG 投影、GLA 混合器以及前馈通道混合器。
归一化与卷积。 每个 token 的残差流向量 xt ∈ ℝD (D = 3, 072) 会先经过一个 LayerNorm 和一个逐深度（depth-wise）因果 Conv1d，其宽度为 k = 4 且带 SiLU 的深度可分离因果 Conv1d。这使每个位置在进入混合器之前获得一个 4-gram 摘要。GLA 混合器随后通过其衰减状态对整个历史进行基于内容的检索。
QKVG 投影。 卷积后的输出经过一个融合矩阵 Wqkvg ∈ ℝ9, 216 × 3, 072，其中 9, 216 = H × (2dk + 2dv）。它被拆分为四个逐 token 的流：查询 qt、键 kt、值 vt，以及一个门控 gt (H = 16, dk = 96, dv = 192）。门控在混合器之后使用。
图 2：Laplace 解码器层架构。绿色 = 固定（无参数），灰色 = 可学习，黄色 = 循环状态，蓝色 = 投影。
GLA 混合器。 每个头 h 维护一个大小为 St(h) 的矩阵状态 dk × dv （96 × 192），这是过去键值对的压缩且衰减的记录。在每个 token 处，它按如下方式更新：
GLA 循环状态更新方程
这里 λh 是固定的每头衰减系数，而 ktvt⊤ 是将当前键值对写入状态的秩-1 外积。读取通过查询状态完成：
GLA 状态读取方程
接下来，一个自我增强步骤向每个头的输出添加一个门控的值副本：
GLA 自我增强方程
这里 σ 是 sigmoid 函数，⊙ 是逐元素乘积。门控是查询与对键做可学习的逐通道缩放后的点积，其中 a 是可学习的缩放因子（aug_balance）。在初始化时， a = 0，这给出了值的一半强度的跳跃连接。每个头的输出在各自衰减层级内经过 Head RMSNorm，然后被拼接，并由 SiLU(gt), 并由 Wout ∈ ℝD × D 投影回残差流。
衰减调度。 每个头的衰减率 λh 控制一个头遗忘过去 token 的速度。它在初始化时固定（不学习），并由几何调度设定。该调度为 16 个头提供了一系列不同的记忆长度，并定义每个头的对数衰减率 γh:
GLA 衰减调度公式
衰减调度的范围。 对于此模型， Tmin = 1 且 Tmax = 2, 048。因此，头 0 具有 λ0 = 0.5（半衰期为一个 token），而头 15 具有 λ15 ≈ 0.9997(半衰期为 2,048 个 token)。各个头被分组为层(tier),每层共享一个 Head RMSNorm。
通道混合器(FFN)。 在注意力块及其残差之后,第二个 LayerNorm 输入到一个内层维度为 8,192 的 GLU (SiLU) FFN。它的三个矩阵是两个上投影 W1, W2 ∈ ℝ8, 192 × D 和一个下投影 W3 ∈ ℝD × 8, 192。随后是第二个残差。
状态 S(h) 固定为 dk × dv = 96 × 192 每头,与序列长度无关。因此,推理是每个 token 的 O(1),训练是 O(T) 而不是 O(T2).
| 参数 | 值 |
| 参数量 | 总计 3.028B(2.718B Muon + 309.6M AdamW) |
| 层数 | 24 个解码器层， D = 3, 072 隐藏维度 |
| 标题 | 16个头, dk = 96(半宽 Q/K), dv = 192(全角 V) |
| Mixer | 采用固定几何衰减的简单 GLA（T分钟 = 1, Tmax = 2, 048) |
| Conv | 深度逐通道因果Conv1d，核宽 4，SiLU |
| FFN | GLU (SiLU)，内部维度 8,192 |
| Norm | LayerNorm pre-attn + Head RMSNorm（每层）+ LayerNorm pre-FFN |
| Vocab | 50,304（填充至 64 的倍数） |
| Sequence | 2,048 |
| 训练数据 | 100B tokens，FineWeb-Edu |
| Optimizer | 2D权重矩阵使用Muon；嵌入和归一化使用AdamW |
| Batch size | 524,288 tokens/步，余弦学习率，3,815步预热 |
工程挑战
这三项改动各自带来了不同的工程挑战。
挑战 1：在 Trainium 上使用原生 PyTorch 实现 GLA
在 GPU 上，GLA 以 flash-linear-attention 库中的融合 Triton 内核 fused_chunk_simple_gla 运行。GenAIIC 团队将其改写为通过 TorchNeuron Native运行的标准 PyTorch 算子，以对固定大小分块的显式循环实现。分块数量在编译时固定，并决定了反向展开深度。该团队使用 fla_num_chunks = 16，因此每个分块包含 C = 128 个 token。
在每个分块 c内部，分块内输出为带掩码的线性注意力，运行状态 S 累加之前的上下文：
分块 GLA 块内输出
随后状态推进到下一个块：
分块 GLA 状态更新
此处 λC 是按头计算并在块上复合的衰减。16 次块内计算并行执行，只有状态更新是串行的（图 3）。
图 3：GLA 分块循环，具有并行的块内注意力和串行的状态链。
由于 SBUF 有 128 个分区，每个 128 token 的块在每个分区上正好映射一个 token，无需填充或跨分区移动。这种对齐使实测吞吐量提升了 14%（图 4）。
图 4：块大小 C = 128 与 128 个 SBUF 分区一一对应。
挑战 2：FSDP2 下的分布式 Muon 正交化
Muon 在更新之前对每个梯度进行正交化。对于 96 个二维权重矩阵（2.718 billion 个参数）中的每一个，它运行五次 Newton–Schulz 迭代（算法 1）。
Newton-Schulz 算法
分布式问题。 Gram 矩阵 A = XX⊤ 将权重的每一行与其他所有行相关联。在全分片数据并行（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，减少 24 倍。吞吐量从 73,378 提升到 119,934 token 每秒，提高 1.63 倍，且损失曲线逐位一致（图 6）。
图 6：每步通信量：all_gather 为 10.87 GB，而 all_to_all 为 0.45 GB。
通信开销。 Newton–Schulz 仅占训练步 FLOPs 的约 2%（1.95 × 1014 ，总量为 9.53 × 1015），因此通信才是瓶颈。在带宽为 1.28 TB/s 的 NeuronLink-v3 上，0.45 GB 的交换在约 4.4 秒的步长中占用不到 0.4 ms。
挑战 3：顺序递归的反向传播展开
反向传播需要对全部 16 次顺序状态更新求导。在 GPU 上，autograd 在运行时构建该图。在 Trainium2 上，编译器必须预先看到整个展开的 16 步图，这带来了三个约束。
图深度随分块数量增长：在 32 个分块（C = 64）时，图超出了调度器容量。每分块内存随分块大小增长：在 8 个分块（C = 256）时，中间结果溢出到 HBM。每个形状都在追踪时固定，因此每个循环中间量都需要自己的静态分配。
16 个分块（C = 128）同时满足全部三个条件：展开图符合调度器容量，中间结果保留在片上，且分块与 128 个 SBUF 分区对齐。
训练结果
GenAIIC 团队在一个 trn2.48xlarge 实例上，使用全部 64 个逻辑 NeuronCore 和 all_to_all Muon 路径训练了 Laplace 3B。完整配方为 190,735 步（1000 亿 FineWeb-Edu 词元）。以下结果是到第 84,210 步、约 440 亿词元时的数据。
吞吐量。 编译预热完成后，吞吐量稳定在约 120,000 词元每秒，并保持到第 84,210 步。
损失与梯度稳定性。 损失（图 7a）从约 11.4 平滑下降到第 84,210 步的 2.3729。梯度范数（图 7b）呈现出预期的早期尖峰，随后稳步下降至约 0.05，预热后没有出现尖峰、发散或 NaN/Inf 值。
图 7：训练动态。(a) 损失 11.4 → 2.3729。(b) 来自 TensorBoard 数据的梯度范数。
Trainium2 上的推理优化
通过未经修改的 vLLM 引擎，GenAIIC 团队在一个 trn2.48xlarge 实例上，以 1,280 个并发序列实现了每秒 11,544 个输出词元的服务。
恒定的服务状态。 每个序列 29.5 MB 的循环状态大约比 8,192 词元的同等 KV 缓存小 82 倍。将提示从 2,048 词元增加到 8,192 词元，仅使每个输出词元的成本增加了 0.5%。
未经修改的 vLLM。 一个包装器替换每个解码器层的一个属性（即词元混合器），以读写 vLLM 池中的循环状态。由于模型声明自身不使用注意力机制，vLLM 为每个序列分配一个固定页面。输出与 i14 参考逐词元一致（在 BF16 下 3/3 个提示上 48/48 步全部匹配）。
定位瓶颈。 对解码时间与解码器层数 L 进行线性拟合得到：
在 L = 24 时，即每步 168.1 ms，其中 96.9% 位于解码器层内。大型投影耗时 13.27 ms，接近 12 ms 的 HBM 流式处理下限。大部分时间花在 eager PyTorch 下逐个分发的小算子上。
扩展到完整实例。 在不使用张量并行的情况下，吞吐量通过副本扩展，每个逻辑 NeuronCore 一个引擎和 20 个序列。eager 吞吐量从一个副本的 133 输出词元每秒增长到 64 个副本的 6,772，此时主机 CPU 使用率达到 82%。编译解码步骤消除了逐算子分发，将吞吐量提高了 1.7 倍，达到 11,544（图 8）。
图8：整实例推理吞吐量。(a) eager 与编译版 decode 上的副本扩展。(b) 编译版 decode 将吞吐量从每秒 6,772 个输出 token 提升至 11,544 个。
固定批量基准测试所掩盖的问题。 在真实流量下，每个新的 decode 批次宽度都是一个单独编译的形状，其首次使用耗时约 11 秒。这些步骤仅占总量的 1.1%，却消耗了 decode 时间的 47.5%。填充到固定宽度可以消除这一成本，且结果逐位一致。编译版 decode 随后在错峰到达场景下将每个副本的吞吐量提升了 3.52 倍，而在固定批次上仅提升 1.46 倍。由于 BF16 重结合会使编译输出与 eager 输出略有差异，在需要精确可复现性时应使用 eager decode。
关键要点
该项目在训练和推理两方面都总结出了经验教训：
- 原生 PyTorch 可以在 Trainium2 上支持自定义架构。 i14 保留了 GLA 架构，并在最终实现中使用了 TorchNeuron Native。
- 对于自定义优化器，通信设计至关重要。 Muon 的 all_to_all 设计将每 rank 流量减少了 24 倍，并将训练吞吐量提升了 1.63 倍。
- 固定状态使推理内存可预测。 每个序列无论上下文长度如何都使用 29.5 MB。
- 编译之前先进行性能分析。 占用大部分 decode 时间的是算子分发，而非算术运算或 HBM 带宽。
- 基准测试要使用真实的到达模式。 固定批次掩盖了为变化的批次宽度编译新形状的成本。
结论
Laplace 3B 表明，AWS Trainium2 可以支持标准 transformer 之外的模型架构。通过 TorchNeuron Native 和 AWS Neuron SDK 使用原生 PyTorch，i14 在保持固定循环状态的同时训练并部署了其基于 GLA 的模型。
对于智能体系统，这改变了一项基本的部署约束。更长的历史记录不再需要更大的每序列状态，这可以帮助团队在保持内存使用可预测的同时保留更多上下文。
如果不断增长的上下文正在增加部署成本、降低并发或限制智能体内存，请尝试 i14 或联系您的 AWS 客户团队，以在 AWS Trainium2 上评估 Laplace 3B。
资源
- 试用 i14
- AWS Trainium 和 Trainium2
- AWS Neuron SDK 文档
- 在 AWS Trainium 上通过 TorchNeuron Native 支持原生 PyTorch
- 具有硬件高效训练的门控线性注意力Transformer
- Muon 优化器参考实现
- AWS 生成式 AI 创新中心