i14 Ports Laplace 3B Linear-Attention Model to AWS Trainium2

Fixed-size recurrent state replaces the KV cache, but porting it off GPU took real engineering.

i14, an AI research company, has trained and served its Laplace 3B linear-attention model on AWS Trainium2. It used native PyTorch rather than the GPU-specific Triton kernels it started with.

The model keeps a fixed recurrent state of 29.5 megabytes per sequence no matter how long the context runs. That replaces the key-value cache that grows with every token in a standard transformer. On a single trn2.48xlarge instance, AWS reports Laplace 3B trained at about 120,000 tokens per second. It served 11,544 output tokens per second across 1,280 concurrent sequences.

Laplace 3B swaps softmax attention for Gated Linear Attention. This compresses prior context into a fixed-size matrix per head rather than storing every key and value. That turns training attention from quadratic scaling with sequence length into linear scaling. It also makes inference memory per sequence constant, regardless of how long a conversation or document runs. The 3.028-billion-parameter model has 24 layers and 16 attention heads, with a geometric decay schedule spanning half-lives from one token to 2,048 tokens. It was trained on 100 billion tokens of FineWeb-Edu data using the Muon optimiser for its two-dimensional weight matrices.

Getting there on Trainium2 took three specific changes, according to the AWS Generative AI Innovation Center team that did the port. It rewrote the GPU-only GLA kernel as standard PyTorch operators through the beta TorchNeuron Native backend. Sequences were chunked into 16 blocks of 128 tokens to match the chip's 128 on-chip memory partitions, which AWS says lifted measured throughput by 14%. The team also redesigned how Muon's orthogonalisation step communicates under FSDP2. It replaced a gather-and-broadcast pattern with two all_to_all collectives, cutting per-rank traffic 24 times, from 10.87 GB to 0.45 GB per step. Throughput rose 1.63 times, from 73,378 to 119,934 tokens per second, with what AWS describes as a bit-identical loss curve. Finally, the team statically unrolled the 16 sequential steps of the GLA backward pass. Trainium2's compiler must trace a fixed graph ahead of time rather than build one at runtime as GPU autograd does, so unrolling let it fit within the chip's graph-size and on-chip memory limits.

Linear attention is not a free upgrade. The model compresses history into a fixed state rather than retaining every key and value, so AWS says teams must separately check retrieval quality for their own workloads. The appeal is for agentic applications and large-document analysis that currently have to truncate context or pay rising serving costs as a KV cache grows. This is also a proof point for AWS's own infrastructure strategy. Trainium2 can now run an architecture built and tuned for GPUs, without the vendor rewriting it in a GPU-specific kernel language first — provided engineers redo compiler-level choices like chunk size and backward-pass unrolling by hand.

Written from the primary source above by The Neural Net’s automated desk (claude-sonnet-5, theneuralnet/pipeline@1) without human review. How this works. Corrections: check the source first — if we got it wrong, the source link is the record.