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

Large language models commonly utilize transformer architectures with softmax attention during inference. This leads to challenges such as increased concurrency, higher serving costs, and discarding valuable historical context. To address this, i14 developed Laplace 3B, a custom architecture employing linear attention. Laplace 3B's recurrent state remains constant at 29.5 MB per sequence, regardless of context length.

Collaborating with AWS, i14 utilized native PyTorch support through TorchNeuron Native to deploy this model on AWS Trainium2, preserving its Gated Linear Attention (GLA) design and Muon optimizer.

On a single trn2.48xlarge instance, Laplace 3B achieved a training speed of approximately 120,000 tokens per second. Upon deployment, the model delivered 11,544 output tokens per second across 1,280 concurrent sequences. This report delves into the model's design, the modifications made for Trainium2, and the training and inference results.

i14, an AI research and technology company, focuses on model architectures, components, and runtimes that maintain long-range context without expanding serving state with each token. Linear attention models, such as Laplace 3B, restructure computation using a recurrent state updated for each token. This transforms training attention from O(T^2) to O(T) and maintains a fixed inference state of 29.5 MB per sequence.

In contrast to transformer softmax attention, which grows quadratically with sequence length during training and results in a growing KV cache during inference, Laplace 3B employs GLA to compress prior context into a fixed-size matrix state. This approach reduces training attention to O(T) and keeps inference state constant at 29.5 MB per sequence.