ATOM Meets Mooncake: A Shared KV Pool for Agentic Serving on AMD Instinct GPUs

Introduction

A coding agent rarely asks a new question. It sends the same long conversation back with a little more at the end. As agentic coding becomes everyday work, serving it means long sessions that pause for tools and come back expecting the server to still hold their history. SemiAnalysis now measures this workload with InferenceX AgentX.

We serve it with prefill and decode on separate GPUs, which keeps long prefills from stalling the agents that are generating. The split has a side effect that is easy to miss: it splits the server’s host memory in two as well, and the half next to decode sits nearly idle. Mooncake Store pools both halves back into one KV pool. With it, ATOM serves around a hundred concurrent agent sessions on one 8-GPU AMD Instinct MI355X server at about 44K tokens/s per GPU, and reuses more than 95% of prompt tokens. On four MI350X servers at NeuroSpark AI, one shared pool reaches up to about twice the throughput of per-server pools.

One Agent Session, Up Close

AgentX replays recorded Claude Code sessions against a live server. Figure 1 follows one of them, session ebaea288 from the public cc-traces-weka-062126 corpus, chosen because its 65 main-thread requests match the corpus median. Over almost two hours, the main thread’s context grows from 40K to 173K tokens, and at one point the agent fans out four subagents in parallel. Times are as recorded with the original model; the one long pause, 87 minutes, is drawn as 300 seconds, the cap AgentX replay applies.

Three things in this picture decide how the server must treat KV.

The history is almost all reused. Each request resends the same conversation with a little added at the end. Across the corpus, the median request carries 142K input tokens, and with a perfect prefix cache only 1.7K of them need prefill. Whether a turn computes 1.7K tokens or 142K depends on one thing: whether the server still holds the history as KV. Some turns no cache can save: after the long pause in Figure 1, only the first 32,704 tokens of the next prompt match what was sent before, so the rest is prefilled again (hatched).

The history is heavy. GLM-5.3 with an FP8 KV cache stores 48,672 bytes per token, so a 100K-token conversation takes about 5 GB.

The history spends most of its life waiting. The amber and light-green blocks in the top panel are time when the main thread has no request running, because a tool, the user or its subagents are working. During the 182-second wait marked in Figure 1, 7.2 GB of KV has nothing to do, but the next turn needs all of it. When the four subagents run, each builds its own context, and the session’s live KV jumps to 26 GB. The server therefore has to hold the history of every live session, not only of the requests running right now.

Why GPU Memory Is Not Enough, and Why Disaggregation Makes It Worse

When the 182-second wait in Figure 1 ends, the main thread sends 147,904 tokens, 146,944 of which it has sent before. If the history survived the wait, prefill computes 960 tokens. If not, it computes all of them (Figure 2, panel A).

One session fits easily; every live session at once does not. By our estimate, each live session holds about 10 GiB of history that a later turn will reuse, counting its main thread and running subagents, so around a hundred sessions need about 1 TiB. The four prefill GPUs have about 0.4 TiB left for KV after the model weights. Under load, other sessions’ KV fills that space, the waiting history is evicted, and the next turn prefills it again.

Host DRAM has the room, but disaggregation splits it. Each half of the server has its own NUMA node, and prefill and decode each take one. A CPU cache attached to prefill sees only prefill’s half. The decode half has as much DRAM again, and almost none of it holds history, because decode keeps its working KV in GPU memory. The capacity exists, but on the wrong side of the split.

So we pool both halves into one KV pool that prefill can save to and restore from (panel B). The history is saved there during the wait and restored for the next request, which then computes only its new tokens. That raises two new problems. A hit in DRAM is not yet a hit in prefill: the history has to reach the prefill GPUs, in prefill’s layout, before the request can start. And every request’s KV still has to cross from prefill to decode before its first token. We set three goals:

• Keep far more agent history than GPU memory can hold.
• Turn a pool hit into a prefill hit, safely and before the request needs it.
• Hand each request’s KV to decode as fast as prefill produces it.

ATOM + Mooncake: Two KV Paths

ATOM is AMD’s open-source inference engine for Instinct GPUs, built on ROCm and the AITER kernel library. On this server, ATOM runs one prefill instance on the four GPUs of one NUMA node and one decode instance on the other four, with ATOM’s router in front. Prefill splits the model by layers: each GPU runs about a quarter of the layers, as one stage of a pipeline. Decode splits each sequence by tokens: each GPU holds every fourth token, for all layers.

