Beyond the KV cache: Training and serving i14’s Laplace 3B linear

Artificial intelligence research firm i14 and Amazon Web Services demonstrated the training and serving of Laplace 3B, a 3.028-billion parameter linear attention language model, on AWS Trainium2 hardware using native PyTorch through TorchNeuron Native. As detailed in AWS Machine Learning Blog's report, the architecture eliminates the traditional expanding key-value cache by compressing past conversational tokens into a recurrent state that stays fixed at 29.5 megabytes per sequence regardless of context length. On a single trn2.48xlarge instance featuring 16 Trainium2 chips and 64 logical NeuronCores, the model achieved a training speed of roughly 120,000 tokens per second on the FineWeb-Edu dataset and served 11,544 output tokens per second across 1,280 concurrent sequences.

This technical report examines how the engineering teams replaced custom GPU Triton kernels with standard PyTorch operators mapped to the NeuronCore on-chip memory hierarchy, redesigned distributed communication for the Muon optimizer, and handled the static unrolling of sequential backward passes. The design targets systems engineers, machine learning practitioners, and agent infrastructure builders who require deterministic accelerator memory consumption and high sequence concurrency for long-horizon context analysis without the memory exhaustion typical of standard softmax attention.

What Beyond the KV cache shipped

The project delivered Laplace 3B, a 24-layer decoder model where every layer swaps conventional quadratic softmax attention for a Gated Linear Attention mechanism. Standard transformer models compute pairwise scores across all sequence tokens, which forces memory use during training to scale quadratically at O(T2) while driving inference key-value cache allocations higher with every newly generated token. In contrast, Laplace 3B organizes processing around an updated matrix state, dropping training attention complexity to linear O(T) and maintaining an exact 29.5 megabyte inference state per sequence regardless of total context size.

Built for the AWS Trainium2 architecture, the release verified model execution on a trn2.48xlarge instance configured with 64 logical NeuronCores linked by NeuronLink-v3. Operating with a batch size of 524,288 tokens per step, a 2,048 sequence length, and a vocabulary size of 50,304 tokens padded to multiples of 64, the training process digested 100 billion tokens from FineWeb-Edu. During inference benchmarks on the same 16-chip Trainium2 node, the fixed-state engine delivered 11,544 output tokens every second across 1,280 simultaneous active sequences.

What changed for builders in Beyond the KV cache

For machine learning engineers, the architectural transition replaces variable-size accelerator allocations with static memory footprints that allow predictable multi-tenant capacity planning. In Laplace 3B, each token vector of hidden dimension 3,072 passes through LayerNorm and a depth-wise causal Conv1d of kernel width 4 with SiLU activation, yielding a four-gram context summary. A fused projection matrix then splits the signal across 16 attention heads into queries, keys, values, and gates, maintaining a per-head matrix state of 96 by 192 dimensions that writes outer products of keys and values adjusted by a static geometric decay schedule spanning token half-lives from 1 to 2,048 tokens.

To execute this architecture on AWS Trainium2, the Generative AI Innovation Center restructured distributed communications within MuonFSDP2. By inserting two all_to_all collectives into the training loop, the engineers reduced per-rank communication volume by 24 times, producing a 1.63 times increase in overall model throughput. Optimization responsibilities were partitioned between Muon, which handles 2.718 billion parameters across two-dimensional weight matrices, and AdamW, which updates 309.6 million parameters across embeddings and normalization layers with a 3,815-step warmup and cosine learning rate schedule.

How to install or upgrade Beyond the KV cache

Deploying Laplace 3B on AWS Trainium hardware requires configuring an environment with the native PyTorch Neuron software development kit and the beta release of TorchNeuron Native. Builders launch an Amazon EC2 trn2.48xlarge instance using the AWS Deep Learning AMI tailored for Neuron, then update the local compiler and runtime packages to ensure PyTorch operations map directly to the NeuronCore-v3 hardware.

pip install --upgrade torch-neuronx neuronx-cc --extra-index-url=https://pip.repos.neuron.amazonaws.com

To configure model execution within application scripts, developers initialize the model definition using PyTorch and specify the native compilation target:

python3 -m pip install git+https://github.com/i14/laplace.git

Environment setups pin the model loading and execution parameters to allocate the fixed recurrent state directly across the 64 logical cores:

export NEURON_RT_NUM_CORES=64
export TORCH_NEURON_NATIVE=1
python3 serve_laplace.py --model i14/laplace-3b --batch-size 1280

Persistent service configurations store the native backend flag and hardware initialization defaults directly inside the deployment environment to prevent automatic fallback to dynamic graph allocators:

echo 'export TORCH_NEURON_NATIVE=1' >> ~/.bashrc
source ~/.bashrc

Gotchas and compatibility in Beyond the KV cache

Adapting the custom linear attention architecture from GPU setups to Trainium2 required abandoning GPU-specific Triton kernels such as fused_chunk_simple_gla from the flash-linear-attention library. Trainium2 relies on a static ahead-of-time compiler and software-managed on-chip memory rather than hardware caches. Each NeuronCore-v3 contains a 28 MiB State Buffer partitioned into 128 segments alongside a 2 MiB partial-sum buffer and a 128 by 128 systolic Tensor Engine, requiring the GLA attention kernel to be rewritten using 16 chunks of 128 tokens to fit hardware dimensions without eviction penalties.

Because linear attention continuously compresses historical tokens into a fixed-size mathematical state rather than retaining raw token key-value pairs, engineers must thoroughly validate retrieval quality on tasks requiring needle-in-a-haystack lookups or verbatim quoting. Furthermore, compiling the sequential backward pass on Trainium2 required developers to statically unroll all 16 chunk steps of the recurrence to avoid exceeding compiler graph size and on-chip memory limits, precluding dynamic chunk lengths during training runs.

What to watch after Beyond the KV cache

The validation of Laplace 3B on Trainium2 provides an empirical baseline for running non-transformer architectures on non-GPU accelerator infrastructure. Because the recurrent state never exceeds 29.5 megabytes, development teams are tracking real-world deployments in long-running autonomous agent workflows, persistent workspace assistants, and large document streams that previously experienced out-of-memory serving crashes. The fixed footprint allows infrastructure operators to pack consistent sequence numbers per chip regardless of how long conversations continue.

Future engineering updates will center on compiler maturity and kernel tuning for linear attention mechanisms within the TorchNeuron Native backend. As AWS expands general access to Trainium2 clusters and refines static graph compilation for recurrent topologies, observers are watching whether hybrid models combining periodic softmax layers with linear attention can sustain Laplace 3B's 11,544 tokens per second serving density while matching full-attention benchmarks on precision-heavy data extraction tasks.