i14将Laplace 3B线性注意力模型移植至AWS Trainium2，推理显存固定为29.5MB

AI研究公司i14与亚马逊云服务（AWS）合作，成功将其开发的Laplace 3B线性注意力模型在AWS Trainium2芯片上进行训练和服务。该模型采用门控线性注意力（GLA）设计，其核心优势在于推理时的循环状态大小固定为每个序列29.5 MB，不随对话或文档上下文的长度增加而膨胀。这一特性旨在解决传统Transformer架构中键值（KV）缓存随上下文增长导致并发降低、服务成本上升的问题。

在单台trn2.48xlarge实例上，Laplace 3B的训练吞吐量约为每秒120,000个token，并在1,280个并发序列下实现了每秒11,544个输出token的服务性能。这些数据由i14和AWS共同披露，属于厂商自测结果，尚未经过独立的第三方基准验证。

为了在Trainium2上运行这一原本依赖GPU优化Triton内核的模型，工程团队进行了三项关键改动：首先，使用标准PyTorch操作符重写GLA注意力内核，以匹配NeuronCore片上缓冲区的128个分区；其次，重新设计Muon优化器的通信方式，通过两次all_to_all集合通信将每卡流量减少24倍，吞吐量提升1.63倍；最后，静态展开GLA递归的16个步骤，以适应Trainium2静态编译器的图大小和片上内存限制。

线性注意力通过将历史压缩到固定大小的矩阵状态中，将训练复杂度从O(T²)降至O(T)。虽然这带来了可预测的服务需求，适合长运行智能体和大型文档分析，但文章也指出，这种压缩意味着必须针对具体工作负载评估检索质量，因为模型不再保留每一个键值对。