TL;DR: We optimized three models (Qwen3.5-4B, Qwen3-30B-A3B, and Gemma4-31B) on Amazon’s Trainium2 accelerator and achieved a 1.4–6.8x speedup in all models. Our optimizations cover framework-level and kernel-level improvements that collectively make inference serving cost-competitive with NVIDIA H100s.

Amazon Web Service’s (AWS) Trainium accelerator is a systolic array-based chip meant for improving the efficiency of AI workflows. Over the past few weeks, the Makora team has been investigating the feasibility of inference serving on Trainium2, the second generation of AWS' chips. In this post, we detail aspects of our serving stack, optimization strategies, and achieved throughput gains on Trainium2. Crucially, we show that in certain regimes, we are able to serve models at a lower cost per million output tokens compared to NVIDIA H100 GPUs in some settings. Our performance optimizations make Trainium2 a meaningful alternative to GPUs in production inference workflows.

Trainium2 Serving Stack Primer

We highlight three aspects of serving models on Trainium2: the chip itself, AWS Neuron, and the vLLM-Neuron plug-in.

1. Trainium2 Chip - The Trainium2 chip is a specialized accelerator that is comprised of eight NeuronCore-V3 cores which function as the compute units of the chip. Each NeuronCore-V3 contains four engines and an on-chip SRAM that is split between SBUF, a working scratchpad, and PSUM, an accumulation buffer. The tensor engine is a 128x128 systolic array specialized for tensor operations such as matrix multiplications and convolutions. The vector engine is built for optimized vector computations such as LayerNorm while the scalar engine supports element-wise operations. Finally, the GPSIMD engine can be used for programming and executing custom operators. Beyond the eight NeuronCore-V3, a Trainium2 chip contains 96 GB of HBM and DMA engines to move data between off-chip HBM and on-chip SRAM. A depiction of the architecture is included below*.

   

2. AWS Neuron - AWS Neuron is the SDK for the entire line of Trainium chips. To understand Neuron, it is important to note that the Trainium stack employs ahead-of-time (AoT) compilation. Unlike GPU stacks, where kernels are dispatched dynamically by a hardware scheduler at runtime, Neuron resolves the complete execution schedule before the first token is ever served. A model enters the stack through a framework frontend (PyTorch via torch-neuronx), which traces the model into a computation graph. The Neuron Compiler (neuronx-cc) consumes this graph, together with any hand-written Neuron Kernel Interface (NKI) kernels, and lowers it into a Neuron Executable File Format (NEFF). The NEFF emits a static schedule that includes instruction streams and cross-engine dependencies for model execution. Finally, the Neuron Runtime loads and executes the NEFF, allowing end-to-end serving of the model with a static schedule.

   

3. vLLM-Neuron (Beta) - Our serving engine is vLLM-Neuron, AWS's beta vLLM plug-in for Trainium2 and Trainium3. From upstream vLLM, the engine inherits the OpenAI-compatible serving and the continuous-batching machinery. However, the plug-in integrates Neuron-specific machinery shaped by the constraints of AoT compilation. For example, the scheduler is aware that each batch shape corresponds to a pre-compiled executable, and the KV cache is managed to match static compiled layouts rather than PagedAttention's dynamic block tables. As of right now, the plug-in ships two supported model families, namely GPT-OSS and Qwen3VL. Other models must be onboarded by implementing the model and its components before registering the model with vLLM-Neuron.

Performance Optimizations on Trainium2

We now detail the three models that we served on Trainium2 and the optimizations we implemented to improve inference efficiency. Note that in each case, baseline performance refers to the out-of-box performance on vllm-Neuron for Trainium2.

Qwen3.5-4B

Qwen3.5-4B is a dense, hybrid model that interleaves full attention layers with Gated Delta Net (GDN) layers. Our serving configuration uses a single Trainium2 chip and tensor-parallel configuration TP=4 (the eight NeuronCore-V3s are grouped into sets of two). When registering the model weights, the GDN layers specifically require a custom implementation as it is not a component of the currently shipped models on vLLM-Neuron. We register a sharded version of GDN that ensures that each rank only receives a fraction of the heads in GDN rather than replicating the whole layer on every rank. This default, naive implementation of GDN would result in significant wasted compute as each rank would have identical output.

We additionally implement three custom kernels for GDN on Trainium2 that improve end-to-end inference efficiency.

