Training a linear attention model on AWS Trainium2 with native PyTorch

Researchers at i14 and the AWS Generative AI Innovation Center have successfully trained and served the Laplace 3B model using AWS Trainium2 accelerators. Published in October 2026, this work demonstrates how linear attention architectures can be optimized for specialized hardware without relying on custom GPU kernels.

The collaboration highlights a shift away from traditional transformer softmax attention, which suffers from quadratic scaling during training and growing memory costs during inference. By leveraging native PyTorch support through TorchNeuron Native, the team achieved significant performance metrics while maintaining a fixed memory footprint per sequence.

What happened

Large language models typically rely on transformer architectures that use softmax attention. During inference, these models maintain a key-value (KV) cache that grows linearly with the length of the conversation or document. This growth limits concurrency, increases serving costs, and often forces applications to discard historical context. To address these limitations, i14 developed Laplace 3B, a model based on linear attention principles.

Unlike standard transformers, Laplace 3B uses a recurrent state that remains constant at 29.5 MB per sequence, regardless of context length. The team worked with AWS to adapt this architecture for Trainium2 hardware. They utilized TorchNeuron Native, a beta feature that allows native PyTorch code to run directly on AWS Trainium accelerators. This approach preserved the model’s Gated Linear Attention (GLA) design and its use of the Muon optimizer.

The results showed substantial efficiency gains. On a single trn2.48xlarge instance, the model trained at approximately 120,000 tokens per second. During inference, it served 11,544 output tokens per second across 1,280 concurrent sequences. These figures demonstrate that linear attention models can achieve high throughput on specialized hardware when the software stack is correctly aligned with the underlying architecture.

How it works

Linear attention models reorganize computation around a recurrent state rather than pairwise token comparisons. In traditional transformers, attention scores grow quadratically with sequence length during training. Laplace 3B replaces this with a mechanism that updates a fixed-size matrix state for each token. This changes training complexity from O(T^2) to O(T) and keeps inference memory constant. The model compresses history into this state, which requires careful evaluation of retrieval quality for specific workloads.

Adapting this model for Trainium2 required three specific engineering changes. First, the team replaced GPU-optimized Triton kernels with standard PyTorch operators. They structured the GLA kernel to use 16 chunks of 128 tokens, aligning with the 128 partitions of the NeuronCore on-chip buffer. Second, they redesigned communication for the Muon optimizer. By using two all-to-all collectives instead of gather-and-broadcast methods, they reduced per-rank traffic by 24 times and increased throughput by 1.63 times. Third, they statically unrolled the sequential backward pass to fit the compiler’s graph-size and memory limits.

The Trainium2 architecture played a crucial role in these optimizations. Each NeuronCore features a software-managed State Buffer (SBUF) of 28 MiB with 128 partitions. Unlike a cache, the compiler schedules every data transfer in advance, ensuring no evictions or misses occur. By mapping each 128-token chunk to one partition, the team achieved a 14 percent increase in measured throughput. This alignment between the model’s chunk size and the hardware’s physical partitions was key to unlocking performance.

Key details

• Laplace 3B maintains a fixed recurrent state of 29.5 MB per sequence, independent of context length.
• Training on one trn2.48xlarge instance reached approximately 120,000 tokens per second.
• Inference served 11,544 output tokens per second across 1,280 concurrent sequences.
• The implementation used TorchNeuron Native to run standard PyTorch operators on Trainium2.
• Communication traffic for the Muon optimizer was reduced by 24 times using all-to-all collectives.
• The model uses 16 chunks of 128 tokens to align with the 128 SBUF partitions on NeuronCores.

Why it matters

For engineers building AI products, the growing cost of KV caches is a major bottleneck. As contexts get longer, memory usage spikes, limiting the number of concurrent users a server can handle. Linear attention models like Laplace 3B offer a path to predictable serving costs. Because the state size does not grow with context, applications can retain long histories without proportional increases in memory allocation. This is particularly valuable for agentic workflows and large-document analysis where persistent memory is critical.

Furthermore, this work demonstrates that specialized hardware like Trainium2 can support novel architectures without requiring low-level kernel rewrites. By using native PyTorch through TorchNeuron Native, developers can leverage hardware-specific optimizations while staying within a familiar framework. The success of the all-to-all communication strategy for the Muon optimizer also provides a blueprint for distributed training on sharded systems. It shows that reducing communication overhead can yield significant throughput gains, even when computational FLOPs remain unchanged.

What you can do

• Evaluate linear attention models for use cases requiring long context windows and predictable memory usage.
• Test TorchNeuron Native if you are deploying custom architectures on AWS Trainium instances.
• Align data chunk sizes with hardware buffer partitions to maximize throughput on specialized accelerators.
• Consider all-to-all communication patterns for distributed optimizers to reduce network traffic bottlenecks.
• Monitor retrieval quality when switching from softmax attention to state-compression mechanisms.
• Profile your inference engine to identify if KV cache growth is limiting your concurrency levels.