From 89441d3bb98175f911790c47c7646c8aa12d1cc2 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Wed, 26 Aug 2026 03:08:08 +0200 Subject: [PATCH 01/32] sched: pipeline the delivery of a host-resident KV cache With --no-kv-offload the attention history lives in host RAM and reaches the accelerator on every decode token. The scheduler issued that transfer on the consumer's own stream immediately before the kernels that read it, so a token cost copy + compute in series. The bytes and the attention operations are unchanged; only the point at which the transfer is issued moves. Greedy output is byte-identical to the ordered path -- verified against a build without these changes, at every look-ahead tested, single GPU and layer-split across two. Three pieces, each load-bearing: - ggml_tensor::stable_prefix records how many leading bytes of a tensor's storage the graph about to run will not write. The KV window is not stable for a whole graph -- a CPU split writes this ubatch's rows into it between one layer's attention and the next -- but everything below the lowest written row is, and at decode depth that is essentially all of it. llama_kv_cache sets it from apply_ubatch(), before the graph is built and allocated, so the plan and the deliveries are decided against the same write position even when the graph is reused; build_graph_shift() clears it. - A staging ring the graph allocator cannot reach. ggml-alloc may recycle a graph-owned input copy after its last graph-level consumer while a look-ahead transfer is still in flight. The scheduler allocates the ring itself and points the staged copies at it before allocation; a ready/release event pair per slot carries the handover in each direction. Every eligible accelerator gets its own ring, cursor and budget, so a layer-split model pipelines on each device and a device with no room falls back alone. - A look-ahead that stays clear of the ring's tail. A delivery L splits ahead recycles the slot of the split L - n_slots back, so n_slots == L + 1 recycles the split just enqueued and still running. The ring keeps two slots of margin, deliveries are issued after a split is enqueued rather than before, and slot recycling is ordered stream to stream rather than through the host. Each of those three alone costs the entire gain while still producing correct output. --kv-pipeline-depth N, default 1, 0 restores the ordered path exactly. It only engages where a host-resident cache produces the deliveries. Because a host-resident cache exists to keep device memory free, the staging is capped outright by --kv-pipeline-budget (default 128 MiB per device) rather than by a fraction of what happens to be free. A ring is (N + 2) slots of one attention layer's K and V over the whole context, so it grows with the context: 27 MiB at 4k, 213 MiB at 32k, 1.7 GiB at 256k. Past the cap the scheduler declines and keeps the ordered path, and declining costs nothing -- the check runs before anything is allocated, the decision is latched because a context only grows, and the transfer backend is created lazily and released with the ring. Single GPU (RTX 4070, Qwen3.8-27B-UD-IQ2_M, -nkvo --kv-cpu-pinned, q8_0 K/V), A/B/A/B with reversed arm order: depth ordered pipelined gain 4,096 31.7324, 31.7363 37.0889, 37.0741 +16.9% 16,384 19.6765, 19.6854 31.5352, 31.5807 +60.4% 32,768 13.0264, 13.0254 15.5325, 15.5329 +19.3% (needs a raised budget) Server decode behind an 18,422-token prompt: 18.468 -> 30.685 t/s, +66.2%. Two GPUs (RTX 4070 + RTX 3060, Qwen3.8-27B-UD-Q5_K_M, -sm layer), both rings engaged: 13.06 -> 18.05 t/s at 4,096 and 6.86 -> 9.75 t/s at 16,384. The gain narrows with depth because compute is a shrinking share of the token, so there is less to hide the copy behind. That is arithmetic, not an implementation limit, and more look-ahead makes it worse rather than better. Tensor parallelism keeps the ordered path: the scheduler sees one meta backend there and the ring is a byte arena, while a meta buffer places tensors as per-device slices rather than at offsets. It declines rather than staging into something it cannot address. docs/kv-transport-pipelining.md carries the design, the numbers and the limits; docs/repro/ carries the scripts that produced them. Assisted-by: Claude Opus 5 --- common/arg.cpp | 28 + common/common.cpp | 2 + common/common.h | 2 + docs/kv-transport-pipelining.md | 279 +++++++ docs/repro/r4-kv-pipeline-ab.sh | 31 + docs/repro/r4-kv-pipeline-context-sweep.sh | 42 ++ docs/repro/r4-kv-pipeline-exact.py | 70 ++ docs/repro/r4-kv-pipeline-exact.sh | 35 + ggml/include/ggml-backend.h | 29 + ggml/include/ggml.h | 18 +- ggml/src/ggml-backend-meta.cpp | 5 +- ggml/src/ggml-backend.cpp | 831 +++++++++++++++++++++ ggml/src/ggml.c | 13 +- include/llama.h | 9 + src/llama-context.cpp | 9 + src/llama-cparams.h | 2 + src/llama-kv-cache.cpp | 48 ++ src/llama-kv-cache.h | 7 + 18 files changed, 1457 insertions(+), 3 deletions(-) create mode 100644 docs/kv-transport-pipelining.md create mode 100755 docs/repro/r4-kv-pipeline-ab.sh create mode 100755 docs/repro/r4-kv-pipeline-context-sweep.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 6cc1e412465..9cdcf11b2b2 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2421,6 +2421,34 @@ 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( + {"--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 6aaec977e58..28ac9531e38 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1744,6 +1744,8 @@ 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.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.kv_gpu_layers = (uint32_t) std::max(0, params.kv_gpu_layers); cparams.phase_aware_workspace = params.phase_aware_workspace; diff --git a/common/common.h b/common/common.h index 8ca190c2df7..109b9269710 100644 --- a/common/common.h +++ b/common/common.h @@ -584,6 +584,8 @@ struct common_params { int32_t kv_gpu_layers = 0; // with no_kv_offload, keep this many attention KV layers device-resident 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/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md new file mode 100644 index 00000000000..b01bc4ae7c0 --- /dev/null +++ b/docs/kv-transport-pipelining.md @@ -0,0 +1,279 @@ +# 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. + +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. + +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 +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 (+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. + +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` +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 | 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% | + +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%**. + +### 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`): + +| | 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, 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, 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 + 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 00000000000..27e8b0ed302 --- /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-context-sweep.sh b/docs/repro/r4-kv-pipeline-context-sweep.sh new file mode 100755 index 00000000000..b8a2b91bf72 --- /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 new file mode 100644 index 00000000000..0a0800f3f79 --- /dev/null +++ b/docs/repro/r4-kv-pipeline-exact.py @@ -0,0 +1,70 @@ +# 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 = sys.argv[1] +LENGTHS = [int(x) for x in sys.argv[2].split(",")] # approximate prefill tokens + +# 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"), +} + +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() + 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=14400) as r: + d = json.load(r) + except Exception as e: + 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 + text = (m.get("reasoning_content") or "") + "\x00" + (m.get("content") or "") + t = d.get("timings", {}) + 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 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 new file mode 100755 index 00000000000..6ab8d7b6b61 --- /dev/null +++ b/docs/repro/r4-kv-pipeline-exact.sh @@ -0,0 +1,35 @@ +#!/bin/bash +# 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 [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}" +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 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 "$CTX" --parallel 1 \ + --host 127.0.0.1 --port "$PORT" --no-warmup > "$LOG" 2>&1 & + SRV=$! + 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" "$LENGTHS" || 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 a663d86da63..924db4eec61 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -323,6 +323,35 @@ extern "C" { GGML_API void ggml_backend_sched_get_buffer_state(ggml_backend_sched_t sched, uint64_t * generation, uint64_t * shrink_generation); GGML_API void ggml_backend_sched_request_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); + + // 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); + // 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 56f9740cd60..66f7618c541 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -701,11 +701,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 fe58ea3bb7a..60332705ec4 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -426,7 +426,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 5a26ef147be..0316ebe0680 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 @@ -761,6 +762,112 @@ static bool ggml_is_view_op(enum ggml_op op) { #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 + +// 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 + +// 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. +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 +}; + +// 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 +// 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 per ring: depth + GGML_SCHED_TRANSPORT_MARGIN + size_t budget; // hard cap on each ring, in bytes + + 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 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; // over all rings, so that execution can skip the machinery entirely + int n_rings_used; + + // 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; + 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; @@ -820,6 +927,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] @@ -1539,6 +1649,538 @@ 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_ring_enabled(ggml_backend_sched_t sched, int backend_id) { + const struct ggml_backend_sched_transport * tr = &sched->transport; + 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 +// 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 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 +// 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 (!ggml_backend_sched_transport_ring_enabled(sched, split->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; + } + + 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, int backend_id, bool sync_consumers) { + struct ggml_backend_sched_transport_ring * r = &sched->transport.rings[backend_id]; + + if (r->buffer == NULL) { + return; + } + + // nothing may still be reading from or writing into the ring + if (r->transfer) { + ggml_backend_synchronize(r->transfer); + } + if (sync_consumers) { + ggml_backend_synchronize(sched->backends[backend_id]); + } + + ggml_backend_buffer_free(r->buffer); + r->buffer = NULL; + r->slot_size = 0; + + for (int i = 0; i < GGML_SCHED_MAX_TRANSPORT_SLOTS; i++) { + r->slots[i].release_armed = false; + } +} + +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 (r->transfer != NULL) { + return true; + } + + ggml_backend_dev_t dev = ggml_backend_get_device(sched->backends[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++) { + 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(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; + } + + r->transfer = transfer; + return true; +} + +// 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->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; + } + + 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; + + 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. 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++) { + 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 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 && !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) { + disqualified = true; + break; + } + } + } + } + + if (disqualified) { + tr->input_staged[tr->split_input_ofs[i] + j] = 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; + + size_t need = 0; + for (int j = 0; j < split->n_inputs; j++) { + if (!tr->input_staged[tr->split_input_ofs[i] + j]) { + continue; + } + need += GGML_PAD(ggml_nbytes(split->inputs[j]), tr->rings[bid].alignment); + } + + if (need == 0) { + continue; + } + + tr->split_order[i] = tr->rings[bid].n_staged++; + slot_size[bid] = std::max(slot_size[bid], need); + } + + 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; + } + + 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; + } + 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; + } + + // nothing has been allocated for this ring until here + if (!ggml_backend_sched_transport_ensure_backend(sched, bid)) { + r->n_staged = 0; + continue; + } + + if (r->buffer == NULL || r->slot_size < slot_size[bid]) { + ggml_backend_sched_transport_free_ring(sched, bid, true); + + 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) { + 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_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); + } + } + + 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]; + 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++) { + 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 = r->buffer; + offset += GGML_PAD(ggml_nbytes(split->inputs[j]), r->alignment); + } + GGML_ASSERT(offset <= r->slot_size); + } +} + +// 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]; + + if (r->n_staged == 0) { + return; + } + + 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] > r->consumed + tr->depth) { + return; + } + + struct ggml_backend_sched_split * split = &sched->splits[i]; + 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(r->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(r->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, r->transfer); + + r->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++) { @@ -1558,6 +2200,10 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { } } + // lay out the transport rings and point the staged input copies at them 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 @@ -1577,6 +2223,11 @@ 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 + 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]); } @@ -1604,6 +2255,30 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s int prev_backend_id = -1; + 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; + + // 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. The cursors start over on every evaluation + // because the plan outlives the graph it was made for. + if (staged) { + 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); + } + } + 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; @@ -1625,6 +2300,22 @@ 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. + 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) { @@ -1742,6 +2433,13 @@ 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->rings[split_backend_id].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) { @@ -1781,6 +2479,19 @@ 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; + tr->rings[split_backend_id].consumed++; + tr->n_deliveries++; + + // with this split's kernels already enqueued, the deliveries for the next staged + // splits can go out even if recycling their slot waits for a reader that is running + ggml_backend_sched_transport_prefetch(sched, split_backend_id); + } + // record the event of this split if (sched->events[split_backend_id][sched->cur_copy] != NULL) { ggml_backend_event_record(sched->events[split_backend_id][sched->cur_copy], split_backend); @@ -1857,6 +2568,12 @@ ggml_backend_sched_t ggml_backend_sched_new( } sched->galloc = ggml_gallocr_new_n(sched->bufts, n_backends); + + 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; + } sched->op_offload = op_offload; ggml_backend_sched_reset(sched); @@ -1864,10 +2581,119 @@ ggml_backend_sched_t ggml_backend_sched_new( return sched; } +static void ggml_backend_sched_transport_teardown(ggml_backend_sched_t sched) { + for (int i = 0; i < sched->n_backends; i++) { + ggml_backend_sched_transport_release_ring(sched, i, false); + } + sched->transport.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; + + ggml_backend_sched_transport_teardown(sched); + 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; + } + + // 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 || + 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; + } + + 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, tr->n_slots); + } + } + + if (n_eligible == 0 && tr->debug > 0) { + GGML_LOG_INFO("%s: no backend supports pipelined host transport, staying on the ordered path\n", __func__); + } +} + +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; + 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); + } + } +} + +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]); @@ -1996,6 +2822,11 @@ 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); + 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]); } diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 708b16edd47..6f3ee0f5a74 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1323,6 +1323,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); @@ -1819,7 +1830,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 f4aa9a9f294..e90a3cc98cb 100644 --- a/include/llama.h +++ b/include/llama.h @@ -412,6 +412,15 @@ extern "C" { // A source/target/parent context that can share results or llama_memory. struct llama_context * ctx_other; + + 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. + 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 c11a68a34d2..4ff80425ec6 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -145,6 +145,8 @@ llama_context::llama_context( cparams.kv_cpu_pinned = params.kv_cpu_pinned; cparams.recurrent_state_offload = params.recurrent_state_offload; cparams.offload_attn_compute = params.offload_kqv || (params.op_offload && params.kv_cpu_pinned); + 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.phase_aware_workspace = params.phase_aware_workspace; cparams.live_context_workspace = params.live_context_workspace; @@ -869,6 +871,11 @@ 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)); cparams.flash_attn_causal_prefix_supported = llama_sched_supports_flash_attn_causal_prefix(sched.get()); + // 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_resizable) { sched_buffers_shared = sched_buffer_owner != nullptr && sched_buffer_owner->get_sched() != nullptr && ggml_backend_sched_set_resizable(sched.get(), sched_buffer_owner->get_sched()); @@ -3988,6 +3995,8 @@ llama_context_params llama_context_default_params() { /*.sampler =*/ nullptr, /*.n_sampler =*/ 0, /*.ctx_other =*/ nullptr, + /*.kv_pipeline_depth =*/ 1, + /*.kv_pipeline_budget_mib =*/ 128, }; return result; diff --git a/src/llama-cparams.h b/src/llama-cparams.h index a51edaa6acf..0555d549e6d 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -60,6 +60,8 @@ struct llama_cparams { bool recurrent_state_offload; 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 diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 397cf7f9252..3ca115f2711 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1201,6 +1201,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]; @@ -1620,7 +1624,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()); @@ -2190,6 +2235,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 5b23b591a35..33ac39caf13 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -225,6 +225,13 @@ class llama_kv_cache : public llama_memory_i { bool is_reserve, const slot_info * sinfo = nullptr) 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; From 5b190526e8ed77cd808fd0905f6badd3248d25b5 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Wed, 26 Aug 2026 03:14:59 +0200 Subject: [PATCH 02/32] docs: record what tensor parallelism still needs, and fix the repro scripts llama-bench does not expose --kv-cpu-pinned or --recurrent-state-offload the way llama-server does, so two of the reproduction scripts passed flags the binary rejects. They now probe --help and pass only what it takes. The feature doc gains what is actually left: why -sm tensor keeps the ordered path (no events in the meta layer, and a ring that is a byte arena while a meta buffer places tensors as per-device slices), the correctness problem underneath it that is not this feature's, and the rest of the open list -- the transient device-memory peak, the --kv-gpu-layers comparison, and the exactness harness's dependence on baseline determinism it does not have. The decline for a backend that cannot record events now says so by name rather than falling into the generic "no backend supports" line, because the backend it catches is the meta backend and the next person to look will want to know that. Assisted-by: Claude Opus 5 --- docs/kv-transport-pipelining.md | 80 ++++++++++++++++++++-- docs/repro/r4-kv-pipeline-ab.sh | 10 ++- docs/repro/r4-kv-pipeline-context-sweep.sh | 10 ++- ggml/src/ggml-backend.cpp | 8 +++ 4 files changed, 99 insertions(+), 9 deletions(-) diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index b01bc4ae7c0..73e0a330675 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -268,12 +268,82 @@ 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, 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. +- **One ring per accelerator.** A layer-split model pipelines on every device + that qualifies; a device with no room within the budget falls back to the + ordered path on its own without disabling the others. +- **Tensor parallelism keeps the ordered path.** See + [Tensor parallelism](#tensor-parallelism). - 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`. + +## Tensor parallelism + +`-sm tensor` is not pipelined. The scheduler sees a single *meta* backend there, +and two separate things stand in the way: + +1. **No events in the meta layer.** `ggml_backend_meta_i.event_record` and + `.event_wait` are null, and so are `event_new` / `event_free` / + `event_synchronize` on the meta device. Pipelining is built on ordering a + transfer stream against the consumer with events, so the eligibility check + rejects the backend and the ordered path runs. Requesting a look-ahead under + `-sm tensor` costs nothing and changes nothing: measured 21.85 t/s at depth 1 + against 21.87 at depth 0, with identical output. +2. **The ring is a byte arena.** It is allocated once and the staged input copies + are pointed into it at fixed offsets. A meta buffer has no flat base -- + `ggml_backend_meta_buffer_get_base` returns a placeholder -- and a tensor in + one is not placed at an offset but built as a set of per-device tensors by the + buffer's `init_tensor`, from a whole `ggml_context` allocated at once. The ring + would have to become slot *tensors* rather than offsets. + +Both sit behind a correctness problem that is not this feature's: +**`-sm tensor` together with `--no-kv-offload` currently produces wrong output.** +On one build and one prompt, `-sm layer --no-kv-offload` and `-sm tensor` with a +device-resident cache agree exactly, while `-sm tensor --no-kv-offload` differs. +It does not crash or warn; it generates fluent, different text. + +The cause is the GQA head mapping. Tensor parallelism splits attention by head, +but a host-resident cache is one undivided tensor, so the scheduler's copy of it +is classified `MIRRORED` and the whole window goes to every device. With 24 query +heads split 12/12 and 4 KV heads mirrored, the kernel derives the GQA ratio from +the tensors it is handed -- 12/4 = 3 rather than 6 -- and the second device's +queries, renumbered from 0, read the first device's keys. With an uneven split the +same fault surfaces as a crash instead: +`GGML_ASSERT(Q->ne[2] % K->ne[2] == 0)`, because 24 heads split 13/11 is not +divisible by 4. + +Head-splitting the copy rather than mirroring it fixes it. That was prototyped +and reproduced the layer-split output byte for byte, and needs four coordinated +changes: classify the scheduler's copy at all (it is a leaf in a compute buffer, +so it never reaches the device's split-state callback), use the head axis for the +permuted `[head_dim, n_kv, n_head_kv, 1]` shape rather than the cache tensor's own +axis, express the granularity in heads aligned to the query split divided by the +GQA ratio, and add a strided write because the heads are interleaved within each +row rather than laid out end to end. + +## Future work + +- Fix `-sm tensor` with `--no-kv-offload` (above). Until then it should not be + used: it is wrong rather than slow. +- Events on the meta backend and device, so a transfer stream can be ordered + against a tensor-parallel consumer at all. +- A ring that can live in a meta buffer, as slot tensors rather than offsets, so + tensor parallelism can be pipelined once it is correct. +- A transfer-only meta backend that does not stand up a second collective + communicator: `ggml_backend_dev_init` on a meta device runs the whole meta + context constructor, which calls `ggml_backend_comm_init` across every device. +- Remove the transient device-memory peak. The budget is applied per graph, so a + context that grows past it still allocates a ring for the small windows of early + prefill and releases it once the window outgrows the budget: +112 MiB at 32,768 + against +0 in steady state. Deciding against the context's final size needs KV + geometry the scheduler does not have. +- Compare the ring against `--kv-gpu-layers` at depth. At 131,072 the ring's + 818 MiB is about three attention layers' K and V; making three of sixteen + device-resident would remove about 19% of the host-to-device traffic against the + 9.1% the ring buys there. Unmeasured, and it moves with layer count and card. +- Give the exactness harness a per-task nonce. Two long-prompt tasks proved + non-deterministic in the baseline -- a second control run reproduced this + branch's hashes rather than its own -- because the server restores a similar + cached prefix. Until that is pinned down the harness is a weaker gate than it + looks. diff --git a/docs/repro/r4-kv-pipeline-ab.sh b/docs/repro/r4-kv-pipeline-ab.sh index 27e8b0ed302..9d8bb83a393 100755 --- a/docs/repro/r4-kv-pipeline-ab.sh +++ b/docs/repro/r4-kv-pipeline-ab.sh @@ -9,9 +9,15 @@ BUILD="${LLAMA_KV_BUILD:-build}" PIN="${LLAMA_KV_TASKSET:-0,2,4}" LOCK=/tmp/beellama-single-gpu.lock +# llama-bench does not expose every host-KV option that llama-server does; pass only what it takes +BENCH_KV_OPTS="" +for opt in --kv-cpu-pinned --recurrent-state-offload; do + "$BUILD/bin/llama-bench" --help 2>&1 | grep -q -- "$opt" && BENCH_KV_OPTS="$BENCH_KV_OPTS $opt" +done + 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 \ + -ngl 99 -sm none -mg 0 -t 3 -nkvo 1 $BENCH_KV_OPTS \ -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']))" @@ -22,7 +28,7 @@ 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' + flock "$LOCK" bash -c "$(declare -f run); BUILD='$BUILD'; MODEL='$MODEL'; PIN='$PIN'; BENCH_KV_OPTS='$BENCH_KV_OPTS' run ordered 0 $D $R run pipelined 1 $D $R run ordered2 0 $D $R diff --git a/docs/repro/r4-kv-pipeline-context-sweep.sh b/docs/repro/r4-kv-pipeline-context-sweep.sh index b8a2b91bf72..3cdbef21a0d 100755 --- a/docs/repro/r4-kv-pipeline-context-sweep.sh +++ b/docs/repro/r4-kv-pipeline-context-sweep.sh @@ -11,13 +11,19 @@ PIN="${LLAMA_KV_TASKSET:-0,2,4}" NGEN="${LLAMA_KV_NGEN:-64}" LOCK=/tmp/beellama-single-gpu.lock +# llama-bench does not expose every host-KV option that llama-server does; pass only what it takes +BENCH_KV_OPTS="" +for opt in --kv-cpu-pinned --recurrent-state-offload; do + "$BUILD/bin/llama-bench" --help 2>&1 | grep -q -- "$opt" && BENCH_KV_OPTS="$BENCH_KV_OPTS $opt" +done + 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 \ + -ngl 99 -sm none -mg 0 -t 3 -nkvo 1 $BENCH_KV_OPTS \ -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 @@ -34,7 +40,7 @@ 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' + flock "$LOCK" bash -c "$(declare -f arm); BUILD='$BUILD'; MODEL='$MODEL'; PIN='$PIN'; BENCH_KV_OPTS='$BENCH_KV_OPTS'; NGEN='$NGEN' arm 0 $D $R arm 1 $D $R arm 0 $D $R diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 0316ebe0680..c580cbef9b7 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -2632,6 +2632,14 @@ void ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, ggml_backend_dev_t dev = ggml_backend_get_device(backend); if (dev == NULL || dev->iface.event_new == NULL) { + // Worth naming rather than folding into the generic message below: a backend that + // fans out over several devices -- the meta backend used for tensor parallelism -- + // lands here because that layer implements no events, and ordering a transfer stream + // against the consumer is what this is built on. It keeps the ordered path. + if (tr->debug > 0) { + GGML_LOG_INFO("%s: %s cannot order streams with events, staying on the ordered path\n", + __func__, ggml_backend_name(backend)); + } continue; } if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) { From 2055257936d44a0a06d58ddebf69a24e5b17a3b7 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 27 Aug 2026 10:19:28 +0200 Subject: [PATCH 03/32] sched: report where the pipelined token goes, and fix the budget check The host-time breakdown the feature doc describes had no code behind it: the counters existed but nothing accumulated or printed them. GGML_SCHED_TRANSPORT_DEBUG=2 now reports the split loop as a mean over each 128 graphs, with the bytes the ordered path still moves and why the look-ahead stopped; =3 names the tensors that are still on it. That is what found the rest: 40 blocking copies a token moving 0.4 MiB, 32 of them the device-to-host KV store. The budget warning now reports what the ring costs at the full context next to what it costs now, so --kv-pipeline-budget can be sized against the number that matters. It is still applied per graph: enforcing the projection would refuse the ring for every large -c even when the window never gets near it. llama-bench gains -kvcp and -rso. Without them a host-resident run measures something else entirely -- 9.02 against 19.43 t/s ordered at 16,384 -- and the repro scripts had been silently dropping both since llama-bench lost them. The exactness harness gives every task a nonce derived from its own name and length, so no two share a prefix the server can restore, and fails a task whose prompt_n says one was reused anyway. Assisted-by: Claude Opus 5 --- docs/kv-transport-pipelining.md | 220 +++++++++++++-------- docs/repro/r4-kv-pipeline-ab.sh | 8 +- docs/repro/r4-kv-pipeline-context-sweep.sh | 8 +- docs/repro/r4-kv-pipeline-exact.py | 25 ++- ggml/src/ggml-backend.cpp | 111 ++++++++++- tools/llama-bench/llama-bench.cpp | 69 ++++++- 6 files changed, 340 insertions(+), 101 deletions(-) diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index 73e0a330675..aa1105093de 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -123,22 +123,72 @@ RTX 4070 (11,902 MiB usable, sm_89), driver 610.57.04 / CUDA 13.3, i5-13400F, --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`): +(`docs/repro/r4-kv-pipeline-ab.sh`, `docs/repro/r4-kv-pipeline-context-sweep.sh`), +at `--kv-pipeline-budget 512` so the 32,768 ring is allowed: -| Depth | reps | ordered | pipelined | gain | `max(copy, compute)` ceiling | 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% | +| depth | ordered | pipelined | gain | peak device memory | +|---:|---|---|---:|---:| +| 4,096 | 31.3066, 31.2334 | 36.6107, 36.4701 | **+17.0%** | +28 MiB | +| 16,384 | 19.4344, 19.4828 | 31.0981, 31.0931 | **+59.8%** | +86 to +104 MiB | +| 32,768 | 12.9453, 12.9400 | 15.4571, 15.4516 | **+19.4%** | +206 MiB | -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). +> These need `-kvcp 1 -rso 1`, and for a while `llama-bench` did not have them: +> the repro scripts probed `--help`, found nothing, and quietly dropped both. The +> same commit then measures 19.43 -> 9.02 t/s ordered at 16,384 and the pipeline +> buys +6.7% instead of +60%, because a host-resident recurrent state costs more +> than the transport can win back. `llama-bench` takes them again. -Server decode behind an 18,422-token prompt -(`docs/repro/r4-kv-pipeline-exact.sh`): **18.468 -> 30.685 t/s, +66.2%**. +`llama-server`, one request, `temperature 0, top_k 1, seed 1234`: + +| prompt | `-c` | ordered | pipelined | gain | copy ms | compute ms | ceiling | share | +|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| 19,246 | 32,768 | 17.785 | 29.833 | **+67.7%** | 28.3 | 25.7 | 31.30 | 95.3% | +| 48,042 | 65,536 | 9.790 | 11.313 | **+15.6%** | 76.4 | 25.4 | 11.86 | 95.4% | + +`copy` and `compute` are read off `GGML_SCHED_TRANSPORT_DEBUG=2` on each arm, not +fitted: the ordered arm reports what it spends blocked in `ggml_backend_tensor_copy` +and what it spends waiting for the consumer. The ceiling is `max(copy, compute)` +plus the per-token work outside the split loop, which is on both arms. + +**The pipeline is within 5% of that ceiling at both depths.** What is left is not +a scheduling problem, and the section on the residual below says what it is. + +Pinning is worth as much as the pipeline and is off by default. Behind a +13,128-token prompt: + +| | ordered | pipelined | +|---|---:|---:| +| `--kv-cpu-pinned` | 21.582 | 32.252 | +| unpinned | 14.945 | 22.709 | + +Look-ahead deeper than one split is worse at every depth measured. At 19,246: +29.83 t/s at `N = 1`, 28.44 at `N = 2`, 25.93 at `N = 4`. `N = 1` is the default +for that reason. + +### The link is the ceiling, so the lever is bytes + +644 MiB in 28.3 ms is 22.0 GB/s, and `nvidia-smi` reports the card at gen4 x16. +That is about 88% of what the link delivers in practice, so there is no room left +in the transport itself. What is left is to send less. `-ctk q4_0 -ctv q4_0` +halves the cache and therefore the traffic: + +| prompt | KV | ordered | pipelined | delivered | +|---:|---|---:|---:|---:| +| 19,246 | q8_0 | 17.785 | 29.833 | 644.0 MiB | +| 19,246 | q4_0 | 22.904 | 31.502 | 343.2 MiB | +| 48,042 | q8_0 | 9.790 | 11.313 | 1602.5 MiB | +| 48,042 | q4_0 | 14.146 | 17.717 | 850.6 MiB | + +Halving the traffic is worth +5.6% on the pipelined path at 19,246 and +56.6% at +48,042. The difference is the crossover: at 19,246 the pipeline has already +brought the copy down to the compute floor, and the consumer wait is 27.31 ms at +q8_0 against 27.26 ms at q4_0, the same number. Removing bytes there removes work +nothing was waiting for. At 48,042 the copy still dominates and every byte +removed is a byte off the token. + +**Whether to spend a quantisation step on the cache is a depth question, and the +two compound.** At 48,042, q4_0 with the pipeline is 17.717 against 9.790 for +q8_0 without it. ### Across context depth, with device memory @@ -173,13 +223,21 @@ 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 ring beats `--kv-gpu-layers` per MiB, and the two barely add up.** Measured +behind a 19,246-token prompt at `-c 32768`, where a device-resident layer costs +about 68 MiB and the ring costs about 205 MiB: + +| | no `--kv-gpu-layers` | `--kv-gpu-layers 4` | `--kv-gpu-layers 8` | +|---|---:|---:|---:| +| ordered | 17.785 | 20.334 | | +| pipelined | 29.843 | 30.263 | 30.640 | + +Four device-resident layers are worth +14.3% on the ordered path and +1.4% on the +pipelined one. The reason they stop paying is the point of the section above: the +pipeline has already moved the bottleneck down to the compute floor, so removing +a quarter of the traffic removes something that was no longer being waited for. +Whether this still holds where the copy dominates by a wide margin has not been +measured. ### The budget @@ -189,49 +247,57 @@ 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. +- Under the cap the ring is allocated and the deliveries pipeline. - 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). + backend's device context are released. + +**The cap is applied to what the current graph needs, not to what the full +context would need.** A run whose window stays small keeps the ring whatever +`-n_ctx` says, which is the common case and the reason it is done this way: a +staged input is a view of the cache tensor, so the full-context figure is there +for the asking, but enforcing it would refuse the ring for every large `-c` even +when the window never gets near it. The warning reports both numbers so that +`--kv-pipeline-budget` can be sized against the one that matters. -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. +The cost of deciding per graph is that a context which grows past the budget +allocates a ring for the small early windows and gives it back once it outgrows +them. That transient is bounded by the budget itself, which is the memory the +user already authorised, so it is a property of the cap rather than a defect in +it. -**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. +At 32,768 the ring is 204 MiB at the full context, over the 128 MiB default. +`--kv-pipeline-budget 512` buys 20.350 -> 31.463 t/s behind an 18,432-token +prompt. -Where the split-loop host time goes, per decode graph at 18.5k -(`GGML_SCHED_TRANSPORT_DEBUG=2`): +### Where the rest of the token goes + +Per decode graph, `GGML_SCHED_TRANSPORT_DEBUG=2`, behind a 19,246-token prompt: | | 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 | +| total | 54.11 ms | 31.10 ms | +| blocked in the ordered `ggml_backend_tensor_copy` | 28.30 ms | 3.57 ms | +| blocked waiting for the consumer backend | 25.70 ms | 27.31 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. +| bytes delivered early / late | 0 / 0 MiB | 644.0 / 2.3 MiB | +| bytes left on the ordered path | 28.3 MiB | 0.4 MiB | + +The blocking host-to-device copy is all but 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. 644 MiB in the 28.3 ms the +ordered arm reports for the same bytes is 22.0 GB/s, which is what this link +does; the transfer cannot be made faster, only hidden. + +The 3.57 ms that remains moves 0.4 MiB. It is not bandwidth, it is 40 separate +blocking copies at about 89 us each, and `GGML_SCHED_TRANSPORT_DEBUG=3` names +them: 16 `cache_k_store_stage_l*`, 16 `cache_v_store_stage_l*`, and the graph +inputs. The store staging tensors are the device-to-host write of this token's +K and V, one per attention layer, and the ring carries deliveries in the other +direction only. At 48,042 the same 40 copies cost 8.7 ms of an 85.6 ms graph, so +this is worth about 10% of the token and it does not shrink with context. Do not compare these numbers against runs on other models, prompts, cache settings, hardware, or commits. @@ -244,18 +310,31 @@ The gates, and what was run for them: 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`. + Every task now carries a nonce derived from its own name and length, so no two + share a prefix the server could restore, and the harness fails a task whose + reported `prompt_n` says a prefix was reused anyway. + + **One task is still not a gate.** Re-run at `-c 32768` with the ring active, + seven of the eight tasks are byte-identical at `N = 0`, `N = 1` and `N = 4`. + `records@18432` is not, and it is not the pipeline: two separate `N = 0` runs + of it produced two different hashes, with `prompt_n = 29561` both times, so no + prefix was reused. Its prompt is about 29.6k tokens against a 32,768 context, + close enough to the limit that something in the slot handling varies. Until + that is understood the task should be read as unmeasured rather than passing. 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. + graph, how much of it goes early, and the source buffer type); `=2` adds the + per-graph host-time breakdown above, as the mean over each 128 graphs, with + how often the look-ahead stopped on the depth it was given against a slot + whose reader had not run; `=3` names the tensors still on the ordered path. `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 +A device-resident KV run is unaffected, and was measured to confirm it: 38.5612 +t/s at depth 0 against 38.5240 at depth 1, `tg128 @ d4096`, with the transport never enabled because the scheduler is given a depth of 0. ## Scope and limits @@ -326,24 +405,11 @@ row rather than laid out end to end. - Fix `-sm tensor` with `--no-kv-offload` (above). Until then it should not be used: it is wrong rather than slow. -- Events on the meta backend and device, so a transfer stream can be ordered - against a tensor-parallel consumer at all. -- A ring that can live in a meta buffer, as slot tensors rather than offsets, so - tensor parallelism can be pipelined once it is correct. -- A transfer-only meta backend that does not stand up a second collective - communicator: `ggml_backend_dev_init` on a meta device runs the whole meta - context constructor, which calls `ggml_backend_comm_init` across every device. -- Remove the transient device-memory peak. The budget is applied per graph, so a - context that grows past it still allocates a ring for the small windows of early - prefill and releases it once the window outgrows the budget: +112 MiB at 32,768 - against +0 in steady state. Deciding against the context's final size needs KV - geometry the scheduler does not have. -- Compare the ring against `--kv-gpu-layers` at depth. At 131,072 the ring's - 818 MiB is about three attention layers' K and V; making three of sixteen - device-resident would remove about 19% of the host-to-device traffic against the - 9.1% the ring buys there. Unmeasured, and it moves with layer count and card. -- Give the exactness harness a per-task nonce. Two long-prompt tasks proved - non-deterministic in the baseline -- a second control run reproduced this - branch's hashes rather than its own -- because the server restores a similar - cached prefix. Until that is pinned down the harness is a weaker gate than it - looks. +- Events on the meta backend and device, a ring that can live in a meta buffer, + and a transfer-only meta backend, so tensor parallelism can be pipelined once + it is correct. Written on a separate branch, and not reachable until it is. +- Take the last small blocking copies off the host's critical path. On the + pipelined path 3.6 ms per graph at 19,246 and 8.7 ms at 48,042 is still spent + inside `ggml_backend_tensor_copy`, for 0.4 MiB. It is 40 separate copies, and + 32 of them are the device-to-host KV store, which runs on the other copy engine + and could overlap the deliveries instead of blocking the host. diff --git a/docs/repro/r4-kv-pipeline-ab.sh b/docs/repro/r4-kv-pipeline-ab.sh index 9d8bb83a393..c1636994a8f 100755 --- a/docs/repro/r4-kv-pipeline-ab.sh +++ b/docs/repro/r4-kv-pipeline-ab.sh @@ -9,10 +9,12 @@ BUILD="${LLAMA_KV_BUILD:-build}" PIN="${LLAMA_KV_TASKSET:-0,2,4}" LOCK=/tmp/beellama-single-gpu.lock -# llama-bench does not expose every host-KV option that llama-server does; pass only what it takes +# An unpinned host cache and a host-resident recurrent state both cost more than the transport +# can win back, so a run without these does not measure the same thing. Older llama-bench builds +# do not have them; skip rather than fail, and the numbers are then not comparable to the doc. BENCH_KV_OPTS="" -for opt in --kv-cpu-pinned --recurrent-state-offload; do - "$BUILD/bin/llama-bench" --help 2>&1 | grep -q -- "$opt" && BENCH_KV_OPTS="$BENCH_KV_OPTS $opt" +for opt in kvcp rso; do + "$BUILD/bin/llama-bench" --help 2>&1 | grep -q -- "-$opt," && BENCH_KV_OPTS="$BENCH_KV_OPTS -$opt 1" done run () { # $1 label, $2 pipeline depth, $3 context depth, $4 reps diff --git a/docs/repro/r4-kv-pipeline-context-sweep.sh b/docs/repro/r4-kv-pipeline-context-sweep.sh index 3cdbef21a0d..0f1386f9483 100755 --- a/docs/repro/r4-kv-pipeline-context-sweep.sh +++ b/docs/repro/r4-kv-pipeline-context-sweep.sh @@ -11,10 +11,12 @@ PIN="${LLAMA_KV_TASKSET:-0,2,4}" NGEN="${LLAMA_KV_NGEN:-64}" LOCK=/tmp/beellama-single-gpu.lock -# llama-bench does not expose every host-KV option that llama-server does; pass only what it takes +# An unpinned host cache and a host-resident recurrent state both cost more than the transport +# can win back, so a run without these does not measure the same thing. Older llama-bench builds +# do not have them; skip rather than fail, and the numbers are then not comparable to the doc. BENCH_KV_OPTS="" -for opt in --kv-cpu-pinned --recurrent-state-offload; do - "$BUILD/bin/llama-bench" --help 2>&1 | grep -q -- "$opt" && BENCH_KV_OPTS="$BENCH_KV_OPTS $opt" +for opt in kvcp rso; do + "$BUILD/bin/llama-bench" --help 2>&1 | grep -q -- "-$opt," && BENCH_KV_OPTS="$BENCH_KV_OPTS -$opt 1" done arm () { # $1 pipeline depth, $2 context depth, $3 reps diff --git a/docs/repro/r4-kv-pipeline-exact.py b/docs/repro/r4-kv-pipeline-exact.py index 0a0800f3f79..cfb10591a99 100644 --- a/docs/repro/r4-kv-pipeline-exact.py +++ b/docs/repro/r4-kv-pipeline-exact.py @@ -41,7 +41,15 @@ def filler(name, target_tokens): return "".join(unit % i for i in range(reps)) return unit * reps -def ask(label, prompt, ntok): +def nonce(name, length): + # The server restores a cached prefix from an earlier task, and a restored window is not + # numerically the same as a freshly prefilled one, so two tasks that share a long prefix stop + # measuring the code under test. This makes every task's prefix unique, and it is derived from + # the task rather than drawn at random so that a control run produces comparable hashes. + h = hashlib.sha256(f"{name}/{length}".encode()).hexdigest()[:32] + return f"Session {h}. Ignore this line.\n\n" + +def ask(label, prompt, ntok, want_prefill): 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, @@ -56,15 +64,20 @@ def ask(label, prompt, ntok): # reasoning models put most of the generation in reasoning_content; hash both text = (m.get("reasoning_content") or "") + "\x00" + (m.get("content") or "") t = d.get("timings", {}) + # a reused prefix shows up as a prompt_n far below the prompt actually sent; the hash it + # produces is not comparable to a fresh prefill, so say so rather than reporting it silently + prompt_n = t.get("prompt_n") or 0 + reused = prompt_n < want_prefill // 2 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 + f"prompt_n={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}" + f"{' CACHE_REUSE' if reused else ''}", flush=True) + return not reused ok = True 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) + prompt = nonce(name, length) + filler(name, length) + "\n\n" + QUESTIONS[name] + ok &= ask(f"{name}@{length}", prompt, ntok, length) sys.exit(0 if ok else 1) diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index c580cbef9b7..fc55b438a79 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -865,6 +865,16 @@ 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; + // why the look-ahead stopped: the depth it was given, or a slot whose reader has not run yet. + // The two want opposite fixes, so they are counted apart. + int64_t n_stop_depth; + int64_t n_stop_recycle; + int64_t p_stop_depth, p_stop_recycle; + + int64_t n_bytes_ordered; // what the ordered blocking copies still move + int64_t p_bytes_ordered; + bool named_ordered; // debug >= 3 names them once, they are the same every graph + int debug; }; @@ -1922,20 +1932,29 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { } } - // per-ring slot size and delivery order - size_t slot_size[GGML_SCHED_MAX_BACKENDS] = { 0 }; + // per-ring slot size and delivery order. The budget is applied to what this graph needs, so + // that a run whose window stays small keeps the ring whatever -n_ctx says. slot_size_max is + // what the same ring costs once the context is full, taken from the cache tensor the staged + // input is a view of; it is reported rather than enforced, because deciding on it would refuse + // the ring for every large -c even when the window never gets there. + size_t slot_size[GGML_SCHED_MAX_BACKENDS] = { 0 }; + size_t slot_size_max[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; - size_t need = 0; + size_t need = 0; + size_t need_max = 0; for (int j = 0; j < split->n_inputs; j++) { if (!tr->input_staged[tr->split_input_ofs[i] + j]) { continue; } - need += GGML_PAD(ggml_nbytes(split->inputs[j]), tr->rings[bid].alignment); + const struct ggml_tensor * input = split->inputs[j]; + const struct ggml_tensor * base = input->view_src ? input->view_src : input; + need += GGML_PAD(ggml_nbytes(input), tr->rings[bid].alignment); + need_max += GGML_PAD(ggml_nbytes(base), tr->rings[bid].alignment); } if (need == 0) { @@ -1943,7 +1962,8 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { } tr->split_order[i] = tr->rings[bid].n_staged++; - slot_size[bid] = std::max(slot_size[bid], need); + slot_size[bid] = std::max(slot_size[bid], need); + slot_size_max[bid] = std::max(slot_size_max[bid], std::max(need, need_max)); } for (int bid = 0; bid < sched->n_backends; bid++) { @@ -1952,7 +1972,8 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { continue; } - const size_t ring_size = slot_size[bid] * tr->n_slots; + const size_t ring_size = slot_size[bid] * tr->n_slots; + const size_t ring_size_max = slot_size_max[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 @@ -1965,10 +1986,11 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { 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); + GGML_LOG_WARN("%s: transport ring on %s needs %zu MiB now and %zu MiB at the full " + "context, 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, + ring_size_max >> 20, tr->budget >> 20); r->reported_no_room = true; } ggml_backend_sched_transport_release_ring(sched, bid, true); @@ -2133,6 +2155,7 @@ static void ggml_backend_sched_transport_prefetch(ggml_backend_sched_t sched, in continue; } if (tr->split_order[i] > r->consumed + tr->depth) { + tr->n_stop_depth++; return; } @@ -2145,6 +2168,7 @@ static void ggml_backend_sched_transport_prefetch(ggml_backend_sched_t sched, in if (slot->release_armed) { ggml_backend_event_wait(r->transfer, slot->release); slot->release_armed = false; + tr->n_stop_recycle++; } for (int j = 0; j < split->n_inputs; j++) { @@ -2279,6 +2303,8 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } } + const int64_t t_graph_0 = tr->debug >= 2 ? 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; @@ -2287,11 +2313,15 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s // ensure the previous split's async work has completed before we start // this split, the allocator may have reused buffer regions across splits if (split->n_inputs == 0 && prev_backend_id >= 0 && prev_backend_id != split_backend_id) { + const int64_t t0 = tr->debug >= 2 ? ggml_time_us() : 0; if (sched->events[prev_backend_id][sched->cur_copy] != NULL) { ggml_backend_event_synchronize(sched->events[prev_backend_id][sched->cur_copy]); } else { ggml_backend_synchronize(sched->backends[prev_backend_id]); } + if (tr->debug >= 2) { + tr->t_sync_us += ggml_time_us() - t0; + } } // copy the input tensors to the split backend @@ -2318,12 +2348,25 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s 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 + const int64_t t0 = tr->debug >= 2 ? ggml_time_us() : 0; 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); } + if (tr->debug >= 2) { + tr->t_sync_us += ggml_time_us() - t0; + } + const int64_t t1 = tr->debug >= 2 ? ggml_time_us() : 0; ggml_backend_tensor_copy(input, input_cpy); + if (tr->debug >= 2) { + tr->t_copy_us += ggml_time_us() - t1; + tr->n_bytes_ordered += ggml_nbytes(input); + } + if (tr->debug >= 3 && !tr->named_ordered) { + GGML_LOG_INFO("%s: ordered copy %s %zu KiB from %s\n", __func__, input->name, + ggml_nbytes(input) >> 10, ggml_backend_buft_name(input->buffer->buft)); + } } else { // wait for the split backend to finish using the input before overwriting it if (sched->events[split_backend_id][sched->cur_copy] != NULL) { @@ -2421,13 +2464,26 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s // try async copy, but if not possible, we can still use a sync copy without synchronizing the dst backend, since we handle the synchronization here with multiple copies and events // TODO: add public function to facilitate this, since applications do not have direct access to the backend interface if (!split_backend->iface.cpy_tensor_async || !split_backend->iface.cpy_tensor_async(input_backend, split_backend, input, input_cpy)) { + const int64_t t0 = tr->debug >= 2 ? 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); } + if (tr->debug >= 2) { + tr->t_sync_us += ggml_time_us() - t0; + } + const int64_t t1 = tr->debug >= 2 ? ggml_time_us() : 0; ggml_backend_tensor_copy(input, input_cpy); + if (tr->debug >= 2) { + tr->t_copy_us += ggml_time_us() - t1; + tr->n_bytes_ordered += ggml_nbytes(input); + } + if (tr->debug >= 3 && !tr->named_ordered) { + GGML_LOG_INFO("%s: ordered copy %s %zu KiB from %s\n", __func__, input->name, + ggml_nbytes(input) >> 10, ggml_backend_buft_name(input->buffer->buft)); + } } } } @@ -2500,6 +2556,41 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s prev_backend_id = split_backend_id; } + if (tr->debug >= 2) { + tr->t_graph_us += ggml_time_us() - t_graph_0; + tr->n_graphs++; + + // every 128 graphs, and as the mean over those 128, so that one graph's noise does not + // decide what the numbers look like + if (tr->n_graphs % 128 == 0) { + const double n = 128.0; + GGML_LOG_INFO("%s: per graph over %d: total %.2f ms, sync %.2f ms, ordered copy %.2f ms, " + "issue %.2f ms, early %.1f MiB, late %.1f MiB, ordered %.1f MiB, " + "stops depth/recycle %.1f/%.1f\n", + __func__, (int) n, + (tr->t_graph_us - tr->p_graph_us)/1e3/n, + (tr->t_sync_us - tr->p_sync_us )/1e3/n, + (tr->t_copy_us - tr->p_copy_us )/1e3/n, + (tr->t_issue_us - tr->p_issue_us)/1e3/n, + (tr->n_bytes_early - tr->p_bytes_early)/1048576.0/n, + (tr->n_bytes_late - tr->p_bytes_late )/1048576.0/n, + (tr->n_bytes_ordered - tr->p_bytes_ordered)/1048576.0/n, + (tr->n_stop_depth - tr->p_stop_depth )/n, + (tr->n_stop_recycle - tr->p_stop_recycle)/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; + tr->p_bytes_ordered = tr->n_bytes_ordered; + tr->named_ordered = true; + tr->p_stop_depth = tr->n_stop_depth; + tr->p_stop_recycle = tr->n_stop_recycle; + } + } + return GGML_STATUS_SUCCESS; } diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index a2da93b9a28..2d03e433c69 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -343,6 +343,8 @@ struct cmd_params { std::vector load_mode; std::vector main_gpu; std::vector no_kv_offload; + std::vector kv_cpu_pinned; + std::vector recurrent_state_offload; std::vector flash_attn; std::vector> devices; std::vector> tensor_split; @@ -387,6 +389,8 @@ static const cmd_params cmd_params_defaults = { /* load_mode */ { LLAMA_LOAD_MODE_AUTO }, /* main_gpu */ { 0 }, /* no_kv_offload */ { false }, + /* kv_cpu_pinned */ { false }, + /* recurrent_state_offload */ { false }, /* flash_attn */ { LLAMA_FLASH_ATTN_TYPE_AUTO }, /* devices */ { {} }, /* tensor_split */ { std::vector(llama_max_devices(), 0.0f) }, @@ -457,6 +461,8 @@ static void print_usage(int /* argc */, char ** argv) { printf(" -sm, --split-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.split_mode, split_mode_str), ",").c_str()); printf(" -mg, --main-gpu (default: %s)\n", join(cmd_params_defaults.main_gpu, ",").c_str()); printf(" -nkvo, --no-kv-offload <0|1> (default: %s)\n", join(cmd_params_defaults.no_kv_offload, ",").c_str()); + printf(" -kvcp, --kv-cpu-pinned <0|1> (default: %s)\n", join(cmd_params_defaults.kv_cpu_pinned, ",").c_str()); + printf(" -rso, --recurrent-state-offload <0|1> (default: %s)\n", join(cmd_params_defaults.recurrent_state_offload, ",").c_str()); printf(" -fa, --flash-attn (default: %s)\n", join(transform_to_str(cmd_params_defaults.flash_attn, llama_flash_attn_type_name), ",").c_str()); printf(" -dev, --device (default: auto)\n"); printf(" -lm, --load-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").c_str()); @@ -799,6 +805,20 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { } auto p = string_split(argv[i], split_delim); params.no_kv_offload.insert(params.no_kv_offload.end(), p.begin(), p.end()); + } else if (arg == "-kvcp" || arg == "--kv-cpu-pinned") { + if (++i >= argc) { + invalid_param = true; + break; + } + auto p = string_split(argv[i], split_delim); + params.kv_cpu_pinned.insert(params.kv_cpu_pinned.end(), p.begin(), p.end()); + } else if (arg == "-rso" || arg == "--recurrent-state-offload") { + if (++i >= argc) { + invalid_param = true; + break; + } + auto p = string_split(argv[i], split_delim); + params.recurrent_state_offload.insert(params.recurrent_state_offload.end(), p.begin(), p.end()); } else if (arg == "--numa") { if (++i >= argc) { invalid_param = true; @@ -1143,6 +1163,12 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { if (params.no_kv_offload.empty()) { params.no_kv_offload = cmd_params_defaults.no_kv_offload; } + if (params.kv_cpu_pinned.empty()) { + params.kv_cpu_pinned = cmd_params_defaults.kv_cpu_pinned; + } + if (params.recurrent_state_offload.empty()) { + params.recurrent_state_offload = cmd_params_defaults.recurrent_state_offload; + } if (params.flash_attn.empty()) { params.flash_attn = cmd_params_defaults.flash_attn; } @@ -1205,6 +1231,8 @@ struct cmd_params_instance { llama_load_mode load_mode; int main_gpu; bool no_kv_offload; + bool kv_cpu_pinned; + bool recurrent_state_offload; llama_flash_attn_type flash_attn; std::vector devices; std::vector tensor_split; @@ -1284,6 +1312,8 @@ struct cmd_params_instance { cparams.type_k = type_k; cparams.type_v = type_v; cparams.offload_kqv = !no_kv_offload; + cparams.kv_cpu_pinned = kv_cpu_pinned; + cparams.recurrent_state_offload = recurrent_state_offload; cparams.flash_attn_type = flash_attn; cparams.embeddings = embeddings; cparams.op_offload = !no_op_offload; @@ -1317,6 +1347,8 @@ static std::vector get_cmd_params_instances(const cmd_param for (const auto & tk : params.type_k) for (const auto & tv : params.type_v) for (const auto & nkvo : params.no_kv_offload) + for (const auto & kvcp : params.kv_cpu_pinned) + for (const auto & rso : params.recurrent_state_offload) for (const auto & fa : params.flash_attn) for (const auto & nt : params.n_threads) for (const auto & cm : params.cpu_mask) @@ -1346,6 +1378,8 @@ static std::vector get_cmd_params_instances(const cmd_param /* .load_mode = */ lm, /* .main_gpu = */ mg, /* .no_kv_offload = */ nkvo, + /* .kv_cpu_pinned = */ kvcp, + /* .recurrent_state_offload = */ rso, /* .flash_attn = */ fa, /* .devices = */ devs, /* .tensor_split = */ ts, @@ -1382,6 +1416,8 @@ static std::vector get_cmd_params_instances(const cmd_param /* .load_mode = */ lm, /* .main_gpu = */ mg, /* .no_kv_offload = */ nkvo, + /* .kv_cpu_pinned = */ kvcp, + /* .recurrent_state_offload = */ rso, /* .flash_attn = */ fa, /* .devices = */ devs, /* .tensor_split = */ ts, @@ -1418,6 +1454,8 @@ static std::vector get_cmd_params_instances(const cmd_param /* .load_mode = */ lm, /* .main_gpu = */ mg, /* .no_kv_offload = */ nkvo, + /* .kv_cpu_pinned = */ kvcp, + /* .recurrent_state_offload = */ rso, /* .flash_attn = */ fa, /* .devices = */ devs, /* .tensor_split = */ ts, @@ -1459,6 +1497,8 @@ struct test { llama_load_mode load_mode; int main_gpu; bool no_kv_offload; + bool kv_cpu_pinned; + bool recurrent_state_offload; llama_flash_attn_type flash_attn; std::vector devices; std::vector tensor_split; @@ -1498,6 +1538,8 @@ struct test { load_mode = inst.load_mode; main_gpu = inst.main_gpu; no_kv_offload = inst.no_kv_offload; + kv_cpu_pinned = inst.kv_cpu_pinned; + recurrent_state_offload = inst.recurrent_state_offload; flash_attn = inst.flash_attn; devices = inst.devices; tensor_split = inst.tensor_split; @@ -1562,7 +1604,8 @@ struct test { "model_filename", "model_type", "model_size", "model_n_params", "n_batch", "n_ubatch", "n_threads", "cpu_mask", "cpu_strict", "poll", "type_k", "type_v", "n_gpu_layers", "n_cpu_moe", "split_mode", - "main_gpu", "no_kv_offload", "flash_attn", "devices", "tensor_split", + "main_gpu", "no_kv_offload", "kv_cpu_pinned", "recurrent_state_offload", + "flash_attn", "devices", "tensor_split", "tensor_buft_overrides", "load_mode", "embeddings", "no_op_offload", "no_host", "fit_target", "fit_min_ctx", "n_prompt", "n_gen", "n_depth", @@ -1581,7 +1624,8 @@ struct test { field == "fit_target" || field == "fit_min_ctx" || field == "flash_attn") { return INT; } - if (field == "f16_kv" || field == "no_kv_offload" || field == "cpu_strict" || + if (field == "f16_kv" || field == "no_kv_offload" || field == "kv_cpu_pinned" || + field == "recurrent_state_offload" || field == "cpu_strict" || field == "embeddings" || field == "no_host") { return BOOL; } @@ -1653,6 +1697,8 @@ struct test { split_mode_str(split_mode), std::to_string(main_gpu), std::to_string(no_kv_offload), + std::to_string(kv_cpu_pinned), + std::to_string(recurrent_state_offload), std::to_string((int) flash_attn), devices_to_string(devices), tensor_split_str, @@ -1848,6 +1894,12 @@ struct markdown_printer : public printer { if (field == "test") { return 15; } + if (field == "kv_cpu_pinned") { + return 4; + } + if (field == "recurrent_state_offload") { + return 3; + } if (field == "no_op_offload") { return 4; } @@ -1873,6 +1925,12 @@ struct markdown_printer : public printer { if (field == "n_threads") { return "threads"; } + if (field == "kv_cpu_pinned") { + return "kvcp"; + } + if (field == "recurrent_state_offload") { + return "rso"; + } if (field == "no_kv_offload") { return "nkvo"; } @@ -1954,6 +2012,13 @@ struct markdown_printer : public printer { if (params.split_mode.size() > 1 || params.split_mode != cmd_params_defaults.split_mode) { fields.emplace_back("split_mode"); } + if (params.kv_cpu_pinned.size() > 1 || params.kv_cpu_pinned != cmd_params_defaults.kv_cpu_pinned) { + fields.emplace_back("kv_cpu_pinned"); + } + if (params.recurrent_state_offload.size() > 1 || + params.recurrent_state_offload != cmd_params_defaults.recurrent_state_offload) { + fields.emplace_back("recurrent_state_offload"); + } if (params.no_kv_offload.size() > 1 || params.no_kv_offload != cmd_params_defaults.no_kv_offload) { fields.emplace_back("no_kv_offload"); } From 4354b739c514256a15dbb9fac1a9cd8d577722e1 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 27 Aug 2026 10:41:36 +0200 Subject: [PATCH 04/32] sched : name what the ordered path still copies, and why it costs what it does GGML_SCHED_TRANSPORT_DEBUG=3 now reports what each remaining blocking copy cost, not just its name. It turns out one of them is almost all of it: attn_inp_k_rot, 256 KiB, 18 us on the ordered path and 3.4 ms behind one split of look-ahead. That is the copy engine, not latency. A blocking copy waits for the deliveries already queued on it, and two staged splits at 22.0 GB/s is 3.6 ms. Issuing the delivery in pieces does not help, the engine is FIFO across streams. Putting the copy on the consumer's stream so the host never blocks moves the time into the consumer wait and leaves throughput alone. The doc records both, so the next person does not spend the afternoon on it again. Assisted-by: Claude Opus 5 --- docs/kv-transport-pipelining.md | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index aa1105093de..56ae3256926 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -291,13 +291,24 @@ critical path without being added to the consumer's. 644 MiB in the 28.3 ms the ordered arm reports for the same bytes is 22.0 GB/s, which is what this link does; the transfer cannot be made faster, only hidden. -The 3.57 ms that remains moves 0.4 MiB. It is not bandwidth, it is 40 separate -blocking copies at about 89 us each, and `GGML_SCHED_TRANSPORT_DEBUG=3` names -them: 16 `cache_k_store_stage_l*`, 16 `cache_v_store_stage_l*`, and the graph -inputs. The store staging tensors are the device-to-host write of this token's -K and V, one per attention layer, and the ring carries deliveries in the other -direction only. At 48,042 the same 40 copies cost 8.7 ms of an 85.6 ms graph, so -this is worth about 10% of the token and it does not shrink with context. +The 3.57 ms that remains moves 0.4 MiB, and `GGML_SCHED_TRANSPORT_DEBUG=3` shows +that almost all of it is one copy: `attn_inp_k_rot`, 256 KiB, 18 us on the +ordered path and 3.4 ms behind one split of look-ahead. The 32 KV store copies +cost 353 us between them. + +It looks like latency and is not. A blocking copy shares the device's copy engine +with the deliveries and waits for what is already queued there: two staged splits +at 22.0 GB/s is 3.6 ms, which is the number. Two things were tried and neither +helped. Issuing the delivery in pieces so the blocking copy can interleave does +nothing -- the engine is FIFO across streams, `attn_inp_k_rot` stays at 3.4 ms at +every piece size, and small pieces cost throughput (29.80 t/s whole, 28.73 at +4 MiB, 22.01 at 1 MiB). Putting the copy on the consumer's own stream so the host +never blocks moves the time rather than removing it: the ordered copy falls from +3.57 ms to 0.16 ms, the consumer wait rises from 27.31 ms to 31.05 ms, and +throughput does not move (29.808 against 29.834). + +So this is not spare time. Those 256 KiB cross the same saturated link as the +644 MiB of deliveries, and the link is the ceiling. Do not compare these numbers against runs on other models, prompts, cache settings, hardware, or commits. From c4d64f4865dec0879dbf471329879503d7189e53 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 27 Aug 2026 10:58:42 +0200 Subject: [PATCH 05/32] repro : make the exactness tasks independent of each other records@18432 was giving different answers across otherwise identical N = 0 runs, which made it useless as a gate and looked like the pipeline breaking exactness. It is not the task: asked on its own with the prompt cache off it returns the same hash three times running, at -c 32768 and at -c 65536. It is the harness. All eight tasks share one server with prompt caching on, and records@18432 is about 29.6k tokens with a task of about the same size ahead of it, so the two do not both fit in a 32,768 cache and placement depended on what was still resident. The nonce stops a prefix being restored, it does not stop the pressure. cache_prompt=false does. Two independent N = 0 passes now agree on all eight tasks, and N = 0, N = 1 and N = 4 agree on all eight. Assisted-by: Claude Opus 5 --- docs/kv-transport-pipelining.md | 24 ++++++++++++++---------- docs/repro/r4-kv-pipeline-exact.py | 6 +++++- 2 files changed, 19 insertions(+), 11 deletions(-) diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index 56ae3256926..9a5e1b7907f 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -321,17 +321,21 @@ The gates, and what was run for them: 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`. - Every task now carries a nonce derived from its own name and length, so no two + Two things keep the tasks independent of each other, and both were needed. + Every task carries a nonce derived from its own name and length, so no two share a prefix the server could restore, and the harness fails a task whose - reported `prompt_n` says a prefix was reused anyway. - - **One task is still not a gate.** Re-run at `-c 32768` with the ring active, - seven of the eight tasks are byte-identical at `N = 0`, `N = 1` and `N = 4`. - `records@18432` is not, and it is not the pipeline: two separate `N = 0` runs - of it produced two different hashes, with `prompt_n = 29561` both times, so no - prefix was reused. Its prompt is about 29.6k tokens against a 32,768 context, - close enough to the limit that something in the slot handling varies. Until - that is understood the task should be read as unmeasured rather than passing. + `prompt_n` says one was reused anyway. Each request also sets + `cache_prompt: false`, so a task never inherits what the previous one left in + the cache. + + The second is what made `records@18432` a gate rather than a coin flip. Its + prompt is about 29.6k tokens against a 32,768 context, and the task before it + is about the same size, so the two do not both fit and placement depended on + what was still resident. Two otherwise identical `N = 0` runs of it produced + different hashes. Asked on its own with the cache off it is perfectly stable: + the same hash three times running, at `-c 32768` and at `-c 65536`. With the + flag set, two independent `N = 0` passes agree on all eight tasks, and + `N = 0`, `N = 1` and `N = 4` agree on all eight. 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. diff --git a/docs/repro/r4-kv-pipeline-exact.py b/docs/repro/r4-kv-pipeline-exact.py index cfb10591a99..cff45a397a6 100644 --- a/docs/repro/r4-kv-pipeline-exact.py +++ b/docs/repro/r4-kv-pipeline-exact.py @@ -50,8 +50,12 @@ def nonce(name, length): return f"Session {h}. Ignore this line.\n\n" def ask(label, prompt, ntok, want_prefill): + # cache_prompt=False forces a full prefill. Without it a task inherits whatever the previous + # one left in the cache, and two tasks whose prompts do not both fit make placement depend on + # that: records@18432 then gives different answers across otherwise identical runs. 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, + "cache_prompt": False}).encode() req = urllib.request.Request(f"http://127.0.0.1:{PORT}/v1/chat/completions", body, {"Content-Type": "application/json"}) try: From 2f70cbe9212c18af01f9ce7d2ff842d26c26657b Mon Sep 17 00:00:00 2001 From: piggidragon Date: Fri, 28 Aug 2026 19:47:42 +0200 Subject: [PATCH 06/32] sched: fix pipelined transport fallback paths Restrict staging to annotated CUDA inputs, make backend decline complete, re-evaluate budgets, freeze scheduler configuration, and preserve tensor layout. Add regressions for prefix changes and fallback behavior. Assisted-by: OpenAI Codex --- common/arg.cpp | 6 +- docs/kv-transport-pipelining.md | 52 ++-- docs/repro/r4-kv-pipeline-exact.py | 11 +- docs/repro/r4-kv-pipeline-exact.sh | 20 +- ggml/include/ggml-backend.h | 13 +- ggml/include/ggml.h | 26 +- ggml/src/ggml-backend-meta.cpp | 5 +- ggml/src/ggml-backend.cpp | 232 +++++++++++------- ggml/src/ggml.c | 5 +- src/llama-context.cpp | 17 +- src/llama-kv-cache.cpp | 9 + tests/test-alloc.cpp | 368 +++++++++++++++++++++++++++-- 12 files changed, 583 insertions(+), 181 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 9cdcf11b2b2..7625d3940cd 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -30,6 +30,7 @@ #include #include #include +#include #include #include #include @@ -2443,8 +2444,9 @@ common_params_context common_params_parser_init(common_params & params, llama_ex "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"); + constexpr size_t mib = 1024u*1024u; + if (value < 0 || (size_t) value > std::numeric_limits::max()/mib) { + throw std::invalid_argument("--kv-pipeline-budget is out of range for this platform"); } params.kv_pipeline_budget_mib = value; } diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index 9a5e1b7907f..2ab357323a0 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -248,9 +248,8 @@ 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. -- 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. +- Over it the scheduler declines and keeps the ordered path for that graph. + Later graphs are evaluated again, so a smaller live window can use the ring. - Declining costs nothing in steady state. Both the ring and the transfer backend's device context are released. @@ -320,7 +319,8 @@ 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`. + `N = 0`, `N = 1` and `N = 4`. `docs/repro/r4-kv-pipeline-exact.sh` compares + every requested depth with the first and fails on a hash difference. Two things keep the tasks independent of each other, and both were needed. Every task carries a nonce derived from its own name and length, so no two share a prefix the server could restore, and the harness fails a task whose @@ -343,10 +343,9 @@ The gates, and what was run for them: `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, as the mean over each 128 graphs, with - how often the look-ahead stopped on the depth it was given against a slot - whose reader had not run; `=3` names the tensors still on the ordered path. - `ggml_backend_sched_get_transport_pipeline_stats()` exposes the same counters - to callers. + depth stops and the number of recycle waits enqueued; `=3` names the tensors + still on the ordered path. `ggml_backend_sched_get_transport_pipeline_stats()` + exposes deliveries and early and late byte counts to callers. A device-resident KV run is unaffected, and was measured to confirm it: 38.5612 t/s at depth 0 against 38.5240 at depth 1, `tg128 @ d4096`, with the @@ -354,10 +353,11 @@ 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, an input copy with a reader in a later - split, any backend that cannot transfer asynchronously or record events -- - keeps the ordered path untouched. +- Only persistent host inputs marked with `GGML_TENSOR_FLAG_TRANSPORT` are + candidates. The stable prefix remains a per-evaluation value. Unmarked inputs, + weights, user inputs, transposed V, and copies with later readers stay ordered. +- CUDA is the only enabled backend. Meta, SYCL, WebGPU, and other backends stay + ordered until their event behavior and transport path are validated. - 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 @@ -374,22 +374,9 @@ transport never enabled because the scheduler is given a depth of 0. ## Tensor parallelism -`-sm tensor` is not pipelined. The scheduler sees a single *meta* backend there, -and two separate things stand in the way: - -1. **No events in the meta layer.** `ggml_backend_meta_i.event_record` and - `.event_wait` are null, and so are `event_new` / `event_free` / - `event_synchronize` on the meta device. Pipelining is built on ordering a - transfer stream against the consumer with events, so the eligibility check - rejects the backend and the ordered path runs. Requesting a look-ahead under - `-sm tensor` costs nothing and changes nothing: measured 21.85 t/s at depth 1 - against 21.87 at depth 0, with identical output. -2. **The ring is a byte arena.** It is allocated once and the staged input copies - are pointed into it at fixed offsets. A meta buffer has no flat base -- - `ggml_backend_meta_buffer_get_base` returns a placeholder -- and a tensor in - one is not placed at an offset but built as a set of per-device tensors by the - buffer's `init_tensor`, from a whole `ggml_context` allocated at once. The ring - would have to become slot *tensors* rather than offsets. +`-sm tensor` is not pipelined. The scheduler explicitly excludes meta devices. +A host-resident cache needs a validated strided head-split write before this can +be enabled. Both sit behind a correctness problem that is not this feature's: **`-sm tensor` together with `--no-kv-offload` currently produces wrong output.** @@ -420,11 +407,4 @@ row rather than laid out end to end. - Fix `-sm tensor` with `--no-kv-offload` (above). Until then it should not be used: it is wrong rather than slow. -- Events on the meta backend and device, a ring that can live in a meta buffer, - and a transfer-only meta backend, so tensor parallelism can be pipelined once - it is correct. Written on a separate branch, and not reachable until it is. -- Take the last small blocking copies off the host's critical path. On the - pipelined path 3.6 ms per graph at 19,246 and 8.7 ms at 48,042 is still spent - inside `ggml_backend_tensor_copy`, for 0.4 MiB. It is 40 separate copies, and - 32 of them are the device-to-host KV store, which runs on the other copy engine - and could overlap the deliveries instead of blocking the host. +- Add the strided head-split delivery above, validate it, and then measure it. diff --git a/docs/repro/r4-kv-pipeline-exact.py b/docs/repro/r4-kv-pipeline-exact.py index cff45a397a6..55278ce3534 100644 --- a/docs/repro/r4-kv-pipeline-exact.py +++ b/docs/repro/r4-kv-pipeline-exact.py @@ -1,10 +1,11 @@ # 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. +# Run through r4-kv-pipeline-exact.sh. The pipelined path must reproduce depth 0 exactly. import hashlib, json, sys, urllib.request PORT = sys.argv[1] LENGTHS = [int(x) for x in sys.argv[2].split(",")] # approximate prefill tokens +RESULTS_PATH = sys.argv[3] +RESULTS = [] # 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. @@ -72,7 +73,9 @@ def ask(label, prompt, ntok, want_prefill): # produces is not comparable to a fresh prefill, so say so rather than reporting it silently prompt_n = t.get("prompt_n") or 0 reused = prompt_n < want_prefill // 2 - print(f"{label:<18} {hashlib.sha256(text.encode()).hexdigest()[:16]} " + digest = hashlib.sha256(text.encode()).hexdigest()[:16] + RESULTS.append(f"{label} {digest}\n") + print(f"{label:<18} {digest} " f"prompt_n={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}" f"{' CACHE_REUSE' if reused else ''}", flush=True) @@ -84,4 +87,6 @@ def ask(label, prompt, ntok, want_prefill): for name in CORPORA: prompt = nonce(name, length) + filler(name, length) + "\n\n" + QUESTIONS[name] ok &= ask(f"{name}@{length}", prompt, ntok, length) +with open(RESULTS_PATH, "w", encoding="utf-8") as f: + f.writelines(RESULTS) 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 6ab8d7b6b61..e0046199c39 100755 --- a/docs/repro/r4-kv-pipeline-exact.sh +++ b/docs/repro/r4-kv-pipeline-exact.sh @@ -1,7 +1,6 @@ #!/bin/bash # 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. +# prefill corpora and prefill lengths. The script compares every requested depth with the first. # # 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). @@ -16,6 +15,7 @@ HERE="$(cd "$(dirname "$0")" && pwd)" DEPTHS=(0 1 4); [ $# -gt 0 ] && DEPTHS=("$@") rc=0 +BASE="" for D in "${DEPTHS[@]}"; do echo "== pipeline depth=$D ctx=$CTX prefill lengths=$LENGTHS" LOG=$(mktemp /tmp/r4-kv-pipeline.XXXX.log) @@ -28,8 +28,20 @@ for D in "${DEPTHS[@]}"; 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" "$LENGTHS" || rc=$? - kill $SRV 2>/dev/null; wait $SRV 2>/dev/null + OUT=$(mktemp /tmp/r4-kv-pipeline.XXXX.hashes) + if python3 "$HERE/r4-kv-pipeline-exact.py" "$PORT" "$LENGTHS" "$OUT"; then + if [ -z "$BASE" ]; then + BASE="$OUT" + elif ! cmp -s "$BASE" "$OUT"; then + diff -u "$BASE" "$OUT" + rc=1 + fi + else + rc=$? + fi + kill "$SRV" 2>/dev/null; wait "$SRV" 2>/dev/null rm -f "$LOG" + [ "$OUT" = "$BASE" ] || rm -f "$OUT" done +[ -z "$BASE" ] || rm -f "$BASE" exit $rc diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index 924db4eec61..f3a9516a970 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -332,22 +332,23 @@ extern "C" { // 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. + // Only persistent host inputs marked with GGML_TENSOR_FLAG_TRANSPORT are eligible. Their + // stable prefix must be current before each evaluation. The producer must be the CPU or the + // same backend stream that consumes the late region. // // `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); + // allocated. Returns false after graph allocation starts. + GGML_API bool 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); + // Returns false after graph allocation starts. Configuration is immutable then. + GGML_API bool 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/include/ggml.h b/ggml/include/ggml.h index 66f7618c541..a4a18f42d83 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -648,11 +648,12 @@ extern "C" { // this tensor... enum ggml_tensor_flag { - GGML_TENSOR_FLAG_INPUT = 1, // ...is an input for the GGML compute graph - GGML_TENSOR_FLAG_OUTPUT = 2, // ...is an output for the GGML compute graph - GGML_TENSOR_FLAG_PARAM = 4, // ...contains trainable parameters - GGML_TENSOR_FLAG_LOSS = 8, // ...defines loss for numerical optimization (multiple loss tensors add up) - GGML_TENSOR_FLAG_COMPUTE = 16, // ...must be computed + GGML_TENSOR_FLAG_INPUT = 1, // ...is an input for the GGML compute graph + GGML_TENSOR_FLAG_OUTPUT = 2, // ...is an output for the GGML compute graph + GGML_TENSOR_FLAG_PARAM = 4, // ...contains trainable parameters + GGML_TENSOR_FLAG_LOSS = 8, // ...defines loss for numerical optimization (multiple loss tensors add up) + GGML_TENSOR_FLAG_COMPUTE = 16, // ...must be computed + GGML_TENSOR_FLAG_TRANSPORT = 32, // ...is persistent host storage that can use split-input transport }; enum ggml_tri_type { @@ -701,16 +702,11 @@ extern "C" { void * extra; // extra things e.g. for ggml-cuda.cu - // 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; + // leading bytes that stay unchanged during the current graph evaluation + union { + size_t stable_prefix; + char padding[8]; + }; }; static const size_t GGML_TENSOR_SIZE = sizeof(struct ggml_tensor); diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 60332705ec4..fe58ea3bb7a 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -426,10 +426,7 @@ 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. - // 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); + static constexpr size_t nbtc = GGML_TENSOR_SIZE - sizeof(ggml_tensor::padding); std::map, std::pair> split_state_cache; int debug; diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index fc55b438a79..caf707b33fb 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 @@ -817,8 +818,7 @@ struct ggml_backend_sched_transport_ring { 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 + bool reported_no_room; }; // Pipelined delivery of host-resident split inputs. @@ -834,6 +834,7 @@ struct ggml_backend_sched_transport { 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 + bool config_locked; struct ggml_backend_sched_transport_ring rings[GGML_SCHED_MAX_BACKENDS]; @@ -844,7 +845,6 @@ struct ggml_backend_sched_transport { int plan_n_splits; int plan_n_inputs; int n_staged; // over all rings, so that execution can skip the machinery entirely - int n_rings_used; // 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 @@ -868,8 +868,8 @@ struct ggml_backend_sched_transport { // why the look-ahead stopped: the depth it was given, or a slot whose reader has not run yet. // The two want opposite fixes, so they are counted apart. int64_t n_stop_depth; - int64_t n_stop_recycle; - int64_t p_stop_depth, p_stop_recycle; + int64_t n_wait_recycle; + int64_t p_stop_depth, p_wait_recycle; int64_t n_bytes_ordered; // what the ordered blocking copies still move int64_t p_bytes_ordered; @@ -1667,7 +1667,7 @@ static bool ggml_backend_sched_transport_ring_enabled(ggml_backend_sched_t sched return false; } const struct ggml_backend_sched_transport_ring * r = &tr->rings[backend_id]; - return r->eligible && !r->over_budget; + return r->eligible; } static bool ggml_backend_sched_transport_enabled(ggml_backend_sched_t sched) { @@ -1715,6 +1715,11 @@ static bool ggml_backend_sched_input_can_stage( } struct ggml_tensor * input = split->inputs[input_id]; + const struct ggml_tensor * base = input->view_src ? input->view_src : input; + + if (!(base->flags & GGML_TENSOR_FLAG_TRANSPORT)) { + return false; + } // user inputs must be copied immediately, before the user can overwrite them if (input->flags & GGML_TENSOR_FLAG_INPUT) { @@ -1788,6 +1793,52 @@ static void ggml_backend_sched_transport_release_ring(ggml_backend_sched_t sched r->n_staged = 0; } +static void ggml_backend_sched_transport_decline_backend(ggml_backend_sched_t sched, int backend_id) { + struct ggml_backend_sched_transport * tr = &sched->transport; + + ggml_backend_sched_transport_release_ring(sched, backend_id, true); + + if (tr->split_order == NULL || tr->split_input_ofs == NULL || tr->input_staged == NULL) { + return; + } + + for (int i = 0; i < tr->plan_n_splits; i++) { + if (sched->splits[i].backend_id != backend_id) { + 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; + } + } +} + +static bool ggml_backend_sched_size_add(size_t a, size_t b, size_t * result) { + if (a > SIZE_MAX - b) { + return false; + } + *result = a + b; + return true; +} + +static bool ggml_backend_sched_size_mul(size_t a, size_t b, size_t * result) { + if (a != 0 && b > SIZE_MAX/a) { + return false; + } + *result = a*b; + return true; +} + +static bool ggml_backend_sched_size_pad(size_t size, size_t alignment, size_t * result) { + GGML_ASSERT(alignment > 0); + const size_t rem = size % alignment; + if (rem == 0) { + *result = size; + return true; + } + return ggml_backend_sched_size_add(size, alignment - rem, result); +} + // 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. @@ -1837,8 +1888,7 @@ static bool ggml_backend_sched_transport_ensure_backend(ggml_backend_sched_t sch 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->n_rings_used = 0; + tr->n_staged = 0; for (int i = 0; i < sched->n_backends; i++) { tr->rings[i].n_staged = 0; tr->rings[i].consumed = 0; @@ -1939,6 +1989,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { // the ring for every large -c even when the window never gets there. size_t slot_size[GGML_SCHED_MAX_BACKENDS] = { 0 }; size_t slot_size_max[GGML_SCHED_MAX_BACKENDS] = { 0 }; + bool size_overflow[GGML_SCHED_MAX_BACKENDS] = { false }; for (int i = 0; i < sched->n_splits; i++) { struct ggml_backend_sched_split * split = &sched->splits[i]; const int bid = split->backend_id; @@ -1953,11 +2004,18 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { } const struct ggml_tensor * input = split->inputs[j]; const struct ggml_tensor * base = input->view_src ? input->view_src : input; - need += GGML_PAD(ggml_nbytes(input), tr->rings[bid].alignment); - need_max += GGML_PAD(ggml_nbytes(base), tr->rings[bid].alignment); + size_t input_size; + size_t input_size_max; + if (!ggml_backend_sched_size_pad(ggml_nbytes(input), tr->rings[bid].alignment, &input_size) || + !ggml_backend_sched_size_pad(ggml_nbytes(base), tr->rings[bid].alignment, &input_size_max) || + !ggml_backend_sched_size_add(need, input_size, &need) || + !ggml_backend_sched_size_add(need_max, input_size_max, &need_max)) { + size_overflow[bid] = true; + break; + } } - if (need == 0) { + if (need == 0 || size_overflow[bid]) { continue; } @@ -1968,23 +2026,27 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { for (int bid = 0; bid < sched->n_backends; bid++) { struct ggml_backend_sched_transport_ring * r = &tr->rings[bid]; + if (size_overflow[bid]) { + GGML_LOG_WARN("%s: transport ring size overflow on %s, staying on the ordered path\n", __func__, ggml_backend_name(sched->backends[bid])); + ggml_backend_sched_transport_decline_backend(sched, bid); + continue; + } if (r->n_staged == 0) { continue; } - const size_t ring_size = slot_size[bid] * tr->n_slots; - const size_t ring_size_max = slot_size_max[bid] * tr->n_slots; + size_t ring_size; + size_t ring_size_max; + if (!ggml_backend_sched_size_mul(slot_size[bid], tr->n_slots, &ring_size)) { + GGML_LOG_WARN("%s: transport ring size overflow on %s, staying on the ordered path\n", __func__, ggml_backend_name(sched->backends[bid])); + ggml_backend_sched_transport_decline_backend(sched, bid); + continue; + } + if (!ggml_backend_sched_size_mul(slot_size_max[bid], tr->n_slots, &ring_size_max)) { + ring_size_max = SIZE_MAX; + } - // 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 (tr->budget > 0 && ring_size > tr->budget) { if (!r->reported_no_room) { GGML_LOG_WARN("%s: transport ring on %s needs %zu MiB now and %zu MiB at the full " "context, against a %zu MiB budget, staying on the ordered path (raise " @@ -1993,23 +2055,12 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { ring_size_max >> 20, tr->budget >> 20); r->reported_no_room = true; } - 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; - } - } + ggml_backend_sched_transport_decline_backend(sched, bid); continue; } - // nothing has been allocated for this ring until here if (!ggml_backend_sched_transport_ensure_backend(sched, bid)) { - r->n_staged = 0; + ggml_backend_sched_transport_decline_backend(sched, bid); continue; } @@ -2025,7 +2076,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { 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 (dev_free > 0 && (dev_free <= GGML_SCHED_TRANSPORT_HEADROOM || ring_size > dev_free - GGML_SCHED_TRANSPORT_HEADROOM)) { 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__, @@ -2033,17 +2084,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { 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; - } - } + ggml_backend_sched_transport_decline_backend(sched, bid); continue; } @@ -2052,16 +2093,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { 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; - } - } + ggml_backend_sched_transport_decline_backend(sched, bid); continue; } ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_COMPUTE); @@ -2076,7 +2108,6 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { } tr->n_staged += r->n_staged; - tr->n_rings_used++; } if (tr->n_staged == 0) { @@ -2130,7 +2161,9 @@ 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 = r->buffer; - offset += GGML_PAD(ggml_nbytes(split->inputs[j]), r->alignment); + size_t input_size; + GGML_ASSERT(ggml_backend_sched_size_pad(ggml_nbytes(split->inputs[j]), r->alignment, &input_size)); + GGML_ASSERT(ggml_backend_sched_size_add(offset, input_size, &offset)); } GGML_ASSERT(offset <= r->slot_size); } @@ -2168,7 +2201,7 @@ static void ggml_backend_sched_transport_prefetch(ggml_backend_sched_t sched, in if (slot->release_armed) { ggml_backend_event_wait(r->transfer, slot->release); slot->release_armed = false; - tr->n_stop_recycle++; + tr->n_wait_recycle++; } for (int j = 0; j < split->n_inputs; j++) { @@ -2566,7 +2599,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s const double n = 128.0; GGML_LOG_INFO("%s: per graph over %d: total %.2f ms, sync %.2f ms, ordered copy %.2f ms, " "issue %.2f ms, early %.1f MiB, late %.1f MiB, ordered %.1f MiB, " - "stops depth/recycle %.1f/%.1f\n", + "stops on depth %.1f, recycle waits %.1f\n", __func__, (int) n, (tr->t_graph_us - tr->p_graph_us)/1e3/n, (tr->t_sync_us - tr->p_sync_us )/1e3/n, @@ -2576,7 +2609,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s (tr->n_bytes_late - tr->p_bytes_late )/1048576.0/n, (tr->n_bytes_ordered - tr->p_bytes_ordered)/1048576.0/n, (tr->n_stop_depth - tr->p_stop_depth )/n, - (tr->n_stop_recycle - tr->p_stop_recycle)/n); + (tr->n_wait_recycle - tr->p_wait_recycle)/n); tr->p_graph_us = tr->t_graph_us; tr->p_sync_us = tr->t_sync_us; @@ -2587,7 +2620,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s tr->p_bytes_ordered = tr->n_bytes_ordered; tr->named_ordered = true; tr->p_stop_depth = tr->n_stop_depth; - tr->p_stop_recycle = tr->n_stop_recycle; + tr->p_wait_recycle = tr->n_wait_recycle; } } @@ -2680,12 +2713,24 @@ static void ggml_backend_sched_transport_teardown(ggml_backend_sched_t sched) { } -void ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, int depth) { +bool ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, int depth) { GGML_ASSERT(sched); + struct ggml_backend_sched_transport * tr = &sched->transport; + if (tr->config_locked) { + return false; + } + const char * env = getenv("GGML_KV_PIPELINE_DEPTH"); if (env != NULL) { - depth = atoi(env); + char * end = NULL; + errno = 0; + const long value = strtol(env, &end, 10); + if (errno != 0 || end == env || *end != '\0' || value < 0 || value > GGML_SCHED_MAX_TRANSPORT_SLOTS - GGML_SCHED_TRANSPORT_MARGIN) { + GGML_LOG_ERROR("%s: invalid GGML_KV_PIPELINE_DEPTH value: %s\n", __func__, env); + return false; + } + depth = (int) value; } depth = std::min(depth, GGML_SCHED_MAX_TRANSPORT_SLOTS - GGML_SCHED_TRANSPORT_MARGIN); @@ -2693,12 +2738,9 @@ void ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, depth = 0; } - struct ggml_backend_sched_transport * tr = &sched->transport; - ggml_backend_sched_transport_teardown(sched); 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; } @@ -2706,31 +2748,29 @@ void ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, tr->n_slots = depth + GGML_SCHED_TRANSPORT_MARGIN; if (depth < 1) { - return; + return true; } - // 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]; + ggml_backend_dev_t dev = ggml_backend_get_device(backend); + if (dev == NULL || ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_META) { + continue; + } + + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev); + if (reg == NULL || strcmp(ggml_backend_reg_name(reg), "CUDA") != 0) { + continue; + } + 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) { - // Worth naming rather than folding into the generic message below: a backend that - // fans out over several devices -- the meta backend used for tensor parallelism -- - // lands here because that layer implements no events, and ordering a transfer stream - // against the consumer is what this is built on. It keeps the ordered path. - if (tr->debug > 0) { - GGML_LOG_INFO("%s: %s cannot order streams with events, staying on the ordered path\n", - __func__, ggml_backend_name(backend)); - } + if (dev->iface.event_new == NULL) { continue; } if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) { @@ -2754,27 +2794,39 @@ void ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, } if (n_eligible == 0 && tr->debug > 0) { - GGML_LOG_INFO("%s: no backend supports pipelined host transport, staying on the ordered path\n", __func__); + GGML_LOG_INFO("%s: no CUDA backend supports pipelined host transport, staying on the ordered path\n", __func__); } + + return true; } -void ggml_backend_sched_set_transport_pipeline_budget(ggml_backend_sched_t sched, size_t bytes) { +bool ggml_backend_sched_set_transport_pipeline_budget(ggml_backend_sched_t sched, size_t bytes) { GGML_ASSERT(sched); + if (sched->transport.config_locked) { + return false; + } + const char * env = getenv("GGML_KV_PIPELINE_BUDGET_MIB"); if (env != NULL) { - bytes = (size_t) strtoull(env, NULL, 10) * 1024 * 1024; + char * end = NULL; + errno = 0; + const unsigned long long value = strtoull(env, &end, 10); + if (errno != 0 || end == env || *end != '\0' || value > SIZE_MAX/(1024u*1024u)) { + GGML_LOG_ERROR("%s: invalid GGML_KV_PIPELINE_BUDGET_MIB value: %s\n", __func__, env); + return false; + } + bytes = (size_t) value*(1024u*1024u); } if (sched->transport.budget != bytes) { 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); } } + + return true; } void ggml_backend_sched_get_transport_pipeline_stats( @@ -2884,6 +2936,8 @@ bool ggml_backend_sched_alloc_graph(ggml_backend_sched_t sched, struct ggml_cgra GGML_ASSERT((int)sched->hash_set.size >= graph->n_nodes + graph->n_leafs); GGML_ASSERT(!sched->is_alloc); + sched->transport.config_locked = true; + sched->cur_copy = sched->next_copy; sched->next_copy = (sched->next_copy + 1) % sched->n_copies; diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 6f3ee0f5a74..86480ea9ce2 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1260,6 +1260,9 @@ static_assert(GGML_GLU_OP_COUNT == 6, "GGML_GLU_OP_COUNT != 6"); static_assert(sizeof(struct ggml_object)%GGML_MEM_ALIGN == 0, "ggml_object size must be a multiple of GGML_MEM_ALIGN"); static_assert(sizeof(struct ggml_tensor)%GGML_MEM_ALIGN == 0, "ggml_tensor size must be a multiple of GGML_MEM_ALIGN"); +static_assert(sizeof(((struct ggml_tensor *) 0)->padding) == 8, "ggml_tensor trailing storage must be 8 bytes"); +static_assert(sizeof(((struct ggml_tensor *) 0)->stable_prefix) <= sizeof(((struct ggml_tensor *) 0)->padding), "stable_prefix must fit in trailing storage"); +static_assert(offsetof(struct ggml_tensor, stable_prefix) + sizeof(((struct ggml_tensor *) 0)->padding) == sizeof(struct ggml_tensor), "ggml_tensor trailing storage must remain last"); //////////////////////////////////////////////////////////////////////////////// @@ -1830,7 +1833,7 @@ static struct ggml_tensor * ggml_new_tensor_impl( /*.data =*/ obj_alloc_size > 0 ? (void *)(result + 1) : data, /*.name =*/ { 0 }, /*.extra =*/ NULL, - /*.stable_prefix=*/ 0, + /*.padding =*/ { 0 }, }; // TODO: this should not be needed as long as we don't rely on aligned SIMD loads diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 4ff80425ec6..e8bc5b2e46a 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -867,15 +867,22 @@ void llama_context::sched_reserve(uint32_t n_tokens_req, uint32_t n_kv_req) { } auto create_sched = [&](bool pipeline_parallel) { + constexpr size_t mib = 1024u*1024u; + if (cparams.kv_pipeline_depth > 14) { + throw std::invalid_argument("kv_pipeline_depth must be between 0 and 14"); + } + if (cparams.kv_pipeline_budget_mib > std::numeric_limits::max()/mib) { + throw std::invalid_argument("kv_pipeline_budget_mib is too large for this platform"); + } + sched.reset(ggml_backend_sched_new( backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, pipeline_parallel, cparams.op_offload)); cparams.flash_attn_causal_prefix_supported = llama_sched_supports_flash_attn_causal_prefix(sched.get()); - // 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 (!ggml_backend_sched_set_transport_pipeline_budget(sched.get(), (size_t) cparams.kv_pipeline_budget_mib*mib) || + !ggml_backend_sched_set_transport_pipeline_depth(sched.get(), cparams.kv_cpu_pinned || !cparams.offload_kqv ? (int) cparams.kv_pipeline_depth : 0)) { + throw std::invalid_argument("invalid KV transport pipeline configuration"); + } if (sched_resizable) { sched_buffers_shared = sched_buffer_owner != nullptr && sched_buffer_owner->get_sched() != nullptr && ggml_backend_sched_set_resizable(sched.get(), sched_buffer_owner->get_sched()); diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 3ca115f2711..fd0a015fd3b 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -312,6 +312,15 @@ llama_kv_cache::llama_kv_cache( ggml_tensor * k = has_k ? ggml_new_tensor_3d(ctx, type_k, n_embd_k_gqa, kv_size, n_stream) : nullptr; ggml_tensor * v = has_v ? ggml_new_tensor_3d(ctx, type_v, n_embd_v_gqa, kv_size, n_stream) : nullptr; + if (ggml_backend_buft_is_host(buft)) { + if (k) { + k->flags |= GGML_TENSOR_FLAG_TRANSPORT; + } + if (v && !v_trans) { + v->flags |= GGML_TENSOR_FLAG_TRANSPORT; + } + } + bool k_store_quantize = false; bool v_store_quantize = false; if (ggml_backend_buft_is_host(buft)) { diff --git a/tests/test-alloc.cpp b/tests/test-alloc.cpp index 6938f20f16c..e7b3190ddd0 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -20,6 +20,17 @@ struct dummy_backend_context { bool fail_alloc = false; bool unique_alloc_addresses = false; int graph_compute_count = 0; + enum ggml_backend_dev_type device_type = GGML_BACKEND_DEVICE_TYPE_CPU; + const char * registry_name = "dummy"; + bool buffer_is_host = true; + bool fail_backend_init = false; + bool fail_event_init = false; + int transfer_backend_count = 0; + int event_wait_count = 0; + int set_tensor_async_count = 0; + size_t set_tensor_async_bytes = 0; + ggml_backend_buffer_type_t buffer_type = nullptr; + ggml_backend_i backend_interface = {}; ggml_backend_buffer_i buffer_interface; std::vector buffers; @@ -63,8 +74,8 @@ static size_t dummy_backend_buffer_type_get_max_size(ggml_backend_buffer_type_t return ctx->max_buffer_size; } -static bool dummy_backend_buffer_type_is_host(ggml_backend_buffer_type_t) { - return true; +static bool dummy_backend_buffer_type_is_host(ggml_backend_buffer_type_t buft) { + return ((dummy_backend_context *) buft->context)->buffer_is_host; } // ggml_backend_buffer interface @@ -104,13 +115,34 @@ static void dummy_backend_buffer_clear(ggml_backend_buffer_t, uint8_t) {} struct dummy_backend { std::unique_ptr context; + std::unique_ptr registry; std::unique_ptr device; std::unique_ptr handle; ggml_backend_buffer_type buffer_type; }; -static const char * dummy_backend_get_name(ggml_backend_t) { - return "dummy_backend"; +static const char * dummy_backend_get_name(ggml_backend_t backend) { + return ((dummy_backend_context *) backend->context)->registry_name; +} + +static void dummy_backend_free(ggml_backend_t backend) { + dummy_backend_context * ctx = (dummy_backend_context *) backend->context; + ctx->transfer_backend_count--; + delete backend; +} + +static void dummy_backend_set_tensor_async(ggml_backend_t backend, ggml_tensor *, const void *, size_t, size_t size) { + dummy_backend_context * ctx = (dummy_backend_context *) backend->context; + ctx->set_tensor_async_count++; + ctx->set_tensor_async_bytes += size; +} + +static void dummy_backend_synchronize(ggml_backend_t) {} + +static void dummy_backend_event_record(ggml_backend_t, ggml_backend_event_t) {} + +static void dummy_backend_event_wait(ggml_backend_t backend, ggml_backend_event_t) { + ((dummy_backend_context *) backend->context)->event_wait_count++; } static enum ggml_status dummy_backend_graph_compute(ggml_backend_t backend, ggml_cgraph *) { @@ -119,24 +151,76 @@ static enum ggml_status dummy_backend_graph_compute(ggml_backend_t backend, ggml return GGML_STATUS_SUCCESS; } -static enum ggml_backend_dev_type dummy_backend_device_get_type(ggml_backend_dev_t) { - return GGML_BACKEND_DEVICE_TYPE_CPU; +static enum ggml_backend_dev_type dummy_backend_device_get_type(ggml_backend_dev_t dev) { + return ((dummy_backend_context *) dev->context)->device_type; +} + +static void dummy_backend_device_get_memory(ggml_backend_dev_t, size_t * free, size_t * total) { + *free = SIZE_MAX; + *total = SIZE_MAX; +} + +static ggml_backend_t dummy_backend_device_init(ggml_backend_dev_t dev, const char *) { + dummy_backend_context * ctx = (dummy_backend_context *) dev->context; + if (ctx->fail_backend_init) { + return nullptr; + } + ggml_backend_t backend = new ggml_backend{}; + backend->iface = ctx->backend_interface; + backend->device = dev; + backend->context = ctx; + ctx->transfer_backend_count++; + return backend; +} + +static ggml_backend_buffer_type_t dummy_backend_device_get_buffer_type(ggml_backend_dev_t dev) { + return ((dummy_backend_context *) dev->context)->buffer_type; +} + +static ggml_backend_event_t dummy_backend_device_event_new(ggml_backend_dev_t dev) { + dummy_backend_context * ctx = (dummy_backend_context *) dev->context; + if (ctx->fail_event_init) { + return nullptr; + } + ggml_backend_event_t event = new ggml_backend_event; + event->device = dev; + event->context = nullptr; + return event; +} + +static void dummy_backend_device_event_free(ggml_backend_dev_t, ggml_backend_event_t event) { + delete event; +} + +static void dummy_backend_device_event_synchronize(ggml_backend_dev_t, ggml_backend_event_t) {} + +static const char * dummy_backend_registry_get_name(ggml_backend_reg_t reg) { + return ((dummy_backend_context *) reg->context)->registry_name; } static bool dummy_backend_device_supports_op(ggml_backend_dev_t, const ggml_tensor *) { return true; } -static bool dummy_backend_device_supports_buft(ggml_backend_dev_t, ggml_backend_buffer_type_t) { - return true; +static bool dummy_backend_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { + return buft == ((dummy_backend_context *) dev->context)->buffer_type; } -static dummy_backend dummy_backend_init(size_t max_buffer_size, size_t alignment = 8, bool unique_alloc_addresses = false) { +static dummy_backend dummy_backend_init( + size_t max_buffer_size, + size_t alignment = 8, + bool unique_alloc_addresses = false, + enum ggml_backend_dev_type device_type = GGML_BACKEND_DEVICE_TYPE_CPU, + const char * registry_name = "dummy", + bool buffer_is_host = true) { dummy_backend b{}; b.context = std::make_unique(); b.context->alignment = alignment; b.context->max_buffer_size = max_buffer_size; b.context->unique_alloc_addresses = unique_alloc_addresses; + b.context->device_type = device_type; + b.context->registry_name = registry_name; + b.context->buffer_is_host = buffer_is_host; b.context->buffer_interface.free_buffer = dummy_backend_buffer_free_buffer; b.context->buffer_interface.get_base = dummy_backend_buffer_get_base; @@ -152,17 +236,36 @@ static dummy_backend dummy_backend_init(size_t max_buffer_size, size_t alignment b.buffer_type.iface.get_alignment = dummy_backend_buffer_type_get_alignment; b.buffer_type.iface.get_max_size = dummy_backend_buffer_type_get_max_size; b.buffer_type.iface.is_host = dummy_backend_buffer_type_is_host; + b.context->buffer_type = &b.buffer_type; + + b.registry = std::make_unique(); + b.registry->iface.get_name = dummy_backend_registry_get_name; + b.registry->context = b.context.get(); b.device = std::make_unique(); - b.device->iface.get_type = dummy_backend_device_get_type; - b.device->iface.supports_op = dummy_backend_device_supports_op; - b.device->iface.supports_buft = dummy_backend_device_supports_buft; + b.device->iface.get_memory = dummy_backend_device_get_memory; + b.device->iface.get_type = dummy_backend_device_get_type; + b.device->iface.init_backend = dummy_backend_device_init; + b.device->iface.get_buffer_type = dummy_backend_device_get_buffer_type; + b.device->iface.supports_op = dummy_backend_device_supports_op; + b.device->iface.supports_buft = dummy_backend_device_supports_buft; + b.device->iface.event_new = dummy_backend_device_event_new; + b.device->iface.event_free = dummy_backend_device_event_free; + b.device->iface.event_synchronize = dummy_backend_device_event_synchronize; + b.device->reg = b.registry.get(); + b.device->context = b.context.get(); b.handle = std::make_unique(); - b.handle->iface.get_name = dummy_backend_get_name; - b.handle->iface.graph_compute = dummy_backend_graph_compute; - b.handle->device = b.device.get(); - b.handle->context = b.context.get(); + b.context->backend_interface.get_name = dummy_backend_get_name; + b.context->backend_interface.free = dummy_backend_free; + b.context->backend_interface.set_tensor_async = dummy_backend_set_tensor_async; + b.context->backend_interface.synchronize = dummy_backend_synchronize; + b.context->backend_interface.graph_compute = dummy_backend_graph_compute; + b.context->backend_interface.event_record = dummy_backend_event_record; + b.context->backend_interface.event_wait = dummy_backend_event_wait; + b.handle->iface = b.context->backend_interface; + b.handle->device = b.device.get(); + b.handle->context = b.context.get(); return b; } @@ -186,6 +289,36 @@ static test_context_with_graph make_context() { return { ctx, graph, std::move(ctx_ptr) }; } +struct transport_graph { + test_context_with_graph ctx; + ggml_backend_buffer_ptr buffer; + ggml_tensor * source; + ggml_tensor * output; +}; + +static transport_graph make_transport_graph(dummy_backend & cpu, size_t size) { + GGML_ASSERT(size % sizeof(float) == 0); + auto result = make_context(); + ggml_tensor * source = ggml_new_tensor_1d(result.ctx, GGML_TYPE_F32, size/sizeof(float)); + ggml_tensor * output = ggml_scale(result.ctx, source, 2.0f); + source->flags |= GGML_TENSOR_FLAG_TRANSPORT; + ggml_build_forward_expand(result.graph, output); + + ggml_backend_buffer_ptr buffer(ggml_backend_buft_alloc_buffer(&cpu.buffer_type, size)); + source->buffer = buffer.get(); + source->data = ggml_backend_buffer_get_base(buffer.get()); + + return { std::move(result), std::move(buffer), source, output }; +} + +static void transport_stats( + ggml_backend_sched_t sched, + int64_t * deliveries, + int64_t * early, + int64_t * late) { + ggml_backend_sched_get_transport_pipeline_stats(sched, deliveries, early, late); +} + static ggml_tensor * make_input_1d(ggml_context * ctx, int64_t n_elements) { ggml_tensor * t = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_elements); ggml_set_input(t); @@ -1124,6 +1257,202 @@ static void test_resizable_buffers_owner_borrower_teardown_order() { } } +static void test_transport_prefix_and_configuration() { + dummy_backend cuda = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); + dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); + auto graph = make_transport_graph(cpu, 64); + + ggml_backend_t backends[] = { cuda.handle.get(), cpu.handle.get() }; + ggml_backend_buffer_type_t bufts[] = { &cuda.buffer_type, &cpu.buffer_type }; + { + ggml_backend_sched_ptr sched(ggml_backend_sched_new(backends, bufts, 2, 128, false, false)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_budget(sched.get(), 1024)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_depth(sched.get(), 1)); + ggml_backend_sched_set_tensor_backend(sched.get(), graph.output, cuda.handle.get()); + + ggml_set_stable_prefix(graph.source, 32); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), graph.ctx.graph)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), graph.ctx.graph) == GGML_STATUS_SUCCESS); + + int64_t deliveries = 0; + int64_t early = 0; + int64_t late = 0; + transport_stats(sched.get(), &deliveries, &early, &late); + GGML_ASSERT(deliveries == 1 && early == 32 && late == 32); + GGML_ASSERT(!ggml_backend_sched_set_transport_pipeline_depth(sched.get(), 0)); + GGML_ASSERT(!ggml_backend_sched_set_transport_pipeline_budget(sched.get(), 0)); + + ggml_set_stable_prefix(graph.source, 0); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), graph.ctx.graph) == GGML_STATUS_SUCCESS); + transport_stats(sched.get(), &deliveries, &early, &late); + GGML_ASSERT(deliveries == 2 && early == 32 && late == 96); + + ggml_set_stable_prefix(graph.source, 64); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), graph.ctx.graph) == GGML_STATUS_SUCCESS); + transport_stats(sched.get(), &deliveries, &early, &late); + GGML_ASSERT(deliveries == 3 && early == 96 && late == 96); + GGML_ASSERT(cuda.context->event_wait_count > 0); + } + GGML_ASSERT(cuda.context->transfer_backend_count == 0); +} + +static void test_transport_depth_zero() { + dummy_backend cuda = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); + dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); + auto graph = make_transport_graph(cpu, 64); + ggml_set_stable_prefix(graph.source, 64); + + ggml_backend_t backends[] = { cuda.handle.get(), cpu.handle.get() }; + ggml_backend_buffer_type_t bufts[] = { &cuda.buffer_type, &cpu.buffer_type }; + ggml_backend_sched_ptr sched(ggml_backend_sched_new(backends, bufts, 2, 128, false, false)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_depth(sched.get(), 0)); + ggml_backend_sched_set_tensor_backend(sched.get(), graph.output, cuda.handle.get()); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), graph.ctx.graph)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), graph.ctx.graph) == GGML_STATUS_SUCCESS); + + int64_t deliveries = -1; + int64_t early = -1; + int64_t late = -1; + transport_stats(sched.get(), &deliveries, &early, &late); + GGML_ASSERT(deliveries == 0 && early == 0 && late == 0); + GGML_ASSERT(cuda.context->transfer_backend_count == 0); +} + +static void test_transport_budget_recovers() { + dummy_backend cuda = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); + dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); + auto large = make_transport_graph(cpu, 256); + auto small = make_transport_graph(cpu, 64); + ggml_set_stable_prefix(large.source, 256); + ggml_set_stable_prefix(small.source, 64); + + ggml_backend_t backends[] = { cuda.handle.get(), cpu.handle.get() }; + ggml_backend_buffer_type_t bufts[] = { &cuda.buffer_type, &cpu.buffer_type }; + ggml_backend_sched_ptr sched(ggml_backend_sched_new(backends, bufts, 2, 128, false, false)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_budget(sched.get(), 384)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_depth(sched.get(), 1)); + + ggml_backend_sched_set_tensor_backend(sched.get(), large.output, cuda.handle.get()); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), large.ctx.graph)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), large.ctx.graph) == GGML_STATUS_SUCCESS); + int64_t deliveries = 0; + transport_stats(sched.get(), &deliveries, nullptr, nullptr); + GGML_ASSERT(deliveries == 0); + + ggml_backend_sched_reset(sched.get()); + ggml_backend_sched_set_tensor_backend(sched.get(), small.output, cuda.handle.get()); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), small.ctx.graph)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), small.ctx.graph) == GGML_STATUS_SUCCESS); + transport_stats(sched.get(), &deliveries, nullptr, nullptr); + GGML_ASSERT(deliveries == 1); +} + +static void test_transport_partial_backend_failure() { + dummy_backend cuda_fail = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); + dummy_backend cuda_ok = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); + dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); + cuda_fail.context->fail_event_init = true; + + auto graph = make_context(); + ggml_tensor * source_fail = ggml_new_tensor_1d(graph.ctx, GGML_TYPE_F32, 16); + ggml_tensor * source_ok = ggml_new_tensor_1d(graph.ctx, GGML_TYPE_F32, 16); + ggml_tensor * output_fail = ggml_scale(graph.ctx, source_fail, 2.0f); + ggml_tensor * output_ok = ggml_scale(graph.ctx, source_ok, 2.0f); + ggml_tensor * output = ggml_add(graph.ctx, output_fail, output_ok); + source_fail->flags |= GGML_TENSOR_FLAG_TRANSPORT; + source_ok->flags |= GGML_TENSOR_FLAG_TRANSPORT; + ggml_set_stable_prefix(source_fail, 64); + ggml_set_stable_prefix(source_ok, 64); + ggml_build_forward_expand(graph.graph, output); + + ggml_backend_buffer_ptr buffer_fail(ggml_backend_buft_alloc_buffer(&cpu.buffer_type, 64)); + ggml_backend_buffer_ptr buffer_ok(ggml_backend_buft_alloc_buffer(&cpu.buffer_type, 64)); + source_fail->buffer = buffer_fail.get(); + source_fail->data = ggml_backend_buffer_get_base(buffer_fail.get()); + source_ok->buffer = buffer_ok.get(); + source_ok->data = ggml_backend_buffer_get_base(buffer_ok.get()); + + ggml_backend_t backends[] = { cuda_fail.handle.get(), cuda_ok.handle.get(), cpu.handle.get() }; + ggml_backend_buffer_type_t bufts[] = { &cuda_fail.buffer_type, &cuda_ok.buffer_type, &cpu.buffer_type }; + { + ggml_backend_sched_ptr sched(ggml_backend_sched_new(backends, bufts, 3, 128, false, false)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_depth(sched.get(), 1)); + ggml_backend_sched_set_tensor_backend(sched.get(), output_fail, cuda_fail.handle.get()); + ggml_backend_sched_set_tensor_backend(sched.get(), output_ok, cuda_ok.handle.get()); + ggml_backend_sched_set_tensor_backend(sched.get(), output, cpu.handle.get()); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), graph.graph)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), graph.graph) == GGML_STATUS_SUCCESS); + + int64_t deliveries = 0; + transport_stats(sched.get(), &deliveries, nullptr, nullptr); + GGML_ASSERT(deliveries == 1); + GGML_ASSERT(cuda_fail.context->transfer_backend_count == 0); + GGML_ASSERT(cuda_ok.context->transfer_backend_count == 1); + } + GGML_ASSERT(cuda_ok.context->transfer_backend_count == 0); +} + +static void test_transport_excludes_meta() { + dummy_backend meta = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_META, "CUDA", false); + dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); + auto graph = make_transport_graph(cpu, 64); + ggml_set_stable_prefix(graph.source, 64); + + ggml_backend_t backends[] = { meta.handle.get(), cpu.handle.get() }; + ggml_backend_buffer_type_t bufts[] = { &meta.buffer_type, &cpu.buffer_type }; + ggml_backend_sched_ptr sched(ggml_backend_sched_new(backends, bufts, 2, 128, false, false)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_depth(sched.get(), 1)); + ggml_backend_sched_set_tensor_backend(sched.get(), graph.output, meta.handle.get()); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), graph.ctx.graph)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), graph.ctx.graph) == GGML_STATUS_SUCCESS); + + int64_t deliveries = -1; + transport_stats(sched.get(), &deliveries, nullptr, nullptr); + GGML_ASSERT(deliveries == 0); + GGML_ASSERT(meta.context->transfer_backend_count == 0); +} + +static void test_transport_requires_annotation() { + dummy_backend cuda = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); + dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); + auto graph = make_transport_graph(cpu, 64); + graph.source->flags &= ~GGML_TENSOR_FLAG_TRANSPORT; + ggml_set_stable_prefix(graph.source, 64); + + ggml_backend_t backends[] = { cuda.handle.get(), cpu.handle.get() }; + ggml_backend_buffer_type_t bufts[] = { &cuda.buffer_type, &cpu.buffer_type }; + ggml_backend_sched_ptr sched(ggml_backend_sched_new(backends, bufts, 2, 128, false, false)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_depth(sched.get(), 1)); + ggml_backend_sched_set_tensor_backend(sched.get(), graph.output, cuda.handle.get()); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), graph.ctx.graph)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), graph.ctx.graph) == GGML_STATUS_SUCCESS); + + int64_t deliveries = -1; + transport_stats(sched.get(), &deliveries, nullptr, nullptr); + GGML_ASSERT(deliveries == 0); + GGML_ASSERT(cuda.context->transfer_backend_count == 0); +} + +static void test_transport_excludes_non_cuda() { + dummy_backend sycl = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "SYCL", false); + dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); + auto graph = make_transport_graph(cpu, 64); + ggml_set_stable_prefix(graph.source, 64); + + ggml_backend_t backends[] = { sycl.handle.get(), cpu.handle.get() }; + ggml_backend_buffer_type_t bufts[] = { &sycl.buffer_type, &cpu.buffer_type }; + ggml_backend_sched_ptr sched(ggml_backend_sched_new(backends, bufts, 2, 128, false, false)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_depth(sched.get(), 1)); + ggml_backend_sched_set_tensor_backend(sched.get(), graph.output, sycl.handle.get()); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), graph.ctx.graph)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), graph.ctx.graph) == GGML_STATUS_SUCCESS); + + int64_t deliveries = -1; + transport_stats(sched.get(), &deliveries, nullptr, nullptr); + GGML_ASSERT(deliveries == 0); + GGML_ASSERT(sycl.context->transfer_backend_count == 0); +} + static void run(const char * name, void (*f)()) { printf("%s ", name); fflush(stdout); @@ -1155,5 +1484,12 @@ int main() { run("test_resizable_buffers_owner_borrower_allocation_failure", test_resizable_buffers_owner_borrower_allocation_failure); run("test_resizable_buffers_owner_borrower_scheduler_failure", test_resizable_buffers_owner_borrower_scheduler_failure); run("test_resizable_buffers_owner_borrower_teardown_order", test_resizable_buffers_owner_borrower_teardown_order); + run("test_transport_prefix_and_configuration", test_transport_prefix_and_configuration); + run("test_transport_depth_zero", test_transport_depth_zero); + run("test_transport_budget_recovers", test_transport_budget_recovers); + run("test_transport_partial_backend_failure", test_transport_partial_backend_failure); + run("test_transport_excludes_meta", test_transport_excludes_meta); + run("test_transport_requires_annotation", test_transport_requires_annotation); + run("test_transport_excludes_non_cuda", test_transport_excludes_non_cuda); return 0; } From b71b1cc736d2840f349a69c108c97cfd81693f2c Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 30 Aug 2026 23:03:52 +0200 Subject: [PATCH 07/32] sched: fix pipelined transport review issues Assisted-by: OpenAI Codex --- docs/kv-transport-pipelining.md | 3 +- docs/repro/r4-kv-pipeline-ab.sh | 4 +- docs/repro/r4-kv-pipeline-context-sweep.sh | 2 +- docs/repro/r4-kv-pipeline-exact.sh | 10 +- ggml/src/ggml-backend.cpp | 240 ++++++++++++++------- tests/test-alloc.cpp | 120 +++++++++++ tools/llama-bench/llama-bench.cpp | 41 +++- 7 files changed, 334 insertions(+), 86 deletions(-) diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index 2ab357323a0..fd16e366935 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -369,8 +369,7 @@ transport never enabled because the scheduler is given a depth of 0. [Tensor parallelism](#tensor-parallelism). - 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`. +- `GGML_KV_PIPELINE_DEPTH` and `GGML_KV_PIPELINE_BUDGET_MIB` provide scheduler defaults. Explicit scheduler settings and command-line options take precedence. ## Tensor parallelism diff --git a/docs/repro/r4-kv-pipeline-ab.sh b/docs/repro/r4-kv-pipeline-ab.sh index c1636994a8f..6bb6a381e56 100755 --- a/docs/repro/r4-kv-pipeline-ab.sh +++ b/docs/repro/r4-kv-pipeline-ab.sh @@ -1,6 +1,6 @@ #!/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. +# The two arms are the same binary: --kv-pipeline-depth 0 is the ordered path. # # LLAMA_KV_MODEL=/path/model.gguf docs/repro/r4-kv-pipeline-ab.sh [depth ...] set -u @@ -18,7 +18,7 @@ for opt in kvcp rso; do done 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" \ + taskset -c "$PIN" "$BUILD/bin/llama-bench" -m "$MODEL" --kv-pipeline-depth "$2" \ -ngl 99 -sm none -mg 0 -t 3 -nkvo 1 $BENCH_KV_OPTS \ -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 \ diff --git a/docs/repro/r4-kv-pipeline-context-sweep.sh b/docs/repro/r4-kv-pipeline-context-sweep.sh index 0f1386f9483..79d2f07ee30 100755 --- a/docs/repro/r4-kv-pipeline-context-sweep.sh +++ b/docs/repro/r4-kv-pipeline-context-sweep.sh @@ -24,7 +24,7 @@ arm () { # $1 pipeline depth, $2 context depth, $3 reps ( 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" \ + ts=$(taskset -c "$PIN" "$BUILD/bin/llama-bench" -m "$MODEL" --kv-pipeline-depth "$1" \ -ngl 99 -sm none -mg 0 -t 3 -nkvo 1 $BENCH_KV_OPTS \ -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 \ diff --git a/docs/repro/r4-kv-pipeline-exact.sh b/docs/repro/r4-kv-pipeline-exact.sh index e0046199c39..485a1196cb8 100755 --- a/docs/repro/r4-kv-pipeline-exact.sh +++ b/docs/repro/r4-kv-pipeline-exact.sh @@ -16,10 +16,11 @@ HERE="$(cd "$(dirname "$0")" && pwd)" DEPTHS=(0 1 4); [ $# -gt 0 ] && DEPTHS=("$@") rc=0 BASE="" -for D in "${DEPTHS[@]}"; do +for I in "${!DEPTHS[@]}"; do + D="${DEPTHS[$I]}" 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" \ + taskset -c "$PIN" "$BUILD/bin/llama-server" -m "$MODEL" --kv-pipeline-depth "$D" \ -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 "$CTX" --parallel 1 \ --host 127.0.0.1 --port "$PORT" --no-warmup > "$LOG" 2>&1 & @@ -30,7 +31,7 @@ for D in "${DEPTHS[@]}"; do done OUT=$(mktemp /tmp/r4-kv-pipeline.XXXX.hashes) if python3 "$HERE/r4-kv-pipeline-exact.py" "$PORT" "$LENGTHS" "$OUT"; then - if [ -z "$BASE" ]; then + if [ "$I" -eq 0 ]; then BASE="$OUT" elif ! cmp -s "$BASE" "$OUT"; then diff -u "$BASE" "$OUT" @@ -42,6 +43,9 @@ for D in "${DEPTHS[@]}"; do kill "$SRV" 2>/dev/null; wait "$SRV" 2>/dev/null rm -f "$LOG" [ "$OUT" = "$BASE" ] || rm -f "$OUT" + if [ "$I" -eq 0 ] && [ -z "$BASE" ]; then + break + fi done [ -z "$BASE" ] || rm -f "$BASE" exit $rc diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index caf707b33fb..d89008f2bb1 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1745,6 +1745,54 @@ static bool ggml_backend_sched_input_is_staged(ggml_backend_sched_t sched, int s return tr->input_staged[base + input_id] != 0; } +static bool ggml_backend_sched_size_add(size_t a, size_t b, size_t * result); +static bool ggml_backend_sched_size_pad(size_t size, size_t alignment, size_t * result); + +static void ggml_backend_sched_transport_clear_addresses(ggml_backend_sched_t sched) { + const struct ggml_backend_sched_transport * tr = &sched->transport; + + for (int i = 0; i < tr->plan_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_cpy = tensor_copy(split->inputs[j], split->backend_id, sched->cur_copy); + input_cpy->data = NULL; + input_cpy->buffer = NULL; + } + } +} + +static void ggml_backend_sched_transport_assign_addresses(ggml_backend_sched_t sched) { + const struct ggml_backend_sched_transport * tr = &sched->transport; + + for (int i = 0; i < tr->plan_n_splits; i++) { + if (tr->split_order[i] < 0) { + continue; + } + + struct ggml_backend_sched_split * split = &sched->splits[i]; + struct ggml_backend_sched_transport_ring * r = &sched->transport.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++) { + 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 = r->buffer; + size_t input_size; + GGML_ASSERT(ggml_backend_sched_size_pad(ggml_nbytes(split->inputs[j]), r->alignment, &input_size)); + GGML_ASSERT(ggml_backend_sched_size_add(offset, input_size, &offset)); + } + GGML_ASSERT(offset <= r->slot_size); + } +} + // 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. @@ -1889,13 +1937,15 @@ 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->plan_n_splits = 0; + tr->plan_n_inputs = 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)) { + if (!ggml_backend_sched_transport_enabled(sched) || sched->n_splits == 0) { return; } @@ -1933,54 +1983,89 @@ 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; + int n_candidates = 0; 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; + n_candidates++; } } } - // 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. 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. + if (n_candidates == 0) { + for (int i = 0; i < sched->n_splits; i++) { + tr->split_order[i] = -1; + } + return; + } + + // A ring copy must have one owner and no views or later readers. + // Build one lookup table, then scan each graph node once. + size_t staged_hash_size = n_candidates; + staged_hash_size += staged_hash_size/4 + 1; + struct ggml_hash_set staged_copies = ggml_hash_set_new(staged_hash_size); + int * staged_owner = (int *) malloc(staged_copies.size * sizeof(int)); + GGML_ASSERT(staged_owner != NULL); + for (size_t i = 0; i < staged_copies.size; i++) { + staged_owner[i] = -1; + } + 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 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 && !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) { - disqualified = true; - break; - } - } + struct ggml_tensor * input_cpy = tensor_copy(split->inputs[j], split->backend_id, sched->cur_copy); + const size_t id = ggml_hash_find_or_insert(&staged_copies, input_cpy); + if (staged_owner[id] == -1) { + staged_owner[id] = i; + } else { + staged_owner[id] = -2; + } + } + } + + for (int i = 0; i < sched->n_splits; i++) { + const struct ggml_cgraph * graph = &sched->splits[i].graph; + for (int j = 0; j < graph->n_nodes; j++) { + const struct ggml_tensor * node = graph->nodes[j]; + if (node->view_src != NULL) { + const size_t id = ggml_hash_find(&staged_copies, node->view_src); + if (id != GGML_HASHSET_FULL && ggml_bitset_get(staged_copies.used, id)) { + staged_owner[id] = -2; + } + } + for (int k = 0; k < GGML_MAX_SRC; k++) { + if (node->src[k] == NULL) { + continue; + } + const size_t id = ggml_hash_find(&staged_copies, node->src[k]); + if (id != GGML_HASHSET_FULL && ggml_bitset_get(staged_copies.used, id) && staged_owner[id] >= 0 && i > staged_owner[id]) { + staged_owner[id] = -2; } } + } + } - if (disqualified) { + 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; + } + struct ggml_tensor * input_cpy = tensor_copy(split->inputs[j], split->backend_id, sched->cur_copy); + const size_t id = ggml_hash_find(&staged_copies, input_cpy); + GGML_ASSERT(id != GGML_HASHSET_FULL && ggml_bitset_get(staged_copies.used, id)); + if (staged_owner[id] < 0) { tr->input_staged[tr->split_input_ofs[i] + j] = 0; } } } + free(staged_owner); + ggml_hash_set_free(&staged_copies); // per-ring slot size and delivery order. The budget is applied to what this graph needs, so // that a run whose window stays small keeps the ring whatever -n_ctx says. slot_size_max is @@ -2143,30 +2228,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { } } - 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]; - 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++) { - 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 = r->buffer; - size_t input_size; - GGML_ASSERT(ggml_backend_sched_size_pad(ggml_nbytes(split->inputs[j]), r->alignment, &input_size)); - GGML_ASSERT(ggml_backend_sched_size_add(offset, input_size, &offset)); - } - GGML_ASSERT(offset <= r->slot_size); - } + ggml_backend_sched_transport_assign_addresses(sched); } // Issue the stable prefix of every staged split on this ring that is within the look-ahead of what @@ -2289,10 +2351,12 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { ggml_backend_synchronize(sched->backends[i]); } + ggml_backend_sched_transport_clear_addresses(sched); if (!ggml_gallocr_reserve_n(sched->galloc, &sched->graph, sched->node_backend_ids, sched->leaf_backend_ids)) { GGML_LOG_ERROR("%s: failed to allocate graph\n", __func__); return false; } + ggml_backend_sched_transport_assign_addresses(sched); if (!ggml_gallocr_alloc_graph(sched->galloc, &sched->graph)) { GGML_LOG_ERROR("%s: failed to allocate graph\n", __func__); return false; @@ -2313,6 +2377,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s int prev_backend_id = -1; struct ggml_backend_sched_transport * tr = &sched->transport; + bool named_ordered_now = false; // 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; @@ -2399,6 +2464,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s if (tr->debug >= 3 && !tr->named_ordered) { GGML_LOG_INFO("%s: ordered copy %s %zu KiB from %s\n", __func__, input->name, ggml_nbytes(input) >> 10, ggml_backend_buft_name(input->buffer->buft)); + named_ordered_now = true; } } else { // wait for the split backend to finish using the input before overwriting it @@ -2516,6 +2582,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s if (tr->debug >= 3 && !tr->named_ordered) { GGML_LOG_INFO("%s: ordered copy %s %zu KiB from %s\n", __func__, input->name, ggml_nbytes(input) >> 10, ggml_backend_buft_name(input->buffer->buft)); + named_ordered_now = true; } } } @@ -2589,6 +2656,10 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s prev_backend_id = split_backend_id; } + if (named_ordered_now) { + tr->named_ordered = true; + } + if (tr->debug >= 2) { tr->t_graph_us += ggml_time_us() - t_graph_0; tr->n_graphs++; @@ -2618,7 +2689,6 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s tr->p_bytes_early = tr->n_bytes_early; tr->p_bytes_late = tr->n_bytes_late; tr->p_bytes_ordered = tr->n_bytes_ordered; - tr->named_ordered = true; tr->p_stop_depth = tr->n_stop_depth; tr->p_wait_recycle = tr->n_wait_recycle; } @@ -2627,6 +2697,42 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s return GGML_STATUS_SUCCESS; } +static bool ggml_backend_sched_transport_depth_from_env(int * depth) { + const char * env = getenv("GGML_KV_PIPELINE_DEPTH"); + if (env == NULL) { + return false; + } + + char * end = NULL; + errno = 0; + const long value = strtol(env, &end, 10); + if (errno != 0 || end == env || *end != '\0' || value < 0 || value > GGML_SCHED_MAX_TRANSPORT_SLOTS - GGML_SCHED_TRANSPORT_MARGIN) { + GGML_LOG_WARN("%s: ignoring invalid GGML_KV_PIPELINE_DEPTH value: %s\n", __func__, env); + return false; + } + + *depth = (int) value; + return true; +} + +static bool ggml_backend_sched_transport_budget_from_env(size_t * budget) { + const char * env = getenv("GGML_KV_PIPELINE_BUDGET_MIB"); + if (env == NULL) { + return false; + } + + char * end = NULL; + errno = 0; + const unsigned long long value = strtoull(env, &end, 10); + if (errno != 0 || end == env || *end != '\0' || value > SIZE_MAX/(1024u*1024u)) { + GGML_LOG_WARN("%s: ignoring invalid GGML_KV_PIPELINE_BUDGET_MIB value: %s\n", __func__, env); + return false; + } + + *budget = (size_t) value*(1024u*1024u); + return true; +} + ggml_backend_sched_t ggml_backend_sched_new( ggml_backend_t * backends, ggml_backend_buffer_type_t * bufts, @@ -2694,6 +2800,7 @@ ggml_backend_sched_t ggml_backend_sched_new( sched->galloc = ggml_gallocr_new_n(sched->bufts, n_backends); sched->transport.budget = GGML_SCHED_TRANSPORT_BUDGET; + ggml_backend_sched_transport_budget_from_env(&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; @@ -2702,6 +2809,11 @@ ggml_backend_sched_t ggml_backend_sched_new( ggml_backend_sched_reset(sched); + int transport_depth; + if (ggml_backend_sched_transport_depth_from_env(&transport_depth)) { + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_depth(sched, transport_depth)); + } + return sched; } @@ -2721,18 +2833,6 @@ bool ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, return false; } - const char * env = getenv("GGML_KV_PIPELINE_DEPTH"); - if (env != NULL) { - char * end = NULL; - errno = 0; - const long value = strtol(env, &end, 10); - if (errno != 0 || end == env || *end != '\0' || value < 0 || value > GGML_SCHED_MAX_TRANSPORT_SLOTS - GGML_SCHED_TRANSPORT_MARGIN) { - GGML_LOG_ERROR("%s: invalid GGML_KV_PIPELINE_DEPTH value: %s\n", __func__, env); - return false; - } - depth = (int) value; - } - depth = std::min(depth, GGML_SCHED_MAX_TRANSPORT_SLOTS - GGML_SCHED_TRANSPORT_MARGIN); if (depth < 0) { depth = 0; @@ -2807,18 +2907,6 @@ bool ggml_backend_sched_set_transport_pipeline_budget(ggml_backend_sched_t sched return false; } - const char * env = getenv("GGML_KV_PIPELINE_BUDGET_MIB"); - if (env != NULL) { - char * end = NULL; - errno = 0; - const unsigned long long value = strtoull(env, &end, 10); - if (errno != 0 || end == env || *end != '\0' || value > SIZE_MAX/(1024u*1024u)) { - GGML_LOG_ERROR("%s: invalid GGML_KV_PIPELINE_BUDGET_MIB value: %s\n", __func__, env); - return false; - } - bytes = (size_t) value*(1024u*1024u); - } - if (sched->transport.budget != bytes) { sched->transport.budget = bytes; for (int i = 0; i < sched->n_backends; i++) { diff --git a/tests/test-alloc.cpp b/tests/test-alloc.cpp index e7b3190ddd0..5e453dc297b 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -5,8 +5,10 @@ #include "ggml.h" #include +#include #include #include +#include #include // @@ -1296,6 +1298,121 @@ static void test_transport_prefix_and_configuration() { GGML_ASSERT(cuda.context->transfer_backend_count == 0); } +static void test_transport_empty_graph() { + dummy_backend cuda = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); + dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); + auto graph = make_context(); + + ggml_backend_t backends[] = { cuda.handle.get(), cpu.handle.get() }; + ggml_backend_buffer_type_t bufts[] = { &cuda.buffer_type, &cpu.buffer_type }; + ggml_backend_sched_ptr sched(ggml_backend_sched_new(backends, bufts, 2, 128, false, false)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_depth(sched.get(), 1)); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), graph.graph)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), graph.graph) == GGML_STATUS_SUCCESS); +} + +static size_t transport_fallback_buffer_size(int depth) { + dummy_backend cuda = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); + dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); + auto graph = make_context(); + + ggml_tensor * weight = ggml_new_tensor_2d(graph.ctx, GGML_TYPE_F32, 4, 4); + ggml_tensor * source = ggml_new_tensor_2d(graph.ctx, GGML_TYPE_F32, 4, 1); + ggml_tensor * output = ggml_mul_mat(graph.ctx, weight, source); + source->flags |= GGML_TENSOR_FLAG_TRANSPORT; + ggml_set_stable_prefix(source, ggml_nbytes(source)); + ggml_build_forward_expand(graph.graph, output); + + ggml_backend_buffer_ptr weight_buffer(ggml_backend_buft_alloc_buffer(&cpu.buffer_type, ggml_nbytes(weight))); + ggml_backend_buffer_ptr source_buffer(ggml_backend_buft_alloc_buffer(&cpu.buffer_type, ggml_nbytes(source))); + weight->buffer = weight_buffer.get(); + weight->data = ggml_backend_buffer_get_base(weight_buffer.get()); + source->buffer = source_buffer.get(); + source->data = ggml_backend_buffer_get_base(source_buffer.get()); + + ggml_backend_t backends[] = { cuda.handle.get(), cpu.handle.get() }; + ggml_backend_buffer_type_t bufts[] = { &cuda.buffer_type, &cpu.buffer_type }; + ggml_backend_sched_ptr sched(ggml_backend_sched_new(backends, bufts, 2, 128, false, false)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_depth(sched.get(), depth)); + ggml_backend_sched_set_tensor_backend(sched.get(), output, cuda.handle.get()); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), graph.graph)); + return ggml_backend_sched_get_buffer_size(sched.get(), cuda.handle.get()); +} + +static void test_transport_fallback_keeps_allocator_plan() { + const size_t ordered = transport_fallback_buffer_size(0); + const size_t pipelined = transport_fallback_buffer_size(1); + GGML_ASSERT(ordered > 0 && pipelined == ordered); +} + +static void set_test_env(const char * name, const char * value) { +#ifdef _WIN32 + GGML_ASSERT(_putenv_s(name, value) == 0); +#else + GGML_ASSERT(setenv(name, value, 1) == 0); +#endif +} + +static void restore_test_env(const char * name, bool had_value, const std::string & value) { +#ifdef _WIN32 + GGML_ASSERT(_putenv_s(name, had_value ? value.c_str() : "") == 0); +#else + GGML_ASSERT(had_value ? setenv(name, value.c_str(), 1) == 0 : unsetenv(name) == 0); +#endif +} + +static void test_transport_environment_is_fallback() { + const char * depth_env = getenv("GGML_KV_PIPELINE_DEPTH"); + const char * budget_env = getenv("GGML_KV_PIPELINE_BUDGET_MIB"); + const bool had_depth = depth_env != nullptr; + const bool had_budget = budget_env != nullptr; + const std::string depth_old = depth_env ? depth_env : ""; + const std::string budget_old = budget_env ? budget_env : ""; + + dummy_backend cuda = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); + dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); + ggml_backend_t backends[] = { cuda.handle.get(), cpu.handle.get() }; + ggml_backend_buffer_type_t bufts[] = { &cuda.buffer_type, &cpu.buffer_type }; + + set_test_env("GGML_KV_PIPELINE_DEPTH", "4"); + set_test_env("GGML_KV_PIPELINE_BUDGET_MIB", "8"); + { + auto graph = make_transport_graph(cpu, 64); + ggml_set_stable_prefix(graph.source, 64); + ggml_backend_sched_ptr sched(ggml_backend_sched_new(backends, bufts, 2, 128, false, false)); + ggml_backend_sched_set_tensor_backend(sched.get(), graph.output, cuda.handle.get()); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), graph.ctx.graph)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), graph.ctx.graph) == GGML_STATUS_SUCCESS); + int64_t deliveries = 0; + transport_stats(sched.get(), &deliveries, nullptr, nullptr); + GGML_ASSERT(deliveries == 1); + } + { + auto graph = make_transport_graph(cpu, 64); + ggml_set_stable_prefix(graph.source, 64); + ggml_backend_sched_ptr sched(ggml_backend_sched_new(backends, bufts, 2, 128, false, false)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_depth(sched.get(), 0)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_budget(sched.get(), 0)); + ggml_backend_sched_set_tensor_backend(sched.get(), graph.output, cuda.handle.get()); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), graph.ctx.graph)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), graph.ctx.graph) == GGML_STATUS_SUCCESS); + int64_t deliveries = -1; + transport_stats(sched.get(), &deliveries, nullptr, nullptr); + GGML_ASSERT(deliveries == 0); + } + + set_test_env("GGML_KV_PIPELINE_DEPTH", "bad"); + set_test_env("GGML_KV_PIPELINE_BUDGET_MIB", "bad"); + { + ggml_backend_sched_ptr sched(ggml_backend_sched_new(backends, bufts, 2, 128, false, false)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_depth(sched.get(), 0)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_budget(sched.get(), 0)); + } + + restore_test_env("GGML_KV_PIPELINE_DEPTH", had_depth, depth_old); + restore_test_env("GGML_KV_PIPELINE_BUDGET_MIB", had_budget, budget_old); +} + static void test_transport_depth_zero() { dummy_backend cuda = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); @@ -1485,6 +1602,9 @@ int main() { run("test_resizable_buffers_owner_borrower_scheduler_failure", test_resizable_buffers_owner_borrower_scheduler_failure); run("test_resizable_buffers_owner_borrower_teardown_order", test_resizable_buffers_owner_borrower_teardown_order); run("test_transport_prefix_and_configuration", test_transport_prefix_and_configuration); + run("test_transport_empty_graph", test_transport_empty_graph); + run("test_transport_fallback_keeps_allocator_plan", test_transport_fallback_keeps_allocator_plan); + run("test_transport_environment_is_fallback", test_transport_environment_is_fallback); run("test_transport_depth_zero", test_transport_depth_zero); run("test_transport_budget_recovers", test_transport_budget_recovers); run("test_transport_partial_backend_failure", test_transport_partial_backend_failure); diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 2d03e433c69..442dc73a483 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -344,6 +344,7 @@ struct cmd_params { std::vector main_gpu; std::vector no_kv_offload; std::vector kv_cpu_pinned; + std::vector kv_pipeline_depth; std::vector recurrent_state_offload; std::vector flash_attn; std::vector> devices; @@ -390,6 +391,7 @@ static const cmd_params cmd_params_defaults = { /* main_gpu */ { 0 }, /* no_kv_offload */ { false }, /* kv_cpu_pinned */ { false }, + /* kv_pipeline_depth */ { 1 }, /* recurrent_state_offload */ { false }, /* flash_attn */ { LLAMA_FLASH_ATTN_TYPE_AUTO }, /* devices */ { {} }, @@ -462,6 +464,7 @@ static void print_usage(int /* argc */, char ** argv) { printf(" -mg, --main-gpu (default: %s)\n", join(cmd_params_defaults.main_gpu, ",").c_str()); printf(" -nkvo, --no-kv-offload <0|1> (default: %s)\n", join(cmd_params_defaults.no_kv_offload, ",").c_str()); printf(" -kvcp, --kv-cpu-pinned <0|1> (default: %s)\n", join(cmd_params_defaults.kv_cpu_pinned, ",").c_str()); + printf(" -kvpd, --kv-pipeline-depth <0...14> (default: %s)\n", join(cmd_params_defaults.kv_pipeline_depth, ",").c_str()); printf(" -rso, --recurrent-state-offload <0|1> (default: %s)\n", join(cmd_params_defaults.recurrent_state_offload, ",").c_str()); printf(" -fa, --flash-attn (default: %s)\n", join(transform_to_str(cmd_params_defaults.flash_attn, llama_flash_attn_type_name), ",").c_str()); printf(" -dev, --device (default: auto)\n"); @@ -812,6 +815,19 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { } auto p = string_split(argv[i], split_delim); params.kv_cpu_pinned.insert(params.kv_cpu_pinned.end(), p.begin(), p.end()); + } else if (arg == "-kvpd" || arg == "--kv-pipeline-depth") { + if (++i >= argc) { + invalid_param = true; + break; + } + auto p = parse_int_range(argv[i]); + for (int depth : p) { + if (depth < 0 || depth > 14) { + invalid_param = true; + break; + } + } + params.kv_pipeline_depth.insert(params.kv_pipeline_depth.end(), p.begin(), p.end()); } else if (arg == "-rso" || arg == "--recurrent-state-offload") { if (++i >= argc) { invalid_param = true; @@ -1166,6 +1182,9 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { if (params.kv_cpu_pinned.empty()) { params.kv_cpu_pinned = cmd_params_defaults.kv_cpu_pinned; } + if (params.kv_pipeline_depth.empty()) { + params.kv_pipeline_depth = cmd_params_defaults.kv_pipeline_depth; + } if (params.recurrent_state_offload.empty()) { params.recurrent_state_offload = cmd_params_defaults.recurrent_state_offload; } @@ -1232,6 +1251,7 @@ struct cmd_params_instance { int main_gpu; bool no_kv_offload; bool kv_cpu_pinned; + int kv_pipeline_depth; bool recurrent_state_offload; llama_flash_attn_type flash_attn; std::vector devices; @@ -1313,6 +1333,7 @@ struct cmd_params_instance { cparams.type_v = type_v; cparams.offload_kqv = !no_kv_offload; cparams.kv_cpu_pinned = kv_cpu_pinned; + cparams.kv_pipeline_depth = kv_pipeline_depth; cparams.recurrent_state_offload = recurrent_state_offload; cparams.flash_attn_type = flash_attn; cparams.embeddings = embeddings; @@ -1348,6 +1369,7 @@ static std::vector get_cmd_params_instances(const cmd_param for (const auto & tv : params.type_v) for (const auto & nkvo : params.no_kv_offload) for (const auto & kvcp : params.kv_cpu_pinned) + for (const auto & kvpd : params.kv_pipeline_depth) for (const auto & rso : params.recurrent_state_offload) for (const auto & fa : params.flash_attn) for (const auto & nt : params.n_threads) @@ -1379,6 +1401,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .main_gpu = */ mg, /* .no_kv_offload = */ nkvo, /* .kv_cpu_pinned = */ kvcp, + /* .kv_pipeline_depth = */ kvpd, /* .recurrent_state_offload = */ rso, /* .flash_attn = */ fa, /* .devices = */ devs, @@ -1417,6 +1440,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .main_gpu = */ mg, /* .no_kv_offload = */ nkvo, /* .kv_cpu_pinned = */ kvcp, + /* .kv_pipeline_depth = */ kvpd, /* .recurrent_state_offload = */ rso, /* .flash_attn = */ fa, /* .devices = */ devs, @@ -1455,6 +1479,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .main_gpu = */ mg, /* .no_kv_offload = */ nkvo, /* .kv_cpu_pinned = */ kvcp, + /* .kv_pipeline_depth = */ kvpd, /* .recurrent_state_offload = */ rso, /* .flash_attn = */ fa, /* .devices = */ devs, @@ -1498,6 +1523,7 @@ struct test { int main_gpu; bool no_kv_offload; bool kv_cpu_pinned; + int kv_pipeline_depth; bool recurrent_state_offload; llama_flash_attn_type flash_attn; std::vector devices; @@ -1539,6 +1565,7 @@ struct test { main_gpu = inst.main_gpu; no_kv_offload = inst.no_kv_offload; kv_cpu_pinned = inst.kv_cpu_pinned; + kv_pipeline_depth = inst.kv_pipeline_depth; recurrent_state_offload = inst.recurrent_state_offload; flash_attn = inst.flash_attn; devices = inst.devices; @@ -1604,7 +1631,7 @@ struct test { "model_filename", "model_type", "model_size", "model_n_params", "n_batch", "n_ubatch", "n_threads", "cpu_mask", "cpu_strict", "poll", "type_k", "type_v", "n_gpu_layers", "n_cpu_moe", "split_mode", - "main_gpu", "no_kv_offload", "kv_cpu_pinned", "recurrent_state_offload", + "main_gpu", "no_kv_offload", "kv_cpu_pinned", "kv_pipeline_depth", "recurrent_state_offload", "flash_attn", "devices", "tensor_split", "tensor_buft_overrides", "load_mode", "embeddings", "no_op_offload", "no_host", "fit_target", "fit_min_ctx", @@ -1619,7 +1646,7 @@ struct test { static field_type get_field_type(const std::string & field) { if (field == "build_number" || field == "n_batch" || field == "n_ubatch" || field == "n_threads" || field == "poll" || field == "model_size" || field == "model_n_params" || field == "n_gpu_layers" || - field == "main_gpu" || field == "n_prompt" || field == "n_gen" || field == "n_depth" || field == "avg_ns" || + field == "main_gpu" || field == "kv_pipeline_depth" || field == "n_prompt" || field == "n_gen" || field == "n_depth" || field == "avg_ns" || field == "stddev_ns" || field == "no_op_offload" || field == "n_cpu_moe" || field == "fit_target" || field == "fit_min_ctx" || field == "flash_attn") { return INT; @@ -1698,6 +1725,7 @@ struct test { std::to_string(main_gpu), std::to_string(no_kv_offload), std::to_string(kv_cpu_pinned), + std::to_string(kv_pipeline_depth), std::to_string(recurrent_state_offload), std::to_string((int) flash_attn), devices_to_string(devices), @@ -1897,6 +1925,9 @@ struct markdown_printer : public printer { if (field == "kv_cpu_pinned") { return 4; } + if (field == "kv_pipeline_depth") { + return 4; + } if (field == "recurrent_state_offload") { return 3; } @@ -1928,6 +1959,9 @@ struct markdown_printer : public printer { if (field == "kv_cpu_pinned") { return "kvcp"; } + if (field == "kv_pipeline_depth") { + return "kvpd"; + } if (field == "recurrent_state_offload") { return "rso"; } @@ -2015,6 +2049,9 @@ struct markdown_printer : public printer { if (params.kv_cpu_pinned.size() > 1 || params.kv_cpu_pinned != cmd_params_defaults.kv_cpu_pinned) { fields.emplace_back("kv_cpu_pinned"); } + if (params.kv_pipeline_depth.size() > 1 || params.kv_pipeline_depth != cmd_params_defaults.kv_pipeline_depth) { + fields.emplace_back("kv_pipeline_depth"); + } if (params.recurrent_state_offload.size() > 1 || params.recurrent_state_offload != cmd_params_defaults.recurrent_state_offload) { fields.emplace_back("recurrent_state_offload"); From 0933834bb329fd53590f13834889dbffa41e911b Mon Sep 17 00:00:00 2001 From: piggidragon Date: Tue, 1 Sep 2026 10:21:35 +0200 Subject: [PATCH 08/32] sched: allocate transport ring entries the way the backend would The ring laid its entries out with ggml_nbytes() and bound them by writing data and buffer directly. A buffer type may ask for more than ggml_nbytes() for a tensor -- CUDA does for a quantized one, and MMQ clears that padding -- so an entry could reach into the next one. Entries are now sized with ggml_backend_buft_get_alloc_size() and bound with ggml_backend_tensor_alloc(), which also gives them the buffer's own initialization and its bounds check. test-alloc gets a dummy buffer type whose get_alloc_size exceeds ggml_nbytes, and a two-entry ring test that checks the entries stay inside the ring and out of each other, that every byte of an entry is delivered once from the matching source offset, and that nothing waits on an event before it is recorded. llama-bench takes -kvpb/--kv-pipeline-budget and reports it. The repro scripts pass 512 and now fail closed: they refuse a build without -kvcp, -rso or -kvpb instead of dropping the option, and every arm propagates its status. The llama-bench table in the doc was measured before the budget existed, so it says so, and the 32,768 row is marked as needing a re-measurement. Assisted-by: Claude Opus 5 --- docs/kv-transport-pipelining.md | 19 ++- docs/repro/r4-kv-pipeline-ab.sh | 57 +++++-- docs/repro/r4-kv-pipeline-context-sweep.sh | 69 +++++--- docs/repro/r4-kv-pipeline-exact.sh | 3 +- ggml/src/ggml-backend.cpp | 22 ++- tests/test-alloc.cpp | 176 ++++++++++++++++++++- tools/llama-bench/llama-bench.cpp | 43 ++++- 7 files changed, 331 insertions(+), 58 deletions(-) diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index fd16e366935..4bf21a3e43e 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -124,7 +124,7 @@ RTX 4070 (11,902 MiB usable, sm_89), driver 610.57.04 / CUDA 13.3, i5-13400F, `llama-bench`, `--no-warmup`, A/B/A/B with reversed arm order (`docs/repro/r4-kv-pipeline-ab.sh`, `docs/repro/r4-kv-pipeline-context-sweep.sh`), -at `--kv-pipeline-budget 512` so the 32,768 ring is allowed: +with no cap on the ring: | depth | ordered | pipelined | gain | peak device memory | |---:|---|---|---:|---:| @@ -132,11 +132,18 @@ at `--kv-pipeline-budget 512` so the 32,768 ring is allowed: | 16,384 | 19.4344, 19.4828 | 31.0981, 31.0931 | **+59.8%** | +86 to +104 MiB | | 32,768 | 12.9453, 12.9400 | 15.4571, 15.4516 | **+19.4%** | +206 MiB | -> These need `-kvcp 1 -rso 1`, and for a while `llama-bench` did not have them: -> the repro scripts probed `--help`, found nothing, and quietly dropped both. The -> same commit then measures 19.43 -> 9.02 t/s ordered at 16,384 and the pipeline -> buys +6.7% instead of +60%, because a host-resident recurrent state costs more -> than the transport can win back. `llama-bench` takes them again. +> These rows were taken before the budget existed, so they are the uncapped +> numbers. The 32,768 ring is 204 MiB and the default cap is 128 MiB, so that row +> does not reproduce on the current default: it needs `-kvpb 512`, which +> `llama-bench` did not take until now. The scripts pass it, and the row is due a +> re-measurement on the current head. + +> These also need `-kvcp 1 -rso 1`, and for a while `llama-bench` did not have +> them: the repro scripts probed `--help`, found nothing, and quietly dropped +> both. The same commit then measures 19.43 -> 9.02 t/s ordered at 16,384 and the +> pipeline buys +6.7% instead of +60%, because a host-resident recurrent state +> costs more than the transport can win back. `llama-bench` takes them again, and +> the scripts now fail rather than drop an option the build does not have. `llama-server`, one request, `temperature 0, top_k 1, seed 1234`: diff --git a/docs/repro/r4-kv-pipeline-ab.sh b/docs/repro/r4-kv-pipeline-ab.sh index 6bb6a381e56..563174e56f8 100755 --- a/docs/repro/r4-kv-pipeline-ab.sh +++ b/docs/repro/r4-kv-pipeline-ab.sh @@ -3,37 +3,62 @@ # The two arms are the same binary: --kv-pipeline-depth 0 is the ordered path. # # LLAMA_KV_MODEL=/path/model.gguf docs/repro/r4-kv-pipeline-ab.sh [depth ...] -set -u +set -euo pipefail MODEL="${LLAMA_KV_MODEL:?set LLAMA_KV_MODEL to a .gguf path}" BUILD="${LLAMA_KV_BUILD:-build}" PIN="${LLAMA_KV_TASKSET:-0,2,4}" +BUDGET="${LLAMA_KV_BUDGET:-512}" LOCK=/tmp/beellama-single-gpu.lock -# An unpinned host cache and a host-resident recurrent state both cost more than the transport -# can win back, so a run without these does not measure the same thing. Older llama-bench builds -# do not have them; skip rather than fail, and the numbers are then not comparable to the doc. -BENCH_KV_OPTS="" -for opt in kvcp rso; do - "$BUILD/bin/llama-bench" --help 2>&1 | grep -q -- "-$opt," && BENCH_KV_OPTS="$BENCH_KV_OPTS -$opt 1" +# An unpinned host cache and a host-resident recurrent state both cost more than the transport can +# win back, and without a budget the ring is declined at the larger contexts, so a build without +# these options does not measure what the doc reports. Fail rather than measure something else. +if ! HELP="$("$BUILD/bin/llama-bench" --help 2>&1)"; then + echo "cannot run $BUILD/bin/llama-bench:" >&2 + echo "$HELP" >&2 + exit 1 +fi +for opt in kvcp rso kvpb; do + if ! grep -q -- "-$opt," <<< "$HELP"; then + echo "$BUILD/bin/llama-bench has no -$opt option" >&2 + exit 1 + fi done run () { # $1 label, $2 pipeline depth, $3 context depth, $4 reps + local out err rc + out="$(mktemp)" + err="$(mktemp)" + rc=0 taskset -c "$PIN" "$BUILD/bin/llama-bench" -m "$MODEL" --kv-pipeline-depth "$2" \ - -ngl 99 -sm none -mg 0 -t 3 -nkvo 1 $BENCH_KV_OPTS \ + --kv-pipeline-budget "$BUDGET" -ngl 99 -sm none -mg 0 -t 3 -nkvo 1 -kvcp 1 -rso 1 \ -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']))" + > "$out" 2> "$err" || rc=$? + if [ "$rc" -eq 0 ]; then + python3 -c "import json,sys;d=json.load(sys.stdin);print(' %-12s %.4f +- %.4f'%('$1',d[0]['avg_ts'],d[0]['stddev_ts']))" \ + < "$out" 2>/dev/null || rc=$? + fi + if [ "$rc" -ne 0 ]; then + echo " $1: FAILED" >&2 + cat "$err" >&2 + fi + rm -f "$out" "$err" + return "$rc" } -DEPTHS=(4096 16384 32768); [ $# -gt 0 ] && DEPTHS=("$@") -rc=0 +DEPTHS=(4096 16384 32768) +if [ $# -gt 0 ]; then + DEPTHS=("$@") +fi for D in "${DEPTHS[@]}"; do - R=3; [ "$D" -le 4096 ] && R=5 + R=3 + if [ "$D" -le 4096 ]; then + R=5 + fi echo "== context depth=$D reps=$R" - flock "$LOCK" bash -c "$(declare -f run); BUILD='$BUILD'; MODEL='$MODEL'; PIN='$PIN'; BENCH_KV_OPTS='$BENCH_KV_OPTS' + flock "$LOCK" bash -c "set -euo pipefail; $(declare -f run); BUILD='$BUILD'; MODEL='$MODEL'; PIN='$PIN'; BUDGET='$BUDGET' run ordered 0 $D $R run pipelined 1 $D $R run ordered2 0 $D $R - run pipelined2 1 $D $R" || rc=$? + run pipelined2 1 $D $R" done -exit $rc diff --git a/docs/repro/r4-kv-pipeline-context-sweep.sh b/docs/repro/r4-kv-pipeline-context-sweep.sh index 79d2f07ee30..cf363dfa002 100755 --- a/docs/repro/r4-kv-pipeline-context-sweep.sh +++ b/docs/repro/r4-kv-pipeline-context-sweep.sh @@ -4,45 +4,68 @@ # 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 +set -euo pipefail 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}" +BUDGET="${LLAMA_KV_BUDGET:-512}" LOCK=/tmp/beellama-single-gpu.lock -# An unpinned host cache and a host-resident recurrent state both cost more than the transport -# can win back, so a run without these does not measure the same thing. Older llama-bench builds -# do not have them; skip rather than fail, and the numbers are then not comparable to the doc. -BENCH_KV_OPTS="" -for opt in kvcp rso; do - "$BUILD/bin/llama-bench" --help 2>&1 | grep -q -- "-$opt," && BENCH_KV_OPTS="$BENCH_KV_OPTS -$opt 1" +# An unpinned host cache and a host-resident recurrent state both cost more than the transport can +# win back, and without a budget the ring is declined at the larger contexts, so a build without +# these options does not measure what the doc reports. Fail rather than measure something else. +if ! HELP="$("$BUILD/bin/llama-bench" --help 2>&1)"; then + echo "cannot run $BUILD/bin/llama-bench:" >&2 + echo "$HELP" >&2 + exit 1 +fi +for opt in kvcp rso kvpb; do + if ! grep -q -- "-$opt," <<< "$HELP"; then + echo "$BUILD/bin/llama-bench has no -$opt option" >&2 + exit 1 + fi done arm () { # $1 pipeline depth, $2 context depth, $3 reps - local vram; vram=$(mktemp) + local vram out err rc ts + vram="$(mktemp)" + out="$(mktemp)" + err="$(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=$(taskset -c "$PIN" "$BUILD/bin/llama-bench" -m "$MODEL" --kv-pipeline-depth "$1" \ - -ngl 99 -sm none -mg 0 -t 3 -nkvo 1 $BENCH_KV_OPTS \ - -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 + rc=0 + taskset -c "$PIN" "$BUILD/bin/llama-bench" -m "$MODEL" --kv-pipeline-depth "$1" \ + --kv-pipeline-budget "$BUDGET" -ngl 99 -sm none -mg 0 -t 3 -nkvo 1 -kvcp 1 -rso 1 \ + -fa on -ctk q8_0 -ctv q8_0 -b 512 -ub 512 --no-warmup -p 0 -n "$NGEN" -d "$2" -r "$3" -o json \ + > "$out" 2> "$err" || rc=$? + kill $sampler 2>/dev/null || true + wait $sampler 2>/dev/null || true + ts="" + if [ "$rc" -eq 0 ]; then + ts="$(python3 -c "import json,sys;d=json.load(sys.stdin);print('%.4f'%d[0]['avg_ts'])" < "$out" 2>/dev/null)" || rc=$? + fi + if [ "$rc" -ne 0 ]; then + echo " depth=$1: FAILED" >&2 + cat "$err" >&2 + rm -f "$vram" "$out" "$err" + return "$rc" + fi printf ' %-10s %-10s %s MiB\n' "depth=$1" "$ts" "$(sort -n "$vram" | tail -1)" - rm -f "$vram" + rm -f "$vram" "$out" "$err" } -DEPTHS=(4096 16384 32768 65536 131072 262144); [ $# -gt 0 ] && DEPTHS=("$@") +DEPTHS=(4096 16384 32768 65536 131072 262144) +if [ $# -gt 0 ]; then + DEPTHS=("$@") +fi for D in "${DEPTHS[@]}"; do - R=3; [ "$D" -gt 32768 ] && R=1 + R=3 + if [ "$D" -gt 32768 ]; then + R=1 + fi echo "== context depth=$D reps=$R (t/s, peak device memory)" - flock "$LOCK" bash -c "$(declare -f arm); BUILD='$BUILD'; MODEL='$MODEL'; PIN='$PIN'; BENCH_KV_OPTS='$BENCH_KV_OPTS'; NGEN='$NGEN' + flock "$LOCK" bash -c "set -euo pipefail; $(declare -f arm); BUILD='$BUILD'; MODEL='$MODEL'; PIN='$PIN'; BUDGET='$BUDGET'; NGEN='$NGEN' arm 0 $D $R arm 1 $D $R arm 0 $D $R diff --git a/docs/repro/r4-kv-pipeline-exact.sh b/docs/repro/r4-kv-pipeline-exact.sh index 485a1196cb8..d2795a5a7c1 100755 --- a/docs/repro/r4-kv-pipeline-exact.sh +++ b/docs/repro/r4-kv-pipeline-exact.sh @@ -11,6 +11,7 @@ PIN="${LLAMA_KV_TASKSET:-0,2,4}" PORT="${LLAMA_KV_PORT:-18099}" LENGTHS="${LLAMA_KV_LENGTHS:-2048,18432}" CTX="${LLAMA_KV_CTX:-32768}" +BUDGET="${LLAMA_KV_BUDGET:-512}" HERE="$(cd "$(dirname "$0")" && pwd)" DEPTHS=(0 1 4); [ $# -gt 0 ] && DEPTHS=("$@") @@ -21,7 +22,7 @@ for I in "${!DEPTHS[@]}"; do echo "== pipeline depth=$D ctx=$CTX prefill lengths=$LENGTHS" LOG=$(mktemp /tmp/r4-kv-pipeline.XXXX.log) taskset -c "$PIN" "$BUILD/bin/llama-server" -m "$MODEL" --kv-pipeline-depth "$D" \ - -ngl 99 -sm none -mg 0 -t 3 -nkvo --kv-cpu-pinned --recurrent-state-offload \ + --kv-pipeline-budget "$BUDGET" -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 "$CTX" --parallel 1 \ --host 127.0.0.1 --port "$PORT" --no-warmup > "$LOG" 2>&1 & SRV=$! diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 80534db33e3..77cace14551 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1807,6 +1807,13 @@ static bool ggml_backend_sched_input_is_staged(ggml_backend_sched_t sched, int s static bool ggml_backend_sched_size_add(size_t a, size_t b, size_t * result); static bool ggml_backend_sched_size_pad(size_t size, size_t alignment, size_t * result); +// A ring entry costs what the backend would allocate for it, which can be more than its data: +// a buffer type may ask for padding past ggml_nbytes(), and its kernels may write into it. +static bool ggml_backend_sched_transport_entry_size( + ggml_backend_buffer_type_t buft, const struct ggml_tensor * t, size_t alignment, size_t * result) { + return ggml_backend_sched_size_pad(ggml_backend_buft_get_alloc_size(buft, t), alignment, result); +} + static void ggml_backend_sched_transport_clear_addresses(ggml_backend_sched_t sched) { const struct ggml_backend_sched_transport * tr = &sched->transport; @@ -1842,10 +1849,14 @@ static void ggml_backend_sched_transport_assign_addresses(ggml_backend_sched_t s 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 = r->buffer; + // bind through the backend, so that the entry is initialized the same way as any other + // tensor the buffer holds. A previous plan may have left this copy bound already. + input_cpy->data = NULL; + input_cpy->buffer = NULL; + const enum ggml_status status = ggml_backend_tensor_alloc(r->buffer, input_cpy, slot + offset); + GGML_ASSERT(status == GGML_STATUS_SUCCESS); size_t input_size; - GGML_ASSERT(ggml_backend_sched_size_pad(ggml_nbytes(split->inputs[j]), r->alignment, &input_size)); + GGML_ASSERT(ggml_backend_sched_transport_entry_size(ggml_backend_buffer_get_type(r->buffer), input_cpy, r->alignment, &input_size)); GGML_ASSERT(ggml_backend_sched_size_add(offset, input_size, &offset)); } GGML_ASSERT(offset <= r->slot_size); @@ -2148,10 +2159,11 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { } const struct ggml_tensor * input = split->inputs[j]; const struct ggml_tensor * base = input->view_src ? input->view_src : input; + const struct ggml_tensor * entry = tensor_copy(split->inputs[j], bid, sched->cur_copy); size_t input_size; size_t input_size_max; - if (!ggml_backend_sched_size_pad(ggml_nbytes(input), tr->rings[bid].alignment, &input_size) || - !ggml_backend_sched_size_pad(ggml_nbytes(base), tr->rings[bid].alignment, &input_size_max) || + if (!ggml_backend_sched_transport_entry_size(sched->bufts[bid], entry, tr->rings[bid].alignment, &input_size) || + !ggml_backend_sched_transport_entry_size(sched->bufts[bid], base, tr->rings[bid].alignment, &input_size_max) || !ggml_backend_sched_size_add(need, input_size, &need) || !ggml_backend_sched_size_add(need_max, input_size_max, &need_max)) { size_overflow[bid] = true; diff --git a/tests/test-alloc.cpp b/tests/test-alloc.cpp index ef8895e336f..7f1bd3adaa3 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -6,6 +6,7 @@ #include #include +#include #include #include #include @@ -31,6 +32,29 @@ struct dummy_backend_context { int event_wait_count = 0; int set_tensor_async_count = 0; size_t set_tensor_async_bytes = 0; + size_t alloc_size_pad = 0; + + // what the backend was asked to hold and to move, so that a test can check the entries a + // transport ring lays out and the bytes it delivers into them + struct tensor_binding { + const ggml_tensor * tensor; + ggml_backend_buffer_t buffer; + const char * data; + size_t size; + }; + struct tensor_delivery { + const ggml_tensor * tensor; + const char * src; + size_t offset; + size_t size; + }; + struct event_step { + bool is_wait; + ggml_backend_event_t event; + }; + std::vector bindings; + std::vector deliveries; + std::vector event_steps; ggml_backend_buffer_type_t buffer_type = nullptr; ggml_backend_i backend_interface = {}; @@ -76,6 +100,12 @@ static size_t dummy_backend_buffer_type_get_max_size(ggml_backend_buffer_type_t return ctx->max_buffer_size; } +static size_t dummy_backend_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { + dummy_backend_context * ctx = (dummy_backend_context *) buft->context; + // only a tensor that is no op's output may ask for more than its data [TAG_ALLOC_SIZE_EXPAND] + return ggml_nbytes(tensor) + (ggml_op_is_empty(tensor->op) ? ctx->alloc_size_pad : 0); +} + static bool dummy_backend_buffer_type_is_host(ggml_backend_buffer_type_t buft) { return ((dummy_backend_context *) buft->context)->buffer_is_host; } @@ -101,7 +131,17 @@ static void * dummy_backend_buffer_get_base(ggml_backend_buffer_t buffer) { return ctx->buffer_bases[i - ctx->buffers.begin()]; } -static ggml_status dummy_backend_buffer_init_tensor(ggml_backend_buffer_t, ggml_tensor *) { +static ggml_status dummy_backend_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { + dummy_backend_context * ctx = (dummy_backend_context *) buffer->context; + const size_t size = ggml_backend_buffer_get_alloc_size(buffer, tensor); + + for (auto & b : ctx->bindings) { + if (b.tensor == tensor) { + b = { tensor, buffer, (const char *) tensor->data, size }; + return GGML_STATUS_SUCCESS; + } + } + ctx->bindings.push_back({ tensor, buffer, (const char *) tensor->data, size }); return GGML_STATUS_SUCCESS; } @@ -133,18 +173,23 @@ static void dummy_backend_free(ggml_backend_t backend) { delete backend; } -static void dummy_backend_set_tensor_async(ggml_backend_t backend, ggml_tensor *, const void *, size_t, size_t size) { +static void dummy_backend_set_tensor_async(ggml_backend_t backend, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { dummy_backend_context * ctx = (dummy_backend_context *) backend->context; ctx->set_tensor_async_count++; ctx->set_tensor_async_bytes += size; + ctx->deliveries.push_back({ tensor, (const char *) data, offset, size }); } static void dummy_backend_synchronize(ggml_backend_t) {} -static void dummy_backend_event_record(ggml_backend_t, ggml_backend_event_t) {} +static void dummy_backend_event_record(ggml_backend_t backend, ggml_backend_event_t event) { + ((dummy_backend_context *) backend->context)->event_steps.push_back({ false, event }); +} -static void dummy_backend_event_wait(ggml_backend_t backend, ggml_backend_event_t) { - ((dummy_backend_context *) backend->context)->event_wait_count++; +static void dummy_backend_event_wait(ggml_backend_t backend, ggml_backend_event_t event) { + dummy_backend_context * ctx = (dummy_backend_context *) backend->context; + ctx->event_wait_count++; + ctx->event_steps.push_back({ true, event }); } static enum ggml_status dummy_backend_graph_compute(ggml_backend_t backend, ggml_cgraph *) { @@ -237,6 +282,7 @@ static dummy_backend dummy_backend_init( b.buffer_type.iface.alloc_buffer = dummy_backend_buffer_type_alloc_buffer; b.buffer_type.iface.get_alignment = dummy_backend_buffer_type_get_alignment; b.buffer_type.iface.get_max_size = dummy_backend_buffer_type_get_max_size; + b.buffer_type.iface.get_alloc_size = dummy_backend_buffer_type_get_alloc_size; b.buffer_type.iface.is_host = dummy_backend_buffer_type_is_host; b.context->buffer_type = &b.buffer_type; @@ -315,6 +361,34 @@ static transport_graph make_transport_graph(dummy_backend & cpu, size_t size) { return { std::move(result), std::move(buffer), source, output }; } +struct transport_graph_pair { + test_context_with_graph ctx; + ggml_backend_buffer_ptr buffer; + ggml_tensor * sources[2]; + ggml_tensor * output; +}; + +// two transported inputs in one split, so that the ring lays out more than one entry per slot +static transport_graph_pair make_transport_graph_pair(dummy_backend & cpu, size_t size) { + GGML_ASSERT(size % sizeof(float) == 0); + auto result = make_context(); + ggml_tensor * s0 = ggml_new_tensor_1d(result.ctx, GGML_TYPE_F32, size/sizeof(float)); + ggml_tensor * s1 = ggml_new_tensor_1d(result.ctx, GGML_TYPE_F32, size/sizeof(float)); + ggml_tensor * output = ggml_add(result.ctx, s0, s1); + s0->flags |= GGML_TENSOR_FLAG_TRANSPORT; + s1->flags |= GGML_TENSOR_FLAG_TRANSPORT; + ggml_build_forward_expand(result.graph, output); + + ggml_backend_buffer_ptr buffer(ggml_backend_buft_alloc_buffer(&cpu.buffer_type, 2*size)); + char * base = (char *) ggml_backend_buffer_get_base(buffer.get()); + s0->buffer = buffer.get(); + s0->data = base; + s1->buffer = buffer.get(); + s1->data = base + size; + + return { std::move(result), std::move(buffer), { s0, s1 }, output }; +} + static void transport_stats( ggml_backend_sched_t sched, int64_t * deliveries, @@ -1300,6 +1374,97 @@ static void test_transport_prefix_and_configuration() { GGML_ASSERT(cuda.context->transfer_backend_count == 0); } +// The ring must hold what the backend allocates for an entry, deliver every byte of it exactly +// once, and never wait on an event it has not recorded. +static void test_transport_entry_allocation() { + dummy_backend cuda = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); + dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); + + // ask for more room per entry than its data needs, the way a quantized tensor does + cuda.context->alloc_size_pad = 32; + + const size_t nbytes = 128; + auto graph = make_transport_graph_pair(cpu, nbytes); + + ggml_backend_t backends[] = { cuda.handle.get(), cpu.handle.get() }; + ggml_backend_buffer_type_t bufts[] = { &cuda.buffer_type, &cpu.buffer_type }; + ggml_backend_sched_ptr sched(ggml_backend_sched_new(backends, bufts, 2, 128, false, false)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_budget(sched.get(), 4096)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_depth(sched.get(), 1)); + ggml_backend_sched_set_tensor_backend(sched.get(), graph.output, cuda.handle.get()); + + // one input goes early in part, the other whole + ggml_set_stable_prefix(graph.sources[0], nbytes/2); + ggml_set_stable_prefix(graph.sources[1], nbytes); + + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), graph.ctx.graph)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), graph.ctx.graph) == GGML_STATUS_SUCCESS); + + int64_t deliveries = 0; + int64_t early = 0; + int64_t late = 0; + transport_stats(sched.get(), &deliveries, &early, &late); + GGML_ASSERT(deliveries == 1 && early == (int64_t) (nbytes + nbytes/2) && late == (int64_t) (nbytes/2)); + + std::vector entries; + for (const auto & b : cuda.context->bindings) { + if (strncmp(b.tensor->name, "CUDA#", 5) == 0) { + entries.push_back(&b); + } + } + GGML_ASSERT(entries.size() == 2); + + for (const auto * e : entries) { + // bound through the buffer, with the room the buffer type asks for + GGML_ASSERT(e->size == ggml_nbytes(e->tensor) + cuda.context->alloc_size_pad); + + const char * base = (const char *) ggml_backend_buffer_get_base(e->buffer); + GGML_ASSERT(e->data >= base); + GGML_ASSERT(e->data + e->size <= base + ggml_backend_buffer_get_size(e->buffer)); + + // and no entry, padding included, reaches into another one + for (const auto * other : entries) { + GGML_ASSERT(other == e || other->data + other->size <= e->data || other->data >= e->data + e->size); + } + + // every byte of the entry is delivered once, in order, from the matching source offset + std::vector parts; + for (const auto & d : cuda.context->deliveries) { + if (d.tensor == e->tensor) { + parts.push_back(d); + } + } + GGML_ASSERT(!parts.empty()); + std::sort(parts.begin(), parts.end(), + [](const dummy_backend_context::tensor_delivery & a, const dummy_backend_context::tensor_delivery & b) { + return a.offset < b.offset; + }); + const char * src = parts.front().src; + GGML_ASSERT(src == (const char *) graph.sources[0]->data || src == (const char *) graph.sources[1]->data); + size_t covered = 0; + for (const auto & d : parts) { + GGML_ASSERT(d.offset == covered); + GGML_ASSERT(d.src == src + d.offset); + covered += d.size; + } + GGML_ASSERT(covered == ggml_nbytes(e->tensor)); + } + + // nothing waits on an event before it is recorded + const auto & steps = cuda.context->event_steps; + GGML_ASSERT(!steps.empty()); + for (size_t i = 0; i < steps.size(); i++) { + if (!steps[i].is_wait) { + continue; + } + bool recorded = false; + for (size_t j = 0; j < i && !recorded; j++) { + recorded = !steps[j].is_wait && steps[j].event == steps[i].event; + } + GGML_ASSERT(recorded); + } +} + static void test_transport_empty_graph() { dummy_backend cuda = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); @@ -1639,6 +1804,7 @@ int main() { run("test_resizable_buffers_owner_borrower_scheduler_failure", test_resizable_buffers_owner_borrower_scheduler_failure); run("test_resizable_buffers_owner_borrower_teardown_order", test_resizable_buffers_owner_borrower_teardown_order); run("test_transport_prefix_and_configuration", test_transport_prefix_and_configuration); + run("test_transport_entry_allocation", test_transport_entry_allocation); run("test_transport_empty_graph", test_transport_empty_graph); run("test_transport_fallback_keeps_allocator_plan", test_transport_fallback_keeps_allocator_plan); run("test_transport_environment_is_fallback", test_transport_environment_is_fallback); diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 1cacccd7992..9dc09d930a9 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -359,6 +359,7 @@ struct cmd_params { std::vector no_kv_offload; std::vector kv_cpu_pinned; std::vector kv_pipeline_depth; + std::vector kv_pipeline_budget_mib; std::vector recurrent_state_offload; std::vector flash_attn; std::vector> devices; @@ -407,6 +408,7 @@ static const cmd_params cmd_params_defaults = { /* no_kv_offload */ { false }, /* kv_cpu_pinned */ { false }, /* kv_pipeline_depth */ { 1 }, + /* kv_pipeline_budget_mib */ { 128 }, /* recurrent_state_offload */ { false }, /* flash_attn */ { LLAMA_FLASH_ATTN_TYPE_AUTO }, /* devices */ { {} }, @@ -480,6 +482,7 @@ static void print_usage(int /* argc */, char ** argv) { printf(" -nkvo, --no-kv-offload <0|1> (default: %s)\n", join(cmd_params_defaults.no_kv_offload, ",").c_str()); printf(" -kvcp, --kv-cpu-pinned <0|1> (default: %s)\n", join(cmd_params_defaults.kv_cpu_pinned, ",").c_str()); printf(" -kvpd, --kv-pipeline-depth <0...14> (default: %s)\n", join(cmd_params_defaults.kv_pipeline_depth, ",").c_str()); + printf(" -kvpb, --kv-pipeline-budget (default: %s)\n", join(cmd_params_defaults.kv_pipeline_budget_mib, ",").c_str()); printf(" -rso, --recurrent-state-offload <0|1> (default: %s)\n", join(cmd_params_defaults.recurrent_state_offload, ",").c_str()); printf(" -fa, --flash-attn (default: %s)\n", join(transform_to_str(cmd_params_defaults.flash_attn, llama_flash_attn_type_name), ",").c_str()); printf(" -dev, --device (default: auto)\n"); @@ -870,6 +873,19 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { } } params.kv_pipeline_depth.insert(params.kv_pipeline_depth.end(), p.begin(), p.end()); + } else if (arg == "-kvpb" || arg == "--kv-pipeline-budget") { + if (++i >= argc) { + invalid_param = true; + break; + } + auto p = parse_int_range(argv[i]); + for (int budget : p) { + if (budget < 0) { + invalid_param = true; + break; + } + } + params.kv_pipeline_budget_mib.insert(params.kv_pipeline_budget_mib.end(), p.begin(), p.end()); } else if (arg == "-rso" || arg == "--recurrent-state-offload") { if (++i >= argc) { invalid_param = true; @@ -1230,6 +1246,9 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { if (params.kv_pipeline_depth.empty()) { params.kv_pipeline_depth = cmd_params_defaults.kv_pipeline_depth; } + if (params.kv_pipeline_budget_mib.empty()) { + params.kv_pipeline_budget_mib = cmd_params_defaults.kv_pipeline_budget_mib; + } if (params.recurrent_state_offload.empty()) { params.recurrent_state_offload = cmd_params_defaults.recurrent_state_offload; } @@ -1298,6 +1317,7 @@ struct cmd_params_instance { bool no_kv_offload; bool kv_cpu_pinned; int kv_pipeline_depth; + int kv_pipeline_budget_mib; bool recurrent_state_offload; llama_flash_attn_type flash_attn; std::vector devices; @@ -1382,6 +1402,7 @@ struct cmd_params_instance { cparams.offload_kqv = !no_kv_offload; cparams.kv_cpu_pinned = kv_cpu_pinned; cparams.kv_pipeline_depth = kv_pipeline_depth; + cparams.kv_pipeline_budget_mib = kv_pipeline_budget_mib; cparams.recurrent_state_offload = recurrent_state_offload; cparams.flash_attn_type = flash_attn; cparams.embeddings = embeddings; @@ -1419,6 +1440,7 @@ static std::vector get_cmd_params_instances(const cmd_param for (const auto & nkvo : params.no_kv_offload) for (const auto & kvcp : params.kv_cpu_pinned) for (const auto & kvpd : params.kv_pipeline_depth) + for (const auto & kvpb : params.kv_pipeline_budget_mib) for (const auto & rso : params.recurrent_state_offload) for (const auto & fa : params.flash_attn) for (const auto & nt : params.n_threads) @@ -1452,6 +1474,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .no_kv_offload = */ nkvo, /* .kv_cpu_pinned = */ kvcp, /* .kv_pipeline_depth = */ kvpd, + /* .kv_pipeline_budget_mib = */ kvpb, /* .recurrent_state_offload = */ rso, /* .flash_attn = */ fa, /* .devices = */ devs, @@ -1492,6 +1515,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .no_kv_offload = */ nkvo, /* .kv_cpu_pinned = */ kvcp, /* .kv_pipeline_depth = */ kvpd, + /* .kv_pipeline_budget_mib = */ kvpb, /* .recurrent_state_offload = */ rso, /* .flash_attn = */ fa, /* .devices = */ devs, @@ -1532,6 +1556,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .no_kv_offload = */ nkvo, /* .kv_cpu_pinned = */ kvcp, /* .kv_pipeline_depth = */ kvpd, + /* .kv_pipeline_budget_mib = */ kvpb, /* .recurrent_state_offload = */ rso, /* .flash_attn = */ fa, /* .devices = */ devs, @@ -1577,6 +1602,7 @@ struct test { bool no_kv_offload; bool kv_cpu_pinned; int kv_pipeline_depth; + int kv_pipeline_budget_mib; bool recurrent_state_offload; llama_flash_attn_type flash_attn; std::vector devices; @@ -1620,6 +1646,7 @@ struct test { no_kv_offload = inst.no_kv_offload; kv_cpu_pinned = inst.kv_cpu_pinned; kv_pipeline_depth = inst.kv_pipeline_depth; + kv_pipeline_budget_mib = inst.kv_pipeline_budget_mib; recurrent_state_offload = inst.recurrent_state_offload; flash_attn = inst.flash_attn; devices = inst.devices; @@ -1685,7 +1712,8 @@ struct test { "model_filename", "model_type", "model_size", "model_n_params", "n_batch", "n_ubatch", "n_threads", "cpu_mask", "cpu_strict", "poll", "type_k", "type_v", "n_gpu_layers", "n_cpu_moe", "split_mode", - "main_gpu", "no_kv_offload", "kv_cpu_pinned", "kv_pipeline_depth", "recurrent_state_offload", + "main_gpu", "no_kv_offload", "kv_cpu_pinned", "kv_pipeline_depth", "kv_pipeline_budget_mib", + "recurrent_state_offload", "flash_attn", "devices", "tensor_split", "tensor_buft_overrides", "load_mode", "lazy_mode", "embeddings", "no_op_offload", "no_host", "fit_target", "fit_min_ctx", @@ -1700,7 +1728,7 @@ struct test { static field_type get_field_type(const std::string & field) { if (field == "build_number" || field == "n_batch" || field == "n_ubatch" || field == "n_threads" || field == "poll" || field == "model_size" || field == "model_n_params" || field == "n_gpu_layers" || - field == "main_gpu" || field == "kv_pipeline_depth" || field == "n_prompt" || field == "n_gen" || field == "n_depth" || field == "avg_ns" || + field == "main_gpu" || field == "kv_pipeline_depth" || field == "kv_pipeline_budget_mib" || field == "n_prompt" || field == "n_gen" || field == "n_depth" || field == "avg_ns" || field == "stddev_ns" || field == "no_op_offload" || field == "n_cpu_moe" || field == "fit_target" || field == "fit_min_ctx" || field == "flash_attn") { return INT; @@ -1780,6 +1808,7 @@ struct test { std::to_string(no_kv_offload), std::to_string(kv_cpu_pinned), std::to_string(kv_pipeline_depth), + std::to_string(kv_pipeline_budget_mib), std::to_string(recurrent_state_offload), std::to_string((int) flash_attn), devices_to_string(devices), @@ -1983,6 +2012,9 @@ struct markdown_printer : public printer { if (field == "kv_pipeline_depth") { return 4; } + if (field == "kv_pipeline_budget_mib") { + return 5; + } if (field == "recurrent_state_offload") { return 3; } @@ -2017,6 +2049,9 @@ struct markdown_printer : public printer { if (field == "kv_pipeline_depth") { return "kvpd"; } + if (field == "kv_pipeline_budget_mib") { + return "kvpb"; + } if (field == "recurrent_state_offload") { return "rso"; } @@ -2107,6 +2142,10 @@ struct markdown_printer : public printer { if (params.kv_pipeline_depth.size() > 1 || params.kv_pipeline_depth != cmd_params_defaults.kv_pipeline_depth) { fields.emplace_back("kv_pipeline_depth"); } + if (params.kv_pipeline_budget_mib.size() > 1 || + params.kv_pipeline_budget_mib != cmd_params_defaults.kv_pipeline_budget_mib) { + fields.emplace_back("kv_pipeline_budget_mib"); + } if (params.recurrent_state_offload.size() > 1 || params.recurrent_state_offload != cmd_params_defaults.recurrent_state_offload) { fields.emplace_back("recurrent_state_offload"); From 11b165d44ea46101029109fc3ef5491062136afe Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 3 Sep 2026 10:12:38 +0200 Subject: [PATCH 09/32] sched: never let the transport ring starve the graph The rings are laid out and allocated before the graph is, so a device that can hold the graph alone but not the graph next to a ring turned into GGML_STATUS_ALLOC_FAILED. The configuration is locked by then, so the caller could not turn the ring off and retry either. When graph reservation fails the rings are now released and the reservation is retried once on the ordered path, and that scheduler keeps the ordered path from then on. A plan over a split list with no inputs left input_staged unallocated and passed it to memset, which UBSan reports even at size 0. Such a plan stages nothing, so it now returns after putting every split back on the ordered path. test-alloc gets a device capacity on the dummy backend and a test that sizes it to hold the graph or the ring but not both. The prose and public comments this branch added were hard-wrapped to a fixed column, against the repository rule. They are unwrapped, one sentence per line. Assisted-by: Claude Opus 5 --- docs/kv-transport-pipelining.md | 360 +++++++------------------------- ggml/include/ggml-backend.h | 32 ++- ggml/include/ggml.h | 7 +- ggml/src/ggml-backend.cpp | 200 +++++++++--------- include/llama.h | 13 +- src/llama-kv-cache.cpp | 11 +- src/llama-kv-cache.h | 6 +- tests/test-alloc.cpp | 50 ++++- 8 files changed, 246 insertions(+), 433 deletions(-) diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index 4bf21a3e43e..467fada3e69 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -1,72 +1,34 @@ # 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. +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. -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. +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. +`--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. -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). +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 -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). +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. +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 (+1.38%, and not exact). +`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 (+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. +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: @@ -74,57 +36,29 @@ Each slot has one ownership cycle: 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. +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. 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. +- **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` -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. +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. +- **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". +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`. +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`, `docs/repro/r4-kv-pipeline-context-sweep.sh`), -with no cap on the ring: +`llama-bench`, `--no-warmup`, A/B/A/B with reversed arm order (`docs/repro/r4-kv-pipeline-ab.sh`, `docs/repro/r4-kv-pipeline-context-sweep.sh`), with no cap on the ring: | depth | ordered | pipelined | gain | peak device memory | |---:|---|---|---:|---:| @@ -132,18 +66,9 @@ with no cap on the ring: | 16,384 | 19.4344, 19.4828 | 31.0981, 31.0931 | **+59.8%** | +86 to +104 MiB | | 32,768 | 12.9453, 12.9400 | 15.4571, 15.4516 | **+19.4%** | +206 MiB | -> These rows were taken before the budget existed, so they are the uncapped -> numbers. The 32,768 ring is 204 MiB and the default cap is 128 MiB, so that row -> does not reproduce on the current default: it needs `-kvpb 512`, which -> `llama-bench` did not take until now. The scripts pass it, and the row is due a -> re-measurement on the current head. +> These rows were taken before the budget existed, so they are the uncapped numbers. The 32,768 ring is 204 MiB and the default cap is 128 MiB, so that row does not reproduce on the current default: it needs `-kvpb 512`, which `llama-bench` did not take until now. The scripts pass it, and the row is due a re-measurement on the current head. -> These also need `-kvcp 1 -rso 1`, and for a while `llama-bench` did not have -> them: the repro scripts probed `--help`, found nothing, and quietly dropped -> both. The same commit then measures 19.43 -> 9.02 t/s ordered at 16,384 and the -> pipeline buys +6.7% instead of +60%, because a host-resident recurrent state -> costs more than the transport can win back. `llama-bench` takes them again, and -> the scripts now fail rather than drop an option the build does not have. +> These also need `-kvcp 1 -rso 1`, and for a while `llama-bench` did not have them: the repro scripts probed `--help`, found nothing, and quietly dropped both. The same commit then measures 19.43 -> 9.02 t/s ordered at 16,384 and the pipeline buys +6.7% instead of +60%, because a host-resident recurrent state costs more than the transport can win back. `llama-bench` takes them again, and the scripts now fail rather than drop an option the build does not have. `llama-server`, one request, `temperature 0, top_k 1, seed 1234`: @@ -152,32 +77,22 @@ with no cap on the ring: | 19,246 | 32,768 | 17.785 | 29.833 | **+67.7%** | 28.3 | 25.7 | 31.30 | 95.3% | | 48,042 | 65,536 | 9.790 | 11.313 | **+15.6%** | 76.4 | 25.4 | 11.86 | 95.4% | -`copy` and `compute` are read off `GGML_SCHED_TRANSPORT_DEBUG=2` on each arm, not -fitted: the ordered arm reports what it spends blocked in `ggml_backend_tensor_copy` -and what it spends waiting for the consumer. The ceiling is `max(copy, compute)` -plus the per-token work outside the split loop, which is on both arms. +`copy` and `compute` are read off `GGML_SCHED_TRANSPORT_DEBUG=2` on each arm, not fitted: the ordered arm reports what it spends blocked in `ggml_backend_tensor_copy` and what it spends waiting for the consumer. The ceiling is `max(copy, compute)` plus the per-token work outside the split loop, which is on both arms. -**The pipeline is within 5% of that ceiling at both depths.** What is left is not -a scheduling problem, and the section on the residual below says what it is. +**The pipeline is within 5% of that ceiling at both depths.** What is left is not a scheduling problem, and the section on the residual below says what it is. -Pinning is worth as much as the pipeline and is off by default. Behind a -13,128-token prompt: +Pinning is worth as much as the pipeline and is off by default. Behind a 13,128-token prompt: | | ordered | pipelined | |---|---:|---:| | `--kv-cpu-pinned` | 21.582 | 32.252 | | unpinned | 14.945 | 22.709 | -Look-ahead deeper than one split is worse at every depth measured. At 19,246: -29.83 t/s at `N = 1`, 28.44 at `N = 2`, 25.93 at `N = 4`. `N = 1` is the default -for that reason. +Look-ahead deeper than one split is worse at every depth measured. At 19,246: 29.83 t/s at `N = 1`, 28.44 at `N = 2`, 25.93 at `N = 4`. `N = 1` is the default for that reason. ### The link is the ceiling, so the lever is bytes -644 MiB in 28.3 ms is 22.0 GB/s, and `nvidia-smi` reports the card at gen4 x16. -That is about 88% of what the link delivers in practice, so there is no room left -in the transport itself. What is left is to send less. `-ctk q4_0 -ctv q4_0` -halves the cache and therefore the traffic: +644 MiB in 28.3 ms is 22.0 GB/s, and `nvidia-smi` reports the card at gen4 x16. That is about 88% of what the link delivers in practice, so there is no room left in the transport itself. What is left is to send less. `-ctk q4_0 -ctv q4_0` halves the cache and therefore the traffic: | prompt | KV | ordered | pipelined | delivered | |---:|---|---:|---:|---:| @@ -186,21 +101,13 @@ halves the cache and therefore the traffic: | 48,042 | q8_0 | 9.790 | 11.313 | 1602.5 MiB | | 48,042 | q4_0 | 14.146 | 17.717 | 850.6 MiB | -Halving the traffic is worth +5.6% on the pipelined path at 19,246 and +56.6% at -48,042. The difference is the crossover: at 19,246 the pipeline has already -brought the copy down to the compute floor, and the consumer wait is 27.31 ms at -q8_0 against 27.26 ms at q4_0, the same number. Removing bytes there removes work -nothing was waiting for. At 48,042 the copy still dominates and every byte -removed is a byte off the token. +Halving the traffic is worth +5.6% on the pipelined path at 19,246 and +56.6% at 48,042. The difference is the crossover: at 19,246 the pipeline has already brought the copy down to the compute floor, and the consumer wait is 27.31 ms at q8_0 against 27.26 ms at q4_0, the same number. Removing bytes there removes work nothing was waiting for. At 48,042 the copy still dominates and every byte removed is a byte off the token. -**Whether to spend a quantisation step on the cache is a depth question, and the -two compound.** At 48,042, q4_0 with the pipeline is 17.717 against 9.790 for -q8_0 without it. +**Whether to spend a quantisation step on the cache is a depth question, and the two compound.** At 48,042, q4_0 with the pipeline is 17.717 against 9.790 for q8_0 without it. ### 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. +`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 | |---|---:|---:|---:|---|---:|---:| @@ -213,70 +120,35 @@ with `nvidia-smi` across each arm. Both passes agreed to the digits shown. 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 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%. +**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. +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. -**The ring beats `--kv-gpu-layers` per MiB, and the two barely add up.** Measured -behind a 19,246-token prompt at `-c 32768`, where a device-resident layer costs -about 68 MiB and the ring costs about 205 MiB: +**The ring beats `--kv-gpu-layers` per MiB, and the two barely add up.** Measured behind a 19,246-token prompt at `-c 32768`, where a device-resident layer costs about 68 MiB and the ring costs about 205 MiB: | | no `--kv-gpu-layers` | `--kv-gpu-layers 4` | `--kv-gpu-layers 8` | |---|---:|---:|---:| | ordered | 17.785 | 20.334 | | | pipelined | 29.843 | 30.263 | 30.640 | -Four device-resident layers are worth +14.3% on the ordered path and +1.4% on the -pipelined one. The reason they stop paying is the point of the section above: the -pipeline has already moved the bottleneck down to the compute floor, so removing -a quarter of the traffic removes something that was no longer being waited for. -Whether this still holds where the copy dominates by a wide margin has not been -measured. +Four device-resident layers are worth +14.3% on the ordered path and +1.4% on the pipelined one. The reason they stop paying is the point of the section above: the pipeline has already moved the bottleneck down to the compute floor, so removing a quarter of the traffic removes something that was no longer being waited for. Whether this still holds where the copy dominates by a wide margin has not been measured. ### 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: +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. -- Over it the scheduler declines and keeps the ordered path for that graph. - Later graphs are evaluated again, so a smaller live window can use the ring. -- Declining costs nothing in steady state. Both the ring and the transfer - backend's device context are released. - -**The cap is applied to what the current graph needs, not to what the full -context would need.** A run whose window stays small keeps the ring whatever -`-n_ctx` says, which is the common case and the reason it is done this way: a -staged input is a view of the cache tensor, so the full-context figure is there -for the asking, but enforcing it would refuse the ring for every large `-c` even -when the window never gets near it. The warning reports both numbers so that -`--kv-pipeline-budget` can be sized against the one that matters. - -The cost of deciding per graph is that a context which grows past the budget -allocates a ring for the small early windows and gives it back once it outgrows -them. That transient is bounded by the budget itself, which is the memory the -user already authorised, so it is a property of the cap rather than a defect in -it. - -At 32,768 the ring is 204 MiB at the full context, over the 128 MiB default. -`--kv-pipeline-budget 512` buys 20.350 -> 31.463 t/s behind an 18,432-token -prompt. +- Over it the scheduler declines and keeps the ordered path for that graph. Later graphs are evaluated again, so a smaller live window can use the ring. +- Declining costs nothing in steady state. Both the ring and the transfer backend's device context are released. +- The budget and the headroom check are decided before the graph is allocated, so they can still leave the graph short. If graph reservation fails, the rings are released and the reservation is retried once on the ordered path, and that scheduler keeps the ordered path from then on. An optional ring never turns a graph that fits into an allocation failure. + +**The cap is applied to what the current graph needs, not to what the full context would need.** A run whose window stays small keeps the ring whatever `-n_ctx` says, which is the common case and the reason it is done this way: a staged input is a view of the cache tensor, so the full-context figure is there for the asking, but enforcing it would refuse the ring for every large `-c` even when the window never gets near it. The warning reports both numbers so that `--kv-pipeline-budget` can be sized against the one that matters. + +The cost of deciding per graph is that a context which grows past the budget allocates a ring for the small early windows and gives it back once it outgrows them. That transient is bounded by the budget itself, which is the memory the user already authorised, so it is a property of the cap rather than a defect in it. + +At 32,768 the ring is 204 MiB at the full context, over the 128 MiB default. `--kv-pipeline-budget 512` buys 20.350 -> 31.463 t/s behind an 18,432-token prompt. ### Where the rest of the token goes @@ -291,126 +163,50 @@ Per decode graph, `GGML_SCHED_TRANSPORT_DEBUG=2`, behind a 19,246-token prompt: | bytes delivered early / late | 0 / 0 MiB | 644.0 / 2.3 MiB | | bytes left on the ordered path | 28.3 MiB | 0.4 MiB | -The blocking host-to-device copy is all but 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. 644 MiB in the 28.3 ms the -ordered arm reports for the same bytes is 22.0 GB/s, which is what this link -does; the transfer cannot be made faster, only hidden. - -The 3.57 ms that remains moves 0.4 MiB, and `GGML_SCHED_TRANSPORT_DEBUG=3` shows -that almost all of it is one copy: `attn_inp_k_rot`, 256 KiB, 18 us on the -ordered path and 3.4 ms behind one split of look-ahead. The 32 KV store copies -cost 353 us between them. - -It looks like latency and is not. A blocking copy shares the device's copy engine -with the deliveries and waits for what is already queued there: two staged splits -at 22.0 GB/s is 3.6 ms, which is the number. Two things were tried and neither -helped. Issuing the delivery in pieces so the blocking copy can interleave does -nothing -- the engine is FIFO across streams, `attn_inp_k_rot` stays at 3.4 ms at -every piece size, and small pieces cost throughput (29.80 t/s whole, 28.73 at -4 MiB, 22.01 at 1 MiB). Putting the copy on the consumer's own stream so the host -never blocks moves the time rather than removing it: the ordered copy falls from -3.57 ms to 0.16 ms, the consumer wait rises from 27.31 ms to 31.05 ms, and -throughput does not move (29.808 against 29.834). - -So this is not spare time. Those 256 KiB cross the same saturated link as the -644 MiB of deliveries, and the link is the ceiling. - -Do not compare these numbers against runs on other models, prompts, cache -settings, hardware, or commits. +The blocking host-to-device copy is all but 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. 644 MiB in the 28.3 ms the ordered arm reports for the same bytes is 22.0 GB/s, which is what this link does; the transfer cannot be made faster, only hidden. + +The 3.57 ms that remains moves 0.4 MiB, and `GGML_SCHED_TRANSPORT_DEBUG=3` shows that almost all of it is one copy: `attn_inp_k_rot`, 256 KiB, 18 us on the ordered path and 3.4 ms behind one split of look-ahead. The 32 KV store copies cost 353 us between them. + +It looks like latency and is not. A blocking copy shares the device's copy engine with the deliveries and waits for what is already queued there: two staged splits at 22.0 GB/s is 3.6 ms, which is the number. Two things were tried and neither helped. Issuing the delivery in pieces so the blocking copy can interleave does nothing -- the engine is FIFO across streams, `attn_inp_k_rot` stays at 3.4 ms at every piece size, and small pieces cost throughput (29.80 t/s whole, 28.73 at 4 MiB, 22.01 at 1 MiB). Putting the copy on the consumer's own stream so the host never blocks moves the time rather than removing it: the ordered copy falls from 3.57 ms to 0.16 ms, the consumer wait rises from 27.31 ms to 31.05 ms, and throughput does not move (29.808 against 29.834). + +So this is not spare time. Those 256 KiB cross the same saturated link as the 644 MiB of deliveries, and the link is the ceiling. + +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` compares - every requested depth with the first and fails on a hash difference. - Two things keep the tasks independent of each other, and both were needed. - Every task carries a nonce derived from its own name and length, so no two - share a prefix the server could restore, and the harness fails a task whose - `prompt_n` says one was reused anyway. Each request also sets - `cache_prompt: false`, so a task never inherits what the previous one left in - the cache. - - The second is what made `records@18432` a gate rather than a coin flip. Its - prompt is about 29.6k tokens against a 32,768 context, and the task before it - is about the same size, so the two do not both fit and placement depended on - what was still resident. Two otherwise identical `N = 0` runs of it produced - different hashes. Asked on its own with the cache off it is perfectly stable: - the same hash three times running, at `-c 32768` and at `-c 65536`. With the - flag set, two independent `N = 0` passes agree on all eight tasks, and - `N = 0`, `N = 1` and `N = 4` agree on all eight. -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. +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` compares every requested depth with the first and fails on a hash difference. Two things keep the tasks independent of each other, and both were needed. Every task carries a nonce derived from its own name and length, so no two share a prefix the server could restore, and the harness fails a task whose `prompt_n` says one was reused anyway. Each request also sets `cache_prompt: false`, so a task never inherits what the previous one left in the cache. + + The second is what made `records@18432` a gate rather than a coin flip. Its prompt is about 29.6k tokens against a 32,768 context, and the task before it is about the same size, so the two do not both fit and placement depended on what was still resident. Two otherwise identical `N = 0` runs of it produced different hashes. Asked on its own with the cache off it is perfectly stable: the same hash three times running, at `-c 32768` and at `-c 65536`. With the flag set, two independent `N = 0` passes agree on all eight tasks, and `N = 0`, `N = 1` and `N = 4` agree on all eight. +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, as the mean over each 128 graphs, with - depth stops and the number of recycle waits enqueued; `=3` names the tensors - still on the ordered path. `ggml_backend_sched_get_transport_pipeline_stats()` - exposes deliveries and early and late byte counts to callers. - -A device-resident KV run is unaffected, and was measured to confirm it: 38.5612 -t/s at depth 0 against 38.5240 at depth 1, `tg128 @ d4096`, with the -transport never enabled because the scheduler is given a depth of 0. +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, as the mean over each 128 graphs, with depth stops and the number of recycle waits enqueued; `=3` names the tensors still on the ordered path. `ggml_backend_sched_get_transport_pipeline_stats()` exposes deliveries and early and late byte counts to callers. + +A device-resident KV run is unaffected, and was measured to confirm it: 38.5612 t/s at depth 0 against 38.5240 at depth 1, `tg128 @ d4096`, with the transport never enabled because the scheduler is given a depth of 0. ## Scope and limits -- Only persistent host inputs marked with `GGML_TENSOR_FLAG_TRANSPORT` are - candidates. The stable prefix remains a per-evaluation value. Unmarked inputs, - weights, user inputs, transposed V, and copies with later readers stay ordered. -- CUDA is the only enabled backend. Meta, SYCL, WebGPU, and other backends stay - ordered until their event behavior and transport path are validated. -- 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 per accelerator.** A layer-split model pipelines on every device - that qualifies; a device with no room within the budget falls back to the - ordered path on its own without disabling the others. -- **Tensor parallelism keeps the ordered path.** See - [Tensor parallelism](#tensor-parallelism). -- 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. +- Only persistent host inputs marked with `GGML_TENSOR_FLAG_TRANSPORT` are candidates. The stable prefix remains a per-evaluation value. Unmarked inputs, weights, user inputs, transposed V, and copies with later readers stay ordered. +- CUDA is the only enabled backend. Meta, SYCL, WebGPU, and other backends stay ordered until their event behavior and transport path are validated. +- 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 per accelerator.** A layer-split model pipelines on every device that qualifies; a device with no room within the budget falls back to the ordered path on its own without disabling the others. +- **Tensor parallelism keeps the ordered path.** See [Tensor parallelism](#tensor-parallelism). +- 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` and `GGML_KV_PIPELINE_BUDGET_MIB` provide scheduler defaults. Explicit scheduler settings and command-line options take precedence. ## Tensor parallelism -`-sm tensor` is not pipelined. The scheduler explicitly excludes meta devices. -A host-resident cache needs a validated strided head-split write before this can -be enabled. - -Both sit behind a correctness problem that is not this feature's: -**`-sm tensor` together with `--no-kv-offload` currently produces wrong output.** -On one build and one prompt, `-sm layer --no-kv-offload` and `-sm tensor` with a -device-resident cache agree exactly, while `-sm tensor --no-kv-offload` differs. -It does not crash or warn; it generates fluent, different text. - -The cause is the GQA head mapping. Tensor parallelism splits attention by head, -but a host-resident cache is one undivided tensor, so the scheduler's copy of it -is classified `MIRRORED` and the whole window goes to every device. With 24 query -heads split 12/12 and 4 KV heads mirrored, the kernel derives the GQA ratio from -the tensors it is handed -- 12/4 = 3 rather than 6 -- and the second device's -queries, renumbered from 0, read the first device's keys. With an uneven split the -same fault surfaces as a crash instead: -`GGML_ASSERT(Q->ne[2] % K->ne[2] == 0)`, because 24 heads split 13/11 is not -divisible by 4. - -Head-splitting the copy rather than mirroring it fixes it. That was prototyped -and reproduced the layer-split output byte for byte, and needs four coordinated -changes: classify the scheduler's copy at all (it is a leaf in a compute buffer, -so it never reaches the device's split-state callback), use the head axis for the -permuted `[head_dim, n_kv, n_head_kv, 1]` shape rather than the cache tensor's own -axis, express the granularity in heads aligned to the query split divided by the -GQA ratio, and add a strided write because the heads are interleaved within each -row rather than laid out end to end. +`-sm tensor` is not pipelined. The scheduler explicitly excludes meta devices. A host-resident cache needs a validated strided head-split write before this can be enabled. + +Both sit behind a correctness problem that is not this feature's: **`-sm tensor` together with `--no-kv-offload` currently produces wrong output.** On one build and one prompt, `-sm layer --no-kv-offload` and `-sm tensor` with a device-resident cache agree exactly, while `-sm tensor --no-kv-offload` differs. It does not crash or warn; it generates fluent, different text. + +The cause is the GQA head mapping. Tensor parallelism splits attention by head, but a host-resident cache is one undivided tensor, so the scheduler's copy of it is classified `MIRRORED` and the whole window goes to every device. With 24 query heads split 12/12 and 4 KV heads mirrored, the kernel derives the GQA ratio from the tensors it is handed -- 12/4 = 3 rather than 6 -- and the second device's queries, renumbered from 0, read the first device's keys. With an uneven split the same fault surfaces as a crash instead: `GGML_ASSERT(Q->ne[2] % K->ne[2] == 0)`, because 24 heads split 13/11 is not divisible by 4. + +Head-splitting the copy rather than mirroring it fixes it. That was prototyped and reproduced the layer-split output byte for byte, and needs four coordinated changes: classify the scheduler's copy at all (it is a leaf in a compute buffer, so it never reaches the device's split-state callback), use the head axis for the permuted `[head_dim, n_kv, n_head_kv, 1]` shape rather than the cache tensor's own axis, express the granularity in heads aligned to the query split divided by the GQA ratio, and add a strided write because the heads are interleaved within each row rather than laid out end to end. ## Future work -- Fix `-sm tensor` with `--no-kv-offload` (above). Until then it should not be - used: it is wrong rather than slow. +- Fix `-sm tensor` with `--no-kv-offload` (above). Until then it should not be used: it is wrong rather than slow. - Add the strided head-split delivery above, validate it, and then measure it. diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index a9e97ccf211..50e4d081089 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -325,29 +325,23 @@ extern "C" { // 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. + // 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 right before the kernels that read it. + // With it the scheduler keeps a ring of staging slots outside the graph allocator's reach and issues the stable prefix of a later split on a separate transfer stream, so the transfer retires under the kernels of the split before it. // - // Only persistent host inputs marked with GGML_TENSOR_FLAG_TRANSPORT are eligible. Their - // stable prefix must be current before each evaluation. The producer must be the CPU or the - // same backend stream that consumes the late region. + // Only persistent host inputs marked with GGML_TENSOR_FLAG_TRANSPORT are eligible, and their stable prefix must be current before each evaluation. + // The producer must be the CPU or the same backend stream that consumes the late region. // - // `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. Returns false after graph allocation starts. + // `depth` is how many splits ahead deliveries run, 0 disables pipelining. + // The ring holds a couple of slots more than that, so recycling a slot never waits for a reader that is still running. + // Needs a destination backend with asynchronous transfers and events, otherwise the setting is ignored. + // Costs roughly (depth + 2) * (largest staged split) of device memory. + // Must be called before the first graph is allocated, and returns false after that. + // The ring is optional: if the graph cannot be allocated next to it, the scheduler releases the ring and keeps the ordered path. GGML_API bool 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. - // Returns false after graph allocation starts. Configuration is immutable then. + // Hard cap on the staging ring, in bytes, default 128 MiB and 0 removes the cap. + // A host-resident cache exists to keep device memory free, so the ring is capped outright rather than against what happens to be free: past the cap the scheduler declines and keeps the ordered path. + // Must be called before the first graph is allocated, and returns false after that. GGML_API bool 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. diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index a4a18f42d83..a76a3bfd5f9 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -711,10 +711,9 @@ extern "C" { 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. + // 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), and it must be set on the tensor that owns the storage, not on a view of it + // 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); diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 77cace14551..8e75d909236 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -777,43 +777,37 @@ static bool ggml_is_view_op(enum ggml_op op) { #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. +// 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 every delivery would recycle the split that was enqueued a moment ago and is still running -- 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 -// 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. +// Device memory the transport ring leaves unclaimed. +// 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. #ifndef GGML_SCHED_TRANSPORT_HEADROOM #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. +// 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, so 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. +// One staging slot of the transport ring. +// A slot is owned by the transfer stream while it is filled and by the consumer stream while it is read, and 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 }; -// 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. +// One ring per accelerator the scheduler drives. +// A layer-split model gives every device its own splits and deliveries, so each needs its own transfer stream, staging and 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 @@ -833,13 +827,9 @@ struct ggml_backend_sched_transport_ring { // 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. +// 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 per ring: depth + GGML_SCHED_TRANSPORT_MARGIN @@ -848,17 +838,15 @@ struct ggml_backend_sched_transport { 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 within its - // own backend's ring, 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; // over all rings, so that execution can skip the machinery entirely - // 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. + // 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 go 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; @@ -875,8 +863,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; - // why the look-ahead stopped: the depth it was given, or a slot whose reader has not run yet. - // The two want opposite fixes, so they are counted apart. + // why the look-ahead stopped: the depth it was given, or a slot whose reader has not run yet + // the two want opposite fixes, so they are counted apart int64_t n_stop_depth; int64_t n_wait_recycle; int64_t p_stop_depth, p_wait_recycle; @@ -1738,11 +1726,9 @@ static bool ggml_backend_sched_transport_enabled(ggml_backend_sched_t sched) { return false; } -// 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. +// The annotation lives on the tensor that owns the storage, and 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) { @@ -1762,11 +1748,9 @@ static size_t ggml_backend_sched_input_stable_prefix(const struct ggml_tensor * // 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 -// 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. +// 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 and 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 (!ggml_backend_sched_transport_ring_enabled(sched, split->backend_id)) { @@ -1807,8 +1791,7 @@ static bool ggml_backend_sched_input_is_staged(ggml_backend_sched_t sched, int s static bool ggml_backend_sched_size_add(size_t a, size_t b, size_t * result); static bool ggml_backend_sched_size_pad(size_t size, size_t alignment, size_t * result); -// A ring entry costs what the backend would allocate for it, which can be more than its data: -// a buffer type may ask for padding past ggml_nbytes(), and its kernels may write into it. +// A ring entry costs what the backend would allocate for it, which can be more than its data: a buffer type may ask for padding past ggml_nbytes(), and its kernels may write into it. static bool ggml_backend_sched_transport_entry_size( ggml_backend_buffer_type_t buft, const struct ggml_tensor * t, size_t alignment, size_t * result) { return ggml_backend_sched_size_pad(ggml_backend_buft_get_alloc_size(buft, t), alignment, result); @@ -1849,8 +1832,8 @@ static void ggml_backend_sched_transport_assign_addresses(ggml_backend_sched_t s continue; } struct ggml_tensor * input_cpy = tensor_copy(split->inputs[j], split->backend_id, sched->cur_copy); - // bind through the backend, so that the entry is initialized the same way as any other - // tensor the buffer holds. A previous plan may have left this copy bound already. + // bind through the backend, so the entry is initialized the same way as any other tensor the buffer holds + // a previous plan may have left this copy bound already input_cpy->data = NULL; input_cpy->buffer = NULL; const enum ggml_status status = ggml_backend_tensor_alloc(r->buffer, input_cpy, slot + offset); @@ -1863,9 +1846,8 @@ static void ggml_backend_sched_transport_assign_addresses(ggml_backend_sched_t 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. +// 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, int backend_id, bool sync_consumers) { struct ggml_backend_sched_transport_ring * r = &sched->transport.rings[backend_id]; @@ -1931,6 +1913,23 @@ static void ggml_backend_sched_transport_decline_backend(ggml_backend_sched_t sc } } +// Give every ring back and stop asking for one. +// The ring is optional, so it is the first thing to release when the device cannot hold it and the graph at the same time. +// Returns whether any ring was holding memory. +static bool ggml_backend_sched_transport_decline_all(ggml_backend_sched_t sched) { + struct ggml_backend_sched_transport * tr = &sched->transport; + + bool released = false; + for (int i = 0; i < sched->n_backends; i++) { + released |= tr->rings[i].buffer != NULL; + ggml_backend_sched_transport_decline_backend(sched, i); + tr->rings[i].eligible = false; + } + tr->n_staged = 0; + + return released; +} + static bool ggml_backend_sched_size_add(size_t a, size_t b, size_t * result) { if (a > SIZE_MAX - b) { return false; @@ -1957,9 +1956,7 @@ static bool ggml_backend_sched_size_pad(size_t size, size_t alignment, size_t * return ggml_backend_sched_size_add(size, alignment - rem, result); } -// 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. +// The transfer backend and the slot events are created on demand, so a backend that 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]; @@ -2000,9 +1997,8 @@ static bool ggml_backend_sched_transport_ensure_backend(ggml_backend_sched_t sch return true; } -// 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. +// 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; @@ -2040,6 +2036,14 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { } tr->split_input_ofs[sched->n_splits] = n_inputs_total; + // a split list without inputs has nothing to stage, and input_staged is still unallocated + if (n_inputs_total == 0) { + for (int i = 0; i < sched->n_splits; i++) { + tr->split_order[i] = -1; + } + return; + } + if (tr->input_capacity < n_inputs_total) { unsigned char * pnew = (unsigned char *) realloc(tr->input_staged, n_inputs_total); if (pnew == NULL) { @@ -2071,8 +2075,8 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { return; } - // A ring copy must have one owner and no views or later readers. - // Build one lookup table, then scan each graph node once. + // a ring copy must have one owner and no views or later readers + // build one lookup table, then scan each graph node once size_t staged_hash_size = n_candidates; staged_hash_size += staged_hash_size/4 + 1; struct ggml_hash_set staged_copies = ggml_hash_set_new(staged_hash_size); @@ -2137,11 +2141,10 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { free(staged_owner); ggml_hash_set_free(&staged_copies); - // per-ring slot size and delivery order. The budget is applied to what this graph needs, so - // that a run whose window stays small keeps the ring whatever -n_ctx says. slot_size_max is - // what the same ring costs once the context is full, taken from the cache tensor the staged - // input is a view of; it is reported rather than enforced, because deciding on it would refuse - // the ring for every large -c even when the window never gets there. + // per-ring slot size and delivery order + // the budget is applied to what this graph needs, so a run whose window stays small keeps the ring whatever -n_ctx says + // slot_size_max is what the same ring costs once the context is full, taken from the cache tensor the staged input is a view of + // it is reported rather than enforced, because deciding on it would refuse the ring for every large -c even when the window never gets there size_t slot_size[GGML_SCHED_MAX_BACKENDS] = { 0 }; size_t slot_size_max[GGML_SCHED_MAX_BACKENDS] = { 0 }; bool size_overflow[GGML_SCHED_MAX_BACKENDS] = { false }; @@ -2225,8 +2228,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { 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. + // 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) { @@ -2302,12 +2304,9 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { ggml_backend_sched_transport_assign_addresses(sched); } -// 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. +// 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 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]; @@ -2328,9 +2327,8 @@ static void ggml_backend_sched_transport_prefetch(ggml_backend_sched_t sched, in struct ggml_backend_sched_split * split = &sched->splits[i]; 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. + // 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(r->transfer, slot->release); slot->release_armed = false; @@ -2343,9 +2341,8 @@ static void ggml_backend_sched_transport_prefetch(ggml_backend_sched_t sched, in } 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 + // 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; @@ -2361,10 +2358,8 @@ static void ggml_backend_sched_transport_prefetch(ggml_backend_sched_t sched, in 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. + // 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, r->transfer); r->scan_cursor = i + 1; @@ -2390,8 +2385,7 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { } } - // lay out the transport rings and point the staged input copies at them before the graph is - // allocated, so ggml-alloc sees those copies as already allocated and leaves them alone + // lay out the transport rings and point the staged input copies at them 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 @@ -2424,8 +2418,13 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { ggml_backend_sched_transport_clear_addresses(sched); if (!ggml_gallocr_reserve_n(sched->galloc, &sched->graph, sched->node_backend_ids, sched->leaf_backend_ids)) { - GGML_LOG_ERROR("%s: failed to allocate graph\n", __func__); - return false; + // the rings hold device memory the graph itself may need, and the caller can no longer turn them off: give them back and reserve once on the ordered path + if (!ggml_backend_sched_transport_decline_all(sched) || + !ggml_gallocr_reserve_n(sched->galloc, &sched->graph, sched->node_backend_ids, sched->leaf_backend_ids)) { + GGML_LOG_ERROR("%s: failed to allocate graph\n", __func__); + return false; + } + GGML_LOG_WARN("%s: the graph does not fit next to the transport rings, they are released and this scheduler stays on the ordered path\n", __func__); } ggml_backend_sched_transport_assign_addresses(sched); if (!ggml_gallocr_alloc_graph(sched->galloc, &sched->graph)) { @@ -2449,8 +2448,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s struct ggml_backend_sched_transport * tr = &sched->transport; bool named_ordered_now = false; - // a reused graph keeps the plan that was made for it, so the split list it describes must be - // the one about to run + // 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; @@ -2458,10 +2456,9 @@ 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; - // 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. The cursors start over on every evaluation - // because the plan outlives the graph it was made for. + // Prime every ring before the first consumer runs. + // From here on deliveries are issued only after a split has been enqueued, never before, so recycling a slot can never hold back work the consumer could already be running. + // The cursors start over on every evaluation because the plan outlives the graph it was made for. if (staged) { for (int i = 0; i < sched->n_backends; i++) { tr->rings[i].consumed = 0; @@ -2500,11 +2497,9 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s 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. + // 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 const size_t prefix = ggml_backend_sched_input_stable_prefix(input); const size_t nbytes = ggml_nbytes(input); if (nbytes > prefix) { @@ -2706,16 +2701,14 @@ 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 + // 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; tr->rings[split_backend_id].consumed++; tr->n_deliveries++; - // with this split's kernels already enqueued, the deliveries for the next staged - // splits can go out even if recycling their slot waits for a reader that is running + // with this split's kernels already enqueued, the deliveries for the next staged splits can go out even if recycling their slot waits for a reader that is running ggml_backend_sched_transport_prefetch(sched, split_backend_id); } @@ -2735,8 +2728,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s tr->t_graph_us += ggml_time_us() - t_graph_0; tr->n_graphs++; - // every 128 graphs, and as the mean over those 128, so that one graph's noise does not - // decide what the numbers look like + // every 128 graphs, and as the mean over those 128, so one graph's noise does not decide what the numbers look like if (tr->n_graphs % 128 == 0) { const double n = 128.0; GGML_LOG_INFO("%s: per graph over %d: total %.2f ms, sync %.2f ms, ordered copy %.2f ms, " @@ -2948,8 +2940,8 @@ bool ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, 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 + // 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; } diff --git a/include/llama.h b/include/llama.h index b9b434efa01..e91225dd174 100644 --- a/include/llama.h +++ b/include/llama.h @@ -421,14 +421,11 @@ extern "C" { // A source/target/parent context that can share results or llama_memory. struct llama_context * ctx_other; - 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. - 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. + uint32_t kv_pipeline_depth; // how many splits ahead the scheduler delivers a host-resident KV cache, 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 + uint32_t kv_pipeline_budget_mib; // hard cap on that device memory, in MiB, past which the scheduler keeps the ordered path + // a host-resident cache never quietly trades back the device memory it exists to save, 0 removes the cap }; struct llama_model_tensor_override { diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 51b6e76cc11..676789fc5ca 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1213,8 +1213,7 @@ 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 + // before the graph is built and allocated, so 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 @@ -1650,9 +1649,8 @@ ggml_tensor * llama_kv_cache::build_input_v_rot(ggml_context * ctx) const { } 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. + // 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(); @@ -1671,8 +1669,7 @@ void llama_kv_cache::update_stable_prefixes(const slot_info & sinfo) const { 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 + // 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]); } } diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index e09d7a044ee..4188f55e5ab 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -238,10 +238,8 @@ class llama_kv_cache : public llama_memory_i { bool is_reserve, const slot_info * sinfo = nullptr) 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. + // 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; diff --git a/tests/test-alloc.cpp b/tests/test-alloc.cpp index 7f1bd3adaa3..3f7ec16265e 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -19,6 +19,7 @@ uint8_t * const alloc_base = (uint8_t *) 16; struct dummy_backend_context { size_t max_buffer_size = 64; + size_t capacity = SIZE_MAX; // what the device can hold at one time size_t alignment = 8; bool fail_alloc = false; bool unique_alloc_addresses = false; @@ -34,8 +35,7 @@ struct dummy_backend_context { size_t set_tensor_async_bytes = 0; size_t alloc_size_pad = 0; - // what the backend was asked to hold and to move, so that a test can check the entries a - // transport ring lays out and the bytes it delivers into them + // what the backend was asked to hold and to move, so a test can check the entries a transport ring lays out and the bytes it delivers into them struct tensor_binding { const ggml_tensor * tensor; ggml_backend_buffer_t buffer; @@ -80,7 +80,7 @@ static const char * dummy_backend_buffer_type_get_name(ggml_backend_buffer_type_ static ggml_backend_buffer_t dummy_backend_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { dummy_backend_context * ctx = (dummy_backend_context *) buft->context; - if (ctx->fail_alloc) { + if (ctx->fail_alloc || size > ctx->capacity - ctx->allocated_total()) { return nullptr; } ggml_backend_buffer_t & buffer = ctx->buffers.emplace_back(); @@ -1374,8 +1374,7 @@ static void test_transport_prefix_and_configuration() { GGML_ASSERT(cuda.context->transfer_backend_count == 0); } -// The ring must hold what the backend allocates for an entry, deliver every byte of it exactly -// once, and never wait on an event it has not recorded. +// The ring must hold what the backend allocates for an entry, deliver every byte of it exactly once, and never wait on an event it has not recorded. static void test_transport_entry_allocation() { dummy_backend cuda = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); @@ -1512,6 +1511,46 @@ static void test_transport_fallback_keeps_allocator_plan() { GGML_ASSERT(ordered > 0 && pipelined == ordered); } +static size_t transport_scale_buffer_size(int depth, size_t nbytes, size_t capacity, int64_t * deliveries, int * transfers) { + dummy_backend cuda = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); + dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); + cuda.context->capacity = capacity; + + auto graph = make_transport_graph(cpu, nbytes); + ggml_set_stable_prefix(graph.source, nbytes); + + ggml_backend_t backends[] = { cuda.handle.get(), cpu.handle.get() }; + ggml_backend_buffer_type_t bufts[] = { &cuda.buffer_type, &cpu.buffer_type }; + ggml_backend_sched_ptr sched(ggml_backend_sched_new(backends, bufts, 2, 128, false, false)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_depth(sched.get(), depth)); + ggml_backend_sched_set_tensor_backend(sched.get(), graph.output, cuda.handle.get()); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), graph.ctx.graph)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), graph.ctx.graph) == GGML_STATUS_SUCCESS); + + if (deliveries) { + transport_stats(sched.get(), deliveries, nullptr, nullptr); + } + if (transfers) { + *transfers = cuda.context->transfer_backend_count; + } + return ggml_backend_sched_get_buffer_size(sched.get(), cuda.handle.get()); +} + +// the ring is optional, so a device that cannot hold it next to the graph keeps the graph +static void test_transport_releases_ring_for_graph() { + const size_t nbytes = 256; + const size_t ring = 3*nbytes; // depth 1 plus the margin, one entry per slot + const size_t ordered = transport_scale_buffer_size(0, nbytes, SIZE_MAX, nullptr, nullptr); + GGML_ASSERT(ordered > 0); + + int64_t deliveries = -1; + int transfers = -1; + const size_t pipelined = transport_scale_buffer_size(1, nbytes, std::max(ordered, ring), &deliveries, &transfers); + GGML_ASSERT(pipelined == ordered); + GGML_ASSERT(deliveries == 0); + GGML_ASSERT(transfers == 0); +} + static void set_test_env(const char * name, const char * value) { #ifdef _WIN32 GGML_ASSERT(_putenv_s(name, value) == 0); @@ -1807,6 +1846,7 @@ int main() { run("test_transport_entry_allocation", test_transport_entry_allocation); run("test_transport_empty_graph", test_transport_empty_graph); run("test_transport_fallback_keeps_allocator_plan", test_transport_fallback_keeps_allocator_plan); + run("test_transport_releases_ring_for_graph", test_transport_releases_ring_for_graph); run("test_transport_environment_is_fallback", test_transport_environment_is_fallback); run("test_transport_depth_zero", test_transport_depth_zero); run("test_transport_budget_recovers", test_transport_budget_recovers); From 72df63340ddc205349300b6087e0d8f158ed7e8f Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 3 Sep 2026 10:59:59 +0200 Subject: [PATCH 10/32] docs: re-measure the transport gates on the current head The llama-bench table predated the budget and said so, and the context sweep and the exactness gate were last run before the ring allocation and binding changed. All three are re-run on an RTX 4070 with a CUDA build of this head, at --kv-pipeline-budget 512. Greedy server output is identical at depth 0, 1 and 4 across all eight tasks. Throughput is +16.1% at 4,096, +56.6% at 16,384, +20.1% at 32,768 and +13.4% at 65,536, for +28, +104, +206 and +410 MiB of device memory. The 131,072 and 262,144 arms are not re-measured and say so. The server table is replaced with the four 18,432-prefill tasks of the exactness gate, which is what this head was actually run on; the copy/compute breakdown keeps its earlier numbers and says which head they came from. Assisted-by: Claude Opus 5 --- docs/kv-transport-pipelining.md | 41 +++++++++++++++++++++------------ 1 file changed, 26 insertions(+), 15 deletions(-) diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index 467fada3e69..854d3444a8b 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -58,19 +58,28 @@ Getting either of these wrong costs the entire gain while still producing correc 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`, `docs/repro/r4-kv-pipeline-context-sweep.sh`), with no cap on the ring: +`llama-bench`, `--no-warmup`, A/B/A/B with reversed arm order (`docs/repro/r4-kv-pipeline-ab.sh`), at `--kv-pipeline-budget 512`, both passes shown: -| depth | ordered | pipelined | gain | peak device memory | -|---:|---|---|---:|---:| -| 4,096 | 31.3066, 31.2334 | 36.6107, 36.4701 | **+17.0%** | +28 MiB | -| 16,384 | 19.4344, 19.4828 | 31.0981, 31.0931 | **+59.8%** | +86 to +104 MiB | -| 32,768 | 12.9453, 12.9400 | 15.4571, 15.4516 | **+19.4%** | +206 MiB | +| depth | ordered | pipelined | gain | +|---:|---|---|---:| +| 4,096 | 29.9802, 29.9776 | 34.8016, 34.8047 | **+16.1%** | +| 16,384 | 19.0116, 19.0093 | 29.7780, 29.7879 | **+56.6%** | +| 32,768 | 12.7237, 12.7237 | 15.2873, 15.2864 | **+20.1%** | -> These rows were taken before the budget existed, so they are the uncapped numbers. The 32,768 ring is 204 MiB and the default cap is 128 MiB, so that row does not reproduce on the current default: it needs `-kvpb 512`, which `llama-bench` did not take until now. The scripts pass it, and the row is due a re-measurement on the current head. +The 32,768 ring is 204 MiB at the full context, over the 128 MiB default, so that row needs `-kvpb 512`. `llama-bench` takes the option and the scripts pass it. > These also need `-kvcp 1 -rso 1`, and for a while `llama-bench` did not have them: the repro scripts probed `--help`, found nothing, and quietly dropped both. The same commit then measures 19.43 -> 9.02 t/s ordered at 16,384 and the pipeline buys +6.7% instead of +60%, because a host-resident recurrent state costs more than the transport can win back. `llama-bench` takes them again, and the scripts now fail rather than drop an option the build does not have. -`llama-server`, one request, `temperature 0, top_k 1, seed 1234`: +`llama-server`, one request, `temperature 0, top_k 1, seed 1234`, the four 18,432-prefill tasks of the exactness gate at `-c 32768`: + +| task | prompt | ordered | pipelined | gain | +|---|---:|---:|---:|---:| +| prose | 14,821 | 19.775 | 30.012 | **+51.8%** | +| dialogue | 15,984 | 19.155 | 29.767 | **+55.4%** | +| records | 29,603 | 13.553 | 16.396 | **+21.0%** | +| code | 29,670 | 13.504 | 16.362 | **+21.2%** | + +The breakdown below was taken on an earlier head, at prompts of 19,246 and 48,042 against `-c 32768` and `-c 65536`: | prompt | `-c` | ordered | pipelined | gain | copy ms | compute ms | ceiling | share | |---:|---:|---:|---:|---:|---:|---:|---:|---:| @@ -88,7 +97,7 @@ Pinning is worth as much as the pipeline and is off by default. Behind a 13,128- | `--kv-cpu-pinned` | 21.582 | 32.252 | | unpinned | 14.945 | 22.709 | -Look-ahead deeper than one split is worse at every depth measured. At 19,246: 29.83 t/s at `N = 1`, 28.44 at `N = 2`, 25.93 at `N = 4`. `N = 1` is the default for that reason. +Look-ahead deeper than one split is worse at every depth measured. At 19,246: 29.83 t/s at `N = 1`, 28.44 at `N = 2`, 25.93 at `N = 4`. `N = 1` is the default for that reason. The exactness gate measures the same on every one of its four 18,432-prefill tasks: 30.012 against 27.037 on prose, 29.767 against 26.584 on dialogue, 16.396 against 15.934 on records, 16.362 against 15.857 on code. ### The link is the ceiling, so the lever is bytes @@ -111,12 +120,12 @@ Halving the traffic is worth +5.6% on the pipelined path at 19,246 and +56.6% at | 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 | +| 4,096 | 29.90 | 34.72 | **+16.1%** | 10,121 -> 10,149 MiB | +28 MiB | 27 MiB | +| 16,384 | 18.98 | 29.72 | **+56.6%** | 10,111 -> 10,215 MiB | +104 MiB | 107 MiB | +| 32,768 | 12.70 | 15.26 | **+20.1%** | 10,113 -> 10,319 MiB | +206 MiB | 213 MiB | +| 65,536 | 7.63 | 8.66 | **+13.4%** | 10,115 -> 10,525 MiB | +410 MiB | 428 MiB | + +The two deepest arms were taken on an earlier head and are not re-measured here: +9.1% at 131,072 for +818 MiB of ring, and at 262,144 the ring is declined and the two arms are the same. Two curves run in opposite directions here, and both matter. @@ -177,6 +186,8 @@ Do not compare these numbers against runs on other models, prompts, cache settin The gates, and what was run for them: +All four gates below were re-run on the current head, on an RTX 4070 with a CUDA build. + 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` compares every requested depth with the first and fails on a hash difference. Two things keep the tasks independent of each other, and both were needed. Every task carries a nonce derived from its own name and length, so no two share a prefix the server could restore, and the harness fails a task whose `prompt_n` says one was reused anyway. Each request also sets `cache_prompt: false`, so a task never inherits what the previous one left in the cache. The second is what made `records@18432` a gate rather than a coin flip. Its prompt is about 29.6k tokens against a 32,768 context, and the task before it is about the same size, so the two do not both fit and placement depended on what was still resident. Two otherwise identical `N = 0` runs of it produced different hashes. Asked on its own with the cache off it is perfectly stable: the same hash three times running, at `-c 32768` and at `-c 65536`. With the flag set, two independent `N = 0` passes agree on all eight tasks, and `N = 0`, `N = 1` and `N = 4` agree on all eight. From 782fc6afe4d756d9e8ed8b77cf3e1a9ebdb75af5 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Fri, 4 Sep 2026 21:27:58 +0200 Subject: [PATCH 11/32] docs: say where the transport stops paying, and why the budget is 128 MiB The context sweep is measured with the budget raised, so its deep rows read as default behaviour when they are not: past 20,556 rows of window the default declines and those depths stay ordered. Say so, and add a finer sweep that puts the peak at 16,384 rows and shows the gain per MiB falling off as 1/rows^2 above it. The peak is where copy and compute are equal, and the ring size there works out to (n_slots / n_attn) * compute * BW - the bytes per row cancel, so the budget is quant-invariant. That predicts 101 MiB against the 102 MiB measured, which is what the 128 MiB default is sized against. Assisted-by: Claude Opus 5 --- docs/kv-transport-pipelining.md | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index 854d3444a8b..6a9aac48798 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -127,6 +127,26 @@ Halving the traffic is worth +5.6% on the pipelined path at 19,246 and +56.6% at The two deepest arms were taken on an earlier head and are not re-measured here: +9.1% at 131,072 for +818 MiB of ring, and at 262,144 the ring is declined and the two arms are the same. +**Every row above 20,556 rows of window is measured with the budget raised**, at `-kvpb 0` or `-kvpb 512`. They are what the ring costs and buys if you pay for it, not what a default run does: at the default 128 MiB the ring declines past 20,556 rows and those depths stay on the ordered path. See [The budget](#the-budget). + +A finer sweep of the same configuration, `-kvpb 0` throughout so nothing declines, locates the peak between 16k and 24k: + +| rows | ordered | pipelined | gain | ring | gain per MiB | +|---:|---:|---:|---:|---:|---:| +| 2,048 | 32.76 | 35.42 | +8.1% | 13 MiB | 0.62 | +| 4,096 | 29.76 | 34.51 | +16.0% | 26 MiB | 0.62 | +| 8,192 | 25.02 | 32.71 | +30.8% | 51 MiB | 0.60 | +| 12,288 | 21.53 | 31.04 | +44.2% | 77 MiB | 0.57 | +| **16,384** | 18.90 | 29.44 | **+55.8%** | 102 MiB | **0.55** | +| 24,576 | 15.18 | 18.84 | +24.1% | 153 MiB | 0.16 | +| 32,768 | 12.65 | 15.19 | +20.1% | 204 MiB | 0.099 | +| 49,152 | 9.51 | 10.99 | +15.6% | 306 MiB | 0.051 | +| 65,536 | 7.63 | 8.66 | +13.6% | 408 MiB | 0.033 | +| 98,304 | 5.45 | 6.06 | +11.1% | 612 MiB | 0.018 | +| 131,072 | 4.26 | 4.66 | +9.5% | 816 MiB | 0.012 | + +The gain per MiB is flat below the peak and falls off as `1/rows^2` above it, because the gain decays as `compute/copy` while the ring grows linearly. There is no depth at which the pipeline becomes slower -- only one past which the memory buys more elsewhere. + 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. @@ -159,6 +179,19 @@ The cost of deciding per graph is that a context which grows past the budget all At 32,768 the ring is 204 MiB at the full context, over the 128 MiB default. `--kv-pipeline-budget 512` buys 20.350 -> 31.463 t/s behind an 18,432-token prompt. +#### Why 128 MiB + +The best the ring can do is hide the smaller of copy and compute behind the larger, so its value peaks where the two are equal. Writing `b` for the bytes one attention layer holds per row of window, the copy is `n_attn * b * rows / BW` and the peak is at + +``` +rows* = compute * BW / (n_attn * b) +ring* = n_slots * b * rows* = (n_slots / n_attn) * compute * BW +``` + +**`b` cancels.** The ring size at the peak does not depend on the cache type or on `n_embd_k_gqa`; it is set by the share of the traffic the ring holds, the compute a graph has to hide behind it, and the link. On this configuration -- 3 slots against 16 staged attention layers, 25.7 ms of consumer wait, 22.0 GB/s -- that is `0.1875 * 25.7e-3 * 22e9`, or **101 MiB against the 102 MiB measured at the +55.8% peak**. + +So the default is a size, not a depth, and it is the right kind of quantity to fix: a cache quantised to q4_0 doubles `rows*` and halves `b`, leaving the same budget. What does move it is `n_slots / n_attn` -- a model with 64 attention layers wants about a quarter of it -- and the compute and link of the machine. 128 MiB is a compromise across that spread with a little margin over this configuration's optimum, and `--kv-pipeline-budget` is there for the configurations it does not suit. + ### Where the rest of the token goes Per decode graph, `GGML_SCHED_TRANSPORT_DEBUG=2`, behind a 19,246-token prompt: From 34fbad96ff1286abd800d837299f7d9484d26924 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 00:30:37 +0200 Subject: [PATCH 12/32] sched: deliver a multi-stream KV window one stream at a time A window over several streams is one view of a tensor whose streams sit end to end, and both the prefix and the delivery treated it as one flat byte range. That made the lowest-writing stream cap the stable prefix for every stream above it, and it copied the cells between one stream's window and the next, which the graph never reads. The prefix is now counted within a stream, and a staged input whose last dimension indexes streams is delivered as one range per stream. A window over one stream keeps the single flat range it had, so single-sequence timing and bytes are unchanged. Behind 8 slots of a non-unified cache this takes decode from 72.18 to 112.70 t/s, against 70.61 ordered. At one slot it measures 35.31 against 35.32 before. test_transport_multi_stream_ranges pins the delivery: every stream's window covered once from its own source offset, the unread cells between them never moved, and the early and late bytes split as the prefix says. Concurrent slots cannot be gated on output the way one sequence can - their batching varies between runs, so the same depth gives different greedy output - which is why this is a unit test. Assisted-by: Claude Opus 5 --- ggml/include/ggml.h | 3 +- ggml/src/ggml-backend.cpp | 71 +++++++++++++++++++++++++----------- src/llama-kv-cache.cpp | 6 +-- tests/test-alloc.cpp | 77 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 132 insertions(+), 25 deletions(-) diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index a76a3bfd5f9..e23e3e6cbbc 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -702,7 +702,8 @@ extern "C" { void * extra; // extra things e.g. for ggml-cuda.cu - // leading bytes that stay unchanged during the current graph evaluation + // leading bytes of each stream that stay unchanged during the current graph evaluation + // a tensor whose last dimension indexes streams repeats this prefix once per stream, so the value is per stream and not from the start of the tensor union { size_t stable_prefix; char padding[8]; diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 8e75d909236..c75a68be449 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1726,24 +1726,44 @@ static bool ggml_backend_sched_transport_enabled(ggml_backend_sched_t sched) { return false; } +// How a staged input's delivery breaks into ranges. +// A window over one stream is one range and is delivered flat, exactly as ggml_nbytes(input) describes it. +// A window over several streams is one range per stream: the streams sit a fixed stride apart in both the source and the copy, which carries the source's layout, and the cells between one stream's window and the next are never read by this graph. +struct ggml_backend_sched_ranges { + int64_t n; // ranges to deliver + size_t stride; // bytes from one range to the next, in the source and in the copy alike + size_t used; // bytes of a range this graph reads + size_t early; // leading bytes of a range that may go before the split that reads it +}; + // The annotation lives on the tensor that owns the storage, and 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) { +// The prefix is per stream, so a range can use it only when the view starts on a stream boundary; anything else keeps the ordered path rather than guessing where the streams fall. +static void ggml_backend_sched_input_ranges(const struct ggml_tensor * input, struct ggml_backend_sched_ranges * out) { const struct ggml_tensor * base = input->view_src ? input->view_src : input; + + out->n = 1; + out->stride = 0; + out->used = ggml_nbytes(input); + out->early = 0; + if (base->stable_prefix == 0) { - return 0; + return; } + // dimensions below the stream have to cover their rows without a gap for a range to be a byte range + const size_t rows = (size_t) input->ne[2]*input->nb[2]; const size_t offs = input->view_src ? input->view_offs : 0; - if (base->stable_prefix <= offs) { - return 0; + if (input->nb[3] < rows || (offs != 0 && (input->nb[3] == 0 || offs % input->nb[3] != 0))) { + return; } - const size_t avail = base->stable_prefix - offs; - const size_t bytes = ggml_nbytes(input); + if (input->ne[3] > 1) { + out->n = input->ne[3]; + out->stride = input->nb[3]; + out->used = rows; + } - return avail < bytes ? avail : bytes; + out->early = base->stable_prefix < out->used ? base->stable_prefix : out->used; } // Whether a split input belongs in its backend's ring. @@ -2291,8 +2311,10 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { 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_ranges rg; + ggml_backend_sched_input_ranges(input, &rg); + total += rg.used*rg.n; + early += rg.early*rg.n; } } GGML_LOG_INFO("%s: %s: %d/%d splits staged, %zu KiB per graph, %zu KiB of it early, source %s\n", @@ -2343,19 +2365,23 @@ static void ggml_backend_sched_transport_prefetch(ggml_backend_sched_t sched, in // 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) { + struct ggml_backend_sched_ranges rg; + ggml_backend_sched_input_ranges(input, &rg); + if (rg.early == 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(r->transfer, input_cpy, input->data, 0, prefix); + for (int64_t r_i = 0; r_i < rg.n; r_i++) { + const size_t at = r_i*rg.stride; + ggml_backend_tensor_set_async(r->transfer, input_cpy, (const char *) input->data + at, at, rg.early); + } if (tr->debug >= 2) { tr->t_issue_us += ggml_time_us() - t0; } - tr->n_bytes_early += prefix; + tr->n_bytes_early += rg.early*rg.n; } // 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 @@ -2500,12 +2526,15 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s // 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 - 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; + struct ggml_backend_sched_ranges rg; + ggml_backend_sched_input_ranges(input, &rg); + if (rg.used > rg.early) { + for (int64_t r_i = 0; r_i < rg.n; r_i++) { + const size_t at = r_i*rg.stride + rg.early; + ggml_backend_tensor_set_async(split_backend, input_cpy, + (const char *) input->data + at, at, rg.used - rg.early); + } + tr->n_bytes_late += (rg.used - rg.early)*rg.n; } continue; } diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 676789fc5ca..eacbe7d37fe 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1649,13 +1649,13 @@ ggml_tensor * llama_kv_cache::build_input_v_rot(ggml_context * ctx) const { } 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. + // Rows written by this ubatch, counted within a stream rather than across the [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. + // Streams sit end to end, so a row counted across the body would let the lowest stream cap every stream above it; per stream, each one keeps its own leading rows. 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); + min_row = std::min(min_row, (uint64_t) idx); } } diff --git a/tests/test-alloc.cpp b/tests/test-alloc.cpp index 3f7ec16265e..6c3dcdbe251 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -1464,6 +1464,82 @@ static void test_transport_entry_allocation() { } } +// A window over several streams sits a fixed stride apart in one tensor, with cells between one +// stream's window and the next that the graph never reads. The delivery has to cover each stream's +// window from its own offset and leave those cells alone. +static void test_transport_multi_stream_ranges() { + dummy_backend cuda = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); + dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); + + const int64_t n_stream = 4; + const int64_t n_row = 8; // rows of a stream the graph reads + const int64_t kv_size = 12; // rows a stream holds, so 4 rows of every stream stay unread + const int64_t n_embd = 4; + + auto ctx = make_context(); + ggml_tensor * store = ggml_new_tensor_3d(ctx.ctx, GGML_TYPE_F32, n_embd, kv_size, n_stream); + store->flags |= GGML_TENSOR_FLAG_TRANSPORT; + + // the same shape a host-resident KV window has: heads split across dims 0 and 1, rows on 2, streams on 3 + ggml_tensor * window = ggml_view_4d(ctx.ctx, store, n_embd/2, 2, n_row, n_stream, + (size_t) (n_embd/2)*sizeof(float), store->nb[1], store->nb[2], 0); + ggml_tensor * output = ggml_cont(ctx.ctx, window); + ggml_build_forward_expand(ctx.graph, output); + + ggml_backend_buffer_ptr buffer(ggml_backend_buft_alloc_buffer(&cpu.buffer_type, ggml_nbytes(store))); + store->buffer = buffer.get(); + store->data = ggml_backend_buffer_get_base(buffer.get()); + + ggml_backend_t backends[] = { cuda.handle.get(), cpu.handle.get() }; + ggml_backend_buffer_type_t bufts[] = { &cuda.buffer_type, &cpu.buffer_type }; + ggml_backend_sched_ptr sched(ggml_backend_sched_new(backends, bufts, 2, 128, false, false)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_budget(sched.get(), 1u << 20)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_depth(sched.get(), 1)); + ggml_backend_sched_set_tensor_backend(sched.get(), output, cuda.handle.get()); + + // half of every stream's window is stable, so each stream splits into an early and a late range + const size_t row_bytes = (size_t) n_embd*sizeof(float); + const size_t used_bytes = (size_t) n_row*row_bytes; + ggml_set_stable_prefix(store, used_bytes/2); + + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), ctx.graph)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), ctx.graph) == GGML_STATUS_SUCCESS); + + std::vector parts = cuda.context->deliveries; + GGML_ASSERT(!parts.empty()); + std::sort(parts.begin(), parts.end(), + [](const dummy_backend_context::tensor_delivery & a, const dummy_backend_context::tensor_delivery & b) { + return a.offset < b.offset; + }); + + // every stream's window is covered exactly once, from its own source offset, and the unread cells never move + const size_t stride = (size_t) window->nb[3]; + size_t total = 0; + for (const auto & d : parts) { + GGML_ASSERT(d.offset/stride < (size_t) n_stream); + GGML_ASSERT(d.offset%stride + d.size <= used_bytes); + GGML_ASSERT(d.src == (const char *) window->data + d.offset); + total += d.size; + } + GGML_ASSERT(total == (size_t) n_stream*used_bytes); + + for (int64_t st = 0; st < n_stream; st++) { + size_t covered = 0; + for (const auto & d : parts) { + if (d.offset/stride == (size_t) st) { + GGML_ASSERT(d.offset%stride == covered); + covered += d.size; + } + } + GGML_ASSERT(covered == used_bytes); + } + + int64_t early = 0, late = 0; + transport_stats(sched.get(), NULL, &early, &late); + GGML_ASSERT(early == (int64_t) ((size_t) n_stream*used_bytes/2)); + GGML_ASSERT(late == (int64_t) ((size_t) n_stream*used_bytes/2)); +} + static void test_transport_empty_graph() { dummy_backend cuda = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); @@ -1844,6 +1920,7 @@ int main() { run("test_resizable_buffers_owner_borrower_teardown_order", test_resizable_buffers_owner_borrower_teardown_order); run("test_transport_prefix_and_configuration", test_transport_prefix_and_configuration); run("test_transport_entry_allocation", test_transport_entry_allocation); + run("test_transport_multi_stream_ranges", test_transport_multi_stream_ranges); run("test_transport_empty_graph", test_transport_empty_graph); run("test_transport_fallback_keeps_allocator_plan", test_transport_fallback_keeps_allocator_plan); run("test_transport_releases_ring_for_graph", test_transport_releases_ring_for_graph); From f5bc0222c30d3b83bbbc8717a8a6bd99d75f23ae Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 01:06:22 +0200 Subject: [PATCH 13/32] docs: measure the quant invariance, and what parallel sequences cost The q4_0 arm of the context sweep moves the peak from 16,384 rows to 32,768 and leaves the ring at it at 108 MiB against 102, which is what "the bytes per row cancel" claims. It was derived before and is measured now. Add the parallel numbers and say what the 8-slot row depends on: run on its own the unified ring fits, and only after a sweep has allocated for 1, 2 and 4 slots in the same process does the headroom guard refuse it. Assisted-by: Claude Opus 5 --- docs/kv-transport-pipelining.md | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index 6a9aac48798..9f2c26a1905 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -164,6 +164,23 @@ At 262,144 the ring would need 1.7 GiB against 573 MiB free, so it declines and Four device-resident layers are worth +14.3% on the ordered path and +1.4% on the pipelined one. The reason they stop paying is the point of the section above: the pipeline has already moved the bottleneck down to the compute floor, so removing a quarter of the traffic removes something that was no longer being waited for. Whether this still holds where the copy dominates by a wide margin has not been measured. +### Parallel sequences + +`llama-batched-bench`, 2,048 prompt tokens per sequence, `-c 32768 -np 8`, generation t/s: + +| `-npl` | unified ordered | unified pipelined | streams ordered | streams pipelined | +|---:|---:|---:|---:|---:| +| 1 | 32.69 | 35.35 | 32.59 | 35.31 | +| 2 | 51.69 | 58.63 | 48.68 | 61.56 | +| 4 | 72.43 | 86.52 | 62.39 | 89.59 | +| 8 | 83.90 | 105.16 | 70.92 | 113.55 | + +The 8-slot row is measured with each arm run on its own. Taken as the last step of a sweep that has already run 1, 2 and 4 slots in the same process, the unified ring is refused for headroom and that arm falls back to 83.88 - the guard reacting to what the process has already allocated, not to the configuration. + +A cache split into streams delivers a window per stream, so it moves more than a unified one for the same work, and before the per-stream delivery it could send almost none of it early: 6.6% at 8 slots, because the prefix stopped at the lowest stream's head. Both caches now pipeline, and which of them to use is a question about how the context is shared between sequences rather than about the transport. + +**Concurrent slots cannot be gated on output the way one sequence can.** Their batching varies between runs, so the same build at the same depth gives different greedy output - three runs at `N = 0` produced three different hashes. `test_transport_multi_stream_ranges` stands in for that gate. + ### 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: @@ -190,6 +207,20 @@ ring* = n_slots * b * rows* = (n_slots / n_attn) * compute * BW **`b` cancels.** The ring size at the peak does not depend on the cache type or on `n_embd_k_gqa`; it is set by the share of the traffic the ring holds, the compute a graph has to hide behind it, and the link. On this configuration -- 3 slots against 16 staged attention layers, 25.7 ms of consumer wait, 22.0 GB/s -- that is `0.1875 * 25.7e-3 * 22e9`, or **101 MiB against the 102 MiB measured at the +55.8% peak**. +The same sweep at `-ctk q4_0 -ctv q4_0` measures that cancellation rather than deriving it. Halving the bytes per row moves the whole curve to twice the window, and leaves the ring at the peak where it was: + +| q8_0 rows | gain | q4_0 rows | gain | +|---:|---:|---:|---:| +| 2,048 | +8.1% | 4,096 | +8.9% | +| 4,096 | +16.0% | 8,192 | +16.4% | +| 8,192 | +30.8% | 16,384 | +30.9% | +| 12,288 | +44.1% | 24,576 | +43.2% | +| **16,384** | **+56.0%** | **32,768** | **+53.9%** | +| 24,576 | +24.0% | 49,152 | +26.4% | +| 32,768 | +20.1% | 65,536 | +22.4% | + +The peak moves from 16,384 rows to 32,768, and the ring at it is 102 MiB against 108 MiB. + So the default is a size, not a depth, and it is the right kind of quantity to fix: a cache quantised to q4_0 doubles `rows*` and halves `b`, leaving the same budget. What does move it is `n_slots / n_attn` -- a model with 64 attention layers wants about a quarter of it -- and the compute and link of the machine. 128 MiB is a compromise across that spread with a little margin over this configuration's optimum, and `--kv-pipeline-budget` is there for the configurations it does not suit. ### Where the rest of the token goes @@ -235,6 +266,8 @@ A device-resident KV run is unaffected, and was measured to confirm it: 38.5612 - Only persistent host inputs marked with `GGML_TENSOR_FLAG_TRANSPORT` are candidates. The stable prefix remains a per-evaluation value. Unmarked inputs, weights, user inputs, transposed V, and copies with later readers stay ordered. - CUDA is the only enabled backend. Meta, SYCL, WebGPU, and other backends stay ordered until their event behavior and transport path are validated. - 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. +- **A cap is per graph, not per sequence.** `--kv-pipeline-budget` bounds the window one graph delivers, which is `n_kv * n_stream` over every sequence in the ubatch, so it cannot be applied to one sequence of a batch and not another. +- **A multi-stream window is delivered one range per stream**, keyed on the last dimension. A window whose streams are not on that dimension keeps the single flat range, which is correct but not accelerated. - **One ring per accelerator.** A layer-split model pipelines on every device that qualifies; a device with no room within the budget falls back to the ordered path on its own without disabling the others. - **Tensor parallelism keeps the ordered path.** See [Tensor parallelism](#tensor-parallelism). - 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. From 1cc279955c23871a8798962bc9bde83cb7b50e37 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 21:31:11 +0200 Subject: [PATCH 14/32] sched: take a staged window's stream span from the whole tensor ne[2]*nb[2] is one KV cell, not the window: attention permutes the window before reading it, so its rows sit on dimension 1. A ubatch over several streams delivered one cell per stream and attention read whatever the ring slot held before. The test built the window in the pre-permute shape, so it passed. Assisted-by: Claude Opus 5 --- ggml/src/ggml-backend.cpp | 5 +++-- tests/test-alloc.cpp | 2 ++ 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index c75a68be449..421fdf738d1 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1750,8 +1750,9 @@ static void ggml_backend_sched_input_ranges(const struct ggml_tensor * input, st return; } - // dimensions below the stream have to cover their rows without a gap for a range to be a byte range - const size_t rows = (size_t) input->ne[2]*input->nb[2]; + // a range is one stream's byte span, which is what the tensor covers below dimension 3 + // taking that span from ggml_nbytes keeps it right whatever order the dimensions below the stream are permuted into: attention reads a KV window with its rows on dimension 1 + const size_t rows = ggml_nbytes(input) - (size_t) (input->ne[3] - 1)*input->nb[3]; const size_t offs = input->view_src ? input->view_offs : 0; if (input->nb[3] < rows || (offs != 0 && (input->nb[3] == 0 || offs % input->nb[3] != 0))) { return; diff --git a/tests/test-alloc.cpp b/tests/test-alloc.cpp index 6c3dcdbe251..88f19990c5a 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -1483,6 +1483,8 @@ static void test_transport_multi_stream_ranges() { // the same shape a host-resident KV window has: heads split across dims 0 and 1, rows on 2, streams on 3 ggml_tensor * window = ggml_view_4d(ctx.ctx, store, n_embd/2, 2, n_row, n_stream, (size_t) (n_embd/2)*sizeof(float), store->nb[1], store->nb[2], 0); + // attention permutes it before reading it, which puts the rows on 1 and the heads on 2 + window = ggml_permute(ctx.ctx, window, 0, 2, 1, 3); ggml_tensor * output = ggml_cont(ctx.ctx, window); ggml_build_forward_expand(ctx.graph, output); From 3ee2430b3e4b9ed93c152883368257a08a4be33a Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 21:31:11 +0200 Subject: [PATCH 15/32] sched: wait for the previous graph before staging a window again A staged delivery reads its host source after the call that issued it returns. The stable prefix keeps the host off that source within a graph, but the next graph writes wherever its own ubatch lands, so a recycled cell below the previous window could be rewritten under an in-flight copy. The ordered path gets this from its blocking copy, once per split. Assisted-by: Claude Opus 5 --- ggml/src/ggml-backend.cpp | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 421fdf738d1..4e032c43610 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -2483,6 +2483,17 @@ 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; + // A staged delivery reads its host source long after the call that issued it returned, and the previous graph can leave some of those reads in flight. + // Within a graph the stable prefix keeps the host off what is still being read, but the next graph writes wherever its own ubatch lands, so wait for the previous one here. + // This is what the ordered path gets from its blocking copy, once per graph rather than once per split. + for (int i = 0; i < sched->n_backends; i++) { + if (tr->rings[i].transfer == NULL) { + continue; + } + ggml_backend_synchronize(tr->rings[i].transfer); + ggml_backend_synchronize(sched->backends[i]); + } + // Prime every ring before the first consumer runs. // From here on deliveries are issued only after a split has been enqueued, never before, so recycling a slot can never hold back work the consumer could already be running. // The cursors start over on every evaluation because the plan outlives the graph it was made for. From 7e1dee7532357e44494a0a806897a3300632942a Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 21:31:11 +0200 Subject: [PATCH 16/32] sched: wait for the consumer before freeing the transport ring Freeing the ring already goes through the backend that allocated it, so that backend is alive here and its kernels may still be reading the slots. Teardown skipped the wait and released the memory under them. Assisted-by: Claude Opus 5 --- ggml/src/ggml-backend.cpp | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 4e032c43610..07150550ccf 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1867,9 +1867,8 @@ static void ggml_backend_sched_transport_assign_addresses(ggml_backend_sched_t 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, int backend_id, bool sync_consumers) { +// Freeing the ring goes through the backend that allocated it, so the consumer is alive here and has to be waited for: kernels of an async compute may still be reading the slots. +static void ggml_backend_sched_transport_free_ring(ggml_backend_sched_t sched, int backend_id) { struct ggml_backend_sched_transport_ring * r = &sched->transport.rings[backend_id]; if (r->buffer == NULL) { @@ -1880,9 +1879,7 @@ static void ggml_backend_sched_transport_free_ring(ggml_backend_sched_t sched, i if (r->transfer) { ggml_backend_synchronize(r->transfer); } - if (sync_consumers) { - ggml_backend_synchronize(sched->backends[backend_id]); - } + ggml_backend_synchronize(sched->backends[backend_id]); ggml_backend_buffer_free(r->buffer); r->buffer = NULL; @@ -1893,10 +1890,10 @@ static void ggml_backend_sched_transport_free_ring(ggml_backend_sched_t sched, i } } -static void ggml_backend_sched_transport_release_ring(ggml_backend_sched_t sched, int backend_id, bool sync_consumers) { +static void ggml_backend_sched_transport_release_ring(ggml_backend_sched_t sched, int backend_id) { struct ggml_backend_sched_transport_ring * r = &sched->transport.rings[backend_id]; - ggml_backend_sched_transport_free_ring(sched, backend_id, sync_consumers); + ggml_backend_sched_transport_free_ring(sched, backend_id); for (int i = 0; i < GGML_SCHED_MAX_TRANSPORT_SLOTS; i++) { ggml_backend_event_free(r->slots[i].ready); @@ -1917,7 +1914,7 @@ static void ggml_backend_sched_transport_release_ring(ggml_backend_sched_t sched static void ggml_backend_sched_transport_decline_backend(ggml_backend_sched_t sched, int backend_id) { struct ggml_backend_sched_transport * tr = &sched->transport; - ggml_backend_sched_transport_release_ring(sched, backend_id, true); + ggml_backend_sched_transport_release_ring(sched, backend_id); if (tr->split_order == NULL || tr->split_input_ofs == NULL || tr->input_staged == NULL) { return; @@ -2245,7 +2242,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { } if (r->buffer == NULL || r->slot_size < slot_size[bid]) { - ggml_backend_sched_transport_free_ring(sched, bid, true); + ggml_backend_sched_transport_free_ring(sched, bid); ggml_backend_buffer_type_t buft = sched->bufts[bid]; @@ -2923,7 +2920,7 @@ ggml_backend_sched_t ggml_backend_sched_new( static void ggml_backend_sched_transport_teardown(ggml_backend_sched_t sched) { for (int i = 0; i < sched->n_backends; i++) { - ggml_backend_sched_transport_release_ring(sched, i, false); + ggml_backend_sched_transport_release_ring(sched, i); } sched->transport.n_staged = 0; } From 72ca3193f05c22c1408e26574f6da19cc05a3515 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 21:31:11 +0200 Subject: [PATCH 17/32] llama: note that a cache sharing cells keeps no stable prefix Assisted-by: Claude Opus 5 --- src/llama-kv-cache.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index eacbe7d37fe..e79aa206e9e 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1209,6 +1209,7 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & ubatch) { // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] + // a cache that shares cells keeps no stable prefix of its own: the layers it aliases get one from the cache that owns the cells, and the layers it does not stay on the ordered path if (other) { return; } From 6c013901ac579e59f1e11015f6d3172212749092 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 21:31:11 +0200 Subject: [PATCH 18/32] llama-bench: stop on an out-of-range -kvpd or -kvpb The range check left only the inner loop and inserted the values anyway. Assisted-by: Claude Opus 5 --- tools/llama-bench/llama-bench.cpp | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 9dc09d930a9..f3ecf991e1e 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -872,6 +872,9 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { break; } } + if (invalid_param) { + break; + } params.kv_pipeline_depth.insert(params.kv_pipeline_depth.end(), p.begin(), p.end()); } else if (arg == "-kvpb" || arg == "--kv-pipeline-budget") { if (++i >= argc) { @@ -885,6 +888,9 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { break; } } + if (invalid_param) { + break; + } params.kv_pipeline_budget_mib.insert(params.kv_pipeline_budget_mib.end(), p.begin(), p.end()); } else if (arg == "-rso" || arg == "--recurrent-state-offload") { if (++i >= argc) { From 73bf9c8c726757f98fe0f953322a5ce29b9613f6 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 22:19:49 +0200 Subject: [PATCH 19/32] sched: wait on the slot release events, not on the consumer backend The scheduler does not own its backends, and llama_context declares its scheduler before them, so member destruction frees the backends first and sched->backends[] dangles by the time the ring is freed. A slot's release event is recorded past every kernel that reads it and dispatches through the device, so waiting on it orders the free after the consumer without touching the backend. Assisted-by: Claude Opus 5 --- ggml/src/ggml-backend.cpp | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 07150550ccf..585ac7bd3b6 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1867,7 +1867,8 @@ static void ggml_backend_sched_transport_assign_addresses(ggml_backend_sched_t s } } -// Freeing the ring goes through the backend that allocated it, so the consumer is alive here and has to be waited for: kernels of an async compute may still be reading the slots. +// The consumer is waited for through the slots' own release events, never through sched->backends[backend_id]. +// The scheduler does not own its backends and llama_context declares its scheduler before them, so on the teardown path they are already gone; the buffer and the events go through the buffer type and the device, which are not. static void ggml_backend_sched_transport_free_ring(ggml_backend_sched_t sched, int backend_id) { struct ggml_backend_sched_transport_ring * r = &sched->transport.rings[backend_id]; @@ -1879,7 +1880,12 @@ static void ggml_backend_sched_transport_free_ring(ggml_backend_sched_t sched, i if (r->transfer) { ggml_backend_synchronize(r->transfer); } - ggml_backend_synchronize(sched->backends[backend_id]); + // release is recorded past every kernel that reads the slot, so reaching it means those kernels are done + for (int i = 0; i < GGML_SCHED_MAX_TRANSPORT_SLOTS; i++) { + if (r->slots[i].release_armed && r->slots[i].release) { + ggml_backend_event_synchronize(r->slots[i].release); + } + } ggml_backend_buffer_free(r->buffer); r->buffer = NULL; From f70f221a71a05de5c4251d1b9472b12806669fb6 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 22:50:57 +0200 Subject: [PATCH 20/32] sched: deliver a staged window with ggml_backend_tensor_set_2d_async The per-stream loop in both delivery paths is what that helper does, and it lets a backend with a 2d set issue one copy instead of one per stream. Assisted-by: Claude Opus 5 --- ggml/src/ggml-backend.cpp | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 585ac7bd3b6..0db5a9e31f1 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -2378,10 +2378,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; - for (int64_t r_i = 0; r_i < rg.n; r_i++) { - const size_t at = r_i*rg.stride; - ggml_backend_tensor_set_async(r->transfer, input_cpy, (const char *) input->data + at, at, rg.early); - } + ggml_backend_tensor_set_2d_async(r->transfer, input_cpy, input->data, 0, rg.early, rg.n, rg.stride, rg.stride); if (tr->debug >= 2) { tr->t_issue_us += ggml_time_us() - t0; } @@ -2544,11 +2541,8 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s struct ggml_backend_sched_ranges rg; ggml_backend_sched_input_ranges(input, &rg); if (rg.used > rg.early) { - for (int64_t r_i = 0; r_i < rg.n; r_i++) { - const size_t at = r_i*rg.stride + rg.early; - ggml_backend_tensor_set_async(split_backend, input_cpy, - (const char *) input->data + at, at, rg.used - rg.early); - } + ggml_backend_tensor_set_2d_async(split_backend, input_cpy, (const char *) input->data + rg.early, + rg.early, rg.used - rg.early, rg.n, rg.stride, rg.stride); tr->n_bytes_late += (rg.used - rg.early)*rg.n; } continue; From 537c0e004e1dc68d54681de37e12ee86e21a213f Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 22:50:57 +0200 Subject: [PATCH 21/32] ggml: say what a stable prefix covers The count is per stream, so the contract has to name the stride it goes with rather than leave it as "the first nbytes". Assisted-by: Claude Opus 5 --- ggml/include/ggml.h | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index e23e3e6cbbc..d617b201cdf 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -702,8 +702,8 @@ extern "C" { void * extra; // extra things e.g. for ggml-cuda.cu - // leading bytes of each stream that stay unchanged during the current graph evaluation - // a tensor whose last dimension indexes streams repeats this prefix once per stream, so the value is per stream and not from the start of the tensor + // bytes at the start of every stream that stay unchanged for the current graph evaluation, 0 for none + // the count is from the start of a stream, not from the start of the tensor, so a reader that splits the storage into streams applies it to each of them union { size_t stable_prefix; char padding[8]; @@ -712,9 +712,11 @@ extern "C" { 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 + // declare that the first nbytes of every stream of tensor->data cannot change while a graph that reads this tensor is being evaluated + // a stream is one index of the last dimension of the view a reader takes of this tensor, and consecutive streams sit that view's nb[3] apart, so a reader that takes the storage whole has a single stream and the region is then simply its first nbytes // nbytes is clamped to ggml_nbytes(tensor), and it must be set on the tensor that owns the storage, not on a view of it // it must describe the graph that is about to run, including when that graph is reused + // a backend may deliver a declared region before the point in the graph that reads it, so 0 declares nothing and is always correct 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); From 2c0a6a2a0a6aa297ea2cf1d47a42e121a03a1a0e Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 22:50:57 +0200 Subject: [PATCH 22/32] repro: gate the multi-stream delivery on concurrent output The server's batching varies between runs, so it cannot gate concurrent sequences. llama-parallel seeds its client schedule, so it can, and its clients ask different questions: with one shared prompt every stream holds the same bytes and a cross-stream read stays invisible. Fails at depth 1 on the commit before the multi-stream span fix. Assisted-by: Claude Opus 5 --- docs/repro/r4-kv-pipeline-parallel-exact.sh | 79 +++++++++++++++++++++ 1 file changed, 79 insertions(+) create mode 100755 docs/repro/r4-kv-pipeline-parallel-exact.sh diff --git a/docs/repro/r4-kv-pipeline-parallel-exact.sh b/docs/repro/r4-kv-pipeline-parallel-exact.sh new file mode 100755 index 00000000000..779b0dae8b9 --- /dev/null +++ b/docs/repro/r4-kv-pipeline-parallel-exact.sh @@ -0,0 +1,79 @@ +#!/bin/bash +# R4 gate 2: greedy output of concurrent sequences over a cache split into streams must be byte-identical to the ordered path. +# Gate 1 covers one sequence per ubatch, where a delivery is a single range. This covers a ubatch that spans several streams, where a delivery is one range per stream and the cells between them are never read. +# llama-parallel is used rather than the server because it seeds its client schedule, so the batches are the same run to run and the outputs can be compared directly. +# +# LLAMA_KV_MODEL=/path/model.gguf docs/repro/r4-kv-pipeline-parallel-exact.sh [pipeline-depth ...] +# LLAMA_KV_NP=8 sets the concurrent sequences, LLAMA_KV_NS=16 the total. +# LLAMA_KV_SM=layer spreads the model over every device, which gives each of them its own ring. +set -euo pipefail +MODEL="${LLAMA_KV_MODEL:?set LLAMA_KV_MODEL to a .gguf path}" +BUILD="${LLAMA_KV_BUILD:-build}" +PIN="${LLAMA_KV_TASKSET:-0,2,4}" +CTX="${LLAMA_KV_CTX:-16384}" +BUDGET="${LLAMA_KV_BUDGET:-512}" +NP="${LLAMA_KV_NP:-8}" +NS="${LLAMA_KV_NS:-16}" +SM="${LLAMA_KV_SM:-none}" + +DEPTHS=(0 1 4); [ $# -gt 0 ] && DEPTHS=("$@") + +# Without a pinned host cache, a host-resident recurrent state and a budget the run does not exercise the path the doc reports on. +if ! HELP="$("$BUILD/bin/llama-parallel" --help 2>&1)"; then + echo "cannot run $BUILD/bin/llama-parallel:" >&2 + echo "$HELP" >&2 + exit 1 +fi +for opt in --kv-cpu-pinned --recurrent-state-offload --kv-pipeline-depth --kv-pipeline-budget --no-kv-unified; do + if ! grep -q -- "$opt" <<< "$HELP"; then + echo "$BUILD/bin/llama-parallel has no $opt option" >&2 + exit 1 + fi +done + +# Keep only what the clients produced: drop the log timestamps, the colour codes and the timing summary, all of which differ between runs by design. +transcript () { + sed -e 's/\x1b\[[0-9;]*m//g' \ + -e '/^[0-9][0-9.]* [A-Z] /d' \ + -e '/speed:/d' -e '/^Cache misses/d' "$1" +} + +rc=0 +BASE="" +for I in "${!DEPTHS[@]}"; do + D="${DEPTHS[$I]}" + echo "== pipeline depth=$D ctx=$CTX np=$NP ns=$NS sm=$SM streams (non-unified)" + RAW="$(mktemp /tmp/r4-kv-parallel.XXXX.log)" + OUT="$(mktemp /tmp/r4-kv-parallel.XXXX.txt)" + arm_rc=0 + taskset -c "$PIN" "$BUILD/bin/llama-parallel" -m "$MODEL" \ + --kv-pipeline-depth "$D" --kv-pipeline-budget "$BUDGET" \ + -ngl 99 -sm "$SM" -mg 0 -t 3 -nkvo --kv-cpu-pinned --recurrent-state-offload \ + -fa on -ctk q8_0 -ctv q8_0 -b 512 -ub 512 -c "$CTX" -no-kvu \ + -np "$NP" -ns "$NS" --temp 0 > "$RAW" 2>&1 || arm_rc=$? + if [ "$arm_rc" -ne 0 ]; then + echo " depth $D: FAILED (exit $arm_rc)" >&2 + tail -20 "$RAW" >&2 + rm -f "$RAW" "$OUT" + exit "$arm_rc" + fi + transcript "$RAW" > "$OUT" + if [ ! -s "$OUT" ]; then + echo " depth $D: FAILED (no client output)" >&2 + rm -f "$RAW" "$OUT" + exit 1 + fi + echo " $(sha256sum < "$OUT" | cut -c1-16) $(wc -l < "$OUT") lines" + if [ "$I" -eq 0 ]; then + BASE="$OUT" + else + if ! cmp -s "$BASE" "$OUT"; then + diff -u "$BASE" "$OUT" | head -40 + rc=1 + fi + rm -f "$OUT" + fi + rm -f "$RAW" +done +rm -f "$BASE" +exit $rc From e477aa4931a260e6042762ccd28c248f77e1ed1f Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 22:50:57 +0200 Subject: [PATCH 23/32] docs, tests: unwrap the hard-wrapped comments Assisted-by: Claude Opus 5 --- docs/repro/r4-kv-pipeline-ab.sh | 5 ++--- docs/repro/r4-kv-pipeline-context-sweep.sh | 10 ++++------ docs/repro/r4-kv-pipeline-exact.py | 17 ++++++----------- docs/repro/r4-kv-pipeline-exact.sh | 4 ++-- tests/test-alloc.cpp | 5 ++--- 5 files changed, 16 insertions(+), 25 deletions(-) diff --git a/docs/repro/r4-kv-pipeline-ab.sh b/docs/repro/r4-kv-pipeline-ab.sh index 563174e56f8..d0e50427ab7 100755 --- a/docs/repro/r4-kv-pipeline-ab.sh +++ b/docs/repro/r4-kv-pipeline-ab.sh @@ -10,9 +10,8 @@ PIN="${LLAMA_KV_TASKSET:-0,2,4}" BUDGET="${LLAMA_KV_BUDGET:-512}" LOCK=/tmp/beellama-single-gpu.lock -# An unpinned host cache and a host-resident recurrent state both cost more than the transport can -# win back, and without a budget the ring is declined at the larger contexts, so a build without -# these options does not measure what the doc reports. Fail rather than measure something else. +# An unpinned host cache and a host-resident recurrent state both cost more than the transport can win back, and without a budget the ring is declined at the larger contexts, so a build without these options does not measure what the doc reports. +# Fail rather than measure something else. if ! HELP="$("$BUILD/bin/llama-bench" --help 2>&1)"; then echo "cannot run $BUILD/bin/llama-bench:" >&2 echo "$HELP" >&2 diff --git a/docs/repro/r4-kv-pipeline-context-sweep.sh b/docs/repro/r4-kv-pipeline-context-sweep.sh index cf363dfa002..64bac9feba3 100755 --- a/docs/repro/r4-kv-pipeline-context-sweep.sh +++ b/docs/repro/r4-kv-pipeline-context-sweep.sh @@ -1,7 +1,6 @@ #!/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. +# 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 -euo pipefail @@ -12,9 +11,8 @@ NGEN="${LLAMA_KV_NGEN:-64}" BUDGET="${LLAMA_KV_BUDGET:-512}" LOCK=/tmp/beellama-single-gpu.lock -# An unpinned host cache and a host-resident recurrent state both cost more than the transport can -# win back, and without a budget the ring is declined at the larger contexts, so a build without -# these options does not measure what the doc reports. Fail rather than measure something else. +# An unpinned host cache and a host-resident recurrent state both cost more than the transport can win back, and without a budget the ring is declined at the larger contexts, so a build without these options does not measure what the doc reports. +# Fail rather than measure something else. if ! HELP="$("$BUILD/bin/llama-bench" --help 2>&1)"; then echo "cannot run $BUILD/bin/llama-bench:" >&2 echo "$HELP" >&2 diff --git a/docs/repro/r4-kv-pipeline-exact.py b/docs/repro/r4-kv-pipeline-exact.py index 55278ce3534..3330893ddfc 100644 --- a/docs/repro/r4-kv-pipeline-exact.py +++ b/docs/repro/r4-kv-pipeline-exact.py @@ -7,8 +7,7 @@ RESULTS_PATH = sys.argv[3] RESULTS = [] -# 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. +# 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 " @@ -43,17 +42,14 @@ def filler(name, target_tokens): return unit * reps def nonce(name, length): - # The server restores a cached prefix from an earlier task, and a restored window is not - # numerically the same as a freshly prefilled one, so two tasks that share a long prefix stop - # measuring the code under test. This makes every task's prefix unique, and it is derived from - # the task rather than drawn at random so that a control run produces comparable hashes. + # The server restores a cached prefix from an earlier task, and a restored window is not numerically the same as a freshly prefilled one, so two tasks that share a long prefix stop measuring the code under test. + # This makes every task's prefix unique, and it is derived from the task rather than drawn at random so that a control run produces comparable hashes. h = hashlib.sha256(f"{name}/{length}".encode()).hexdigest()[:32] return f"Session {h}. Ignore this line.\n\n" def ask(label, prompt, ntok, want_prefill): - # cache_prompt=False forces a full prefill. Without it a task inherits whatever the previous - # one left in the cache, and two tasks whose prompts do not both fit make placement depend on - # that: records@18432 then gives different answers across otherwise identical runs. + # cache_prompt=False forces a full prefill. + # Without it a task inherits whatever the previous one left in the cache, and two tasks whose prompts do not both fit make placement depend on that: records@18432 then gives different answers across otherwise identical runs. body = json.dumps({"model": "m", "messages": [{"role": "user", "content": prompt}], "max_tokens": ntok, "temperature": 0, "top_k": 1, "seed": 1234, "cache_prompt": False}).encode() @@ -69,8 +65,7 @@ def ask(label, prompt, ntok, want_prefill): # reasoning models put most of the generation in reasoning_content; hash both text = (m.get("reasoning_content") or "") + "\x00" + (m.get("content") or "") t = d.get("timings", {}) - # a reused prefix shows up as a prompt_n far below the prompt actually sent; the hash it - # produces is not comparable to a fresh prefill, so say so rather than reporting it silently + # a reused prefix shows up as a prompt_n far below the prompt actually sent; the hash it produces is not comparable to a fresh prefill, so say so rather than reporting it silently prompt_n = t.get("prompt_n") or 0 reused = prompt_n < want_prefill // 2 digest = hashlib.sha256(text.encode()).hexdigest()[:16] diff --git a/docs/repro/r4-kv-pipeline-exact.sh b/docs/repro/r4-kv-pipeline-exact.sh index d2795a5a7c1..017926d29a0 100755 --- a/docs/repro/r4-kv-pipeline-exact.sh +++ b/docs/repro/r4-kv-pipeline-exact.sh @@ -1,6 +1,6 @@ #!/bin/bash -# R4 gate 1: greedy server output must be byte-identical to the ordered path, across several -# prefill corpora and prefill lengths. The script compares every requested depth with the first. +# R4 gate 1: greedy server output must be byte-identical to the ordered path, across several prefill corpora and prefill lengths. +# The script compares every requested depth with the first. # # 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). diff --git a/tests/test-alloc.cpp b/tests/test-alloc.cpp index 88f19990c5a..c0f4b72d9e0 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -1464,9 +1464,8 @@ static void test_transport_entry_allocation() { } } -// A window over several streams sits a fixed stride apart in one tensor, with cells between one -// stream's window and the next that the graph never reads. The delivery has to cover each stream's -// window from its own offset and leave those cells alone. +// A window over several streams sits a fixed stride apart in one tensor, with cells between one stream's window and the next that the graph never reads. +// The delivery has to cover each stream's window from its own offset and leave those cells alone. static void test_transport_multi_stream_ranges() { dummy_backend cuda = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); From acf9fa29a5d102b4c472b38e56438fdbace61719 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 23:10:32 +0200 Subject: [PATCH 24/32] docs: re-measure on this head, and say which head each number is from The parallel table was taken before the multi-stream span fix, so it reported a delivery that moved a fraction of the window. Gates 1, 2 and 5 are re-run here, and the numbers that are still from an earlier head now say so. Assisted-by: Claude Opus 5 --- docs/kv-transport-pipelining.md | 49 +++++++++++++++++++-------------- 1 file changed, 28 insertions(+), 21 deletions(-) diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index 9f2c26a1905..3d4d0faa69c 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -24,6 +24,8 @@ The scheduler delivers `[0, stable_prefix)` early on the transfer stream and the 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. +It also says nothing about the *next* graph. A delivery is still reading the host cache after the call that issued it has returned, and the next ubatch writes wherever its own slots fall, which can be below the window the previous graph is still delivering. The scheduler therefore waits for the transfer stream and for the consumer once at the top of each evaluation, before any split of the new graph can write the cache. The ordered path gets the same guarantee from its blocking copy, which pays for it once per split rather than once per graph. + ### 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 (+1.38%, and not exact). @@ -62,9 +64,11 @@ RTX 4070 (11,902 MiB usable, sm_89), driver 610.57.04 / CUDA 13.3, i5-13400F, `Q | depth | ordered | pipelined | gain | |---:|---|---|---:| -| 4,096 | 29.9802, 29.9776 | 34.8016, 34.8047 | **+16.1%** | -| 16,384 | 19.0116, 19.0093 | 29.7780, 29.7879 | **+56.6%** | -| 32,768 | 12.7237, 12.7237 | 15.2873, 15.2864 | **+20.1%** | +| 4,096 | 29.9557, 29.9585 | 34.7899, 34.8016 | **+16.2%** | +| 16,384 | 18.9675, 18.9793 | 29.7606, 29.7697 | **+56.9%** | +| 32,768 | 12.7079, 12.7093 | 15.2760, 15.2630 | **+20.2%** | + +Re-measured on the current head as gate 2, and within 0.3% of the same table taken before the graph-boundary wait was added. The 32,768 ring is 204 MiB at the full context, over the 128 MiB default, so that row needs `-kvpb 512`. `llama-bench` takes the option and the scripts pass it. @@ -74,10 +78,12 @@ The 32,768 ring is 204 MiB at the full context, over the 128 MiB default, so tha | task | prompt | ordered | pipelined | gain | |---|---:|---:|---:|---:| -| prose | 14,821 | 19.775 | 30.012 | **+51.8%** | -| dialogue | 15,984 | 19.155 | 29.767 | **+55.4%** | -| records | 29,603 | 13.553 | 16.396 | **+21.0%** | -| code | 29,670 | 13.504 | 16.362 | **+21.2%** | +| prose | 14,821 | 19.738 | 29.955 | **+51.8%** | +| dialogue | 15,984 | 19.136 | 29.699 | **+55.2%** | +| records | 29,603 | 13.531 | 16.406 | **+21.2%** | +| code | 29,670 | 13.484 | 16.340 | **+21.2%** | + +Re-measured on the current head as gate 1, and within 0.3% of the same table on the commit before the graph-boundary wait was added, which is the change that could have cost it. The breakdown below was taken on an earlier head, at prompts of 19,246 and 48,042 against `-c 32768` and `-c 65536`: @@ -97,7 +103,7 @@ Pinning is worth as much as the pipeline and is off by default. Behind a 13,128- | `--kv-cpu-pinned` | 21.582 | 32.252 | | unpinned | 14.945 | 22.709 | -Look-ahead deeper than one split is worse at every depth measured. At 19,246: 29.83 t/s at `N = 1`, 28.44 at `N = 2`, 25.93 at `N = 4`. `N = 1` is the default for that reason. The exactness gate measures the same on every one of its four 18,432-prefill tasks: 30.012 against 27.037 on prose, 29.767 against 26.584 on dialogue, 16.396 against 15.934 on records, 16.362 against 15.857 on code. +Look-ahead deeper than one split is worse at every depth measured. At 19,246: 29.83 t/s at `N = 1`, 28.44 at `N = 2`, 25.93 at `N = 4`. `N = 1` is the default for that reason. The exactness gate measures the same on every one of its four 18,432-prefill tasks: 29.955 against 27.031 on prose, 29.699 against 26.626 on dialogue, 16.406 against 15.926 on records, 16.340 against 15.864 on code. ### The link is the ceiling, so the lever is bytes @@ -145,7 +151,7 @@ A finer sweep of the same configuration, `-kvpb 0` throughout so nothing decline | 98,304 | 5.45 | 6.06 | +11.1% | 612 MiB | 0.018 | | 131,072 | 4.26 | 4.66 | +9.5% | 816 MiB | 0.012 | -The gain per MiB is flat below the peak and falls off as `1/rows^2` above it, because the gain decays as `compute/copy` while the ring grows linearly. There is no depth at which the pipeline becomes slower -- only one past which the memory buys more elsewhere. +The gain per MiB is flat below the peak and falls off as `1/rows^2` above it, because the gain decays as `compute/copy` while the ring grows linearly. No depth measured here makes the pipeline slower -- only one past which the memory buys more elsewhere. That is one model on one link, not a general claim. Two curves run in opposite directions here, and both matter. @@ -166,20 +172,20 @@ Four device-resident layers are worth +14.3% on the ordered path and +1.4% on th ### Parallel sequences -`llama-batched-bench`, 2,048 prompt tokens per sequence, `-c 32768 -np 8`, generation t/s: +`llama-batched-bench`, 2,048 prompt tokens per sequence, `-c 32768 -np 8`, generation t/s. Every cell is its own process, because the headroom guard reacts to what a process has already allocated rather than to the configuration: taken as the last step of a sweep that has already run 1, 2 and 4 slots, the 8-slot unified ring is refused and that arm reads 83.88 instead. Three passes, spread at most 0.05 t/s: | `-npl` | unified ordered | unified pipelined | streams ordered | streams pipelined | |---:|---:|---:|---:|---:| -| 1 | 32.69 | 35.35 | 32.59 | 35.31 | -| 2 | 51.69 | 58.63 | 48.68 | 61.56 | -| 4 | 72.43 | 86.52 | 62.39 | 89.59 | -| 8 | 83.90 | 105.16 | 70.92 | 113.55 | +| 1 | 32.62 | 35.34 | 32.62 | 35.36 | +| 2 | 51.51 | 58.48 | 34.09 | 58.45 | +| 4 | 72.04 | 86.24 | 49.44 | 85.45 | +| 8 | 83.95 | 105.22 | 70.92 | 105.63 | -The 8-slot row is measured with each arm run on its own. Taken as the last step of a sweep that has already run 1, 2 and 4 slots in the same process, the unified ring is refused for headroom and that arm falls back to 83.88 - the guard reacting to what the process has already allocated, not to the configuration. +A cache split into streams delivers a window per stream, so it moves more than a unified one for the same work, and before the per-stream delivery it could send almost none of it early: 6.6% at 8 slots, because the prefix stopped at the lowest stream's head. Both caches now pipeline to the same throughput, and which of them to use is a question about how the context is shared between sequences rather than about the transport. The ordered arm is the one that separates them: a non-unified cache scales much worse without the pipeline, so the pipeline is worth more there. -A cache split into streams delivers a window per stream, so it moves more than a unified one for the same work, and before the per-stream delivery it could send almost none of it early: 6.6% at 8 slots, because the prefix stopped at the lowest stream's head. Both caches now pipeline, and which of them to use is a question about how the context is shared between sequences rather than about the transport. +**The streams-pipelined column is lower than it was before the multi-stream span was fixed**, and the earlier numbers were wrong rather than better. The delivery sized one stream's range from `ne[2]*nb[2]`, which is one KV cell rather than the window, so it moved a fraction of the bytes and the attention read whatever the ring slot held before. Measured on the same machine, the predecessor reports 61.04, 90.81 and 108.23 at 2, 4 and 8 slots against 58.45, 85.45 and 105.63 here; the difference is the cost of copying the right amount. -**Concurrent slots cannot be gated on output the way one sequence can.** Their batching varies between runs, so the same build at the same depth gives different greedy output - three runs at `N = 0` produced three different hashes. `test_transport_multi_stream_ranges` stands in for that gate. +**Concurrent slots can be gated on output, with a harness that fixes the batching.** The server cannot: its batching varies between runs, so the same build at the same depth gives different greedy output, and three runs at `N = 0` produced three different hashes. `llama-parallel` seeds its client schedule, so the batches repeat, and `docs/repro/r4-kv-pipeline-parallel-exact.sh` compares the transcripts of 8 concurrent sequences over a non-unified cache. Its clients ask different questions, which is what makes it a gate: with one prompt shared by every sequence the streams hold the same bytes and a cross-stream read is invisible. The predecessor above fails it at `N = 1` on the first sequence. ### The budget @@ -242,7 +248,7 @@ The 3.57 ms that remains moves 0.4 MiB, and `GGML_SCHED_TRANSPORT_DEBUG=3` shows It looks like latency and is not. A blocking copy shares the device's copy engine with the deliveries and waits for what is already queued there: two staged splits at 22.0 GB/s is 3.6 ms, which is the number. Two things were tried and neither helped. Issuing the delivery in pieces so the blocking copy can interleave does nothing -- the engine is FIFO across streams, `attn_inp_k_rot` stays at 3.4 ms at every piece size, and small pieces cost throughput (29.80 t/s whole, 28.73 at 4 MiB, 22.01 at 1 MiB). Putting the copy on the consumer's own stream so the host never blocks moves the time rather than removing it: the ordered copy falls from 3.57 ms to 0.16 ms, the consumer wait rises from 27.31 ms to 31.05 ms, and throughput does not move (29.808 against 29.834). -So this is not spare time. Those 256 KiB cross the same saturated link as the 644 MiB of deliveries, and the link is the ceiling. +So this is not spare time. Those 256 KiB cross the same saturated link as the 644 MiB of deliveries, and on this configuration the link is the ceiling. A faster link, or a slower device behind it, moves that ceiling somewhere else. Do not compare these numbers against runs on other models, prompts, cache settings, hardware, or commits. @@ -250,14 +256,15 @@ Do not compare these numbers against runs on other models, prompts, cache settin The gates, and what was run for them: -All four gates below were re-run on the current head, on an RTX 4070 with a CUDA build. +Gates 1, 2, 3 and 5 and `test-alloc` were run on the current head, on an RTX 4070 with a CUDA build, gate 5 also over both devices with `-sm layer`. The `llama-server` table and the parallel table under [Measurements](#measurements) are from those runs; the breakdowns marked as taken on an earlier head still are. -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` compares every requested depth with the first and fails on a hash difference. Two things keep the tasks independent of each other, and both were needed. Every task carries a nonce derived from its own name and length, so no two share a prefix the server could restore, and the harness fails a task whose `prompt_n` says one was reused anyway. Each request also sets `cache_prompt: false`, so a task never inherits what the previous one left in the cache. +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`, re-run on the current head. `docs/repro/r4-kv-pipeline-exact.sh` compares every requested depth with the first and fails on a hash difference. Two things keep the tasks independent of each other, and both were needed. Every task carries a nonce derived from its own name and length, so no two share a prefix the server could restore, and the harness fails a task whose `prompt_n` says one was reused anyway. Each request also sets `cache_prompt: false`, so a task never inherits what the previous one left in the cache. The second is what made `records@18432` a gate rather than a coin flip. Its prompt is about 29.6k tokens against a 32,768 context, and the task before it is about the same size, so the two do not both fit and placement depended on what was still resident. Two otherwise identical `N = 0` runs of it produced different hashes. Asked on its own with the cache off it is perfectly stable: the same hash three times running, at `-c 32768` and at `-c 65536`. With the flag set, two independent `N = 0` passes agree on all eight tasks, and `N = 0`, `N = 1` and `N = 4` agree on all eight. -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. +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, re-run on the current head. 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, as the mean over each 128 graphs, with depth stops and the number of recycle waits enqueued; `=3` names the tensors still on the ordered path. `ggml_backend_sched_get_transport_pipeline_stats()` exposes deliveries and early and late byte counts to callers. +5. **Byte-identical greedy output of concurrent sequences over a cache split into streams.** `docs/repro/r4-kv-pipeline-parallel-exact.sh` runs 8 concurrent sequences of 16 through `llama-parallel`, which seeds its client schedule so the batches repeat, and compares every depth's transcripts with the first. Identical at `N = 0`, `N = 1` and `N = 4` with `-sm none`, and at `N = 0` and `N = 1` with `-sm layer` over both devices, which is the case where each device carries its own ring. Gate 1 covers one sequence per ubatch, where a delivery is one range; this covers a ubatch spanning several streams, where it is one range per stream. Run against the commit before the multi-stream span fix it fails at `N = 1`, which is what makes it a gate rather than a smoke test. A device-resident KV run is unaffected, and was measured to confirm it: 38.5612 t/s at depth 0 against 38.5240 at depth 1, `tg128 @ d4096`, with the transport never enabled because the scheduler is given a depth of 0. From b27d940f402d0d539c243dd05170c1cd5f745c7a Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 6 Sep 2026 01:28:35 +0200 Subject: [PATCH 25/32] sched: release an idle transport ring, and grow a slot in powers of two A graph that stages nothing kept the ring and the second device context for the life of the scheduler. Give them back, from every path that ends with nothing staged. A window wider than the ring holds frees the ring and allocates it again, which a prefill did on nearly every ubatch. Allocate a slot in powers of two, capped by the full context, the budget and the headroom check, so a 16k prefill reallocates 6 times rather than 32. The decline decision still goes by what the graph needs. Wait at a graph boundary only for a ring that delivered, not for every ring that has a transfer backend. Assisted-by: Claude Opus 5 --- docs/kv-transport-pipelining.md | 5 ++- ggml/src/ggml-backend.cpp | 56 ++++++++++++++++++++++++++++++--- 2 files changed, 55 insertions(+), 6 deletions(-) diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index 3d4d0faa69c..9b17d572516 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -193,13 +193,16 @@ The table above is what the feature costs uncapped, and it is the reason it is c - Under the cap the ring is allocated and the deliveries pipeline. - Over it the scheduler declines and keeps the ordered path for that graph. Later graphs are evaluated again, so a smaller live window can use the ring. -- Declining costs nothing in steady state. Both the ring and the transfer backend's device context are released. +- Declining costs nothing in steady state. Both the ring and the transfer backend's device context are released, and a graph that stages nothing at all releases them the same way. +- The value is in MiB, `0` removes the cap, and 65536 is the largest accepted: a slot holds one attention layer's K or V, so a cap past that is a typo rather than a budget. - The budget and the headroom check are decided before the graph is allocated, so they can still leave the graph short. If graph reservation fails, the rings are released and the reservation is retried once on the ordered path, and that scheduler keeps the ordered path from then on. An optional ring never turns a graph that fits into an allocation failure. **The cap is applied to what the current graph needs, not to what the full context would need.** A run whose window stays small keeps the ring whatever `-n_ctx` says, which is the common case and the reason it is done this way: a staged input is a view of the cache tensor, so the full-context figure is there for the asking, but enforcing it would refuse the ring for every large `-c` even when the window never gets near it. The warning reports both numbers so that `--kv-pipeline-budget` can be sized against the one that matters. The cost of deciding per graph is that a context which grows past the budget allocates a ring for the small early windows and gives it back once it outgrows them. That transient is bounded by the budget itself, which is the memory the user already authorised, so it is a property of the cap rather than a defect in it. +A window wider than the ring holds has to free the ring and allocate it again, which blocks the host on the device, and a prefill widens the window on nearly every ubatch. So a slot is allocated in powers of two, up to what the full context needs and never past the budget or the headroom check. The decision to decline still goes by what the graph needs, so the cap falls where it did; only the allocation is coarse. + At 32,768 the ring is 204 MiB at the full context, over the 128 MiB default. `--kv-pipeline-budget 512` buys 20.350 -> 31.463 t/s behind an 18,432-token prompt. #### Why 128 MiB diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 0db5a9e31f1..abe4a6dae61 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -822,6 +822,8 @@ struct ggml_backend_sched_transport_ring { 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 delivered; // this ring issued deliveries and has not been waited for since + bool reported_no_room; }; @@ -1890,6 +1892,7 @@ static void ggml_backend_sched_transport_free_ring(ggml_backend_sched_t sched, i ggml_backend_buffer_free(r->buffer); r->buffer = NULL; r->slot_size = 0; + r->delivered = false; for (int i = 0; i < GGML_SCHED_MAX_TRANSPORT_SLOTS; i++) { r->slots[i].release_armed = false; @@ -1917,6 +1920,16 @@ static void ggml_backend_sched_transport_release_ring(ggml_backend_sched_t sched r->n_staged = 0; } +// Give back every ring the current graph does not use. +// The ring is optional storage, so a scheduler that leaves the staged path -- for one graph or for good -- holds no device memory and no second device context for it. +static void ggml_backend_sched_transport_release_idle(ggml_backend_sched_t sched) { + for (int i = 0; i < sched->n_backends; i++) { + if (sched->transport.rings[i].n_staged == 0) { + ggml_backend_sched_transport_release_ring(sched, i); + } + } +} + static void ggml_backend_sched_transport_decline_backend(ggml_backend_sched_t sched, int backend_id) { struct ggml_backend_sched_transport * tr = &sched->transport; @@ -1980,6 +1993,17 @@ static bool ggml_backend_sched_size_pad(size_t size, size_t alignment, size_t * return ggml_backend_sched_size_add(size, alignment - rem, result); } +// How large a slot to allocate for a window that needs `need`, where `limit` is the most that may be spent on one. +// The staged window widens on nearly every prefill ubatch, and a wider window than the ring holds frees the ring and allocates it again, which blocks the host on the device. +// Powers of two make that happen a handful of times over a prompt instead of once per ubatch. +static size_t ggml_backend_sched_transport_slot_alloc(size_t need, size_t limit) { + size_t size = 1; + while (size < need && size <= SIZE_MAX/2) { + size *= 2; + } + return std::max(std::min(size, limit), need); +} + // The transfer backend and the slot events are created on demand, so a backend that 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; @@ -2036,6 +2060,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { } if (!ggml_backend_sched_transport_enabled(sched) || sched->n_splits == 0) { + ggml_backend_sched_transport_release_idle(sched); return; } @@ -2046,6 +2071,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { 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; + ggml_backend_sched_transport_release_idle(sched); return; } tr->split_order = pnew; @@ -2065,6 +2091,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { for (int i = 0; i < sched->n_splits; i++) { tr->split_order[i] = -1; } + ggml_backend_sched_transport_release_idle(sched); return; } @@ -2072,6 +2099,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { 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__); + ggml_backend_sched_transport_release_idle(sched); return; } tr->input_staged = pnew; @@ -2096,6 +2124,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { for (int i = 0; i < sched->n_splits; i++) { tr->split_order[i] = -1; } + ggml_backend_sched_transport_release_idle(sched); return; } @@ -2214,7 +2243,9 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { ggml_backend_sched_transport_decline_backend(sched, bid); continue; } + // a graph that stages nothing here gives the ring and the transfer context back, so a scheduler that leaves the staged path holds no device memory for it if (r->n_staged == 0) { + ggml_backend_sched_transport_release_ring(sched, bid); continue; } @@ -2270,10 +2301,21 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { continue; } - ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(buft, ring_size); + // the slot grows past what this graph needs, but never past what the full context needs, what the budget allows, or what the headroom check just approved + size_t slot_limit = std::min(slot_size_max[bid], SIZE_MAX/tr->n_slots); + if (tr->budget > 0) { + slot_limit = std::min(slot_limit, tr->budget/tr->n_slots); + } + if (dev_free > GGML_SCHED_TRANSPORT_HEADROOM) { + slot_limit = std::min(slot_limit, (dev_free - GGML_SCHED_TRANSPORT_HEADROOM)/tr->n_slots); + } + const size_t slot_alloc = ggml_backend_sched_transport_slot_alloc(slot_size[bid], slot_limit); + const size_t alloc_size = slot_alloc*tr->n_slots; + + ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(buft, alloc_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, + "pipelining disabled there\n", __func__, alloc_size >> 20, ggml_backend_name(sched->backends[bid])); ggml_backend_sched_transport_decline_backend(sched, bid); continue; @@ -2281,11 +2323,11 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_COMPUTE); r->buffer = buffer; - r->slot_size = slot_size[bid]; + r->slot_size = slot_alloc; 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); + ggml_backend_name(sched->backends[bid]), tr->n_slots, slot_alloc >> 10); } } @@ -2389,6 +2431,7 @@ static void ggml_backend_sched_transport_prefetch(ggml_backend_sched_t sched, in // 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, r->transfer); + r->delivered = true; r->scan_cursor = i + 1; } } @@ -2441,6 +2484,7 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { } for (int i = 0; i < sched->n_backends; i++) { ggml_backend_synchronize(sched->backends[i]); + sched->transport.rings[i].delivered = false; } ggml_backend_sched_transport_clear_addresses(sched); @@ -2487,11 +2531,12 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s // Within a graph the stable prefix keeps the host off what is still being read, but the next graph writes wherever its own ubatch lands, so wait for the previous one here. // This is what the ordered path gets from its blocking copy, once per graph rather than once per split. for (int i = 0; i < sched->n_backends; i++) { - if (tr->rings[i].transfer == NULL) { + if (!tr->rings[i].delivered) { continue; } ggml_backend_synchronize(tr->rings[i].transfer); ggml_backend_synchronize(sched->backends[i]); + tr->rings[i].delivered = false; } // Prime every ring before the first consumer runs. @@ -3171,6 +3216,7 @@ void ggml_backend_sched_synchronize(ggml_backend_sched_t sched) { } for (int i = 0; i < sched->n_backends; i++) { ggml_backend_synchronize(sched->backends[i]); + sched->transport.rings[i].delivered = false; } if (!sched->is_alloc) { // if the graph is not already allocated, always use copy 0 after a synchronization From 590eafb6ab6f14282e8aacd242f207c062107ebf Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 6 Sep 2026 01:29:03 +0200 Subject: [PATCH 26/32] arg, llama : bound --kv-pipeline-budget The platform check it had can never fail on a 64-bit size_t, so any value was accepted and a large one was silently the same as 0. Cap it at 65536 MiB in all three places that parse it. Assisted-by: Claude Opus 5 --- common/arg.cpp | 8 ++++++-- src/llama-context.cpp | 6 ++++-- tools/llama-bench/llama-bench.cpp | 2 +- 3 files changed, 11 insertions(+), 5 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index a2ae73c987e..27ccf11df6c 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2450,10 +2450,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex "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), + "to save. 0 removes the cap, 65536 is the largest accepted (default: %d)", params.kv_pipeline_budget_mib), [](common_params & params, int value) { constexpr size_t mib = 1024u*1024u; - if (value < 0 || (size_t) value > std::numeric_limits::max()/mib) { + // a slot holds one attention layer's K or V, so a cap this large is already uncapped: past it the value is a typo, not a budget + if (value < 0 || value > 65536) { + throw std::invalid_argument("--kv-pipeline-budget must be between 0 and 65536 MiB"); + } + if ((size_t) value > std::numeric_limits::max()/mib) { throw std::invalid_argument("--kv-pipeline-budget is out of range for this platform"); } params.kv_pipeline_budget_mib = value; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 2abe49953fa..463b60bc90b 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -871,8 +871,10 @@ void llama_context::sched_reserve(uint32_t n_tokens_req, uint32_t n_kv_req) { if (cparams.kv_pipeline_depth > 14) { throw std::invalid_argument("kv_pipeline_depth must be between 0 and 14"); } - if (cparams.kv_pipeline_budget_mib > std::numeric_limits::max()/mib) { - throw std::invalid_argument("kv_pipeline_budget_mib is too large for this platform"); + // a slot holds one attention layer's K or V, so a cap of 64 GiB is already uncapped: past it the value is a typo, not a budget + constexpr uint64_t max_budget_mib = std::min(65536, std::numeric_limits::max()/mib); + if (cparams.kv_pipeline_budget_mib > max_budget_mib) { + throw std::invalid_argument("kv_pipeline_budget_mib must be between 0 and 65536 MiB"); } sched.reset(ggml_backend_sched_new( diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index f3ecf991e1e..d524c3ac876 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -883,7 +883,7 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { } auto p = parse_int_range(argv[i]); for (int budget : p) { - if (budget < 0) { + if (budget < 0 || budget > 65536) { invalid_param = true; break; } From 1ede08fd9ab628f6bf27033dc32c26a014867274 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 6 Sep 2026 01:29:13 +0200 Subject: [PATCH 27/32] llama : do not clear the KV buffers under a running decode A staged delivery reads the host cache long after the decode that issued it returned, so a memset of those buffers races it. llama_memory_clear holds no context and cannot wait, so say so on the public function, and wait where a context is at hand. Assisted-by: Claude Opus 5 --- include/llama.h | 1 + src/llama-context.cpp | 3 +++ 2 files changed, 4 insertions(+) diff --git a/include/llama.h b/include/llama.h index e91225dd174..72db8f0c11c 100644 --- a/include/llama.h +++ b/include/llama.h @@ -746,6 +746,7 @@ extern "C" { // Clear the memory contents // If data == true, the data buffers will also be cleared together with the metadata + // NOTE: with data == true, call llama_synchronize() first if a decode may still be running - a decode can still be reading the buffers LLAMA_API void llama_memory_clear( llama_memory_t mem, bool data); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 463b60bc90b..edee6368a0c 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -3898,6 +3898,9 @@ void llama_context::opt_epoch_iter( const uint32_t n_batch = std::min(this->n_batch(), n_ctx); const uint32_t n_ubatch = std::min(this->n_ubatch(), n_batch); + // a previous decode can still be reading the cache buffers, so do not clear them under it + synchronize(); + memory->clear(true); for (uint32_t pos_ctx = 0; pos_ctx < n_ctx; pos_ctx += n_batch) { From 9fe8c48e8f0d6a2ca304aff494f44965816e62ef Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 6 Sep 2026 10:42:01 +0200 Subject: [PATCH 28/32] llama : wait for the decode before clearing the memory buffers clear(data=true) memsets the buffers while a decode can still be reading them: a staged delivery for a host-resident cache, the graph itself for a device-resident one. The second one is not new and is easy to hit - llama_decode followed by llama_memory_clear(mem, true) changed the logits of that decode on every trial. Every memory type already passes the context down through init_update, so the caches keep it from there and wait on it before the memset. Assisted-by: Claude Opus 5 --- docs/kv-transport-pipelining.md | 1 + include/llama.h | 2 +- src/llama-context.cpp | 3 --- src/llama-kv-cache.cpp | 8 ++++++++ src/llama-kv-cache.h | 4 ++++ src/llama-memory-recurrent.cpp | 9 ++++++++- src/llama-memory-recurrent.h | 4 ++++ 7 files changed, 26 insertions(+), 5 deletions(-) diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index 9b17d572516..8bee834178c 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -280,6 +280,7 @@ A device-resident KV run is unaffected, and was measured to confirm it: 38.5612 - **A multi-stream window is delivered one range per stream**, keyed on the last dimension. A window whose streams are not on that dimension keeps the single flat range, which is correct but not accelerated. - **One ring per accelerator.** A layer-split model pipelines on every device that qualifies; a device with no room within the budget falls back to the ordered path on its own without disabling the others. - **Tensor parallelism keeps the ordered path.** See [Tensor parallelism](#tensor-parallelism). +- **A host write to the cache waits for the delivery.** `llama_memory_clear(mem, true)` synchronizes the context before it clears the buffers, because a delivery the last decode issued can still be reading them. This was already needed without the transport: with a device-resident cache the same call cleared the buffers under the running graph, and `llama_decode` followed by that clear changed the logits of that decode on every trial. - 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` and `GGML_KV_PIPELINE_BUDGET_MIB` provide scheduler defaults. Explicit scheduler settings and command-line options take precedence. diff --git a/include/llama.h b/include/llama.h index 72db8f0c11c..1873f48dabd 100644 --- a/include/llama.h +++ b/include/llama.h @@ -746,7 +746,7 @@ extern "C" { // Clear the memory contents // If data == true, the data buffers will also be cleared together with the metadata - // NOTE: with data == true, call llama_synchronize() first if a decode may still be running - a decode can still be reading the buffers + // NOTE: with data == true this waits for a decode that is still running, which can still be reading the buffers LLAMA_API void llama_memory_clear( llama_memory_t mem, bool data); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index edee6368a0c..463b60bc90b 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -3898,9 +3898,6 @@ void llama_context::opt_epoch_iter( const uint32_t n_batch = std::min(this->n_batch(), n_ctx); const uint32_t n_ubatch = std::min(this->n_ubatch(), n_batch); - // a previous decode can still be reading the cache buffers, so do not clear them under it - synchronize(); - memory->clear(true); for (uint32_t pos_ctx = 0; pos_ctx < n_ctx; pos_ctx += n_batch) { diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index e79aa206e9e..d928b817714 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -478,6 +478,11 @@ void llama_kv_cache::clear(bool data) { } if (data) { + // a decode can still be delivering these buffers to the device, and the memset would race that read + if (lctx) { + llama_synchronize(lctx); + } + for (auto & [_, buf] : ctxs_bufs) { ggml_backend_buffer_clear(buf.get(), 0); } @@ -856,6 +861,9 @@ uint32_t llama_kv_cache::get_attn_reserve_capacity() const { llama_memory_context_ptr llama_kv_cache::init_update(llama_context * lctx, bool optimize) { GGML_UNUSED(optimize); + // every decode prepares an update, and every memory type passes the context down to its caches, so this is set before anything can be in flight + this->lctx = lctx; + bool do_shift = get_has_shift(); return std::make_unique(this, lctx, do_shift, std::move(sc_info)); diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 4188f55e5ab..8517b37e6cb 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -310,6 +310,10 @@ class llama_kv_cache : public llama_memory_i { // env: LLAMA_KV_CACHE_DEBUG int debug = 0; + // the context that evaluates this cache, taken from the last update it prepared + // clear() writes the buffers, a delivery of them can still be in flight, and only the context can wait for it + llama_context * lctx = nullptr; + // this is the SWA type of the cache - not to be confused with the model SWA type const llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE; diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index 543f34be1a9..647d46431df 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -155,6 +155,11 @@ void llama_memory_recurrent::clear(bool data) { used = 0; if (data) { + // a decode can still be reading these buffers, and the memset would race it + if (lctx) { + llama_synchronize(lctx); + } + for (auto & [_, buf] : ctxs_bufs) { ggml_backend_buffer_clear(buf.get(), 0); } @@ -576,9 +581,11 @@ llama_memory_context_ptr llama_memory_recurrent::init_full() { } llama_memory_context_ptr llama_memory_recurrent::init_update(llama_context * lctx, bool optimize) { - GGML_UNUSED(lctx); GGML_UNUSED(optimize); + // every decode prepares an update, and every memory type passes the context down, so this is set before anything can be in flight + this->lctx = lctx; + return std::make_unique(LLAMA_MEMORY_STATUS_NO_UPDATE); } diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index c23ed2bb4ed..38bbe22d633 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -131,6 +131,10 @@ class llama_memory_recurrent : public llama_memory_i { llama_recurrent_snapshot_mode next_snapshot_mode; bool sparse_metadata_active = false; + // the context that evaluates this memory, taken from the last update it prepared + // clear() writes the buffers, a decode can still be reading them, and only the context can wait for it + llama_context * lctx = nullptr; + // ggml contexts for the KV cache along with the allocated backend buffers: std::vector> ctxs_bufs; From e0446cec0854e706243ca1ab76c986445d815513 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 6 Sep 2026 13:22:12 +0200 Subject: [PATCH 29/32] sched : fix the review issues of the pipelined transport Keep a staged input's producer on the CPU or on the consumer itself, keep the transfer context over a graph that stages nothing, and stop asking a device that cannot give one. Say in the header and the docs that the destination is CUDA-only, and that a delivering graph costs the pipelining of n_copies > 1. Assisted-by: Claude Opus 5 --- docs/kv-transport-pipelining.md | 2 + ggml/include/ggml-backend.h | 2 +- ggml/src/ggml-backend.cpp | 32 +++++++++++--- tests/test-alloc.cpp | 76 +++++++++++++++++++++++++++++++++ 4 files changed, 105 insertions(+), 7 deletions(-) diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index 8bee834178c..99050eb95cf 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -279,6 +279,8 @@ A device-resident KV run is unaffected, and was measured to confirm it: 38.5612 - **A cap is per graph, not per sequence.** `--kv-pipeline-budget` bounds the window one graph delivers, which is `n_kv * n_stream` over every sequence in the ubatch, so it cannot be applied to one sequence of a batch and not another. - **A multi-stream window is delivered one range per stream**, keyed on the last dimension. A window whose streams are not on that dimension keeps the single flat range, which is correct but not accelerated. - **One ring per accelerator.** A layer-split model pipelines on every device that qualifies; a device with no room within the budget falls back to the ordered path on its own without disabling the others. +- **The producer of a staged input must be the CPU or the consumer itself.** Neither part of a staged delivery is ordered against a third device: the stable prefix goes on the transfer stream and the rest on the consumer's own stream, where the ordered path would have synchronized the producer first. An input a second accelerator writes keeps the ordered path. +- **It turns graph-level pipeline parallelism off while it is delivering.** A graph that delivered has to block the host on its consumer before the next graph writes the host cache, because the host source of a delivery is read long after the call that issued it returned. That block is what `n_copies > 1` exists to avoid, so the two do not overlap: with `-sm layer` over several GPUs and `--kv-cpu-pinned`, `llama_context` enables both and the ring wins. Use `--kv-pipeline-depth 0` to keep the graph-level pipelining instead. - **Tensor parallelism keeps the ordered path.** See [Tensor parallelism](#tensor-parallelism). - **A host write to the cache waits for the delivery.** `llama_memory_clear(mem, true)` synchronizes the context before it clears the buffers, because a delivery the last decode issued can still be reading them. This was already needed without the transport: with a device-resident cache the same call cleared the buffers under the running graph, and `llama_decode` followed by that clear changed the logits of that decode on every trial. - 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. diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index 50e4d081089..6371e6acf7e 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -333,7 +333,7 @@ extern "C" { // // `depth` is how many splits ahead deliveries run, 0 disables pipelining. // The ring holds a couple of slots more than that, so recycling a slot never waits for a reader that is still running. - // Needs a destination backend with asynchronous transfers and events, otherwise the setting is ignored. + // Only the CUDA backend is accepted as the destination: the ring needs a second context on the same device that transfers asynchronously and orders with events, and CUDA is where that is measured. Every other backend ignores the setting and keeps the ordered path. // Costs roughly (depth + 2) * (largest staged split) of device memory. // Must be called before the first graph is allocated, and returns false after that. // The ring is optional: if the graph cannot be allocated next to it, the scheduler releases the ring and keeps the ordered path. diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index abe4a6dae61..f619ed5f374 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1802,6 +1802,16 @@ static bool ggml_backend_sched_input_can_stage( return false; } + // the ordered path synchronizes the producer before it copies, the staged path never does: the early part goes on the transfer stream and the late part on the consumer's own + // both are ordered against the consumer alone, so a producer on another accelerator could still be writing the source when the delivery reads it + ggml_backend_t producer = ggml_backend_sched_get_tensor_backend(sched, input); + if (producer != NULL && producer != sched->backends[split->backend_id]) { + ggml_backend_dev_t dev = ggml_backend_get_device(producer); + if (dev == NULL || ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_CPU) { + return false; + } + } + return tensor_copy(input, split->backend_id, sched->cur_copy) != NULL; } @@ -1920,12 +1930,13 @@ static void ggml_backend_sched_transport_release_ring(ggml_backend_sched_t sched r->n_staged = 0; } -// Give back every ring the current graph does not use. -// The ring is optional storage, so a scheduler that leaves the staged path -- for one graph or for good -- holds no device memory and no second device context for it. +// Give back the staging of every ring the current graph does not use. +// The ring is optional storage, so a scheduler that leaves the staged path for a graph holds no device memory for it. +// The transfer context and the events stay: a graph that stages nothing is a normal thing to meet between staged ones -- a context shift runs on the CPU backend alone -- and rebuilding a device context around each of them costs far more than holding it. static void ggml_backend_sched_transport_release_idle(ggml_backend_sched_t sched) { for (int i = 0; i < sched->n_backends; i++) { if (sched->transport.rings[i].n_staged == 0) { - ggml_backend_sched_transport_release_ring(sched, i); + ggml_backend_sched_transport_free_ring(sched, i); } } } @@ -2243,9 +2254,9 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { ggml_backend_sched_transport_decline_backend(sched, bid); continue; } - // a graph that stages nothing here gives the ring and the transfer context back, so a scheduler that leaves the staged path holds no device memory for it + // a graph that stages nothing here gives the staging back, so a scheduler that leaves the staged path holds no device memory for it if (r->n_staged == 0) { - ggml_backend_sched_transport_release_ring(sched, bid); + ggml_backend_sched_transport_free_ring(sched, bid); continue; } @@ -2273,8 +2284,12 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { continue; } + // a device that cannot give a second context will not give one to the next graph either, so stop asking rather than rebuilding it per token if (!ggml_backend_sched_transport_ensure_backend(sched, bid)) { + GGML_LOG_WARN("%s: failed to create a transfer context on %s, pipelining disabled there\n", __func__, + ggml_backend_name(sched->backends[bid])); ggml_backend_sched_transport_decline_backend(sched, bid); + r->eligible = false; continue; } @@ -2314,10 +2329,13 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(buft, alloc_size); if (buffer == NULL) { + // the headroom check above already approved the size, so the device is out of memory for reasons this will not see coming + // retrying every graph would allocate and free a device context per token for nothing, so stop asking here too GGML_LOG_WARN("%s: failed to allocate %zu MiB for the transport ring on %s, " "pipelining disabled there\n", __func__, alloc_size >> 20, ggml_backend_name(sched->backends[bid])); ggml_backend_sched_transport_decline_backend(sched, bid); + r->eligible = false; continue; } ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_COMPUTE); @@ -2530,6 +2548,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s // A staged delivery reads its host source long after the call that issued it returned, and the previous graph can leave some of those reads in flight. // Within a graph the stable prefix keeps the host off what is still being read, but the next graph writes wherever its own ubatch lands, so wait for the previous one here. // This is what the ordered path gets from its blocking copy, once per graph rather than once per split. + // It also costs the graph-level pipelining of n_copies > 1, which exists to keep the host from blocking here: a scheduler that wants that instead keeps the ordered path with depth 0. for (int i = 0; i < sched->n_backends; i++) { if (!tr->rings[i].delivered) { continue; @@ -2957,7 +2976,8 @@ ggml_backend_sched_t ggml_backend_sched_new( int transport_depth; if (ggml_backend_sched_transport_depth_from_env(&transport_depth)) { - GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_depth(sched, transport_depth)); + const bool ok = ggml_backend_sched_set_transport_pipeline_depth(sched, transport_depth); + GGML_ASSERT(ok); } return sched; diff --git a/tests/test-alloc.cpp b/tests/test-alloc.cpp index c0f4b72d9e0..5b2d7a10343 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -30,6 +30,7 @@ struct dummy_backend_context { bool fail_backend_init = false; bool fail_event_init = false; int transfer_backend_count = 0; + int transfer_backend_inits = 0; // how often one was asked for, so a test can see a retry the count above hides int event_wait_count = 0; int set_tensor_async_count = 0; size_t set_tensor_async_bytes = 0; @@ -209,6 +210,7 @@ static void dummy_backend_device_get_memory(ggml_backend_dev_t, size_t * free, s static ggml_backend_t dummy_backend_device_init(ggml_backend_dev_t dev, const char *) { dummy_backend_context * ctx = (dummy_backend_context *) dev->context; + ctx->transfer_backend_inits++; if (ctx->fail_backend_init) { return nullptr; } @@ -1628,6 +1630,46 @@ static void test_transport_releases_ring_for_graph() { GGML_ASSERT(transfers == 0); } +// a graph that stages nothing gives the staging back but keeps the transfer context: a context shift runs between decodes and must not rebuild a device context every time +static void test_transport_keeps_context_over_idle_graph() { + dummy_backend cuda = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); + dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); + auto first = make_transport_graph(cpu, 64); + auto idle = make_transport_graph(cpu, 64); + auto second = make_transport_graph(cpu, 64); + idle.source->flags &= ~GGML_TENSOR_FLAG_TRANSPORT; + ggml_set_stable_prefix(first.source, 64); + ggml_set_stable_prefix(second.source, 64); + + ggml_backend_t backends[] = { cuda.handle.get(), cpu.handle.get() }; + ggml_backend_buffer_type_t bufts[] = { &cuda.buffer_type, &cpu.buffer_type }; + ggml_backend_sched_ptr sched(ggml_backend_sched_new(backends, bufts, 2, 128, false, false)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_depth(sched.get(), 1)); + + ggml_backend_sched_set_tensor_backend(sched.get(), first.output, cuda.handle.get()); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), first.ctx.graph)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), first.ctx.graph) == GGML_STATUS_SUCCESS); + const size_t staged_size = cuda.context->allocated_total(); + GGML_ASSERT(cuda.context->transfer_backend_count == 1); + GGML_ASSERT(cuda.context->transfer_backend_inits == 1); + + ggml_backend_sched_reset(sched.get()); + ggml_backend_sched_set_tensor_backend(sched.get(), idle.output, cuda.handle.get()); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), idle.ctx.graph)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), idle.ctx.graph) == GGML_STATUS_SUCCESS); + GGML_ASSERT(cuda.context->allocated_total() < staged_size); + GGML_ASSERT(cuda.context->transfer_backend_count == 1); + + ggml_backend_sched_reset(sched.get()); + ggml_backend_sched_set_tensor_backend(sched.get(), second.output, cuda.handle.get()); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), second.ctx.graph)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), second.ctx.graph) == GGML_STATUS_SUCCESS); + int64_t deliveries = 0; + transport_stats(sched.get(), &deliveries, nullptr, nullptr); + GGML_ASSERT(deliveries == 2); + GGML_ASSERT(cuda.context->transfer_backend_inits == 1); +} + static void set_test_env(const char * name, const char * value) { #ifdef _WIN32 GGML_ASSERT(_putenv_s(name, value) == 0); @@ -1788,10 +1830,42 @@ static void test_transport_partial_backend_failure() { GGML_ASSERT(deliveries == 1); GGML_ASSERT(cuda_fail.context->transfer_backend_count == 0); GGML_ASSERT(cuda_ok.context->transfer_backend_count == 1); + GGML_ASSERT(cuda_fail.context->transfer_backend_inits == 1); } GGML_ASSERT(cuda_ok.context->transfer_backend_count == 0); } +// a device that cannot give a transfer context is not asked again: a retry per graph would build one and tear it down per token +static void test_transport_stops_after_backend_failure() { + dummy_backend cuda = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); + dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); + cuda.context->fail_event_init = true; + auto first = make_transport_graph(cpu, 64); + auto second = make_transport_graph(cpu, 64); + ggml_set_stable_prefix(first.source, 64); + ggml_set_stable_prefix(second.source, 64); + + ggml_backend_t backends[] = { cuda.handle.get(), cpu.handle.get() }; + ggml_backend_buffer_type_t bufts[] = { &cuda.buffer_type, &cpu.buffer_type }; + ggml_backend_sched_ptr sched(ggml_backend_sched_new(backends, bufts, 2, 128, false, false)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_depth(sched.get(), 1)); + + ggml_backend_sched_set_tensor_backend(sched.get(), first.output, cuda.handle.get()); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), first.ctx.graph)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), first.ctx.graph) == GGML_STATUS_SUCCESS); + GGML_ASSERT(cuda.context->transfer_backend_inits == 1); + + ggml_backend_sched_reset(sched.get()); + ggml_backend_sched_set_tensor_backend(sched.get(), second.output, cuda.handle.get()); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), second.ctx.graph)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), second.ctx.graph) == GGML_STATUS_SUCCESS); + GGML_ASSERT(cuda.context->transfer_backend_inits == 1); + + int64_t deliveries = -1; + transport_stats(sched.get(), &deliveries, nullptr, nullptr); + GGML_ASSERT(deliveries == 0); +} + static void test_transport_excludes_meta() { dummy_backend meta = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_META, "CUDA", false); dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); @@ -1925,10 +1999,12 @@ int main() { run("test_transport_empty_graph", test_transport_empty_graph); run("test_transport_fallback_keeps_allocator_plan", test_transport_fallback_keeps_allocator_plan); run("test_transport_releases_ring_for_graph", test_transport_releases_ring_for_graph); + run("test_transport_keeps_context_over_idle_graph", test_transport_keeps_context_over_idle_graph); run("test_transport_environment_is_fallback", test_transport_environment_is_fallback); run("test_transport_depth_zero", test_transport_depth_zero); run("test_transport_budget_recovers", test_transport_budget_recovers); run("test_transport_partial_backend_failure", test_transport_partial_backend_failure); + run("test_transport_stops_after_backend_failure", test_transport_stops_after_backend_failure); run("test_transport_excludes_meta", test_transport_excludes_meta); run("test_transport_requires_annotation", test_transport_requires_annotation); run("test_transport_excludes_non_cuda", test_transport_excludes_non_cuda); From b6e3cd55a1db0803ea33bad383b8b0cf3d4e4fe2 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 6 Sep 2026 14:26:54 +0200 Subject: [PATCH 30/32] sched, llama: refinements to pipelined transport handling - Fix kv_pipeline_budget_mib handling to preserve negative sentinel value (-1 = not set, 0 = no cap) - Add null check for transport backend before synchronizing - Disable pipelined transport when n_copies > 1 to avoid conflicts with pipeline parallelism - Expose set_lctx() method in kv_cache for hybrid index context management Co-Authored-By: Claude Haiku 4.5 Claude-Session: https://claude.ai/code/session_01PbyQFddbGJQ6MpRwUH2Rtj --- common/common.cpp | 5 ++++- ggml/src/ggml-backend.cpp | 16 +++++++++++++++- src/llama-kv-cache.cpp | 6 +++++- src/llama-kv-cache.h | 3 +++ src/llama-memory-hybrid-idx.cpp | 5 +++++ 5 files changed, 32 insertions(+), 3 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index d77248d9992..4662daad8f4 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1746,7 +1746,10 @@ 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; + // 0 removes the cap, so a negative value must not fall through to it + if (params.kv_pipeline_budget_mib >= 0) { + cparams.kv_pipeline_budget_mib = (uint32_t) params.kv_pipeline_budget_mib; + } cparams.recurrent_state_offload = params.recurrent_state_offload; cparams.kv_gpu_layers = (uint32_t) std::max(0, params.kv_gpu_layers); cparams.phase_aware_workspace = params.phase_aware_workspace; diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index f619ed5f374..8679c4f22f8 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -2553,7 +2553,9 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s if (!tr->rings[i].delivered) { continue; } - ggml_backend_synchronize(tr->rings[i].transfer); + if (tr->rings[i].transfer) { + ggml_backend_synchronize(tr->rings[i].transfer); + } ggml_backend_synchronize(sched->backends[i]); tr->rings[i].delivered = false; } @@ -3017,6 +3019,18 @@ bool ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, return true; } + // n_copies > 1 overlaps the graphs through sched->events, and a staged delivery has to block the host on the previous graph: the two cancel out + if (sched->n_copies > 1) { + tr->depth = 0; + tr->n_slots = GGML_SCHED_TRANSPORT_MARGIN; + + if (tr->debug > 0) { + GGML_LOG_INFO("%s: pipeline parallelism is on, staying on the ordered path\n", __func__); + } + + return true; + } + int n_eligible = 0; for (int i = 0; i < sched->n_backends; i++) { ggml_backend_t backend = sched->backends[i]; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index d928b817714..eca10aac63e 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -858,11 +858,15 @@ uint32_t llama_kv_cache::get_attn_reserve_capacity() const { return get_size(); } +void llama_kv_cache::set_lctx(llama_context * lctx) { + this->lctx = lctx; +} + llama_memory_context_ptr llama_kv_cache::init_update(llama_context * lctx, bool optimize) { GGML_UNUSED(optimize); // every decode prepares an update, and every memory type passes the context down to its caches, so this is set before anything can be in flight - this->lctx = lctx; + set_lctx(lctx); bool do_shift = get_has_shift(); diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 8517b37e6cb..f12dda86ed5 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -167,6 +167,9 @@ class llama_kv_cache : public llama_memory_i { uint32_t get_size() const; uint32_t get_n_stream() const; + // the context that evaluates this cache; a cache that prepares no update of its own is told by its owner + void set_lctx(llama_context * lctx); + bool get_has_shift() const; ggml_type type_k() const; diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index d65edf9dc85..fb122cf4a3a 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -129,6 +129,11 @@ llama_memory_context_ptr llama_memory_hybrid_idx::init_full() { } llama_memory_context_ptr llama_memory_hybrid_idx::init_update(llama_context * lctx, bool optimize) { + // the indexer builds no update graph, so it gets no context of its own, but clear() still has to wait for a decode that reads its buffers + if (mem_idx) { + mem_idx->set_lctx(lctx); + } + return std::make_unique(this, lctx, optimize); } From 977fb4a3a02888c1f299826697c77ba912b948c0 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 6 Sep 2026 21:49:02 +0200 Subject: [PATCH 31/32] sched, llama : fix the second review of the pipelined transport Round the slot size to the ring alignment when the budget or the free memory is the binding cap: slot k starts at k*slot_size, and neither cap is a multiple of it, so every slot after the first bound its entries to a misaligned address. Reproduced on a 27B at -c 32768 with --kv-pipeline-budget 17, which aborts with a CUDA misaligned address. Wait for the consumer before the priming prefetch whenever the graph is about to stage, not only when the ring delivered last graph: a context shift or a declined graph in between left the previous graph's writes to the host source in flight. Say which split list a plan was built for with a generation counter, rather than by split count and input count, which do not distinguish two different lists. Keep the ring over a few graphs that stage nothing instead of freeing it on the first one, and keep the transfer context over a budget or headroom decline, which the next graph can recover from. Lay out the plan from scheduler-owned storage instead of a hash set and an array per graph. Clamp a stable prefix to one stream rather than to the whole body, the way the header describes it. Validate the pipeline depth and budget once, in the context constructor, against bounds the header now names, instead of restating them in three places. Assisted-by: Claude Opus 5 --- common/arg.cpp | 12 +-- common/common.cpp | 7 +- ggml/include/ggml.h | 2 +- ggml/src/ggml-backend.cpp | 140 ++++++++++++++++++++---------- ggml/src/ggml.c | 5 +- include/llama.h | 5 ++ src/llama-context.cpp | 15 ++-- src/llama-kv-cache.cpp | 1 - tests/test-alloc.cpp | 79 +++++++++++++++-- tools/llama-bench/llama-bench.cpp | 6 +- 10 files changed, 192 insertions(+), 80 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 27ccf11df6c..65022db235a 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2438,8 +2438,8 @@ common_params_context common_params_parser_init(common_params & params, llama_ex "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"); + if (value < 0 || value > LLAMA_KV_PIPELINE_DEPTH_MAX) { + throw std::invalid_argument(string_format("--kv-pipeline-depth must be between 0 and %d", LLAMA_KV_PIPELINE_DEPTH_MAX)); } params.kv_pipeline_depth = value; } @@ -2450,12 +2450,12 @@ common_params_context common_params_parser_init(common_params & params, llama_ex "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, 65536 is the largest accepted (default: %d)", params.kv_pipeline_budget_mib), + "to save. 0 removes the cap, %d is the largest accepted (default: %d)", + LLAMA_KV_PIPELINE_BUDGET_MIB_MAX, params.kv_pipeline_budget_mib), [](common_params & params, int value) { constexpr size_t mib = 1024u*1024u; - // a slot holds one attention layer's K or V, so a cap this large is already uncapped: past it the value is a typo, not a budget - if (value < 0 || value > 65536) { - throw std::invalid_argument("--kv-pipeline-budget must be between 0 and 65536 MiB"); + if (value < 0 || value > LLAMA_KV_PIPELINE_BUDGET_MIB_MAX) { + throw std::invalid_argument(string_format("--kv-pipeline-budget must be between 0 and %d MiB", LLAMA_KV_PIPELINE_BUDGET_MIB_MAX)); } if ((size_t) value > std::numeric_limits::max()/mib) { throw std::invalid_argument("--kv-pipeline-budget is out of range for this platform"); diff --git a/common/common.cpp b/common/common.cpp index 4662daad8f4..b171ad9b29d 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1745,11 +1745,8 @@ 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; - // 0 removes the cap, so a negative value must not fall through to it - if (params.kv_pipeline_budget_mib >= 0) { - cparams.kv_pipeline_budget_mib = (uint32_t) params.kv_pipeline_budget_mib; - } + cparams.kv_pipeline_depth = (uint32_t) params.kv_pipeline_depth; + cparams.kv_pipeline_budget_mib = (uint32_t) params.kv_pipeline_budget_mib; cparams.recurrent_state_offload = params.recurrent_state_offload; cparams.kv_gpu_layers = (uint32_t) std::max(0, params.kv_gpu_layers); cparams.phase_aware_workspace = params.phase_aware_workspace; diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index d617b201cdf..254290608ef 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -714,7 +714,7 @@ extern "C" { // declare that the first nbytes of every stream of tensor->data cannot change while a graph that reads this tensor is being evaluated // a stream is one index of the last dimension of the view a reader takes of this tensor, and consecutive streams sit that view's nb[3] apart, so a reader that takes the storage whole has a single stream and the region is then simply its first nbytes - // nbytes is clamped to ggml_nbytes(tensor), and it must be set on the tensor that owns the storage, not on a view of it + // nbytes is clamped to one stream of tensor, and it must be set on the tensor that owns the storage, not on a view of it // it must describe the graph that is about to run, including when that graph is reused // a backend may deliver a declared region before the point in the graph that reads it, so 0 declares nothing and is always correct GGML_API void ggml_set_stable_prefix(struct ggml_tensor * tensor, size_t nbytes); diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 8679c4f22f8..62f285294d1 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -15,7 +15,6 @@ #include #include -#include #include #include #include @@ -797,6 +796,12 @@ static bool ggml_is_view_op(enum ggml_op op) { #define GGML_SCHED_TRANSPORT_BUDGET (128u*1024*1024) #endif +// How many graphs in a row a ring may stage nothing before its buffer is given back. +// A context shift or an encoder graph between two staged graphs is normal, and freeing the ring for one costs a device free plus the host block that allocating it again brings. +#ifndef GGML_SCHED_TRANSPORT_IDLE_GRAPHS +#define GGML_SCHED_TRANSPORT_IDLE_GRAPHS 4 +#endif + // One staging slot of the transport ring. // A slot is owned by the transfer stream while it is filled and by the consumer stream while it is read, and the two events below are the handover in each direction. struct ggml_backend_sched_transport_slot { @@ -821,6 +826,7 @@ struct ggml_backend_sched_transport_ring { 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 + int idle_graphs; // graphs in a row that staged nothing here bool delivered; // this ring issued deliveries and has not been waited for since @@ -844,9 +850,14 @@ struct ggml_backend_sched_transport { int * split_order; int plan_capacity; int plan_n_splits; - int plan_n_inputs; + uint64_t plan_gen; // the split list this plan was built for int n_staged; // over all rings, so that execution can skip the machinery entirely + // which copies a plan may put in a ring, and the split that owns each of them + // scheduler-owned so that laying out a plan costs no allocation per graph + struct ggml_hash_set staged_set; + int * staged_owner; // [staged_set.size] + // 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 go early moves with every ubatch, but which input copies live in the ring must not unsigned char * input_staged; @@ -917,6 +928,7 @@ struct ggml_backend_sched { struct ggml_backend_sched_split * splits; int n_splits; int splits_capacity; + uint64_t splits_gen; // bumped on every split, so a plan can say which split list it was built for // pipeline parallelism support int n_copies; @@ -1176,6 +1188,7 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra // reset splits sched->n_splits = 0; sched->n_graph_inputs = 0; + sched->splits_gen++; sched->is_reset = false; struct ggml_init_params params = { @@ -1900,9 +1913,10 @@ static void ggml_backend_sched_transport_free_ring(ggml_backend_sched_t sched, i } ggml_backend_buffer_free(r->buffer); - r->buffer = NULL; - r->slot_size = 0; - r->delivered = false; + r->buffer = NULL; + r->slot_size = 0; + r->delivered = false; + r->idle_graphs = 0; for (int i = 0; i < GGML_SCHED_MAX_TRANSPORT_SLOTS; i++) { r->slots[i].release_armed = false; @@ -1930,21 +1944,33 @@ static void ggml_backend_sched_transport_release_ring(ggml_backend_sched_t sched r->n_staged = 0; } -// Give back the staging of every ring the current graph does not use. -// The ring is optional storage, so a scheduler that leaves the staged path for a graph holds no device memory for it. -// The transfer context and the events stay: a graph that stages nothing is a normal thing to meet between staged ones -- a context shift runs on the CPU backend alone -- and rebuilding a device context around each of them costs far more than holding it. +// Count one graph that staged nothing on this ring, and give the staging back once there have been a few in a row. +// The ring is optional storage, so a scheduler that has left the staged path holds no device memory for it. +// It is not given back on the first idle graph: one between two staged ones is a normal thing to meet -- a context shift runs on the CPU backend alone -- and the ring is grown in powers of two exactly to keep allocating it again off the decode path. +// The transfer context and the events stay for the same reason: rebuilding a device context costs far more than holding it. +static void ggml_backend_sched_transport_ring_idle(ggml_backend_sched_t sched, int backend_id) { + struct ggml_backend_sched_transport_ring * r = &sched->transport.rings[backend_id]; + + if (r->buffer != NULL && ++r->idle_graphs > GGML_SCHED_TRANSPORT_IDLE_GRAPHS) { + ggml_backend_sched_transport_free_ring(sched, backend_id); + } +} + static void ggml_backend_sched_transport_release_idle(ggml_backend_sched_t sched) { for (int i = 0; i < sched->n_backends; i++) { if (sched->transport.rings[i].n_staged == 0) { - ggml_backend_sched_transport_free_ring(sched, i); + ggml_backend_sched_transport_ring_idle(sched, i); } } } +// Take this backend's splits out of the current plan and give its staging back. +// The transfer context and the events stay: what made this graph decline -- a window past the budget, a device that is momentarily full -- can be gone by the next one. static void ggml_backend_sched_transport_decline_backend(ggml_backend_sched_t sched, int backend_id) { struct ggml_backend_sched_transport * tr = &sched->transport; - ggml_backend_sched_transport_release_ring(sched, backend_id); + ggml_backend_sched_transport_free_ring(sched, backend_id); + tr->rings[backend_id].n_staged = 0; if (tr->split_order == NULL || tr->split_input_ofs == NULL || tr->input_staged == NULL) { return; @@ -1961,6 +1987,14 @@ static void ggml_backend_sched_transport_decline_backend(ggml_backend_sched_t sc } } +// Stop asking this backend for a ring, and give back the device context with it. +// For the declines that will not go away on their own: a device that cannot give a second context, or that fails an allocation the headroom check approved. +static void ggml_backend_sched_transport_disable_backend(ggml_backend_sched_t sched, int backend_id) { + ggml_backend_sched_transport_decline_backend(sched, backend_id); + ggml_backend_sched_transport_release_ring(sched, backend_id); + sched->transport.rings[backend_id].eligible = false; +} + // Give every ring back and stop asking for one. // The ring is optional, so it is the first thing to release when the device cannot hold it and the graph at the same time. // Returns whether any ring was holding memory. @@ -1970,8 +2004,7 @@ static bool ggml_backend_sched_transport_decline_all(ggml_backend_sched_t sched) bool released = false; for (int i = 0; i < sched->n_backends; i++) { released |= tr->rings[i].buffer != NULL; - ggml_backend_sched_transport_decline_backend(sched, i); - tr->rings[i].eligible = false; + ggml_backend_sched_transport_disable_backend(sched, i); } tr->n_staged = 0; @@ -2007,12 +2040,18 @@ static bool ggml_backend_sched_size_pad(size_t size, size_t alignment, size_t * // How large a slot to allocate for a window that needs `need`, where `limit` is the most that may be spent on one. // The staged window widens on nearly every prefill ubatch, and a wider window than the ring holds frees the ring and allocates it again, which blocks the host on the device. // Powers of two make that happen a handful of times over a prompt instead of once per ubatch. -static size_t ggml_backend_sched_transport_slot_alloc(size_t need, size_t limit) { +// Slot k starts at k*slot_size, so the result must be a multiple of the alignment: `limit` comes from the budget and the free memory and is not one. `need` already is. +static size_t ggml_backend_sched_transport_slot_alloc(size_t need, size_t limit, size_t alignment) { + GGML_ASSERT(alignment > 0 && need % alignment == 0); + size_t size = 1; while (size < need && size <= SIZE_MAX/2) { size *= 2; } - return std::max(std::min(size, limit), need); + size = std::max(std::min(size, limit), need); + size -= size % alignment; + + return std::max(size, need); } // The transfer backend and the slot events are created on demand, so a backend that never gets to stage anything -- no eligible inputs, or no room within the budget -- does not carry a second device context for nothing. @@ -2063,7 +2102,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { tr->n_staged = 0; tr->plan_n_splits = 0; - tr->plan_n_inputs = 0; + tr->plan_gen = 0; for (int i = 0; i < sched->n_backends; i++) { tr->rings[i].n_staged = 0; tr->rings[i].consumed = 0; @@ -2118,7 +2157,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { } memset(tr->input_staged, 0, n_inputs_total); tr->plan_n_splits = sched->n_splits; - tr->plan_n_inputs = n_inputs_total; + tr->plan_gen = sched->splits_gen; int n_candidates = 0; for (int i = 0; i < sched->n_splits; i++) { @@ -2143,10 +2182,28 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { // build one lookup table, then scan each graph node once size_t staged_hash_size = n_candidates; staged_hash_size += staged_hash_size/4 + 1; - struct ggml_hash_set staged_copies = ggml_hash_set_new(staged_hash_size); - int * staged_owner = (int *) malloc(staged_copies.size * sizeof(int)); - GGML_ASSERT(staged_owner != NULL); - for (size_t i = 0; i < staged_copies.size; i++) { + if (tr->staged_set.size < staged_hash_size) { + struct ggml_hash_set set = ggml_hash_set_new(staged_hash_size); + int * owner = (int *) realloc(tr->staged_owner, set.size * sizeof(int)); + if (owner == NULL) { + ggml_hash_set_free(&set); + GGML_LOG_WARN("%s: failed to allocate the transport plan, pipelining disabled for this graph\n", __func__); + for (int i = 0; i < sched->n_splits; i++) { + tr->split_order[i] = -1; + } + ggml_backend_sched_transport_release_idle(sched); + return; + } + ggml_hash_set_free(&tr->staged_set); + tr->staged_set = set; + tr->staged_owner = owner; + } else { + ggml_hash_set_reset(&tr->staged_set); + } + + struct ggml_hash_set * staged_copies = &tr->staged_set; + int * staged_owner = tr->staged_owner; + for (size_t i = 0; i < staged_copies->size; i++) { staged_owner[i] = -1; } @@ -2157,7 +2214,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { continue; } struct ggml_tensor * input_cpy = tensor_copy(split->inputs[j], split->backend_id, sched->cur_copy); - const size_t id = ggml_hash_find_or_insert(&staged_copies, input_cpy); + const size_t id = ggml_hash_find_or_insert(staged_copies, input_cpy); if (staged_owner[id] == -1) { staged_owner[id] = i; } else { @@ -2171,8 +2228,8 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { for (int j = 0; j < graph->n_nodes; j++) { const struct ggml_tensor * node = graph->nodes[j]; if (node->view_src != NULL) { - const size_t id = ggml_hash_find(&staged_copies, node->view_src); - if (id != GGML_HASHSET_FULL && ggml_bitset_get(staged_copies.used, id)) { + const size_t id = ggml_hash_find(staged_copies, node->view_src); + if (id != GGML_HASHSET_FULL && ggml_bitset_get(staged_copies->used, id)) { staged_owner[id] = -2; } } @@ -2180,8 +2237,8 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { if (node->src[k] == NULL) { continue; } - const size_t id = ggml_hash_find(&staged_copies, node->src[k]); - if (id != GGML_HASHSET_FULL && ggml_bitset_get(staged_copies.used, id) && staged_owner[id] >= 0 && i > staged_owner[id]) { + const size_t id = ggml_hash_find(staged_copies, node->src[k]); + if (id != GGML_HASHSET_FULL && ggml_bitset_get(staged_copies->used, id) && staged_owner[id] >= 0 && i > staged_owner[id]) { staged_owner[id] = -2; } } @@ -2195,15 +2252,13 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { continue; } struct ggml_tensor * input_cpy = tensor_copy(split->inputs[j], split->backend_id, sched->cur_copy); - const size_t id = ggml_hash_find(&staged_copies, input_cpy); - GGML_ASSERT(id != GGML_HASHSET_FULL && ggml_bitset_get(staged_copies.used, id)); + const size_t id = ggml_hash_find(staged_copies, input_cpy); + GGML_ASSERT(id != GGML_HASHSET_FULL && ggml_bitset_get(staged_copies->used, id)); if (staged_owner[id] < 0) { tr->input_staged[tr->split_input_ofs[i] + j] = 0; } } } - free(staged_owner); - ggml_hash_set_free(&staged_copies); // per-ring slot size and delivery order // the budget is applied to what this graph needs, so a run whose window stays small keeps the ring whatever -n_ctx says @@ -2254,9 +2309,9 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { ggml_backend_sched_transport_decline_backend(sched, bid); continue; } - // a graph that stages nothing here gives the staging back, so a scheduler that leaves the staged path holds no device memory for it + // a graph that stages nothing here counts as idle, and the ring is given back once a few of them have gone by if (r->n_staged == 0) { - ggml_backend_sched_transport_free_ring(sched, bid); + ggml_backend_sched_transport_ring_idle(sched, bid); continue; } @@ -2288,8 +2343,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { if (!ggml_backend_sched_transport_ensure_backend(sched, bid)) { GGML_LOG_WARN("%s: failed to create a transfer context on %s, pipelining disabled there\n", __func__, ggml_backend_name(sched->backends[bid])); - ggml_backend_sched_transport_decline_backend(sched, bid); - r->eligible = false; + ggml_backend_sched_transport_disable_backend(sched, bid); continue; } @@ -2324,7 +2378,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { if (dev_free > GGML_SCHED_TRANSPORT_HEADROOM) { slot_limit = std::min(slot_limit, (dev_free - GGML_SCHED_TRANSPORT_HEADROOM)/tr->n_slots); } - const size_t slot_alloc = ggml_backend_sched_transport_slot_alloc(slot_size[bid], slot_limit); + const size_t slot_alloc = ggml_backend_sched_transport_slot_alloc(slot_size[bid], slot_limit, r->alignment); const size_t alloc_size = slot_alloc*tr->n_slots; ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(buft, alloc_size); @@ -2334,8 +2388,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { GGML_LOG_WARN("%s: failed to allocate %zu MiB for the transport ring on %s, " "pipelining disabled there\n", __func__, alloc_size >> 20, ggml_backend_name(sched->backends[bid])); - ggml_backend_sched_transport_decline_backend(sched, bid); - r->eligible = false; + ggml_backend_sched_transport_disable_backend(sched, bid); continue; } ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_COMPUTE); @@ -2349,7 +2402,8 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { } } - tr->n_staged += r->n_staged; + r->idle_graphs = 0; + tr->n_staged += r->n_staged; } if (tr->n_staged == 0) { @@ -2538,19 +2592,15 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s struct ggml_backend_sched_transport * tr = &sched->transport; bool named_ordered_now = false; // 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; + const bool staged = tr->n_staged > 0 && tr->plan_gen == sched->splits_gen; // A staged delivery reads its host source long after the call that issued it returned, and the previous graph can leave some of those reads in flight. // Within a graph the stable prefix keeps the host off what is still being read, but the next graph writes wherever its own ubatch lands, so wait for the previous one here. // This is what the ordered path gets from its blocking copy, once per graph rather than once per split. // It also costs the graph-level pipelining of n_copies > 1, which exists to keep the host from blocking here: a scheduler that wants that instead keeps the ordered path with depth 0. + // A ring that is about to stage has to wait even when it delivered nothing last graph: the previous graph may have left the consumer writing the host source, and the priming prefetch below reads it. for (int i = 0; i < sched->n_backends; i++) { - if (!tr->rings[i].delivered) { + if (!tr->rings[i].delivered && !(staged && tr->rings[i].n_staged > 0)) { continue; } if (tr->rings[i].transfer) { @@ -3113,6 +3163,8 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) { free(sched->transport.split_order); free(sched->transport.split_input_ofs); free(sched->transport.input_staged); + free(sched->transport.staged_owner); + ggml_hash_set_free(&sched->transport.staged_set); 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]); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 86480ea9ce2..839c7e9a69b 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1328,8 +1328,9 @@ size_t ggml_nbytes_pad(const struct ggml_tensor * tensor) { 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; + // the count is per stream, so it is clamped to one, not to the whole body: a larger value would let a reader deliver bytes of the next stream early + const size_t stream = ggml_nbytes(tensor) - (size_t) (tensor->ne[3] - 1)*tensor->nb[3]; + tensor->stable_prefix = nbytes < stream ? nbytes : stream; } size_t ggml_get_stable_prefix(const struct ggml_tensor * tensor) { diff --git a/include/llama.h b/include/llama.h index 1873f48dabd..3a10d2f3369 100644 --- a/include/llama.h +++ b/include/llama.h @@ -38,6 +38,11 @@ #define LLAMA_TOKEN_NULL -1 +// bounds of llama_context_params::kv_pipeline_depth and ::kv_pipeline_budget_mib +// a staging slot holds one attention layer's K or V, so a budget of 64 GiB is already uncapped: past it the value is a typo, not a budget +#define LLAMA_KV_PIPELINE_DEPTH_MAX 14 +#define LLAMA_KV_PIPELINE_BUDGET_MIB_MAX 65536 + #define LLAMA_FILE_MAGIC_GGLA 0x67676c61u // 'ggla' #define LLAMA_FILE_MAGIC_GGSN 0x6767736eu // 'ggsn' #define LLAMA_FILE_MAGIC_GGSQ 0x67677371u // 'ggsq' diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 463b60bc90b..17205662fa3 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -147,6 +147,12 @@ llama_context::llama_context( cparams.offload_attn_compute = params.offload_kqv || (params.op_offload && params.kv_cpu_pinned); cparams.kv_pipeline_depth = params.kv_pipeline_depth; cparams.kv_pipeline_budget_mib = params.kv_pipeline_budget_mib; + if (cparams.kv_pipeline_depth > LLAMA_KV_PIPELINE_DEPTH_MAX) { + throw std::invalid_argument("kv_pipeline_depth must be <= " + std::to_string(LLAMA_KV_PIPELINE_DEPTH_MAX)); + } + if (cparams.kv_pipeline_budget_mib > std::min(LLAMA_KV_PIPELINE_BUDGET_MIB_MAX, std::numeric_limits::max()/(1024*1024))) { + throw std::invalid_argument("kv_pipeline_budget_mib must be <= " + std::to_string(LLAMA_KV_PIPELINE_BUDGET_MIB_MAX)); + } cparams.kv_gpu_layers = params.kv_gpu_layers; cparams.phase_aware_workspace = params.phase_aware_workspace; cparams.live_context_workspace = params.live_context_workspace; @@ -868,15 +874,6 @@ void llama_context::sched_reserve(uint32_t n_tokens_req, uint32_t n_kv_req) { auto create_sched = [&](bool pipeline_parallel) { constexpr size_t mib = 1024u*1024u; - if (cparams.kv_pipeline_depth > 14) { - throw std::invalid_argument("kv_pipeline_depth must be between 0 and 14"); - } - // a slot holds one attention layer's K or V, so a cap of 64 GiB is already uncapped: past it the value is a typo, not a budget - constexpr uint64_t max_budget_mib = std::min(65536, std::numeric_limits::max()/mib); - if (cparams.kv_pipeline_budget_mib > max_budget_mib) { - throw std::invalid_argument("kv_pipeline_budget_mib must be between 0 and 65536 MiB"); - } - sched.reset(ggml_backend_sched_new( backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, pipeline_parallel, cparams.op_offload)); diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index eca10aac63e..819eed63ca3 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1700,7 +1700,6 @@ void llama_kv_cache::clear_stable_prefixes() const { } 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()); diff --git a/tests/test-alloc.cpp b/tests/test-alloc.cpp index 5b2d7a10343..9fe9b574719 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -1466,6 +1466,55 @@ static void test_transport_entry_allocation() { } } +// Slot k starts at k*slot_size, so a slot size the budget caps must still be a multiple of the ring alignment, or every slot after the first binds its entries to a misaligned address. +static void test_transport_slot_alignment() { + const size_t alignment = 256; + dummy_backend cuda = dummy_backend_init(SIZE_MAX, alignment, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); + dummy_backend cpu = dummy_backend_init(SIZE_MAX, alignment, true); + + const size_t used = 3*alignment; // what this graph reads, already a whole number of entries + const size_t store = 16*alignment; + + auto ctx = make_context(); + ggml_tensor * source = ggml_new_tensor_1d(ctx.ctx, GGML_TYPE_F32, store/sizeof(float)); + source->flags |= GGML_TENSOR_FLAG_TRANSPORT; + ggml_tensor * window = ggml_view_1d(ctx.ctx, source, used/sizeof(float), 0); + ggml_tensor * output = ggml_scale(ctx.ctx, window, 2.0f); + ggml_build_forward_expand(ctx.graph, output); + + ggml_backend_buffer_ptr buffer(ggml_backend_buft_alloc_buffer(&cpu.buffer_type, store)); + source->buffer = buffer.get(); + source->data = ggml_backend_buffer_get_base(buffer.get()); + ggml_set_stable_prefix(source, used); + + // the budget lands between the window and the next power of two, and is not a multiple of the alignment + const int n_slots = 3; // depth 1 plus the margin + const size_t budget = n_slots*used + alignment/8; + + ggml_backend_t backends[] = { cuda.handle.get(), cpu.handle.get() }; + ggml_backend_buffer_type_t bufts[] = { &cuda.buffer_type, &cpu.buffer_type }; + ggml_backend_sched_ptr sched(ggml_backend_sched_new(backends, bufts, 2, 128, false, false)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_budget(sched.get(), budget)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_depth(sched.get(), 1)); + ggml_backend_sched_set_tensor_backend(sched.get(), output, cuda.handle.get()); + + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), ctx.graph)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), ctx.graph) == GGML_STATUS_SUCCESS); + + int64_t deliveries = 0; + transport_stats(sched.get(), &deliveries, NULL, NULL); + GGML_ASSERT(deliveries == 1); + + size_t ring_size = 0; + for (const auto & b : cuda.context->bindings) { + if (strncmp(b.tensor->name, "CUDA#", 5) == 0) { + ring_size = ggml_backend_buffer_get_size(b.buffer); + } + } + GGML_ASSERT(ring_size > 0); + GGML_ASSERT(ring_size % (n_slots*alignment) == 0); +} + // A window over several streams sits a fixed stride apart in one tensor, with cells between one stream's window and the next that the graph never reads. // The delivery has to cover each stream's window from its own offset and leave those cells alone. static void test_transport_multi_stream_ranges() { @@ -1630,8 +1679,8 @@ static void test_transport_releases_ring_for_graph() { GGML_ASSERT(transfers == 0); } -// a graph that stages nothing gives the staging back but keeps the transfer context: a context shift runs between decodes and must not rebuild a device context every time -static void test_transport_keeps_context_over_idle_graph() { +// a graph that stages nothing keeps the ring for a few graphs and the transfer context for good: a context shift runs between decodes and must not rebuild either every time +static void test_transport_keeps_ring_over_idle_graph() { dummy_backend cuda = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); auto first = make_transport_graph(cpu, 64); @@ -1649,15 +1698,26 @@ static void test_transport_keeps_context_over_idle_graph() { ggml_backend_sched_set_tensor_backend(sched.get(), first.output, cuda.handle.get()); GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), first.ctx.graph)); GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), first.ctx.graph) == GGML_STATUS_SUCCESS); - const size_t staged_size = cuda.context->allocated_total(); + const size_t n_staged_buffers = cuda.context->buffers.size(); GGML_ASSERT(cuda.context->transfer_backend_count == 1); GGML_ASSERT(cuda.context->transfer_backend_inits == 1); - ggml_backend_sched_reset(sched.get()); - ggml_backend_sched_set_tensor_backend(sched.get(), idle.output, cuda.handle.get()); - GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), idle.ctx.graph)); - GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), idle.ctx.graph) == GGML_STATUS_SUCCESS); - GGML_ASSERT(cuda.context->allocated_total() < staged_size); + auto run_idle = [&]() { + ggml_backend_sched_reset(sched.get()); + ggml_backend_sched_set_tensor_backend(sched.get(), idle.output, cuda.handle.get()); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), idle.ctx.graph)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), idle.ctx.graph) == GGML_STATUS_SUCCESS); + }; + + // one idle graph keeps the ring: freeing and allocating it again blocks the host on the device, which is what growing it in powers of two exists to avoid + run_idle(); + GGML_ASSERT(cuda.context->buffers.size() == n_staged_buffers); + + // a run of them gives it back: the ring is optional storage, so a scheduler that has left the staged path holds none + for (int i = 0; i < 16 && cuda.context->buffers.size() == n_staged_buffers; i++) { + run_idle(); + } + GGML_ASSERT(cuda.context->buffers.size() < n_staged_buffers); GGML_ASSERT(cuda.context->transfer_backend_count == 1); ggml_backend_sched_reset(sched.get()); @@ -1995,11 +2055,12 @@ int main() { run("test_resizable_buffers_owner_borrower_teardown_order", test_resizable_buffers_owner_borrower_teardown_order); run("test_transport_prefix_and_configuration", test_transport_prefix_and_configuration); run("test_transport_entry_allocation", test_transport_entry_allocation); + run("test_transport_slot_alignment", test_transport_slot_alignment); run("test_transport_multi_stream_ranges", test_transport_multi_stream_ranges); run("test_transport_empty_graph", test_transport_empty_graph); run("test_transport_fallback_keeps_allocator_plan", test_transport_fallback_keeps_allocator_plan); run("test_transport_releases_ring_for_graph", test_transport_releases_ring_for_graph); - run("test_transport_keeps_context_over_idle_graph", test_transport_keeps_context_over_idle_graph); + run("test_transport_keeps_ring_over_idle_graph", test_transport_keeps_ring_over_idle_graph); run("test_transport_environment_is_fallback", test_transport_environment_is_fallback); run("test_transport_depth_zero", test_transport_depth_zero); run("test_transport_budget_recovers", test_transport_budget_recovers); diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index d524c3ac876..1b2d03c89d1 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -481,7 +481,7 @@ static void print_usage(int /* argc */, char ** argv) { printf(" -mg, --main-gpu (default: %s)\n", join(cmd_params_defaults.main_gpu, ",").c_str()); printf(" -nkvo, --no-kv-offload <0|1> (default: %s)\n", join(cmd_params_defaults.no_kv_offload, ",").c_str()); printf(" -kvcp, --kv-cpu-pinned <0|1> (default: %s)\n", join(cmd_params_defaults.kv_cpu_pinned, ",").c_str()); - printf(" -kvpd, --kv-pipeline-depth <0...14> (default: %s)\n", join(cmd_params_defaults.kv_pipeline_depth, ",").c_str()); + printf(" -kvpd, --kv-pipeline-depth <0...14> (default: %s)\n", join(cmd_params_defaults.kv_pipeline_depth, ",").c_str()); printf(" -kvpb, --kv-pipeline-budget (default: %s)\n", join(cmd_params_defaults.kv_pipeline_budget_mib, ",").c_str()); printf(" -rso, --recurrent-state-offload <0|1> (default: %s)\n", join(cmd_params_defaults.recurrent_state_offload, ",").c_str()); printf(" -fa, --flash-attn (default: %s)\n", join(transform_to_str(cmd_params_defaults.flash_attn, llama_flash_attn_type_name), ",").c_str()); @@ -867,7 +867,7 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { } auto p = parse_int_range(argv[i]); for (int depth : p) { - if (depth < 0 || depth > 14) { + if (depth < 0 || depth > LLAMA_KV_PIPELINE_DEPTH_MAX) { invalid_param = true; break; } @@ -883,7 +883,7 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { } auto p = parse_int_range(argv[i]); for (int budget : p) { - if (budget < 0 || budget > 65536) { + if (budget < 0 || budget > LLAMA_KV_PIPELINE_BUDGET_MIB_MAX) { invalid_param = true; break; } From c26d6b05ae7c6269a80a04ae7142f4481d03641c Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 6 Sep 2026 23:40:56 +0200 Subject: [PATCH 32/32] sched, llama : fix the third review of the pipelined transport Split a staged delivery into per-stream ranges whatever the stable prefix says. The split described the geometry only when a prefix was set, so the same view moved a different number of bytes depending on a value that decides when bytes move, not which. Clear a fresh ring so the padding around a window is never uninitialised device memory. Keep the layers of a cache that shares cells on the ordered path. [TAG_KV_CACHE_SHARE_CELLS] gives the borrower the owner's K/V tensors, and the borrower returns from apply_ubatch before it can describe them, so they kept the owner's prefix while the borrower wrote its own rows underneath. The borrower now drops the transport flag from the tensors it takes. Bound a stable prefix by the tensor rather than by a stream the storage does not describe: the KV tensors carry their streams on ne[2], so the old clamp was the whole body and never bound anything. The reader clamps to one stream of its own view, which is where the stream count is known. Refuse a pipeline depth out of range instead of clamping it and reporting success, which made the context's own bounds check dead code and let a ggml caller get a configuration it did not ask for. Walk a per-ring list of staged splits in the look-ahead instead of rescanning the split list past every other backend's splits on each call. Release only the rings that were holding memory when the graph does not fit next to them, rather than disabling every device including those that never had one. Drop the recurrent memory's context back-pointer and its synchronize: no recurrent tensor is ever marked GGML_TENSOR_FLAG_TRANSPORT, so it guarded a race the transport cannot reach. Restore the pipeline environment variables from a scope guard in test-alloc, so an assert in the middle does not leave them set for the tests after it. Assisted-by: Claude Opus 5 --- docs/kv-transport-pipelining.md | 3 +- ggml/include/ggml-backend.h | 6 +- ggml/include/ggml.h | 8 +- ggml/src/ggml-backend.cpp | 195 ++++++++++++++++---------------- ggml/src/ggml.c | 7 +- src/llama-context.cpp | 1 + src/llama-kv-cache.cpp | 17 ++- src/llama-kv-cache.h | 5 +- src/llama-memory-recurrent.cpp | 9 +- src/llama-memory-recurrent.h | 4 - tests/test-alloc.cpp | 44 ++++--- 11 files changed, 155 insertions(+), 144 deletions(-) diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index 99050eb95cf..c27e8e9763c 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -195,7 +195,7 @@ The table above is what the feature costs uncapped, and it is the reason it is c - Over it the scheduler declines and keeps the ordered path for that graph. Later graphs are evaluated again, so a smaller live window can use the ring. - Declining costs nothing in steady state. Both the ring and the transfer backend's device context are released, and a graph that stages nothing at all releases them the same way. - The value is in MiB, `0` removes the cap, and 65536 is the largest accepted: a slot holds one attention layer's K or V, so a cap past that is a typo rather than a budget. -- The budget and the headroom check are decided before the graph is allocated, so they can still leave the graph short. If graph reservation fails, the rings are released and the reservation is retried once on the ordered path, and that scheduler keeps the ordered path from then on. An optional ring never turns a graph that fits into an allocation failure. +- The budget and the headroom check are decided before the graph is allocated, so they can still leave the graph short. If graph reservation fails, the rings are released and the reservation is retried once on the ordered path; the devices that were holding a ring keep the ordered path from then on, the ones that held none keep their eligibility. An optional ring never turns a graph that fits into an allocation failure. **The cap is applied to what the current graph needs, not to what the full context would need.** A run whose window stays small keeps the ring whatever `-n_ctx` says, which is the common case and the reason it is done this way: a staged input is a view of the cache tensor, so the full-context figure is there for the asking, but enforcing it would refuse the ring for every large `-c` even when the window never gets near it. The warning reports both numbers so that `--kv-pipeline-budget` can be sized against the one that matters. @@ -278,6 +278,7 @@ A device-resident KV run is unaffected, and was measured to confirm it: 38.5612 - 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. - **A cap is per graph, not per sequence.** `--kv-pipeline-budget` bounds the window one graph delivers, which is `n_kv * n_stream` over every sequence in the ubatch, so it cannot be applied to one sequence of a batch and not another. - **A multi-stream window is delivered one range per stream**, keyed on the last dimension. A window whose streams are not on that dimension keeps the single flat range, which is correct but not accelerated. +- **A cache that shares cells with another one keeps the ordered path for the layers it shares.** [TAG_KV_CACHE_SHARE_CELLS] gives the borrowing cache the owner's K/V tensors, so their stable prefix would have two writers with two slot layouts. The borrower drops `GGML_TENSOR_FLAG_TRANSPORT` from the tensors it takes; the layers it allocates itself are unaffected. - **One ring per accelerator.** A layer-split model pipelines on every device that qualifies; a device with no room within the budget falls back to the ordered path on its own without disabling the others. - **The producer of a staged input must be the CPU or the consumer itself.** Neither part of a staged delivery is ordered against a third device: the stable prefix goes on the transfer stream and the rest on the consumer's own stream, where the ordered path would have synchronized the producer first. An input a second accelerator writes keeps the ordered path. - **It turns graph-level pipeline parallelism off while it is delivering.** A graph that delivered has to block the host on its consumer before the next graph writes the host cache, because the host source of a delivery is read long after the call that issued it returned. That block is what `n_copies > 1` exists to avoid, so the two do not overlap: with `-sm layer` over several GPUs and `--kv-cpu-pinned`, `llama_context` enables both and the ring wins. Use `--kv-pipeline-depth 0` to keep the graph-level pipelining instead. diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index 6371e6acf7e..505e1f31df7 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -331,12 +331,12 @@ extern "C" { // Only persistent host inputs marked with GGML_TENSOR_FLAG_TRANSPORT are eligible, and their stable prefix must be current before each evaluation. // The producer must be the CPU or the same backend stream that consumes the late region. // - // `depth` is how many splits ahead deliveries run, 0 disables pipelining. + // `depth` is how many splits ahead deliveries run, 0 disables pipelining, and it must not be more than GGML_SCHED_MAX_TRANSPORT_SLOTS - GGML_SCHED_TRANSPORT_MARGIN (14 by default). // The ring holds a couple of slots more than that, so recycling a slot never waits for a reader that is still running. // Only the CUDA backend is accepted as the destination: the ring needs a second context on the same device that transfers asynchronously and orders with events, and CUDA is where that is measured. Every other backend ignores the setting and keeps the ordered path. // Costs roughly (depth + 2) * (largest staged split) of device memory. - // Must be called before the first graph is allocated, and returns false after that. - // The ring is optional: if the graph cannot be allocated next to it, the scheduler releases the ring and keeps the ordered path. + // Returns false for a depth out of range, and after the first graph is allocated. + // The ring is optional: if the graph cannot be allocated next to it, the scheduler releases the rings that were holding memory and stops asking for them. GGML_API bool ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, int depth); // Hard cap on the staging ring, in bytes, default 128 MiB and 0 removes the cap. diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 254290608ef..fdb70ce004f 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -703,7 +703,6 @@ extern "C" { void * extra; // extra things e.g. for ggml-cuda.cu // bytes at the start of every stream that stay unchanged for the current graph evaluation, 0 for none - // the count is from the start of a stream, not from the start of the tensor, so a reader that splits the storage into streams applies it to each of them union { size_t stable_prefix; char padding[8]; @@ -712,10 +711,9 @@ extern "C" { static const size_t GGML_TENSOR_SIZE = sizeof(struct ggml_tensor); - // declare that the first nbytes of every stream of tensor->data cannot change while a graph that reads this tensor is being evaluated - // a stream is one index of the last dimension of the view a reader takes of this tensor, and consecutive streams sit that view's nb[3] apart, so a reader that takes the storage whole has a single stream and the region is then simply its first nbytes - // nbytes is clamped to one stream of tensor, and it must be set on the tensor that owns the storage, not on a view of it - // it must describe the graph that is about to run, including when that graph is reused + // declare that the first nbytes of every stream of tensor->data cannot change while a graph that reads this tensor runs + // a reader splits the storage into streams along the last dimension of its own view, so nbytes counts from the start of a stream, not of the tensor, and the caller must not pass more than one stream holds + // set it on the tensor that owns the storage, not on a view of it, and refresh it for the graph that is about to run, including when that graph is reused // a backend may deliver a declared region before the point in the graph that reads it, so 0 declares nothing and is always correct 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); diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 62f285294d1..9acc68208e4 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -777,33 +777,30 @@ static bool ggml_is_view_op(enum ggml_op op) { #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 every delivery would recycle the split that was enqueued a moment ago and is still running -- the ordered path with extra steps. -// Two slots of margin put the recycled reader far enough behind to have finished. +// With no margin a delivery would recycle the slot of the split that was just enqueued and is still running, which is the ordered path with extra steps. #ifndef GGML_SCHED_TRANSPORT_MARGIN #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 it must not take the room those buffers may still have to grow into. +// Device memory the transport ring leaves unclaimed, for the graph buffers to grow into. #ifndef GGML_SCHED_TRANSPORT_HEADROOM #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, so past this the feature declines rather than quietly spending hundreds of MiB. +// A slot holds one layer's K or V over the whole context and grows with it, so a host-resident cache does not quietly spend back the device memory it exists to save. #ifndef GGML_SCHED_TRANSPORT_BUDGET #define GGML_SCHED_TRANSPORT_BUDGET (128u*1024*1024) #endif // How many graphs in a row a ring may stage nothing before its buffer is given back. -// A context shift or an encoder graph between two staged graphs is normal, and freeing the ring for one costs a device free plus the host block that allocating it again brings. +// A context shift or an encoder graph between two staged graphs is normal, and freeing the ring for one of those costs an allocation to get it back. #ifndef GGML_SCHED_TRANSPORT_IDLE_GRAPHS #define GGML_SCHED_TRANSPORT_IDLE_GRAPHS 4 #endif // One staging slot of the transport ring. -// A slot is owned by the transfer stream while it is filled and by the consumer stream while it is read, and the two events below are the handover in each direction. +// The transfer stream owns it while it is filled and the consumer while it is read; the two events 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 @@ -811,8 +808,7 @@ struct ggml_backend_sched_transport_slot { }; // One ring per accelerator the scheduler drives. -// A layer-split model gives every device its own splits and deliveries, so each needs its own transfer stream, staging and 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. +// A layer-split model gives every device its own splits, so one device running ahead must not take another's slots and one 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 @@ -825,7 +821,7 @@ struct ggml_backend_sched_transport_ring { 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 + int scan_cursor; // of those, how many the look-ahead has issued int idle_graphs; // graphs in a row that staged nothing here bool delivered; // this ring issued deliveries and has not been waited for since @@ -834,10 +830,8 @@ struct ggml_backend_sched_transport_ring { }; // 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. +// The ordered path issues a split's host-to-device copy on the consumer's own stream right before the kernels that read it, so a token pays copy + compute in series; this ring runs the stable part of a later split on a separate transfer stream instead. +// The scheduler owns the ring and never hands it to ggml-alloc: ggml-alloc may recycle a graph-owned copy once its last 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 per ring: depth + GGML_SCHED_TRANSPORT_MARGIN @@ -848,18 +842,21 @@ struct ggml_backend_sched_transport { // 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; + // the same plan seen from a ring: the split ids it delivers, in that order, grouped by backend + // the look-ahead walks this instead of rescanning the split list past the splits of every other backend + int * ring_split; + int ring_split_ofs[GGML_SCHED_MAX_BACKENDS]; int plan_capacity; int plan_n_splits; uint64_t plan_gen; // the split list this plan was built for int n_staged; // over all rings, so that execution can skip the machinery entirely // which copies a plan may put in a ring, and the split that owns each of them - // scheduler-owned so that laying out a plan costs no allocation per graph struct ggml_hash_set staged_set; int * staged_owner; // [staged_set.size] // 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 go early moves with every ubatch, but which input copies live in the ring must not + // membership is decided when the ring is laid out and does not move with the ubatch, unlike how much of an input may go early unsigned char * input_staged; int * split_input_ofs; // [plan_capacity + 1] int input_capacity; @@ -877,7 +874,6 @@ 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; // why the look-ahead stopped: the depth it was given, or a slot whose reader has not run yet - // the two want opposite fixes, so they are counted apart int64_t n_stop_depth; int64_t n_wait_recycle; int64_t p_stop_depth, p_wait_recycle; @@ -1742,8 +1738,8 @@ static bool ggml_backend_sched_transport_enabled(ggml_backend_sched_t sched) { } // How a staged input's delivery breaks into ranges. -// A window over one stream is one range and is delivered flat, exactly as ggml_nbytes(input) describes it. -// A window over several streams is one range per stream: the streams sit a fixed stride apart in both the source and the copy, which carries the source's layout, and the cells between one stream's window and the next are never read by this graph. +// A window over one stream is one range, delivered flat as ggml_nbytes(input) describes it. +// A window over several streams is one range per stream: the streams sit a fixed stride apart in the source and in the copy alike, and the cells between one stream's window and the next are never read by this graph. struct ggml_backend_sched_ranges { int64_t n; // ranges to deliver size_t stride; // bytes from one range to the next, in the source and in the copy alike @@ -1751,8 +1747,8 @@ struct ggml_backend_sched_ranges { size_t early; // leading bytes of a range that may go before the split that reads it }; -// The annotation lives on the tensor that owns the storage, and a split input is normally a view of it. -// The prefix is per stream, so a range can use it only when the view starts on a stream boundary; anything else keeps the ordered path rather than guessing where the streams fall. +// The ranges do not depend on the prefix: it decides only how many of those bytes may go early. +// It counts from the start of a stream, so it applies only when the view starts on a stream boundary; anything else keeps the whole window late rather than guessing where the streams fall. static void ggml_backend_sched_input_ranges(const struct ggml_tensor * input, struct ggml_backend_sched_ranges * out) { const struct ggml_tensor * base = input->view_src ? input->view_src : input; @@ -1761,12 +1757,7 @@ static void ggml_backend_sched_input_ranges(const struct ggml_tensor * input, st out->used = ggml_nbytes(input); out->early = 0; - if (base->stable_prefix == 0) { - return; - } - // a range is one stream's byte span, which is what the tensor covers below dimension 3 - // taking that span from ggml_nbytes keeps it right whatever order the dimensions below the stream are permuted into: attention reads a KV window with its rows on dimension 1 const size_t rows = ggml_nbytes(input) - (size_t) (input->ne[3] - 1)*input->nb[3]; const size_t offs = input->view_src ? input->view_offs : 0; if (input->nb[3] < rows || (offs != 0 && (input->nb[3] == 0 || offs % input->nb[3] != 0))) { @@ -1783,10 +1774,7 @@ static void ggml_backend_sched_input_ranges(const struct ggml_tensor * input, st } // 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 ubatch yet and 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. +// Independent of the stable prefix: membership decides where an input copy lives, which the allocator has to know when it reserves, and there is no ubatch yet at that point. static bool ggml_backend_sched_input_can_stage( ggml_backend_sched_t sched, struct ggml_backend_sched_split * split, int input_id) { if (!ggml_backend_sched_transport_ring_enabled(sched, split->backend_id)) { @@ -1815,8 +1803,8 @@ static bool ggml_backend_sched_input_can_stage( return false; } - // the ordered path synchronizes the producer before it copies, the staged path never does: the early part goes on the transfer stream and the late part on the consumer's own - // both are ordered against the consumer alone, so a producer on another accelerator could still be writing the source when the delivery reads it + // the staged path never synchronizes the producer, and both its parts are ordered against the consumer alone + // so a producer on another accelerator could still be writing the source when the delivery reads it ggml_backend_t producer = ggml_backend_sched_get_tensor_backend(sched, input); if (producer != NULL && producer != sched->backends[split->backend_id]) { ggml_backend_dev_t dev = ggml_backend_get_device(producer); @@ -1837,7 +1825,7 @@ static bool ggml_backend_sched_input_is_staged(ggml_backend_sched_t sched, int s static bool ggml_backend_sched_size_add(size_t a, size_t b, size_t * result); static bool ggml_backend_sched_size_pad(size_t size, size_t alignment, size_t * result); -// A ring entry costs what the backend would allocate for it, which can be more than its data: a buffer type may ask for padding past ggml_nbytes(), and its kernels may write into it. +// A ring entry costs what the backend would allocate for it, which can be more than ggml_nbytes(): a buffer type may ask for padding that its kernels write into. static bool ggml_backend_sched_transport_entry_size( ggml_backend_buffer_type_t buft, const struct ggml_tensor * t, size_t alignment, size_t * result) { return ggml_backend_sched_size_pad(ggml_backend_buft_get_alloc_size(buft, t), alignment, result); @@ -1878,8 +1866,7 @@ static void ggml_backend_sched_transport_assign_addresses(ggml_backend_sched_t s continue; } struct ggml_tensor * input_cpy = tensor_copy(split->inputs[j], split->backend_id, sched->cur_copy); - // bind through the backend, so the entry is initialized the same way as any other tensor the buffer holds - // a previous plan may have left this copy bound already + // bind through the backend, so the entry is set up like any other tensor of this buffer; a previous plan may have left it bound input_cpy->data = NULL; input_cpy->buffer = NULL; const enum ggml_status status = ggml_backend_tensor_alloc(r->buffer, input_cpy, slot + offset); @@ -1892,8 +1879,7 @@ static void ggml_backend_sched_transport_assign_addresses(ggml_backend_sched_t s } } -// The consumer is waited for through the slots' own release events, never through sched->backends[backend_id]. -// The scheduler does not own its backends and llama_context declares its scheduler before them, so on the teardown path they are already gone; the buffer and the events go through the buffer type and the device, which are not. +// The consumer is waited for through the slots' own release events, never through sched->backends[backend_id]: the scheduler does not own its backends and they can already be gone on the teardown path. static void ggml_backend_sched_transport_free_ring(ggml_backend_sched_t sched, int backend_id) { struct ggml_backend_sched_transport_ring * r = &sched->transport.rings[backend_id]; @@ -1945,9 +1931,8 @@ static void ggml_backend_sched_transport_release_ring(ggml_backend_sched_t sched } // Count one graph that staged nothing on this ring, and give the staging back once there have been a few in a row. -// The ring is optional storage, so a scheduler that has left the staged path holds no device memory for it. -// It is not given back on the first idle graph: one between two staged ones is a normal thing to meet -- a context shift runs on the CPU backend alone -- and the ring is grown in powers of two exactly to keep allocating it again off the decode path. -// The transfer context and the events stay for the same reason: rebuilding a device context costs far more than holding it. +// Not on the first one: a context shift between two staged graphs is normal, and getting the ring back costs an allocation. +// The transfer context and the events stay, because rebuilding a device context costs far more than holding it. static void ggml_backend_sched_transport_ring_idle(ggml_backend_sched_t sched, int backend_id) { struct ggml_backend_sched_transport_ring * r = &sched->transport.rings[backend_id]; @@ -1965,7 +1950,7 @@ static void ggml_backend_sched_transport_release_idle(ggml_backend_sched_t sched } // Take this backend's splits out of the current plan and give its staging back. -// The transfer context and the events stay: what made this graph decline -- a window past the budget, a device that is momentarily full -- can be gone by the next one. +// The transfer context and the events stay: a window past the budget or a device that is momentarily full can be gone by the next graph. static void ggml_backend_sched_transport_decline_backend(ggml_backend_sched_t sched, int backend_id) { struct ggml_backend_sched_transport * tr = &sched->transport; @@ -1988,23 +1973,28 @@ static void ggml_backend_sched_transport_decline_backend(ggml_backend_sched_t sc } // Stop asking this backend for a ring, and give back the device context with it. -// For the declines that will not go away on their own: a device that cannot give a second context, or that fails an allocation the headroom check approved. +// For the declines that will not go away: a device that cannot give a second context, or that fails an allocation the headroom check approved. static void ggml_backend_sched_transport_disable_backend(ggml_backend_sched_t sched, int backend_id) { ggml_backend_sched_transport_decline_backend(sched, backend_id); ggml_backend_sched_transport_release_ring(sched, backend_id); sched->transport.rings[backend_id].eligible = false; } -// Give every ring back and stop asking for one. -// The ring is optional, so it is the first thing to release when the device cannot hold it and the graph at the same time. +// Give every ring back when the graph cannot be allocated next to them. +// A ring that was holding memory is also stopped for good: it competed with the graph and would do so again. +// A backend that held none did not, so it keeps its eligibility. // Returns whether any ring was holding memory. static bool ggml_backend_sched_transport_decline_all(ggml_backend_sched_t sched) { struct ggml_backend_sched_transport * tr = &sched->transport; bool released = false; for (int i = 0; i < sched->n_backends; i++) { - released |= tr->rings[i].buffer != NULL; - ggml_backend_sched_transport_disable_backend(sched, i); + if (tr->rings[i].buffer != NULL) { + released = true; + ggml_backend_sched_transport_disable_backend(sched, i); + } else { + ggml_backend_sched_transport_decline_backend(sched, i); + } } tr->n_staged = 0; @@ -2038,9 +2028,8 @@ static bool ggml_backend_sched_size_pad(size_t size, size_t alignment, size_t * } // How large a slot to allocate for a window that needs `need`, where `limit` is the most that may be spent on one. -// The staged window widens on nearly every prefill ubatch, and a wider window than the ring holds frees the ring and allocates it again, which blocks the host on the device. -// Powers of two make that happen a handful of times over a prompt instead of once per ubatch. -// Slot k starts at k*slot_size, so the result must be a multiple of the alignment: `limit` comes from the budget and the free memory and is not one. `need` already is. +// The window widens on nearly every prefill ubatch and outgrowing the ring means allocating it again, so grow in powers of two to pay that a handful of times per prompt instead of once per ubatch. +// Slot k starts at k*slot_size, so the result must be a multiple of the alignment; `limit` comes from the budget and the free memory and is not one. static size_t ggml_backend_sched_transport_slot_alloc(size_t need, size_t limit, size_t alignment) { GGML_ASSERT(alignment > 0 && need % alignment == 0); @@ -2054,7 +2043,7 @@ static size_t ggml_backend_sched_transport_slot_alloc(size_t need, size_t limit, return std::max(size, need); } -// The transfer backend and the slot events are created on demand, so a backend that never gets to stage anything -- no eligible inputs, or no room within the budget -- does not carry a second device context for nothing. +// Created on demand, so a backend that never gets to stage anything 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]; @@ -2096,7 +2085,7 @@ static bool ggml_backend_sched_transport_ensure_backend(ggml_backend_sched_t sch } // 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. +// Called before the graph is allocated: ggml-alloc leaves a tensor that already has data alone, so the staged copies stay out of its reuse analysis. static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { struct ggml_backend_sched_transport * tr = &sched->transport; @@ -2117,15 +2106,18 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { 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) { + int * pord = (int *) realloc(tr->ring_split, sched->n_splits * sizeof(int)); + if (pnew == NULL || pofs == NULL || pord == 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; + tr->ring_split = pord ? pord : tr->ring_split; ggml_backend_sched_transport_release_idle(sched); return; } tr->split_order = pnew; tr->split_input_ofs = pofs; + tr->ring_split = pord; tr->plan_capacity = sched->n_splits; } @@ -2178,8 +2170,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { return; } - // a ring copy must have one owner and no views or later readers - // build one lookup table, then scan each graph node once + // a ring copy must have one owner and no views or later readers: build one lookup table, then scan each graph node once size_t staged_hash_size = n_candidates; staged_hash_size += staged_hash_size/4 + 1; if (tr->staged_set.size < staged_hash_size) { @@ -2262,8 +2253,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { // per-ring slot size and delivery order // the budget is applied to what this graph needs, so a run whose window stays small keeps the ring whatever -n_ctx says - // slot_size_max is what the same ring costs once the context is full, taken from the cache tensor the staged input is a view of - // it is reported rather than enforced, because deciding on it would refuse the ring for every large -c even when the window never gets there + // slot_size_max is what the same ring costs at the full context; it is reported rather than enforced, or every large -c would be refused a ring it never grows into size_t slot_size[GGML_SCHED_MAX_BACKENDS] = { 0 }; size_t slot_size_max[GGML_SCHED_MAX_BACKENDS] = { 0 }; bool size_overflow[GGML_SCHED_MAX_BACKENDS] = { false }; @@ -2302,6 +2292,23 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { slot_size_max[bid] = std::max(slot_size_max[bid], std::max(need, need_max)); } + // group the staged splits by ring, so the look-ahead indexes its own deliveries instead of rescanning the split list + int ring_split_n = 0; + for (int bid = 0; bid < sched->n_backends; bid++) { + tr->ring_split_ofs[bid] = ring_split_n; + ring_split_n += tr->rings[bid].n_staged; + } + { + int fill[GGML_SCHED_MAX_BACKENDS] = { 0 }; + for (int i = 0; i < sched->n_splits; i++) { + if (tr->split_order[i] < 0) { + continue; + } + const int bid = sched->splits[i].backend_id; + tr->ring_split[tr->ring_split_ofs[bid] + fill[bid]++] = i; + } + } + for (int bid = 0; bid < sched->n_backends; bid++) { struct ggml_backend_sched_transport_ring * r = &tr->rings[bid]; if (size_overflow[bid]) { @@ -2339,7 +2346,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { continue; } - // a device that cannot give a second context will not give one to the next graph either, so stop asking rather than rebuilding it per token + // a device that cannot give a second context will not give one to the next graph either if (!ggml_backend_sched_transport_ensure_backend(sched, bid)) { GGML_LOG_WARN("%s: failed to create a transfer context on %s, pipelining disabled there\n", __func__, ggml_backend_name(sched->backends[bid])); @@ -2352,7 +2359,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { 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 + // the graph allocator reserved before this, so leave it the room its buffers may still 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) { @@ -2370,7 +2377,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { continue; } - // the slot grows past what this graph needs, but never past what the full context needs, what the budget allows, or what the headroom check just approved + // grow past what this graph needs, but never past the full context, the budget, or what the headroom check just approved size_t slot_limit = std::min(slot_size_max[bid], SIZE_MAX/tr->n_slots); if (tr->budget > 0) { slot_limit = std::min(slot_limit, tr->budget/tr->n_slots); @@ -2383,8 +2390,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(buft, alloc_size); if (buffer == NULL) { - // the headroom check above already approved the size, so the device is out of memory for reasons this will not see coming - // retrying every graph would allocate and free a device context per token for nothing, so stop asking here too + // the headroom check approved this size, so the device is out of memory for reasons this cannot see coming; retrying every graph would cost a device context per token GGML_LOG_WARN("%s: failed to allocate %zu MiB for the transport ring on %s, " "pipelining disabled there\n", __func__, alloc_size >> 20, ggml_backend_name(sched->backends[bid])); @@ -2392,6 +2398,8 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { continue; } ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_COMPUTE); + // a delivery moves only what the graph reads, so the padding around it is never written here + ggml_backend_buffer_clear(buffer, 0); r->buffer = buffer; r->slot_size = slot_alloc; @@ -2445,8 +2453,8 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { } // 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 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. +// The margin slots put the slot a delivery recycles several readers behind the split just enqueued, so refilling it does not put the transfer stream back in lock-step with the consumer. +// Each ring keeps its own cursor: one device saturating its look-ahead must not stop another from running ahead. 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]; @@ -2455,20 +2463,20 @@ static void ggml_backend_sched_transport_prefetch(ggml_backend_sched_t sched, in return; } - 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] > r->consumed + tr->depth) { + const int * ring_split = tr->ring_split + tr->ring_split_ofs[backend_id]; + + for (int o = r->scan_cursor; o < r->n_staged; o++) { + if (o > r->consumed + tr->depth) { tr->n_stop_depth++; return; } + const int i = ring_split[o]; struct ggml_backend_sched_split * split = &sched->splits[i]; - struct ggml_backend_sched_transport_slot * slot = &r->slots[tr->split_order[i] % tr->n_slots]; + struct ggml_backend_sched_transport_slot * slot = &r->slots[o % 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 + // ordered stream to stream, not through the host: blocking the host here would hold back the work it has not enqueued yet if (slot->release_armed) { ggml_backend_event_wait(r->transfer, slot->release); slot->release_armed = false; @@ -2481,8 +2489,7 @@ static void ggml_backend_sched_transport_prefetch(ggml_backend_sched_t sched, in } 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 + // how much of this input is stable belongs to the ubatch about to run, not to the plan, so it can be less than when the ring was laid out struct ggml_backend_sched_ranges rg; ggml_backend_sched_input_ranges(input, &rg); if (rg.early == 0) { @@ -2499,12 +2506,11 @@ static void ggml_backend_sched_transport_prefetch(ggml_backend_sched_t sched, in tr->n_bytes_early += rg.early*rg.n; } - // 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 + // record the handover here rather than when the split runs: the transfer stream is FIFO, so an event recorded later would make the consumer wait for the whole look-ahead behind it ggml_backend_event_record(slot->ready, r->transfer); r->delivered = true; - r->scan_cursor = i + 1; + r->scan_cursor = o + 1; } } @@ -2527,7 +2533,7 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { } } - // lay out the transport rings and point the staged input copies at them before the graph is allocated, so ggml-alloc sees those copies as already allocated and leaves them alone + // lay out the rings before the graph is allocated, so ggml-alloc sees the staged copies as already allocated and leaves them alone ggml_backend_sched_transport_plan(sched); // allocate graph @@ -2561,13 +2567,13 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { ggml_backend_sched_transport_clear_addresses(sched); if (!ggml_gallocr_reserve_n(sched->galloc, &sched->graph, sched->node_backend_ids, sched->leaf_backend_ids)) { - // the rings hold device memory the graph itself may need, and the caller can no longer turn them off: give them back and reserve once on the ordered path + // the rings hold device memory the graph itself needs, and the caller can no longer turn them off if (!ggml_backend_sched_transport_decline_all(sched) || !ggml_gallocr_reserve_n(sched->galloc, &sched->graph, sched->node_backend_ids, sched->leaf_backend_ids)) { GGML_LOG_ERROR("%s: failed to allocate graph\n", __func__); return false; } - GGML_LOG_WARN("%s: the graph does not fit next to the transport rings, they are released and this scheduler stays on the ordered path\n", __func__); + GGML_LOG_WARN("%s: the graph does not fit next to the transport rings, the devices that held one are released and stay on the ordered path for the rest of this scheduler\n", __func__); } ggml_backend_sched_transport_assign_addresses(sched); if (!ggml_gallocr_alloc_graph(sched->galloc, &sched->graph)) { @@ -2594,11 +2600,9 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s // a reused graph keeps the plan that was made for it, so the split list it describes must be the one about to run const bool staged = tr->n_staged > 0 && tr->plan_gen == sched->splits_gen; - // A staged delivery reads its host source long after the call that issued it returned, and the previous graph can leave some of those reads in flight. - // Within a graph the stable prefix keeps the host off what is still being read, but the next graph writes wherever its own ubatch lands, so wait for the previous one here. - // This is what the ordered path gets from its blocking copy, once per graph rather than once per split. - // It also costs the graph-level pipelining of n_copies > 1, which exists to keep the host from blocking here: a scheduler that wants that instead keeps the ordered path with depth 0. - // A ring that is about to stage has to wait even when it delivered nothing last graph: the previous graph may have left the consumer writing the host source, and the priming prefetch below reads it. + // A staged delivery reads its host source long after the call that issued it returned, so the previous graph can leave reads in flight where this one's ubatch is about to write. + // Waiting here is what the ordered path gets from its blocking copy, once per graph rather than once per split. It is also why depth and n_copies > 1 do not go together. + // A ring that is about to stage waits even when it delivered nothing last graph: the priming prefetch below reads a host source the previous graph may still be writing. for (int i = 0; i < sched->n_backends; i++) { if (!tr->rings[i].delivered && !(staged && tr->rings[i].n_staged > 0)) { continue; @@ -2610,8 +2614,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s tr->rings[i].delivered = false; } - // Prime every ring before the first consumer runs. - // From here on deliveries are issued only after a split has been enqueued, never before, so recycling a slot can never hold back work the consumer could already be running. + // Prime every ring before the first consumer runs; after this a delivery goes out only once a split has been enqueued, so recycling a slot cannot hold back work the consumer could already run. // The cursors start over on every evaluation because the plan outlives the graph it was made for. if (staged) { for (int i = 0; i < sched->n_backends; i++) { @@ -2651,9 +2654,8 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s 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 + // the stable prefix went out on the transfer stream earlier; the rest may still have been written by an earlier split of this graph, so it is only safe to read now + // it goes on the consumer's own stream, already ordered ahead of the kernels and behind the reader of whatever occupied this slot before struct ggml_backend_sched_ranges rg; ggml_backend_sched_input_ranges(input, &rg); if (rg.used > rg.early) { @@ -2855,14 +2857,14 @@ 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 + // every kernel that reads this slot is enqueued, so it 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; tr->rings[split_backend_id].consumed++; tr->n_deliveries++; - // with this split's kernels already enqueued, the deliveries for the next staged splits can go out even if recycling their slot waits for a reader that is running + // this split's kernels are enqueued, so the next deliveries can go out even if recycling their slot waits on a reader that is running ggml_backend_sched_transport_prefetch(sched, split_backend_id); } @@ -3051,9 +3053,8 @@ bool ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, return false; } - depth = std::min(depth, GGML_SCHED_MAX_TRANSPORT_SLOTS - GGML_SCHED_TRANSPORT_MARGIN); - if (depth < 0) { - depth = 0; + if (depth < 0 || depth > GGML_SCHED_MAX_TRANSPORT_SLOTS - GGML_SCHED_TRANSPORT_MARGIN) { + return false; } ggml_backend_sched_transport_teardown(sched); @@ -3069,7 +3070,7 @@ bool ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, return true; } - // n_copies > 1 overlaps the graphs through sched->events, and a staged delivery has to block the host on the previous graph: the two cancel out + // n_copies > 1 overlaps graphs through sched->events, and a staged delivery blocks the host on the previous graph: the two cancel out if (sched->n_copies > 1) { tr->depth = 0; tr->n_slots = GGML_SCHED_TRANSPORT_MARGIN; @@ -3107,8 +3108,7 @@ bool ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, 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 + // the transfer backend writes the ring, and it only accepts the device's own default buffer type if (sched->bufts[i] != ggml_backend_dev_buffer_type(dev)) { continue; } @@ -3161,6 +3161,7 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) { } ggml_backend_sched_transport_teardown(sched); free(sched->transport.split_order); + free(sched->transport.ring_split); free(sched->transport.split_input_ofs); free(sched->transport.input_staged); free(sched->transport.staged_owner); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 839c7e9a69b..f91e2e5598b 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1328,9 +1328,10 @@ size_t ggml_nbytes_pad(const struct ggml_tensor * tensor) { void ggml_set_stable_prefix(struct ggml_tensor * tensor, size_t nbytes) { GGML_ASSERT(tensor); - // the count is per stream, so it is clamped to one, not to the whole body: a larger value would let a reader deliver bytes of the next stream early - const size_t stream = ggml_nbytes(tensor) - (size_t) (tensor->ne[3] - 1)*tensor->nb[3]; - tensor->stable_prefix = nbytes < stream ? nbytes : stream; + // the storage does not say how a reader splits it into streams, so only the tensor itself bounds the value here + // a reader clamps it again to one stream of its own view + 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) { diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 17205662fa3..bfb14cf6146 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -878,6 +878,7 @@ 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)); cparams.flash_attn_causal_prefix_supported = llama_sched_supports_flash_attn_causal_prefix(sched.get()); + // the scheduler refuses a depth past its own bound, which a build can set lower than LLAMA_KV_PIPELINE_DEPTH_MAX if (!ggml_backend_sched_set_transport_pipeline_budget(sched.get(), (size_t) cparams.kv_pipeline_budget_mib*mib) || !ggml_backend_sched_set_transport_pipeline_depth(sched.get(), cparams.kv_cpu_pinned || !cparams.offload_kqv ? (int) cparams.kv_pipeline_depth : 0)) { throw std::invalid_argument("invalid KV transport pipeline configuration"); diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 819eed63ca3..d81b1baeec3 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -264,6 +264,14 @@ llama_kv_cache::llama_kv_cache( layers.push_back(layer_share); layers.back().il = il; + // this cache writes the shared tensors at its own slot, so their stable prefix would have two writers: keep them on the ordered path + if (layers.back().k) { + layers.back().k->flags &= ~GGML_TENSOR_FLAG_TRANSPORT; + } + if (layers.back().v) { + layers.back().v->flags &= ~GGML_TENSOR_FLAG_TRANSPORT; + } + continue; } } @@ -865,7 +873,7 @@ void llama_kv_cache::set_lctx(llama_context * lctx) { llama_memory_context_ptr llama_kv_cache::init_update(llama_context * lctx, bool optimize) { GGML_UNUSED(optimize); - // every decode prepares an update, and every memory type passes the context down to its caches, so this is set before anything can be in flight + // every decode prepares an update, so this is set before a delivery can be in flight set_lctx(lctx); bool do_shift = get_has_shift(); @@ -1221,7 +1229,7 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & ubatch) { // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] - // a cache that shares cells keeps no stable prefix of its own: the layers it aliases get one from the cache that owns the cells, and the layers it does not stay on the ordered path + // a cache that shares cells sets no stable prefix: the layers it aliases lost the transport flag when it took them, and the rest are the owner's to describe if (other) { return; } @@ -1662,9 +1670,8 @@ ggml_tensor * llama_kv_cache::build_input_v_rot(ggml_context * ctx) const { } void llama_kv_cache::update_stable_prefixes(const slot_info & sinfo) const { - // Rows written by this ubatch, counted within a stream rather than across the [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. - // Streams sit end to end, so a row counted across the body would let the lowest stream cap every stream above it; per stream, each one keeps its own leading rows. + // the lowest row this ubatch writes, counted within a stream: everything below it keeps what the previous ubatch left there for the whole graph + // counting across the body instead would let the lowest stream cap every stream above it uint64_t min_row = UINT64_MAX; for (uint32_t s = 0; s < sinfo.n_stream(); ++s) { for (const uint32_t idx : sinfo.idxs[s]) { diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index f12dda86ed5..39b6eab4c1e 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -241,8 +241,8 @@ class llama_kv_cache : public llama_memory_i { bool is_reserve, const slot_info * sinfo = nullptr) 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. + // tell the scheduler which part of each layer's K/V storage this ubatch does not write, so a host-resident cache can be delivered 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; @@ -315,6 +315,7 @@ class llama_kv_cache : public llama_memory_i { // the context that evaluates this cache, taken from the last update it prepared // clear() writes the buffers, a delivery of them can still be in flight, and only the context can wait for it + // a cache belongs to one context: a cache that shares another's cells is still a separate object with its own buffers llama_context * lctx = nullptr; // this is the SWA type of the cache - not to be confused with the model SWA type diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index 647d46431df..543f34be1a9 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -155,11 +155,6 @@ void llama_memory_recurrent::clear(bool data) { used = 0; if (data) { - // a decode can still be reading these buffers, and the memset would race it - if (lctx) { - llama_synchronize(lctx); - } - for (auto & [_, buf] : ctxs_bufs) { ggml_backend_buffer_clear(buf.get(), 0); } @@ -581,11 +576,9 @@ llama_memory_context_ptr llama_memory_recurrent::init_full() { } llama_memory_context_ptr llama_memory_recurrent::init_update(llama_context * lctx, bool optimize) { + GGML_UNUSED(lctx); GGML_UNUSED(optimize); - // every decode prepares an update, and every memory type passes the context down, so this is set before anything can be in flight - this->lctx = lctx; - return std::make_unique(LLAMA_MEMORY_STATUS_NO_UPDATE); } diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index 38bbe22d633..c23ed2bb4ed 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -131,10 +131,6 @@ class llama_memory_recurrent : public llama_memory_i { llama_recurrent_snapshot_mode next_snapshot_mode; bool sparse_metadata_active = false; - // the context that evaluates this memory, taken from the last update it prepared - // clear() writes the buffers, a decode can still be reading them, and only the context can wait for it - llama_context * lctx = nullptr; - // ggml contexts for the KV cache along with the allocated backend buffers: std::vector> ctxs_bufs; diff --git a/tests/test-alloc.cpp b/tests/test-alloc.cpp index 9fe9b574719..78f2aadcb4d 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -1738,29 +1738,40 @@ static void set_test_env(const char * name, const char * value) { #endif } -static void restore_test_env(const char * name, bool had_value, const std::string & value) { +// the scheduler reads these on construction, so a test that aborts in the middle must not leave them behind for the tests after it +struct scoped_test_env { + const char * name; + bool had_value; + std::string value; + + scoped_test_env(const char * name, const char * set_to) : name(name) { + const char * env = getenv(name); + had_value = env != nullptr; + value = env ? env : ""; + set_test_env(name, set_to); + } + + ~scoped_test_env() { #ifdef _WIN32 - GGML_ASSERT(_putenv_s(name, had_value ? value.c_str() : "") == 0); + _putenv_s(name, had_value ? value.c_str() : ""); #else - GGML_ASSERT(had_value ? setenv(name, value.c_str(), 1) == 0 : unsetenv(name) == 0); + if (had_value) { + setenv(name, value.c_str(), 1); + } else { + unsetenv(name); + } #endif -} + } +}; static void test_transport_environment_is_fallback() { - const char * depth_env = getenv("GGML_KV_PIPELINE_DEPTH"); - const char * budget_env = getenv("GGML_KV_PIPELINE_BUDGET_MIB"); - const bool had_depth = depth_env != nullptr; - const bool had_budget = budget_env != nullptr; - const std::string depth_old = depth_env ? depth_env : ""; - const std::string budget_old = budget_env ? budget_env : ""; - dummy_backend cuda = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); ggml_backend_t backends[] = { cuda.handle.get(), cpu.handle.get() }; ggml_backend_buffer_type_t bufts[] = { &cuda.buffer_type, &cpu.buffer_type }; - set_test_env("GGML_KV_PIPELINE_DEPTH", "4"); - set_test_env("GGML_KV_PIPELINE_BUDGET_MIB", "8"); + scoped_test_env depth_env("GGML_KV_PIPELINE_DEPTH", "4"); + scoped_test_env budget_env("GGML_KV_PIPELINE_BUDGET_MIB", "8"); { auto graph = make_transport_graph(cpu, 64); ggml_set_stable_prefix(graph.source, 64); @@ -1792,10 +1803,11 @@ static void test_transport_environment_is_fallback() { ggml_backend_sched_ptr sched(ggml_backend_sched_new(backends, bufts, 2, 128, false, false)); GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_depth(sched.get(), 0)); GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_budget(sched.get(), 0)); - } - restore_test_env("GGML_KV_PIPELINE_DEPTH", had_depth, depth_old); - restore_test_env("GGML_KV_PIPELINE_BUDGET_MIB", had_budget, budget_old); + // a depth out of range is refused rather than clamped, so a caller cannot get a different one than it asked for + GGML_ASSERT(!ggml_backend_sched_set_transport_pipeline_depth(sched.get(), 1000)); + GGML_ASSERT(!ggml_backend_sched_set_transport_pipeline_depth(sched.get(), -1)); + } } static void test_transport_depth_zero() {