When a turn arrives, prefill reuses the longest prefix still in its GPU memory and loads the part beyond it from Mooncake Store, over RDMA straight into GPU memory. It computes only the tokens that neither tier holds and saves the new KV to the Store, for the session’s next turn. The router then hands the request to decode, Mooncake Transfer Engine carries its KV from the prefill GPUs to the decode GPUs, and decode generates the answer.

The two Mooncake components handle KV with different lifetimes. Transfer Engine moves KV that a request needs now. Mooncake Store keeps KV that a session will need later, after the request that produced it has finished, and decode never reads it. ATOM decides which KV is reusable and when it is safe to use; Mooncake decides where it lives and how it moves.

Mooncake Store separates the services that provide memory from the clients that use it, which suits a disaggregated server. We run one Store service per NUMA node, each contributing 768 GiB of DRAM, plus a Mooncake Master that tracks where objects live and evicts the least recently used ones when the pool is 90% full. Only the prefill GPUs read and write the pool, yet decode’s NUMA node supplies half of it. The result is a 1.5 TiB KV pool, nearly four times the KV space of the prefill GPUs, built from memory that would otherwise sit idle. The same design scales out across servers: Store services on each added server join the same pool, which grows with every server, and prefill GPUs read and write it over RDMA just as they do here.

The pool also needed huge pages. All of its memory is registered with the RDMA NICs, and registration fails if any part of a region falls back to 4 KiB pages. Transparent huge pages could not guarantee that on busy hosts, so we reserve huge pages before the Store services start. On the four NeuroSpark servers below, dropping caches and compacting memory before the Store starts was enough for transparent huge pages.

From a Pool Hit to the First Token

Capacity is the easy part. A pool hit saves work only if all four prefill GPUs have their share before the request starts, and the handoff to decode has to keep pace with prefill.

Stage-Aligned Restore: Turning a Pool Hit into a Prefill Hit

Figure 4 opens up the green restore arrow of Figure 2. The history is far: it can sit in either half’s DRAM and must cross the network into GPU memory. It is packed: the pool holds plain bytes, which must be unpacked into the paged KV that attention reads. And it is all or nothing: prefill is split by layers, so a prefix is usable only when every stage has its layers.

ATOM stores KV the way prefill is split. Each pipeline stage saves its own layers of each 256-token chunk as a separate Store object, and a prefix counts as a hit only as far as all four stages hold it. Mooncake keeps a chunk’s four objects in one group and evicts them together, so eviction never leaves a chunk that only some stages can load. Small chunks also mean that a turn that changes its history partway through still reuses everything up to the last 256-token boundary before the change.

Each prefill GPU registers one buffer with Mooncake. KV moves between that buffer and Store DRAM over RDMA, and the GPU packs and unpacks it itself, so no CPU copy sits on the path. A request that needs Store data waits until every stage has its part, while other requests keep running.

Handing KV to Decode in Large Writes

Prefill (pipeline parallelism, PP4) and decode (tensor parallelism with decode context parallelism, TP4 + DCP4) keep KV in different layouts: a prefill GPU holds every token for its own layers, while a decode GPU holds every fourth token for all layers. A direct copy therefore becomes one tiny write per token per layer. For one 8,192-token chunk of one 20-layer stage, that is 163,840 writes of 576 bytes each, and posting them, not the network, becomes the bottleneck.

ATOM instead gathers the KV bound for each decode GPU into contiguous buffers on the prefill GPU. Mooncake Transfer Engine moves each buffer in RDMA writes of up to 8 MiB into a landing buffer on the decode GPU, and decode scatters it into its KV pages. The same chunk, about 94 MB, now takes 12 writes.

The trade is GPU work and memory for far fewer network operations. Gather and scatter run on the GPUs, the staging and landing buffers take GPU memory that could otherwise hold KV, and a request is ready for decode only after its scatter finishes.

Results: Mooncake Store Carries the Reuse as Sessions Grow

We ran InferenceX AgentX against this deployment for an hour at each load, from 56 to 112 concurrent agent sessions, in one InferenceX sweep. The sweep measured GLM-5.2, which has the same architecture and KV size per token as GLM-5.3. It runs with MXFP4 weights, an FP8 KV cache and multi-token prediction, with the draft acceptance length pinned to the value InferenceX prescribes. Throughput per GPU counts every prompt token, cached or not, plus generated tokens, divided by all eight GPUs.

Throughput. Throughput per GPU rises from about 31K tokens/s at 56 sessions to about 44K at 96 and 112 sessions, where the median generation speed per user is about 46 tokens/s.

