i14 and AWS Train Laplace 3B Linear Attention Model on Trainium2

i14 and AWS trained and ran the Laplace 3B linear attention model, which offers a solution to the key-value cache growth problem in transformer architectures, on AWS Trainium2 with native PyTorch support.

Most large language models use transformer architectures and softmax attention mechanisms where the key-value (KV) cache grows as context length increases. Designed by i14 to overcome this problem, Laplace 3B was developed as a linear attention alternative that keeps the recurrent state at a constant 29.5 MB per sequence. As a result of the collaboration with AWS, this custom architecture was adapted to AWS Trainium2 hardware with native PyTorch support delivered via TorchNeuron Native, preserving the Gated Linear Attention design and the Muon optimizer.

In tests conducted on a single trn2.48xlarge instance, Laplace 3B achieved training speeds of approximately 120,000 tokens per second and delivered 11,544 output tokens per second across 1,280 concurrent sequences. During this integration process, it was noted that adapting the GLA attention kernel to standard PyTorch operators, redesigning Muon communication with all_to_all collectives to reduce network traffic, and statically enabling the backward pass to comply with compiler boundaries were required.