Beyond the KV cache: Training and serving i14’s Laplace 3B linear attention model on AWS Trainium2 with native PyTorch
aws ·
This article details how i14 trained and served its Laplace 3B linear attention model on AWS Trainium2 using native PyTorch via TorchNeuron Native. Unlike traditional transformers that rely on a growing key-value cache, Laplace 3B uses Gated Linear Attention to maintain a fixed recurrent state of 29.5 MB per sequence, regardless of context length. The implementation required three major engineering adaptations: rewriting the GLA kernel with standard PyTorch operators, optimizing distributed communication for the Muon optimizer, and statically unrolling the backward pass for the compiler. Training on a single trn2.48xlarge instance achieved approximately 120,000 tokens per second with stable loss convergence. Inference performance reached 11,544 output tokens per second across 1,280 concurrent sequences by compiling the decode step to reduce operator dispatch overhead. These results demonstrate that custom architectures can achieve high efficiency on AWS hardware while offering predictable memory usage for long-context applications.
aws