Where the prompt tokens come from. Reuse stays high because Mooncake Store picks up what GPU memory can no longer keep. At 56 sessions, prefill GPU memory serves 81% of prompt tokens and the Store 16%. As sessions grow, GPU memory’s share shrinks and the Store’s grows, to 46% of all prompt tokens at 112 sessions. Total reuse stays above 95% at every load, close to the 96–97% that the traces allow.

What the Store saves. If the Store’s share had to be computed instead, prefill at 112 sessions would compute about half of all prompt tokens rather than 4.6%, roughly 11 times the prefill work; at 96 sessions, about 9 times. This is an estimate from the measured shares, not a run without the Store.

How fast a restore is. A typical restore brings back 110K–130K tokens of history. Each prefill GPU reads and unpacks its share in a median of 150–220 ms. Because the four stages do not load at the same moment, a request waits on average 0.4–1.2 s from its first stage starting to its last stage finishing, rising with load. No load failed in any of the measured runs. At 112 sessions, Mooncake moved about 66 TB of KV per hour inside the server, between the Store and prefill and from prefill to decode.

Where it stops. Up to 96 sessions, median time to first token stays at 2.5 s or below. At 112 sessions throughput still holds at 44K, but the server has run out of headroom: decode’s KV memory is a median 91% full, requests start to queue in front of prefill, and median time to first token reaches about 7 s.

Scaling Out: GLM-5.3 on Four MI350X Servers (NeuroSpark)

NeuroSpark serves GLM-5.3 (MXFP4 weights, FP8 KV cache) with ATOM and Mooncake Transfer Engine on four 8-GPU AMD Instinct MI350X servers. We added Mooncake Store to that deployment and ran two experiments, each comparing a split pool with a shared pool: one on replayed production traffic, and one with the same AgentX sweep as above, with the same prefill and decode parallelism.

The deployment. The 32 GPUs run six prefill instances and two decode instances (6P2D), arranged as two groups of three prefills and one decode behind a KV admission gate. Each decode instance runs TP4 with decode context parallelism (DCP4) and receives KV through ATOM’s gather / large-write / scatter path described above. For the AgentX sweep, each prefill instance runs PP4, as in the single-server setup, so the Store keeps one copy of every token. The production-traffic comparison ran on NeuroSpark’s production recipe, whose prefill instances run TP4: at production load, TP4 prefill is not the bottleneck, and both pool arms use it, so the comparison isolates pooling.

The pool. The Store has no NICs of its own; it shares each NIC with engine traffic. The pool is split by RoCE rail into four parts, each with its own Mooncake Master, and rank i of every prefill instance reads and writes part i. With PP4 prefill, each stage’s objects therefore sit in a different part and are evicted by a different Master; a prefix still counts as a hit only as far as all four stages hold it. Each server, decode servers included, contributes eight 96 GiB Store services: 768 GiB per server, about 3 TiB in all. Only prefill instances read and write it.

Split Pool vs Shared Pool at Production Load

NeuroSpark’s production GLM-5.3 traffic averages about 0.8 requests per second, so we replayed coding-agent sessions reconstructed from production traffic at 0.6, 0.8 and 1.0 requests per second, 20 minutes per run, with production-style requests (up to 32K output tokens, reasoning on). Both arms use the same deployment and the same 768 GiB of Store memory per server; only sharing differs. In the split pool, each server’s memory serves only the prefill instances on that server. In the shared pool, all four servers’ memory forms one pool that every prefill instance reads and writes. Each replay lane starts partway into a session, so its first turn is a cold prefill in every arm; the statistics cover the turns after it.

Across servers, the pool pays off only when it is shared. With the shared pool, Mooncake Store serves 58% of prompt tokens at every load and KV cache reuse is 95–96%: prefill computes only 4–5% of prompt tokens, and time to first token stays at about 1.0 s at the median and 1.7–1.8 s at p90. With the split pool, the Store serves only about 11% and reuse falls to 51–55%. The gate and routers spread a session’s turns over six prefill instances on four servers, so the next turn usually lands on a server whose memory does not hold that session’s history. Prefill then computes about half of every prompt, and time to first token rises to 7–10 s at p90. This is the multi-server version of the NUMA split above: memory that only one side can reach does little.

AgentX on 32 GPUs

