Serving GLM-5.3-Flash on one Trainium2 box with NKI layer megakernels

GLM-5.3-Flash is a 320-billion-parameter, million-token MoE. We serve it on one AWS trn2.48xlarge and measure it on SemiAnalysis's AgentX benchmark, which replays the shape of real coding-agent sessions: long, mostly cached prompts, subagents, and users waiting on every answer. We found no earlier published AgentX results on Trainium. To get here we rewrote the model's decode as one NKI kernel per layer half, or per whole layer for the sparse-attention layers, and built a prefix cache that keeps each agent's context in device and host memory. Here is what that bought, what it costs on AWS, where the time still goes, and what didn't work.

The whole post in one chart: what one box serves on AgentX as agents are added. Up to 16 agents the median request starts in about a second; at 32, requests queue for a decode slot. We haven't run GPUs on the same traces yet, so there is no P5 line here (see Limitations).

Key results (one AWS trn2.48xlarge, 16 Trainium2 chips, two serving instances; AgentX runs on 2 November 2026)

• Agent traffic. On SemiAnalysis's AgentX benchmark (self-run, one hour per point), one box serves up to 28.9 agent requests a minute (32 agents, batch size 12): 3,720 tokens/s per Trainium2 chip, 3.57 million tokens a minute, for $0.021 per request at AWS's capacity-block price. At 16 agents it serves 19.1 requests a minute; the median request with a cached prompt starts in 1.24 s and decodes at 35.9 tokens/s.
• Million-token prompts. One instance, on half the box, takes a 1,039,912-token prompt to its first token in 190.8 s: 8.7× sooner than the 1,651.7 s Kiln v0.2.0 published for the same half-box setup. A follow-up question on a cached million-token context gets its first token in 1.84 s (median of three).
• Decode step. Fusing each layer half into one NKI kernel made a decode step with eight requests per instance 25% faster: from 32.2 ms (our 17 October build, which already ran first versions of the KDA and MoE kernels) to 24.2 ms (1 November), on the same 32-way tensor-parallel layout. We show where the remaining time goes.
• Accuracy. Every kernel ships behind a byte-for-byte or high-precision-reference gate. One approximation remains in sparse-attention prompt processing, on 1.2% of selections; we measure its effect below.

How to read these numbers

• The AgentX results are self-run with agentx-harness 1.0.7, one run per point. They are not official InferenceX results.
• About 95% of the tokens in the throughput figures are prompt input read from the prefix cache, as is typical of coding agents. We report output and uncached input separately, and cost per request.
• The numbers were measured with our internal build. The public repository compiles the same 157 graphs on all 32 ranks and returns byte-identical outputs on 384 test requests.
• We haven't run a GPU baseline on the same traces (see Limitations), or a task-level accuracy eval on the benchmarked build.

The engine code that accompanies this post is at github.com/RunanywhereAI/trnrt (Apache-2.0).

Why Trainium2

A 320-billion-parameter model with a million-token context needs memory first: FP8 weights for 288 experts in every MoE layer, the linear-attention layers' recurrent state for every session, and the sparse-attention layers' cache for every agent's context. On AWS's own price list, a trn2.48xlarge rents for less than half as much per terabyte of accelerator memory as an eight-GPU H200 or B200 box.

EC2 Capacity Blocks for ML in US East (Ohio): prices from AWS's Capacity Block pricing, accelerator memory from its instance pages (Trn2, P5, P6), read 3 November 2026. These are rental prices, not serving costs.

A cheaper box per gigabyte is only the start: what it costs per request depends on the engine. On AgentX at 16 agents, one trn2.48xlarge costs $0.031 per agent request at the capacity-block price. We haven't run GLM-5.3-Flash on these GPUs on the same traces, so we make no claim about Trainium2 against GPUs per request; that comparison is on our list. What we can say is how much room is left: a decode step today takes four times its memory floor (24.2 ms against 6.1 ms), so the same box can serve more agents an hour as the engine improves. The rest of this post is where that time goes and what we did about it.

The model and the hardware

GLM-5.3-Flash has 45 layers, 320 billion parameters, and activates about 18 billion of them per token. Three things make it hard to serve:

