The Pulse
AWS Serves i14’s Laplace 3B With a Fixed 29.5 MB State
AWS says i14’s Laplace 3B model keeps its recurrent inference state at 29.5 MB per sequence on Trainium2, regardless of context length. The company reports 11,544 output tokens per second across 1,280 concurrent sequences on one instance.
AI.info Team · Written with AI · fact-checked ·
Linear attention compresses history rather than retaining every key and value.
Kevin Shabahang, Iman Abbasnejad, Massimo Bardetti, Daniel Wirjo and Xiaodong Yu, AWS Machine Learning Blog
Laplace 3B keeps its serving state fixed
AWS and AI research company i14 say they trained and served Laplace 3B, a linear-attention model, on an AWS Trainium2 instance while keeping its recurrent inference state at 29.5 megabytes per sequence. Unlike a transformer’s key-value cache, which grows as more context enters a conversation, Laplace’s state stays the same size as context length increases.
The companies report 11,544 output tokens per second across 1,280 concurrent sequences on a single trn2.48xlarge instance. AWS says the model ran through an unmodified vLLM engine, with a wrapper connecting its recurrent state to vLLM’s memory pool. The result is a technical demonstration, not a published comparison of answer quality or serving cost against a transformer model.
GLA trades stored tokens for a recurrent record
Laplace 3B has 24 decoder layers and 3.028 billion parameters. Each layer replaces softmax attention with Simple Gated Linear Attention, which updates a fixed-size matrix state as it processes each token. AWS describes training attention as scaling linearly with sequence length, rather than quadratically, and inference state as independent of context length.
That compression changes what the model retains: it summarizes prior context instead of keeping every key and value. AWS cautions that teams should assess retrieval quality for their own workloads. The post points to long-running agents and large-document analysis as possible uses, but it does not report a task-level evaluation showing how Laplace performs on those workloads.
Trainium2 required changes to kernels and communication
The original Laplace implementation used GPU-optimized Triton kernels. The AWS Generative AI Innovation Center team replaced the Gated Linear Attention kernel with standard PyTorch operations, arranged as 16 chunks of 128 tokens, and used the beta TorchNeuron Native backend for Trainium. The chunk size matches the 128 partitions in each NeuronCore’s on-chip state buffer.
Training also depended on changing how the Muon optimizer communicated across the 64 logical NeuronCores. AWS says two all-to-all operations cut per-rank data traffic 24-fold and raised throughput 1.63 times, from 73,378 to 119,934 tokens per second. The team compiled the sequential backward pass by unrolling its 16 recurrent steps to fit the compiler’s graph and memory limits.
The reported results have defined limits
The training recipe targets 100 billion FineWeb-Edu tokens, but the results in the post cover 84,210 steps, or about 44 billion tokens. AWS reports throughput of roughly 120,000 tokens per second during that run. The reported serving state is about 82 times smaller than a comparable KV cache at 8,192 tokens; increasing the prompt from 2,048 to 8,192 tokens added 0.5% to the cost per output token in the team’s test.
Those figures describe one model, one accelerator instance and AWS’s own measurements; the post does not supply a direct transformer baseline or pricing comparison. It also identifies a practical serving wrinkle: when incoming batch widths change, new shapes can require compilation, which the team says takes about 11 seconds on first use. For Laplace 3B, fixed memory removes context growth from the per-sequence state calculation, while retrieval behavior and variable traffic still need workload-specific testing.