From 658bb4ccea7057e262f0e21f70ec11344feb537b Mon Sep 17 00:00:00 2001 From: piggidragon Date: Tue, 25 Aug 2026 19:20:41 +0200 Subject: [PATCH 1/6] sched: pipeline the delivery of a host-resident KV cache MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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: 26.80 ms of blocking host-to-device against 25.14 ms of consumer wait, per decode graph at 18.5k tokens of context. This is R4 of the KV-offload roadmap (PR #31). The bytes and the attention operations are unchanged; only the point at which the transfer is issued moves. Greedy server output is byte-identical to the ordered path, at every look-ahead tested and against a build of the parent commit. Three pieces: - 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 the 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 619 of 620 MiB. llama_kv_cache sets it from apply_ubatch(), before the graph is built and allocated, so the scheduler's plan and the deliveries it issues 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, and a look-ahead transfer is still in flight outside that lifetime; that is what made the earlier cross-layer prefetch experiment non-exact. The scheduler allocates the ring itself and points the staged input copies at it before allocation, so ggml-alloc leaves them alone. One ready event and one release event per slot carry the handover in each direction. - 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 that was just enqueued and is still running, and every delivery waits for the consumer -- the ordered path with extra steps. 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; a device-resident run never creates the transport at all and measures unchanged (39.13 -> 39.10 t/s at tg128 @ d4096). A/B/A/B on an RTX 4070, Qwen3.8-27B-UD-IQ2_M, -nkvo --kv-cpu-pinned, q8_0 K/V: depth ordered pipelined gain ceiling (PR #31 §3) 4,096 31.7324, 31.7363 37.0889, 37.0741 +16.9% 96.4% of 38.49 16,384 19.6765, 19.6854 31.5352, 31.5807 +60.4% 90.4% of 34.88 32,768 13.0264, 13.0254 15.5325, 15.5329 +19.3% 74.6% of 20.83 Server decode behind an 18,422-token prompt: 18.468 -> 30.685 t/s, +66.2%. Device allocation high-water at 32,768: 10,271 -> 10,477 MiB, +206 MiB. 32,768 is the weak point and is documented as such: the copy per staged split exceeds the compute between staged splits there, and one split of look-ahead cannot cover it. Raising the look-ahead is not the fix -- N = 1 measured best at every depth, which is why it is the default. docs/kv-transport-pipelining.md carries the design, the numbers and the limits; docs/repro/ carries the two scripts that produced them. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01SRa5ZTc1bpnk1fE3mrDcSh --- common/arg.cpp | 14 + common/common.cpp | 1 + common/common.h | 1 + docs/kv-transport-pipelining.md | 183 +++++++++ docs/repro/r4-kv-pipeline-ab.sh | 31 ++ docs/repro/r4-kv-pipeline-exact.py | 46 +++ docs/repro/r4-kv-pipeline-exact.sh | 33 ++ ggml/include/ggml-backend.h | 24 ++ ggml/include/ggml.h | 18 +- ggml/src/ggml-backend-meta.cpp | 5 +- ggml/src/ggml-backend.cpp | 619 +++++++++++++++++++++++++++++ ggml/src/ggml.c | 13 +- include/llama.h | 5 + src/llama-context.cpp | 5 + src/llama-cparams.h | 1 + src/llama-kv-cache.cpp | 48 +++ src/llama-kv-cache.h | 7 + 17 files changed, 1051 insertions(+), 3 deletions(-) create mode 100644 docs/kv-transport-pipelining.md create mode 100755 docs/repro/r4-kv-pipeline-ab.sh create mode 100644 docs/repro/r4-kv-pipeline-exact.py create mode 100755 docs/repro/r4-kv-pipeline-exact.sh diff --git a/common/arg.cpp b/common/arg.cpp index 8d277f345aa6..13f12028604f 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2587,6 +2587,20 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.kv_cpu_pinned = value; } ).set_env("LLAMA_ARG_KV_CPU_PINNED")); + add_opt(common_arg( + {"--kv-pipeline-depth"}, "N", + string_format("how many splits ahead the scheduler delivers a host-resident KV cache to the accelerator, so " + "that the transfer runs while the previous split computes; 0 keeps the ordered path, where a " + "decode token pays the transfer and the attention kernels in series. Only takes effect with a " + "host-resident cache, e.g. --no-kv-offload or --kv-cpu-pinned, and costs (N + 2) * (largest " + "staged split) of device memory (default: %d)", params.kv_pipeline_depth), + [](common_params & params, int value) { + if (value < 0 || value > 14) { + throw std::invalid_argument("--kv-pipeline-depth must be between 0 and 14"); + } + params.kv_pipeline_depth = value; + } + ).set_env("LLAMA_ARG_KV_PIPELINE_DEPTH")); add_opt(common_arg( {"--recurrent-state-offload"}, {"--no-recurrent-state-offload"}, diff --git a/common/common.cpp b/common/common.cpp index 8dba7b6cbad0..f35677be8c51 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1890,6 +1890,7 @@ struct llama_context_params common_context_params_to_llama(const common_params & cparams.cb_eval_user_data = params.cb_eval_user_data; cparams.offload_kqv = !params.no_kv_offload; cparams.kv_cpu_pinned = params.kv_cpu_pinned; + cparams.kv_pipeline_depth = params.kv_pipeline_depth < 0 ? 0 : (uint32_t) params.kv_pipeline_depth; cparams.recurrent_state_offload = params.recurrent_state_offload; cparams.phase_aware_workspace = params.phase_aware_workspace; cparams.live_context_workspace = params.live_context_workspace; diff --git a/common/common.h b/common/common.h index 5f545cd7710b..86dbf1563e6a 100644 --- a/common/common.h +++ b/common/common.h @@ -640,6 +640,7 @@ struct common_params { bool recurrent_state_offload = false; // for hybrid models, keep recurrent state GPU-resident even with no_kv_offload bool phase_aware_workspace = false; // resize compute schedulers between prompt and generation phases bool live_context_workspace = false; // size supported attention workspaces from the padded live KV extent + int32_t kv_pipeline_depth = 1; // splits of look-ahead for pipelined delivery of a host-resident KV cache (0 = off) bool warmup = true; // warmup run bool check_tensors = false; // validate tensor data bool no_op_offload = false; // globally disable offload host tensor operations to device diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md new file mode 100644 index 000000000000..c145359db455 --- /dev/null +++ b/docs/kv-transport-pipelining.md @@ -0,0 +1,183 @@ +# Pipelined delivery of a host-resident KV cache + +With `--no-kv-offload` (optionally with `--kv-cpu-pinned`), the attention history +lives in host RAM and has to reach the accelerator on every decode token. The +backend scheduler used to issue that transfer on the consumer's own stream, right +before the kernels that read it, so a token cost `copy + compute` in series. + +This is the R4 item of the KV-offload roadmap +([PR #31](https://github.com/GenerelSchwerz/llama.cpp/pull/31)): the transfer and +the attention arithmetic are the same as before, but the transfer is issued one +split ahead, on a stream of its own, so the copy engine retires it underneath the +kernels of the split before it. + +`--kv-pipeline-depth N` controls it. It is on by default at `N = 1` and only has +an effect where a host-resident cache produces the deliveries; `0` restores the +ordered path exactly. + +## What it changes, and what it must not + +R4 pipelines *deliveries*, not attention. Every byte and every attention +operation is the same as on the ordered path; only the point at which the +transfer is issued moves. Greedy server output is byte-identical, and that is a +gate, not an aspiration -- see [Validation](#validation). + +Three pieces make it work. + +### 1. A stable prefix, so there is something safe to send early + +The KV window a split reads is not stable for the whole graph: the same graph +writes this ubatch's rows into it, and on a host-resident cache that write is a +CPU split that runs *between* the attention of one layer and the attention of the +next. Delivering the whole window ahead of that split would send rows that have +not been written yet. + +What *is* stable is everything below the lowest row this ubatch writes, which at +decode depth is essentially the whole window. `ggml_tensor::stable_prefix` records +that, in bytes, on the tensor that owns the storage; a view inherits the part of +it that its own byte window covers. `llama_kv_cache::update_stable_prefixes()` +sets it from the slot info in `apply_ubatch()` -- before the graph is built and +allocated, so the scheduler's plan and the deliveries it then issues are decided +against the same write position -- and `build_graph_shift()` clears it, because a +shift rewrites the body in place. + +The scheduler delivers `[0, stable_prefix)` early on the transfer stream and the +remainder at the split, once every earlier split of the graph has run. At 18k +tokens of context that split is about 620 MiB early against 1-5 MiB late. + +The prefix is a hint about *this* graph. It has to be refreshed for every ubatch +even when the graph is reused, which is why it is set from `apply_ubatch()` and +not from graph construction. Where it cannot be established -- a transposed V +cache, whose ubatch writes are scattered across the whole tensor -- it stays 0 and +the input keeps the ordered path. + +### 2. A ring the graph allocator cannot reach + +`ggml-alloc` is free to recycle a graph-owned input copy once its last graph-level +consumer is done, and a look-ahead transfer is still in flight outside that +lifetime. Writing split `k + 1`'s delivery into the scheduler's own input copies +corrupts the split still reading them; that is the defect class the earlier +cross-layer prefetch experiment hit (measurements §5 of PR #31: `+1.38%`, and not +exact). + +So the scheduler allocates its own ring and points the staged input copies at it +before the graph is allocated. A tensor that already has `data` is left alone by +`ggml_gallocr_init_tensor`, so the ring sits outside the allocator's reuse +analysis rather than competing with it. + +Each slot has one ownership cycle: + +1. the transfer stream owns an idle slot and writes one future split's prefix into it; +2. it records the slot's `ready` event, which the consumer stream waits for before launching the split that reads the slot; +3. the consumer records `release` once every kernel that reads the slot has been enqueued, and the transfer stream waits for that before overwriting the slot for a later split. + +Membership in the ring is decided once, when the ring is laid out, and execution +goes by the recorded answer. How much of a staged input can go early moves with +every ubatch; *which* input copies live in the ring must not, because their +addresses were handed out at allocation time. + +### 3. A look-ahead that stays clear of the ring's tail + +A delivery running `L` splits ahead recycles the slot of the split `L - n_slots` +back. With `n_slots == L + 1` the ring is exactly full, so every delivery has to +recycle the split that was enqueued a moment ago and is still running -- the +ordered path with extra steps. The ring therefore keeps +`GGML_SCHED_TRANSPORT_MARGIN` (2) slots behind the look-ahead, and +`--kv-pipeline-depth N` allocates `N + 2` slots. + +Two details matter as much as the margin: + +- **Deliveries are issued after a split is enqueued, never before.** Issuing them + first means the host can block on slot recycling while holding back work the + consumer could already be running. +- **Slot recycling is ordered stream to stream, not through the host.** A host + wait empties the transfer queue for as long as it blocks. + +Getting either of these wrong costs the entire gain while still producing correct +output, which is the failure mode worth knowing about: on this configuration the +first attempt measured `+0.5%` and looked like "the copy simply does not overlap". + +## Measurements + +RTX 4070 (11,902 MiB usable, sm_89), 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`, everything under `taskset -c 0,2,4`. + +`llama-bench`, `--no-warmup`, A/B/A/B with reversed arm order +(`docs/repro/r4-kv-pipeline-ab.sh`): + +| Depth | reps | ordered | pipelined | gain | `max(copy, compute)` ceiling from PR #31 §3 | share | +|---|---:|---|---|---:|---:|---:| +| 4,096 | 5 | 31.7324, 31.7363 | 37.0889, 37.0741 | **+16.9%** | 38.49 | 96.4% | +| 16,384 | 3 | 19.6765, 19.6854 | 31.5352, 31.5807 | **+60.4%** | 34.88 | 90.4% | +| 32,768 | 3 | 13.0264, 13.0254 | 15.5325, 15.5329 | **+19.3%** | 20.83 | 74.6% | + +Server decode behind an 18,422-token prompt +(`docs/repro/r4-kv-pipeline-exact.sh`): **18.468 -> 30.685 t/s, +66.2%**. + +Device allocation high-water at 32,768, sampled with `nvidia-smi` across a +`tg64 @ d32768` run: **10,271 -> 10,477 MiB, +206 MiB**, against a ring of +3 x 71 MiB. The ring competes with the model, and at 32,768 it is the largest +single new allocation the feature makes. + +Where the split-loop host time goes, per decode graph at 18.5k +(`GGML_SCHED_TRANSPORT_DEBUG=2`): + +| | ordered | pipelined | +|---|---:|---:| +| total | 52.11 ms | 30.37 ms | +| blocked in the ordered `ggml_backend_tensor_copy` | 26.80 ms | 0.15 ms | +| blocked waiting for the consumer backend | 25.14 ms | 26.69 ms | +| issuing early deliveries | 0.00 ms | 0.04 ms | +| bytes delivered early / late | 0 / 0 MiB | 619.2 / 1.3 MiB | + +The blocking host-to-device copy is gone and the consumer wait is unchanged, +which is the shape a working overlap has: the transfer left the host's critical +path without being added to the consumer's. + +**32,768 is the weak point and is reported as such.** The gain there is 74.6% of +the probe ceiling, against 90-96% at shallower depths. At that depth the copy per +staged split (about 2.9 ms) exceeds the compute between staged splits (about +1.9 ms), so one split of look-ahead cannot cover it. Raising the look-ahead does +not help: at 32,768 `N = 2` measured 15.5274 and `N = 3` measured 15.1114 against +15.5273 for `N = 1`, and at 16,384 the same sweep gave 30.34 and 29.15 against +31.56. `N = 1` is the best setting at every depth measured, which is why it is +the default. Closing the 32,768 gap is a separate piece of work, not a knob. + +Do not compare these numbers against runs on other models, prompts, cache +settings, hardware, or commits. + +## Validation + +The gates, and what was run for them: + +1. **Byte-identical greedy server output against the control.** Four fixed tasks + at `temperature 0, top_k 1, seed 1234`, plus two tasks behind an 18,422-token + prompt, hashed and compared against a build of the parent commit. Identical at + `N = 0`, `N = 1` and `N = 4`. `docs/repro/r4-kv-pipeline-exact.sh`. +2. **A/B/A/B at 4,096 / 16,384 / 32,768 with reversed arm order.** + `docs/repro/r4-kv-pipeline-ab.sh`; the table above is its output. +3. **Device allocation high-water reported.** Above. +4. **Telemetry showing the deliveries actually converted.** + `GGML_SCHED_TRANSPORT_DEBUG=1` reports the plan (staged splits, bytes per + graph, how much of it goes early, and the source buffer type); + `=2` adds the per-graph host-time breakdown above. + `ggml_backend_sched_get_transport_pipeline_stats()` exposes the same counters + to callers. + +A device-resident KV run is unaffected, and was measured to confirm it: 39.13 t/s +on the parent commit against 39.10 t/s here at `tg128 @ d4096`, with the +transport never enabled because the scheduler is given a depth of 0. + +## Scope and limits + +- Only inputs carrying a stable prefix are eligible. Everything else -- weights, + user inputs, a transposed V cache, any backend that cannot transfer + asynchronously or record events -- keeps the ordered path untouched. +- One ring, on the first backend that qualifies. Additional accelerators keep the + ordered path; multi-GPU host-KV is not covered. +- The scheduler must be configured with the device's own default buffer type. A + scheduler built on a split or host buffer type keeps the ordered path. +- `GGML_KV_PIPELINE_DEPTH` overrides the depth for tools that do not expose the + command-line option, such as `llama-bench`. diff --git a/docs/repro/r4-kv-pipeline-ab.sh b/docs/repro/r4-kv-pipeline-ab.sh new file mode 100755 index 000000000000..27e8b0ed302a --- /dev/null +++ b/docs/repro/r4-kv-pipeline-ab.sh @@ -0,0 +1,31 @@ +#!/bin/bash +# R4: pipelined delivery of a host-resident KV cache, A/B/A/B with reversed arm order. +# The two arms are the same binary: GGML_KV_PIPELINE_DEPTH=0 is the ordered path. +# +# LLAMA_KV_MODEL=/path/model.gguf docs/repro/r4-kv-pipeline-ab.sh [depth ...] +set -u +MODEL="${LLAMA_KV_MODEL:?set LLAMA_KV_MODEL to a .gguf path}" +BUILD="${LLAMA_KV_BUILD:-build}" +PIN="${LLAMA_KV_TASKSET:-0,2,4}" +LOCK=/tmp/beellama-single-gpu.lock + +run () { # $1 label, $2 pipeline depth, $3 context depth, $4 reps + GGML_KV_PIPELINE_DEPTH=$2 taskset -c "$PIN" "$BUILD/bin/llama-bench" -m "$MODEL" \ + -ngl 99 -sm none -mg 0 -t 3 -nkvo 1 --kv-cpu-pinned --recurrent-state-offload \ + -fa on -ctk q8_0 -ctv q8_0 -b 512 -ub 512 --no-warmup -p 0 -n 128 -d "$3" -r "$4" -o json \ + 2>/dev/null \ + | python3 -c "import json,sys;d=json.load(sys.stdin);print(' %-12s %.4f +- %.4f'%('$1',d[0]['avg_ts'],d[0]['stddev_ts']))" +} + +DEPTHS=(4096 16384 32768); [ $# -gt 0 ] && DEPTHS=("$@") +rc=0 +for D in "${DEPTHS[@]}"; do + R=3; [ "$D" -le 4096 ] && R=5 + echo "== context depth=$D reps=$R" + flock "$LOCK" bash -c "$(declare -f run); BUILD='$BUILD'; MODEL='$MODEL'; PIN='$PIN' + run ordered 0 $D $R + run pipelined 1 $D $R + run ordered2 0 $D $R + run pipelined2 1 $D $R" || rc=$? +done +exit $rc diff --git a/docs/repro/r4-kv-pipeline-exact.py b/docs/repro/r4-kv-pipeline-exact.py new file mode 100644 index 000000000000..0d160374d50f --- /dev/null +++ b/docs/repro/r4-kv-pipeline-exact.py @@ -0,0 +1,46 @@ +# Greedy server output, hashed. Run through r4-kv-pipeline-exact.sh. +import hashlib, json, sys, urllib.request + +PORT, NTOK, NPARA = sys.argv[1], int(sys.argv[2]), int(sys.argv[3]) + +SHORT = [ + "Explain how a B-tree index speeds up range queries in a relational database, step by step.", + "Write a Python function that merges two sorted lists, then explain its complexity.", + "Summarise the tradeoffs between optimistic and pessimistic concurrency control.", + "Describe, in order, what happens when a CPU takes a page fault on a memory-mapped file.", +] +# a long, deterministic filler so that decode runs against a deep KV window +PARA = ("A B-tree index stores keys in sorted order across a shallow, balanced tree. " + "Range queries descend once to the first qualifying leaf and then walk the leaf " + "chain sequentially, so the cost is one descent plus the size of the range. ") +DEEP = [ + "Summarise the text above in exactly five sentences.", + "List three claims the text above makes, then say which is weakest and why.", +] + +def ask(label, prompt): + body = json.dumps({"model": "m", "messages": [{"role": "user", "content": prompt}], + "max_tokens": NTOK, "temperature": 0, "top_k": 1, "seed": 1234}).encode() + req = urllib.request.Request(f"http://127.0.0.1:{PORT}/v1/chat/completions", body, + {"Content-Type": "application/json"}) + try: + with urllib.request.urlopen(req, timeout=7200) as r: + d = json.load(r) + except Exception as e: + print(f"{label} REQUEST_FAILED {type(e).__name__}") + return False + m = d["choices"][0]["message"] + # reasoning models put most of the generation in reasoning_content; hash both + body_text = (m.get("reasoning_content") or "") + "\x00" + (m.get("content") or "") + t = d.get("timings", {}) + print(f"{label} {hashlib.sha256(body_text.encode()).hexdigest()[:16]} " + f"prompt_n={t.get('prompt_n')} n={t.get('predicted_n')} " + f"tg={t.get('predicted_per_second'):.3f}") + return True + +ok = True +for i, q in enumerate(SHORT): + ok &= ask(f"short{i}", q) +for i, q in enumerate(DEEP): + ok &= ask(f"deep{i} ", (PARA * NPARA) + "\n\n" + q) +sys.exit(0 if ok else 1) diff --git a/docs/repro/r4-kv-pipeline-exact.sh b/docs/repro/r4-kv-pipeline-exact.sh new file mode 100755 index 000000000000..74b308eab470 --- /dev/null +++ b/docs/repro/r4-kv-pipeline-exact.sh @@ -0,0 +1,33 @@ +#!/bin/bash +# R4 gate 1: greedy server output must be byte-identical to the ordered path. +# Compare the hashes across pipeline depths, and against a build of the parent commit. +# +# LLAMA_KV_MODEL=/path/model.gguf docs/repro/r4-kv-pipeline-exact.sh [depth ...] +set -u +MODEL="${LLAMA_KV_MODEL:?set LLAMA_KV_MODEL to a .gguf path}" +BUILD="${LLAMA_KV_BUILD:-build}" +PIN="${LLAMA_KV_TASKSET:-0,2,4}" +PORT="${LLAMA_KV_PORT:-18099}" +NTOK="${LLAMA_KV_NTOK:-256}" +NPARA="${LLAMA_KV_NPARA:-400}" +HERE="$(cd "$(dirname "$0")" && pwd)" + +DEPTHS=(0 1 4); [ $# -gt 0 ] && DEPTHS=("$@") +rc=0 +for D in "${DEPTHS[@]}"; do + echo "== pipeline depth=$D" + LOG=$(mktemp /tmp/r4-kv-pipeline.XXXX.log) + GGML_KV_PIPELINE_DEPTH=$D taskset -c "$PIN" "$BUILD/bin/llama-server" -m "$MODEL" \ + -ngl 99 -sm none -mg 0 -t 3 -nkvo --kv-cpu-pinned --recurrent-state-offload \ + -fa on -ctk q8_0 -ctv q8_0 -b 512 -ub 512 -c 32768 --parallel 1 \ + --host 127.0.0.1 --port "$PORT" --no-warmup > "$LOG" 2>&1 & + SRV=$! + for _ in $(seq 1 300); do + curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1 && break + sleep 1 + done + python3 "$HERE/r4-kv-pipeline-exact.py" "$PORT" "$NTOK" "$NPARA" || rc=$? + kill $SRV 2>/dev/null; wait $SRV 2>/dev/null + rm -f "$LOG" +done +exit $rc diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index 7ff367079cf1..a706268882ac 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -371,6 +371,30 @@ extern "C" { GGML_API void ggml_backend_sched_request_shared_buffer_shrink( ggml_backend_sched_t sched); + // Pipelined delivery of host-resident split inputs. + // + // Without it, a split that reads a host-resident input pays copy + compute in series: + // the transfer is issued on the consumer's own stream immediately before the kernels + // that read it. With it, the scheduler keeps a ring of `depth` staging slots outside + // the graph allocator's reach and issues the stable prefix of a later split's inputs on + // a separate transfer stream while the current split computes, so the transfer retires + // underneath the kernels. + // + // Only inputs that carry a stable prefix (ggml_set_stable_prefix) are eligible: without + // one, the scheduler cannot know that an earlier split of the same graph will not still + // write the bytes it would deliver ahead of time. Everything else keeps the ordered path. + // + // `depth` is how many splits ahead deliveries run; 0 disables pipelining. The ring holds a + // couple of slots more than that, so that recycling a slot never has to wait for a reader + // that is still running. Requires a destination backend with asynchronous transfers and + // events; where that is missing the setting is ignored. Costs roughly (depth + 2) * + // (largest staged split) of device memory. Must be called before the first graph is + // allocated. + GGML_API void ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, int depth); + + // Number of staged deliveries and staged bytes issued since the scheduler was created. + GGML_API void ggml_backend_sched_get_transport_pipeline_stats(ggml_backend_sched_t sched, int64_t * n_deliveries, int64_t * n_bytes_early, int64_t * n_bytes_late); + // Initialize backend buffers from a measure graph GGML_API void ggml_backend_sched_reserve_size(ggml_backend_sched_t sched, struct ggml_cgraph * measure_graph, size_t * sizes); GGML_API bool ggml_backend_sched_reserve(ggml_backend_sched_t sched, struct ggml_cgraph * measure_graph); // returns success diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index ce1a4d3d5997..213dcff2f37e 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -732,11 +732,27 @@ extern "C" { void * extra; // extra things e.g. for ggml-cuda.cu - char padding[8]; + // number of leading bytes of this tensor's storage that are guaranteed not to be + // written during a single graph evaluation. 0 means "not known". + // set on the tensor that owns the storage, by whoever knows what the graph will write; + // a view inherits the part of it that its own byte window covers. read by the backend + // scheduler, which may use it to deliver a host-resident split input to an accelerator + // before the split that reads it runs, see + // ggml_backend_sched_set_transport_pipeline_depth(). + // (kept last, in place of the former trailing padding, so that sizeof(struct ggml_tensor) + // does not change) + size_t stable_prefix; }; static const size_t GGML_TENSOR_SIZE = sizeof(struct ggml_tensor); + // declare that the first nbytes bytes of tensor->data cannot change while a graph that + // reads this tensor is being evaluated. nbytes is clamped to ggml_nbytes(tensor). + // set it on the tensor that owns the storage, not on a view of it, and keep it current: + // it must describe the graph that is about to run, including when that graph is reused. + GGML_API void ggml_set_stable_prefix(struct ggml_tensor * tensor, size_t nbytes); + GGML_API size_t ggml_get_stable_prefix(const struct ggml_tensor * tensor); + // Abort callback // If not NULL, called before ggml computation // If it returns true, the computation is aborted diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 7edce750b588..c9573df8e5b3 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -434,7 +434,10 @@ struct ggml_backend_meta_buffer_context { // FIXME // The size of the split state cache is unbounded and can theoretically grow infinitely large. // However, it is also expensive to build and clearing it on every rebuild in ggml_backend_meta_graph_compute is too expensive. - static constexpr size_t nbtc = GGML_TENSOR_SIZE - sizeof(ggml_tensor::padding); + // ggml_tensor::stable_prefix is a hint for the backend scheduler that this backend does + // not consume, and it changes from graph to graph, so keep it out of the compared image + // together with the trailing padding. + static constexpr size_t nbtc = offsetof(ggml_tensor, stable_prefix); std::map, std::pair> split_state_cache; int debug; diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 6b114b403a31..d6891d779da6 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -14,6 +14,7 @@ #include "ggml-impl.h" #include +#include #include #include #include @@ -786,6 +787,80 @@ static bool ggml_backend_sched_allows_bufferless_kvarn_src( #define GGML_SCHED_MAX_COPIES 4 #endif +#ifndef GGML_SCHED_MAX_TRANSPORT_SLOTS +#define GGML_SCHED_MAX_TRANSPORT_SLOTS 16 +#endif + +// How many slots the transport ring keeps behind the look-ahead. A delivery that runs L splits +// ahead recycles the slot of the split L - n_slots back, so with n_slots == L + 1 it would +// recycle the split that was enqueued a moment ago and is still running, and every delivery +// would have to wait for the consumer to catch up -- the ordered path with extra steps. Two +// slots of margin put the recycled reader far enough behind to have finished. +#ifndef GGML_SCHED_TRANSPORT_MARGIN +#define GGML_SCHED_TRANSPORT_MARGIN 2 +#endif + +// One staging slot of the transport ring. A slot is owned by the transfer stream while it is +// being filled and by the consumer stream while it is being read; the two events below are the +// handover in each direction. +struct ggml_backend_sched_transport_slot { + ggml_backend_event_t ready; // recorded on the transfer backend once the slot is fully delivered + ggml_backend_event_t release; // recorded on the consumer backend once the reader was enqueued + bool release_armed; // a reader was enqueued and has not been waited for yet +}; + +// Pipelined delivery of host-resident split inputs. +// +// The ordered path issues a split's host-to-device delivery on the consumer's own stream right +// before the kernels that read it, so a token costs copy + compute in series. This ring lets the +// stable part of a later split's delivery run on a separate transfer stream while the current +// split computes. The ring is allocated by the scheduler and never handed to ggml-alloc, which +// is what makes writing ahead safe: ggml-alloc is free to recycle a graph-owned input copy once +// its last graph-level consumer is done, and a look-ahead transfer is still in flight outside +// that lifetime. +struct ggml_backend_sched_transport { + int depth; // how many splits ahead deliveries run; 0 disables pipelining + int n_slots; // slots actually created: depth + GGML_SCHED_TRANSPORT_MARGIN + int backend_id; // the consumer backend this ring serves, -1 if none + + ggml_backend_t transfer; // second backend on the same device: owns the transfer stream + ggml_backend_buffer_t buffer; // the ring itself + size_t slot_size; + size_t alignment; + + struct ggml_backend_sched_transport_slot slots[GGML_SCHED_MAX_TRANSPORT_SLOTS]; + + // plan for the current graph, indexed by split id: the delivery order of the split, or -1 + // when the split stages nothing + int * split_order; + int plan_capacity; + int plan_n_splits; + int plan_n_inputs; + int n_staged; + int scan_cursor; + + // which inputs the plan put in the ring, flattened over splits. Membership is decided once, + // when the ring is laid out, and is what execution goes by: the amount that can be delivered + // early moves with every ubatch, but which input copies live in the ring must not. + unsigned char * input_staged; + int * split_input_ofs; // [plan_capacity + 1] + int input_capacity; + + int64_t n_deliveries; + int64_t n_bytes_early; + int64_t n_bytes_late; + + // debug >= 2: where the host's time in the split loop goes, and the values at the last report + int64_t t_issue_us; // issuing early deliveries + int64_t t_sync_us; // blocked in ggml_backend_synchronize / event_synchronize + int64_t t_copy_us; // blocked in the ordered ggml_backend_tensor_copy + int64_t t_graph_us; + int64_t n_graphs; + int64_t p_graph_us, p_sync_us, p_copy_us, p_issue_us, p_bytes_early, p_bytes_late; + + int debug; +}; + struct ggml_backend_sched_split { int backend_id; int i_start; @@ -846,6 +921,9 @@ struct ggml_backend_sched { bool op_offload; + // pipelined delivery of host-resident split inputs + struct ggml_backend_sched_transport transport; + int debug; // used for debugging graph reallocations [GGML_SCHED_DEBUG_REALLOC] @@ -1633,6 +1711,309 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra } } +// +// pipelined transport of host-resident split inputs +// + +static bool ggml_backend_sched_transport_enabled(ggml_backend_sched_t sched) { + return sched->transport.depth >= 1 && sched->transport.n_slots >= 2 && sched->transport.transfer != NULL; +} + +// The annotation lives on the tensor that owns the storage; a split input is normally a view of +// it. ggml keeps view_src pointing at the root tensor and view_offs absolute, so the byte window +// a delivery reads is [view_offs, view_offs + nbytes) of the root, and the stable part of it is +// whatever that window shares with the root's stable prefix. This holds whatever the view's +// shape and strides are, because the delivery is a flat copy of ggml_nbytes(input) bytes. +static size_t ggml_backend_sched_input_stable_prefix(const struct ggml_tensor * input) { + const struct ggml_tensor * base = input->view_src ? input->view_src : input; + if (base->stable_prefix == 0) { + return 0; + } + + const size_t offs = input->view_src ? input->view_offs : 0; + if (base->stable_prefix <= offs) { + return 0; + } + + const size_t avail = base->stable_prefix - offs; + const size_t bytes = ggml_nbytes(input); + + return avail < bytes ? avail : bytes; +} + +// Whether a split input belongs in the ring. Decided once per plan; from then on execution +// consults the recorded answer rather than asking again, because the write position this depends +// on moves between ubatches while the plan and the addresses it handed out do not. +static bool ggml_backend_sched_input_can_stage( + ggml_backend_sched_t sched, struct ggml_backend_sched_split * split, int input_id) { + if (split->backend_id != sched->transport.backend_id) { + return false; + } + + struct ggml_tensor * input = split->inputs[input_id]; + + // user inputs must be copied immediately, before the user can overwrite them + if (input->flags & GGML_TENSOR_FLAG_INPUT) { + return false; + } + + // an input with no stable prefix has nothing that could be delivered ahead of its split + if (ggml_backend_sched_input_stable_prefix(input) == 0) { + return false; + } + + ggml_backend_buffer_t buf = input->view_src ? input->view_src->buffer : input->buffer; + if (buf == NULL || !ggml_backend_buffer_is_host(buf)) { + return false; + } + + // weights take the used-experts path in the split loop, which delivers a subset of the bytes + if (ggml_backend_buffer_get_usage(buf) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { + return false; + } + + return tensor_copy(input, split->backend_id, sched->cur_copy) != NULL; +} + +static bool ggml_backend_sched_input_is_staged(ggml_backend_sched_t sched, int split_id, int input_id) { + const struct ggml_backend_sched_transport * tr = &sched->transport; + const int base = tr->split_input_ofs[split_id]; + return tr->input_staged[base + input_id] != 0; +} + +// sync_consumers must be false once the scheduler's backends may already be gone, which is the +// case on the teardown path: llama_context and other owners outlive the scheduler only by +// declaration order, and the backends it points at are not the scheduler's to keep alive. +static void ggml_backend_sched_transport_free_ring(ggml_backend_sched_t sched, bool sync_consumers) { + struct ggml_backend_sched_transport * tr = &sched->transport; + + if (tr->buffer == NULL) { + return; + } + + // nothing may still be reading from or writing into the ring + if (tr->transfer) { + ggml_backend_synchronize(tr->transfer); + } + if (sync_consumers) { + for (int i = 0; i < sched->n_backends; i++) { + ggml_backend_synchronize(sched->backends[i]); + } + } + + ggml_backend_buffer_free(tr->buffer); + tr->buffer = NULL; + tr->slot_size = 0; + + for (int i = 0; i < GGML_SCHED_MAX_TRANSPORT_SLOTS; i++) { + tr->slots[i].release_armed = false; + } +} + +// Lay the ring out over the current split list and point the staged input copies at it. Called +// before the graph is allocated: ggml-alloc leaves a tensor that already has data alone, so the +// staged copies are excluded from its reuse analysis instead of competing with it. +static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { + struct ggml_backend_sched_transport * tr = &sched->transport; + + tr->n_staged = 0; + tr->scan_cursor = 0; + + if (!ggml_backend_sched_transport_enabled(sched)) { + return; + } + + if (tr->plan_capacity < sched->n_splits) { + int * pnew = (int *) realloc(tr->split_order, sched->n_splits * sizeof(int)); + int * pofs = (int *) realloc(tr->split_input_ofs, (sched->n_splits + 1) * sizeof(int)); + if (pnew == NULL || pofs == NULL) { + GGML_LOG_WARN("%s: failed to allocate the transport plan, pipelining disabled for this graph\n", __func__); + tr->split_order = pnew ? pnew : tr->split_order; + tr->split_input_ofs = pofs ? pofs : tr->split_input_ofs; + return; + } + tr->split_order = pnew; + tr->split_input_ofs = pofs; + tr->plan_capacity = sched->n_splits; + } + + int n_inputs_total = 0; + for (int i = 0; i < sched->n_splits; i++) { + tr->split_input_ofs[i] = n_inputs_total; + n_inputs_total += sched->splits[i].n_inputs; + } + tr->split_input_ofs[sched->n_splits] = n_inputs_total; + + if (tr->input_capacity < n_inputs_total) { + unsigned char * pnew = (unsigned char *) realloc(tr->input_staged, n_inputs_total); + if (pnew == NULL) { + GGML_LOG_WARN("%s: failed to allocate the transport plan, pipelining disabled for this graph\n", __func__); + return; + } + tr->input_staged = pnew; + tr->input_capacity = n_inputs_total; + } + memset(tr->input_staged, 0, n_inputs_total); + tr->plan_n_splits = sched->n_splits; + tr->plan_n_inputs = n_inputs_total; + + size_t slot_size = 0; + for (int i = 0; i < sched->n_splits; i++) { + struct ggml_backend_sched_split * split = &sched->splits[i]; + + tr->split_order[i] = -1; + + size_t need = 0; + for (int j = 0; j < split->n_inputs; j++) { + if (!ggml_backend_sched_input_can_stage(sched, split, j)) { + continue; + } + tr->input_staged[tr->split_input_ofs[i] + j] = 1; + need += GGML_PAD(ggml_nbytes(split->inputs[j]), tr->alignment); + } + + if (need == 0) { + continue; + } + + tr->split_order[i] = tr->n_staged++; + slot_size = std::max(slot_size, need); + } + + if (tr->n_staged == 0) { + return; + } + + if (tr->buffer == NULL || tr->slot_size < slot_size) { + ggml_backend_sched_transport_free_ring(sched, true); + + ggml_backend_buffer_type_t buft = sched->bufts[tr->backend_id]; + ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(buft, slot_size * tr->n_slots); + if (buffer == NULL) { + GGML_LOG_WARN("%s: failed to allocate %zu MiB for the transport ring, pipelining disabled\n", + __func__, (slot_size * tr->n_slots) >> 20); + tr->n_staged = 0; + return; + } + ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_COMPUTE); + + tr->buffer = buffer; + tr->slot_size = slot_size; + + if (tr->debug > 0) { + GGML_LOG_INFO("%s: transport ring: %d slots x %zu KiB on %s\n", __func__, + tr->n_slots, slot_size >> 10, ggml_backend_buft_name(buft)); + } + } + + if (tr->debug > 0) { + size_t total = 0; + size_t early = 0; + const char * src_buft = "?"; + for (int i = 0; i < sched->n_splits; i++) { + struct ggml_backend_sched_split * split = &sched->splits[i]; + for (int j = 0; j < split->n_inputs; j++) { + if (!ggml_backend_sched_input_is_staged(sched, i, j)) { + continue; + } + struct ggml_tensor * input = split->inputs[j]; + ggml_backend_buffer_t buf = input->view_src ? input->view_src->buffer : input->buffer; + src_buft = ggml_backend_buft_name(buf->buft); + total += ggml_nbytes(input); + early += ggml_backend_sched_input_stable_prefix(input); + } + } + GGML_LOG_INFO("%s: %d/%d splits staged, %zu KiB per graph, %zu KiB of it early, source %s\n", + __func__, tr->n_staged, sched->n_splits, total >> 10, early >> 10, src_buft); + } + + char * const ring = (char *) ggml_backend_buffer_get_base(tr->buffer); + + for (int i = 0; i < sched->n_splits; i++) { + if (tr->split_order[i] < 0) { + continue; + } + + struct ggml_backend_sched_split * split = &sched->splits[i]; + char * slot = ring + (size_t)(tr->split_order[i] % tr->n_slots) * tr->slot_size; + + size_t offset = 0; + for (int j = 0; j < split->n_inputs; j++) { + if (!ggml_backend_sched_input_is_staged(sched, i, j)) { + continue; + } + struct ggml_tensor * input_cpy = tensor_copy(split->inputs[j], split->backend_id, sched->cur_copy); + input_cpy->data = slot + offset; + input_cpy->buffer = tr->buffer; + offset += GGML_PAD(ggml_nbytes(split->inputs[j]), tr->alignment); + } + GGML_ASSERT(offset <= tr->slot_size); + } +} + +// Issue the stable prefix of every staged split within the look-ahead of what has already been +// enqueued. The ring carries GGML_SCHED_TRANSPORT_MARGIN slots more than the look-ahead, so the +// slot a delivery writes into belongs to a split that is several readers behind the one just +// enqueued, and recycling it does not put the transfer stream back in lock-step with the +// consumer. +static void ggml_backend_sched_transport_prefetch(ggml_backend_sched_t sched, int consumed) { + struct ggml_backend_sched_transport * tr = &sched->transport; + + const int lookahead = tr->depth; + + for (int i = tr->scan_cursor; i < sched->n_splits; i++) { + if (tr->split_order[i] < 0) { + continue; + } + if (tr->split_order[i] > consumed + lookahead) { + return; + } + + struct ggml_backend_sched_split * split = &sched->splits[i]; + struct ggml_backend_sched_transport_slot * slot = &tr->slots[tr->split_order[i] % tr->n_slots]; + + // the previous occupant of this slot must be read before the slot is overwritten. This is + // ordered stream to stream rather than through the host: blocking the host here would + // hold back the work it has not enqueued yet, which is what the margin exists to avoid. + if (slot->release_armed) { + ggml_backend_event_wait(tr->transfer, slot->release); + slot->release_armed = false; + } + + for (int j = 0; j < split->n_inputs; j++) { + if (!ggml_backend_sched_input_is_staged(sched, i, j)) { + continue; + } + struct ggml_tensor * input = split->inputs[j]; + + // how much of this input is stable is a property of the ubatch about to run, not of + // the plan: it can be less than when the ring was laid out, and then only the + // remainder moves and the rest waits for the split, exactly as before + const size_t prefix = ggml_backend_sched_input_stable_prefix(input); + if (prefix == 0) { + continue; + } + + struct ggml_tensor * input_cpy = tensor_copy(input, split->backend_id, sched->cur_copy); + GGML_ASSERT(input->data != NULL && input_cpy->data != NULL); + const int64_t t0 = tr->debug >= 2 ? ggml_time_us() : 0; + ggml_backend_tensor_set_async(tr->transfer, input_cpy, input->data, 0, prefix); + if (tr->debug >= 2) { + tr->t_issue_us += ggml_time_us() - t0; + } + tr->n_bytes_early += prefix; + } + + // record the handover here rather than when the split runs: the transfer stream is FIFO, + // and by then the deliveries for the splits after this one are already queued behind it. + // Waiting on an event recorded after those would make the consumer wait for the whole + // look-ahead, which is the ordered path again with extra steps. + ggml_backend_event_record(slot->ready, tr->transfer); + + tr->scan_cursor = i + 1; + } +} + static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { bool backend_ids_changed = false; for (int i = 0; i < sched->graph.n_nodes; i++) { @@ -1652,6 +2033,10 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { } } + // lay out the transport ring and point the staged input copies at it before the graph is + // allocated, so ggml-alloc sees those copies as already allocated and leaves them alone + ggml_backend_sched_transport_plan(sched); + // allocate graph if (backend_ids_changed || !ggml_gallocr_alloc_graph(sched->galloc, &sched->graph)) { #ifndef NDEBUG @@ -1671,6 +2056,9 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { // the re-allocation may cause the split inputs to be moved to a different address // synchronize without ggml_backend_sched_synchronize to avoid changing cur_copy + if (sched->transport.transfer) { + ggml_backend_synchronize(sched->transport.transfer); + } for (int i = 0; i < sched->n_backends; i++) { ggml_backend_synchronize(sched->backends[i]); } @@ -1693,6 +2081,31 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s std::vector ids; std::vector used_ids; + struct ggml_backend_sched_transport * tr = &sched->transport; + // a reused graph keeps the plan that was made for it, so the split list it describes must be + // the one about to run + int n_inputs_now = 0; + for (int i = 0; i < sched->n_splits; i++) { + n_inputs_now += splits[i].n_inputs; + } + const bool staged = tr->n_staged > 0 && tr->plan_n_splits == sched->n_splits && + tr->plan_n_inputs == n_inputs_now; + int consumed = 0; // staged splits whose reader has been enqueued + + // the plan outlives the graph it was made for -- a reused graph keeps it -- so the cursor + // that walks it has to start over on every evaluation, not only when the plan is rebuilt + tr->scan_cursor = 0; + + // prime the ring before the first consumer runs; from here on deliveries are issued only + // after a split has been enqueued, never before, so that recycling a slot can never hold + // back work the consumer could already be running + if (staged) { + ggml_backend_sched_transport_prefetch(sched, consumed); + } + + const bool timing = tr->debug >= 2; + const int64_t t_graph_start = timing ? ggml_time_us() : 0; + for (int split_id = 0; split_id < sched->n_splits; split_id++) { struct ggml_backend_sched_split * split = &splits[split_id]; int split_backend_id = split->backend_id; @@ -1704,6 +2117,23 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s struct ggml_tensor * input = split->inputs[input_id]; struct ggml_tensor * input_cpy = tensor_copy(input, split_backend_id, sched->cur_copy); + if (staged && ggml_backend_sched_input_is_staged(sched, split_id, input_id)) { + // whatever prefix was stable went out on the transfer stream earlier; the rest is + // what an earlier split of this graph may still have written, and it is only safe + // to read now that every earlier split has run. It goes on the consumer's own + // stream, where it is already ordered ahead of the kernels and behind the reader + // of whatever occupied this slot before, and where it cannot push the transfer + // stream's look-ahead behind it. + const size_t prefix = ggml_backend_sched_input_stable_prefix(input); + const size_t nbytes = ggml_nbytes(input); + if (nbytes > prefix) { + ggml_backend_tensor_set_async(split_backend, input_cpy, + (const char *) input->data + prefix, prefix, nbytes - prefix); + tr->n_bytes_late += nbytes - prefix; + } + continue; + } + if (input->flags & GGML_TENSOR_FLAG_INPUT) { // inputs from the user must be copied immediately to prevent the user overwriting the data before the copy is done if (sched->events[split_backend_id][sched->cur_copy] != NULL) { @@ -1714,11 +2144,15 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s ggml_backend_tensor_copy(input, input_cpy); } else { // wait for the split backend to finish using the input before overwriting it + const int64_t t_sync0 = timing ? ggml_time_us() : 0; if (sched->events[split_backend_id][sched->cur_copy] != NULL) { ggml_backend_event_wait(split_backend, sched->events[split_backend_id][sched->cur_copy]); } else { ggml_backend_synchronize(split_backend); } + if (timing) { + tr->t_sync_us += ggml_time_us() - t_sync0; + } // when offloading MoE weights, we can reduce the amount of data copied by copying only the experts that are used ggml_tensor * node = split->graph.nodes[0]; @@ -1810,18 +2244,31 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s // TODO: add public function to facilitate this, since applications do not have direct access to the backend interface const bool async_ok = split_backend->iface.cpy_tensor_async && split_backend->iface.cpy_tensor_async(input_backend, split_backend, input, input_cpy); if (!async_ok) { + const int64_t t_sync1 = timing ? ggml_time_us() : 0; ggml_backend_synchronize(input_backend); if (sched->events[split_backend_id][sched->cur_copy] != NULL) { ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]); } else { ggml_backend_synchronize(split_backend); } + const int64_t t_copy0 = timing ? ggml_time_us() : 0; ggml_backend_tensor_copy(input, input_cpy); + if (timing) { + tr->t_sync_us += t_copy0 - t_sync1; + tr->t_copy_us += ggml_time_us() - t_copy0; + } } } } } + // order the consumer behind this split's early deliveries + struct ggml_backend_sched_transport_slot * slot = NULL; + if (staged && tr->split_order[split_id] >= 0) { + slot = &tr->slots[tr->split_order[split_id] % tr->n_slots]; + ggml_backend_event_wait(split_backend, slot->ready); + } + if (!sched->callback_eval) { enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &split->graph); if (ec != GGML_STATUS_SUCCESS) { @@ -1863,6 +2310,20 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } } + // every kernel that reads this split's slot is enqueued, so the slot may be refilled + // once the consumer stream reaches this point + if (slot != NULL) { + ggml_backend_event_record(slot->release, split_backend); + slot->release_armed = true; + consumed++; + tr->n_deliveries++; + + // run the transfer stream ahead of the consumer: with this split's kernels already + // enqueued, the deliveries for the next few staged splits can go out even if + // recycling their slot has to wait for a reader that is now running + ggml_backend_sched_transport_prefetch(sched, consumed); + } + // record the event of this copy if (split->n_inputs > 0) { if (sched->events[split_backend_id][sched->cur_copy] != NULL) { @@ -1871,6 +2332,28 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } } + if (timing) { + tr->t_graph_us += ggml_time_us() - t_graph_start; + if (++tr->n_graphs % 64 == 0) { + const double n = 64.0; + const double d_graph = (tr->t_graph_us - tr->p_graph_us)/1e3/n; + const double d_sync = (tr->t_sync_us - tr->p_sync_us) /1e3/n; + const double d_copy = (tr->t_copy_us - tr->p_copy_us) /1e3/n; + const double d_issue = (tr->t_issue_us - tr->p_issue_us)/1e3/n; + GGML_LOG_INFO("%s: %" PRId64 " graphs, last 64: %.2f ms/graph = sync %.2f + ordered copy %.2f + " + "issue %.2f + rest %.2f; %.1f MiB early, %.1f MiB late\n", __func__, tr->n_graphs, + d_graph, d_sync, d_copy, d_issue, d_graph - d_sync - d_copy - d_issue, + (tr->n_bytes_early - tr->p_bytes_early)/1048576.0/n, + (tr->n_bytes_late - tr->p_bytes_late) /1048576.0/n); + tr->p_graph_us = tr->t_graph_us; + tr->p_sync_us = tr->t_sync_us; + tr->p_copy_us = tr->t_copy_us; + tr->p_issue_us = tr->t_issue_us; + tr->p_bytes_early = tr->n_bytes_early; + tr->p_bytes_late = tr->n_bytes_late; + } + } + return GGML_STATUS_SUCCESS; } @@ -1941,15 +2424,148 @@ ggml_backend_sched_t ggml_backend_sched_new( sched->galloc = ggml_gallocr_new_n(sched->bufts, n_backends); sched->op_offload = op_offload; + sched->transport.backend_id = -1; + { + const char * GGML_SCHED_TRANSPORT_DEBUG = getenv("GGML_SCHED_TRANSPORT_DEBUG"); + sched->transport.debug = GGML_SCHED_TRANSPORT_DEBUG ? atoi(GGML_SCHED_TRANSPORT_DEBUG) : 0; + } + ggml_backend_sched_reset(sched); return sched; } +static void ggml_backend_sched_transport_teardown(ggml_backend_sched_t sched) { + struct ggml_backend_sched_transport * tr = &sched->transport; + + ggml_backend_sched_transport_free_ring(sched, false); + + for (int i = 0; i < GGML_SCHED_MAX_TRANSPORT_SLOTS; i++) { + ggml_backend_event_free(tr->slots[i].ready); + ggml_backend_event_free(tr->slots[i].release); + tr->slots[i].ready = NULL; + tr->slots[i].release = NULL; + tr->slots[i].release_armed = false; + } + + if (tr->transfer) { + ggml_backend_free(tr->transfer); + tr->transfer = NULL; + } + + tr->n_slots = 0; + tr->backend_id = -1; + tr->n_staged = 0; +} + +void ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, int depth) { + GGML_ASSERT(sched); + + const char * env = getenv("GGML_KV_PIPELINE_DEPTH"); + if (env != NULL) { + depth = atoi(env); + } + + depth = std::min(depth, GGML_SCHED_MAX_TRANSPORT_SLOTS - GGML_SCHED_TRANSPORT_MARGIN); + if (depth < 0) { + depth = 0; + } + + struct ggml_backend_sched_transport * tr = &sched->transport; + if (tr->depth == depth && (depth < 1 || tr->transfer != NULL)) { + return; + } + + ggml_backend_sched_transport_teardown(sched); + tr->depth = depth; + + if (depth < 1) { + return; + } + + const int n_slots = depth + GGML_SCHED_TRANSPORT_MARGIN; + + // one ring, on the first backend that can transfer asynchronously and order streams with + // events. Additional accelerators keep the ordered path. + for (int i = 0; i < sched->n_backends; i++) { + ggml_backend_t backend = sched->backends[i]; + if (backend->iface.set_tensor_async == NULL || + backend->iface.event_record == NULL || + backend->iface.event_wait == NULL) { + continue; + } + + ggml_backend_dev_t dev = ggml_backend_get_device(backend); + if (dev == NULL || dev->iface.event_new == NULL) { + continue; + } + if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) { + continue; + } + + // the ring is written through the transfer backend, which only accepts the device's own + // default buffer type; a scheduler configured with anything else keeps the ordered path + if (sched->bufts[i] != ggml_backend_dev_buffer_type(dev)) { + continue; + } + + // a second backend on the same device gives the deliveries a stream of their own + ggml_backend_t transfer = ggml_backend_dev_init(dev, NULL); + if (transfer == NULL) { + continue; + } + + bool ok = true; + for (int slot = 0; slot < n_slots && ok; slot++) { + tr->slots[slot].ready = ggml_backend_event_new(dev); + tr->slots[slot].release = ggml_backend_event_new(dev); + ok = tr->slots[slot].ready != NULL && tr->slots[slot].release != NULL; + } + + if (!ok) { + for (int slot = 0; slot < n_slots; slot++) { + ggml_backend_event_free(tr->slots[slot].ready); + ggml_backend_event_free(tr->slots[slot].release); + tr->slots[slot].ready = NULL; + tr->slots[slot].release = NULL; + } + ggml_backend_free(transfer); + continue; + } + + tr->transfer = transfer; + tr->backend_id = i; + tr->n_slots = n_slots; + tr->alignment = std::max(ggml_backend_buft_get_alignment(sched->bufts[i]), 128); + + if (tr->debug > 0) { + GGML_LOG_INFO("%s: pipelined host transport enabled on %s, %d splits ahead, %d slots\n", + __func__, ggml_backend_name(backend), depth, n_slots); + } + return; + } + + if (tr->debug > 0) { + GGML_LOG_INFO("%s: no backend supports pipelined host transport, staying on the ordered path\n", __func__); + } +} + +void ggml_backend_sched_get_transport_pipeline_stats( + ggml_backend_sched_t sched, int64_t * n_deliveries, int64_t * n_bytes_early, int64_t * n_bytes_late) { + GGML_ASSERT(sched); + if (n_deliveries) { *n_deliveries = sched->transport.n_deliveries; } + if (n_bytes_early) { *n_bytes_early = sched->transport.n_bytes_early; } + if (n_bytes_late) { *n_bytes_late = sched->transport.n_bytes_late; } +} + void ggml_backend_sched_free(ggml_backend_sched_t sched) { if (sched == NULL) { return; } + ggml_backend_sched_transport_teardown(sched); + free(sched->transport.split_order); + free(sched->transport.split_input_ofs); + free(sched->transport.input_staged); for (int b = 0; b < sched->n_backends; b++) { for (int c = 0; c < sched->n_copies; c++) { ggml_backend_event_free(sched->events[b][c]); @@ -2086,6 +2702,9 @@ enum ggml_status ggml_backend_sched_graph_compute_async(ggml_backend_sched_t sch void ggml_backend_sched_synchronize(ggml_backend_sched_t sched) { GGML_ASSERT(sched); + if (sched->transport.transfer) { + ggml_backend_synchronize(sched->transport.transfer); + } for (int i = 0; i < sched->n_backends; i++) { ggml_backend_synchronize(sched->backends[i]); } diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index bf7c3dca0951..2cbd6217f80d 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1379,6 +1379,17 @@ size_t ggml_nbytes_pad(const struct ggml_tensor * tensor) { return GGML_PAD(ggml_nbytes(tensor), GGML_MEM_ALIGN); } +void ggml_set_stable_prefix(struct ggml_tensor * tensor, size_t nbytes) { + GGML_ASSERT(tensor); + const size_t total = ggml_nbytes(tensor); + tensor->stable_prefix = nbytes < total ? nbytes : total; +} + +size_t ggml_get_stable_prefix(const struct ggml_tensor * tensor) { + GGML_ASSERT(tensor); + return tensor->stable_prefix; +} + int64_t ggml_blck_size(enum ggml_type type) { assert(type >= 0); assert(type < GGML_TYPE_COUNT); @@ -1875,7 +1886,7 @@ static struct ggml_tensor * ggml_new_tensor_impl( /*.data =*/ obj_alloc_size > 0 ? (void *)(result + 1) : data, /*.name =*/ { 0 }, /*.extra =*/ NULL, - /*.padding =*/ { 0 }, + /*.stable_prefix=*/ 0, }; // TODO: this should not be needed as long as we don't rely on aligned SIMD loads diff --git a/include/llama.h b/include/llama.h index 2ffe20d15b5b..dcbae788e177 100644 --- a/include/llama.h +++ b/include/llama.h @@ -512,6 +512,11 @@ extern "C" { // Fit/no-alloc contexts still measure the full prompt reservation. bool live_context_workspace; // grow supported attention workspace plans with the padded live physical KV extent. // false preserves the full-context reservation and initialization order. + uint32_t kv_pipeline_depth; // how many splits ahead the scheduler delivers a host-resident KV cache to the + // accelerator, so that the transfer runs while the previous split computes. + // 0 keeps the ordered path, where a decode token pays the transfer and the + // attention kernels in series. Costs (kv_pipeline_depth + 2) * (largest staged + // split) of device memory. }; struct llama_model_tensor_override { diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 8119012cc273..c51ba5fe9ed6 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -302,6 +302,7 @@ llama_context::llama_context( cparams.recurrent_state_offload = params.recurrent_state_offload; cparams.phase_aware_workspace = params.phase_aware_workspace; cparams.live_context_workspace = params.live_context_workspace; + cparams.kv_pipeline_depth = params.kv_pipeline_depth; cparams.kv_gpu_layers = params.kv_gpu_layers; cparams.offload_attn_compute = params.offload_kqv || (params.op_offload && (params.kv_cpu_pinned || params.kv_gpu_layers > 0)); @@ -1118,6 +1119,9 @@ void llama_context::sched_reserve(uint32_t n_tokens_req, uint32_t n_kv_req) { sched.reset(ggml_backend_sched_new( backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, pipeline_parallel, cparams.op_offload)); + // only a host-resident KV cache produces the deliveries this pipelines + ggml_backend_sched_set_transport_pipeline_depth(sched.get(), + cparams.kv_cpu_pinned || !cparams.offload_kqv ? (int) cparams.kv_pipeline_depth : 0); if (sched_shared_buffers) { ggml_backend_sched_set_shared_buffers(sched.get(), sched_shared_buffers.get()); sched_shared_generation = @@ -4559,6 +4563,7 @@ llama_context_params llama_context_default_params() { /*.kv_gpu_layers =*/ 0, /*.phase_aware_workspace =*/ false, /*.live_context_workspace =*/ false, + /*.kv_pipeline_depth =*/ 1, }; return result; diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 3ae2c32085fb..c3ece1e2b0ed 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -64,6 +64,7 @@ struct llama_cparams { bool recurrent_state_offload; bool phase_aware_workspace; bool live_context_workspace; + uint32_t kv_pipeline_depth; std::vector embeddings_layer_inp; // [n_layer()] extract input embeddings for layer diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 97007f30ae88..8e44499727c4 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -2803,6 +2803,10 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & return; } + // before the graph is built and allocated, so that the scheduler's delivery plan and the + // deliveries it then issues are decided against the same write position + update_stable_prefixes(sinfo); + // keep track of the max sequence position that we would overwrite with this ubatch // for non-SWA cache, this would be always empty llama_seq_id seq_pos_max_rm[LLAMA_MAX_SEQ]; @@ -3552,7 +3556,48 @@ ggml_tensor * llama_kv_cache::build_input_v_rot(ggml_context * ctx) const { return res; } +void llama_kv_cache::update_stable_prefixes(const slot_info & sinfo) const { + // Rows written by this ubatch, in the flattened [n_embd_gqa, kv_size*n_stream] body: every + // byte below the lowest of them keeps whatever the previous ubatch left there for the whole + // graph, so a delivery of that region may be issued before the split that reads it. + uint64_t min_row = UINT64_MAX; + for (uint32_t s = 0; s < sinfo.n_stream(); ++s) { + const uint64_t offs = (uint64_t) sinfo.strm[s]*get_size(); + for (const uint32_t idx : sinfo.idxs[s]) { + min_row = std::min(min_row, offs + idx); + } + } + + if (min_row == UINT64_MAX) { + clear_stable_prefixes(); + return; + } + + for (const auto & layer : layers) { + if (layer.k) { + ggml_set_stable_prefix(layer.k, min_row*layer.k->nb[1]); + } + if (layer.v) { + // the transposed V cache scatters each ubatch across the whole tensor, so there is + // no leading region that this ubatch leaves alone + ggml_set_stable_prefix(layer.v, v_trans ? 0 : min_row*layer.v->nb[1]); + } + } +} + +void llama_kv_cache::clear_stable_prefixes() const { + for (const auto & layer : layers) { + if (layer.k) { + ggml_set_stable_prefix(layer.k, 0); + } + if (layer.v) { + ggml_set_stable_prefix(layer.v, 0); + } + } +} + void llama_kv_cache::set_input_k_idxs(ggml_tensor * dst, const llama_ubatch * ubatch, const slot_info & sinfo) const { + const uint32_t n_tokens = ubatch->n_tokens; GGML_ASSERT(n_tokens == (int64_t) sinfo.size()*sinfo.n_stream()); @@ -4117,6 +4162,9 @@ ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_co // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] GGML_ASSERT(!other); + // this graph rewrites the whole body in place, so nothing in it may be delivered early + clear_stable_prefixes(); + auto * ctx = res->get_ctx(); auto * gf = res->get_gf(); diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 892e7c10542e..2ca262af3dd5 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -328,6 +328,13 @@ class llama_kv_cache : public llama_memory_i { ggml_tensor * build_input_k_rot(ggml_context * ctx) const; ggml_tensor * build_input_v_rot(ggml_context * ctx) const; + // Tell the backend scheduler which part of each layer's persistent K/V storage this ubatch + // will not write, so a host-resident cache can be delivered to the accelerator ahead of the + // attention that reads it. Must be refreshed for every ubatch, including when the graph is + // reused, because the write position moves while the graph does not. + void update_stable_prefixes(const slot_info & sinfo) const; + void clear_stable_prefixes() const; + void set_input_k_idxs(ggml_tensor * dst, const llama_ubatch * ubatch, const slot_info & sinfo) const; void set_input_v_idxs(ggml_tensor * dst, const llama_ubatch * ubatch, const slot_info & sinfo) const; void set_input_tail_idxs(ggml_tensor * dst, const llama_ubatch * ubatch) const; From 16225f1cf4baf4482720b14b6161b73076f14555 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Tue, 25 Aug 2026 19:33:22 +0200 Subject: [PATCH 2/6] docs: document --kv-pipeline-depth in the fork's argument and development guides Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01SRa5ZTc1bpnk1fE3mrDcSh --- CLAUDE.md | 10 ++++++++++ docs/beellama-args.md | 1 + 2 files changed, 11 insertions(+) diff --git a/CLAUDE.md b/CLAUDE.md index 51831408b3d6..c8a879fa7b49 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -8,6 +8,7 @@ keeps a small set of maintained extensions: `q6_1`). - Quantized-native CUDA MMA FlashAttention (`--flash-attn-native-quants`). - Profit-only adaptive draft depth for upstream `draft-dflash` speculation. +- Pipelined delivery of a host-resident KV cache (`--kv-pipeline-depth`). - Server-side reasoning-loop protection. - KLD save/load support in `llama-perplexity` for KVarN validation. @@ -73,6 +74,15 @@ pairs; combining it with `GGML_CUDA_FA_ALL_QUANTS=ON` selects all 36. Every vali KVarN bit pair remains available through descriptor-native MMA when it is outside the fast matrix. Existing CMake caches retain earlier selections. +`--kv-pipeline-depth N` issues the host-to-device delivery of a host-resident KV +cache N splits ahead of the split that reads it, on a transfer stream of its own, +so a decode token stops paying transfer and attention in series. It is on by +default at `N = 1`, `0` restores the ordered path exactly, and it only engages +where a host-resident cache produces the deliveries (`--no-kv-offload`, with or +without `--kv-cpu-pinned`). Output is byte-identical either way. See +[`docs/kv-transport-pipelining.md`](docs/kv-transport-pipelining.md) for the +design, the numbers, the limits, and the two reproduction scripts. + ## Layout - `src/llama-kvarn.{h,cpp}`: KVarN descriptors, layouts, type parsing, and diff --git a/docs/beellama-args.md b/docs/beellama-args.md index 2ba91e02373d..9bd23dab7f2c 100644 --- a/docs/beellama-args.md +++ b/docs/beellama-args.md @@ -202,6 +202,7 @@ context-dependent H2D traffic; it is not equivalent to CPU attention. | `--recurrent-state-offload`, `--no-recurrent-state-offload` | `LLAMA_ARG_RECURRENT_STATE_OFFLOAD` | Disabled | For hybrid models under `--no-kv-offload`, keeps the fixed recurrent R/S state on the GPU while long attention KV remains in host memory. With ordinary KV offload enabled, recurrent state is already device-resident. | | `--kv-gpu-layers N` | `LLAMA_ARG_KV_GPU_LAYERS` | `0` | Under `--no-kv-offload`, retains the first `N` attention-KV layers on the device. This trades VRAM for lower repeated host-to-device traffic; a value at or above the attention-layer count matches full KV offload for those layers. | | `--spec-draft-kv-gpu-layers N`, `--kv-gpu-layers-draft N` | `LLAMA_ARG_SPEC_DRAFT_KV_GPU_LAYERS` | Inherit target | Overrides the target KV policy for a separate draft context and keeps the first `N` independently owned draft attention-KV layers device-resident. Shared KV layers follow their owner. The INI key is `spec-draft-kv-gpu-layers`. | +| `--kv-pipeline-depth N` | `LLAMA_ARG_KV_PIPELINE_DEPTH` | `1` | How many splits ahead the scheduler delivers a host-resident KV cache to the accelerator, so that the transfer runs while the previous split computes instead of in series with it. `0` restores the ordered path. Only engages where a host-resident cache produces the deliveries, and costs `(N + 2) x (largest staged split)` of device memory. See [`kv-transport-pipelining.md`](kv-transport-pipelining.md). | | `--op-offload`, `--no-op-offload` | — | Enabled | Controls whether operations over host tensors may execute on the accelerator. The negated form selects the separately tested CPU execution route; it is not the serving-performance path measured below. | Target and draft attention-KV placement is independent of model-weight From 52b8d68807003256e38357ae1e2c358b27a11bd3 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Tue, 25 Aug 2026 20:34:38 +0200 Subject: [PATCH 3/6] sched: bound the transport ring, and keep it away from copies with a later reader Two things the first pass did not cover, both found while measuring the feature from 4k to 256k of context. A slot holds one split's whole delivery, so the ring is linear in context length: 27 MiB at 4k, 213 MiB at 32k, 855 MiB at 128k, 1.7 GiB at 256k. It is allocated after the graph allocator has reserved its buffers, so taking whatever is left can starve a later reallocation. It now declines unless it can leave 512 MiB of device memory free, says so once, and stays on the ordered path. The scheduler also creates one input copy per (tensor, backend) rather than per split, so a later split can be pointed at the same copy without appearing in its input list -- and by then the ring may have recycled the slot. The plan now scans the splits after the owner for such a reader and puts those inputs back on the ordered path. Attention does not produce this shape, but nothing in the scheduler forbids it, and the failure would be silently wrong output. Measured on the same RTX 4070 / Qwen3.8-27B-IQ2_M / -nkvo --kv-cpu-pinned setup, A/B/A/B with peak device memory sampled across each arm: context ordered pipelined gain peak VRAM ring 4,096 31.66 37.01 +16.9% 10169 -> 10197 27 MiB 16,384 19.64 31.43 +60.1% 10159 -> 10263 107 MiB 32,768 12.99 15.49 +19.2% 10161 -> 10367 213 MiB 65,536 7.74 8.74 +12.9% 10163 -> 10573 428 MiB 131,072 4.29 4.68 +9.1% 10537 -> 11355 855 MiB The gain narrows with depth because compute is a shrinking share of the token, so there is less to hide the copy behind -- arithmetic, not an implementation limit. What the table also shows is that the ring's cost does not narrow with it, and at 128k it already claims 818 MiB of an 11.9 GiB card. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01SRa5ZTc1bpnk1fE3mrDcSh --- docs/kv-transport-pipelining.md | 24 ++++++- docs/repro/r4-kv-pipeline-context-sweep.sh | 42 +++++++++++ docs/repro/r4-kv-pipeline-exact.py | 80 +++++++++++++-------- docs/repro/r4-kv-pipeline-exact.sh | 20 +++--- ggml/src/ggml-backend.cpp | 82 ++++++++++++++++++++-- 5 files changed, 205 insertions(+), 43 deletions(-) create mode 100755 docs/repro/r4-kv-pipeline-context-sweep.sh diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index c145359db455..24bf2050aa77 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -76,6 +76,21 @@ goes by the recorded answer. How much of a staged input can go early moves with every ubatch; *which* input copies live in the ring must not, because their addresses were handed out at allocation time. +Two things disqualify an input that otherwise looks eligible: + +- **A reader further down the graph.** The scheduler creates one input copy per + (tensor, backend), not per split, so a later split can be pointed at the same + copy without appearing to consume it -- and by then the ring may have recycled + the slot. The plan scans the splits after the owner for such a reader and puts + those inputs back on the ordered path. Attention does not produce this shape, + but nothing in the scheduler forbids it. +- **No room on the device.** A slot holds one split's whole delivery, so the ring + grows with the context: 27 MiB at 4k, 213 MiB at 32k, 1.7 GiB at 256k. The ring + is allocated after the graph allocator has reserved its buffers, so it must not + take the room those buffers may still have to grow into; it declines unless it + can leave `GGML_SCHED_TRANSPORT_HEADROOM` (512 MiB) free, says so once, and + stays on the ordered path. + ### 3. A look-ahead that stays clear of the ring's tail A delivery running `L` splits ahead recycles the slot of the split `L - n_slots` @@ -173,8 +188,13 @@ transport never enabled because the scheduler is given a depth of 0. ## Scope and limits - Only inputs carrying a stable prefix are eligible. Everything else -- weights, - user inputs, a transposed V cache, any backend that cannot transfer - asynchronously or record events -- keeps the ordered path untouched. + user inputs, a transposed V cache, an input copy with a reader in a later + split, any backend that cannot transfer asynchronously or record events -- + keeps the ordered path untouched. +- The ring costs `(depth + 2) x (largest staged split)` of device memory, and a + staged split is both K and V of one attention layer over the whole context. That + is linear in context length, and it is what bounds the feature at depth rather + than anything about the transfer itself. - One ring, on the first backend that qualifies. Additional accelerators keep the ordered path; multi-GPU host-KV is not covered. - The scheduler must be configured with the device's own default buffer type. A diff --git a/docs/repro/r4-kv-pipeline-context-sweep.sh b/docs/repro/r4-kv-pipeline-context-sweep.sh new file mode 100755 index 000000000000..b8a2b91bf72b --- /dev/null +++ b/docs/repro/r4-kv-pipeline-context-sweep.sh @@ -0,0 +1,42 @@ +#!/bin/bash +# R4 across context depth: throughput and device allocation high-water, ordered against +# pipelined, on the same binary. The ring holds one split's whole delivery per slot, so its +# cost grows with the context; this is what measures where that stops being affordable. +# +# LLAMA_KV_MODEL=/path/model.gguf docs/repro/r4-kv-pipeline-context-sweep.sh [depth ...] +set -u +MODEL="${LLAMA_KV_MODEL:?set LLAMA_KV_MODEL to a .gguf path}" +BUILD="${LLAMA_KV_BUILD:-build}" +PIN="${LLAMA_KV_TASKSET:-0,2,4}" +NGEN="${LLAMA_KV_NGEN:-64}" +LOCK=/tmp/beellama-single-gpu.lock + +arm () { # $1 pipeline depth, $2 context depth, $3 reps + local vram; vram=$(mktemp) + ( while true; do nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits; sleep 0.25; done ) > "$vram" 2>/dev/null & + local sampler=$! + local ts + ts=$(GGML_KV_PIPELINE_DEPTH=$1 taskset -c "$PIN" "$BUILD/bin/llama-bench" -m "$MODEL" \ + -ngl 99 -sm none -mg 0 -t 3 -nkvo 1 --kv-cpu-pinned --recurrent-state-offload \ + -fa on -ctk q8_0 -ctv q8_0 -b 512 -ub 512 --no-warmup -p 0 -n "$NGEN" -d "$2" -r "$3" -o json \ + 2>/dev/null \ + | python3 -c "import json,sys +try: + d=json.load(sys.stdin); print('%.4f'%d[0]['avg_ts']) +except Exception: + print('FAILED')") + kill $sampler 2>/dev/null; wait $sampler 2>/dev/null + printf ' %-10s %-10s %s MiB\n' "depth=$1" "$ts" "$(sort -n "$vram" | tail -1)" + rm -f "$vram" +} + +DEPTHS=(4096 16384 32768 65536 131072 262144); [ $# -gt 0 ] && DEPTHS=("$@") +for D in "${DEPTHS[@]}"; do + R=3; [ "$D" -gt 32768 ] && R=1 + echo "== context depth=$D reps=$R (t/s, peak device memory)" + flock "$LOCK" bash -c "$(declare -f arm); BUILD='$BUILD'; MODEL='$MODEL'; PIN='$PIN'; NGEN='$NGEN' + arm 0 $D $R + arm 1 $D $R + arm 0 $D $R + arm 1 $D $R" +done diff --git a/docs/repro/r4-kv-pipeline-exact.py b/docs/repro/r4-kv-pipeline-exact.py index 0d160374d50f..0a0800f3f795 100644 --- a/docs/repro/r4-kv-pipeline-exact.py +++ b/docs/repro/r4-kv-pipeline-exact.py @@ -1,46 +1,70 @@ -# Greedy server output, hashed. Run through r4-kv-pipeline-exact.sh. +# Greedy server output, hashed, over several prefill corpora and prefill lengths. +# Run through r4-kv-pipeline-exact.sh. Compare the hashes across pipeline depths and against a +# build of the parent commit: the pipelined path must reproduce the ordered path exactly. import hashlib, json, sys, urllib.request -PORT, NTOK, NPARA = sys.argv[1], int(sys.argv[2]), int(sys.argv[3]) +PORT = sys.argv[1] +LENGTHS = [int(x) for x in sys.argv[2].split(",")] # approximate prefill tokens -SHORT = [ - "Explain how a B-tree index speeds up range queries in a relational database, step by step.", - "Write a Python function that merges two sorted lists, then explain its complexity.", - "Summarise the tradeoffs between optimistic and pessimistic concurrency control.", - "Describe, in order, what happens when a CPU takes a page fault on a memory-mapped file.", -] -# a long, deterministic filler so that decode runs against a deep KV window -PARA = ("A B-tree index stores keys in sorted order across a shallow, balanced tree. " - "Range queries descend once to the first qualifying leaf and then walk the leaf " - "chain sequentially, so the cost is one descent plus the size of the range. ") -DEEP = [ - "Summarise the text above in exactly five sentences.", - "List three claims the text above makes, then say which is weakest and why.", -] +# Four corpora with different token statistics, so that the deliveries being pipelined are not +# always the same shape of content: prose, source code, structured records, and dialogue. +CORPORA = { + "prose": ("A B-tree index stores keys in sorted order across a shallow, balanced tree. " + "Range queries descend once to the first qualifying leaf and then walk the leaf " + "chain sequentially, so the cost is one descent plus the size of the range. "), + "code": ("static int walk_leaf_chain(struct btree *t, uint64_t lo, uint64_t hi, " + "int (*cb)(void *, uint64_t), void *ctx) {\n" + " struct leaf *l = btree_descend(t, lo);\n" + " while (l && l->keys[0] <= hi) {\n" + " for (int i = 0; i < l->n; i++) { if (l->keys[i] > hi) return 0; " + "cb(ctx, l->keys[i]); }\n" + " l = l->next;\n }\n return 0;\n}\n"), + "records": ('{"id":%d,"region":"eu-central","bytes":918273,"status":"ok",' + '"latency_ms":12.75,"tags":["index","range","btree"]}\n'), + "dialogue": ("Q: Why does the planner prefer a sequential scan here?\n" + "A: Because the predicate matches most of the table, and random leaf access " + "would cost more than reading every page once.\n"), +} -def ask(label, prompt): +QUESTIONS = { + "prose": "Summarise the text above in exactly five sentences.", + "code": "Describe what the function above does, then name one bug it could hide.", + "records": "How many distinct fields does each record above have, and what are they?", + "dialogue": "State the single claim the answers above keep returning to.", +} + +def filler(name, target_tokens): + unit = CORPORA[name] + # roughly four characters to the token; the exact prefill length is reported per task + reps = max(1, (target_tokens * 4) // len(unit % 0 if "%d" in unit else unit)) + if "%d" in unit: + return "".join(unit % i for i in range(reps)) + return unit * reps + +def ask(label, prompt, ntok): body = json.dumps({"model": "m", "messages": [{"role": "user", "content": prompt}], - "max_tokens": NTOK, "temperature": 0, "top_k": 1, "seed": 1234}).encode() + "max_tokens": ntok, "temperature": 0, "top_k": 1, "seed": 1234}).encode() req = urllib.request.Request(f"http://127.0.0.1:{PORT}/v1/chat/completions", body, {"Content-Type": "application/json"}) try: - with urllib.request.urlopen(req, timeout=7200) as r: + with urllib.request.urlopen(req, timeout=14400) as r: d = json.load(r) except Exception as e: - print(f"{label} REQUEST_FAILED {type(e).__name__}") + print(f"{label} REQUEST_FAILED {type(e).__name__}", flush=True) return False m = d["choices"][0]["message"] # reasoning models put most of the generation in reasoning_content; hash both - body_text = (m.get("reasoning_content") or "") + "\x00" + (m.get("content") or "") + text = (m.get("reasoning_content") or "") + "\x00" + (m.get("content") or "") t = d.get("timings", {}) - print(f"{label} {hashlib.sha256(body_text.encode()).hexdigest()[:16]} " - f"prompt_n={t.get('prompt_n')} n={t.get('predicted_n')} " - f"tg={t.get('predicted_per_second'):.3f}") + print(f"{label:<18} {hashlib.sha256(text.encode()).hexdigest()[:16]} " + f"prompt_n={t.get('prompt_n'):<7} n={t.get('predicted_n'):<4} " + f"pp={t.get('prompt_per_second'):8.2f} tg={t.get('predicted_per_second'):7.3f}", flush=True) return True ok = True -for i, q in enumerate(SHORT): - ok &= ask(f"short{i}", q) -for i, q in enumerate(DEEP): - ok &= ask(f"deep{i} ", (PARA * NPARA) + "\n\n" + q) +for length in LENGTHS: + ntok = 256 if length <= 4096 else 128 + for name in CORPORA: + prompt = filler(name, length) + "\n\n" + QUESTIONS[name] + ok &= ask(f"{name}@{length}", prompt, ntok) sys.exit(0 if ok else 1) diff --git a/docs/repro/r4-kv-pipeline-exact.sh b/docs/repro/r4-kv-pipeline-exact.sh index 74b308eab470..6ab8d7b6b619 100755 --- a/docs/repro/r4-kv-pipeline-exact.sh +++ b/docs/repro/r4-kv-pipeline-exact.sh @@ -1,32 +1,34 @@ #!/bin/bash -# R4 gate 1: greedy server output must be byte-identical to the ordered path. -# Compare the hashes across pipeline depths, and against a build of the parent commit. +# R4 gate 1: greedy server output must be byte-identical to the ordered path, across several +# prefill corpora and prefill lengths. Compare the hashes across pipeline depths, and against a +# build of the parent commit. # -# LLAMA_KV_MODEL=/path/model.gguf docs/repro/r4-kv-pipeline-exact.sh [depth ...] +# LLAMA_KV_MODEL=/path/model.gguf docs/repro/r4-kv-pipeline-exact.sh [pipeline-depth ...] +# LLAMA_KV_LENGTHS=2048,18432,65536 selects the prefill lengths (default 2048,18432). set -u MODEL="${LLAMA_KV_MODEL:?set LLAMA_KV_MODEL to a .gguf path}" BUILD="${LLAMA_KV_BUILD:-build}" PIN="${LLAMA_KV_TASKSET:-0,2,4}" PORT="${LLAMA_KV_PORT:-18099}" -NTOK="${LLAMA_KV_NTOK:-256}" -NPARA="${LLAMA_KV_NPARA:-400}" +LENGTHS="${LLAMA_KV_LENGTHS:-2048,18432}" +CTX="${LLAMA_KV_CTX:-32768}" HERE="$(cd "$(dirname "$0")" && pwd)" DEPTHS=(0 1 4); [ $# -gt 0 ] && DEPTHS=("$@") rc=0 for D in "${DEPTHS[@]}"; do - echo "== pipeline depth=$D" + echo "== pipeline depth=$D ctx=$CTX prefill lengths=$LENGTHS" LOG=$(mktemp /tmp/r4-kv-pipeline.XXXX.log) GGML_KV_PIPELINE_DEPTH=$D taskset -c "$PIN" "$BUILD/bin/llama-server" -m "$MODEL" \ -ngl 99 -sm none -mg 0 -t 3 -nkvo --kv-cpu-pinned --recurrent-state-offload \ - -fa on -ctk q8_0 -ctv q8_0 -b 512 -ub 512 -c 32768 --parallel 1 \ + -fa on -ctk q8_0 -ctv q8_0 -b 512 -ub 512 -c "$CTX" --parallel 1 \ --host 127.0.0.1 --port "$PORT" --no-warmup > "$LOG" 2>&1 & SRV=$! - for _ in $(seq 1 300); do + for _ in $(seq 1 600); do curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1 && break sleep 1 done - python3 "$HERE/r4-kv-pipeline-exact.py" "$PORT" "$NTOK" "$NPARA" || rc=$? + python3 "$HERE/r4-kv-pipeline-exact.py" "$PORT" "$LENGTHS" || rc=$? kill $SRV 2>/dev/null; wait $SRV 2>/dev/null rm -f "$LOG" done diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index d6891d779da6..ec9d0aab9cc2 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -800,6 +800,13 @@ static bool ggml_backend_sched_allows_bufferless_kvarn_src( #define GGML_SCHED_TRANSPORT_MARGIN 2 #endif +// Device memory the transport ring leaves unclaimed. The ring is allocated after the graph +// allocator has reserved its buffers, so what it must not do is take the room those buffers may +// still have to grow into. +#ifndef GGML_SCHED_TRANSPORT_HEADROOM +#define GGML_SCHED_TRANSPORT_HEADROOM (512u*1024*1024) +#endif + // One staging slot of the transport ring. A slot is owned by the transfer stream while it is // being filled and by the consumer stream while it is being read; the two events below are the // handover in each direction. @@ -858,6 +865,8 @@ struct ggml_backend_sched_transport { int64_t n_graphs; int64_t p_graph_us, p_sync_us, p_copy_us, p_issue_us, p_bytes_early, p_bytes_late; + bool reported_no_room; // the "no room for the ring" warning is worth saying once, not per graph + int debug; }; @@ -1857,6 +1866,48 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { tr->plan_n_splits = sched->n_splits; tr->plan_n_inputs = n_inputs_total; + for (int i = 0; i < sched->n_splits; i++) { + struct ggml_backend_sched_split * split = &sched->splits[i]; + for (int j = 0; j < split->n_inputs; j++) { + if (ggml_backend_sched_input_can_stage(sched, split, j)) { + tr->input_staged[tr->split_input_ofs[i] + j] = 1; + } + } + } + + // A staged input copy may only be read by the split that owns it. The scheduler creates one + // copy per (tensor, backend) rather than per split, so a later split can be pointed at the + // same copy without it appearing in that split's input list -- and by then the ring may have + // recycled the slot for a different delivery. Splits whose copy has a reader further down the + // graph go back on the ordered path. + for (int i = 0; i < sched->n_splits; i++) { + struct ggml_backend_sched_split * split = &sched->splits[i]; + for (int j = 0; j < split->n_inputs; j++) { + if (!tr->input_staged[tr->split_input_ofs[i] + j]) { + continue; + } + const struct ggml_tensor * input_cpy = + tensor_copy(split->inputs[j], split->backend_id, sched->cur_copy); + + bool late_reader = false; + for (int k = i + 1; k < sched->n_splits && !late_reader; k++) { + const struct ggml_cgraph * g = &sched->splits[k].graph; + for (int n = 0; n < g->n_nodes && !late_reader; n++) { + for (int sr = 0; sr < GGML_MAX_SRC; sr++) { + if (g->nodes[n]->src[sr] == input_cpy) { + late_reader = true; + break; + } + } + } + } + + if (late_reader) { + tr->input_staged[tr->split_input_ofs[i] + j] = 0; + } + } + } + size_t slot_size = 0; for (int i = 0; i < sched->n_splits; i++) { struct ggml_backend_sched_split * split = &sched->splits[i]; @@ -1865,10 +1916,9 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { size_t need = 0; for (int j = 0; j < split->n_inputs; j++) { - if (!ggml_backend_sched_input_can_stage(sched, split, j)) { + if (!tr->input_staged[tr->split_input_ofs[i] + j]) { continue; } - tr->input_staged[tr->split_input_ofs[i] + j] = 1; need += GGML_PAD(ggml_nbytes(split->inputs[j]), tr->alignment); } @@ -1888,10 +1938,34 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { ggml_backend_sched_transport_free_ring(sched, true); ggml_backend_buffer_type_t buft = sched->bufts[tr->backend_id]; - ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(buft, slot_size * tr->n_slots); + const size_t ring_size = slot_size * tr->n_slots; + + // A slot has to hold one split's whole delivery, so the ring grows with the context and + // at some depth it stops fitting. The graph allocator has already reserved its buffers by + // now, but it may still have to grow one, and an allocation that succeeds here and + // starves that later is worse than not pipelining at all -- so leave it room rather than + // taking whatever is left. + ggml_backend_dev_t dev = ggml_backend_get_device(sched->backends[tr->backend_id]); + size_t dev_free = 0, dev_total = 0; + if (dev != NULL) { + ggml_backend_dev_memory(dev, &dev_free, &dev_total); + } + if (dev_free > 0 && ring_size + GGML_SCHED_TRANSPORT_HEADROOM > dev_free) { + if (!tr->reported_no_room) { + GGML_LOG_WARN("%s: transport ring would need %zu MiB and leave less than %u MiB " + "of the %zu MiB free on %s, staying on the ordered path\n", __func__, + ring_size >> 20, GGML_SCHED_TRANSPORT_HEADROOM >> 20, dev_free >> 20, + ggml_backend_buft_name(buft)); + tr->reported_no_room = true; + } + tr->n_staged = 0; + return; + } + + ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(buft, ring_size); if (buffer == NULL) { GGML_LOG_WARN("%s: failed to allocate %zu MiB for the transport ring, pipelining disabled\n", - __func__, (slot_size * tr->n_slots) >> 20); + __func__, ring_size >> 20); tr->n_staged = 0; return; } From 2703d04d43697161250797937428ab3e1b648784 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Tue, 25 Aug 2026 22:39:12 +0200 Subject: [PATCH 4/6] sched: cap the transport ring, and charge nothing when it declines A host-resident KV cache exists to keep device memory free. Measuring the transport from 4k to 256k showed it spending that memory back at a rate that grows with the context, which is working against the thing it accelerates: context gain ring peak device memory 4,096 +16.9% 27 MiB 10,169 -> 10,197 16,384 +60.1% 107 MiB 10,159 -> 10,263 32,768 +19.2% 213 MiB 10,161 -> 10,367 65,536 +12.9% 428 MiB 10,163 -> 10,573 131,072 +9.1% 855 MiB 10,537 -> 11,355 The gain narrows with depth and the cost does not, so past about 32k this is a bad trade on a card where the cache is on the host precisely because memory is short. --kv-pipeline-budget, default 128 MiB, is an absolute cap on the ring rather than a fraction of what happens to be free. Over it the scheduler declines and keeps the ordered path. Three things make declining actually free: - The cap is checked before anything is allocated, on every plan rather than only when the ring has to grow. - The decision is latched. A context only grows, so a ring that fits the small windows of early prefill would only have to be given back later. - The transfer backend is created on the first graph that stages something, not when the depth is set, and is released along with the ring. It is a second device context and it 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 against the ordered path's 10,161 MiB and throughput matches (12.965 vs 12.984 t/s). --kv-pipeline-budget 512 buys 15.487 t/s there for 206 MiB, for whoever wants that trade. Also tried and reverted: planning the ring during reserve, so that ggml-alloc would not budget blocks for the staged input copies as well. It reclaims nothing. The compute buffer is sized by a reserve that runs before the KV cache exists -- 342.27 MiB either way at 32,768 -- and the copies then fit inside space already reserved for other transients. The ordered path's copies are free, so the ring is additive and the only lever is its size. Known limitation: the cap is applied per graph, so a run whose context grows past it still shows a transient peak (+112 MiB at 32,768) before the latch trips, even though the steady state is +0. Deciding against the context's final size needs KV geometry the scheduler does not have. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01SRa5ZTc1bpnk1fE3mrDcSh --- common/arg.cpp | 14 +++ common/common.cpp | 1 + common/common.h | 1 + docs/beellama-args.md | 1 + docs/kv-transport-pipelining.md | 77 +++++++++++++- ggml/include/ggml-backend.h | 5 + ggml/src/ggml-backend.cpp | 175 ++++++++++++++++++++++++-------- include/llama.h | 3 + src/llama-context.cpp | 4 + src/llama-cparams.h | 1 + 10 files changed, 235 insertions(+), 47 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index c12d079b26f3..5f62cef069a1 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2614,6 +2614,20 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.kv_pipeline_depth = value; } ).set_env("LLAMA_ARG_KV_PIPELINE_DEPTH")); + add_opt(common_arg( + {"--kv-pipeline-budget"}, "N", + string_format("hard cap, in MiB, on the device memory that pipelined delivery of a host-resident KV cache " + "may use. A staging slot holds one attention layer's K or V over the whole context, so the " + "requirement grows with the context; past this cap the scheduler declines and keeps the " + "ordered path, so a host-resident cache never quietly trades away the device memory it exists " + "to save. 0 removes the cap (default: %d)", params.kv_pipeline_budget_mib), + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("--kv-pipeline-budget must not be negative"); + } + params.kv_pipeline_budget_mib = value; + } + ).set_env("LLAMA_ARG_KV_PIPELINE_BUDGET")); add_opt(common_arg( {"--recurrent-state-offload"}, {"--no-recurrent-state-offload"}, diff --git a/common/common.cpp b/common/common.cpp index 5bb20560589b..cab11aa47b4f 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1891,6 +1891,7 @@ struct llama_context_params common_context_params_to_llama(const common_params & cparams.offload_kqv = !params.no_kv_offload; cparams.kv_cpu_pinned = params.kv_cpu_pinned; cparams.kv_pipeline_depth = params.kv_pipeline_depth < 0 ? 0 : (uint32_t) params.kv_pipeline_depth; + cparams.kv_pipeline_budget_mib = params.kv_pipeline_budget_mib < 0 ? 0 : (uint32_t) params.kv_pipeline_budget_mib; cparams.recurrent_state_offload = params.recurrent_state_offload; cparams.phase_aware_workspace = params.phase_aware_workspace; cparams.live_context_workspace = params.live_context_workspace; diff --git a/common/common.h b/common/common.h index f332afd0e440..1b487bd8d503 100644 --- a/common/common.h +++ b/common/common.h @@ -643,6 +643,7 @@ struct common_params { bool phase_aware_workspace = false; // resize compute schedulers between prompt and generation phases bool live_context_workspace = false; // size supported attention workspaces from the padded live KV extent int32_t kv_pipeline_depth = 1; // splits of look-ahead for pipelined delivery of a host-resident KV cache (0 = off) + int32_t kv_pipeline_budget_mib = 128; // hard cap on the device memory that delivery may use (0 = uncapped) bool warmup = true; // warmup run bool check_tensors = false; // validate tensor data bool no_op_offload = false; // globally disable offload host tensor operations to device diff --git a/docs/beellama-args.md b/docs/beellama-args.md index 9bd23dab7f2c..feb7c364a7cc 100644 --- a/docs/beellama-args.md +++ b/docs/beellama-args.md @@ -203,6 +203,7 @@ context-dependent H2D traffic; it is not equivalent to CPU attention. | `--kv-gpu-layers N` | `LLAMA_ARG_KV_GPU_LAYERS` | `0` | Under `--no-kv-offload`, retains the first `N` attention-KV layers on the device. This trades VRAM for lower repeated host-to-device traffic; a value at or above the attention-layer count matches full KV offload for those layers. | | `--spec-draft-kv-gpu-layers N`, `--kv-gpu-layers-draft N` | `LLAMA_ARG_SPEC_DRAFT_KV_GPU_LAYERS` | Inherit target | Overrides the target KV policy for a separate draft context and keeps the first `N` independently owned draft attention-KV layers device-resident. Shared KV layers follow their owner. The INI key is `spec-draft-kv-gpu-layers`. | | `--kv-pipeline-depth N` | `LLAMA_ARG_KV_PIPELINE_DEPTH` | `1` | How many splits ahead the scheduler delivers a host-resident KV cache to the accelerator, so that the transfer runs while the previous split computes instead of in series with it. `0` restores the ordered path. Only engages where a host-resident cache produces the deliveries, and costs `(N + 2) x (largest staged split)` of device memory. See [`kv-transport-pipelining.md`](kv-transport-pipelining.md). | +| `--kv-pipeline-budget N` | `LLAMA_ARG_KV_PIPELINE_BUDGET` | `128` | Hard cap, in MiB, on the device memory pipelined delivery may use. A staging slot holds one attention layer's K or V over the whole context, so the requirement grows with the context; past this cap the scheduler declines and keeps the ordered path, at no steady-state cost, so a host-resident cache never quietly trades away the device memory it exists to save. `0` removes the cap. | | `--op-offload`, `--no-op-offload` | — | Enabled | Controls whether operations over host tensors may execute on the accelerator. The negated form selects the separately tested CPU execution route; it is not the serving-performance path measured below. | Target and draft attention-KV placement is independent of model-weight diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index 24bf2050aa77..94e94f2fe18c 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -15,6 +15,11 @@ kernels of the split before it. an effect where a host-resident cache produces the deliveries; `0` restores the ordered path exactly. +The staging it needs is bounded by `--kv-pipeline-budget` (default 128 MiB), so +that a cache which lives on the host to keep device memory free never quietly +spends that memory back. Past the cap the scheduler declines and the ordered path +runs, at no cost. See [The budget](#the-budget). + ## What it changes, and what it must not R4 pipelines *deliveries*, not attention. Every byte and every attention @@ -131,10 +136,74 @@ RTX 4070 (11,902 MiB usable, sm_89), driver 610.57.04 / CUDA 13.3, i5-13400F, Server decode behind an 18,422-token prompt (`docs/repro/r4-kv-pipeline-exact.sh`): **18.468 -> 30.685 t/s, +66.2%**. -Device allocation high-water at 32,768, sampled with `nvidia-smi` across a -`tg64 @ d32768` run: **10,271 -> 10,477 MiB, +206 MiB**, against a ring of -3 x 71 MiB. The ring competes with the model, and at 32,768 it is the largest -single new allocation the feature makes. +### Across context depth, with device memory + +`docs/repro/r4-kv-pipeline-context-sweep.sh`, A/B/A/B, peak device memory sampled +with `nvidia-smi` across each arm. Both passes agreed to the digits shown. + +| Context | ordered | pipelined | gain | peak device memory | delta | ring | +|---|---:|---:|---:|---|---:|---:| +| 4,096 | 31.66 | 37.01 | **+16.9%** | 10,169 -> 10,197 MiB | +28 MiB | 27 MiB | +| 16,384 | 19.64 | 31.43 | **+60.1%** | 10,159 -> 10,263 MiB | +104 MiB | 107 MiB | +| 32,768 | 12.99 | 15.49 | **+19.2%** | 10,161 -> 10,367 MiB | +206 MiB | 213 MiB | +| 65,536 | 7.74 | 8.74 | **+12.9%** | 10,163 -> 10,573 MiB | +410 MiB | 428 MiB | +| 131,072 | 4.29 | 4.68 | **+9.1%** | 10,537 -> 11,355 MiB | +818 MiB | 855 MiB | +| 262,144 | 2.25 | 2.24 | **declined** | 11,329 -> 11,391 MiB | +62 MiB | not allocated | + +Two curves run in opposite directions here, and both matter. + +**The gain narrows with depth.** A token is copy plus compute; as the context +grows the copy grows with it while the compute per staged split does not, so the +share of the token that can hide a transfer shrinks. At 16,384 compute still +covers most of the copy; by 131,072 it covers a tenth of it. That is arithmetic, +not an implementation limit, and no amount of look-ahead changes it. + +**The ring's cost does not narrow.** It is `(depth + 2)` slots of one staged +split, and a staged split is K and V of one attention layer over the whole +context: it doubles every time the context doubles. At 131,072 it claims 818 MiB +of an 11,902 MiB card to buy 9.1%. + +At 262,144 the ring would need 1.7 GiB against 573 MiB free, so it declines and +the run stays on the ordered path -- 2.25 against 2.24 t/s, inside the spread of +the ordered arm's own two passes, and 62 MiB of device memory for the transfer +backend's context. Declining is the intended outcome, not a failure: the +62 MiB +and the unchanged throughput are what "the guard did its job" looks like. + +**On a memory-constrained card, past roughly 64k the same device memory is +probably better spent on `--kv-gpu-layers`.** At 131,072 a staged split is +285 MiB, so the 818 MiB the ring takes is about three attention layers' worth of +K and V; making three of sixteen layers device-resident removes about 19% of the +host-to-device traffic against the 9.1% the ring buys. That comparison has not +been measured here and it will move with the model's layer count and the card, so +it is a pointer for whoever tunes a deployment, not a recommendation. + +### The budget + +The table above is what the feature costs uncapped, and it is the reason it is +capped. A host-resident KV cache exists to keep device memory free; a transport +that speeds it up by spending hundreds of MiB of that memory is working against +the thing it is accelerating. `--kv-pipeline-budget` (default 128 MiB) is an +absolute cap on the ring, not a fraction of what happens to be free: + +- Under the cap the ring is allocated and the deliveries pipeline: 4,096 and + 16,384 in the table, at 28 MiB and 104 MiB. +- Over it the scheduler declines and keeps the ordered path, and the decision is + latched, because a context only grows and a ring allocated for the small + windows of early prefill would only have to be given back later. +- Declining costs nothing in steady state. Both the ring and the transfer + backend's device context are released: at 32,768 with the default budget, + device memory settles at 10,161 MiB, the same as the ordered path, and + throughput matches it (12.965 against 12.984 t/s). + +Raising the budget trades that memory back for speed where it is worth it: +`--kv-pipeline-budget 512` at 32,768 gives 15.487 t/s for 206 MiB. + +**Known limitation.** The cap is applied per graph, so a run whose context grows +past it still allocates a ring for the early prefill graphs and releases it once +the window outgrows the budget -- at 32,768 that shows up as a transient peak of ++112 MiB even though the steady state is +0. Deciding against the context's final +size rather than the current graph's would remove it, and needs the KV geometry +the scheduler does not have. Where the split-loop host time goes, per decode graph at 18.5k (`GGML_SCHED_TRANSPORT_DEBUG=2`): diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index a706268882ac..43877fd0794f 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -392,6 +392,11 @@ extern "C" { // allocated. GGML_API void ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, int depth); + // Hard cap on the staging ring, in bytes. A host-resident cache exists to keep device memory + // free, so the ring is capped outright and not merely against what happens to be free: past + // the cap the scheduler declines and keeps the ordered path. 0 removes the cap. Default 128 MiB. + GGML_API void ggml_backend_sched_set_transport_pipeline_budget(ggml_backend_sched_t sched, size_t bytes); + // Number of staged deliveries and staged bytes issued since the scheduler was created. GGML_API void ggml_backend_sched_get_transport_pipeline_stats(ggml_backend_sched_t sched, int64_t * n_deliveries, int64_t * n_bytes_early, int64_t * n_bytes_late); diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index ec9d0aab9cc2..b4ce3840cf3c 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -807,6 +807,14 @@ static bool ggml_backend_sched_allows_bufferless_kvarn_src( #define GGML_SCHED_TRANSPORT_HEADROOM (512u*1024*1024) #endif +// Default cap on the ring itself. A host-resident KV cache exists to keep device memory free, so +// the transport that speeds it up has to stay small whether or not the device has room to spare: +// a slot is one attention layer's K or V over the whole context, which grows without bound as the +// context does. Past this the feature declines rather than quietly spending hundreds of MiB. +#ifndef GGML_SCHED_TRANSPORT_BUDGET +#define GGML_SCHED_TRANSPORT_BUDGET (128u*1024*1024) +#endif + // One staging slot of the transport ring. A slot is owned by the transfer stream while it is // being filled and by the consumer stream while it is being read; the two events below are the // handover in each direction. @@ -834,6 +842,7 @@ struct ggml_backend_sched_transport { ggml_backend_buffer_t buffer; // the ring itself size_t slot_size; size_t alignment; + size_t budget; // hard cap on the ring, in bytes struct ggml_backend_sched_transport_slot slots[GGML_SCHED_MAX_TRANSPORT_SLOTS]; @@ -866,6 +875,7 @@ struct ggml_backend_sched_transport { int64_t p_graph_us, p_sync_us, p_copy_us, p_issue_us, p_bytes_early, p_bytes_late; bool reported_no_room; // the "no room for the ring" warning is worth saying once, not per graph + bool over_budget; // latched: the ring has been asked for more than it may have int debug; }; @@ -1724,8 +1734,11 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra // pipelined transport of host-resident split inputs // +static void ggml_backend_sched_transport_teardown(ggml_backend_sched_t sched); + static bool ggml_backend_sched_transport_enabled(ggml_backend_sched_t sched) { - return sched->transport.depth >= 1 && sched->transport.n_slots >= 2 && sched->transport.transfer != NULL; + const struct ggml_backend_sched_transport * tr = &sched->transport; + return tr->depth >= 1 && tr->n_slots >= 2 && tr->backend_id >= 0 && !tr->over_budget; } // The annotation lives on the tensor that owns the storage; a split input is normally a view of @@ -1750,9 +1763,13 @@ static size_t ggml_backend_sched_input_stable_prefix(const struct ggml_tensor * return avail < bytes ? avail : bytes; } -// Whether a split input belongs in the ring. Decided once per plan; from then on execution -// consults the recorded answer rather than asking again, because the write position this depends -// on moves between ubatches while the plan and the addresses it handed out do not. +// Whether a split input belongs in the ring. +// +// Deliberately independent of the stable prefix. Membership decides where an input copy lives, +// which the graph allocator has to know when it reserves -- and at reserve time there is no +// ubatch yet, so no prefix. The prefix decides only how much of a staged input can go early; +// zero means all of it waits for the split, which is the ordered path's timing with the ring's +// storage, and is still correct. static bool ggml_backend_sched_input_can_stage( ggml_backend_sched_t sched, struct ggml_backend_sched_split * split, int input_id) { if (split->backend_id != sched->transport.backend_id) { @@ -1766,11 +1783,6 @@ static bool ggml_backend_sched_input_can_stage( return false; } - // an input with no stable prefix has nothing that could be delivered ahead of its split - if (ggml_backend_sched_input_stable_prefix(input) == 0) { - return false; - } - ggml_backend_buffer_t buf = input->view_src ? input->view_src->buffer : input->buffer; if (buf == NULL || !ggml_backend_buffer_is_host(buf)) { return false; @@ -1819,6 +1831,48 @@ static void ggml_backend_sched_transport_free_ring(ggml_backend_sched_t sched, b } } +// The transfer backend and the slot events are created on demand, so that a scheduler which +// never gets to stage anything -- no eligible inputs, or no room within the budget -- does not +// carry a second device context for nothing. +static bool ggml_backend_sched_transport_ensure_backend(ggml_backend_sched_t sched) { + struct ggml_backend_sched_transport * tr = &sched->transport; + + if (tr->transfer != NULL) { + return true; + } + + ggml_backend_dev_t dev = ggml_backend_get_device(sched->backends[tr->backend_id]); + if (dev == NULL) { + return false; + } + + ggml_backend_t transfer = ggml_backend_dev_init(dev, NULL); + if (transfer == NULL) { + return false; + } + + bool ok = true; + for (int slot = 0; slot < tr->n_slots && ok; slot++) { + tr->slots[slot].ready = ggml_backend_event_new(dev); + tr->slots[slot].release = ggml_backend_event_new(dev); + ok = tr->slots[slot].ready != NULL && tr->slots[slot].release != NULL; + } + + if (!ok) { + for (int slot = 0; slot < tr->n_slots; slot++) { + ggml_backend_event_free(tr->slots[slot].ready); + ggml_backend_event_free(tr->slots[slot].release); + tr->slots[slot].ready = NULL; + tr->slots[slot].release = NULL; + } + ggml_backend_free(transfer); + return false; + } + + tr->transfer = transfer; + return true; +} + // Lay the ring out over the current split list and point the staged input copies at it. Called // before the graph is allocated: ggml-alloc leaves a tensor that already has data alone, so the // staged copies are excluded from its reuse analysis instead of competing with it. @@ -1902,7 +1956,19 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { } } - if (late_reader) { + // a view of the copy would have had its address resolved from the copy's own, so + // redirecting the copy afterwards would leave the view pointing at the old block + bool viewed = false; + for (int k = 0; k < sched->n_splits && !viewed; k++) { + const struct ggml_cgraph * g = &sched->splits[k].graph; + for (int n = 0; n < g->n_nodes && !viewed; n++) { + if (g->nodes[n]->view_src == input_cpy) { + viewed = true; + } + } + } + + if (late_reader || viewed) { tr->input_staged[tr->split_input_ofs[i] + j] = 0; } } @@ -1934,12 +2000,40 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { return; } + ggml_backend_buffer_type_t buft = sched->bufts[tr->backend_id]; + const size_t ring_size = slot_size * tr->n_slots; + + // Checked before anything is allocated, and on every plan rather than only when the ring has + // to grow: a graph the ring already fits does not license keeping it once a later graph needs + // more than it may have. Declining has to cost nothing, so a ring in hand is given back. + // + // The point of a host-resident cache is to leave device memory alone, so the cap is absolute + // and not merely a fraction of what happens to be free. + // Latched, because a context only grows: the early prefill graphs have a small window and + // would fit, and allocating a ring for them only to give it back once the window outgrows the + // budget claims device memory that a host-resident cache is supposed to be leaving alone. + if (tr->budget > 0 && (ring_size > tr->budget || tr->over_budget)) { + tr->over_budget = true; + if (!tr->reported_no_room) { + GGML_LOG_WARN("%s: transport ring would need %zu MiB against a %zu MiB budget, " + "staying on the ordered path (raise --kv-pipeline-budget to spend more " + "device memory on it)\n", __func__, ring_size >> 20, tr->budget >> 20); + tr->reported_no_room = true; + } + ggml_backend_sched_transport_teardown(sched); + tr->n_staged = 0; + return; + } + + // nothing has been allocated for the transport until here + if (!ggml_backend_sched_transport_ensure_backend(sched)) { + tr->n_staged = 0; + return; + } + if (tr->buffer == NULL || tr->slot_size < slot_size) { ggml_backend_sched_transport_free_ring(sched, true); - ggml_backend_buffer_type_t buft = sched->bufts[tr->backend_id]; - const size_t ring_size = slot_size * tr->n_slots; - // A slot has to hold one split's whole delivery, so the ring grows with the context and // at some depth it stops fitting. The graph allocator has already reserved its buffers by // now, but it may still have to grow one, and an allocation that succeeds here and @@ -2499,6 +2593,7 @@ ggml_backend_sched_t ggml_backend_sched_new( sched->op_offload = op_offload; sched->transport.backend_id = -1; + sched->transport.budget = GGML_SCHED_TRANSPORT_BUDGET; { const char * GGML_SCHED_TRANSPORT_DEBUG = getenv("GGML_SCHED_TRANSPORT_DEBUG"); sched->transport.debug = GGML_SCHED_TRANSPORT_DEBUG ? atoi(GGML_SCHED_TRANSPORT_DEBUG) : 0; @@ -2527,9 +2622,7 @@ static void ggml_backend_sched_transport_teardown(ggml_backend_sched_t sched) { tr->transfer = NULL; } - tr->n_slots = 0; - tr->backend_id = -1; - tr->n_staged = 0; + tr->n_staged = 0; } void ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, int depth) { @@ -2546,7 +2639,7 @@ void ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, } struct ggml_backend_sched_transport * tr = &sched->transport; - if (tr->depth == depth && (depth < 1 || tr->transfer != NULL)) { + if (tr->depth == depth && (depth < 1 || tr->backend_id >= 0)) { return; } @@ -2583,37 +2676,16 @@ void ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, continue; } - // a second backend on the same device gives the deliveries a stream of their own - ggml_backend_t transfer = ggml_backend_dev_init(dev, NULL); - if (transfer == NULL) { - continue; - } - - bool ok = true; - for (int slot = 0; slot < n_slots && ok; slot++) { - tr->slots[slot].ready = ggml_backend_event_new(dev); - tr->slots[slot].release = ggml_backend_event_new(dev); - ok = tr->slots[slot].ready != NULL && tr->slots[slot].release != NULL; - } - - if (!ok) { - for (int slot = 0; slot < n_slots; slot++) { - ggml_backend_event_free(tr->slots[slot].ready); - ggml_backend_event_free(tr->slots[slot].release); - tr->slots[slot].ready = NULL; - tr->slots[slot].release = NULL; - } - ggml_backend_free(transfer); - continue; - } - - tr->transfer = transfer; + // The transfer backend is a second backend on the same device, which gives the deliveries + // a stream of their own -- and costs a device context of its own. It is created on the + // first graph that actually stages something, not here: a scheduler that never has room + // for a ring must not pay for a context it will not use. tr->backend_id = i; tr->n_slots = n_slots; tr->alignment = std::max(ggml_backend_buft_get_alignment(sched->bufts[i]), 128); if (tr->debug > 0) { - GGML_LOG_INFO("%s: pipelined host transport enabled on %s, %d splits ahead, %d slots\n", + GGML_LOG_INFO("%s: pipelined host transport selected %s, %d splits ahead, %d slots\n", __func__, ggml_backend_name(backend), depth, n_slots); } return; @@ -2624,6 +2696,23 @@ void ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, } } +void ggml_backend_sched_set_transport_pipeline_budget(ggml_backend_sched_t sched, size_t bytes) { + GGML_ASSERT(sched); + + const char * env = getenv("GGML_KV_PIPELINE_BUDGET_MIB"); + if (env != NULL) { + bytes = (size_t) strtoull(env, NULL, 10) * 1024 * 1024; + } + + if (sched->transport.budget != bytes) { + sched->transport.budget = bytes; + sched->transport.reported_no_room = false; + sched->transport.over_budget = false; + // the ring in hand may no longer be allowed + ggml_backend_sched_transport_free_ring(sched, true); + } +} + void ggml_backend_sched_get_transport_pipeline_stats( ggml_backend_sched_t sched, int64_t * n_deliveries, int64_t * n_bytes_early, int64_t * n_bytes_late) { GGML_ASSERT(sched); diff --git a/include/llama.h b/include/llama.h index 4e5c55e49a3a..e294a74919e9 100644 --- a/include/llama.h +++ b/include/llama.h @@ -521,6 +521,9 @@ extern "C" { // 0 keeps the ordered path, where a decode token pays the transfer and the // attention kernels in series. Costs (kv_pipeline_depth + 2) * (largest staged // split) of device memory. + uint32_t kv_pipeline_budget_mib; // hard cap on that device memory, in MiB. Past it the scheduler declines and + // keeps the ordered path, so a host-resident cache never quietly trades the + // device memory it exists to save. 0 removes the cap. }; struct llama_model_tensor_override { diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 2562d8773ffc..3b0b16dff2b2 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -303,6 +303,7 @@ llama_context::llama_context( cparams.phase_aware_workspace = params.phase_aware_workspace; cparams.live_context_workspace = params.live_context_workspace; cparams.kv_pipeline_depth = params.kv_pipeline_depth; + cparams.kv_pipeline_budget_mib = params.kv_pipeline_budget_mib; cparams.kv_gpu_layers = params.kv_gpu_layers; cparams.offload_attn_compute = params.offload_kqv || (params.op_offload && (params.kv_cpu_pinned || params.kv_gpu_layers > 0)); @@ -1122,6 +1123,8 @@ void llama_context::sched_reserve(uint32_t n_tokens_req, uint32_t n_kv_req) { backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, pipeline_parallel, cparams.op_offload)); // only a host-resident KV cache produces the deliveries this pipelines + ggml_backend_sched_set_transport_pipeline_budget(sched.get(), + (size_t) cparams.kv_pipeline_budget_mib * 1024 * 1024); ggml_backend_sched_set_transport_pipeline_depth(sched.get(), cparams.kv_cpu_pinned || !cparams.offload_kqv ? (int) cparams.kv_pipeline_depth : 0); if (sched_shared_buffers) { @@ -4569,6 +4572,7 @@ llama_context_params llama_context_default_params() { /*.phase_aware_workspace =*/ false, /*.live_context_workspace =*/ false, /*.kv_pipeline_depth =*/ 1, + /*.kv_pipeline_budget_mib =*/ 128, }; return result; diff --git a/src/llama-cparams.h b/src/llama-cparams.h index af34cd201a62..34ab37519d5a 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -66,6 +66,7 @@ struct llama_cparams { bool phase_aware_workspace; bool live_context_workspace; uint32_t kv_pipeline_depth; + uint32_t kv_pipeline_budget_mib; std::vector embeddings_layer_inp; // [n_layer()] extract input embeddings for layer From 3d001726519b34644b1a9140de3797703c24bd6b Mon Sep 17 00:00:00 2001 From: piggidragon Date: Tue, 25 Aug 2026 23:22:22 +0200 Subject: [PATCH 5/6] docs: mark the uncapped measurements as uncapped, and the single-GPU scope as single-GPU Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01SRa5ZTc1bpnk1fE3mrDcSh --- docs/kv-transport-pipelining.md | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index 94e94f2fe18c..c1f5188b4ef6 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -133,6 +133,12 @@ RTX 4070 (11,902 MiB usable, sm_89), driver 610.57.04 / CUDA 13.3, i5-13400F, | 16,384 | 3 | 19.6765, 19.6854 | 31.5352, 31.5807 | **+60.4%** | 34.88 | 90.4% | | 32,768 | 3 | 13.0264, 13.0254 | 15.5325, 15.5329 | **+19.3%** | 20.83 | 74.6% | +These are the uncapped numbers, measured before `--kv-pipeline-budget` existed; +they are what the ring can buy, and the 32,768 row needs +`--kv-pipeline-budget 512` to reproduce, because 213 MiB is over the 128 MiB +default. At the default the 4,096 and 16,384 rows stand and 32,768 declines to +the ordered path. See [The budget](#the-budget). + Server decode behind an 18,422-token prompt (`docs/repro/r4-kv-pipeline-exact.sh`): **18.468 -> 30.685 t/s, +66.2%**. @@ -264,8 +270,11 @@ transport never enabled because the scheduler is given a depth of 0. staged split is both K and V of one attention layer over the whole context. That is linear in context length, and it is what bounds the feature at depth rather than anything about the transfer itself. -- One ring, on the first backend that qualifies. Additional accelerators keep the - ordered path; multi-GPU host-KV is not covered. +- **One ring, on the first backend that qualifies. Additional accelerators keep + the ordered path, and nothing here has been measured on more than one GPU** -- + every number in this document is `-sm none -mg 0` on a single RTX 4070. A + multi-GPU host-resident cache needs its own validation before any of this is + claimed for it. - The scheduler must be configured with the device's own default buffer type. A scheduler built on a split or host buffer type keeps the ordered path. - `GGML_KV_PIPELINE_DEPTH` overrides the depth for tools that do not expose the From c3f4d2c2a96511f08b46892619d0c2a12cedb6bf Mon Sep 17 00:00:00 2001 From: piggidragon Date: Wed, 26 Aug 2026 01:08:43 +0200 Subject: [PATCH 6/6] sched: give every accelerator its own transport ring The transport picked the first backend that could transfer asynchronously and order streams with events, and stopped there. On a layer-split model that leaves every device after the first paying copy + compute in series while the first one does not, which is the wrong shape for a feature whose whole purpose is to stop paying that. Each eligible backend now carries its own ring: its own transfer backend and stream, its own slots and handover events, its own look-ahead cursor, and its own budget decision. The pieces that had to become per-ring rather than global: - The delivery order is numbered within a ring, not across the graph, so a slot is recycled against its own device's progress. - Each ring walks the split list on its own cursor. One device saturating its look-ahead must not stop another from running ahead on its own. - The budget and the device-memory headroom are decided per ring and latched per ring. A device with no room falls back to the ordered path on its own and un-stages only its own inputs; the others keep pipelining. The cap stays per device, which is the honest reading now that there is more than one: each ring is a claim on its own card, and a second accelerator brings its own memory to spend. Measured on an RTX 4070 + RTX 3060, Qwen3.8-27B-UD-Q5_K_M (18.40 GiB) at -sm layer -ts 55/45, -nkvo --kv-cpu-pinned, q8_0 K/V, tg32 @ d4096: ordered 6.91 t/s peak 9787 / 9069 MiB pipelined 7.52 t/s peak 9815 / 9099 MiB +8.8%, +28 MiB per card with both rings engaged -- CUDA0 staging 9 of 355 splits and CUDA1 staging 7. Single-GPU is unchanged: 19.66 -> 31.42 t/s at d16384 on the 4070, as before. Tensor parallelism (-sm tensor) is still on the ordered path and is not addressed here. The scheduler sees one meta backend there, and ggml-backend-meta implements neither event_record/event_wait on the backend nor event_new/event_free/ event_synchronize on the device, so the transport cannot order a transfer stream against the consumer and correctly declines. Supporting it means adding events to the meta layer, relaxing its set_tensor_async offset == 0 assertion, and giving a transfer-only meta backend a way not to stand up a second NCCL communicator -- all in BeeLlama's tensor-parallel layer rather than the KV-offload line. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01SRa5ZTc1bpnk1fE3mrDcSh --- ggml/src/ggml-backend.cpp | 553 ++++++++++++++++++++++---------------- 1 file changed, 322 insertions(+), 231 deletions(-) diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index b4ce3840cf3c..9b500afdd1b6 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -824,6 +824,28 @@ struct ggml_backend_sched_transport_slot { bool release_armed; // a reader was enqueued and has not been waited for yet }; +// One ring per accelerator the scheduler drives. A layer-split model gives every device its own +// splits and its own host-resident deliveries, so each needs its own transfer stream, its own +// staging, and its own place in the look-ahead: one device running ahead must not consume another +// device's slots, and one device declining for want of memory must not disable the others. +struct ggml_backend_sched_transport_ring { + bool eligible; // this backend can transfer asynchronously and order with events + + ggml_backend_t transfer; // second backend on the same device: owns the transfer stream + ggml_backend_buffer_t buffer; // the ring itself + size_t slot_size; + size_t alignment; + + struct ggml_backend_sched_transport_slot slots[GGML_SCHED_MAX_TRANSPORT_SLOTS]; + + int n_staged; // staged splits on this backend in the current graph + int consumed; // of those, how many readers have been enqueued + int scan_cursor; // how far the look-ahead has walked the split list for this ring + + bool reported_no_room; // the "no room for the ring" warning is worth saying once, not per graph + bool over_budget; // latched: this ring has been asked for more than it may have +}; + // Pipelined delivery of host-resident split inputs. // // The ordered path issues a split's host-to-device delivery on the consumer's own stream right @@ -834,28 +856,22 @@ struct ggml_backend_sched_transport_slot { // its last graph-level consumer is done, and a look-ahead transfer is still in flight outside // that lifetime. struct ggml_backend_sched_transport { - int depth; // how many splits ahead deliveries run; 0 disables pipelining - int n_slots; // slots actually created: depth + GGML_SCHED_TRANSPORT_MARGIN - int backend_id; // the consumer backend this ring serves, -1 if none + int depth; // how many splits ahead deliveries run; 0 disables pipelining + int n_slots; // slots per ring: depth + GGML_SCHED_TRANSPORT_MARGIN + size_t budget; // hard cap on each ring, in bytes - ggml_backend_t transfer; // second backend on the same device: owns the transfer stream - ggml_backend_buffer_t buffer; // the ring itself - size_t slot_size; - size_t alignment; - size_t budget; // hard cap on the ring, in bytes - - struct ggml_backend_sched_transport_slot slots[GGML_SCHED_MAX_TRANSPORT_SLOTS]; + struct ggml_backend_sched_transport_ring rings[GGML_SCHED_MAX_BACKENDS]; - // plan for the current graph, indexed by split id: the delivery order of the split, or -1 - // when the split stages nothing + // plan for the current graph, indexed by split id: the delivery order of the split within its + // own backend's ring, or -1 when the split stages nothing int * split_order; int plan_capacity; int plan_n_splits; int plan_n_inputs; - int n_staged; - int scan_cursor; + int n_staged; // over all rings, so that execution can skip the machinery entirely + int n_rings_used; - // which inputs the plan put in the ring, flattened over splits. Membership is decided once, + // which inputs the plan put in a ring, flattened over splits. Membership is decided once, // when the ring is laid out, and is what execution goes by: the amount that can be delivered // early moves with every ubatch, but which input copies live in the ring must not. unsigned char * input_staged; @@ -874,9 +890,6 @@ struct ggml_backend_sched_transport { int64_t n_graphs; int64_t p_graph_us, p_sync_us, p_copy_us, p_issue_us, p_bytes_early, p_bytes_late; - bool reported_no_room; // the "no room for the ring" warning is worth saying once, not per graph - bool over_budget; // latched: the ring has been asked for more than it may have - int debug; }; @@ -1736,9 +1749,22 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra static void ggml_backend_sched_transport_teardown(ggml_backend_sched_t sched); -static bool ggml_backend_sched_transport_enabled(ggml_backend_sched_t sched) { +static bool ggml_backend_sched_transport_ring_enabled(ggml_backend_sched_t sched, int backend_id) { const struct ggml_backend_sched_transport * tr = &sched->transport; - return tr->depth >= 1 && tr->n_slots >= 2 && tr->backend_id >= 0 && !tr->over_budget; + if (tr->depth < 1 || tr->n_slots < 2 || backend_id < 0) { + return false; + } + const struct ggml_backend_sched_transport_ring * r = &tr->rings[backend_id]; + return r->eligible && !r->over_budget; +} + +static bool ggml_backend_sched_transport_enabled(ggml_backend_sched_t sched) { + for (int i = 0; i < sched->n_backends; i++) { + if (ggml_backend_sched_transport_ring_enabled(sched, i)) { + return true; + } + } + return false; } // The annotation lives on the tensor that owns the storage; a split input is normally a view of @@ -1763,7 +1789,7 @@ static size_t ggml_backend_sched_input_stable_prefix(const struct ggml_tensor * return avail < bytes ? avail : bytes; } -// Whether a split input belongs in the ring. +// Whether a split input belongs in its backend's ring. // // Deliberately independent of the stable prefix. Membership decides where an input copy lives, // which the graph allocator has to know when it reserves -- and at reserve time there is no @@ -1772,7 +1798,7 @@ static size_t ggml_backend_sched_input_stable_prefix(const struct ggml_tensor * // storage, and is still correct. static bool ggml_backend_sched_input_can_stage( ggml_backend_sched_t sched, struct ggml_backend_sched_split * split, int input_id) { - if (split->backend_id != sched->transport.backend_id) { + if (!ggml_backend_sched_transport_ring_enabled(sched, split->backend_id)) { return false; } @@ -1805,43 +1831,63 @@ static bool ggml_backend_sched_input_is_staged(ggml_backend_sched_t sched, int s // sync_consumers must be false once the scheduler's backends may already be gone, which is the // case on the teardown path: llama_context and other owners outlive the scheduler only by // declaration order, and the backends it points at are not the scheduler's to keep alive. -static void ggml_backend_sched_transport_free_ring(ggml_backend_sched_t sched, bool sync_consumers) { - struct ggml_backend_sched_transport * tr = &sched->transport; +static void ggml_backend_sched_transport_free_ring(ggml_backend_sched_t sched, int backend_id, bool sync_consumers) { + struct ggml_backend_sched_transport_ring * r = &sched->transport.rings[backend_id]; - if (tr->buffer == NULL) { + if (r->buffer == NULL) { return; } // nothing may still be reading from or writing into the ring - if (tr->transfer) { - ggml_backend_synchronize(tr->transfer); + if (r->transfer) { + ggml_backend_synchronize(r->transfer); } if (sync_consumers) { - for (int i = 0; i < sched->n_backends; i++) { - ggml_backend_synchronize(sched->backends[i]); - } + ggml_backend_synchronize(sched->backends[backend_id]); } - ggml_backend_buffer_free(tr->buffer); - tr->buffer = NULL; - tr->slot_size = 0; + ggml_backend_buffer_free(r->buffer); + r->buffer = NULL; + r->slot_size = 0; for (int i = 0; i < GGML_SCHED_MAX_TRANSPORT_SLOTS; i++) { - tr->slots[i].release_armed = false; + r->slots[i].release_armed = false; } } -// The transfer backend and the slot events are created on demand, so that a scheduler which -// never gets to stage anything -- no eligible inputs, or no room within the budget -- does not -// carry a second device context for nothing. -static bool ggml_backend_sched_transport_ensure_backend(ggml_backend_sched_t sched) { +static void ggml_backend_sched_transport_release_ring(ggml_backend_sched_t sched, int backend_id, bool sync_consumers) { + struct ggml_backend_sched_transport_ring * r = &sched->transport.rings[backend_id]; + + ggml_backend_sched_transport_free_ring(sched, backend_id, sync_consumers); + + for (int i = 0; i < GGML_SCHED_MAX_TRANSPORT_SLOTS; i++) { + ggml_backend_event_free(r->slots[i].ready); + ggml_backend_event_free(r->slots[i].release); + r->slots[i].ready = NULL; + r->slots[i].release = NULL; + r->slots[i].release_armed = false; + } + + if (r->transfer) { + ggml_backend_free(r->transfer); + r->transfer = NULL; + } + + r->n_staged = 0; +} + +// The transfer backend and the slot events are created on demand, so that a backend which never +// gets to stage anything -- no eligible inputs, or no room within the budget -- does not carry a +// second device context for nothing. +static bool ggml_backend_sched_transport_ensure_backend(ggml_backend_sched_t sched, int backend_id) { struct ggml_backend_sched_transport * tr = &sched->transport; + struct ggml_backend_sched_transport_ring * r = &tr->rings[backend_id]; - if (tr->transfer != NULL) { + if (r->transfer != NULL) { return true; } - ggml_backend_dev_t dev = ggml_backend_get_device(sched->backends[tr->backend_id]); + ggml_backend_dev_t dev = ggml_backend_get_device(sched->backends[backend_id]); if (dev == NULL) { return false; } @@ -1853,34 +1899,39 @@ static bool ggml_backend_sched_transport_ensure_backend(ggml_backend_sched_t sch bool ok = true; for (int slot = 0; slot < tr->n_slots && ok; slot++) { - tr->slots[slot].ready = ggml_backend_event_new(dev); - tr->slots[slot].release = ggml_backend_event_new(dev); - ok = tr->slots[slot].ready != NULL && tr->slots[slot].release != NULL; + r->slots[slot].ready = ggml_backend_event_new(dev); + r->slots[slot].release = ggml_backend_event_new(dev); + ok = r->slots[slot].ready != NULL && r->slots[slot].release != NULL; } if (!ok) { for (int slot = 0; slot < tr->n_slots; slot++) { - ggml_backend_event_free(tr->slots[slot].ready); - ggml_backend_event_free(tr->slots[slot].release); - tr->slots[slot].ready = NULL; - tr->slots[slot].release = NULL; + ggml_backend_event_free(r->slots[slot].ready); + ggml_backend_event_free(r->slots[slot].release); + r->slots[slot].ready = NULL; + r->slots[slot].release = NULL; } ggml_backend_free(transfer); return false; } - tr->transfer = transfer; + r->transfer = transfer; return true; } -// Lay the ring out over the current split list and point the staged input copies at it. Called +// Lay the rings out over the current split list and point the staged input copies at them. Called // before the graph is allocated: ggml-alloc leaves a tensor that already has data alone, so the // staged copies are excluded from its reuse analysis instead of competing with it. static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { struct ggml_backend_sched_transport * tr = &sched->transport; - tr->n_staged = 0; - tr->scan_cursor = 0; + tr->n_staged = 0; + tr->n_rings_used = 0; + for (int i = 0; i < sched->n_backends; i++) { + tr->rings[i].n_staged = 0; + tr->rings[i].consumed = 0; + tr->rings[i].scan_cursor = 0; + } if (!ggml_backend_sched_transport_enabled(sched)) { return; @@ -1932,8 +1983,8 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { // A staged input copy may only be read by the split that owns it. The scheduler creates one // copy per (tensor, backend) rather than per split, so a later split can be pointed at the // same copy without it appearing in that split's input list -- and by then the ring may have - // recycled the slot for a different delivery. Splits whose copy has a reader further down the - // graph go back on the ordered path. + // recycled the slot. A view of the copy is excluded for the same reason: its address was + // resolved from the copy's own, so redirecting the copy afterwards would leave it behind. for (int i = 0; i < sched->n_splits; i++) { struct ggml_backend_sched_split * split = &sched->splits[i]; for (int j = 0; j < split->n_inputs; j++) { @@ -1943,40 +1994,37 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { const struct ggml_tensor * input_cpy = tensor_copy(split->inputs[j], split->backend_id, sched->cur_copy); - bool late_reader = false; - for (int k = i + 1; k < sched->n_splits && !late_reader; k++) { + bool disqualified = false; + for (int k = 0; k < sched->n_splits && !disqualified; k++) { const struct ggml_cgraph * g = &sched->splits[k].graph; - for (int n = 0; n < g->n_nodes && !late_reader; n++) { + for (int n = 0; n < g->n_nodes && !disqualified; n++) { + if (g->nodes[n]->view_src == input_cpy) { + disqualified = true; + break; + } + if (k <= i) { + continue; + } for (int sr = 0; sr < GGML_MAX_SRC; sr++) { if (g->nodes[n]->src[sr] == input_cpy) { - late_reader = true; + disqualified = true; break; } } } } - // a view of the copy would have had its address resolved from the copy's own, so - // redirecting the copy afterwards would leave the view pointing at the old block - bool viewed = false; - for (int k = 0; k < sched->n_splits && !viewed; k++) { - const struct ggml_cgraph * g = &sched->splits[k].graph; - for (int n = 0; n < g->n_nodes && !viewed; n++) { - if (g->nodes[n]->view_src == input_cpy) { - viewed = true; - } - } - } - - if (late_reader || viewed) { + if (disqualified) { tr->input_staged[tr->split_input_ofs[i] + j] = 0; } } } - size_t slot_size = 0; + // per-ring slot size and delivery order + size_t slot_size[GGML_SCHED_MAX_BACKENDS] = { 0 }; for (int i = 0; i < sched->n_splits; i++) { struct ggml_backend_sched_split * split = &sched->splits[i]; + const int bid = split->backend_id; tr->split_order[i] = -1; @@ -1985,125 +2033,170 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { if (!tr->input_staged[tr->split_input_ofs[i] + j]) { continue; } - need += GGML_PAD(ggml_nbytes(split->inputs[j]), tr->alignment); + need += GGML_PAD(ggml_nbytes(split->inputs[j]), tr->rings[bid].alignment); } if (need == 0) { continue; } - tr->split_order[i] = tr->n_staged++; - slot_size = std::max(slot_size, need); - } - - if (tr->n_staged == 0) { - return; - } - - ggml_backend_buffer_type_t buft = sched->bufts[tr->backend_id]; - const size_t ring_size = slot_size * tr->n_slots; - - // Checked before anything is allocated, and on every plan rather than only when the ring has - // to grow: a graph the ring already fits does not license keeping it once a later graph needs - // more than it may have. Declining has to cost nothing, so a ring in hand is given back. - // - // The point of a host-resident cache is to leave device memory alone, so the cap is absolute - // and not merely a fraction of what happens to be free. - // Latched, because a context only grows: the early prefill graphs have a small window and - // would fit, and allocating a ring for them only to give it back once the window outgrows the - // budget claims device memory that a host-resident cache is supposed to be leaving alone. - if (tr->budget > 0 && (ring_size > tr->budget || tr->over_budget)) { - tr->over_budget = true; - if (!tr->reported_no_room) { - GGML_LOG_WARN("%s: transport ring would need %zu MiB against a %zu MiB budget, " - "staying on the ordered path (raise --kv-pipeline-budget to spend more " - "device memory on it)\n", __func__, ring_size >> 20, tr->budget >> 20); - tr->reported_no_room = true; - } - ggml_backend_sched_transport_teardown(sched); - tr->n_staged = 0; - return; + tr->split_order[i] = tr->rings[bid].n_staged++; + slot_size[bid] = std::max(slot_size[bid], need); } - // nothing has been allocated for the transport until here - if (!ggml_backend_sched_transport_ensure_backend(sched)) { - tr->n_staged = 0; - return; - } + for (int bid = 0; bid < sched->n_backends; bid++) { + struct ggml_backend_sched_transport_ring * r = &tr->rings[bid]; + if (r->n_staged == 0) { + continue; + } - if (tr->buffer == NULL || tr->slot_size < slot_size) { - ggml_backend_sched_transport_free_ring(sched, true); - - // A slot has to hold one split's whole delivery, so the ring grows with the context and - // at some depth it stops fitting. The graph allocator has already reserved its buffers by - // now, but it may still have to grow one, and an allocation that succeeds here and - // starves that later is worse than not pipelining at all -- so leave it room rather than - // taking whatever is left. - ggml_backend_dev_t dev = ggml_backend_get_device(sched->backends[tr->backend_id]); - size_t dev_free = 0, dev_total = 0; - if (dev != NULL) { - ggml_backend_dev_memory(dev, &dev_free, &dev_total); - } - if (dev_free > 0 && ring_size + GGML_SCHED_TRANSPORT_HEADROOM > dev_free) { - if (!tr->reported_no_room) { - GGML_LOG_WARN("%s: transport ring would need %zu MiB and leave less than %u MiB " - "of the %zu MiB free on %s, staying on the ordered path\n", __func__, - ring_size >> 20, GGML_SCHED_TRANSPORT_HEADROOM >> 20, dev_free >> 20, - ggml_backend_buft_name(buft)); - tr->reported_no_room = true; + const size_t ring_size = slot_size[bid] * tr->n_slots; + + // Checked before anything is allocated, and on every plan rather than only when the ring + // has to grow. Latched, because a context only grows: the early prefill graphs have a + // small window and would fit, and allocating a ring for them only to give it back once + // the window outgrows the budget claims device memory that a host-resident cache is + // supposed to be leaving alone. + // + // The cap is per device. Each ring is a claim on its own card, and a second accelerator + // brings its own memory to spend. + if (tr->budget > 0 && (ring_size > tr->budget || r->over_budget)) { + r->over_budget = true; + if (!r->reported_no_room) { + GGML_LOG_WARN("%s: transport ring on %s would need %zu MiB against a %zu MiB budget, " + "staying on the ordered path (raise --kv-pipeline-budget to spend more " + "device memory on it)\n", __func__, ggml_backend_name(sched->backends[bid]), + ring_size >> 20, tr->budget >> 20); + r->reported_no_room = true; } - tr->n_staged = 0; - return; + ggml_backend_sched_transport_release_ring(sched, bid, true); + // un-stage this ring's inputs: they have no ring to live in + for (int i = 0; i < sched->n_splits; i++) { + if (sched->splits[i].backend_id != bid) { + continue; + } + tr->split_order[i] = -1; + for (int j = 0; j < sched->splits[i].n_inputs; j++) { + tr->input_staged[tr->split_input_ofs[i] + j] = 0; + } + } + continue; } - ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(buft, ring_size); - if (buffer == NULL) { - GGML_LOG_WARN("%s: failed to allocate %zu MiB for the transport ring, pipelining disabled\n", - __func__, ring_size >> 20); - tr->n_staged = 0; - return; + // nothing has been allocated for this ring until here + if (!ggml_backend_sched_transport_ensure_backend(sched, bid)) { + r->n_staged = 0; + continue; } - ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_COMPUTE); - tr->buffer = buffer; - tr->slot_size = slot_size; + if (r->buffer == NULL || r->slot_size < slot_size[bid]) { + ggml_backend_sched_transport_free_ring(sched, bid, true); - if (tr->debug > 0) { - GGML_LOG_INFO("%s: transport ring: %d slots x %zu KiB on %s\n", __func__, - tr->n_slots, slot_size >> 10, ggml_backend_buft_name(buft)); + ggml_backend_buffer_type_t buft = sched->bufts[bid]; + + // The ring is allocated after the graph allocator has reserved its buffers, so it must + // not take the room those buffers may still have to grow into. + ggml_backend_dev_t dev = ggml_backend_get_device(sched->backends[bid]); + size_t dev_free = 0, dev_total = 0; + if (dev != NULL) { + ggml_backend_dev_memory(dev, &dev_free, &dev_total); + } + if (dev_free > 0 && ring_size + GGML_SCHED_TRANSPORT_HEADROOM > dev_free) { + if (!r->reported_no_room) { + GGML_LOG_WARN("%s: transport ring on %s would need %zu MiB and leave less than " + "%u MiB of the %zu MiB free, staying on the ordered path\n", __func__, + ggml_backend_name(sched->backends[bid]), ring_size >> 20, + GGML_SCHED_TRANSPORT_HEADROOM >> 20, dev_free >> 20); + r->reported_no_room = true; + } + r->over_budget = true; + ggml_backend_sched_transport_release_ring(sched, bid, true); + for (int i = 0; i < sched->n_splits; i++) { + if (sched->splits[i].backend_id != bid) { + continue; + } + tr->split_order[i] = -1; + for (int j = 0; j < sched->splits[i].n_inputs; j++) { + tr->input_staged[tr->split_input_ofs[i] + j] = 0; + } + } + continue; + } + + ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(buft, ring_size); + if (buffer == NULL) { + GGML_LOG_WARN("%s: failed to allocate %zu MiB for the transport ring on %s, " + "pipelining disabled there\n", __func__, ring_size >> 20, + ggml_backend_name(sched->backends[bid])); + ggml_backend_sched_transport_release_ring(sched, bid, true); + for (int i = 0; i < sched->n_splits; i++) { + if (sched->splits[i].backend_id != bid) { + continue; + } + tr->split_order[i] = -1; + for (int j = 0; j < sched->splits[i].n_inputs; j++) { + tr->input_staged[tr->split_input_ofs[i] + j] = 0; + } + } + continue; + } + ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_COMPUTE); + + r->buffer = buffer; + r->slot_size = slot_size[bid]; + + if (tr->debug > 0) { + GGML_LOG_INFO("%s: transport ring on %s: %d slots x %zu KiB\n", __func__, + ggml_backend_name(sched->backends[bid]), tr->n_slots, slot_size[bid] >> 10); + } } + + tr->n_staged += r->n_staged; + tr->n_rings_used++; + } + + if (tr->n_staged == 0) { + return; } if (tr->debug > 0) { - size_t total = 0; - size_t early = 0; - const char * src_buft = "?"; - for (int i = 0; i < sched->n_splits; i++) { - struct ggml_backend_sched_split * split = &sched->splits[i]; - for (int j = 0; j < split->n_inputs; j++) { - if (!ggml_backend_sched_input_is_staged(sched, i, j)) { + for (int bid = 0; bid < sched->n_backends; bid++) { + if (tr->rings[bid].n_staged == 0) { + continue; + } + size_t total = 0, early = 0; + const char * src_buft = "?"; + for (int i = 0; i < sched->n_splits; i++) { + if (sched->splits[i].backend_id != bid || tr->split_order[i] < 0) { continue; } - struct ggml_tensor * input = split->inputs[j]; - ggml_backend_buffer_t buf = input->view_src ? input->view_src->buffer : input->buffer; - src_buft = ggml_backend_buft_name(buf->buft); - total += ggml_nbytes(input); - early += ggml_backend_sched_input_stable_prefix(input); + struct ggml_backend_sched_split * split = &sched->splits[i]; + for (int j = 0; j < split->n_inputs; j++) { + if (!ggml_backend_sched_input_is_staged(sched, i, j)) { + continue; + } + struct ggml_tensor * input = split->inputs[j]; + ggml_backend_buffer_t buf = input->view_src ? input->view_src->buffer : input->buffer; + src_buft = ggml_backend_buft_name(buf->buft); + total += ggml_nbytes(input); + early += ggml_backend_sched_input_stable_prefix(input); + } } + GGML_LOG_INFO("%s: %s: %d/%d splits staged, %zu KiB per graph, %zu KiB of it early, source %s\n", + __func__, ggml_backend_name(sched->backends[bid]), tr->rings[bid].n_staged, + sched->n_splits, total >> 10, early >> 10, src_buft); } - GGML_LOG_INFO("%s: %d/%d splits staged, %zu KiB per graph, %zu KiB of it early, source %s\n", - __func__, tr->n_staged, sched->n_splits, total >> 10, early >> 10, src_buft); } - char * const ring = (char *) ggml_backend_buffer_get_base(tr->buffer); - for (int i = 0; i < sched->n_splits; i++) { if (tr->split_order[i] < 0) { continue; } struct ggml_backend_sched_split * split = &sched->splits[i]; - char * slot = ring + (size_t)(tr->split_order[i] % tr->n_slots) * tr->slot_size; + struct ggml_backend_sched_transport_ring * r = &tr->rings[split->backend_id]; + char * const ring = (char *) ggml_backend_buffer_get_base(r->buffer); + char * slot = ring + (size_t)(tr->split_order[i] % tr->n_slots) * r->slot_size; size_t offset = 0; for (int j = 0; j < split->n_inputs; j++) { @@ -2112,39 +2205,43 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { } struct ggml_tensor * input_cpy = tensor_copy(split->inputs[j], split->backend_id, sched->cur_copy); input_cpy->data = slot + offset; - input_cpy->buffer = tr->buffer; - offset += GGML_PAD(ggml_nbytes(split->inputs[j]), tr->alignment); + input_cpy->buffer = r->buffer; + offset += GGML_PAD(ggml_nbytes(split->inputs[j]), r->alignment); } - GGML_ASSERT(offset <= tr->slot_size); + GGML_ASSERT(offset <= r->slot_size); } } -// Issue the stable prefix of every staged split within the look-ahead of what has already been -// enqueued. The ring carries GGML_SCHED_TRANSPORT_MARGIN slots more than the look-ahead, so the -// slot a delivery writes into belongs to a split that is several readers behind the one just -// enqueued, and recycling it does not put the transfer stream back in lock-step with the -// consumer. -static void ggml_backend_sched_transport_prefetch(ggml_backend_sched_t sched, int consumed) { +// Issue the stable prefix of every staged split on this ring that is within the look-ahead of what +// has already been enqueued on it. The ring carries GGML_SCHED_TRANSPORT_MARGIN slots more than the +// look-ahead, so the slot a delivery writes into belongs to a split that is several readers behind +// the one just enqueued, and recycling it does not put the transfer stream back in lock-step with +// the consumer. Each ring walks the split list on its own cursor: one device saturating its +// look-ahead must not stop another device from running ahead on its own. +static void ggml_backend_sched_transport_prefetch(ggml_backend_sched_t sched, int backend_id) { struct ggml_backend_sched_transport * tr = &sched->transport; + struct ggml_backend_sched_transport_ring * r = &tr->rings[backend_id]; - const int lookahead = tr->depth; + if (r->n_staged == 0) { + return; + } - for (int i = tr->scan_cursor; i < sched->n_splits; i++) { - if (tr->split_order[i] < 0) { + for (int i = r->scan_cursor; i < sched->n_splits; i++) { + if (tr->split_order[i] < 0 || sched->splits[i].backend_id != backend_id) { continue; } - if (tr->split_order[i] > consumed + lookahead) { + if (tr->split_order[i] > r->consumed + tr->depth) { return; } struct ggml_backend_sched_split * split = &sched->splits[i]; - struct ggml_backend_sched_transport_slot * slot = &tr->slots[tr->split_order[i] % tr->n_slots]; + struct ggml_backend_sched_transport_slot * slot = &r->slots[tr->split_order[i] % tr->n_slots]; // the previous occupant of this slot must be read before the slot is overwritten. This is // ordered stream to stream rather than through the host: blocking the host here would // hold back the work it has not enqueued yet, which is what the margin exists to avoid. if (slot->release_armed) { - ggml_backend_event_wait(tr->transfer, slot->release); + ggml_backend_event_wait(r->transfer, slot->release); slot->release_armed = false; } @@ -2165,7 +2262,7 @@ static void ggml_backend_sched_transport_prefetch(ggml_backend_sched_t sched, in struct ggml_tensor * input_cpy = tensor_copy(input, split->backend_id, sched->cur_copy); GGML_ASSERT(input->data != NULL && input_cpy->data != NULL); const int64_t t0 = tr->debug >= 2 ? ggml_time_us() : 0; - ggml_backend_tensor_set_async(tr->transfer, input_cpy, input->data, 0, prefix); + ggml_backend_tensor_set_async(r->transfer, input_cpy, input->data, 0, prefix); if (tr->debug >= 2) { tr->t_issue_us += ggml_time_us() - t0; } @@ -2176,9 +2273,9 @@ static void ggml_backend_sched_transport_prefetch(ggml_backend_sched_t sched, in // and by then the deliveries for the splits after this one are already queued behind it. // Waiting on an event recorded after those would make the consumer wait for the whole // look-ahead, which is the ordered path again with extra steps. - ggml_backend_event_record(slot->ready, tr->transfer); + ggml_backend_event_record(slot->ready, r->transfer); - tr->scan_cursor = i + 1; + r->scan_cursor = i + 1; } } @@ -2224,8 +2321,10 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { // the re-allocation may cause the split inputs to be moved to a different address // synchronize without ggml_backend_sched_synchronize to avoid changing cur_copy - if (sched->transport.transfer) { - ggml_backend_synchronize(sched->transport.transfer); + for (int i = 0; i < sched->n_backends; i++) { + if (sched->transport.rings[i].transfer) { + ggml_backend_synchronize(sched->transport.rings[i].transfer); + } } for (int i = 0; i < sched->n_backends; i++) { ggml_backend_synchronize(sched->backends[i]); @@ -2258,17 +2357,20 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } const bool staged = tr->n_staged > 0 && tr->plan_n_splits == sched->n_splits && tr->plan_n_inputs == n_inputs_now; - int consumed = 0; // staged splits whose reader has been enqueued // the plan outlives the graph it was made for -- a reused graph keeps it -- so the cursor // that walks it has to start over on every evaluation, not only when the plan is rebuilt - tr->scan_cursor = 0; - - // prime the ring before the first consumer runs; from here on deliveries are issued only + // prime every ring before the first consumer runs; from here on deliveries are issued only // after a split has been enqueued, never before, so that recycling a slot can never hold // back work the consumer could already be running if (staged) { - ggml_backend_sched_transport_prefetch(sched, consumed); + for (int i = 0; i < sched->n_backends; i++) { + tr->rings[i].consumed = 0; + tr->rings[i].scan_cursor = 0; + } + for (int i = 0; i < sched->n_backends; i++) { + ggml_backend_sched_transport_prefetch(sched, i); + } } const bool timing = tr->debug >= 2; @@ -2433,7 +2535,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s // order the consumer behind this split's early deliveries struct ggml_backend_sched_transport_slot * slot = NULL; if (staged && tr->split_order[split_id] >= 0) { - slot = &tr->slots[tr->split_order[split_id] % tr->n_slots]; + slot = &tr->rings[split_backend_id].slots[tr->split_order[split_id] % tr->n_slots]; ggml_backend_event_wait(split_backend, slot->ready); } @@ -2483,13 +2585,13 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s if (slot != NULL) { ggml_backend_event_record(slot->release, split_backend); slot->release_armed = true; - consumed++; + tr->rings[split_backend_id].consumed++; tr->n_deliveries++; // run the transfer stream ahead of the consumer: with this split's kernels already // enqueued, the deliveries for the next few staged splits can go out even if // recycling their slot has to wait for a reader that is now running - ggml_backend_sched_transport_prefetch(sched, consumed); + ggml_backend_sched_transport_prefetch(sched, split_backend_id); } // record the event of this copy @@ -2592,8 +2694,7 @@ ggml_backend_sched_t ggml_backend_sched_new( sched->galloc = ggml_gallocr_new_n(sched->bufts, n_backends); sched->op_offload = op_offload; - sched->transport.backend_id = -1; - sched->transport.budget = GGML_SCHED_TRANSPORT_BUDGET; + sched->transport.budget = GGML_SCHED_TRANSPORT_BUDGET; { const char * GGML_SCHED_TRANSPORT_DEBUG = getenv("GGML_SCHED_TRANSPORT_DEBUG"); sched->transport.debug = GGML_SCHED_TRANSPORT_DEBUG ? atoi(GGML_SCHED_TRANSPORT_DEBUG) : 0; @@ -2605,26 +2706,13 @@ ggml_backend_sched_t ggml_backend_sched_new( } static void ggml_backend_sched_transport_teardown(ggml_backend_sched_t sched) { - struct ggml_backend_sched_transport * tr = &sched->transport; - - ggml_backend_sched_transport_free_ring(sched, false); - - for (int i = 0; i < GGML_SCHED_MAX_TRANSPORT_SLOTS; i++) { - ggml_backend_event_free(tr->slots[i].ready); - ggml_backend_event_free(tr->slots[i].release); - tr->slots[i].ready = NULL; - tr->slots[i].release = NULL; - tr->slots[i].release_armed = false; - } - - if (tr->transfer) { - ggml_backend_free(tr->transfer); - tr->transfer = NULL; + for (int i = 0; i < sched->n_backends; i++) { + ggml_backend_sched_transport_release_ring(sched, i, false); } - - tr->n_staged = 0; + sched->transport.n_staged = 0; } + void ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, int depth) { GGML_ASSERT(sched); @@ -2639,21 +2727,25 @@ void ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, } struct ggml_backend_sched_transport * tr = &sched->transport; - if (tr->depth == depth && (depth < 1 || tr->backend_id >= 0)) { - return; - } ggml_backend_sched_transport_teardown(sched); - tr->depth = depth; + for (int i = 0; i < sched->n_backends; i++) { + tr->rings[i].eligible = false; + tr->rings[i].over_budget = false; + tr->rings[i].reported_no_room = false; + } + + tr->depth = depth; + tr->n_slots = depth + GGML_SCHED_TRANSPORT_MARGIN; if (depth < 1) { return; } - const int n_slots = depth + GGML_SCHED_TRANSPORT_MARGIN; - - // one ring, on the first backend that can transfer asynchronously and order streams with - // events. Additional accelerators keep the ordered path. + // Every backend that can transfer asynchronously and order streams with events gets its own + // ring. A layer-split model puts splits on each device, and a device left on the ordered path + // would pay copy + compute in series while the others do not. + int n_eligible = 0; for (int i = 0; i < sched->n_backends; i++) { ggml_backend_t backend = sched->backends[i]; if (backend->iface.set_tensor_async == NULL || @@ -2676,22 +2768,17 @@ void ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, continue; } - // The transfer backend is a second backend on the same device, which gives the deliveries - // a stream of their own -- and costs a device context of its own. It is created on the - // first graph that actually stages something, not here: a scheduler that never has room - // for a ring must not pay for a context it will not use. - tr->backend_id = i; - tr->n_slots = n_slots; - tr->alignment = std::max(ggml_backend_buft_get_alignment(sched->bufts[i]), 128); + tr->rings[i].eligible = true; + tr->rings[i].alignment = std::max(ggml_backend_buft_get_alignment(sched->bufts[i]), 128); + n_eligible++; if (tr->debug > 0) { GGML_LOG_INFO("%s: pipelined host transport selected %s, %d splits ahead, %d slots\n", - __func__, ggml_backend_name(backend), depth, n_slots); + __func__, ggml_backend_name(backend), depth, tr->n_slots); } - return; } - if (tr->debug > 0) { + if (n_eligible == 0 && tr->debug > 0) { GGML_LOG_INFO("%s: no backend supports pipelined host transport, staying on the ordered path\n", __func__); } } @@ -2705,11 +2792,13 @@ void ggml_backend_sched_set_transport_pipeline_budget(ggml_backend_sched_t sched } if (sched->transport.budget != bytes) { - sched->transport.budget = bytes; - sched->transport.reported_no_room = false; - sched->transport.over_budget = false; - // the ring in hand may no longer be allowed - ggml_backend_sched_transport_free_ring(sched, true); + sched->transport.budget = bytes; + for (int i = 0; i < sched->n_backends; i++) { + sched->transport.rings[i].reported_no_room = false; + sched->transport.rings[i].over_budget = false; + // a ring in hand may no longer be allowed + ggml_backend_sched_transport_free_ring(sched, i, true); + } } } @@ -2865,8 +2954,10 @@ enum ggml_status ggml_backend_sched_graph_compute_async(ggml_backend_sched_t sch void ggml_backend_sched_synchronize(ggml_backend_sched_t sched) { GGML_ASSERT(sched); - if (sched->transport.transfer) { - ggml_backend_synchronize(sched->transport.transfer); + for (int i = 0; i < sched->n_backends; i++) { + if (sched->transport.rings[i].transfer) { + ggml_backend_synchronize(sched->transport.rings[i].transfer); + } } for (int i = 0; i < sched->n_backends; i++) { ggml_backend_synchronize(sched->backends[i]);