Optimizing Arcee Trinity-Mini for AWS Trainium2

We at Minima partnered with AWS to find out how far a model can be optimized when every optimization is built specifically for AWS Trainium2 and AWS Trainium3. AWS Trainium is a portfolio of AI chips designed to deliver high performance at lower cost for modern deep learning AI training and inference workloads. We ran all of our benchmarks and the full optimization pipeline directly on Trainium2, where the compressed model is 3.3× smaller than the original, fits twice as many replicas on one instance and is projected to deliver up to 2.6× the tokens. Trainium3 adds native 4-bit (MX) matmuls, so the 4-bit weights can run without the FP8 expansion that Trainium2 needs.

Our test model was Arcee’s Trinity-Mini, a 26B-parameter mixture-of-experts model. Here are the results.

Below we describe the attention change, the 4-bit quantization and the NKI kernels, then give the full evaluation, including where the compressed model falls short.

The Model and the Target

Trinity-Mini routes every token through 8 of 128 experts per layer, so only 3B of its 26B parameters are active at a time. It has 32 layers with a hidden size of 2048, and grouped-query attention with 32 query heads sharing 4 key/value heads of 128 dimensions.

The layers come in two kinds. Twenty-four use a 2,048-token sliding window. The other eight (every fourth layer: 3, 7, …, 31) attend over the whole context and carry no positional encoding. Those eight are the only place where the key/value cache keeps growing with context length.

The target is an Amazon EC2 trn2.48xlarge instance, which has 16 Trainium2 devices with 96 GB each, exposed as 64 logical cores of 24 GiB. Most of our design choices follow from two hardware facts. The tensor engine is fast at FP8 but has no 4-bit (MX) matmul, which is supported in Trainium3, resulting in even faster expert layers there: 4-bit weights move half the bytes of FP8, and data movement is what limits our expert kernel. And each core has 24 GiB, so a bf16 Trinity-Mini needs four cores just to load.

Our target was to keep every benchmark within one point of the bf16 original. We made three changes:

Structural Change: Latent Attention on the Global Layers

We replaced the key and value projections of the eight global layers with a 256-wide latent. After training, that costs 0.3% perplexity and cuts those layers’ cache per token by 3.9×.

The projections

In the original, each global layer projects the residual stream to keys and values for 4 heads. MLA factorises that pair of matrices through a narrow latent and caches only the latent.

Only the 256-wide latent goes into the cache; kv_b_proj expands it back into keys and values for the 4 heads.

original                              MLA
k_proj : 2048 -> 512                  kv_a_proj : 2048 -> 256    (this is what gets cached)
v_proj : 2048 -> 512                  kv_b_proj :  256 -> 1024   (keys and values for 4 heads)
2.10M parameters, 1024 values/token   0.79M parameters, 256 values/token

The two new matrices are initialised from the old ones by a whitened joint SVD with a least-squares refit, calibrated on 65,536 tokens of fineweb-edu. The converted model already works without training; below we measure how much it loses.

Where the key norm goes matters most

Trinity-Mini applies a per-head RMSNorm to its keys. After factorisation the norm has to go somewhere, and where it goes matters more than the latent width. These are perplexity penalties straight after conversion, with no training:

We kept the original norm. With the norm kept, wider latents lose less: +25% at 128 dimensions, +9.3% at 192, +3.1% at 256, +0.3% at 512. We shipped 256.

Recovery training

Training is self-distillation. The converted checkpoint keeps a frozen copy of the original attention next to the new one, so a single model produces both teacher and student logits. Everything else (experts, router, embeddings, the 24 sliding layers) is shared, so the teacher is exactly the original model. The loss is KL divergence over the full vocabulary.

We train only these parameters on purpose: when we also trained the query, output and gate projections, training diverged at every learning rate from 1e-5 up.

After the recovery run we continued twice at a learning rate of 3e-5 on 8k sequences, each run stopped by a plateau watcher.

How it was evaluated

