i14发布Laplace 3B：基于AWS Trainium2的线性注意力模型，突破KV缓存瓶颈

摘要：i14公司推出了一款名为Laplace 3B的线性注意力模型，通过AWS Trainium2进行训练和推理。该模型采用Gated Linear Attention (GLA)架构，将每序列的递归状态固定在29.5MB，避免了传统Transformer模型KV缓存随上下文长度增长的问题，从而实现更高效的内存利用和可预测的推理成本。

---

核心突破与架构设计

i14公司发布的Laplace 3B是一款基于线性注意力机制的模型，旨在解决传统Transformer模型在处理长序列时KV缓存快速增长的问题。其核心架构特点包括：

• 固定递归状态：Laplace 3B的递归状态每序列固定为29.5MB，不随上下文长度变化。这使得在处理长序列时，内存使用更加可控，避免了KV缓存带来的内存开销问题。
• Gated Linear Attention (GLA)：通过GLA机制，Laplace 3B将注意力计算从二次复杂度降低到线性复杂度，同时保持对长程依赖的有效建模能力。

技术机制剖析

Laplace 3B的GLA架构通过以下方式实现高效计算：

1. 状态更新与压缩：每处理一个token，模型会更新一个固定大小的矩阵状态，该状态压缩了之前的上下文信息。
2. 查询-键-值-门控投影：模型通过一个融合矩阵将输入token投影为查询（Q）、键（K）、值（V）和门控（G）四个部分。门控用于控制状态的自适应更新。
3. 状态读取与自增强：在读取状态时，模型通过自增强机制将当前值的一部分添加到输出中，以增强对当前token的响应。

工程权衡与实测表现

在AWS Trainium2上，Laplace 3B的训练和推理性能表现如下：

• 训练速度：在单个trn2.48xlarge实例上，Laplace 3B的训练速度约为120,000 tokens/秒。
• 推理吞吐量：在1,280个并发序列下，模型每秒可生成11,544个输出token。
• 内存效率：由于递归状态固定，Laplace 3B的内存使用量不随上下文长度变化，这使得其在长序列处理任务中具有显著优势。

开发者落地与部署建议

对于希望利用Laplace 3B的开发者，以下几点值得注意：

• 原生PyTorch支持：Laplace 3B通过TorchNeuron Native在AWS Trainium2上实现了原生PyTorch支持，开发者可以无缝集成现有PyTorch代码。
• 推理优化：建议使用编译后的解码步骤以提高吞吐量，同时注意在需要精确再现性时使用eager模式。
• 资源管理：由于Laplace 3B的内存使用量固定，开发者可以更有效地规划资源分配，特别是在多租户或资源受限的环境中。

结论

Laplace 3B展示了AWS Trainium2在支持非传统Transformer架构方面的能力。通过使用GLA机制，Laplace 3B在保持高效训练和推理性能的同时，解决了KV缓存带来的内存瓶颈问题。这为需要处理长序列或对内存使用有严格要求的应用场景提供了新的解决方案。

关键要点

• 原生PyTorch支持：Trainium2支持自定义架构，Laplace 3B通过TorchNeuron Native实现。
• 通信设计的重要性：Muon的all_to_all设计减少了每rank的通信量，提升了训练吞吐量。
• 固定状态确保推理内存可预测：每序列使用29.5MB内存，不随上下文长度变化。
• 编译前进行性能分析：大部分解码时间由操作符调度而非算术或HBM带宽决定。
• 针对真实流量模式进行基准测试：固定批次隐藏了动态批次宽度下编译新形状的成本。

---

消息来源：AWS Machine Learning Blog (2026-10-20)

—— 完 ——

主题标签： #i14 关注 #Laplace 3B 关注 #线性注意力模型 关注 #AWS Trainium2 关注 #GLA架构 关注