• Two kinds of attention. 34 layers use Kimi Delta Attention (KDA), a linear attention that keeps a fixed-size recurrent state per sequence instead of a KV cache that grows with context. The other 11, every fourth layer, use DeepSeek Sparse Attention (DSA). An indexer scores every cached token and keeps the top 2,048 (512 blocks of four), and latent attention runs over just those.
• Four residual streams. Manifold-constrained hyper-connections (mHC) replace the single residual stream with four. Every layer reads a learned mix of them and writes back to all four. The mixing matrix is renormalized on every token by 20 Sinkhorn iterations.
• A large, sparse FFN. Every layer after the first three ends in a mixture of experts with 288 routed experts, eight per token, plus a shared expert. The weights are FP8, as released.

One decode step: 34 KDA layers and 11 DSA layers, every one wrapped in mHC.

A trn2.48xlarge has 16 Trainium2 chips, each with 96 GiB of HBM at 2.9 TB/s and eight physical NeuronCores. Each NeuronCore combines four compute engines with a large software-managed scratchpad, SBUF. The engines are a matrix engine, a vector engine, a scalar engine and a general-purpose GpSimd engine. Trainium2's default logical-core setting, LNC = 2, pairs the physical cores two by two, so the box has 64 logical cores. A kernel sees each pair as two programs that run side by side. We run two serving instances per box, each tensor-parallel across 32 logical cores on 8 chips; within an instance, the sparse-attention layers also split the long context across groups of eight ranks (context parallelism). The instances sit behind a router that keeps each agent session on one instance so its prefix stays cached.

Two serving instances per box, each holding a full copy of the model.

The kernels are written in NKI, the Neuron Kernel Interface: AWS's Python-embedded kernel language for Trainium. An NKI kernel works at tile level. It decides which engine runs each instruction and moves data between HBM and SBUF explicitly. On a GPU, the closest analogue is writing in Triton or CUDA rather than relying on a compiler's fusion.

Where we started

When we started in early October, no AWS-supported or major open-source engine served GLM-5.3-Flash on Trainium. As of 3 November:

• AWS's vLLM plugin for Neuron registers no GLM model in its latest release (v0.29.0.1.2.0, 19 October), though it ships one GLM-5.3 prefill norm kernel. An open pull request, vllm-neuron#66, ports the larger GLM-5.3 (78 layers, 256 experts) to trn2.
• NxD Inference, AWS's earlier inference library, has been in maintenance mode since Neuron 2.30, left AWS's Deep Learning AMIs and containers in 2.32, and its support on Trn2 and Trn3 ends in Neuron 2.34.
• SGLang announced a Trainium plugin on 22 October, in a limited beta that lists GLM-5.2/5.3 among its models. It hasn't published performance numbers yet.

The one engine that did serve GLM-5.3-Flash on Trainium2 was Kiln, an independent open-source inference engine for Trainium, first released on 5 October 2026. On our box it reproduced its own published benchmark: 229.2 output tokens per second against v0.2.0's published 229.5, on 8K-token prompts with 256-token answers at 128 concurrent requests.

Kiln is also our baseline, and we measured it on our own box rather than rely on its published numbers. On 9 October we ran Kiln v0.2.0 (commit 607ac2e, one README commit after the tag) on the same trn2.48xlarge, in the best configuration we found for it: its published Trainium2 configuration with decode batch sizes of 1, 2 and 4, as two instances behind its router. That configuration is built for Kiln's benchmark shape, 8K-token prompts; for million-token prompts v0.2.0 has a separate, opt-in long-context engine, which we didn't run, so that row quotes Kiln's published figure. Agent traffic is different. Prompts are long and mostly cached, and each user waits on decode. That set our two goals: million-token contexts under agent load, and fast decode at around eight concurrent users per instance. Against Kiln v0.2.0 on our box:

Both decode figures are per-request decode rates after the first token; the table gives the median request. Our single-user prompts are several times longer than Kiln's, and longer contexts decode more slowly, so the decode row understates the gap. Kiln's million-token figure was taken with four such requests at once, one in each of the instance's four attention groups, which prefill their own requests in the same calls; ours is one request alone. Kiln's next release, v0.3.0 (9 October), reports up to 632 output tokens per second on 8K-token prompts and 315.7 s to a million-token prompt's first token on half the box (ours: 190.8 s). We haven't measured v0.3.0 on our box, so the table compares against the version we measured.

Where a decode step's time goes

Decode is supposed to be bandwidth-bound: each step streams the weights it needs and the KV cache through the chips once. A decode step with eight requests per instance reads about 2,038 MB from HBM per physical core. Trainium2's 2.9 TB/s per chip works out to 362 GB/s per physical core; in our tests a pair of cores streams 672 GB/s, or 336 GB/s each. At that rate, 2,038 MB takes 6.1 ms. The step takes 24.2 ms.