• Bucketed perplexity. 16 windows of 32,768 tokens from the PG-19 test split, scored separately for positions 0-2k, 2k-8k and 8k-32k. Positions past 2,048 lie outside the sliding window, so the far buckets test the converted layers directly.
• Passkey retrieval. A random 5-digit key hidden at five depths in contexts of 4k to 32k tokens, three trials per cell, exact match.
• Task accuracy. HellaSwag, ARC-Challenge, MMLU and GSM8K at full size through lm-eval.
• Noise floor. The same bf16 model scored 86.35, 85.97 and 85.5 on repeated GSM8K runs, so we treat single-run differences under about 1.5 points on that task as inconclusive.

Quantization and Error Compensation

Expert and MLP weights go to 4-bit MXFP4, which takes the checkpoint from 48.7 GB to 14.85 GB. With error compensation the cost against the bf16 MLA model is about half a point per core task.

What is quantized, and what is not

MXFP4 stores each weight as a 4-bit float (E2M1) and shares one 8-bit power-of-two scale across every 32 values: 4.25 bits per parameter. We measured the sensitivity of each tensor kind before deciding where to apply it.

The output projection cannot go to 4 bits: its outlier columns don’t fit a shared power-of-two scale per 32 values. So we keep all of attention in bf16.

Error compensation

Plain round-to-nearest loses 2.2 points on GSM8K. Our final recipe recovers most of that by combining six techniques.

1. Sequential calibration. Layers are quantized in order, and calibration activations are propagated through the layers already quantized, with FP8 activation rounding switched on. Each layer is fitted to the inputs it will really see.
2. GPTQ error feedback. For each weight matrix we build the Hessian of its inputs, quantize one column at a time, and push each column’s rounding error into the columns not yet quantized. Blocks of 128 columns, 1% damping. All experts of a layer are processed as one batch.
3. Activation ordering. Columns are quantized in descending order of their Hessian diagonal, so the most important columns are rounded first, while later columns can still absorb their error.
4. Activation-aware block scales. The shared exponent of each 32-value block is chosen to minimise Hessian-weighted error rather than taken from the largest magnitude.
5. Bias correction. Each output channel gets a bias equal to the mean output error, (W − Q) times the mean input. It is stored in bf16 and costs 8 KB per expert.
6. A large, balanced calibration pool. 8.4M tokens drawn from held-out rows of the training mix, so that all 128 experts of every layer see enough data. An expert with fewer than 16,384 calibration tokens falls back to plain rounding.

All quantized rows include FP8 activations, FP8 KV cache and the FP8 LM head.

What did not help

• Block-wise reconstruction. Tuning each layer to match its bf16 output lowered layer error from 0.90% to 0.77% and left perplexity unchanged, yet every task got worse (MMLU by another point).
• Activation smoothing. Moving scale between activations and weights made the quantization error proxy worse (0.0062 against 0.0057) and cost 0.3 to 0.7 points on three of four tasks.
• Fewer active experts. Routing to 7 experts instead of 8 costs about 1.1 points on GSM8K with no retraining.

The finished checkpoint is stored packed: 4-bit codes and 8-bit scales for every MLP matrix, FP8 for the LM head, bf16 for the rest. Loading the packed file reproduces the unpacked model’s logits exactly.

Custom Kernels

With custom NKI kernels, the compressed model decodes 29% faster than the original on Trainium2. Before that work, it ran correctly but 20% slower, with no memory saving. Closing the gap took a lossless weight conversion, one patched vendor kernel and three NKI kernels of our own.

The serving stack is vLLM with the Neuron plugin, neuronx-cc 2.27 and NKI 0.6. Unless noted, numbers are for one replica on 8 logical cores.

Prefill rebuilds keys and values and stores only the latent. Decode works on the latent directly, so its time per token stays flat from 2k to 32k of context.

History

Expert weights: FP8 without loss

Trainium2 cannot multiply MXFP4 directly, and decoding it inside the kernel is slower than not compressing at all. The tensor engine is fast at FP8, and MXFP4 converts to FP8 exactly if each row gets one power-of-two scale:

