We brought Mamba to Trainium — and mapped where custom kernels actually help
Non-NVIDIA chips like AWS Trainium can be cheaper to run, but only if someone writes fast code for them, one operation at a time. We had our research agent do exactly that, end to end. It wrote a hand-tuned kernel that beats a naive baseline by 1.57x, showed that the same kernel makes no difference inside a real model (because the compiler already does that job), and then brought up a model type Trainium has never supported and wrote the missing kernel it needed. The honest map of "where custom kernels help" is more useful than any single number.
Why "where kernels help" is the real question
Trainium is attractive on price, but every operation needs code written to its programming model. That is expensive engineering, so you want to spend it where it actually moves the needle. A hand-written kernel can beat a lazy baseline and still lose to the chip's own compiler, which quietly optimizes ordinary code. So the useful deliverable is not a single speedup; it is a map of which operations reward hand-tuning and which the compiler already handles.
We had an agentic research system (Primus, from Transformer Lab) run the whole loop: read the literature, form a hypothesis, write and debug the kernels through a real bring-up on live Trainium hardware, measure every result, and write the paper. Below is what it found.
RMSNorm: a clean win on the operation, and nothing in the model
RMSNorm rescales a vector so its values have a consistent size. It is memory-bound: the chip spends its time moving numbers, not computing, so the way to speed it up is to move fewer numbers. A naive version writes each intermediate result out to memory and reads it back; a fused kernel does the whole thing in one pass on-chip.
On the isolated operation at Llama-3.1-8B width, that fusion pays off cleanly: 112 microseconds versus 176 for the naive version, a 1.57x speedup, with the outputs bit-for-bit identical to a high-precision reference. A simple "how much memory moves" model predicts a ceiling around 2.5x, and the measured 1.57x sits just below it, exactly as you would expect once real overheads are included.
Then we put the same kernel inside a real Llama-3.1-8B block and measured again. Press the button to switch between the two views.
The win vanishes: 0.999x, a tie. Two reasons, both worth knowing. First, RMSNorm is only about 1% of the block's work; the big matrix multiplies dominate. Second, and more important, the Trainium compiler already fuses the standard RMSNorm, so our hand-written kernel matches it rather than beating it. This is not a failure of the kernel. It is a precise finding: when the compiler already does the fusion, doing it by hand buys nothing.
Mamba: bringing an unsupported model to the chip, and writing its missing kernel
Almost everything that runs on Trainium today is a transformer. Mamba is different. It is a "state-space model" whose core is not attention but a selective scan, a running calculation walked along the sequence:
ht = At · ht−1 + Bt · xt (update the running state, each step)
yt = Σn Ct · ht (read the state out)
On GPUs this exact operation needs a bespoke CUDA kernel. On Trainium there is none. So Mamba is the sharp test: can the system make an architecture the chip has never seen actually run?
It runs. A real Mamba block loads and executes on the chip. Out of the box the scan falls back to an unrolled, step-by-step path that does not scale to long sequences, so we wrote the kernel it was missing: a fused NKI selective-scan that keeps the running state on-chip across the whole sequence. It is numerically exact (matching a high-precision reference to about 3e-7, validated at Mamba's true 1536-wide state), and, most tellingly, it runs correctly inside the real model on the chip: swapping it in for the reference scan reproduces the model's output to a relative error of 7e-8.
That is the achievement: Trainium now has a correct, in-model selective-scan primitive, the analogue of the CUDA kernel that makes Mamba practical on GPUs.
We built the parallel scan — and it taught us something sharper
Our first scan kernel is correct but computes the recurrence sequentially, one time step after another, and that serial structure is slow: a strictly step-by-step calculation over tiny pieces underuses a massively parallel chip. So we wrote the fix the literature points to — a parallel scan.
The recurrence is associative, which means the whole sequence can be combined in a tree of log₂(L) steps instead of L sequential ones. The tricky part on Trainium is the "shift" each step needs; the toolchain wouldn't let us express it directly, so we did something a little sneaky: we laid the sequence along the chip's parallel lanes and turned every shift into a matrix multiply against a fixed shift matrix — which runs on the systolic array, the one thing Trainium is fastest at. It's numerically exact, validated at Mamba's true 1536-wide state.
Then we measured all three on the chip (sequence length 128, 1536-wide state):
The parallel scan is 1.14× faster than our own sequential kernel — the algorithmic idea works — but it's still 0.74× the compiler, so it doesn't win outright. And here is the sharper lesson. We tried the obvious speedup, running the shift matmul in the chip's native low precision, and it made no difference at all — proof that the matmul isn't the bottleneck. The real cost is that this style of parallel scan lowers the depth of the computation but raises the total work, and on this problem Trainium is limited by total work, not depth. Beating the compiler needs a work-efficient scan (the chunked approach of Mamba2 and FlashMamba), which Mamba's per-channel structure makes genuinely harder to write. That — not "just parallelize it" — is the real remaining step.
(One honest footnote on method: the low-precision version passed in the simulator and failed on the real chip, because the simulator doesn't model low-precision rounding. We now check reduced-precision kernels against real device output, never the simulator alone.)
The map: where custom kernels earn their keep
These results point the same way. The compiler is good; it already parallelizes and fuses ordinary code. So:
- Custom kernels reliably beat naive baselines (RMSNorm, 1.57x on the operation; the parallel scan, 1.14x over our sequential one).
- Custom kernels are necessary to enable operations the stack does not support (Mamba's scan now runs, correctly, only because we wrote it).
- Custom kernels do not beat the compiler when it already fuses or parallelizes the operation — not a naive fused RMSNorm in the model (0.999x), and not even a genuinely parallel scan (0.74x). Reducing the computation's depth isn't enough; the algorithm also has to reduce total work, which for a state-space scan means the harder chunked (Mamba2-style) formulation.
For anyone planning kernel work on Trainium, that is the actionable guidance: point scarce effort at work-efficient parallel algorithms (fused attention, chunked state-space scans, communication-fused collectives), not at naive fusions the compiler already captures — or at depth-only reformulations that quietly add work.
How we kept it honest and cheap
Two methodology notes. We validate correctness in Trainium's kernel simulator on a normal CPU, with no accelerator needed, which caught every algorithm bug for free. We measure latency on the real chip. We found the on-chip benchmark tool returns a slightly unreliable output (a read-back race) while its timing is stable, so we take correctness from the simulator and timing from the device, which is standard practice. Every number in the paper traces to a specific recorded job on a live trn1.2xlarge. The entire study cost a few dollars of chip time.
Code and the full reproduction package are available on request.