i14 在 AWS Trainium2 上用原生 PyTorch 训练与部署 Laplace 3B 线性注意力模型
原标题:Beyond the KV cache: Training and serving i14’s Laplace 3B linear attention model on AWS Trainium2 with native PyTorch
AI 摘要
i14 公司设计了采用门控线性注意力(GLA)的 Laplace 3B 模型,其循环状态固定为每序列 29.5 MB,不随上下文长度增长。i14 与 AWS 合作,通过 TorchNeuron Native 的原生 PyTorch 支持将该架构移植到 AWS Trainium2,保留了 GLA 设计与 Muon 优化器。在单台 trn2.48xlarge 实例上,训练速度约 12 万 tokens/秒,推理在 1280 个并发序列下达到 11,544 输出 tokens/秒。
正文节选
Beyond the KV cache: Training and serving i14’s Laplace 3B linear attention model on AWS Trainium2 with native PyTorch Most large language models (LLMs) use transformer architectures with softmax attention. During inference, their key-value (KV) cache grows with the conversation or document context. This can reduce concurrency, increase serving cost, and force agentic applications to discard useful history. i14 designed Laplace 3B as a linear attention alternative. Its recurrent state stays at 2