sched: pipeline the delivery of a host-resident KV cache - #39
Piggidragon wants to merge 38 commits into
Conversation
|
Review of head 175236b80f2d9ef6b2b67076fb1d452fbf935d18: changes requested before merge. The CUDA performance result looks credible, and the declining relative gain is consistent with the overlap ceiling: ordered time is approximately copy + compute, while pipelined time approaches max(copy, compute) plus residual overhead. As context grows, KV traffic grows while the amount of compute available to hide it does not grow at the same rate. NCU is not needed to explain that curve. A short nsys trace is useful after the correctness fixes, but it is not a substitute for them. Blocking defects1. The incomplete meta/tensor-parallel path is now reachableThe generic eligibility check selects any non-CPU backend with async tensor set and event interfaces: This PR adds those interfaces to the meta backend: A normal partial-prefix delivery sends the changing tail with a nonzero destination offset: Meta hard-asserts that the offset is zero and also requires whole-chunk granularity: The design document says the strided delivery is still missing and the meta work has not been run: This means the documentation claim that tensor parallelism remains ordered does not match the code. A host-KV tensor-parallel run can enter the pipeline and assert. Fix: explicitly exclude GGML_BACKEND_DEVICE_TYPE_META for now. Prefer moving the untested meta event and transfer groundwork into the later change that implements and validates the strided head-split write. Removing only the offset assertion is not sufficient. 2. Partial per-backend setup failure leaves stale plan membershipTransfer backend or event creation can fail here: The caller only sets r->n_staged to zero: It does not clear split_order and input_staged for that backend, unlike the budget, headroom, and allocation failure paths. If another backend succeeds, global staging remains enabled and address assignment reaches the failed backend: That path calls ggml_backend_buffer_get_base with a null ring buffer. Fix: factor one backend-decline helper that releases the ring and clears n_staged, every affected split_order entry, and every corresponding input_staged entry. Use it for every failure path. Add failure injection for one backend failing while another succeeds. 3. Crossing the budget permanently disables the featureRing eligibility rejects a backend once over_budget is set: The budget and headroom paths set that flag permanently: The assumption that a context only grows is not valid for a long-running server. The KV view size is derived from the highest currently used cell and can shrink after a sequence is removed: Once over_budget is set, the initial eligibility check prevents later small graphs from even reaching the calculation that could reconsider the ring. One long request can therefore force all later short requests on the same context onto the ordered path. Fix: re-evaluate the budget for each new graph. Keep warning suppression separate from eligibility. If allocation churn is a concern, use hysteresis rather than a permanent latch. 4. The public setters can invalidate an allocated graphDepth is documented as pre-allocation, but the restriction is not enforced. Budget has no timing restriction: Both setters can free ring storage: Ring freeing does not invalidate the staged input-copy data and buffer fields: Calling either setter after graph allocation can therefore leave the graph pointing into freed ring storage. Fix: make the configuration immutable once allocation starts and enforce that state. Passing transport configuration at scheduler construction would be safer than mutable exported setters. At minimum, budget needs the same enforced precondition as depth. 5. The claimed ggml_tensor size preservation is false on some 32-bit ABIsThe parent reserves eight trailing bytes: Lines 698 to 705 in 71b2b1b The PR replaces that with size_t while claiming sizeof(ggml_tensor) does not change: That is true on LP64, but on an ILP32 ABI such as i386, size_t is four bytes and the structure can shrink from 256 to 252 bytes. Fix: preserve an eight-byte storage slot, for example with a fixed-size union or representation, and add layout assertions for representative 32-bit and 64-bit targets. The memory-cap arithmetic also needs overflow checks. On 32-bit, sufficiently large MiB values can wrap to zero, which means uncapped: Use checked parsing, checked MiB conversion, and checked ring addition/multiplication. Major merge risks6. Backend eligibility assumes stronger event semantics than the API providesDepth 1 is enabled by default for a host-resident cache: The eligibility test uses function-pointer presence as proof of nonblocking stream-to-stream ordering. That is not true for all selected backends. SYCL blocks the host in event->wait(): WebGPU event wait calls host synchronization: This contradicts the design requirement that recycling not stop the host and can produce serialization or regressions. Fix: add a capability that specifically guarantees nonblocking event waits, or keep the feature opt-in until each backend is validated. For the current fork, a temporary CUDA-only gate is safer than generic pointer-based eligibility. 7. The late-tail path does not order against a general async producerThe original path either uses a backend-aware async copy or synchronizes the producer and consumer before a blocking copy: The staged late-tail path directly calls set_tensor_async from input->data without coordinating with input_backend: The current llama KV producer is CPU-oriented, but the new ggml API is generic. An asynchronous producer can still be writing the tail when the destination starts reading it. Fix: either restrict staging to persistent CPU-produced tensors and document that restriction, or preserve producer ordering for the late region. 8. Inputs with no stable prefix are still redirected into the ringThe public documentation says only inputs carrying a stable prefix are eligible: Actual ring membership deliberately ignores the prefix: This includes transposed V, which is explicitly assigned a zero prefix: Those inputs consume ring capacity and use the new event path while contributing no overlap. This can cause transposed-V configurations to cross the budget earlier despite the documentation saying they stay untouched. Fix: use a static transport-candidate annotation for reserve-time membership and a separate per-evaluation stable byte count. Validation and instrumentationThere are no automated test changes. The exactness harness prints hashes but does not compare them: It only fails on a request error or detected cache reuse. It is a manual measurement harness, not an exactness gate. Reuse the existing test infrastructure rather than adding a new test file. At minimum cover:
Some profiling claims are also not reproducible from this head. The document says debug level 3 showed individual copy costs: The implementation logs only tensor name and size: Also, n_stop_recycle increments when a stream wait is enqueued, not when look-ahead actually stops: The public stats getter exposes only deliveries and early/late bytes, not all the counters described by the document: Restore the instrumentation that produced the documented numbers or narrow the claims. Commit cleanupThe first two commits use Co-Authored-By for Claude. Repository policy requires Assisted-by for agent assistance: Later commits use the required trailer. Rewrite the first two before merge. Recommended development order
|
175236b to
2f70cbe
Compare
|
Addressed on head 2f70cbe.
The existing test-alloc target now covers partial, zero, and changing prefixes on a reused graph; depth 0; budget recovery; post-allocation setters; partial setup failure; meta exclusion; annotation requirement; and non-CUDA exclusion. Local validation on the final commit:
The branch is rebased onto llama/dev at 01b141f. The obsolete meta commit was removed and the first two commit trailers now use Assisted-by. CI is queued on the updated head. I could not capture the requested nsys traces locally because this environment has no working NVIDIA driver. The CUDA CI job is queued, but it does not replace the runtime trace. |
Review findings
🤖 Generated with Claude Code |
|
AI-assisted review of current head
For an eventual pristine llama.cpp port, this also needs to be stacked after the fork's host-KV-storage/accelerator-attention separation. Upstream currently moves the entire attention region to CPU when KV offload is disabled, so this scheduler patch alone has no useful host-to-accelerator path. Local checks passed: |
github-actions
Bot
added
documentation
From the moe-cache-drafting branch: --moe-expert-cache-size, --moe-expert-cache-l2-pinned-mb, --experimental-logs, the automatic grouped decode / prefetch / bias residency behaviour, and the layer-split-only and speculative interactions. Marked as coming from an unmerged branch, like the PR GenerelSchwerz#57 and GenerelSchwerz#39 material. Assisted-by: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_013Bs926e6SdsxnkrqqKq5HP
|
Review of head Blocking1. The optional ring can consume memory required by the graphThe transport planner allocates and retains the ring before graph allocation: llama.cpp/ggml/src/ggml-backend.cpp Lines 2223 to 2258 in 0933834 Only afterward does graph allocation fall back to llama.cpp/ggml/src/ggml-backend.cpp Lines 2393 to 2433 in 0933834 This can turn a graph that fits without pipelining into
Configuration is then locked here, so the caller cannot disable the ring and retry on the same scheduler: llama.cpp/ggml/src/ggml-backend.cpp Lines 3093 to 3107 in 0933834 The graph-allocation failure path should discard the current graph's optional ring bindings and retry reservation/allocation once on the ordered path. 2. The previously reported zero-input defect remains unresolvedWhen a split exists but the total number of split inputs is zero, llama.cpp/ggml/src/ggml-backend.cpp Lines 2036 to 2054 in 0933834 The new llama.cpp/tests/test-alloc.cpp Lines 1468 to 1478 in 0933834 The test passes in Release, but the exact head fails under UBSan: This is follow-up to the existing finding, not independent duplicate feedback: Handle Will slow the review1. The PR description materially contradicts the current headThe description still says that three meta-backend prerequisites ship here and that budget decline is latched. The current implementation explicitly excludes meta and requires the exact llama.cpp/ggml/src/ggml-backend.cpp Lines 2925 to 2955 in 0933834 The current design document instead says smaller later graphs are reconsidered: llama.cpp/docs/kv-transport-pipelining.md Lines 257 to 274 in 0933834 The description's measurement presentation also conflicts with the document, which says the 32k row predates the current budget behavior and needs remeasurement on this head: llama.cpp/docs/kv-transport-pipelining.md Lines 135 to 139 in 0933834 Update the description so reviewers are evaluating the implementation and validation that actually remain in this revision. 2. Changed prose and public comments extensively violate the no-hard-wrap ruleThe repository rule is explicit here: Lines 79 to 90 in 0933834 Examples include:
This is widespread rather than an isolated nit. Existing validation concernI am not duplicating the already-raised request for real CUDA data and ordering validation. It remains unresolved on this exact head: the current CUDA CI job was canceled without running, and commit Canceled job: https://github.com/GenerelSchwerz/llama.cpp/actions/runs/33486643492/job/99788163320 Existing owner feedback: Upstream comparisonNo equivalent implementation exists in current pristine upstream. The closest relevant work supports parts of the approach but does not remove the blockers above:
The scope gate otherwise passes. Fork PR #31 documented R4 and its design decision before implementation, and neither upstream ggml-org#21067 nor ggml-org#27311 is a direct duplicate. Upstream issue ggml-org#27757 concerns the separate tensor-parallel host-KV correctness defect. Verification
No additional security, lifetime, event-order, public-API, or backend-dispatch finding survived verification. |
ab1c7ae to
72df633
Compare
|
Automated preliminary review by Codex; the repository owner plans a separate manual review. Updated verdict: FAIL under pristine-upstream compatibility. The Blocking
Will slow review
Nits
Reviewed head: 6c01390 |
With --no-kv-offload the attention history lives in host RAM and reaches the accelerator on every decode token. The scheduler issued that transfer on the consumer's own stream immediately before the kernels that read it, so a token cost copy + compute in series. The bytes and the attention operations are unchanged; only the point at which the transfer is issued moves. Greedy output is byte-identical to the ordered path -- verified against a build without these changes, at every look-ahead tested, single GPU and layer-split across two. Three pieces, each load-bearing: - ggml_tensor::stable_prefix records how many leading bytes of a tensor's storage the graph about to run will not write. 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 -- but everything below the lowest written row is, and at decode depth that is essentially all of it. llama_kv_cache 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. - A staging ring the graph allocator cannot 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. The scheduler allocates the ring itself and points the staged copies at it before allocation; a ready/release event pair per slot carries the handover in each direction. Every eligible accelerator gets its own ring, cursor and budget, so a layer-split model pipelines on each device and a device with no room falls back alone. - A look-ahead that stays 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 two slots of margin, deliveries are issued after a split is enqueued rather than before, 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. --kv-pipeline-depth N, default 1, 0 restores the ordered path exactly. It only engages where a host-resident cache produces the deliveries. Because a host-resident cache exists to keep device memory free, the staging is capped outright by --kv-pipeline-budget (default 128 MiB per device) rather than by a fraction of what happens to be free. A ring is (N + 2) slots of one attention layer's K and V over the whole context, so it grows with the context: 27 MiB at 4k, 213 MiB at 32k, 1.7 GiB at 256k. Past the cap the scheduler declines and keeps the ordered path, and declining costs nothing -- the check runs before anything is allocated, the decision is latched because a context only grows, and the transfer backend is created lazily and released with the ring. Single GPU (RTX 4070, Qwen3.8-27B-UD-IQ2_M, -nkvo --kv-cpu-pinned, q8_0 K/V), A/B/A/B with reversed arm order: depth ordered pipelined gain 4,096 31.7324, 31.7363 37.0889, 37.0741 +16.9% 16,384 19.6765, 19.6854 31.5352, 31.5807 +60.4% 32,768 13.0264, 13.0254 15.5325, 15.5329 +19.3% (needs a raised budget) Server decode behind an 18,422-token prompt: 18.468 -> 30.685 t/s, +66.2%. Two GPUs (RTX 4070 + RTX 3060, Qwen3.8-27B-UD-Q5_K_M, -sm layer), both rings engaged: 13.06 -> 18.05 t/s at 4,096 and 6.86 -> 9.75 t/s at 16,384. The gain narrows with depth because compute is a shrinking share of the token, so there is less to hide the copy behind. That is arithmetic, not an implementation limit, and more look-ahead makes it worse rather than better. Tensor parallelism keeps the ordered path: the scheduler sees one meta backend there and the ring is a byte arena, while a meta buffer places tensors as per-device slices rather than at offsets. It declines rather than staging into something it cannot address. docs/kv-transport-pipelining.md carries the design, the numbers and the limits; docs/repro/ carries the scripts that produced them. Assisted-by: Claude Opus 5
…cripts llama-bench does not expose --kv-cpu-pinned or --recurrent-state-offload the way llama-server does, so two of the reproduction scripts passed flags the binary rejects. They now probe --help and pass only what it takes. The feature doc gains what is actually left: why -sm tensor keeps the ordered path (no events in the meta layer, and a ring that is a byte arena while a meta buffer places tensors as per-device slices), the correctness problem underneath it that is not this feature's, and the rest of the open list -- the transient device-memory peak, the --kv-gpu-layers comparison, and the exactness harness's dependence on baseline determinism it does not have. The decline for a backend that cannot record events now says so by name rather than falling into the generic "no backend supports" line, because the backend it catches is the meta backend and the next person to look will want to know that. Assisted-by: Claude Opus 5
The host-time breakdown the feature doc describes had no code behind it: the counters existed but nothing accumulated or printed them. GGML_SCHED_TRANSPORT_DEBUG=2 now reports the split loop as a mean over each 128 graphs, with the bytes the ordered path still moves and why the look-ahead stopped; =3 names the tensors that are still on it. That is what found the rest: 40 blocking copies a token moving 0.4 MiB, 32 of them the device-to-host KV store. The budget warning now reports what the ring costs at the full context next to what it costs now, so --kv-pipeline-budget can be sized against the number that matters. It is still applied per graph: enforcing the projection would refuse the ring for every large -c even when the window never gets near it. llama-bench gains -kvcp and -rso. Without them a host-resident run measures something else entirely -- 9.02 against 19.43 t/s ordered at 16,384 -- and the repro scripts had been silently dropping both since llama-bench lost them. The exactness harness gives every task a nonce derived from its own name and length, so no two share a prefix the server can restore, and fails a task whose prompt_n says one was reused anyway. Assisted-by: Claude Opus 5
…t it does GGML_SCHED_TRANSPORT_DEBUG=3 now reports what each remaining blocking copy cost, not just its name. It turns out one of them is almost all of it: attn_inp_k_rot, 256 KiB, 18 us on the ordered path and 3.4 ms behind one split of look-ahead. That is the copy engine, not latency. A blocking copy waits for the deliveries already queued on it, and two staged splits at 22.0 GB/s is 3.6 ms. Issuing the delivery in pieces does not help, the engine is FIFO across streams. Putting the copy on the consumer's stream so the host never blocks moves the time into the consumer wait and leaves throughput alone. The doc records both, so the next person does not spend the afternoon on it again. Assisted-by: Claude Opus 5
records@18432 was giving different answers across otherwise identical N = 0 runs, which made it useless as a gate and looked like the pipeline breaking exactness. It is not the task: asked on its own with the prompt cache off it returns the same hash three times running, at -c 32768 and at -c 65536. It is the harness. All eight tasks share one server with prompt caching on, and records@18432 is about 29.6k tokens with a task of about the same size ahead of it, so the two do not both fit in a 32,768 cache and placement depended on what was still resident. The nonce stops a prefix being restored, it does not stop the pressure. cache_prompt=false does. Two independent N = 0 passes now agree on all eight tasks, and N = 0, N = 1 and N = 4 agree on all eight. Assisted-by: Claude Opus 5
Restrict staging to annotated CUDA inputs, make backend decline complete, re-evaluate budgets, freeze scheduler configuration, and preserve tensor layout. Add regressions for prefix changes and fallback behavior. Assisted-by: OpenAI Codex
Assisted-by: OpenAI Codex
The ring laid its entries out with ggml_nbytes() and bound them by writing data and buffer directly. A buffer type may ask for more than ggml_nbytes() for a tensor -- CUDA does for a quantized one, and MMQ clears that padding -- so an entry could reach into the next one. Entries are now sized with ggml_backend_buft_get_alloc_size() and bound with ggml_backend_tensor_alloc(), which also gives them the buffer's own initialization and its bounds check. test-alloc gets a dummy buffer type whose get_alloc_size exceeds ggml_nbytes, and a two-entry ring test that checks the entries stay inside the ring and out of each other, that every byte of an entry is delivered once from the matching source offset, and that nothing waits on an event before it is recorded. llama-bench takes -kvpb/--kv-pipeline-budget and reports it. The repro scripts pass 512 and now fail closed: they refuse a build without -kvcp, -rso or -kvpb instead of dropping the option, and every arm propagates its status. The llama-bench table in the doc was measured before the budget existed, so it says so, and the 32,768 row is marked as needing a re-measurement. Assisted-by: Claude Opus 5
The rings are laid out and allocated before the graph is, so a device that can hold the graph alone but not the graph next to a ring turned into GGML_STATUS_ALLOC_FAILED. The configuration is locked by then, so the caller could not turn the ring off and retry either. When graph reservation fails the rings are now released and the reservation is retried once on the ordered path, and that scheduler keeps the ordered path from then on. A plan over a split list with no inputs left input_staged unallocated and passed it to memset, which UBSan reports even at size 0. Such a plan stages nothing, so it now returns after putting every split back on the ordered path. test-alloc gets a device capacity on the dummy backend and a test that sizes it to hold the graph or the ring but not both. The prose and public comments this branch added were hard-wrapped to a fixed column, against the repository rule. They are unwrapped, one sentence per line. Assisted-by: Claude Opus 5
The llama-bench table predated the budget and said so, and the context sweep and the exactness gate were last run before the ring allocation and binding changed. All three are re-run on an RTX 4070 with a CUDA build of this head, at --kv-pipeline-budget 512. Greedy server output is identical at depth 0, 1 and 4 across all eight tasks. Throughput is +16.1% at 4,096, +56.6% at 16,384, +20.1% at 32,768 and +13.4% at 65,536, for +28, +104, +206 and +410 MiB of device memory. The 131,072 and 262,144 arms are not re-measured and say so. The server table is replaced with the four 18,432-prefill tasks of the exactness gate, which is what this head was actually run on; the copy/compute breakdown keeps its earlier numbers and says which head they came from. Assisted-by: Claude Opus 5
… MiB The context sweep is measured with the budget raised, so its deep rows read as default behaviour when they are not: past 20,556 rows of window the default declines and those depths stay ordered. Say so, and add a finer sweep that puts the peak at 16,384 rows and shows the gain per MiB falling off as 1/rows^2 above it. The peak is where copy and compute are equal, and the ring size there works out to (n_slots / n_attn) * compute * BW - the bytes per row cancel, so the budget is quant-invariant. That predicts 101 MiB against the 102 MiB measured, which is what the 128 MiB default is sized against. Assisted-by: Claude Opus 5
A window over several streams is one view of a tensor whose streams sit end to end, and both the prefix and the delivery treated it as one flat byte range. That made the lowest-writing stream cap the stable prefix for every stream above it, and it copied the cells between one stream's window and the next, which the graph never reads. The prefix is now counted within a stream, and a staged input whose last dimension indexes streams is delivered as one range per stream. A window over one stream keeps the single flat range it had, so single-sequence timing and bytes are unchanged. Behind 8 slots of a non-unified cache this takes decode from 72.18 to 112.70 t/s, against 70.61 ordered. At one slot it measures 35.31 against 35.32 before. test_transport_multi_stream_ranges pins the delivery: every stream's window covered once from its own source offset, the unread cells between them never moved, and the early and late bytes split as the prefix says. Concurrent slots cannot be gated on output the way one sequence can - their batching varies between runs, so the same depth gives different greedy output - which is why this is a unit test. Assisted-by: Claude Opus 5
The scheduler does not own its backends, and llama_context declares its scheduler before them, so member destruction frees the backends first and sched->backends[] dangles by the time the ring is freed. A slot's release event is recorded past every kernel that reads it and dispatches through the device, so waiting on it orders the free after the consumer without touching the backend. Assisted-by: Claude Opus 5
The per-stream loop in both delivery paths is what that helper does, and it lets a backend with a 2d set issue one copy instead of one per stream. Assisted-by: Claude Opus 5
The count is per stream, so the contract has to name the stride it goes with rather than leave it as "the first nbytes". Assisted-by: Claude Opus 5
The server's batching varies between runs, so it cannot gate concurrent sequences. llama-parallel seeds its client schedule, so it can, and its clients ask different questions: with one shared prompt every stream holds the same bytes and a cross-stream read stays invisible. Fails at depth 1 on the commit before the multi-stream span fix. Assisted-by: Claude Opus 5
Assisted-by: Claude Opus 5
The parallel table was taken before the multi-stream span fix, so it reported a delivery that moved a fraction of the window. Gates 1, 2 and 5 are re-run here, and the numbers that are still from an earlier head now say so. Assisted-by: Claude Opus 5
A graph that stages nothing kept the ring and the second device context for the life of the scheduler. Give them back, from every path that ends with nothing staged. A window wider than the ring holds frees the ring and allocates it again, which a prefill did on nearly every ubatch. Allocate a slot in powers of two, capped by the full context, the budget and the headroom check, so a 16k prefill reallocates 6 times rather than 32. The decline decision still goes by what the graph needs. Wait at a graph boundary only for a ring that delivered, not for every ring that has a transfer backend. Assisted-by: Claude Opus 5
The platform check it had can never fail on a 64-bit size_t, so any value was accepted and a large one was silently the same as 0. Cap it at 65536 MiB in all three places that parse it. Assisted-by: Claude Opus 5
A staged delivery reads the host cache long after the decode that issued it returned, so a memset of those buffers races it. llama_memory_clear holds no context and cannot wait, so say so on the public function, and wait where a context is at hand. Assisted-by: Claude Opus 5
clear(data=true) memsets the buffers while a decode can still be reading them: a staged delivery for a host-resident cache, the graph itself for a device-resident one. The second one is not new and is easy to hit - llama_decode followed by llama_memory_clear(mem, true) changed the logits of that decode on every trial. Every memory type already passes the context down through init_update, so the caches keep it from there and wait on it before the memset. Assisted-by: Claude Opus 5
Keep a staged input's producer on the CPU or on the consumer itself, keep the transfer context over a graph that stages nothing, and stop asking a device that cannot give one. Say in the header and the docs that the destination is CUDA-only, and that a delivering graph costs the pipelining of n_copies > 1. Assisted-by: Claude Opus 5
- Fix kv_pipeline_budget_mib handling to preserve negative sentinel value (-1 = not set, 0 = no cap) - Add null check for transport backend before synchronizing - Disable pipelined transport when n_copies > 1 to avoid conflicts with pipeline parallelism - Expose set_lctx() method in kv_cache for hybrid index context management Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PbyQFddbGJQ6MpRwUH2Rtj
Round the slot size to the ring alignment when the budget or the free memory is the binding cap: slot k starts at k*slot_size, and neither cap is a multiple of it, so every slot after the first bound its entries to a misaligned address. Reproduced on a 27B at -c 32768 with --kv-pipeline-budget 17, which aborts with a CUDA misaligned address. Wait for the consumer before the priming prefetch whenever the graph is about to stage, not only when the ring delivered last graph: a context shift or a declined graph in between left the previous graph's writes to the host source in flight. Say which split list a plan was built for with a generation counter, rather than by split count and input count, which do not distinguish two different lists. Keep the ring over a few graphs that stage nothing instead of freeing it on the first one, and keep the transfer context over a budget or headroom decline, which the next graph can recover from. Lay out the plan from scheduler-owned storage instead of a hash set and an array per graph. Clamp a stable prefix to one stream rather than to the whole body, the way the header describes it. Validate the pipeline depth and budget once, in the context constructor, against bounds the header now names, instead of restating them in three places. Assisted-by: Claude Opus 5
Split a staged delivery into per-stream ranges whatever the stable prefix says. The split described the geometry only when a prefix was set, so the same view moved a different number of bytes depending on a value that decides when bytes move, not which. Clear a fresh ring so the padding around a window is never uninitialised device memory. Keep the layers of a cache that shares cells on the ordered path. [TAG_KV_CACHE_SHARE_CELLS] gives the borrower the owner's K/V tensors, and the borrower returns from apply_ubatch before it can describe them, so they kept the owner's prefix while the borrower wrote its own rows underneath. The borrower now drops the transport flag from the tensors it takes. Bound a stable prefix by the tensor rather than by a stream the storage does not describe: the KV tensors carry their streams on ne[2], so the old clamp was the whole body and never bound anything. The reader clamps to one stream of its own view, which is where the stream count is known. Refuse a pipeline depth out of range instead of clamping it and reporting success, which made the context's own bounds check dead code and let a ggml caller get a configuration it did not ask for. Walk a per-ring list of staged splits in the look-ahead instead of rescanning the split list past every other backend's splits on each call. Release only the rings that were holding memory when the graph does not fit next to them, rather than disabling every device including those that never had one. Drop the recurrent memory's context back-pointer and its synchronize: no recurrent tensor is ever marked GGML_TENSOR_FLAG_TRANSPORT, so it guarded a race the transport cannot reach. Restore the pipeline environment variables from a scope guard in test-alloc, so an assert in the middle does not leave them set for the tests after it. Assisted-by: Claude Opus 5
…opt-in A staged copy kept its source's stride, so a cache split into streams reserved a whole kv_size per stream and the ring grew with n_stream rather than with the window. At the default budget it then declined its own ring: 170 MiB needed for a 42.7 MiB window, and the arm measured the ordered path. The copy packs the ranges now, and the delivery is given the source and the destination stride separately. At -npl 8 -no-kvu -c 32768 that is 70.90 -> 106.19 t/s, and unchanged where the budget was already raised past what the gaps cost. --kv-pipeline-depth defaults to 0. Any value above 0 turns it on, and 1 is the value that measures best. Also from review: - do not compute the ring layout inside GGML_ASSERT - arm a slot's release event even when its split fails, so freeing the ring waits for a late copy that is already on the consumer stream - clear input_staged when the plan allocation fails, so it cannot disagree with split_order - reset reported_no_room once a ring is allocated again - exclude CPU devices where the device type is checked, not after the registry name, where it could never be taken - zero the whole trailing storage of ggml_tensor again on a 32-bit target - fix an unsigned underflow in the test allocator's capacity check Assisted-by: Claude Opus 5
The ordered copy of a multi-stream window landed in llama/dev on its own. It brings its own copy of the range geometry, which this branch already had in a form that also carries the copy's stride and the stable prefix, so drop the duplicate and keep the one the ring uses. The ordered copy now steps by the copy's own stride. A copy is laid out for the ring only while it is staged, so the two agree today, but reading it from the copy is what keeps them agreeing. The dummy backend's set_tensor_async records a delivery and now performs it as well when the backend is asked for real memory, so the test that checks the bytes of the ordered copy still sees them. ggml_new_tensor_impl needs a brace for the array inside the union. Assisted-by: Claude Opus 5
The ordered copy of a multi-stream window landed in llama/dev, so the baseline the parallel table compares against moves. Only the ordered column of a cache split into streams changes; the rest is within noise. The pipeline is worth +8.5%, +13.8%, +19.8% and +26.0% at 1, 2, 4 and 8 slots over a non-unified cache, against +72% and +73% read off the old baseline, and packing the ring is worth +25.7% at the default budget rather than +49.6%. The single-sequence table does not move at all: a unified window is one range and has no gaps to have been wasting. Gate 5 is re-stated against a build of llama/dev rather than an earlier head of this branch, which is the stronger comparison. Assisted-by: Claude Opus 5
4c38596 to
da23680
Compare
Also gives the stable prefix one value per stream, so a slot that was just reset no longer caps the early region of every other stream. Assisted-by: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JycLdWs6KRgizbdnNAfZqM
Assisted-by: Claude Opus 5
2 participants
Pipelines the host-to-device delivery of a host-resident KV cache: the transfer is issued one split ahead, on a transfer stream of its own, so a decode token stops paying transfer and attention in series.
Supersedes #38, which targeted
beellama/dev; this is the same work rebased ontollama/dev.What it does
Four pieces, each load-bearing:
ggml_tensor::stable_prefixrecords it on the tensor that owns the storage;llama_kv_cache::update_stable_prefixes()sets it fromapply_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. The prefix says nothing about the next graph, whose ubatch can land below the window the previous one is still delivering, so the scheduler waits for the transfer stream and the consumer once at the top of each evaluation - the guarantee the ordered path gets from its blocking copy, paid once per graph rather than once per split.ggml-alloc's reach.ggml-allocmay 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 an earlier cross-layer prefetch experiment non-exact. The scheduler allocates its own ring and points the staged copies at it before allocation; aready/releaseevent pair per slot carries the handover each way. Every eligible accelerator gets its own ring, cursor and budget, so a layer-split model pipelines on each device and a device with no room falls back alone.Lsplits ahead recycles the slot of the splitL - n_slotsback, son_slots == L + 1recycles 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, not 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".Turning it on
--kv-pipeline-depthdefaults to 0, so this changes nothing about an existing run until it is set.Any
Nabove 0 turns the pipeline on and0is the ordered path exactly. The same value isllama_context_params::kv_pipeline_depth,-kvpdinllama-bench, andLLAMA_ARG_KV_PIPELINE_DEPTHin the environment;GGML_KV_PIPELINE_DEPTHsets the scheduler default under all of them.N = 1is the recommendation and not just the smallest value that works: deeper look-ahead is worse at every depth measured (at a 19,246-token prompt, 29.83 t/s atN = 1, 28.44 atN = 2, 25.93 atN = 4), and the ring costs(N + 2)slots of device memory.It is off by default because it spends device memory on a cache that is on the host to save it, and because the numbers below are one model on one link. Turn it on, look at
GGML_SCHED_TRANSPORT_DEBUG=1, and keep it if the deliveries convert. It only engages where a host-resident cache produces the deliveries; a device-resident run never creates the transport.CUDA only. The scheduler enables the ring for a device whose backend registry is named
CUDA, that has asynchronous transfers and events, and whose default buffer type the scheduler was configured with. Meta devices are excluded outright, so-sm tensorkeeps the ordered path. Everything else - SYCL, WebGPU, Vulkan, the CPU - is untouched and stays ordered.Measurements
Re-measured on this head, on an RTX 4070 (sm_89, 11,902 MiB usable), driver 610.57.04 / CUDA 13.3, i5-13400F,
Qwen3.8-27B-UD-IQ2_M.gguf,-ngl 99 -sm none -mg 0 -t 3 -fa on -ctk q8_0 -ctv q8_0 -b 512 -ub 512, host residency-nkvo --kv-cpu-pinned --recurrent-state-offload,--kv-pipeline-budget 512, everything undertaskset -c 0,2,4.llama-bench, A/B/A/B with reversed arm order (docs/repro/r4-kv-pipeline-ab.sh), both passes shown, plus peak device memory from the context sweep:The first three rows are re-run on this head; the 65,536 row is from the context sweep on an earlier one.
llama-server, greedy, the four 18,432-prefill tasks of the exactness gate at-c 32768:Depth 4 is slower than depth 1 on all four (27.031, 26.626, 15.926, 15.864), which is why
N = 1is the value to turn it on with.These are within 0.3% of the same table taken before the cross-graph wait was added, which is the change that could have cost them. It costs nothing here or on
llama-bench: with a host-resident cache the host is the bottleneck, so by the time the next graph starts the previous graph's deliveries have already retired and the wait finds nothing to wait for. Timed directly at 16,384 over 128 graphs of steady-state decode, the consumer sync and the transfer sync are both 0.00 ms of a 31.55 ms graph, and 1.39 ms per graph during prefill only. Ordering the two streams with an event instead of blocking the host was measured against this and has nothing to win.Where the token goes, and why the gain narrows
GGML_SCHED_TRANSPORT_DEBUG=2is implemented in this PR (the counters existed but nothing accumulated or printed them). Per decode graph behind a 19,246-token prompt:ggml_backend_tensor_copy644 MiB in the 28.30 ms the ordered arm spends on the same bytes is 22.0 GB/s, and
nvidia-smireports the card at gen4 x16 - about 88% of what the link does in practice. The pipeline came within 5% of themax(copy, compute)ceiling at both 19k and 48k. The narrowing gain is therefore arithmetic: copy grows with context, compute does not, and once copy dominates there is only compute left to hide behind it.=3names the tensors still on the ordered path.Two attempts to recover the last 3.57 ms, both recorded in the doc as negative results: issuing the delivery in pieces so a blocking copy can interleave does nothing (the copy engine is FIFO across streams, and small pieces cost throughput), and putting the copy on the consumer's stream moves the time into the consumer wait without changing throughput. That 3.57 ms is almost entirely one 256 KiB graph input queued behind the deliveries - it is bandwidth, not latency.
The lever is bytes, not the link
-ctk q4_0 -ctv q4_0halves the traffic, and what that is worth depends on which side of the crossover you are:+5.6% at 19,246 (the consumer wait is 27.31 ms at q8_0 and 27.26 at q4_0 - the pipeline had already reached the compute floor, so the removed bytes were bytes nothing waited for) and +56.6% at 48,042, where copy still dominates.
The ring also wins per MiB against
--kv-gpu-layers. At 19,246 with-c 32768, a device-resident layer costs ~68 MiB and the ring ~205 MiB:--kv-gpu-layers48Four device-resident layers are worth +14.3% on the ordered path and +1.4% on the pipelined one, for more memory than the ring costs.
Pinning is worth as much as the pipeline and is off by default. Behind a 13,128-token prompt: 21.582 -> 32.252 pinned, 14.945 -> 22.709 unpinned.
Parallel sequences
Measured with
llama-batched-bench, 2,048 prompt tokens per sequence,-c 32768 -np 8, generation t/s. Each cell is its own process, because the headroom guard reacts to what a process has already allocated: run as the last step of a sweep, the 8-slot unified ring is refused and reads 83.88 instead. Three passes, spread at most 0.05 t/s:-nplThe pipeline is worth +8.5%, +13.8%, +19.8% and +26.0% at 1, 2, 4 and 8 slots over a non-unified cache, and about the same over a unified one. The two caches now measure the same in both arms.
These are lower than this description carried before, and the earlier baseline was wrong rather than the gain being smaller. The ordered copy of a window over several streams went through
ggml_backend_tensor_copy, which movesggml_nbytes()- the span the ranges are cut from - so it copied every gap between them and the streams-ordered column read 34.09, 49.44 and 70.92 at 2, 4 and 8 slots. Fixed in #77, which is merged and which this branch is rebased onto, so the numbers above are what the pipeline is worth against a baseline that moves the right bytes.These numbers are lower than the ones this description carried before, and the earlier ones were wrong rather than better. The multi-stream delivery sized a stream's range from
ne[2]*nb[2], which is one KV cell and not the window, becausebuild_attn_mhapermutes the window before the scheduler sees it and its rows sit on dimension 1. A ubatch spanning several streams therefore delivered one cell per stream and attention read whatever the ring slot held before - silently, with fluent output. On the same machine the predecessor reports 61.04, 90.81 and 108.23 at 2, 4 and 8 slots against 58.45, 85.45 and 105.63 here; the difference is the cost of copying the right amount.A slot holds the window, not the cache it is cut from. A staged copy kept its source's layout, and in a cache split into streams that layout steps a whole
kv_sizefrom one stream to the next while the graph reads onlyn_kvof it. Sizing the slot from that stride reserved every gap the delivery skips, so the ring grew withn_streaminstead of with the window - and the ring is what the budget is applied to, so at the default budget the feature declined its own ring and did nothing. The copy packs the ranges now, withcudaMemcpy2DAsyncgiven the source stride and the packed stride separately; the bytes delivered do not change, only where they land.llama-batched-bench -npp 2048 -ntg 128 -npl 8at-c 32768 -no-kvu, one process per cell, generation t/s:At the default budget the unpacked ring wants 170 MiB for a window worth 42.7 MiB, so it is declined and that arm reads the ordered path's own throughput. Packed, the same run keeps the ring and gains +25.7%. The unpacked arm is a build of the commit before the packing one, not an inference. Where the budget was already raised past what the gaps cost, both are the same speed and the packed ring is slightly smaller. A unified cache is one range per delivery and is unaffected either way, which the single-sequence A/B confirms: 18.976 -> 29.779 packed against 18.988 -> 29.784 before, at 16,384.
Concurrent slots can be gated on output after all. The server cannot do it - its batching varies between runs, and three runs at
N = 0gave three different hashes - butllama-parallelseeds its client schedule, so the batches repeat.docs/repro/r4-kv-pipeline-parallel-exact.shcompares the transcripts of 8 concurrent sequences over a non-unified cache and is identical atN = 0,N = 1andN = 4, and atN = 0andN = 1with-sm layeracross both devices. Its clients ask different questions, which is what makes it a gate: with one shared prompt every stream holds the same bytes and a cross-stream read is invisible. Run against the commit before the fix it fails atN = 1on the first sequence.Device memory
A host-resident cache exists to keep device memory free, so the staging is capped outright by
--kv-pipeline-budget(default 128 MiB per device), not by a fraction of what happens to be free. A ring is(N + 2)slots of one attention layer's K+V over the whole context, so it grows linearly with context.The cap is applied to what the current graph needs, and the warning reports the full-context figure alongside it so the budget can be sized against the number that matters. Enforcing the projection instead was tried and reverted: it refuses the ring for every large
-ceven when the window never gets near it, which at-c 32768turned the feature off by default.A budget or headroom decline is per graph, not latched. The plan is remade for every graph, so a context that outgrows the budget gives the ring back and a later, smaller live window takes it again. The cost is that transient, and it is bounded by the budget the user already authorised. Declining is otherwise free: the ring and the transfer backend's device context are both released, and the transfer backend is created lazily in the first place.
The ring never starves the graph. The rings are laid out and allocated before the graph is, and the configuration is locked by then, so a device that can hold the graph alone but not the graph next to a ring used to fail allocation outright with no way for the caller to retry. If graph reservation fails, the rings are now released and the reservation is retried once on the ordered path, and that scheduler keeps the ordered path from then on.
test_transport_releases_ring_for_graphsizes a dummy device to hold the graph or the ring but not both.The headroom guard was measured and left alone. The ring declines unless it can leave
GGML_SCHED_TRANSPORT_HEADROOM(512 MiB) free. The compute buffer on this configuration is 617 MiB and this fork resizes compute buffers at run time, so the guard is about the size of the thing it exists to leave room for. Two candidate changes were checked and neither survives: sizing it against the compute buffer instead of a constant is stricter here (617 against 512), and lowering it is not supported by the cases that looked like evidence for it - run on their own, both the 8-slot unified ring and-d 131072fit and pipeline at the shipped 512 MiB. The guard refuses only under accumulated pressure inside one process, which is what it is for.Tried and reverted: planning the ring during reserve so
ggml-allocwould not budget blocks for the staged copies. It reclaims nothing, andtest_transport_fallback_keeps_allocator_planpins the invariant: the compute buffer is the same size on both arms.Validation
N = 0,N = 1andN = 4, one sequence - all eight tasks (prose, source code, JSON records, dialogue, at 2k and 18k), re-run on this head with a CUDA build:docs/repro/r4-kv-pipeline-exact.sh 0 1 4exits 0 with every hash matching depth 0.docs/repro/r4-kv-pipeline-parallel-exact.sh 0 1 4exits 0, and again at-sm layerover both devices, where each device carries its own ring. It fails atN = 1on the commit before the multi-stream span fix, so it is a gate and not a smoke test.llama/devitself, after the rebase onto sched : copy the ranges of a multi-stream window, not the span they sit in #77:17f946c340db110bwith-sm noneanddb661b7a08686b97with-sm layerover both devices, the same hash from a build ofllama/devand from this branch atN = 0,1and4. Comparing against the base rather than against an earlier head of this branch is the stronger check. Packing moves where a range lands in the slot and nothing about the bytes, so an unchanged hash is what to expect; the gate is there because a stride mistake would not look like one.test_transport_multi_stream_rangesnow checks the source offset against the source stride and the destination against the packed stride, and asserts the packed stride is the smaller of the two.N = 0passes agree on all eight single-sequence tasks, which was not true before this PR.records@18432used to give different hashes across otherwise identical runs and looked like the pipeline breaking exactness. It is not the task: asked on its own with the prompt cache off it returns the same hash three times running, at-c 32768and-c 65536. It was the harness - all eight tasks share one server with prompt caching on, andrecords@18432is ~29.6k tokens with a task of about the same size ahead of it, so the two do not both fit in a 32,768 cache and placement depended on what was still resident. The harness now setscache_prompt: false, and gives every task a nonce so no two share a restorable prefix.test-alloccovers entry allocation and bounds against a buffer type whoseget_alloc_sizeexceedsggml_nbytes, per-byte delivery coverage from the matching source offset, event ordering (nothing waits on an event before it is recorded), the empty graph, the budget decline and recovery, per-backend partial failure, the meta and non-CUDA exclusions, and the ring-starvation fallback.test_transport_multi_stream_rangesbuilds its window in the shape attention actually hands the scheduler - permuted, rows on dimension 1 - which is what the earlier version of that test got wrong. Clean under UBSan.llama_contextpasses a depth of 0.Relationship to upstream
This is fork-only work and does not stand alone on pristine upstream. The scheduler change here is generic, but the path that makes it useful in
llamadepends on three things upstream does not have:--kv-cpu-pinned,--recurrent-state-offload, and the fork's separation of host KV storage from accelerator attention compute. Upstream moves the whole attention region to the CPU when KV offload is disabled, so there is no host-to-accelerator delivery for this to pipeline. Any upstream port has to be stacked behind that separation; this PR does not attempt it, and the measurements here are not reproducible on upstream.The scope is also large for one PR - a scheduler subsystem, public API,
llamapolicy, CLI plumbing, tests, and benchmark documentation. Splitting it is worth discussing before maintainer review; the pieces are separable, with theggmlscheduler andggml_set_stable_prefixcontract on one side and thellama/CLI/doc work on the other.Tensor parallelism
-sm tensoris not pipelined, and nothing in this PR moves it closer. The scheduler excludes meta devices explicitly and requires theCUDAregistry name, so a tensor-parallel run keeps the ordered path. Enabling it needs a validated strided head-split write for a host-resident cache, which is not here. An earlier revision of this branch carried three meta-backend prerequisites; they are not in this revision.Independently of that,
-sm tensortogether with--no-kv-offloadis currently incorrect, which is a pre-existing fork defect and not caused by this PR. Fixed separately in #48; reported upstream as ggml-org#27757.Same build, same prompt, greedy:
-sm layer+-nkvo3e127464e8b901a7-sm tensor+ device KV3e127464e8b901a7-sm tensor+-nkvo9b82be0158a2fa4dCause. TP splits attention by head, but a host-resident cache is one undivided tensor, so the scheduler's copy is classified
MIRROREDand the whole window goes to both devices. With 24 query heads split 12/12 and 4 KV heads mirrored, the kernel derives the GQA ratio from the tensors it is handed - 12/4 = 3 instead of 6 - and the second device's queries, renumbered from 0, read the first device's keys. With an uneven split the same fault surfaces as a crash:fattn.cu:371 GGML_ASSERT(Q->ne[2] % K->ne[2] == 0), since 24 split 13/11 is not divisible by 4.Still to do
-sm tensor+--no-kv-offload(ggml-meta : split a host-resident KV cache by head #48). Until then it is wrong, not slow.-sm tensorcan be pipelined at all.Rebased onto
llama/devafter #77 merged. Re-run on this head:r4-kv-pipeline-parallel-exact.sh 0 1 4(exit 0, and again at-sm layer, both matching a build ofllama/dev),r4-kv-pipeline-ab.sh 4096 16384 32768(exit 0, within 0.3% of the table above), the parallel sweep, the packed-against-unpacked comparison,test-alloc41/41 in Release against a CUDA build,git diff --check.The rebase needed three reconciliations that the automatic merge did not catch: #77 brings its own copy of the range geometry, which this branch already had in a richer form, so the duplicate is dropped; the ordered copy now steps by the copy's own stride, which is what keeps it correct once a copy can be laid out for the ring; and the dummy backend's
set_tensor_asyncrecords a delivery and now performs it too, so #77's byte-level test still sees the bytes. The context sweep at 65,536 and the telemetry breakdowns are from earlier heads and are marked as such.🤖 Generated with Claude Code
https://claude.ai/code/session_01VKqPr9qVGfqX1wxNfbCQsK