More tokens per dollar on AWS Trainium2: optimizing Qwen3.5-35B-A3B in days
Summary
New open models now arrive every few weeks. Many models are only optimized initially for the accelerators they have been developed on. But inference teams want to serve each one with more tokens per second and a lower cost per token. Together with AWS, yasp showed how fast that can happen on AWS Trainium2: our model optimization platform took Qwen3.5-35B-A3B, a modern 256-expert hybrid mixture-of-experts model, from a first working port to optimized serving in days. As a partner building with AWS’s new native PyTorch backend TorchNeuron, yasp makes sure new models don’t just run on Trainium, but run at their best.
- Up to 83% more output tokens per second from the same Trainium2 chip (37% to 83% across 1 to 8 concurrent users).
- Up to 45% lower cost per token on the same hardware, with each of eight concurrent users receiving 53 instead of 29 output tokens per second, which is well above human reading speed.
- Days, not months: from first port to optimized serving, with every kernel verified for correctness on the chip.
The challenge: models are tuned for the hardware they were built on
For teams running inference, the goal is simple: more tokens, at a lower cost, for every model their users ask for. The catch is that most models are tuned for very few architectures, most often the one they were developed on. On any other accelerator, a new model usually runs soon after release, but running at its best takes specialist kernel work for every model and target.
On AWS Trainium2, this matters in particular: the chip rewards kernels that closely match its architecture, and some operations in new model designs, such as the recurrent state updates in linear attention, do not map to it well out of the box. That is the gap yasp’s platform is built to close.
What we did
We picked Qwen3.5-35B-A3B as a realistic test for what teams want to serve in 2026: 40 layers that interleave Gated DeltaNet linear attention (30x) and grouped-query attention (10x), each followed by a mixture-of-experts block with 256 experts. With 67 GB of bf16 weights, it runs on one Trainium2 chip in a single trn2.3xlarge instance when distributed with tensor parallelism across the chip’s four logical cores, served with vLLM Neuron, which builds on vLLM with TorchNeuron (AWS’s native PyTorch backend for Trainium). We ported the reference implementation from HuggingFace following the official guidelines in order to serve the model on Trainium2 via vLLM. For that initial port, we added support for the recurrent state of hybrid models to vLLM-Neuron. We focused on the decode path as an example and pointed Gaia at the three blocks where most of the time is spent according to profiling. Gaia autonomously generated, tested and verified optimized kernels for each of them directly on Trainium2.
The results: more tokens, lower cost per token
More tokens per second
At eight concurrent users, one Trainium2 chip now serves 427 output tokens per second, 83% more than the initial port. Across 1 to 8 users, the gain ranges from 37% to 83%. Every user benefits as well: at eight concurrent users, each user now receives about 53 output tokens per second instead of 29. Improving throughput and interactivity at the same time means that Gaia moves the Pareto frontier as a whole (Figure 1).
More tokens per dollar
The hardware and its price did not change, only the kernels. So every throughput gain is a price-performance gain: each Trainium2 hour now produces up to 83% more output tokens, and the cost of each token falls accordingly, whatever your pricing.
In other words, the same workload needs up to 45% less Trainium2 time.
How Gaia got there
Gaia worked on the three blocks that dominate decode time in parallel. For one of them, it generated and benchmarked twelve candidate kernels in about eight hours on a single accelerator. Gaia’s kernels plug into the stack through TorchNeuron’s NKI custom operators, the extension point AWS provides for exactly this. Each block showed a different strength of the approach.
Specializing for the model’s shapes. AWS’s kernel library is built to serve many models well. For Qwen3.5’s specific expert shapes, Gaia generated specialized NKI kernels that make the most of Trainium2 for this model.
Finding the right boundary, together with engineers. For the linear-attention block, the best result combined a narrow kernel for the recurrent state update with a whole-layer kernel. One of our engineers did that combination, and what we learned went straight back into Gaia.
Learning from AWS’s reference kernels. Qwen3.5’s attention head dimension falls outside the shapes the library’s attention kernel supports. Gaia learned tiling and work-overlap strategies from AWS’s attention reference kernel and carried them over to this model, which produced the best attention result in the stack.
Same answers, verified on the chip
Speed only counts if the answers stay the same. Gaia’s harness checks every candidate kernel against the reference. At model level, the optimized model passes every hard check of the equivalence procedure in AWS’s Neuron tools, and matches or exceeds the initial port on the GSM8K benchmark within error.
What AWS and yasp say
New model architectures are arriving faster than ever, and our customers want to run them cost-efficiently on AWS Trainium as soon as they are released. This work shows that Trainium’s price performance advantage can reach new models like the one presented within days. yasp’s model optimization platform is a strong complement to the AWS Neuron stack, and a great example of what our partners can build on it. We look forward to continuing the collaboration. – Kamran Khan, Head of Business Development and GTM AWS/Annapurna ML
Teams should not have to choose between the newest model and the best economics. With Gaia, a brand-new model reaches its best on AWS Trainium in days, and our customers get more tokens for every dollar they spend on inference. – Stefan Krassin, CEO, yasp
What this means for you
If you serve open models at scale and want each new release on the hardware with the best economics as soon as it is out, the bottleneck is no longer the hardware. It is the kernels, and that gap now closes in days. yasp is one of the partners building with TorchNeuron to extend the Trainium ecosystem.
Deciding where to run inference? Book a Trainium feasibility check. Bring your model and the traffic you serve, and we will show you the throughput, latency and price-performance it reaches on Trainium2, measured on real hardware, not estimated.
Building on Trainium? Request access to our model optimization platform or book a technical session with our team.
Table of Contents
