Open-source NKI megakernel engine serves 320B GLM-5.3-Flash on a single AWS Trainium2 box

WildCard727 · reddit · 2026-11-05

RunAnywhere open-sourced trnrt, an inference engine that serves GLM-5.3-Flash (320B MoE, 1M context) on a single trn2.48xlarge. Profiling showed the model's actual math was a tiny slice of decode: in one KDA layer at batch 1, the recurrence was just 2% of the attention half, with 36% going to pre-projection and 33% to XLA glue; a batch-8 step took 24.2ms vs a 6.1ms memory floor due to small vector-engine instruction chains, HBM writebacks at every kernel boundary, and 125 cross-chip collectives per step.

The fix: rewrite each layer half as one NKI kernel following three rules — keep working sets on-chip in the 28MiB SBUF, split weight streams across both physical cores (672–708 GB/s), and move blocking collectives inside the kernel. Batch-8 decode went from 32.2ms to 24.2ms.

Lessons in what didn't work: fusing seams the compiler already overlapped saved only 0.85ms; isolated kernel benchmarks misled (one dense FFN kernel got faster alone but slowed the full step); and three kernels passed the NKI simulator but failed on device. Every kernel is verified byte-for-byte or against a high-precision reference. Self-run AgentX results: up to 28.9 agent req/min, 35.9 tok/s median decode at 16 agents. No GPU comparison exists since vLLM's Neuron plugin registers no GLM model.

Original post →

Open-source NKI megakernel engine serves 320B GLM-5.3-Flash on a single AWS Trainium2 box

More from Infra