i14 ports Laplace 3B linear attention model to AWS Trainium2 with fixed 29.5MB inference memory
i14 and AWS ported the Laplace 3B linear attention model to Trainium2, achieving a fixed inference state size regardless of context length.
ImportanceLocalEvidenceE2 unreplicatedWrite-upQuick
AI research company i14 has collaborated with Amazon Web Services (AWS) to successfully train and serve its Laplace 3B linear attention model on AWS Trainium2 chips. The model utilizes Gated Linear Attention (GLA), a design where the recurrent state remains fixed at 29.5 MB per sequence during inference, independent of the conversation or document context length. This approach aims to address the limitations of traditional Transformer architectures, where the key-value (KV) cache grows with context, reducing concurrency and increasing serving costs.
On a single trn2.48xlarge instance, Laplace 3B achieved a training throughput of approximately 120,000 tokens per second and served 11,544 output tokens per second across 1,280 concurrent sequences. These figures are disclosed by i14 and AWS as vendor-reported results and have not yet been independently verified by third-party benchmarks.
To run this model, which originally relied on GPU-optimized Triton kernels, on Trainium2, the engineering team implemented three key changes: rewriting the GLA attention kernel using standard PyTorch operators to match the 128 partitions of the NeuronCore on-chip buffer; redesigning Muon optimizer communication via two all_to_all collectives, which reduced per-rank traffic by 24 times and increased throughput by 1.63 times; and statically unrolling the 16 chunk steps of the GLA recurrence to fit the static compiler's graph-size and on-chip memory limits.