Laplace 3B: A Linear Attention Model With a Fixed 29.5 MB State on Trainium2

AWS has published a detailed account of how a company called i14 trained and served Laplace 3B, a language model that replaces softmax attention with gated linear attention, on AWS Trainium2 chips using native PyTorch. The headline number is the memory: the model's recurrent state is 29.5 MB per sequence and stays that size at any context length, according to the AWS Machine Learning blog. A transformer's key-value cache grows with every token. This one does not.

That is a vendor-published case study about a model most readers cannot download, so this post separates what the write-up reports from what it leaves open. It is useful beyond the model itself: it shows what a non-transformer architecture costs and saves in practice, and it adds evidence to a debate we have covered before, in our look at sliding-window attention versus linear attention after post-training.

Quick facts

What is Laplace 3B and how is it built?

Laplace 3B is a decoder with 24 layers and a hidden size of 3,072. Instead of softmax attention, each layer runs simple gated linear attention, the design from the GLA paper on arXiv. There are 16 heads with a key dimension of 96 and a value dimension of 192, a depth-wise causal convolution of width 4 in each layer, and a gated feed-forward block with an inner size of 8,192. The vocabulary is 50,304 tokens and the training sequence length is 2,048.

One unusual choice is the decay schedule. Each head has a fixed forgetting rate rather than a learned one, spanning timescales from 1 to 2,048 tokens. Head 0 has a decay of 0.5, so it forgets almost immediately. Head 15 has a decay of about 0.9997, so it holds information across the whole window. The idea is that some heads track the last few words while others carry long-range context, all inside one constant-size state.

Why does the state stay at 29.5 MB?

Softmax attention stores a key and a value for every past token, in every layer. The cache grows linearly with context, and it is the main reason long prompts eat accelerator memory. Linear attention instead folds the past into a matrix that is updated token by token. The matrix has a fixed shape, so memory per sequence is a constant.

The AWS post says its 29.5 MB is about 82 times smaller than a comparable KV cache at 8,192 tokens, which implies a cache of roughly 2.4 GB per sequence for a model of this shape. That is our arithmetic from the post's ratio, not a figure it states. As a sanity check, 24 layers times 16 heads times a 96 by 192 state is about 7.1 million values. At four bytes each that is about 28 MB, which is in the neighborhood of the reported figure once the small convolution state is added. We did not verify i14's exact accounting.

A smooth river path carrying green dots, an image for a recurrent state that carries context forward at constant size in a linear attention model

The practical effect shows up in serving. Raising the prompt from 2,048 to 8,192 tokens added only 0.5% to the cost per output token, according to the post. In a transformer server, the same change would shrink how many sequences fit in memory. For background on why the cache is such a bottleneck, see our coverage of DeepSeek V4.1 Flash cutting KV cache memory and of push-based KV transfer in disaggregated inference.

How was it trained on Trainium2?

Training ran on a single trn2.48xlarge instance. The full recipe is 190,735 steps on 100 billion FineWeb-Edu tokens, at 524,288 tokens per step, with a cosine learning-rate schedule and a 3,815-step warmup. Throughput settled at roughly 120,000 tokens per second.

The reported results stop at step 84,210, about 44 billion tokens. Loss fell from about 11.4 to 2.3729, and gradient norm declined to about 0.05 with no post-warmup spikes, divergence, or NaN or infinite values. That is a clean training curve, but it is a partial run. Loss on FineWeb-Edu also cannot be compared with scores on public benchmarks, and the post does not report any.

Two optimizations the post highlights

Chunk size matched to the chip. The GLA kernel processes the sequence in chunks. AWS reports that using 16 chunks of 128 tokens, which lines up with the 128 partitions of the on-chip SBUF memory, improved throughput by 14%. With 32 chunks the scheduler overflowed, and with 8 chunks intermediates spilled to HBM. The backward pass is written with the 16 chunk steps statically unrolled so the compiler can handle it.

A cheaper Muon exchange. Laplace trains its 2D weight matrices with the Muon optimizer, using the reference implementation, and the remaining embeddings and norms with AdamW. Muon needs communication that standard data-parallel code does not. By switching to two all_to_all collectives, per-rank traffic dropped from 10.87 GB to 0.45 GB, a 24 times reduction. Throughput rose from 73,378 to 119,934 tokens per second, a 1.63 times gain, with a bit-identical loss curve. The post's takeaway: "Communication design matters for custom optimizers."

