Inference disaggregation (part 2): moving the cache ·

Part 1 showed why you’d split prefill and decode onto different GPUs: a long prefill on the same GPU makes everyone who is already decoding there wait. Splitting has a price, though. The KV cache that prefill builds has to move to the GPU that does decode. The offload post (part 3 of Watching a KV cache grow) predicted that moving a 32k-token cache would add about 3% to prefill over a 400 Gb/s link and about 12% over 100 Gb/s. This post measures the fastest case, within one machine, on real hardware.

The setup is the simplest version: two GPUs in the same machine, connected by NVLink. Prefill runs on GPU 0, the cache moves to GPU 1, and decode continues there.

The short answer: moving a 1.8 GiB cache over NVLink took 5.7 ms, half a percent of the 1.2 seconds prefill took to build it. Sending the cache layer by layer while prefill runs, which is how the offload post said real systems hide the transfer, didn’t help here.

Setup Link to heading

Machine:      2x NVIDIA H100 SXM (80GB HBM3), NVLink between them (Modal)
Software:     Python 3.11, PyTorch 2.14.0, Transformers 5.17.0
Model:        Qwen2.5-7B-Instruct, one copy on each GPU (57,344 bytes of KV per token)

The new script is kv_disagg.py, run through the same modal_run.py as the offload post.

First, how are the GPUs connected? Link to heading

Before timing anything, I checked what’s between the two GPUs. nvidia-smi shows 18 NVLink links on GPU 0:

$ nvidia-smi nvlink -s -i 0
GPU 0: NVIDIA H100 80GB HBM3
	 Link 0: 26.562 GB/s
	 ...
	 Link 17: 26.562 GB/s
Peer access GPU 0 -> GPU 1: True
Measured GPU 0 -> GPU 1 copy bandwidth: 395 GB/s

On an HGX board those links most likely go to NVSwitch chips rather than straight to GPU 1, and NVIDIA quotes 450 GB/s per direction. The number that matters is the measured one: a 1 GiB copy from GPU 0 to GPU 1 ran at 395 GB/s. For comparison, the offload post’s best host-memory path (pinned RAM over PCIe) ran at about 27 GB/s, and a 400 Gb/s network link is 50 GB/s.

What I measured Link to heading

For each context length (1k, 4k, 16k and 32k tokens), three ways of getting to the first decode step:

Colocated: prefill and decode both on GPU 0. This is the baseline.
Bulk: prefill on GPU 0, then copy the whole cache to GPU 1, then decode there.
Layer by layer: prefill on GPU 0 with a hook on each of the 28 layers. As soon as a layer finishes, its K and V start copying to GPU 1 on a separate stream, while the next layer computes.

The hook is small. Each layer’s K/V are already in the cache by the time the layer returns:

def fn(module, args, output):
    ready = torch.cuda.Event()
    ready.record(torch.cuda.current_stream(PREFILL_GPU))
    with torch.cuda.stream(self.stream):
        self.stream.wait_event(ready)
        layer = self.cache.layers[i]
        self.dst[2 * i].copy_(layer.keys, non_blocking=True)
        self.dst[2 * i + 1].copy_(layer.values, non_blocking=True)

I also timed the bulk copy on its own, from a cache that was already built. My first run computed the transfer cost as “prefill plus copy” minus “prefill alone”, but prefill varies by up to 15 ms between runs, which swamped a copy of a few milliseconds. Timing the copy by itself gives a clean number. Each result is the median of 9 runs.

As in the offload post, the cache has to arrive intact. For every length, 20 tokens decoded on GPU 1 from the transferred cache matched the colocated run on GPU 0 exactly, tokens and logits, for both transfer methods.

Results Link to heading

Cost of moving the cache to the decode GPU, vs prefill (Qwen2.5-7B, median ms)

prefill on GPU 0 measured on 2x H100 predicted (bytes / nominal link speed)

Data table

Drag the slider to change the context length. Bars are scaled to prefill. The dashed rows are the offload post’s network predictions for the same cache, not measurements.

context   KV MiB   prefill   prefill spread   NVLink copy     layer-wise extra
    1k        56     29 ms            2 ms    1.1 ms (3.9%)    0.1 ms (0.3%)
    4k       224    109 ms           15 ms    1.4 ms (1.3%)    3.9 ms (3.6%)
   16k       896    495 ms            9 ms    3.4 ms (0.7%)    8.6 ms (1.7%)
   32k      1792   1167 ms            9 ms    5.7 ms (0.5%)   23.3 ms (2.0%)