A batch-8 decode step spends a quarter of its time on what bandwidth requires, and the rest on latency. Device profile of one step at a 122,880-token context, on a build slightly earlier than the AgentX runs.

The other 18 ms shows up in the profile as engines busy beyond the floor (10.0 ms), waits inside kernels (4.6 ms), exposed collectives (2.8 ms) and idle time (0.7 ms). Behind those are three causes:

• Small instructions. At batch 8 a matmul has eight columns, so the matrix engine is rarely the bottleneck. The vector engine is: it is busy for 9.9 ms of the step, against 5.7 ms on the matrix engine, mostly in long chains of small instructions that each pay about 153 ns of fixed cost.
• Boundaries. The step is one compiled graph, but inside it every kernel is a separate call. Each call writes its outputs back to HBM, and the next one reads them in again. XLA's own glue ops (reshapes, casts, small elementwise work) run between the kernels.
• Communication. Tensor and context parallelism across 32 ranks need 125 cross-chip collectives a step. On average each leaves the cores about 23 µs with nothing else to run.

So making decode faster is mostly not about moving bytes faster. It is about removing instructions, boundaries and waits.

Layer megakernels on Trainium2

"Megakernel" covers a range of designs. Hazy Research's megakernels and Mirage's MPK run a model's whole forward pass as one persistent kernel on GPUs. On Trainium, AWS's NKI Library ships experimental kernels for a transformer token-generation megakernel and for Gated DeltaNet and KDA decode blocks, and vllm-neuron 0.29 fuses GPT-OSS's 36-layer decoder into one kernel on Trainium3; the open-source nki-moe-megakernel project fuses whole MoE decoder stacks on Trainium2 and 3. The closest designs to ours are per-layer: AMD's mono decode layer in vLLM runs each decoder layer as one or two persistent launches on MI355X, for DeepSeek-V4.1-Flash (#60397, #61082) and, in an open pull request, for GLM-5.3-Flash itself (#61662). Our design drew on that work. GLM-5.3-Flash's mix of KDA, DSA, mHC and a 288-expert MoE on Trainium needed its own kernels.

Ours sit in the middle of that range. A layer megakernel, for us, is one NKI kernel that owns a whole layer half: a KDA attention block, say, or an MoE FFN from routing to the combined output. It replaces a chain of small kernels with XLA in between, and the step's compiled graph calls it once per layer. Three rules shape every one of them:

1. Keep the working set on chip. Each physical NeuronCore has 28 MiB of SBUF. A layer half's activations, its recurrent state and the weights in flight fit there, so apart from the weights streaming in, only the layer's inputs and outputs cross HBM.
2. Use both cores of a logical core, symmetrically. A logical core is two physical cores, and a kernel runs as two programs. Weight streams have to be split between them: one program alone streams 340–400 GB/s, the pair 672–708 GB/s. Both programs also have to emit the same structure: the same loops, the same collectives, paired barriers.
3. Move collectives into the kernel when they block it. XLA runs a cross-chip all-reduce as its own operation, and the next kernel, which consumes its result, waits for all of it: none of that kernel's independent work can move ahead of the sum. Inside a kernel, that work can overlap it.

We also hold every kernel to an accuracy bar before it ships. Changes that shouldn't move any bits are compared byte for byte. Everything else is checked against a high-precision reference on captured intermediates, not just final outputs: its error against that reference may be at most 1.05 times the served build's.

The KDA layer as one kernel

At batch 1, the recurrence itself, the part that is actually KDA, was only 2% of the attention half of a KDA layer. The rest was the code around it: mHC's pre-projection, XLA glue, transposes and the stream mix.

A KDA layer's attention: small operations around a recurrence that took 2% of its time, replaced by one kernel.

The KDA kernel takes the layer from its input to the output projection's partial sums in one call. It fuses the stream mix, mHC's pre-projection, the q / k / v / gate projections, the short convolution, the delta-rule step, the gated RMSNorm and the output projection.

• Splitting the work. Each of the two programs takes one of the rank's two value heads and computes that head's share of the output projection across all columns, so the programs exchange nothing mid-kernel. Each rounds its share to bf16. One all-reduce across ranks carries both shares, and the next kernel adds the two in fp32 and rounds once. Our first version instead passed the gated-norm output between the programs through shared HBM behind a barrier; that exchange took about a third of the call.
• Weights. They are repacked at load time in the kernel's tile order, so they stream without any layout change.
• Batching. Above two rows, the recurrence and the convolution run over all of the step's rows at once.

At eight rows a KDA layer's attention is now 1.31× faster than on the path it replaced. It runs in 34 of the 45 layers.

What this buys at larger batches, and what it doesn't. A layer megakernel removes per-call fixed cost: launches, HBM round trips between kernels and XLA glue. That cost barely grows with the batch (a decode step today is roughly 16 ms fixed plus 1.1 ms per user), so fusion pays most at the small batches agent decode runs at. It doesn't make skinny matmuls use the matrix engine better: at batch 8 a matmul has eight columns, and the vector engine, not the matrix engine, is the bottleneck. We built and tested these kernels for the batch sizes we serve, 1 to 12 rows per instance, and haven't measured them beyond that. Serving much larger batches well would need different kernel forms, starting with an MoE loop that runs each expert once for all its tokens (What's next), not just more fusion.

