超越KV缓存：使用原生PyTorch在AWS Trainium2上训练和部署i14的Laplace 3B线性注意力模型

英文原标题：Beyond the KV cache: Training and serving i14’s Laplace 3B linear attention model on AWS Trainium2 with native PyTorch

本文为英文原文的机器翻译摘要，原文见：AWS Machine Learning Blog ↗

• i14 推出 Laplace 3B 线性注意力模型，用 Gated Linear Attention (GLA) 替代 softmax 注意力，将推理状态固定为每序列 29.5 MB，不随上下文长度增长。
• i14 与 AWS 合作，通过 TorchNeuron Native 将自定义架构移植到 AWS Trainium2，保留了 GLA 设计和 Muon 优化器。
• 在单台 trn2.48xlarge 实例上，Laplace 3B 训练速度约 120,000 tokens/s；推理在 1,280 个并发序列下达到 11,544 输出 tokens/s。
• 为适配 Trainium2，团队做了三项关键改动：用标准 PyTorch 算子重写 GLA 内核、重新设计 Muon 通信（all_to_all 集合通信将每 rank 流量降低 24 倍、吞吐提升 1.63 倍）、静态展开 16 个 chunk 步骤以编译顺序反向传播。
• Trainium2 的软件管理片上内存、快速芯片互连和静态编译器是影响实现方式的三个特性；模型每层包含 LayerNorm、深度可分离因果卷积、QKVG 投影、GLA 混合器和 GLU FFN。

关键数据

• 推理状态固定为每序列 29.5 MB
• 训练速度约 120,000 tokens/s（单台 trn2.48xlarge）
• 推理吞吐 11,544 输出 tokens/s，1,280 并发序列
• Muon 通信优化：每 rank 流量降低 24 倍，吞吐提升 1.63 倍
• GLA 内核使用 16 个 chunk、每 chunk 128 tokens，匹配 NeuronCore 片上缓冲的 128 个分区
• 模型 24 层，隐藏维度 D=3,072，16 个注意力头，d_k=96，d_v=192
• 每个 NeuronCore 有 28 MiB SBUF（128 分区）和 2 MiB PSUM
• 衰减调度：T_min=1，T_max=2,048；头 0 的 λ=0.5，头 15 的 λ≈0.9997

为什么值得看：该案例展示了线性注意力模型如何通过原生 PyTorch 在 Trainium2 上高效训练和推理，为长上下文、高并发场景提供了固定内存占用的替代方案。