We used the same AgentX harness and traces as above (cc-traces-weka-062126), run directly rather than through InferenceX, with one 15-minute run at each load. Traces whose peak context exceeds the deployment’s 917,504-token input limit are skipped (10 of 393). Output lengths are pinned to the recorded ones. We report the whole system’s throughput: every prompt token, cached or not, plus generated tokens, per second across all 32 GPUs. In AgentX, each agent thread sends its next request only after its previous one finishes and the recorded wait has passed, so at a fixed number of sessions, throughput follows the sessions’ pace rather than the GPU count. The four servers therefore reach about the same system throughput as the single MI355X server above, 3.1K–4.6K tokens/s per session in both, and their throughput per GPU is not comparable with Figure 6; this deployment also uses MI350X GPUs and runs without multi-token prediction.

Throughput and reuse. Throughput rises from 256K tokens/s at 56 sessions to a peak of 389K at 96 and is 348K at 112, while time to first token stays at about 1.0 s at the median (0.96–1.05 s) and at or below 2.3 s at p90. KV cache reuse stays at 94–96% at every load, against a trace ceiling of about 97%. As in Figure 6, Mooncake Store picks up what GPU memory can no longer keep: prefill GPU memory serves 70% of prompt tokens at 56 sessions and 54% at 96 and 112, while the Store’s share grows from 26% to 40–41%.

A split pool cannot keep up at high concurrency. We repeated the sweep on the same deployment with a split pool, in which each server’s Store serves only its own prefill instances. Throughput levels off at 156K–194K tokens/s, against 256K–389K with the shared pool: the shared pool is 1.6 times faster at 56 sessions and about twice as fast at 96 and 112. With the split pool, the Store serves only 12–15% of prompt tokens and KV cache reuse falls from 79% to 65% as sessions grow, so prefill computes 21–35% of prompt tokens instead of 4–6%. Time to first token at p90 climbs from 3.9 s to 19.3 s, while the shared pool stays at 1.9–2.3 s.

Future Work

ATOM with Mooncake now keeps agent history within reach on one server and across four, but there is room to go further.

Sharing one pool across more servers. On four servers, one shared pool already about doubles throughput over split pools at high concurrency. A wider pool keeps reuse this high for more sessions, at the cost of more cross-server traffic and larger failure domains; measuring that trade-off beyond four servers is next.

Adding an SSD tier. Some agent sessions pause for many minutes. Mooncake’s SSD offloading could keep their history at a fraction of the cost of DRAM.

Making room on decode. Decode keeps each running request’s KV in GPU memory, which was nearly full on the single server at 112 sessions. More decode GPUs per prefill instance and more KV per decode GPU are next.

RoCE congestion control across nodes. Many-to-one and one-to-many RDMA traffic (prefill-to-decode KV transfer and Store restores across servers) needs working congestion control, ECN/DCQCN or PFC, so that bandwidth is shared evenly across GPUs. NeuroSpark is fixing this on its four-server fabric and will rerun the AgentX sweep with it in place.

Acknowledgments

We thank the Mooncake community for Mooncake Store and Transfer Engine and for working with us on this deployment, SemiAnalysis for AgentX and InferenceX, and AMD’s ATOM, AITER, FlyDSL, and ROCm teams for the serving engine, kernels, and software this work builds on. In particular, we are grateful to:

• AMD: Yajie Zhang, Hattie Wu, Lingpeng Jin, Carlus Huang, Yuechao Guo, Honglie Yi, Lirong Zhang, Peng Sun, Emad Barsoum, Anatolii Rozanov*, Andy Luo*
• NeuroSpark AI: Runzhe Hao, Calvin Fei, Ivan Shan, Thomas Ning
• Mooncake: Lzy17, Feng Ren, staryxchen, Wenjie Xu, Teng Ma, Ke Yang

* Anatolii Rozanov and Andy Luo are also maintainers of Mooncake.

ATOM’s Mooncake Store connector (ROCm/ATOM#2464) and its prefill-to-decode transfer path (ROCm/ATOM#2431) are open source.

Related Links

• ATOM: https://github.com/ROCm/ATOM
• Mooncake project: https://github.com/kvcache-ai/Mooncake
• AgentX benchmark: https://github.com/SemiAnalysisAI/agentx-harness
• AgentX trace corpus: https://huggingface.co/datasets/semianalysisai/cc-traces-weka-062126
• InferenceX results for this deployment: https://github.com/SemiAnalysisAI/InferenceX/pull/4045

Appendix: Test Environment and Reproducibility (Four-Node GLM-5.3)

This appendix covers the four-node GLM-5.3 measurements in the NeuroSpark section (Figures 7 and 8).

Hardware and system (identical on all four servers unless noted)

Software

Configuration