sched: pipeline the delivery of a host-resident KV cache (R4)

[Piggidragon]

Implements R4 of the KV-offload roadmap (#31): the host-to-device delivery of a host-resident KV cache is issued one split ahead, on a transfer stream of its own, so a decode token stops paying transfer and attention in series.

Draft. Works on single GPU and on -sm layer multi-GPU. -sm tensor (tensor parallelism) is not supported and cannot be, without changes to the meta backend — see What tensor parallelism would take, written so it can be scoped independently.

What it does

Three pieces, each load-bearing:

1. A stable prefix. The KV window is not stable for a whole graph — a CPU split writes this ubatch's rows into it between one layer's attention and the next. Everything below the lowest written row is stable, and at decode depth that is essentially the whole window (98.5% of the bytes in the dual-GPU run below). ggml_tensor::stable_prefix records it on the tensor that owns the storage; llama_kv_cache::update_stable_prefixes() sets it from apply_ubatch(), before the graph is built and allocated, so the plan and the deliveries are decided against the same write position even when the graph is reused. build_graph_shift() clears it.
2. A staging ring outside ggml-alloc's reach. ggml-alloc may recycle a graph-owned input copy after its last graph-level consumer while a look-ahead transfer is still in flight — that is what made the earlier cross-layer prefetch experiment non-exact (§5 of KV offload: defect register, measurements and roadmap #31: +1.38%, not exact). The scheduler allocates its own ring and points the staged copies at it before allocation; a ready/release event pair per slot carries the handover in each direction.
3. A look-ahead clear of the ring's tail. A delivery L splits ahead recycles the slot of the split L - n_slots back, so n_slots == L + 1 recycles the split just enqueued and still running. The ring keeps 2 slots of margin, deliveries are issued after a split is enqueued, and slot recycling is ordered stream-to-stream rather than through the host. Each of those three, alone, costs the entire gain while still producing correct output — the first working version measured +0.5% and looked like "the copy just doesn't overlap".

Every accelerator the scheduler drives gets its own ring: own transfer stream, own slots, own look-ahead cursor, own budget decision. A device with no room falls back to the ordered path alone; the others keep pipelining.

--kv-pipeline-depth N (default 1, 0 = ordered path exactly). Only engages where a host-resident cache produces the deliveries; a device-resident run never creates the transport at all.

Device memory

A host-resident cache exists to keep device memory free, so what the transport costs matters as much as what it buys. A ring is (N + 2) slots of one attention layer's K+V over the whole context — linear in context length:

context	ordered	pipelined	gain	peak VRAM	Δ	ring
4,096	31.66	37.01	+16.9%	10169 → 10197	+28 MiB	27 MiB
16,384	19.64	31.43	+60.1%	10159 → 10263	+104 MiB	107 MiB
32,768	12.99	15.49	+19.2%	10161 → 10367	+206 MiB	213 MiB
65,536	7.74	8.74	+12.9%	10163 → 10573	+410 MiB	428 MiB
131,072	4.29	4.68	+9.1%	10537 → 11355	+818 MiB	855 MiB
262,144	2.25	2.24	declined	11329 → 11391	+62 MiB	not allocated

Two curves run opposite ways: the gain narrows with depth (compute is a shrinking share of the token, so there is less to hide the copy behind — arithmetic, not an implementation limit, and more look-ahead makes it worse) while the cost doubles with every doubling of context.

So --kv-pipeline-budget (default 128 MiB per device) caps each ring outright, not as a fraction of what happens to be free. Three things make declining actually free rather than nominally free:

• checked before anything is allocated, on every plan, not only when the ring has to grow;
• latched — a context only grows, so a ring that fits early prefill is not allocated only to be handed back;
• the transfer backend is created lazily and released with the ring; it is a second device context and was costing 112 MiB on a scheduler that never staged anything.

At 32,768 with the default budget, device memory settles at 10,161 MiB — identical to the ordered path — and throughput matches it. --kv-pipeline-budget 512 buys 15.49 t/s there for 206 MiB.

Tried and reverted: planning the ring during reserve so ggml-alloc would not budget blocks for the staged copies as well. It reclaims nothing — the compute buffer is 342.27 MiB either way at 32k, because it is sized by a reserve that runs before the KV cache exists and the copies then fit inside space already reserved for other transients. The ordered path's copies are free; the ring is unavoidably additive, and its size is the only lever.

Measurements

Single GPU — RTX 4070, driver 610.57.04 / CUDA 13.3, i5-13400F, Qwen3.8-27B-UD-IQ2_M.gguf, -nkvo --kv-cpu-pinned --recurrent-state-offload, q8_0 K/V, -sm none, taskset -c 0,2,4. A/B/A/B with reversed arm order; table above.

Dual GPU, layer split — RTX 4070 + RTX 3060, Qwen3.8-27B-UD-Q5_K_M.gguf (18.40 GiB), -sm layer -ts 55/45:

context	ordered	pipelined	gain	peak VRAM 4070 / 3060
4,096	6.91	7.52	+8.8%	9787→9815 / 9069→9099 MiB
16,384	4.68	5.46	+16.7%	ring 103 MiB per card

Both rings engage — CUDA0 stages 9 of 355 splits, CUDA1 stages 7 — and 98.5% of the staged bytes go early.

Validation

• Byte-identical greedy server output at depths 0/1/4/8 against a build of the parent commit: four fixed tasks, two behind an 18,422-token prompt, and later four prefill corpora (prose, source code, JSON records, dialogue) at 2k and 18k.
  • Two 18k tasks initially looked divergent. They are non-deterministic in the baseline: a second run of the control build reproduced my build's hashes and not its own. Cause is the server's similarity-based prompt-cache restoration, not this change. The harness needs a unique per-task nonce to be a reliable gate.
• Device-resident KV unaffected: 39.13 t/s on the parent commit vs 39.10 t/s here at tg128 @ d4096; the transport is never created.
• Test suite green apart from three pre-existing failures unrelated to this change: test-kvarn (D4 in KV offload: defect register, measurements and roadmap #31), test-upstream-merge-keepers-static (D10), and test-tokenizers-ggml-vocabs (a vocab GGUF in the repo is an LFS pointer, not a GGUF).

What tensor parallelism would take

Scoped separately because all of it lands in ggml/src/ggml-backend-meta.cpp — BeeLlama's tensor-parallel layer, not the KV-offload line this PR owns.

Why it does not work today. Under -sm tensor, llama.cpp:184-211 builds a single meta device wrapping both GPUs, so ggml_backend_sched sees one backend, not CUDA0/CUDA1. The transport's eligibility test (ggml-backend.cpp:2751) requires set_tensor_async + event_record + event_wait on the backend and event_new on the device. The meta backend has set_tensor_async but:

• ggml_backend_meta_i (ggml-backend-meta.cpp:2515) — .event_record and .event_wait are nullptr (lines 2529-2530)
• ggml_backend_meta_device_iface (line 178) — .event_new, .event_free, .event_synchronize are all nullptr (lines 191-193)

There is no event mechanism in the meta layer at all, and pipelining is built on ordering a transfer stream against the consumer with events. So the transport declines — correctly, rather than racing.

Four work items, in dependency order:

1. Meta device events (~40 lines). Implement event_new / event_free / event_synchronize at ggml-backend-meta.cpp:191-193. A meta event is N per-device events, one per sub-device, created via each sub-device's own event_new. ggml_backend_meta_device_count / ggml_backend_meta_device_get (lines 200, 206) already give the sub-devices.
2. Meta backend events (~40 lines). Implement event_record / event_wait at lines 2529-2530, fanning out over every sub-backend. ggml_backend_meta_n_backends / ggml_backend_meta_simple_backend (lines 2549, 2555) already give the sub-backends. Mostly wiring.
3. set_tensor_async with a non-zero offset. ggml_backend_meta_set_tensor_async opens with GGML_ASSERT(offset == 0) (line 1911; get_tensor_async has the same at 1956). The prefix delivery passes offset 0 and is fine, but the tail delivery writes [prefix, nbytes) and would assert. The function already splices by chunk offsets internally, so relaxing it is plausible rather than deep — but it is a change to a shared function and needs its own test.
4. A transfer backend that does not duplicate the NCCL communicator. This is the design decision, not a wiring job. The transport needs a second backend on the same device for its own stream. For a plain CUDA device ggml_backend_dev_init is cheap. For a meta device it runs the whole ggml_backend_meta_context constructor, which calls ggml_backend_comm_init (line 1872) across all devices — standing up a second NCCL communicator over both cards purely to issue memcpys. Options: a params flag on the meta constructor (it already takes const char * params) for a transfer-only backend that skips comm setup, or a way to share the existing comm.

What is already favourable. The ring lives in a meta buffer, so it is sliced across the GPUs by the same split state as everything else — the per-card VRAM cost is roughly the ring divided by device count, not the whole ring. And nothing in items 1-3 is conceptually hard; item 4 is the only one that needs a decision rather than an implementation.

Validation gates for it should mirror this PR's: byte-identical greedy output against -sm tensor with --kv-pipeline-depth 0, A/B/A/B throughput, and per-card peak device memory.

Known limitations

• -sm row is unrelated to this PR but worth recording: the CUDA backend registers no ggml_backend_split_buffer_type, so -sm row fails at model load in this tree regardless of these changes.
• The budget is applied per graph, so a context that grows past it still shows a transient peak (+112 MiB at 32k) before the latch trips, against +0 in steady state. Deciding against the context's final size needs KV geometry the scheduler does not have.
• At 128k the ring's 818 MiB is about three attention layers' K+V; making three of sixteen device-resident with --kv-gpu-layers would remove ~19% of H2D traffic against the 9.1% the ring buys. Unmeasured, and it moves with layer count and card.

docs/kv-transport-pipelining.md carries the design, the numbers and the limits; docs/repro/ carries the scripts that produced them.

🤖 Generated with Claude Code

https://claude.ai/code/session_01SRa5ZTc1bpnk1fE3mrDcSh

[Piggidragon]

Superseded by #39, which is the same work rebased onto llama/dev instead of beellama/dev.

beellama/dev is 301 commits behind llama/main and is missing five meta-backend commits, so it is not the right base for this. Closing in favour of #39.