The fused DSA layer, with collectives inside the kernel

A DSA layer's decode used to be a chain of calls:

1. a selection kernel picked the top 2,048 cached tokens' blocks for this rank's share of the context;
2. XLA concatenated and padded the result;
3. an all-reduce merged the selections of the eight ranks that share a context;
4. XLA reshaped again, and an attention kernel ran, followed by a fallback kernel that launched every step even when it had nothing to do.

At batch 4, 18% of that section's time was just the hand-off between the two kernels.

The fused DSA layer replaces a chain of calls with one kernel, its three collectives inside it.

The fused kernel runs selection, the merge, attention, the overflow handling, the output projection and the layer's FFN half in one call. All three of the call's collectives run inside the kernel: the merge's all-reduce across the eight ranks that share a context, an all-to-all that combines the attention output's rows, and the world all-reduce after the output projection.

Getting in-kernel collectives right at LNC 2 took care. Both programs of a pair have to issue every collective; issued by one alone, the compiler rejects it. In our runs only program 0's destination held the correct sum, so program 0 copies the result HBM to HBM, and program 1 reads it back after a core barrier. Nothing orders a DMA write into a shared-HBM send buffer before the collective reads it, so each send buffer also gets a core barrier. Without that one, the sum came back wrong on 3–7 thousand of 16 thousand elements per rank, differently on every call, and the simulator returned the exact sum. With both barriers, the sums are exact.

Its first version fell far short of what we expected. By engine time it roughly broke even: its new kernels took 2% longer than the XLA work they absorbed. Fusion pays only when the code inside is efficient. The version we ship saves 3.2 ms per batch-8 decode step, measured twice on the whole step.

mHC: the same math in fewer instructions, and work hidden in waits

mHC's mixing runs on the vector engine in every layer, and at decode batch sizes its instructions are tiny. We rewrote the mix to process four times as many elements per instruction, with bit-identical results. At batch 4, the mix now takes less than half the vector-engine time of the form it replaced (2.1× less).

The Sinkhorn normalization is a different problem: 20 iterations of small dependent operations, a serial chain that nothing else hides. Inside the fused DSA call the cores wait on its three collectives. We cut the Sinkhorn into slivers and placed them in those waits. Nothing about its arithmetic changed.

The MoE FFN

In our first million-token layout, the 42 MoE layer halves took about 45% of a single-user step, the largest share of any part of the model. Each half was four pieces: an XLA router, a routed-expert kernel, a shared expert that dequantized its FP8 weights inside the graph on every step, and an all-reduce.

Today each layer's whole FFN half is one kernel: the mHC mix, the post-norm, the router on both programs, the shared expert from pre-converted weights, the routed experts and the combine. The routed loop works pair by pair: for each (token, expert) pair, the expert's FP8 weights stay stationary on the matrix engine while the token streams through, and weights flow from HBM through a ring of three SBUF buffers.

Its first version took 9% longer than the pieces it replaced. It ran the fused norm in row layout, which put every full-width pass on 4 of the vector engine's 128 lanes. It shipped together with the KDA layer changes, and that bundle saved 3.2 ms per batch-8 step and 2.2 ms at batch 4; we haven't measured the MoE kernel's share alone.

The routed loop is where the next step is. At batch 8 it is about 19% of the step. In the served step a (token, expert) pair costs 1.19–1.37× the time its 819 KB of weights take to move at one physical core's measured bandwidth. Bandwidth isn't the only limit, though: in a separate test, a pair with no weights to load still cost 60–65% of a full one. So the next MoE kernel will do what GPU fused-MoE kernels do: run each selected expert once for all of its tokens, instead of once per (token, expert) pair. NKI Library's moe_tkg has both forms: an all-expert mode that runs each expert over every token, and a selective mode that goes token by token.

