AMD's ATOM inference engine gains a shared Mooncake KV pool, roughly doubling throughput over split pools on four MI350X servers
Synopsis
Coding agents resend the same long conversation, so whether the server still holds that history outside GPU memory decides how many tokens each turn must recompute. AMD's open-source ATOM inference engine, which places prefill and decode on separate GPUs, now connects to Mooncake Store to merge both halves of host memory into one KV pool that prefill saves to and restores from. On four AMD Instinct MI350X servers, the shared pool reached about twice the throughput of per-server pools at 96 and 112 concurrent sessions, with KV reuse held at 94–96%.
Interpretation
Agentic coding requests almost always resend the same conversation, so a server that still holds that history as KV only has to compute the small number of new tokens. Once prefill and decode sit on separate GPUs, the server's host memory is split along NUMA nodes as well, and the decode half holds almost no history, so capacity exists on the side prefill cannot reach. On one 8-GPU MI355X server, ATOM with Mooncake Store reached about 44K tokens/s per GPU at 96 and 112 concurrent sessions, with prompt-token reuse above 95% at every load, close to the 96–97% the traces allow.
ATOM places history that prefill GPU memory cannot hold into Mooncake Store and fetches it back for the next turn, computing only tokens that neither tier retained. Mooncake Store separates the services that provide memory from the clients that use it, running one service per NUMA node at 768 GiB of DRAM each, with a Master tracking object locations and evicting least-recently-used objects when the pool is 90% full. That gives one server a 1.5 TiB KV pool, nearly four times the KV space available on the prefill GPUs; at 112 sessions the pool served 46% of all prompt tokens while GPU memory's share shrank as sessions grew.
A pool hit becomes a prefill hit only when the history returns aligned to the prefill pipeline, since a prefix counts only as far as all four stages hold their layers. Each pipeline stage stores its own layers of each 256-token chunk as a separate object, and Mooncake keeps a chunk's four objects in one group so eviction never leaves a chunk only some stages can load. 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, and a request waits on average 0.4–1.2 s from its first stage starting to its last stage finishing, rising with load.
The prefill-to-decode KV handoff gathers on prefill, moves large RDMA writes, and scatters on decode, cutting the number of network operations sharply. Prefill splits by layers while decode splits each sequence by tokens, so a direct copy degenerates into one tiny write per token per layer. For one 8,192-token chunk of one 20-layer stage, a direct copy needs 163,840 writes of 576 bytes each, while the gathered chunk of about 94 MB takes 12 writes.
Perspective
The demonstrated scope is an ATOM plus Mooncake deployment on one 8-GPU MI355X server and on NeuroSpark's four 8-GPU MI350X servers, under AgentX-replayed coding-agent sessions and production-style replay traffic; the trade-off of sharing one pool across more than four servers is listed as the next thing to measure.
At 112 concurrent sessions the single server still holds 44K tokens/s per GPU, but decode's KV memory is a median 91% full, requests begin to queue in front of prefill, and median time to first token reaches about 7 s; widening the pool beyond four servers would add cross-server traffic and larger failure domains, a trade-off not yet measured.