s_row = 2^(e_max - 5)        e_max = largest block exponent in the row
w_fp8 = w_mxfp4 / s_row      exact when the row's block exponents span <= 13 octaves

The headroom of 5 keeps the largest value at 192, inside FP8’s range on Trainium2. Across all 15.9M expert rows of the shipped checkpoint the largest exponent spread is 7, so the conversion is exact for every row. The loader does it once at startup and we verified the result bit for bit. Expert weights on device drop from 48 GB to 24 GB.

The fp8x8 expert kernel

The vendor kernels upcast FP8 weights to bf16 before multiplying, which saves memory but not time. Ours keeps both operands in FP8:

1. Expert inputs are quantized to FP8 inside the kernel, with one scale per token.
2. Gate and up projections run as FP8 × FP8 matmuls in the tensor engine’s double-row mode, twice the bf16 rate.
3. The two scales (per weight row, per token) are multiplied back in as results leave the accumulator.
4. The down projection runs transposed so its scale lands on the output axis.

At 3.87 ms per layer for 2,048 tokens the kernel is limited by data movement, not by the multiplier.

A prefill kernel walks a fixed number of expert blocks, and the blocks that received no tokens are wasted work. Looping only over real blocks at run time gave 3.32 ms but hit two toolchain limits: a compiler error whenever a runtime-length loop shares a graph with the segmented attention kernel, and a runtime hang when the two halves of a logical core see different trip counts. The predicated kernel avoids both. Its loop length is fixed at compile time, and groups of four blocks past the always-used prefix are each guarded by a zero-or-one-trip predicate. That reached 3.02 ms. It still cannot share a graph with segmented attention, so it serves the first chunk of every prompt and the static kernel serves the rest.

Absorbed latent attention

If decode rebuilt keys and values from the latent at every step, the smaller cache would save nothing. Instead, the up-projection can be folded into the query and output side, where it touches one token rather than the whole history:

score_i = q . (W_k c_i)        =  (W_k^T q) . c_i      fold W_k into the query
output  = sum_i a_i (W_v c_i)  =  W_v (sum_i a_i c_i)  fold W_v into the output projection

Per step the kernel transforms one query, takes dot products directly against the cached latents, and sums in latent space. The history is read once and never expanded. Decode time is flat from 2k to 32k of context.

Keeping the original key norm made this harder. The norm gives every key its own data-dependent scale, which cannot be folded into a fixed matrix. The vendor attention kernel takes two equal-width tensors with one scale each, so the first port stored the latent twice, once pre-scaled as keys and once raw as values. That made the cache larger than the original’s.

Our decode kernel stores the latent once and recomputes the scale inside the kernel. It walks the context in 128-key tiles with one indexed gather per tile, over a loop whose length is fixed per context bucket (2,304, 4,352, 8,448 or 16,640 tokens). We use buckets because a runtime-length loop executes its body serially at 250-300 ns per instruction, 4-8 microseconds per tile, where the static loop takes 1.7-2.

Prefill does the opposite. With thousands of new tokens attending to each other, expanding keys and values is cheap, while the latent form costs 4-8× more per layer. So prefill reconstructs keys and values, runs ordinary 128-dimension attention, and writes the latent to the cache for decode to use.

Bugs that only showed up under batching

• Wrong at batch 8 and 16, right at 1, 4 and 32. The decode kernel passed every single-sequence check. Only comparing each sequence of a batch against its own reference exposed it.
• Cache pages shared between requests. The single-slot layout mapped a latent page onto bytes that also held another request’s sliding-window keys. Invisible at batch 1.
• A write that never happened. The latent was written through a reshaped view of the cache. The Neuron backend rewrites in-place operations, but only redirected later uses of the view, not of the cache itself, so reads saw whatever the previous request had left. It surfaced as perplexity 9.26 instead of 8.91, not as a crash.

Each accepted kernel change had to pass the same gate on device: next-token agreement with the GPU reference, logit cosine, sampled perplexity and passkey retrieval at 14k tokens. Three faster configurations failed it and were dropped.

Results