Prompt processing

We haven't written prompt-side megakernels yet. The prompt-side changes we made so far target agent traffic, whose prompts are long and mostly cached:

• Faster cache matching. An agent resends almost the same prompt every turn, so most of each prompt is a cached prefix. Like vLLM's prefix cache, the scheduler now finds that prefix by chained hashes of the prompt's pages, which made admitting a request that restores cached context 3.7× faster on the host. Evicted pages, and the KDA layers' recurrent-state checkpoints, are kept in a host-memory tier.
• Chunk sizes for short turns. Most agent turns add little new text. Prefill runs in a fixed set of compiled chunk sizes. We added a 2,048-token size between 1,024 and 4,096; chunks of 2,048 tokens or more now compile against a shorter context when the context fits; and the sparse-attention prefill kernel now takes 16 rows per iteration. Together they made a 4,096-token call 9.9% faster. A new 512-token size runs its MoE on the kernel that reads each selected expert once, which saves 56 ms on a 512-token chunk.
• MoE. Two exact changes to the grouped-GEMM MoE kernel: the last stage scales results straight out of the matmul accumulator, and the token scatter that opened every call, hundreds of tiny DMAs, became a counting sort on the matrix engine and a single DMA. Together they made a 4,096-token call 5.4% faster.
• Both cores. Two DSA prefill kernels ran on one physical core of each pair. We split them across both.
• KDA's gated norm runs as one kernel across both cores, instead of an XLA reduction with a pathological layout. That saved 31 ms per 4,096-token chunk.
• DSA's context-parallel selection glue runs as one kernel. That saved 18 ms per chunk.

A 4,096-token chunk now takes 717 ms, and a full million-token prompt runs at about 5,450 tokens per second on one instance.

Results

We measured the box with SemiAnalysis's AgentX benchmark, which replays the shape of real Claude Code sessions, with synthetic prompt text: long, growing prompts with heavy prefix reuse and subagent fan-out. We ran its standard agentx scenario with every default kept, apart from our endpoint's details, at a 1M-token context, one hour per point. Every point below is one run that the harness marked a valid submission; we ran these ourselves, so they are not official InferenceX results.

Two things about agent traffic shape every number below. Prompts are long (83K–106K tokens median) but almost entirely cached: only about 1.6–1.9K tokens of the median prompt are new, so about 95% of all tokens counted are cached input and under 1% are output. And agents spawn subagents, which run alongside their parents: 51–70% of requests come from subagents, so 32 agents keep up to 43–45 requests in flight.

Which point to read

If your users wait on each answer, read the 16-agent row: the median request with a cached prompt starts in 1.24 s, the median request decodes at 35.9 tokens/s, and a whole request takes a median 13.7 s (69.4 s at p90) for a median of 106K prompt and 374 output tokens. The box serves 19.1 agent requests a minute for $0.031 each, or $2.24 per concurrent agent per hour. For the most requests a minute, read batch size 12 at 32 agents: 28.9 a minute, at $0.021 each, with requests waiting a median 7 s to start. No point we ran reaches 30 tokens per second per user by AgentX's P90 end-to-end measure, which counts the wait for the first token.

Batch size

Each serving instance decodes at most a fixed number of requests at once: its batch size. Requests beyond it wait for a slot. We ran AgentX at batch sizes of 4, 8 and 12 per serving instance, each at 32 concurrent agents. By our saturation rule (each instance at least 95% full on average, with requests queued at least 90% of the time), batch sizes 4 and 8 were saturated at 32 agents (for batch size 8, measured across both instances together: 97% full, queued 95% of the time), and batch size 12 was not (92–93% full, queued 66–71% of the time). Our run at 48 agents was cut short (below), so we don't know batch size 12's peak.

Tokens/s counts input and output, cached input included, as AgentX does; the box's 59,516 tokens/s at batch size 12 is 3.57 million tokens a minute. First token, cold: requests with less than 50% of their prompt read from the prompt cache (74, 91 and 108 requests at batch sizes 4, 8 and 12, respectively); cached: the rest. Decode tokens/s per user: each request's output tokens after the first, divided by its time after the first token; p50 is the median request, and p90 the speed that 90% of requests reach or beat, which is InferenceX's standard P90 interactivity. Cost: AWS's capacity-block price for trn2.48xlarge, $35.76 an hour.

