RunAnywhere 开源 AWS Trainium2 上 GLM-5.3-Flash 的 NKI 层 megakernel 推理引擎

We wrote an inference engine for AWS Trainium2 that serves GLM-5.3-Flash (320B MoE, 1M-token context) on one trn2.48xlarge, and open-sourced it. Most of the work went into decode, and the part I found most interesting is how little of a decode step was the model's actual math.

In an early build, in one KDA (linear attention) layer at batch 1, the recurrence was 2% of the attention half. The rest was the mHC pre-projection (36%), XLA glue (33%), transposes (10%), the stream mix (10%) and other small ops (8%). A batch-8 decode step took 24.2 ms against a 6.1 ms memory floor, and the profile showed why: long chains of small vector-engine instructions at about 153 ns of fixed cost each, a write back to HBM at every kernel boundary, and 125 cross-chip collectives per step.

So we rewrote each layer half as one NKI kernel (each whole layer, for the sparse-attention layers). The kernels follow three rules:

Keep the working set on chip. Each physical NeuronCore has 28 MiB of SBUF, enough for a layer half's activations, its recurrent state and the weights in flight, so only the layer's inputs, outputs and streamed weights cross HBM.
Use both physical cores of a logical core. One program streams 340–400 GB/s of weights, the pair 672–708 GB/s, so the weight stream has to be split between them.
Move a collective into the kernel when it blocks the next one. Run by XLA, an all-reduce makes the next kernel wait for all of it. Inside the kernel, independent work can overlap it.

On the same layout, a batch-8 decode step went from 32.2 ms on our 17 October build to 24.2 ms on 1 November. Two changes did most of it, each measured on the whole step: the fused sparse-attention layer saved 3.2 ms per batch-8 step, and the KDA and MoE kernels, which shipped together, saved another 3.2 ms. There's no stock-AWS baseline to compare against: the latest vLLM Neuron plugin release registers no GLM model, and NxD Inference is in maintenance mode. Of the remaining 24.2 ms, 6.1 ms is the memory floor, 10.0 ms is engines busy beyond it (mostly vector-engine instruction chains), 4.6 ms is waits inside kernels, 2.8 ms is exposed collectives and 0.7 ms is idle.

What didn't work:

Fusing seams the compiler already overlapped. We priced it at 2 to 4 ms a step and it saved 0.85 ms at one batch size, because the step was already one compiled graph.
Trusting isolated benchmarks. A dense FFN kernel ran in 84 µs instead of 145 µs on its own, bit-identical, and then made the full step 55 µs slower at batch 2, because the compiler rescheduled everything around the new call.
Trusting the simulator. Three kernels passed NKI's simulator and gave wrong answers on the device: two matmul accumulation groups in slices of one PSUM tile, in-kernel collectives without fences, and one program reading shared HBM that the other program of its pair had just written.

Accuracy: every kernel is checked byte for byte where no bits should change, and otherwise against a high-precision reference, with its error capped at 1.05x the previous build's. One approximation remains in sparse-attention prefill, on 1.2% of selections: against an exact rerun it shifted next-token log-probabilities by 0.012 to 0.021 nats on average and changed the top token at 3 of 448 positions. We haven't run a task-level eval on this build.

End to end, on SemiAnalysis's AgentX benchmark (self-run, one hour per point, not official InferenceX), we measured up to 28.9 agent requests a minute on one box, and at 16 agents the median request decoded at 35.9 tokens/s. We haven't run GPUs on the same traces, so there's no GPU comparison here.

Disclosure: I work at RunAnywhere, and the code is ours.

Code: https://github.com/RunanywhereAI/trnrt (commands to rerun every number are in docs/RESULTS.md) Write-up with the profiles and every number: https://runanywhere.ai/blog/glm-5-3-flash-trainium2-nki-megakernels

Our design drew on AMD's per-layer decode kernels in vLLM and on Hazy Research's and Mirage's megakernel work. Happy to hear criticism, especially from anyone who has fought NKI or XLA on Neuron.

submitted by /u/WildCard727
[link] [留言]