Beyond the KV cache: Training and serving i14’s Laplace 3B linear attention model on AWS Trainium2 with native PyTorch

Most large language models (LLMs) use transformer architectures with softmax attention. During inference, their key-value (KV) cache grows with the conversation or document context. This can reduce concurrency, increase serving cost, and force agentic applications to discard useful history. i14 designed Laplace 3B as a linear attention alternative. Its recurrent state stays at 29.5 MB per sequence, independent of context length. Working with AWS, i14 used native PyTorch support through TorchNeuron Native to bring this custom architecture to AWS Trainium2 . This preserved its Gated Linear Attention (GLA) design and Muon optimizer. On one trn2.48xlarge instance, Laplace 3B trained at about 120,000 tokens per second. The model served 11,544 output tokens per second across 1,280 concurrent sequences. This post explains the model design, the changes i14 made for Trainium2, and the training and inference results. About i14 i14 is an AI research and technology company that develops model architectures, components, and runtimes. Its work draws on AI, computational cognitive science, and research into how biological systems learn over time. i14 applies these ideas to models that can retain long-range context without allowing serving state to grow with every token. Why linear attention models? Transformer softmax attention compares tokens across the context. During training, this computation grows quadratically with sequence length. During inference, the KV cache grows as the model processes more tokens. Linear attention reorganizes the computation around a recurrent state that is updated for each token. Laplace 3B uses GLA to compress prior context into a fixed-size matrix state. This changes training attention from O( T 2) to O( T ) and keeps inference state fixed at 29.5 MB per sequence. Consideration Transformer softmax attention Laplace 3B with GLA Training attention Builds pairwise token scores with O( T 2) scaling Updates chunked recurrent state with O( T ) scaling Inference memory KV cache grows with context Recurrent state remains fixed in size Longer contexts Each sequence consumes more accelerator memory Context length does not increase recurrent-state allocation Implementation Uses widely optimized attention kernels Requires architecture-specific kernels and compiler tuning Linear attention compresses history rather than retaining every key and value. You must therefore evaluate retrieval quality for your workloads. The fixed state can benefit long-running agents, large-document analysis, and applications that need persistent memory with predictable serving requirements. How i14 adapted Laplace 3B for AWS Trainium2 The original implementation of Laplace 3B used GPU-optimized Triton kernels. The AWS Generative AI Innovation Center (GenAIIC) team first rewrote the GLA kernel with standard PyTorch operators for Trainium2. After comparing both approaches, the team chose the beta version of TorchNeuron Native, the native PyTorch backend for AWS Trainium , because it performed better. The implementation required three changes: Using native PyTorch for the GLA attention kernel. The team replaced the GPU-specific fused_chunk_simple_gla Triton kernel with standard PyTorch operators, using 16 chunks of 128 tokens to match the 128 partitions of the NeuronCore on-chip buffer. Redesigning how Muon communicates. Two all_to_all collectives in MuonFSDP2 cut per-rank traffic by 24 times and raised throughput by 1.63 times. Compiling the sequential backward pass. The team statically unrolled the 16 chunk steps of the GLA recurrence to fit the graph-size and on-chip memory limits of the compiler. The engineering challenges section describes each change in detail. Why AWS Trainium accelerators? AWS Trainium is a family of AWS machine learning (ML) accelerators designed for model training and inference. A trn2.48xlarge instance combines 16 Trainium2 chips, 64 logical NeuronCores, high-bandwidth memory, and NeuronLink-v3 connectivity. These resources support the distributed training and inference configurations described in this post. Three properties of Trainium2 shaped the implementation: software-managed on-chip memory (Challenge 1), a fast chip-to-chip interconnect (Challenge 2), and a static compiler (Challenge 3). The compiler traces the GLA recurrence as a fixed graph. Software-managed on-chip memory. Each NeuronCore has a 28 MiB State Buffer (SBUF) with 128 partitions and a 2 MiB partial-sum buffer (PSUM). The PSUM is wired to a 128×128 systolic Tensor Engine, alongside separate vector and scalar units. The SBUF is not a cache: the compiler schedules every HBM-to-SBUF transfer ahead of time, so the data is in place with no evictions or misses (Figure 1). Figure 1: NeuronCore-v3 internal architecture. Software-managed SBUF (28 MiB, 128 partitions), Tensor Engine (128×128 systolic array), Vector Engine, Scalar Engine, and GPSIMD cores. i14 and the Laplace 3B architecture Laplace 3B is the 24-layer i14 decoder, where every layer replaces softmax attention with Simple GLA. Figure 2 follows one layer from the input x t at the top to the output x t + 1 at the bottom. This section explains each module: normalization and convolution, the QKVG projection, the GLA mixer, and the feed-forward channel mixer. Normalization and convolution. Each token’s residual-stream vector x t ∈ R D ( D = 3, 072 ) passes through a LayerNorm and a depth-wise causal Conv1d of width k = 4 with SiLU. This gives every position a 4-gram summary before the mixer. The GLA mixer then does content-based retrieval over the whole history through its decaying state. QKVG projection. The post-conv output goes through one fused matrix W q k v g ∈ R 9, 216 × 3, 072 , where 9, 216 = H × (2 d k + 2 d v ) . It splits into four per-token streams: queries q t , keys k t , values v t , and a gate g t ( H = 16 , d k = 96 , d v = 192 ). The gate is used after the mixer. Figure 2: Laplace decoder layer architecture. Green = fixed (no params), gray = learned, yellow = recurrent state, blue = projection. GLA mixer. Each head h keeps a matrix state S t ( h ) of size d k × d v (96 × 192), a compressed, decaying record of past key-value pairs. At each token, it updates as: GLA recurrent-state update equation Here λ h is the fixed per-head decay, and k t v t ⊤ is the rank-1 outer product that writes the current key-value pair into the state. The readout queries the state: GLA state-readout equation Next, a self-augmentation step adds a gated copy of the value to each head’s output: GLA self-augmentation equation Here σ is the sigmoid and ⊙ is the element-wise product. The gate is the query dotted with a learned per-channel rescaling of the key, where a is the learned scale (aug_balance). At initialization, a = 0 , which gives a half-strength skip of the value. The per-head outputs pass through a Head RMSNorm within each decay tier, are concatenated, gated by S i L U ( g t ) , and projected by W o u t ∈ R D × D back to the residual stream. Decay schedule. The per-head decay rate λ h controls how fast a head forgets past tokens. It is fixed at initialization (not learned) and set by a geometric schedule. The schedule gives the 16 heads a spread of memory lengths and defines each head’s log-decay rate γ h : GLA decay schedule equations Decay schedule’s reach. For this model, T min = 1 and T max = 2, 048 . Head 0 therefore has λ 0 = 0.5 (half-life of one token), and head 15 has λ 15 ≈ 0.9997 (half-life of 2,048 tokens). Heads are grouped into tiers, each sharing a Head RMSNorm. Channel mixer (FFN). After the attention block and its residual, a second LayerNorm feeds a GLU (SiLU) FFN with inner dimension 8,192. Its three matrices are two up-projections W 1 , W 2 ∈ R 8, 192 × D and one down-projection W 3 ∈ R D × 8, 192 . A second residual follows. The state S ( h ) is fixed at d k × d v = 96 × 192 per head regardless of sequence length. As a result, inference is O (1) per token, and training is O ( T ) rather than O ( T 2