1. Packing recurrent and convolutional state - Each GDN layer has both a recurrent and convolutional state that are used in computing token contributions. By default, these states are moved from HBM into SBUF by means of identical pages, a consequence of the vLLM-Neuron plumbing that is used for full attention layers. However, unlike full attention layers where the K and V tensors are identically sized, the convolutional state is an order of magnitude smaller than the recurrent state as it only maintains contributions from the last four tokens at any given step. This results in the convolutional state being moved in a page more than 10x larger than necessary, causing a significant amount of DMA waste. Instead, we extend the size of the page being used to move the recurrent state and pack the recurrent state and convolutional state into a single page. By moving only this page into SBUF, we cut down DMA overhead for every GDN layer throughout the decoding procedure.

2. Zero-pad in-projection - As the tensor engine is a 128x128 systolic array, it performs multiplication in 128-wide tiles. Each GDN layer begins with an in-projection which generates Q, K, and V matrices along with gating parameters z, b, and a for every head. This in-projection is particularly large, even with column-parallel TP=4 for Qwen3.5-4B, each in-projection is 3088x2560. Notably, 3088 leaves a remainder of 16 when divided by 128, which makes the in-projection unfriendly for Trainium2’s tensor engine. We zero-pad the in-projection to a size of 3200x2560, emitting 25 exact 128-row tiles in place of 24 full tiles and one 16-row remainder. Note that this padding does not reduce tensor engine work and actually increases DMA bytes. Instead, it allows for efficient pipelining and does not require a separate instruction for the ragged shape of the original tensor. With 25 identical tiles, every iteration of the projection loop has the same instruction shape, SBUF footprint, and timing, so the compiler can overlap each tile's weight DMA with the previous tile's matmul. In contrast, a trailing 16-row tile breaks the pattern at the loop's edge, forcing a differently-shaped instruction whose loads and compute don't interleave cleanly with the pipeline around it. Optimizations such as these are highly specific to Trainium2’s architecture and compiler stack demonstrating the need for hardware-aware optimizations.

3. Fused GDN Kernel - The GDN block's decode path is a chain of small, dependent operations which, when compiled naively, materializes its intermediate result in HBM and reloads it for the next stage. On Trainium this round trip is expensive as every spill and refill is an explicit, compiler-emitted DMA transfer. Unlike GPUs which operate with a memory coalescer, this regime is dominated by descriptor generation and per-descriptor cost rather than pure memory bandwidth. We instead implement the block as a single fused NKI kernel that holds intermediate states in SBUF. Here, rather than moving the recurrent state as an independent small transfer, the kernel lays the per-head states contiguously and moves them as a small number of large DMAs, replacing many short descriptors with few long ones and keeping the DMA engines in their bandwidth-efficient regime.

We achieve the following speedups relative to the unoptimized serving path. Our optimizations improve throughput by 1.6–6.8× across concurrency levels, with the largest gains at higher concurrency.

Qwen3-30B-A3B

Another optimization target was Qwen3-30B-A3B, a sparse Mixture-of-Experts (MoE) model with 8 activated and 128 total experts per layer. For Qwen3-30B-A3B, we utilize a number of framework-level optimizations to improve performance at low concurrency.

1. Selective expert decode - A statically compiled MoE graph cannot branch on router output, so every decode step loads and computes all 128 experts per layer even though only 8 receive tokens. At low concurrency, decode is dominated by weight movement. In contrast, we only load routed experts, reducing DMA traffic and unnecessary compute.

2. On-device greedy sampling - The baseline ships the full batch × vocab logits tensor to host memory each step, samples there, and returns token IDs, a round trip that is particularly expensive as it requires HBM ↔ DRAM communication. Sampling on-device shrinks the transfer to a few token IDs and removes the host from the critical path.

3. EAGLE-3 Speculative Decoding - EAGLE3 attaches a lightweight draft head that autoregressively proposes several candidate tokens per step from the target model's hidden features. The target model then verifies the whole candidate sequence in a single forward pass, accepting the longest correct prefix. Because low-concurrency decode is memory-bound, verifying k tokens costs nearly the same as generating one. We would like to note that methods like speculative decoding and MTP can cause ragged shapes within a batch at high concurrency due to varied token acceptances, a challenge that may cause slowdowns due to Trainium’s AoT setup. However, at low-concurrency, these methods still prove to be quite useful.