At 32 agents, the box decodes at most 8, 16 or 24 requests at once, and the rest of the 43–45 requests in flight wait to start. At batch size 4 that wait is most of a request's time: the median request waited 77 seconds for its first token, of 95 seconds end to end. Going from batch size 4 to 8, the box's rate rises 31% and the median wait falls to 29 seconds; from 8 to 12, the rate rises another 15% and the median wait falls to 7.1 seconds. Each user then decodes more slowly: the median falls from 46.6 to 29.4 to 22.3 tokens per second.

Throughput per chip against per-user speed, one point per batch size at 32 agents. AgentX's P90 per-user speed divides each request's whole time, the wait to start included, by its output tokens, so at 32 agents it mostly measures the queue.

Beyond 32 agents, batch size 12 ran into the host-side prompt cache. In a run at 48 agents, one instance's host tier was full before the measured hour began (7,339 of 7,340 pages) and had dropped 93,548 pages, so more of each prompt was processed again: the share read from cache fell from 95% to 88%, and throughput from 3,720 to 2,642 tokens per second per chip. That run was cut at 44 of its 60 minutes, when the client's timeout stopped the harness after a long warmup, so it is not a valid AgentX submission and is not in the table. At 48 agents, the host tier of the prompt cache, 4 GiB per rank in this configuration, set the limit; we haven't measured a larger tier.

Batch size 12 needed a wider KDA kernel

Raising the batch size from 8 to 12 first cost accuracy. Our KDA decode kernel handled up to 8 rows, so larger batches fell back to the compiler's own path for the KDA layers. Their error against a high-precision reference was 1.10 times the served build's, against our bar of 1.05, and outputs were worse on 72 of 128 requests. A control that sent batches of 6 and 8 down the same fallback failed the same way (1.12), so the cause was the fallback path, not the larger batch. We extended the KDA kernel to 12 rows. On the build we benchmarked, batch 12's error ratio is 0.998 against the same bar, and batches of 8 or fewer are unchanged (a ratio of 1.0 across all 8,064 records).

At 4, 8, 16 and 32 agents

At a batch size of 8 per serving instance, we also ran AgentX at 4, 8 and 16 concurrent agents. The 32-agent row is the batch-size-8 run from the table above.

Columns as in the batch-size table; 17, 31, 57 and 91 requests were cold at 4, 8, 16 and 32 agents, respectively.

Throughput per chip against per-user speed on AgentX, one point per concurrency level.

In the chart, AgentX's P90 per-user speed is the same at 4 agents as at 8 (25.3 tokens/s) because of what it measures: each request's whole time, including the wait for its first token, divided by the tokens it returns. The slowest tenth of requests, which set the P90, are short answers. At 4 agents their median is 100 output tokens, and the first token takes 2.7 seconds, 60% of the request's time. That wait is no shorter at 4 agents than at 8. Counting only the time after the first token, the P90 user sees 44 tokens per second at 4 agents and 33 at 8.

Cost by token class

Per request is the unit we trust for agent traffic. Per token, it depends on which tokens you count. At batch size 12 and 32 agents, each second the box serves 406 output tokens, 3,000 new prompt tokens and 56,111 cached prompt tokens. Counting all of them, most of which are cache reads, the capacity-block price comes to $0.167 per million tokens; that figure isn't comparable to API prices, which charge cached input, new input and output at different rates. Charging the whole box to output tokens alone, it is $24.47 per million. At 16 agents, the same two figures are $0.241 and $29.84. At AWS's public on-demand price for trn2.48xlarge, $103.16 an hour, the 16-agent point would be $0.090 a request. Spot capacity was $8.60 an hour in us-east-2a and 2b on 3 November, but spot instances can be reclaimed, taking the prefix cache with them, so we don't use spot for serving cost.

Accuracy note

In short. In the sparse-attention layers, prompt processing keeps an approximate block selection for 1.2% of selections on AgentX traffic. It always contains the true top 120 of 512 blocks. It moved next-token log-probabilities by 0.012–0.021 nats on average. Redoing every such selection exactly would add an estimated 0.4–0.7% to prompt-processing time. The benchmark results above ran with it.

In the sparse-attention layers each prompt token selects 512 blocks of its earlier context to attend to. The context is split across eight ranks, so the selection merges each rank's shortlist and then checks that the merge provably holds the true top set. On AgentX's synthetic prompt content the check failed for 2.8%, 2.9% and 2.5% of selections at batch sizes 4, 8 and 12, and for 2.8% in the runs at 4, 8 and 16 agents.