The NVLink copy costs a few milliseconds at every length, and it gets relatively cheaper as the context grows, because prefill grows faster than the cache does. At 32k it’s 0.5% of prefill. That’s 1.88 GB in 5.7 ms, about 330 GB/s, a little under the 395 GB/s of the 1 GiB test copy. The offload post predicted about 3% at 400 Gb/s and 12% at 100 Gb/s for the same cache, so NVLink is several times better than even the fast network case.

At 1k the copy is 3.9% of prefill, higher than at longer lengths. A simple two-part model fits all four lengths within about 0.2 ms: roughly 1 ms of fixed cost for the 56 separate copies (K and V for 28 layers, about 18 µs each), plus the bytes at 395 GB/s. That predicts 1.15 ms at 1k and 5.76 ms at 32k, against 1.1 and 5.7 measured. The fixed part dominates for small caches, and it’s also why the 32k copy averages 330 GB/s rather than 395.

That fixed cost matters more than it looks. As far as I know, real engines like vLLM store the cache in small pages of about 16 tokens, so a 32k cache is thousands of pieces per layer rather than one tensor. Copying those one at a time with plain copy_, at about 18 µs each, would take around 2 seconds (an estimate from the numbers above, not a measurement), longer than the prefill itself. That’s the problem transfer libraries like NVIDIA’s NIXL exist to solve: they batch and combine the copies.

Layer by layer didn’t help Link to heading

Overlapping the transfer with prefill should mean most of the cache is already on GPU 1 by the time prefill finishes. Here it didn’t pay off.

The “prefill spread” column matters for reading this. It’s the difference between the fastest and slowest of the 9 prefill runs, and at 4k and 16k the layer-wise extra (3.9 and 8.6 ms) is inside it. So at those lengths I can’t say layer-wise was worse or better, only that it didn’t help. At 32k it’s clearly worse: 23.3 ms extra against a 9 ms spread, four times the cost of copying everything at the end.

It isn’t the hook bookkeeping. Running the same hooks with the copies switched off stayed within a few milliseconds of plain prefill, inside the noise. So at 32k the copies themselves either slow prefill down or don’t overlap with it the way they should.

I checked two likely causes and ruled them out. The copies run on a separate stream on GPU 0, not on prefill’s own stream. And the cache tensors are contiguous, so the copies should use the GPU’s copy hardware rather than its compute cores. One remaining candidate is the copies competing with prefill for GPU memory, but at about 10% of the H100’s memory bandwidth that isn’t convincing on its own. I don’t know the cause yet. A timeline of what the GPU is doing (NVIDIA’s Nsight Systems) would show it, and I’ll take one when I measure this across two machines.

Whatever the cause, there was little to gain here. Over NVLink the whole transfer is already a small fraction of prefill, so there isn’t much to hide, and any slowdown the overlap causes costs more than it saves. Over a network link between two machines the transfer would be several times bigger and overlapping would have more to hide, which makes it worth measuring properly there.

The first decode step Link to heading

Once the cache is on GPU 1, the first decode step takes the same time as on GPU 0: 14 to 16 ms on both, within a millisecond of each other at every length. Nothing about the cache having moved makes decode slower.

What this means Link to heading

Within one machine, moving a KV cache between GPUs is cheap: 0.5% of prefill at 32k tokens, a few milliseconds at every length I tried. That’s the easy case for disaggregation. The harder one, and the one real deployments care about, is between machines, where the link is a network card instead of NVLink and the offload post’s predictions say the same cache costs 6 to 26 times longer to move.

Caveats Link to heading

One model, one pair of GPUs, one machine rented from Modal. nvidia-smi topo -m doesn’t work inside Modal’s container, so the NVLink evidence is the link listing and the measured bandwidth.
The transfer here is a plain copy_ of 56 whole tensors between GPUs. Real systems move paged caches through transfer libraries like NIXL, which batch the pieces.
cuDNN attention was on for prefill here (the PyTorch default) and off for decode steps, for the reason in the offload post. Part 1 had it off for prefill too, which is why its 32k prefill is 1,456 ms against 1,167 ms here.

Next Link to heading

Parts 1 and 2 together cover the within-one-machine case: splitting removes the stall, and moving the cache over NVLink costs half a percent of prefill. The case I’d like to measure next is two machines over a network link, where that cost should be several times higher.