At concurrency=1, we accomplish the following speedups from each optimization. Our best configuration achieves a 4.5x interactivity gain over the default serving setting.

Gemma4-31B

Finally, we optimize Gemma4-31B, a dense model that utilizes sliding window attention (SWA) in place of certain attention layers with a 1024-token window. Again, we apply framework-level optimizations and bug fixes to demonstrate that serving configuration changes can be used to significantly improve model throughput.

1. Enabling fast attention kernel - The vLLM-Neuron beta uses a fast attention kernel that is gated by a literal MAX_HEAD_DIM set to 128. Gemma4-31B, which has a per-head dimension of 256, trips the gate and defaults to a slower path. Changing the literal enabled Gemma4-31B to use the faster kernel instead.

2. Reducing block size - Reducing the block size of KV cache transfers from 32 → 16 enables transfers to align better with each NeuronCore-V3’s internal SBUF tiling. While it marginally increases DMAs, the tiling advantages and the streaming of transfers still results in faster overall throughput over the default serving setting.

3. SWA redundant computation - Gemma4's local attention layers are defined to attend only over a 1024-token sliding window, but the baseline attention path treated every layer as global. As a consequence, tokens were attending over a >1024 token horizon, increasing memory accesses and slowing down inference. Fixing this error ensured that SWA layers were attending over limited context as intended.

We test multiple serving configurations to display the advantages of these optimizations. Across sequence lengths and concurrency settings, the optimized pipeline delivers roughly 1.4–2.8× higher throughput than the default configuration.

GPU Price Comparison

Our optimizations show promise in making model serving cost effective even relative to highly-optimized H100s. AWS provides on-demand pricing for Trainium2 and H100 instances, with a single accelerator being priced at $2.23/hr and $5.19/hr respectively. Using the measured interactivity numbers and the on-demand pricing, we calculate the cost per million output tokens. The following graphs represent the relative costs of our serving configurations against serving on H100s.

While the unoptimized path can be up to 3x more expensive than the H100 baseline, our optimizations make serving on Trainium2 far more competitive. In the case of Gemma4-31B, the Trainium2 optimized path is cheaper at every concurrency. For Qwen3.5-4B, serving is competitive with the H100 baseline at low concurrency, but the H100 scales to high concurrency more elegantly, maintaining high interactivity.

Tooling

Our development workflow builds on Makora's existing agentic tooling and utilizes MakoraGenerate for implementation and MakoraOptimize for performance iteration. On top of that foundation, we found a set of Trainium-specific skills and tools essential. Some of these tools are provided by the AWS Neuron team and others developed for our custom workflows. We describe the most useful ones here

1. Neuron Agentic Development - We use AWS’s Neuron Agentic Development package to help generate and profile custom NKI kernels. This open source package includes skills for writing, debugging, and profiling and agents that combine these skills for broad NKI development. In particular, we used the neuron-nki-agent and neuron-nki-writer-agent in our kernel writing process.

2. neuron-gdn-profiling Skill - In registering the GDN layer, we create a custom skill that ingests NEFFs, throughput numbers, and kernel design to determine bottlenecks for serving. The skill also correctness gates changes to make sure that kernel changes produce identical hidden states and final logits in multiple serving configurations.

3. dma-traffic-predictor Skill - As mentioned before, memory accesses do not function the same on Trainium and GPUs. As Trainium does not have a memory coalescer, each memory access requires the generation of a descriptor which is often serialized. The result is that many small DMAs can be significantly slower than a few large DMAs even if they carry the same number of bytes. We develop a skill that ingests kernel tiling strategies and the resulting DMA traffic from multiple layers to more effectively utilize memory bandwidth when developing kernels. This was especially useful in writing our GDN kernels as similar tiling strategies often produced vastly different numbers of DMAs even if they carried the same overall numbers of bytes.

Makora is committed to optimizing model inference on a variety of hardware. Our experiments on Trainium2 demonstrate that these accelerators can be a meaningful alternative to serving on GPUs and that our performance engineering strategies are broadly useful. As inference increasingly runs on heterogeneous systems, the techniques developed here form a playbook we intend to carry across the emerging hardware landscape. Look out for new models in our inference service that will be run on AWS Trainium2 servers! app.makora.com

*Taken from AWS Neuron Docs - https://awsdocs-neuron.readthedocs-hosted.com/en/latest/about-neuron/arch/neuron-hardware/trainium2.html