Laplace 3B brings fixed-state linear attention to AWS Trainium2
AWS Machine Learning Blog reports that i14 and AWS adapted the 3B-parameter Laplace 3B model to AWS Trainium2 using native PyTorch through TorchNeuron Native, preserving its Gated Linear Attention architecture and Muon optimizer. On one trn2.48xlarge instance, the implementation trained at about 120,000 tokens per second and served 11,544 output tokens per second across 1,280 concurrent sequences. The blog says the recurrent state remains 29.5 MB per sequence regardless of context length, while warning that fixed-state compression requires workload-specific retrieval-quality evaluation.
Laplace 3B demonstrates that AWS Trainium2 can support a nonstandard recurrent architecture with constant per-sequence inference state, offering practitioners a concrete linear-attention deployment case to evaluate.
Most large language models (LLMs) use transformer architectures with softmax attention. During inference, their key-value (KV) cache grows with the conversation or document context. This can reduce concurrency, increase serving cost, and force agentic applications to discard useful history.
i14 designed Laplace 3B as a linear attention alternative. Its recurrent state stays at 29.5 MB per sequence, independent of context length. Working with AWS, i14 used native PyTorch support through TorchNeuron Native to bring this custom architecture to AWS Trainium2. This preserved its Gated Linear Attention (GLA) design and Muon optimizer.
On one trn2.48xlarge instance, Laplace 3B trained at about 120,000 tokens per second. The model served 11,544 output tokens per second across 1,280 concurrent sequences. This post explains the model design, the changes i14 made for Trainium2, and the training and inference results.
About i14
i14 is an AI research and technology company that develops model architectures, components, and runtimes. Its work draws on AI, computational cognitive science, and research into how biological systems learn over time. i14 applies these ideas to models that can retain long-range context without allowing serving state to grow with every token.
#i14#Laplace 3B#AWS Trainium2#TorchNeuron Native#AWS Neuron SDK#Gated Linear Attention#Muon optimizer#PyTorch