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

- Company: Amazon (amazon.com)
- Announced: 2026-10-20T21:14:10+00:00
- Category: not stated
- Coverage: not counted
- Announcement: no
- Group: routine
- Source: https://aws.amazon.com/blogs/machine-learning/beyond-the-kv-cache-training-and-serving-i14s-laplace-3b-linear-attention-model-on-aws-trainium2-with-native-pytorch/
- Record: https://forck.live/items/25451-beyond-the-kv-cache-training-and-serving-i14-s-laplace-3b-linear-attention
- Subject: Bedrock / Nova
- Models affected: Laplace 3B

i14 and AWS trained and served i14's Laplace 3B linear attention model on Trainium2 using native PyTorch, keeping its Gated Linear Attention design and Muon optimizer. Laplace 3B was built as an alternative to softmax attention, whose key-value cache grows with context and can cut concurrency, raise serving cost and force agentic applications to drop useful history. Its recurrent state stays at 29.5 MB per sequence regardless of context length, changing training attention from O(T²) to O(T). On one trn2.48xlarge instance the model trained at about 120,000 tokens per second and served 11,544 output tokens per second across 1,280 concurrent sequences. The AWS Generative AI Innovation Center rewrote the GPU-optimized Triton GLA kernel with standard PyTorch operators, using 16 chunks of 128 tokens to match the 128 partitions of the NeuronCore on-chip buffer, and chose the beta TorchNeuron Native backend after comparing approaches. Two all_to_all collectives in MuonFSDP2 cut per-rank traffic by 24 times and raised throughput by 1.63 times, and the sequential backward pass was statically unrolled across the 16 chunk steps to fit compiler graph-size and on-chip memory limits. Because linear attention compresses history rather than retaining every key and value, i14 says retrieval quality must be evaluated for each workload.

## Evidence

Verbatim from https://aws.amazon.com/blogs/machine-learning/beyond-the-kv-cache-training-and-serving-i14s-laplace-3b-linear-attention-model-on-aws-trainium2-with-native-pytorch/:

> 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.

---

Record: https://forck.live/items/25451-beyond-the-kv-cache-training-and-serving-i14-s-laplace-3b-linear-attention
Catalogue: https://forck.live/llms.txt
Current issue: https://forck.live/feed.md