i14线性注意力模型Laplace 3B在AWS Trainium2上的训练与推理实践

传统Transformer的softmax注意力机制在训练时计算复杂度随序列长度呈O(T²)增长，推理时KV缓存随上下文持续膨胀，每条序列占用的加速器内存不断增加。这一特性在长上下文场景下会显著压缩并发数量，推高服务成本。i14设计Laplace 3B的核心出发点，正是用固定大小的循环状态替代不断增长的KV缓存，将训练注意力复杂度降至O(T)，同时将推理状态锁定在每序列29.5MB。

Laplace 3B是一个24层解码器，每层用简单门控线性注意力替代softmax注意力。每个注意力头维护一个d_k×d_v（96×192）的矩阵状态，作为过去键值对的压缩衰减记录。每个token到来时，状态通过固定的逐头衰减系数和当前键值对的秩一外积进行更新，查询时直接从该状态中读取。此外，模型还在每个头的输出上叠加了一个门控自增强步骤，进一步提升表达能力。

原始Laplace 3B实现依赖GPU优化的Triton内核，迁移至Trainium2需要解决三项核心工程挑战。首先，团队用标准PyTorch算子替换了GPU专用的fused_chunk_simple_gla Triton内核，将序列划分为16个128-token的块，以匹配NeuronCore片上缓冲区的128个分区。其次，针对Muon优化器的分布式通信，团队在MuonFSDP2中引入两个all_to_all集合操作，将每个rank的通信量降低24倍，吞吐量提升1.63倍。

第三项挑战来自Trainium2的静态编译器。编译器需要将计算图静态追踪为固定结构，而GLA循环的16个块步骤必须被静态展开，才能满足编译器对图规模和片上内存的限制。Trainium2的NeuronCore配备28MiB软件管理的状态缓冲区（SBUF，128个分区）和2MiB部分和缓冲区，编译器提前调度所有HBM到SBUF的数据传输，消除了缓存缺失和驱逐问题，为GLA循环的高效执行提供了硬件基础。

在性能表现上，单台trn2.48xlarge实例（搭载16颗Trainium2芯片、64个逻辑NeuronCore）上，Laplace 3B训练速度约为每秒12万个token。推理阶段，模型在1280条并发序列下实现每秒11544个输出token的吞吐量，且服务内存占用完全可预测。这一架构特别适合长期运行的智能体、大文档分析以及需要持久记忆的应用，但需注意线性注意力对历史信息的压缩特性，在部署前应针对具体检索任务评估召回质量。

要点

• Laplace 3B采用门控线性注意力，每条序列循环状态固定为29.5MB，彻底消除KV缓存随上下文增长的问题，使长上下文服务内存完全可预测。
• 迁移至Trainium2的关键在于三项改造：用标准PyTorch算子替换Triton内核、优化Muon分布式通信降低24倍per-rank流量、静态展开GLA循环以适配静态编译器。
• 单台trn2.48xlarge实例训练吞吐约12万token/秒，推理在1280并发序列下达11544输出token/秒，验证了线性注意力模型在专用AI加速器上的工程可行性。
• 线性注意力通过压缩历史而非保留全部键值对来管理上下文，部署前需针对具体检索任务评估信息召回质量，不能直接等同于标准Transformer的注意力效果。