From the source
i14 and AWS trained and served i14's Laplace 3B linear attention model on Trainium2 using native PyTorch, keeping its Gated Linear Attention design and Muon optimizer.
- Laplace 3B was built as an alternative to softmax attention, whose key-value cache grows with context and can cut concurrency, raise serving cost and force agentic applications to drop useful history.
- Its recurrent state stays at 29.5 MB per sequence regardless of context length, changing training attention from O(T²) to O(T).
- On one trn2.48xlarge instance the model trained at about 120,000 tokens per second and served 11,544 output tokens per second across 1,280 concurrent sequences.
- The AWS Generative AI Innovation Center rewrote the GPU-optimized Triton GLA kernel with standard PyTorch operators, using 16 chunks of 128 tokens to match the 128 partitions of the NeuronCore on-chip buffer, and chose the beta TorchNeuron Native backend after comparing approaches.
- Two all_to_all collectives in MuonFSDP2 cut per-rank traffic by 24 times and raised throughput by 1.63 times, and the sequential backward pass was statically unrolled across the 16 chunk steps to fit compiler graph-size and on-chip memory limits.
- Because linear attention compresses history rather than retaining every key and value, i14 says retrieval quality must be evaluated for each workload.