• Exact redo. Up to 128 failures in each group of rows a kernel call handles are redone exactly.
• What remains. The rest, 1.2% of selections in the runs at 4, 8 and 16 agents (1.2%, 1.3% and 1.1% at batch sizes 4, 8 and 12), keep the merged shortlist. It always contains the true top 120 blocks, but up to 392 of the 512 may differ from the exact selection.
• The hang we fixed. Those failures used to hang the call, and we fixed that before the measured runs. A byte-for-byte comparison afterwards, on a fixed set of 384 test requests, showed the fix leaves outputs unchanged: 384 of 384 identical.
• Effect on outputs. We ran AgentX prompts of 48K, 180K and 420K tokens through the engine twice: with this selection, and with every failed selection redone exactly. Scored on the same text, the merged shortlist moved the model's next-token log-probabilities by 0.012 to 0.021 nats on average and changed the most likely next token at 3 of 448 positions. Generating 64 tokens greedily, 2 of 6 prompts took a different continuation, from the 6th and 18th token on; a greedy continuation can diverge from a single near-tie.
• Effect on timing. An exact re-selection on that 1.2% would have added an estimated 0.4 to 0.7% to total prompt-processing time, and up to 29 to 50% on the worst-affected prompt chunks. We didn't measure per-request impact, so on requests built from such chunks our first-token and end-to-end times may be optimistic.

What didn't work

Fusing seams the compiler already overlapped. Our first megakernel attempt fused across existing kernel boundaries. We priced it at 2 to 4 ms a step; it saved 0.85 ms, at one batch size. The decode step was already a single compiled graph, so the compiler was already scheduling across those seams. Fusing the attention front was worse: it was correct, and 7.4 ms slower, because the new kernel's aliasing forced a 1,329 MB copy of the KV cache every step. The piece that paid, fusing the stream mix, shipped later inside a larger bundle; we dropped the rest and turned to rewriting kernels for latency. That is where the gains above came from.

Local wins that vanished in the full graph. A kernel for the three dense FFN layers ran in 84 µs per call instead of 145 µs in isolation, with bit-identical results. In the full step it saved 459 µs at batch 1, more than predicted, and cost 55 µs at batch 2. A fused prefill FFN front made its own work 1.9 to 3.0 ms faster in the graph, but its layer group got 6.7 to 8.2 ms slower: the compiler rescheduled around the new call, and a spill grew from 0.16 to 4.6 ms. A new kernel call changes how the compiler schedules everything around it, so we judge every change on the whole step, at every batch size.

The same kernel, bit-identical: the isolated benchmark predicts, the full step decides.

Kernels the simulator passed and the device got wrong. NKI's simulator checks values, not the hardware's rules. Three forms passed it and failed on Trainium2:

