AWS and i14 ditch the KV cache with a 3B linear-attention model on Trainium2
In short: AI research company i14 partnered with AWS to port its Laplace 3B model, a 24-layer architecture built on Gated Linear Attention (GLA) instead of standard transformer attention, onto AWS Trainium2 using native PyTorch (TorchNeuron Native). Rather than a KV cache that expands with every token, the model keeps a constant 29.5 MB recurrent state per sequence. On a single trn2.48xlarge instance, the team reached about 120,000 training tokens per second and served 11,544 output tokens per second across 1,280 simultaneous sequences through an unmodified vLLM engine.
This summary was generated automatically by AI from AWS Machine Learning's publication. It is our own text, not a copy of the original — facts, figures and quotes belong to the source, linked above and below.
What changed?
- 1GPU-only Triton kernel for GLA was rewritten in plain PyTorch, splitting sequences into 16 chunks of 128 tokens so each chunk lines up with Trainium2's on-chip SBUF partitions
- 2Muon optimizer's cross-rank communication was restructured with a double all_to_all exchange, cutting per-rank data traffic 24x and lifting training speed 1.63x
- 3The 16-step GLA backward pass was fully unrolled at compile time so the static Trainium2 compiler could handle the sequential recurrence within memory and scheduling limits
- 4Because the recurrent state never grows, pushing prompt length from 2,048 up to 8,192 tokens only added half a percent to the cost of generating each output token, and the fixed state is about 82x smaller than an equivalent KV cache at 8,192 tokens
| Parameter | Before | Now |
|---|---|---|
| Architecture | Transformer softmax attention with growing KV cache | Gated Linear Attention (GLA) with fixed-size recurrent state |
| Parameters | Not specified | 3.028B (2.718B Muon + 309.6M AdamW) |
| Per-sequence memory state | KV cache grows with context length | Fixed 29.5 MB recurrent state regardless of context length |
| Training throughput | 73,378 tokens/sec (gather-broadcast Muon) | 119,934 tokens/sec (all_to_all Muon), ~120,000 tokens/sec overall on trn2.48xlarge |
| Inference throughput | Not specified | 11,544 output tokens/sec across 1,280 concurrent sequences |
| Context window | 2,048 training sequence length | Tested up to 8,192 tokens with only 0.5% extra cost per output token |
Why it matters
Growing KV caches limit how much conversation or document history an LLM can hold cheaply during long-running sessions. A model with a flat memory footprint per sequence could support longer-running agents and document analysis without the usual memory and cost penalties, though retrieval quality compared to full attention still needs case-by-case evaluation.
What it means for AI agents and contact centers
For teams running voice agents or contact-center bots that need to retain long call histories or large amounts of context across a conversation, this architecture suggests a path to keep per-call memory and inference cost flat as context grows, instead of scaling with conversation length; before adopting it, compare retrieval accuracy on your own transcripts against your current transformer-based model, and benchmark latency and cost per minute.
🧪 Worth Testing
Flat memory and cost scaling with context length is directly relevant to long voice/contact-center sessions, but retrieval quality versus full attention should be validated on real call transcripts before switching.
Sources
- AWS Machine LearningOfficialPrimary sourceOriginal article →„Beyond the KV cache: Training and serving i14’s Laplace 3B linear attention model on AWS Trainium2 with native PyTorch“21 Oct 2026, 00:14
- Published by source
- 21 Oct 2026, 00:14
- Found by our system
- 21 Oct 2026, 00:16
- Summary generated
- 21 Oct 2026, 00:19
This article was written by AI from the original source. Facts, numbers and prices come from the source; missing values are marked “Not specified”. Legal notice, copyright and privacy