i14 brings Laplace 3B linear attention model to AWS Trainium2 with native PyTorch

Why i14 moved away from the KV cache

Most large language models use transformer architectures with softmax attention. During inference, their key-value (KV) cache grows with the conversation or document context. According to the AWS post, this can reduce concurrency, increase serving cost, and force agentic applications to discard useful history.

i14 designed Laplace 3B as a linear attention alternative. It uses Gated Linear Attention (GLA) to compress prior context into a fixed-size matrix state. That changes training attention from O(T^2) to O(T) and holds inference state at 29.5 MB per sequence. i14 describes itself as a company that develops model architectures, components and runtimes, drawing on AI, computational cognitive science and research into how biological systems learn over time.

The post is clear about the trade-off: linear attention compresses history rather than retaining every key and value. Users therefore need to evaluate retrieval quality for their own workloads. The fixed state can benefit long-running agents, large-document analysis, and applications that need persistent memory with predictable serving requirements.

The model and its results on Trainium2

Laplace 3B is a 24-layer decoder in which every layer replaces softmax attention with Simple GLA. The model has 3.028B parameters in total, split into 2.718B trained with the Muon optimizer and 309.6M with AdamW. It uses a hidden dimension of 3,072, 16 heads, a sequence length of 2,048 and a vocabulary of 50,304. It was trained on 100B tokens of FineWeb-Edu with a batch size of 524,288 tokens per step.

Each head keeps a matrix state of 96 x 192 that acts as a decaying record of past key-value pairs. The decay rate per head is fixed at initialization and follows a geometric schedule, giving the heads a spread of memory lengths. Head 0 has a half-life of one token, while head 15 has a half-life of 2,048 tokens.

On one trn2.48xlarge instance, Laplace 3B trained at about 120,000 tokens per second, according to the post. The model served 11,544 output tokens per second across 1,280 concurrent sequences. A trn2.48xlarge combines 16 Trainium2 chips, 64 logical NeuronCores, high-bandwidth memory and NeuronLink-v3 connectivity.

Three changes to fit Trainium2

The original Laplace 3B implementation relied on GPU-optimized Triton kernels. The AWS Generative AI Innovation Center (GenAIIC) team first rewrote the GLA kernel with standard PyTorch operators. After comparing approaches, it chose the beta version of TorchNeuron Native, the native PyTorch backend for AWS Trainium, because it performed better. This preserved the GLA design and the Muon optimizer.

• GLA kernel in native PyTorch: the GPU-specific fused_chunk_simple_gla Triton kernel was replaced with standard PyTorch operators, using 16 chunks of 128 tokens to match the 128 partitions of the NeuronCore on-chip buffer.
• Muon communication: two all_to_all collectives in MuonFSDP2 cut per-rank traffic by 24 times and raised throughput by 1.63 times.
• Compiled backward pass: the team statically unrolled the 16 chunk steps of the GLA recurrence to fit the compiler's graph-size and on-chip memory limits.

Three Trainium2 properties shaped this work: software-managed on-chip memory, a fast chip-to-chip interconnect, and a static compiler that traces the GLA recurrence as a fixed graph. Each NeuronCore has a 28 MiB State Buffer with 128 partitions, and the compiler schedules every transfer ahead of time.

What it means for developers

The post shows that a custom linear attention architecture can run on Trainium2 without leaving native PyTorch, though it notes that such models require architecture-specific kernels and compiler tuning, unlike transformers with widely optimized attention kernels. For teams building long-context or agentic applications, the fixed per-sequence state offers predictable memory use. Whether the compressed state retains enough information still has to be tested on each workload.