i14 Laplace 3B on Trainium2: Fixed-Memory Serving Results

AWS and i14 published an engineering report on October 20 showing Laplace 3B training and serving on Trainium2. Its recurrent state occupies 29.5 MB per sequence, independent of context length. That describes its context state, not total model memory.

For infrastructure teams, the interesting question is whether smaller, predictable session memory translates into a useful service. The answer requires looking at model quality, the software needed to run it and the difference between aggregate throughput and a user's waiting time.

Why linear attention changes the memory budget

The underlying idea has a research history. The Gated Linear Attention paper, first submitted in 2023, describes models that can train in parallel while behaving like recurrent networks during inference. Instead of expanding a store of past keys and values, a recurrent implementation updates a matrix-shaped hidden state as tokens arrive.

The paper also explains why attractive complexity figures are insufficient. Moving data around the accelerator can offset theoretical savings, and linear attention has historically faced quality trade-offs against ordinary softmax attention. Its contribution includes an implementation designed around those hardware costs.

This distinction matters when evaluating any fixed-state model. Accepting a long input and reliably recovering a detail from it are separate capabilities. A document-analysis trial should test information near the beginning, middle and end, including questions that require combining distant facts. A memory-allocation number cannot answer those questions.

What the Laplace measurements establish

On one trn2.48xlarge instance with sixteen Trainium2 chips, the team reports 11,544 output tokens per second across 1,280 concurrent sequences through vLLM. Training reached roughly 120,000 tokens per second. The published training results cover about 44 billion tokens, rather than completion of the stated 100-billion-token recipe.

These are team-reported infrastructure measurements. The serving rate is an aggregate across requests, not the speed of one conversation. For a user-facing product, it needs to sit beside time to first token, response completion time and the slowest requests under load.

The training result also answers a different question from model usefulness. Sustained processing speed tells an engineering team about the feasibility of a training run. Choosing a checkpoint for an application still requires evaluation of its outputs.

Native PyTorch still requires adaptation

The port replaced GPU-specific Triton attention kernels with PyTorch operations, reorganized Muon optimizer communication and statically unrolled the recurrent backward pass. Native support did not eliminate architecture-specific work.

The TorchNeuron documentation describes a public beta supporting Trainium2 and Trainium3. It uses PyTorch's device-backend mechanism and supports eager execution, torch.compile and standard distributed APIs such as FSDP and DTensor.

For ordinary PyTorch code, the documented changes include selecting the Neuron device and mixed-precision context, configuring sharding around NeuronCores, and removing CUDA-specific compilation settings. Those instructions provide a starting point. A team with custom GPU kernels should inventory them separately rather than estimate migration effort from a device-name change alone.

Changing batch sizes can hide a deployment problem

In the Laplace tests, new decode batch widths incurred about eleven seconds of compilation on first use. Padding to fixed widths removed that cost.

The Neuron compilation guide explains the mechanism: the backend does not support dynamic shapes, so distinct shapes produce separate compiled graphs. Reusing a cached graph is different from encountering a new shape. Padding or grouping inputs into a limited set of sizes helps reuse compiled work.

A deployment test should therefore include requests arriving and finishing at different times. Run both a cold-start trial and a warmed service, then inspect whether changing traffic introduces new compilations. A benchmark that repeats one comfortable batch shape can miss the delay a real user encounters.

The report directs interested teams to i14 or AWS for evaluation; it does not announce downloadable weights. Treat it as a documented engineering case, with an application-specific trial still needed before making a migration decision.

I'm Paras. I work at the intersection of AI, engineering and distribution.

I help technical founders turn infrastructure results into clear launch stories, engineering explainers and developer content.

Book a 30-minute call.