Beyond the KV cache: Training and serving i14’s Laplace 3B linear attention model on AWS Trainium2 with native PyTorch

This article describes how i14 adapted its Laplace 3B linear attention model for AWS Trainium2, replacing transformer softmax attention with Gated Linear Attention (GLA) to maintain constant-size recurrent state independent of context length.

• Laplace 3B uses GLA to compress context into fixed 29.5 MB recurrent state per sequence, eliminating KV cache growth
• Achieved 120,000 tokens per second training throughput on one trn2.48xlarge instance with 64 NeuronCores
• Served 11,544 output tokens per second across 1,280 concurrent sequences with predictable memory usage
• Replaced GPU-optimized Triton kernels with native PyTorch operators through TorchNeuron Native backend
• Optimized Muon optimizer communication using all-to-all collectives, reducing per-rank traffic by 24x and increasing throughput 1.63x
• Unrolled 16-step backward pass to fit Trainium2 compiler constraints while aligning 128-token chunks with SBUF partitions
• Fixed recurrent state enables long-running agents and large-document analysis without growing serving memory

The project demonstrates that AWS Trainium2 supports custom architectures beyond standard transformers, enabling efficient training and inference for models with alternative attention mechanisms.