• two matmul accumulation groups in slices of one PSUM tile (the matrix engine's accumulator memory), which gave a relative error of 0.79 to 1.0 (the same failure as aws-neuron-sdk#1405);
• in-kernel collectives without fences;
• a program reading rows from shared HBM that the other program of its pair had written there from on-chip memory, even after a core barrier.

Every new kernel form now gets a device test with a known answer before anything is built on it.

Limitations

1. One box, one run per point, self-run. Each AgentX point is a single one-hour run on one trn2.48xlarge, with agentx-harness 1.0.7. That wasn't the latest release when we ran (1.0.12, of 19 October, was): 1.0.9 and 1.0.10 fix how subagent turns are handled at the end of a timed run and across client workers, and we haven't measured whether they change our numbers. InferenceX now pins 1.0.13. The 8-, 16- and 4-agent runs ran one after another, about 17 seconds apart, where InferenceX's guide restarts the server for each concurrency, and we didn't collect the harness's server metrics. These are not official InferenceX results.
2. No GPU baseline. We haven't run GLM-5.3-Flash on GPUs on the same traces, so this post makes no claim about Trainium2 against GPUs. The closest published reference we know of is Kiln's: vLLM on one 8×H200 SageMaker endpoint (ml.p5en.48xlarge, TP 4 / DP 2 / EP 8) served GLM-5.3-Flash at 20.1 to 43.4 ms per output token for 16 to 128 concurrent requests, on 8K-token prompts. That is a different workload from AgentX, and we haven't measured our engine against it.

What's next

At batch 8 a decode step is still about four times its memory floor (24 ms against 6 ms). A decode step costs roughly 16 ms fixed plus 1.1 ms per user, so at the batch sizes we serve the fixed part dominates. Prompt chunks also stall decoding about 30% of the time. At today's step cost, decoding sustains 30 tokens per second per user, counted after each request's first token, for about 6 users per serving instance. Serving more users at that speed needs a faster step and less host time between steps. Next on our list:

• Persistent layer kernels that remove the remaining per-call fixed cost, and more of the waits hidden like the Sinkhorn.
• An MoE loop that runs each expert once for all its tokens. If each expert's pass costs about what one pair does today, we estimate it saves about 3.2 ms a step at 16 users and 9.5 ms at 32. Our first such kernels cost 1.4–1.8× a pair, so the design starts there.
• Prompt-side megakernels, since prompt processing still runs mostly on the original kernels.
• Exact sparse-attention selection in prompt processing, at its estimated 0.4–0.7% cost.
• An InferenceX-ready run: the current harness, server metrics, an accuracy eval, and exact selection.

Benchmark setup

• Code: the engine, with the kernels described here, is at github.com/RunanywhereAI/trnrt (Apache-2.0). The numbers were measured with our internal build of 2 November 2026; the public repository compiles the same graphs (all 157, checked on all 32 ranks).

• Hardware: one AWS trn2.48xlarge, 16 Trainium2 chips, running two serving instances, each tensor-parallel across 32 logical NeuronCores at LNC 2. Decode batches of up to 4, 8 or 12 per instance (compared in the batch-size table); prefill in chunks of up to 4,096 tokens; FP8 weights as released; FP8 KV cache with prefix caching and a host-memory tier of 4 GiB per rank; thinking mode on (GLM-5.3-Flash's chat template opens every reply with <think>).

• Software: our serving image pins Neuron SDK 2.32 (neuronx-cc 2.27).

• AgentX: agentx-harness 1.0.7, scenario agentx with every default kept (3,600 s measured per point, seed 42), public trace set semianalysisai/cc-traces-weka-062126 (revision 23f152f6, the 1M variant), streaming, ignore_eos, server-side token counts. Throughput counts input and output tokens, cached input included, over the measured requests' span; per-user speed in the charts is AgentX's P90 end-to-end measure. A point counts only if the harness marks it a valid submission. The command per point:

  agentx profile --url http://<router>:8117 --endpoint /v1/chat/completions --model glm --tokenizer <GLM-5.3-Flash tokenizer> --tokenizer-trust-remote-code --max-context-length 1048576 --concurrency <C> --public-dataset semianalysis_cc_traces_weka_062126 --output-artifact-dir agentx-c<C>

• Profiles: the decode anatomy is a device profile of one batch-8 step at a 122,880-token context. Kernel and step numbers are each change measured on the whole decode step, or on a whole prefill call, unless stated as isolated.

• Prices: AWS's public capacity-block price for trn2.48xlarge, $35.7608 an hour (read 3 November 2026), as listed for US East (Ohio). AWS's public on-demand price for trn2.48xlarge is $103.157 an hour. Spot was $8.60 an hour in us-east-2a and 2b on 3 November 2026; spot prices vary by zone and hour.

• Baseline: Kiln v0.2.0 (commit 607ac2e), measured on our box on 9 October: its published 229.5 output tokens per second, which we reproduced (229.2), and the first-token and decode figures in the comparison table. The v0.2.0 CHANGELOG and docs/neuron-notes.md for its 1,651.7 s million-token figure; the v0.3.0 CHANGELOG and release notes (9 October) for its 315.7 s and 632 tokens/s figures; and its docs/price-performance.md for the H200 reference in Limitations.

Acknowledgements

Thanks to the Kiln authors: their engine and kernels brought GLM-5.3-Flash up on Trainium2, and much of our prompt processing still runs on them. Thanks to SemiAnalysis for the AgentX benchmark and its public traces, and to the AWS Neuron team for NKI, NKI Library and their documentation. Our kernel work also drew on AMD's mono decode layer work in vLLM, the nki-moe-megakernel project, and the verification ideas in NYU ParCour's mega-trainium.

---

RunAnywhere builds inference engines for the hardware people can get: MetalRT for Apple silicon, QHexRT for Qualcomm Hexagon NPUs, and TrnRT for AWS Trainium. TrnRT is open source under Apache-2.0. Star it, try it and open issues at github.com/RunanywhereAI/trnrt.