This work sits on the PyTorch path AWS has been building for its chips. We covered the claims around TorchNeuron and running PyTorch on Trainium without a rewrite and a claimed 2x speedup for TRL on Trainium. Laplace is a third data point, and a more demanding one because linear attention is not a stock transformer block.

How fast does it serve?

Served through an unmodified vLLM engine on one trn2.48xlarge, Laplace 3B produced 11,544 output tokens per second across 1,280 concurrent sequences. The post says outputs matched i14's reference implementation in BF16 for 48 of 48 steps across three prompts, a small but useful correctness check for a new backend.

Other serving details from the post:

• Decode takes 168.1 ms per step at 24 layers, with 96.9% of that time in the decoder layers.
• Eager decode scaled from 133 tokens per second on one replica to 6,772 on 64 replicas, where host CPU reached 82%.
• Compiling the decode step raised throughput 1.7 times.
• Under staggered arrivals, which resemble real traffic, compiled decode reached 3.52 times the per-replica throughput versus 1.46 times on fixed batches.

The host-CPU point is worth noticing. At 64 replicas the bottleneck moved off the accelerator and onto the CPU that feeds it, which is a common wall in small-model serving.

A row of small rounded slots with one green slot, illustrating the per-token cache slots a transformer needs and a linear attention model avoids

What does the post not tell you?

This is where a careful reader should slow down.

• No quality comparison. There is no benchmark table against a transformer of the same size trained on the same tokens. A smaller cache is only a win if quality holds, and linear attention has a record of struggling on tasks that need exact recall from far back in the context.
• Partial training. Results end at about 44 billion of 100 billion planned tokens.
• No release. No weights, no repository, no license. The post ends by pointing readers to i14 or an AWS account team.
• A vendor's frame. The post is published by AWS to show Trainium2 and TorchNeuron working. That does not make the numbers wrong, but nobody independent has reproduced them.
• Small model. Three billion parameters and a 2,048-token training length say little about how the approach behaves at frontier scale or at one million tokens.

On the quality question, it helps to know where the field currently stands. Our post on sliding-window attention beating linear attention in post-training covers evidence that pure linear attention loses ground once models are fine-tuned for instruction following. Hybrids that keep a few full-attention layers are a common compromise, and NVIDIA's work on caching SSM inputs rather than state in hybrid decoding and on Mamba-2 kernels on B300 points the same way: recurrent layers are getting faster, but they are rarely used alone.

Why this matters for people building with models

Three takeaways hold up even with the caveats.

Memory predictability is a product feature. If your service handles long, uneven prompts, a fixed per-request footprint makes capacity planning simpler. That is the strongest practical argument for architectures like this one, and it is the one the post supports with a measured 0.5% cost change.

Non-NVIDIA training is getting more usable. Training a custom optimizer and a custom attention kernel on Trainium2 with native PyTorch, and then serving through unmodified vLLM, is a lower bar than the older compile-your-graph workflow. Whether it is cheaper depends on pricing that the post does not discuss.

The KV cache is a design choice, not a law. Our note on coding-agent KV cache behavior shows how much cache effects shape agent costs today. Models that do not need one change that math, if they can match quality.

How can you try it?

You cannot download Laplace 3B based on the post. What you can do today:

1. Read the AWS post for the kernel and Muon details, which are the reusable parts.
2. Start from the GLA paper if you want to experiment with simple gated linear attention on your own hardware.
3. Use the chunk-size rule as a general heuristic: align your chunk length with your accelerator's on-chip memory partitions before tuning anything else.
4. Contact i14 or AWS if you want an evaluation of the model itself.

We will update this post if i14 publishes weights or benchmark results.

Related reading

• Sliding-window attention beats linear attention in post-training
• TorchNeuron: PyTorch on Trainium, claimed vs verified
• TRL v1.16 on AWS Trainium: the 2x claim
• vLLM push-based KV transfer on AWS P5en
• DeepSeek V4.1 Flash and KV cache memory
• NVIDIA ReplaySSM for hybrid decoding
• NVIDIA cuTile Mamba-2 kernels on B300

Figures are as published by AWS on October 20, 2026 and may change if i14 releases further results.