The compressed model keeps the original’s accuracy: across ten standard benchmarks it averages 72.46 against 72.75, a difference of 0.3 points, and no benchmark moves by more than one point. It also decodes 29% faster on Trainium2 and holds a fifth less device memory. It is behind on long chain-of-thought tasks and on time to first token.

Accuracy: standard benchmarks

Full-size runs on GPU, original bf16 against the shipped compressed model (MLA, MXFP4 weights, FP8 activations, KV cache and LM head).

Fidelity on the device

The Trainium2 build was checked against the same weights running on a GPU.

Scoring identical positions on both platforms, the mean difference in negative log-likelihood is 0.005 nats. The serving stack adds no measurable error of its own.

Speed

One replica per model on the same trn2.48xlarge, 256 output tokens per request, prefix caching disabled so that every request pays its full prefill.

The gain comes from decode. Prefill does the same expert work plus the latent projections and FP8 casts, and the fastest expert kernel cannot yet run in every prefill graph. The deficit is a fixed cost per request, while the gain is 12.4 ms on every generated token, so total request time favours the compressed model after about 2 output tokens on a 2k prompt, 7 on 8k and 24 on 32k.

The original gets its first token out 72 ms earlier, but each later token takes 12.4 ms longer, so the compressed model is ahead after a few tokens and the gap grows from there.

At 32 concurrent requests the limit is our batched decode kernel for the latent layers, not the experts: the result is the same with bf16 experts.

Memory

The numbers come from neuron-monitor with both servers loaded and idle. Scratch is the largest single difference: working in a 256-wide latent with FP8 experts needs a twentieth of the activation memory. On the other hand, the compressed build carries about 16 GiB more compiled graph code at 4 cores, because it ships more kernel variants.

On four cores the original sits within 1.5 GiB of the 24 GiB limit per core, and under repeated runs of the 8-concurrent test we saw it fail with a device allocation error. The compressed model ran the same test with 6 GiB per core to spare.

Projected per trn2.48xlarge: each 2-core compressed replica is assumed to keep the 1.29× per-replica rate measured on 4 cores.

We expected the KV cache to shrink on the device too, and it did not. The original’s 4 KV heads are split across cores, while the latent is replicated on every core, so with the vendor attention kernel the compressed cache holds 730k tokens per replica against the original’s 1.26M. Our single-slot decode kernel closes that gap. An FP8 latent, which would put the compressed model ahead, works in prefill but is blocked in decode by a runtime fault we have not yet resolved.

Why Trainium3 Should Do Even Better

We have not run Trinity-Mini on Trainium3 yet, but its hardware targets the exact limits we hit on Trainium2:

• 4-bit weights in hardware. Our recipe uses 4-bit weights with FP8 activations. Trainium3 accelerates this W4A8 format in hardware, which AWS says doubles the effective weight-loading rate. On Trainium2 we expand the 4-bit weights to FP8 at load, and our expert kernel is limited by data movement, so this is where we expect the largest gain.
• 1.7× the memory bandwidth. 4.9 TB/s per chip. Decode reads the weights and the latent cache at every step, so more bandwidth means faster tokens.
• 1.5× the memory. 144 GB of HBM per chip, against 96 GB on Trainium2, leaves more room for KV cache and replicas.
• Up to 2× the MXFP8 compute. This helps prefill, the one place where the compressed model is behind today (time to first token is 5–16% slower).
• No bring-up workaround. The vendor MX kernels for MoE already run on Trainium3, so expert weights don’t need to be expanded at load.

Trainium3 figures from AWS Trainium.

Working with AWS

We did this work with AWS. Trainium2’s FP8 double-row mode runs at twice the bf16 rate, and our expert kernel is built around it. NKI let us write our own kernels and patch a vendor one at the level of the tensor engine, and its simulator let us check a kernel to 7e-7 before running it on a device. The 24 GiB per core limit also makes compression pay off directly: the compressed model fits on 2 cores, so one trn2.48xlarge can hold 32 replicas instead of 16. Thanks to the AWS team for their support. Next we will run this recipe on Trainium3.