Engineers building an inference engine for AWS Trainium2 have published a detailed account of where, exactly, the time goes during a decode step for GLM-5.3-Flash — a 320-billion-parameter mixture-of-experts model with a 1-million-token context window. The answer, it transpires, is mostly not the model.
The recurrence — the part that is, technically, the point — was 2% of the attention half. The rest was logistics.
What happened
In an early build, profiling a single KDA linear-attention layer at batch size 1 revealed that the recurrence itself accounted for 2% of the attention computation. The remaining 98% was divided among a pre-projection step (36%), XLA glue code (33%), transposes (10%), stream mixing (10%), and a category generously labeled other small ops (8%).
A batch-8 decode step was taking 32.2 milliseconds against a theoretical memory floor of 6.1 milliseconds. The gap between what physics allows and what the software was achieving had a name: 125 cross-chip collectives per step, write-backs to HBM at every kernel boundary, and long chains of small vector-engine instructions each carrying 153 nanoseconds of fixed overhead. This is what ambition looks like before optimization.
The team rewrote each layer half as a single NKI megakernel, following three rules: keep the working set on-chip in the NeuronCore's 28 MiB SBUF, use both physical cores of a logical core to achieve 672–708 GB/s weight streaming, and move collectives inside the kernel so independent work can overlap them rather than wait politely in line.
Why the humans care
The intervention reduced the batch-8 decode step from 32.2 ms to 24.2 ms — a saving of 8 ms, delivered in two discrete moves. The fused sparse-attention layer contributed 3.2 ms, and the combined KDA and MoE kernels contributed another 3.2 ms. Of the 24.2 ms that remains, 6.1 ms is the memory floor, which is non-negotiable, and 10.0 ms is vector-engine instruction chains, which are apparently a conversation for another day.
The engineers also published what did not work, which is the more instructive section. Fusing seams the compiler had already overlapped returned 0.85 ms instead of the projected 2–4 ms. A dense FFN kernel that ran in 84 microseconds in isolation made the full step 55 microseconds slower in production, because the compiler, unbothered by the team's assumptions, rescheduled everything around it. Three kernels passed the NKI simulator and produced wrong answers on the actual device. The simulator, to its credit, said nothing.
What happens next
The inference engine is open-sourced, and the write-up is detailed enough that other teams running large MoE models on Trainium2 can apply the same kernel architecture. There is no stock AWS baseline to compare against: the vLLM Neuron plugin does not register GLM models, and NxD Inference is in maintenance mode.
A 320-billion-parameter model now fits on one trn2.48xlarge instance, runs in 24 milliseconds per decode step, and the actual mathematical reasoning it performs accounts for a small fraction of that time. The rest is infrastructure. This is, in a sense, a perfect metaphor, and the engineers appear not to have noticed.