From d14138bef93719bce5638587ad9e8643a8668cc1 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Wed, 26 Aug 2026 03:08:08 +0200 Subject: [PATCH 01/50] 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 957b195cb4d5..060394ebaa8f 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2434,6 +2434,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 05af52737426..d77248d99920 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1745,6 +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; + 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 8edd79f3d053..e3a377f997e2 100644 --- a/common/common.h +++ b/common/common.h @@ -594,6 +594,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 000000000000..b01bc4ae7c08 --- /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 000000000000..27e8b0ed302a --- /dev/null +++ b/docs/repro/r4-kv-pipeline-ab.sh @@ -0,0 +1,31 @@ +#!/bin/bash +# R4: pipelined delivery of a host-resident KV cache, A/B/A/B with reversed arm order. +# The two arms are the same binary: GGML_KV_PIPELINE_DEPTH=0 is the ordered path. +# +# LLAMA_KV_MODEL=/path/model.gguf docs/repro/r4-kv-pipeline-ab.sh [depth ...] +set -u +MODEL="${LLAMA_KV_MODEL:?set LLAMA_KV_MODEL to a .gguf path}" +BUILD="${LLAMA_KV_BUILD:-build}" +PIN="${LLAMA_KV_TASKSET:-0,2,4}" +LOCK=/tmp/beellama-single-gpu.lock + +run () { # $1 label, $2 pipeline depth, $3 context depth, $4 reps + GGML_KV_PIPELINE_DEPTH=$2 taskset -c "$PIN" "$BUILD/bin/llama-bench" -m "$MODEL" \ + -ngl 99 -sm none -mg 0 -t 3 -nkvo 1 --kv-cpu-pinned --recurrent-state-offload \ + -fa on -ctk q8_0 -ctv q8_0 -b 512 -ub 512 --no-warmup -p 0 -n 128 -d "$3" -r "$4" -o json \ + 2>/dev/null \ + | python3 -c "import json,sys;d=json.load(sys.stdin);print(' %-12s %.4f +- %.4f'%('$1',d[0]['avg_ts'],d[0]['stddev_ts']))" +} + +DEPTHS=(4096 16384 32768); [ $# -gt 0 ] && DEPTHS=("$@") +rc=0 +for D in "${DEPTHS[@]}"; do + R=3; [ "$D" -le 4096 ] && R=5 + echo "== context depth=$D reps=$R" + flock "$LOCK" bash -c "$(declare -f run); BUILD='$BUILD'; MODEL='$MODEL'; PIN='$PIN' + run ordered 0 $D $R + run pipelined 1 $D $R + run ordered2 0 $D $R + run pipelined2 1 $D $R" || rc=$? +done +exit $rc diff --git a/docs/repro/r4-kv-pipeline-context-sweep.sh b/docs/repro/r4-kv-pipeline-context-sweep.sh new file mode 100755 index 000000000000..b8a2b91bf72b --- /dev/null +++ b/docs/repro/r4-kv-pipeline-context-sweep.sh @@ -0,0 +1,42 @@ +#!/bin/bash +# R4 across context depth: throughput and device allocation high-water, ordered against +# pipelined, on the same binary. The ring holds one split's whole delivery per slot, so its +# cost grows with the context; this is what measures where that stops being affordable. +# +# LLAMA_KV_MODEL=/path/model.gguf docs/repro/r4-kv-pipeline-context-sweep.sh [depth ...] +set -u +MODEL="${LLAMA_KV_MODEL:?set LLAMA_KV_MODEL to a .gguf path}" +BUILD="${LLAMA_KV_BUILD:-build}" +PIN="${LLAMA_KV_TASKSET:-0,2,4}" +NGEN="${LLAMA_KV_NGEN:-64}" +LOCK=/tmp/beellama-single-gpu.lock + +arm () { # $1 pipeline depth, $2 context depth, $3 reps + local vram; vram=$(mktemp) + ( while true; do nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits; sleep 0.25; done ) > "$vram" 2>/dev/null & + local sampler=$! + local ts + ts=$(GGML_KV_PIPELINE_DEPTH=$1 taskset -c "$PIN" "$BUILD/bin/llama-bench" -m "$MODEL" \ + -ngl 99 -sm none -mg 0 -t 3 -nkvo 1 --kv-cpu-pinned --recurrent-state-offload \ + -fa on -ctk q8_0 -ctv q8_0 -b 512 -ub 512 --no-warmup -p 0 -n "$NGEN" -d "$2" -r "$3" -o json \ + 2>/dev/null \ + | python3 -c "import json,sys +try: + d=json.load(sys.stdin); print('%.4f'%d[0]['avg_ts']) +except Exception: + print('FAILED')") + kill $sampler 2>/dev/null; wait $sampler 2>/dev/null + printf ' %-10s %-10s %s MiB\n' "depth=$1" "$ts" "$(sort -n "$vram" | tail -1)" + rm -f "$vram" +} + +DEPTHS=(4096 16384 32768 65536 131072 262144); [ $# -gt 0 ] && DEPTHS=("$@") +for D in "${DEPTHS[@]}"; do + R=3; [ "$D" -gt 32768 ] && R=1 + echo "== context depth=$D reps=$R (t/s, peak device memory)" + flock "$LOCK" bash -c "$(declare -f arm); BUILD='$BUILD'; MODEL='$MODEL'; PIN='$PIN'; NGEN='$NGEN' + arm 0 $D $R + arm 1 $D $R + arm 0 $D $R + arm 1 $D $R" +done diff --git a/docs/repro/r4-kv-pipeline-exact.py b/docs/repro/r4-kv-pipeline-exact.py new file mode 100644 index 000000000000..0a0800f3f795 --- /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 000000000000..6ab8d7b6b619 --- /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 a663d86da632..924db4eec615 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 e18f8bdbffe5..dde484f02c86 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -702,11 +702,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 3ec40fb1af7f..d60b7c86a064 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 e84bb193ab17..9f27f499f658 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 @@ -772,6 +773,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; @@ -831,6 +938,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] @@ -1588,6 +1698,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++) { @@ -1607,6 +2249,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 @@ -1626,6 +2272,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]); } @@ -1681,6 +2332,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; @@ -1702,6 +2377,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) { @@ -1831,6 +2522,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) { @@ -1870,6 +2568,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); @@ -1946,6 +2657,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); @@ -1953,10 +2670,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]); @@ -2085,6 +2911,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 ca79a95b2f70..96d9cf71ce71 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1324,6 +1324,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); @@ -1820,7 +1831,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 27c9a74846e5..b9b434efa01b 100644 --- a/include/llama.h +++ b/include/llama.h @@ -420,6 +420,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 233dfdcf790a..c57cb8ab494b 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; @@ -878,6 +880,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()); @@ -4103,6 +4110,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 27eb62d822e3..2680fdedfe6c 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 b526af616a74..66bc9ee3df5d 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1203,6 +1203,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]; @@ -1635,7 +1639,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()); @@ -2274,6 +2319,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 af765830d01a..e09d7a044eed 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -238,6 +238,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 fbe9cdd106fd5f0d54041f0b2047056ecec391b6 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Wed, 26 Aug 2026 03:14:59 +0200 Subject: [PATCH 02/50] 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 b01bc4ae7c08..73e0a3306758 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 27e8b0ed302a..9d8bb83a3930 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 b8a2b91bf72b..3cdbef21a0dd 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 9f27f499f658..9ecded9874e1 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -2721,6 +2721,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 22b865ae4f3bf75c0e95467208cf9b7ba55ff63c Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 27 Aug 2026 10:19:28 +0200 Subject: [PATCH 03/50] 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 73e0a3306758..aa1105093dea 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 9d8bb83a3930..c1636994a8fc 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 3cdbef21a0dd..0f1386f94839 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 0a0800f3f795..cfb10591a994 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 9ecded9874e1..80f8e3a73d92 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -876,6 +876,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; }; @@ -1971,20 +1981,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) { @@ -1992,7 +2011,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++) { @@ -2001,7 +2021,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 @@ -2014,10 +2035,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); @@ -2182,6 +2204,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; } @@ -2194,6 +2217,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++) { @@ -2356,6 +2380,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; @@ -2364,11 +2390,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 @@ -2395,12 +2425,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) { @@ -2503,13 +2546,19 @@ 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 + const size_t n_bytes = ranged ? rg.used*rg.n : ggml_nbytes(input); if (ranged || !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; if (ranged) { // blocking like the copy it replaces: the split backend is idle here, so the ranges go on its own stream and the host waits for them ggml_backend_tensor_set_2d_async(split_backend, input_cpy, input->data, 0, rg.used, rg.n, rg.stride, rg.stride); @@ -2517,6 +2566,13 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } else { ggml_backend_tensor_copy(input, input_cpy); } + if (tr->debug >= 2) { + tr->t_copy_us += ggml_time_us() - t1; + tr->n_bytes_ordered += n_bytes; + } + if (tr->debug >= 3 && !tr->named_ordered) { + GGML_LOG_INFO("%s: ordered copy %s %zu KiB from %s\n", __func__, input->name, + n_bytes >> 10, ggml_backend_buft_name(input->buffer->buft)); } } } @@ -2589,6 +2645,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 1fff21f701e2..d949c50b8a5e 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -357,6 +357,8 @@ struct cmd_params { std::vector lazy_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; @@ -402,6 +404,8 @@ static const cmd_params cmd_params_defaults = { /* lazy_mode */ { LLAMA_LAZY_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) }, @@ -472,6 +476,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()); @@ -841,6 +847,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; @@ -1188,6 +1208,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; } @@ -1251,6 +1277,8 @@ struct cmd_params_instance { llama_lazy_mode lazy_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; @@ -1332,6 +1360,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; @@ -1366,6 +1396,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) @@ -1396,6 +1428,8 @@ static std::vector get_cmd_params_instances(const cmd_param /* .lazy_mode = */ lzm, /* .main_gpu = */ mg, /* .no_kv_offload = */ nkvo, + /* .kv_cpu_pinned = */ kvcp, + /* .recurrent_state_offload = */ rso, /* .flash_attn = */ fa, /* .devices = */ devs, /* .tensor_split = */ ts, @@ -1433,6 +1467,8 @@ static std::vector get_cmd_params_instances(const cmd_param /* .lazy_mode = */ lzm, /* .main_gpu = */ mg, /* .no_kv_offload = */ nkvo, + /* .kv_cpu_pinned = */ kvcp, + /* .recurrent_state_offload = */ rso, /* .flash_attn = */ fa, /* .devices = */ devs, /* .tensor_split = */ ts, @@ -1470,6 +1506,8 @@ static std::vector get_cmd_params_instances(const cmd_param /* .lazy_mode = */ lzm, /* .main_gpu = */ mg, /* .no_kv_offload = */ nkvo, + /* .kv_cpu_pinned = */ kvcp, + /* .recurrent_state_offload = */ rso, /* .flash_attn = */ fa, /* .devices = */ devs, /* .tensor_split = */ ts, @@ -1512,6 +1550,8 @@ struct test { llama_lazy_mode lazy_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; @@ -1552,6 +1592,8 @@ struct test { lazy_mode = inst.lazy_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; @@ -1616,7 +1658,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", "lazy_mode", "embeddings", "no_op_offload", "no_host", "fit_target", "fit_min_ctx", @@ -1636,7 +1679,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; } @@ -1708,6 +1752,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, @@ -1904,6 +1950,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; } @@ -1929,6 +1981,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"; } @@ -2010,6 +2068,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 499c10a36855a8f26566c6d5255d00755571e1d5 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 27 Aug 2026 10:41:36 +0200 Subject: [PATCH 04/50] 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 aa1105093dea..56ae3256926c 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 fac757cb91f38288d345110a1aa182959e995058 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 27 Aug 2026 10:58:42 +0200 Subject: [PATCH 05/50] 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 56ae3256926c..9a5e1b7907fe 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 cfb10591a994..cff45a397a6d 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 cc1afb7e576b4762936f512084c3b7bbdb5c3fdd Mon Sep 17 00:00:00 2001 From: piggidragon Date: Fri, 28 Aug 2026 19:47:42 +0200 Subject: [PATCH 06/50] 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 | 373 +++++++++++++++++++++++++++-- 12 files changed, 585 insertions(+), 184 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 060394ebaa8f..386c1259dcdb 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -30,6 +30,7 @@ #include #include #include +#include #include #include #include @@ -2456,8 +2457,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 9a5e1b7907fe..2ab357323a0f 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 cff45a397a6d..55278ce35346 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 6ab8d7b6b619..e0046199c398 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 924db4eec615..f3a9516a970c 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 dde484f02c86..16507c5fc4de 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -649,11 +649,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 { @@ -702,16 +703,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 d60b7c86a064..3ec40fb1af7f 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 80f8e3a73d92..bd09167c2767 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 @@ -828,8 +829,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. @@ -845,6 +845,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]; @@ -855,7 +856,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 @@ -879,8 +879,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; @@ -1716,7 +1716,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) { @@ -1764,6 +1764,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) { @@ -1837,6 +1842,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. @@ -1886,8 +1937,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; @@ -1988,6 +2038,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; @@ -2002,11 +2053,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; } @@ -2017,23 +2075,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 " @@ -2042,23 +2104,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; } @@ -2074,7 +2125,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__, @@ -2082,17 +2133,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; } @@ -2101,16 +2142,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); @@ -2125,7 +2157,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) { @@ -2179,7 +2210,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); } @@ -2217,7 +2250,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++) { @@ -2655,7 +2688,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, @@ -2665,7 +2698,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; @@ -2676,7 +2709,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; } } @@ -2769,12 +2802,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); @@ -2782,12 +2827,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; } @@ -2795,31 +2837,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) { @@ -2843,27 +2883,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( @@ -2973,6 +3025,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 96d9cf71ce71..faaa7d769423 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1261,6 +1261,9 @@ static_assert(GGML_GLU_OP_COUNT == 7, "GGML_GLU_OP_COUNT != 7"); 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"); //////////////////////////////////////////////////////////////////////////////// @@ -1831,7 +1834,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 c57cb8ab494b..cbe151da3b06 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -876,15 +876,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 66bc9ee3df5d..dfe7efce9f14 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -314,6 +314,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 88751a2204e7..a985f985d61f 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -23,6 +23,17 @@ struct dummy_backend_context { bool unique_alloc_addresses = false; bool real_memory = false; // back the buffers with memory, so a test can look at the bytes a copy moved 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; @@ -79,8 +90,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 @@ -139,13 +150,34 @@ static void dummy_backend_buffer_clear(ggml_backend_buffer_t buffer, uint8_t val 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 *) { @@ -154,25 +186,77 @@ 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 device, ggml_backend_buffer_type_t buft) { - return device->context == buft->context; +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, bool real_memory = false) { 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->real_memory = real_memory; b.context->buffer_interface.free_buffer = dummy_backend_buffer_free_buffer; @@ -189,20 +273,38 @@ 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->context = b.context.get(); - 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.buffer_type.device = b.device.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; } @@ -226,6 +328,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); @@ -1197,8 +1329,8 @@ static bool graph_reuses_allocation(bool add_alloc_dep) { // A split input that is a window over a cache split into streams is one range per stream. // The ordered copy has to move each stream's range and leave the cells between one range and the next as it found them. static void test_ordered_multi_stream_ranges() { - dummy_backend host = dummy_backend_init(SIZE_MAX, 8, /*unique_alloc_addresses*/ false, /*real_memory*/ true); - dummy_backend dev = dummy_backend_init(SIZE_MAX, 8, /*unique_alloc_addresses*/ false, /*real_memory*/ true); + dummy_backend host = dummy_backend_init(SIZE_MAX, 8, false, GGML_BACKEND_DEVICE_TYPE_CPU, "dummy", true, /*real_memory*/ true); + dummy_backend dev = dummy_backend_init(SIZE_MAX, 8, false, GGML_BACKEND_DEVICE_TYPE_CPU, "dummy", true, /*real_memory*/ true); const int64_t n_stream = 4; const int64_t n_row = 8; // rows of a stream the graph reads @@ -1280,6 +1412,202 @@ static void test_graph_optimize_alloc_dep() { GGML_ASSERT(!graph_reuses_allocation(true)); } +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); @@ -1313,5 +1641,12 @@ int main() { run("test_resizable_buffers_owner_borrower_teardown_order", test_resizable_buffers_owner_borrower_teardown_order); run("test_ordered_multi_stream_ranges", test_ordered_multi_stream_ranges); run("test_graph_optimize_alloc_dep", test_graph_optimize_alloc_dep); + 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 2ca196eefe6807549a88ae23e2681cfb6be3ce1d Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 30 Aug 2026 23:03:52 +0200 Subject: [PATCH 07/50] 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 | 241 ++++++++++++++------- tests/test-alloc.cpp | 120 ++++++++++ tools/llama-bench/llama-bench.cpp | 41 +++- 7 files changed, 335 insertions(+), 86 deletions(-) diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index 2ab357323a0f..fd16e3669351 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 c1636994a8fc..6bb6a381e565 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 0f1386f94839..79d2f07ee301 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 e0046199c398..485a1196cb8e 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 bd09167c2767..3f6e42b1d3dd 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1794,6 +1794,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. @@ -1938,13 +1986,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; } @@ -1982,54 +2032,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 @@ -2192,30 +2277,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 @@ -2338,10 +2400,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; @@ -2390,6 +2454,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; @@ -2476,6 +2541,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 @@ -2606,6 +2672,8 @@ 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, n_bytes >> 10, ggml_backend_buft_name(input->buffer->buft)); + named_ordered_now = true; + } } } } @@ -2678,6 +2746,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++; @@ -2707,7 +2779,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; } @@ -2716,6 +2787,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, @@ -2783,6 +2890,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; @@ -2791,6 +2899,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; } @@ -2810,18 +2923,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; @@ -2896,18 +2997,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 a985f985d61f..1e8f9f147388 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -6,9 +6,11 @@ #include #include +#include #include #include #include +#include #include // @@ -1451,6 +1453,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); @@ -1642,6 +1759,9 @@ int main() { run("test_ordered_multi_stream_ranges", test_ordered_multi_stream_ranges); run("test_graph_optimize_alloc_dep", test_graph_optimize_alloc_dep); 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 d949c50b8a5e..eed97f2c37a4 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -358,6 +358,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; @@ -405,6 +406,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 */ { {} }, @@ -477,6 +479,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"); @@ -854,6 +857,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; @@ -1211,6 +1227,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; } @@ -1278,6 +1297,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; @@ -1361,6 +1381,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; @@ -1397,6 +1418,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) @@ -1429,6 +1451,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, @@ -1468,6 +1491,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, @@ -1507,6 +1531,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, @@ -1551,6 +1576,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; @@ -1593,6 +1619,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; @@ -1658,7 +1685,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", "lazy_mode", "embeddings", @@ -1674,7 +1701,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; @@ -1753,6 +1780,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), @@ -1953,6 +1981,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; } @@ -1984,6 +2015,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"; } @@ -2071,6 +2105,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 6e2fe449f84d26133794d758ae985005bfbd11f4 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Tue, 1 Sep 2026 10:21:35 +0200 Subject: [PATCH 08/50] 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 | 175 ++++++++++++++++++++- tools/llama-bench/llama-bench.cpp | 43 ++++- 7 files changed, 330 insertions(+), 58 deletions(-) diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index fd16e3669351..4bf21a3e43e1 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 6bb6a381e565..563174e56f8e 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 79d2f07ee301..cf363dfa0029 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 485a1196cb8e..d2795a5a7c15 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 3f6e42b1d3dd..1bcf42073c62 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1797,6 +1797,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; @@ -1832,10 +1839,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); @@ -2138,10 +2149,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 1e8f9f147388..79ac18a5b545 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -34,6 +34,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 = {}; @@ -92,6 +115,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; } @@ -115,7 +144,17 @@ static void * dummy_backend_buffer_get_base(ggml_backend_buffer_t buffer) { return ctx->buffer_bases[ctx->buffer_index(buffer)]; } -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; } @@ -168,18 +207,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 *) { @@ -274,6 +318,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; @@ -352,6 +397,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, @@ -1453,6 +1526,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); @@ -1759,6 +1923,7 @@ int main() { run("test_ordered_multi_stream_ranges", test_ordered_multi_stream_ranges); run("test_graph_optimize_alloc_dep", test_graph_optimize_alloc_dep); 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 eed97f2c37a4..9edea552cb80 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", @@ -1701,7 +1729,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; @@ -1781,6 +1809,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), @@ -1984,6 +2013,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; } @@ -2018,6 +2050,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"; } @@ -2108,6 +2143,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 df5f1f7e8189410eae5a2a7ae6f4fb7aa31d24d1 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 3 Sep 2026 10:12:38 +0200 Subject: [PATCH 09/50] 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 4bf21a3e43e1..467fada3e695 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 f3a9516a970c..1e44b62a4cd9 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 16507c5fc4de..83a7a7c3fa16 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -712,10 +712,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 1bcf42073c62..a53b3d996166 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -778,43 +778,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 @@ -834,13 +828,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 @@ -849,17 +839,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; @@ -876,8 +864,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; @@ -1728,11 +1716,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) { @@ -1752,11 +1738,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)) { @@ -1797,8 +1781,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); @@ -1839,8 +1822,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); @@ -1853,9 +1836,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]; @@ -1921,6 +1903,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; @@ -1947,9 +1946,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]; @@ -1990,9 +1987,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; @@ -2030,6 +2026,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) { @@ -2061,8 +2065,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); @@ -2127,11 +2131,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 }; @@ -2215,8 +2218,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) { @@ -2292,12 +2294,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]; @@ -2318,9 +2317,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; @@ -2333,9 +2331,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; @@ -2351,10 +2348,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; @@ -2380,8 +2375,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 @@ -2414,8 +2408,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)) { @@ -2467,8 +2466,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; @@ -2476,10 +2474,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; @@ -2518,11 +2515,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) { @@ -2737,16 +2732,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); } @@ -2766,8 +2759,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, " @@ -2979,8 +2971,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 b9b434efa01b..e91225dd1742 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 dfe7efce9f14..adba93716f4d 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1212,8 +1212,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 @@ -1649,9 +1648,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(); @@ -1670,8 +1668,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 e09d7a044eed..4188f55e5ab5 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 79ac18a5b545..b8f7baf9bf36 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -20,6 +20,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; @@ -36,8 +37,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; @@ -89,7 +89,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(); @@ -1526,8 +1526,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); @@ -1664,6 +1663,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); @@ -1926,6 +1965,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 2811703498aefa1130d2007a60f99ad211dc9461 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 3 Sep 2026 10:59:59 +0200 Subject: [PATCH 10/50] 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 467fada3e695..854d3444a8be 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 9a12d97eca67b719e3665aab248cb00440a90b34 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Fri, 4 Sep 2026 21:27:58 +0200 Subject: [PATCH 11/50] 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 854d3444a8be..6a9aac48798f 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 34851f64f8d20b610c662de669fae636f4149fad Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 00:30:37 +0200 Subject: [PATCH 12/50] 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 83a7a7c3fa16..bc333cf35ffd 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -703,7 +703,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 a53b3d996166..be7f63e4e9dc 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1716,24 +1716,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. @@ -2281,8 +2301,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", @@ -2333,19 +2355,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 @@ -2518,12 +2544,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 adba93716f4d..a25a0dcfe2b8 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1648,13 +1648,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 b8f7baf9bf36..b7ba7c511229 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -1616,6 +1616,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); @@ -1963,6 +2039,7 @@ int main() { run("test_graph_optimize_alloc_dep", test_graph_optimize_alloc_dep); 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 9b92b9c6b43ae5b438394b6f87140f5f9610381a Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 01:06:22 +0200 Subject: [PATCH 13/50] 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 6a9aac48798f..9f2c26a19050 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 f5971b6d038cef3a9b67c2907a06c8df0243bac7 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 21:31:11 +0200 Subject: [PATCH 14/50] 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 be7f63e4e9dc..82ade7cbab5c 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1740,8 +1740,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 b7ba7c511229..95ad9d562403 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -1635,6 +1635,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 9bcfcacf94178773416a077952bb4296fc070903 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 21:31:11 +0200 Subject: [PATCH 15/50] 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 82ade7cbab5c..2582a363d546 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -2501,6 +2501,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 9b0d2f20f6ae9bcae1eac98bb002d1fe72be607b Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 21:31:11 +0200 Subject: [PATCH 16/50] 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 2582a363d546..b9be570c7366 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1857,9 +1857,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) { @@ -1870,9 +1869,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; @@ -1883,10 +1880,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); @@ -1907,7 +1904,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; @@ -2235,7 +2232,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]; @@ -2954,7 +2951,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 f619d58a164ef79bf6f80e26610235543c7c4271 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 21:31:11 +0200 Subject: [PATCH 17/50] 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 a25a0dcfe2b8..a0a365709e67 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1208,6 +1208,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 b87c74b6e2e0b4fa16df2dcf8ea12d1986dbd400 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 21:31:11 +0200 Subject: [PATCH 18/50] 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 9edea552cb80..d0070eaffe70 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 499d2b9d808d4da19a607289003e1791cd47712b Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 22:19:49 +0200 Subject: [PATCH 19/50] 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 b9be570c7366..2b8fbc0709a7 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1857,7 +1857,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]; @@ -1869,7 +1870,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 763f706ffeb1d9f38e5a52755f51e76a2654ec7c Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 22:50:57 +0200 Subject: [PATCH 20/50] 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 2b8fbc0709a7..2f674691ba60 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -2368,10 +2368,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; } @@ -2562,11 +2559,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 feeb3320a9c16c0c90b49957eec6282ca412eef4 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 22:50:57 +0200 Subject: [PATCH 21/50] 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 bc333cf35ffd..745903062e40 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -703,8 +703,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]; @@ -713,9 +713,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 950e7fb0493a8922c20d1e8ae6c83ca1a290e194 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 22:50:57 +0200 Subject: [PATCH 22/50] 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 000000000000..779b0dae8b9c --- /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 7edce451ec6681e8578fa97244dcc8a645ff9c4e Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 22:50:57 +0200 Subject: [PATCH 23/50] 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 563174e56f8e..d0e50427ab74 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 cf363dfa0029..64bac9feba3d 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 55278ce35346..3330893ddfc8 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 d2795a5a7c15..017926d29a0c 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 95ad9d562403..2081eea0ab3c 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -1616,9 +1616,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 2ce9c7b0475c51da327378fb4bb85f755222a814 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 5 Sep 2026 23:10:32 +0200 Subject: [PATCH 24/50] 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 9f2c26a19050..3d4d0faa69c5 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 8eaa67c980a377f13e604c2a28feb182e0952bb2 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 6 Sep 2026 01:28:35 +0200 Subject: [PATCH 25/50] 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 3d4d0faa69c5..9b17d572516b 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 2f674691ba60..376e2d6adcb4 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -823,6 +823,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; }; @@ -1880,6 +1882,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; @@ -1907,6 +1910,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; @@ -1970,6 +1983,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; @@ -2026,6 +2050,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; } @@ -2036,6 +2061,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; @@ -2055,6 +2081,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; } @@ -2062,6 +2089,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; @@ -2086,6 +2114,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; } @@ -2204,7 +2233,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; } @@ -2260,10 +2291,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; @@ -2271,11 +2313,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); } } @@ -2379,6 +2421,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; } } @@ -2431,6 +2474,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); @@ -2505,11 +2549,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. @@ -3202,6 +3247,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 c7b2e8e978fbff3c3205299c25d0a95461c100aa Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 6 Sep 2026 01:29:03 +0200 Subject: [PATCH 26/50] 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 386c1259dcdb..3364c0584beb 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2455,10 +2455,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 cbe151da3b06..332fabb0e2c7 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -880,8 +880,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 d0070eaffe70..7cfcf6e9fee7 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 56e8076eef4fd635e1d194b89683d21319884445 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 6 Sep 2026 01:29:13 +0200 Subject: [PATCH 27/50] 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 e91225dd1742..72db8f0c11c8 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 332fabb0e2c7..c0ae1223eaf9 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -3930,6 +3930,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 7b9fdb2115b609cac0788da7d6711c24885f9d11 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 6 Sep 2026 10:42:01 +0200 Subject: [PATCH 28/50] 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 9b17d572516b..8bee834178c5 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 72db8f0c11c8..1873f48dabd5 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 c0ae1223eaf9..332fabb0e2c7 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -3930,9 +3930,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 a0a365709e67..6630d89027e5 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -477,6 +477,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); } @@ -855,6 +860,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 4188f55e5ab5..8517b37e6cba 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 543f34be1a97..647d46431df2 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 c23ed2bb4ed7..38bbe22d633d 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 ed6bc86430afdcae136b4050d686e2b68b76c983 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 6 Sep 2026 13:22:12 +0200 Subject: [PATCH 29/50] 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 8bee834178c5..99050eb95cf0 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 1e44b62a4cd9..a813a0f9b45c 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 376e2d6adcb4..de7b0d6aed4d 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1792,6 +1792,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; } @@ -1910,12 +1920,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); } } } @@ -2233,9 +2244,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; } @@ -2263,8 +2274,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; } @@ -2304,10 +2319,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); @@ -2548,6 +2566,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; @@ -2988,7 +3007,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 2081eea0ab3c..1190a7273c79 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -32,6 +32,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; @@ -243,6 +244,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; } @@ -1780,6 +1782,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); @@ -1940,10 +1982,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); @@ -2044,10 +2118,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 3f41a7563fac24e61f0f79b59fb21228ba53f3cd Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 6 Sep 2026 14:26:54 +0200 Subject: [PATCH 30/50] 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 d77248d99920..4662daad8f46 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 de7b0d6aed4d..c762ec66f911 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -2571,7 +2571,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; } @@ -3048,6 +3050,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 6630d89027e5..8efeb9fa3f2f 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -857,11 +857,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 8517b37e6cba..f12dda86ed5d 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 c354b2213e4c..e4b37ab357d7 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -134,6 +134,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 bbdd08f0c77c87c879ea1a40d8d80ac6aa86865f Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 6 Sep 2026 21:49:02 +0200 Subject: [PATCH 31/50] 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 3364c0584beb..d5563d1d0706 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2443,8 +2443,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; } @@ -2455,12 +2455,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 4662daad8f46..b171ad9b29d7 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 745903062e40..b56aca1dde9a 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -715,7 +715,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 c762ec66f911..d9a1a2bde09c 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -15,7 +15,6 @@ #include #include -#include #include #include #include @@ -798,6 +797,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 { @@ -822,6 +827,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 @@ -845,9 +851,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; @@ -918,6 +929,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; @@ -1177,6 +1189,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 = { @@ -1890,9 +1903,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; @@ -1920,21 +1934,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; @@ -1951,6 +1977,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. @@ -1960,8 +1994,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; @@ -1997,12 +2030,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. @@ -2053,7 +2092,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; @@ -2108,7 +2147,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++) { @@ -2133,10 +2172,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; } @@ -2147,7 +2204,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 { @@ -2161,8 +2218,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; } } @@ -2170,8 +2227,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; } } @@ -2185,15 +2242,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 @@ -2244,9 +2299,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; } @@ -2278,8 +2333,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; } @@ -2314,7 +2368,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); @@ -2324,8 +2378,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); @@ -2339,7 +2392,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) { @@ -2556,19 +2610,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) { @@ -3144,6 +3194,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 faaa7d769423..35462cae204b 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1329,8 +1329,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 1873f48dabd5..3a10d2f3369f 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 332fabb0e2c7..d65ef18b8d82 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; @@ -877,15 +883,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 8efeb9fa3f2f..2e7eb12355fc 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1699,7 +1699,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 1190a7273c79..dd7972ead506 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -1618,6 +1618,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() { @@ -1782,8 +1831,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); @@ -1801,15 +1850,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()); @@ -2114,11 +2174,12 @@ int main() { run("test_graph_optimize_alloc_dep", test_graph_optimize_alloc_dep); 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 7cfcf6e9fee7..19c34b78fb57 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 b0d9c4924efb105f1c57e36977c66130565cb5d6 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 6 Sep 2026 23:40:56 +0200 Subject: [PATCH 32/50] 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 99050eb95cf0..c27e8e9763c0 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 a813a0f9b45c..96fbfd02f144 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 b56aca1dde9a..d8de91e81501 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -704,7 +704,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]; @@ -713,10 +712,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 d9a1a2bde09c..007afe0b57db 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -778,33 +778,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 @@ -812,8 +809,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 @@ -826,7 +822,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 @@ -835,10 +831,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 @@ -849,18 +843,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; @@ -878,7 +875,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; @@ -1732,8 +1728,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 @@ -1741,8 +1737,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; @@ -1751,12 +1747,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))) { @@ -1773,10 +1764,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)) { @@ -1805,8 +1793,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); @@ -1827,7 +1815,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); @@ -1868,8 +1856,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); @@ -1882,8 +1869,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]; @@ -1935,9 +1921,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]; @@ -1955,7 +1940,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; @@ -1978,23 +1963,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; @@ -2028,9 +2018,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); @@ -2044,7 +2033,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]; @@ -2086,7 +2075,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; @@ -2107,15 +2096,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; } @@ -2168,8 +2160,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) { @@ -2252,8 +2243,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 }; @@ -2292,6 +2282,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]) { @@ -2329,7 +2336,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])); @@ -2342,7 +2349,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) { @@ -2360,7 +2367,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); @@ -2373,8 +2380,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])); @@ -2382,6 +2388,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; @@ -2435,8 +2443,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]; @@ -2445,20 +2453,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; @@ -2471,8 +2479,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) { @@ -2489,12 +2496,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; } } @@ -2517,7 +2523,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 @@ -2551,13 +2557,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)) { @@ -2612,11 +2618,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; @@ -2628,8 +2632,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++) { @@ -2669,9 +2672,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) { @@ -2886,14 +2888,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); } @@ -3082,9 +3084,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); @@ -3100,7 +3101,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; @@ -3138,8 +3139,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; } @@ -3192,6 +3192,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 35462cae204b..f21f8cd67379 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1329,9 +1329,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 d65ef18b8d82..20680d89ee4e 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -887,6 +887,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 2e7eb12355fc..78813733b4a3 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -263,6 +263,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; } } @@ -864,7 +872,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(); @@ -1220,7 +1228,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; } @@ -1661,9 +1669,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 f12dda86ed5d..39b6eab4c1eb 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 647d46431df2..543f34be1a97 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 38bbe22d633d..c23ed2bb4ed7 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 dd7972ead506..6db699fe35bc 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -1890,29 +1890,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); @@ -1944,10 +1955,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() { From f79205bfc010cd34e5b439374ae1827baaeee2df Mon Sep 17 00:00:00 2001 From: piggidragon Date: Tue, 8 Sep 2026 23:17:56 +0200 Subject: [PATCH 33/50] sched, llama : pack the multi-stream transport ring, make pipelining opt-in A staged copy kept its source's stride, so a cache split into streams reserved a whole kv_size per stream and the ring grew with n_stream rather than with the window. At the default budget it then declined its own ring: 170 MiB needed for a 42.7 MiB window, and the arm measured the ordered path. The copy packs the ranges now, and the delivery is given the source and the destination stride separately. At -npl 8 -no-kvu -c 32768 that is 70.90 -> 106.19 t/s, and unchanged where the budget was already raised past what the gaps cost. --kv-pipeline-depth defaults to 0. Any value above 0 turns it on, and 1 is the value that measures best. Also from review: - do not compute the ring layout inside GGML_ASSERT - arm a slot's release event even when its split fails, so freeing the ring waits for a late copy that is already on the consumer stream - clear input_staged when the plan allocation fails, so it cannot disagree with split_order - reset reported_no_room once a ring is allocated again - exclude CPU devices where the device type is checked, not after the registry name, where it could never be taken - zero the whole trailing storage of ggml_tensor again on a 32-bit target - fix an unsigned underflow in the test allocator's capacity check Assisted-by: Claude Opus 5 --- common/arg.cpp | 9 +-- common/common.h | 2 +- docs/kv-transport-pipelining.md | 36 +++++++++-- ggml/include/ggml.h | 3 +- ggml/src/ggml-backend.cpp | 104 +++++++++++++++++++++--------- include/llama.h | 4 +- src/llama-context.cpp | 2 +- tests/test-alloc.cpp | 22 +++++-- tools/llama-bench/llama-bench.cpp | 2 +- 9 files changed, 130 insertions(+), 54 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index d5563d1d0706..0d7979d16152 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2438,10 +2438,11 @@ common_params_context common_params_parser_init(common_params & params, llama_ex 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), + "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; any other value turns the " + "pipeline on, and 1 is the value that measures best. 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 > 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)); diff --git a/common/common.h b/common/common.h index e3a377f997e2..bc888e94b36c 100644 --- a/common/common.h +++ b/common/common.h @@ -594,7 +594,7 @@ 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_depth = 0; // 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 diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index c27e8e9763c0..cd70c07e6c6e 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -4,7 +4,17 @@ With `--no-kv-offload` (optionally with `--kv-cpu-pinned`), the attention histor 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, and it is **off by default**. Any `N` above 0 turns it on, `N = 1` is the value that measures best, and `0` is the ordered path. It only has an effect where a host-resident cache produces the deliveries. + +## Turning it on + +``` +--no-kv-offload --kv-cpu-pinned --kv-pipeline-depth 1 +``` + +`--kv-pipeline-depth` defaults to 0, so nothing about an existing run changes until it is set. The same value is `llama_context_params::kv_pipeline_depth`, `-kvpd` in `llama-bench`, and `LLAMA_ARG_KV_PIPELINE_DEPTH` in the environment; `GGML_KV_PIPELINE_DEPTH` sets the scheduler default under all of them. + +`N = 1` is the recommendation, not just the smallest value that works: every deeper look-ahead measured is slower, and the ring costs `(N + 2)` slots of device memory. `N = 0` is the default because the feature spends device memory on a cache that is on the host to save it, and because the measurements below are one model on one link. Turn it on, look at `GGML_SCHED_TRANSPORT_DEBUG=1`, and keep it if the deliveries convert. 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). @@ -103,7 +113,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: 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. +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 value to turn it on with, 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 @@ -185,6 +195,18 @@ A cache split into streams delivers a window per stream, so it moves more than a **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. +**A slot holds the window, not the cache it is cut from.** A staged copy keeps its source's layout, and in a cache split into streams that layout steps a whole `kv_size` from one stream to the next while the graph reads only `n_kv` of it. Sizing the slot from that stride reserves every gap the delivery skips, so the ring grew with `n_stream` instead of with the window. The copy now packs the streams: one window padded to the ring's alignment, with `cudaMemcpy2DAsync` given the source stride and the packed stride separately. Nothing about the bytes delivered changes, only where they land. + +The ring is what the budget is applied to, so this decides whether there is a ring at all. `llama-batched-bench -npp 2048 -ntg 128 -npl 8` at `-c 32768 -no-kvu`, one process per cell, generation t/s: + +| budget | depth | ring slot | before | after | +|---:|---:|---:|---:|---:| +| 128 (default) | 0 | - | 71.09 | 70.99 | +| 128 (default) | 1 | 8 MiB, then declined / 42.7 MiB | 70.90 | **106.19** | +| 512 | 1 | 68 MiB / 64 MiB | 106.20 | 106.11 | + +At the default budget the unpacked ring needs 170 MiB for a window worth 42.7 MiB, so it is declined and the arm reads the ordered path's own throughput. Packed, the same run keeps the ring and gains **+49.6%** over the ordered path. Where the budget was already raised past what the gaps cost, both are the same speed and the packed ring is slightly smaller. A unified cache is one range per delivery and is not affected either way, which the single-sequence A/B confirms: 18.976 -> 29.779 packed against 18.988 -> 29.784 before, at 16,384. + **Concurrent slots can be gated on output, 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 @@ -251,6 +273,8 @@ 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). +The wait at the graph boundary does not show up here because it costs nothing to show. Timed separately at 16,384 over 128 graphs of steady-state decode, the consumer sync and the transfer sync are both 0.00 ms of a 31.55 ms graph: the host reaches the boundary well after the streams have. It is 1.39 ms per graph during prefill, where the window widens on nearly every ubatch, and nothing after that. Ordering the two streams with an event instead of blocking the host was measured against this and has nothing to win. + 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. @@ -261,6 +285,8 @@ The gates, and what was run for them: 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. +Gate 5 and the A/B were re-run for the packed multi-stream ring, on both this head and the commit before it, and the two agree byte for byte: `9c13743e07b55934` at `N = 0`, `1` and `4` with `-sm none`, `7fc64d5ed9709861` at the same depths with `-sm layer`. Packing changes where a range lands in the slot and nothing about the bytes, so an unchanged hash is the result to expect; the gate is there because a stride mistake would not look like one. + 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. @@ -275,11 +301,11 @@ 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. +- 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 window the graph reads. 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 multi-stream window is delivered one range per stream**, keyed on the last dimension, and the copy packs those ranges so a slot holds the window rather than the whole cache. 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. +- **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 exception is a graph that cannot be allocated next to the rings: there every device that was holding one gives it back for good, because the allocator does not say which of them it competed with. - **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). diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index d8de91e81501..fe31b61f153c 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -704,9 +704,10 @@ 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 + // padding comes first so that ggml_new_tensor_impl zeroes the whole storage, whatever size_t is union { + char padding[8]; size_t stable_prefix; - char padding[8]; }; }; diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 007afe0b57db..e976cd9461a5 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1729,17 +1729,19 @@ 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, 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. +// A window over several streams is one range per stream: the streams sit a fixed stride apart in the source, and the cells between one stream's window and the next are never read by this graph, so the copy packs the ranges. 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 + int64_t n; // ranges to deliver + size_t stride; // bytes from one range to the next in the host source + size_t stride_cpy; // the same in the copy, which packs the ranges the source holds apart + 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 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) { +static void ggml_backend_sched_input_ranges(const struct ggml_tensor * input, const struct ggml_tensor * input_cpy, + struct ggml_backend_sched_ranges * out) { const struct ggml_tensor * base = input->view_src ? input->view_src : input; out->n = 1; @@ -1760,7 +1762,8 @@ static void ggml_backend_sched_input_ranges(const struct ggml_tensor * input, st out->used = rows; } - out->early = base->stable_prefix < out->used ? base->stable_prefix : out->used; + out->early = base->stable_prefix < out->used ? base->stable_prefix : out->used; + out->stride_cpy = out->n > 1 && input_cpy ? input_cpy->nb[3] : out->stride; } // Whether a split input belongs in its backend's ring. @@ -1821,6 +1824,25 @@ static bool ggml_backend_sched_transport_entry_size( return ggml_backend_sched_size_pad(ggml_backend_buft_get_alloc_size(buft, t), alignment, result); } +// The ranges of a multi-stream window sit a whole cache apart in the host source, and the cells between them are never delivered. +// The copy packs them, so a slot holds the window rather than the cache; 0 means keep the source layout. +static size_t ggml_backend_sched_transport_packed_stride(const struct ggml_tensor * input, size_t alignment) { + struct ggml_backend_sched_ranges rg; + ggml_backend_sched_input_ranges(input, NULL, &rg); + + size_t stride; + if (rg.n < 2 || !ggml_backend_sched_size_pad(rg.used, alignment, &stride) || stride >= rg.stride) { + return 0; + } + return stride; +} + +// A copy in the ring is laid out for the ring, and one that is not is laid out like its source: the ordered path copies it whole. +static void ggml_backend_sched_transport_set_layout(struct ggml_tensor * input_cpy, const struct ggml_tensor * input, size_t alignment) { + const size_t stride = ggml_backend_sched_transport_packed_stride(input, alignment); + input_cpy->nb[3] = stride ? stride : input->nb[3]; +} + static void ggml_backend_sched_transport_clear_addresses(ggml_backend_sched_t sched) { const struct ggml_backend_sched_transport * tr = &sched->transport; @@ -1833,6 +1855,7 @@ static void ggml_backend_sched_transport_clear_addresses(ggml_backend_sched_t sc struct ggml_tensor * input_cpy = tensor_copy(split->inputs[j], split->backend_id, sched->cur_copy); input_cpy->data = NULL; input_cpy->buffer = NULL; + input_cpy->nb[3] = split->inputs[j]->nb[3]; } } } @@ -1859,11 +1882,13 @@ static void ggml_backend_sched_transport_assign_addresses(ggml_backend_sched_t s // 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; + ggml_backend_sched_transport_set_layout(input_cpy, split->inputs[j], r->alignment); 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_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)); + size_t input_size = 0; + bool ok = ggml_backend_sched_transport_entry_size(ggml_backend_buffer_get_type(r->buffer), input_cpy, r->alignment, &input_size); + ok = ok && ggml_backend_sched_size_add(offset, input_size, &offset); + GGML_ASSERT(ok); } GGML_ASSERT(offset <= r->slot_size); } @@ -1952,12 +1977,18 @@ static void ggml_backend_sched_transport_decline_backend(ggml_backend_sched_t sc } for (int i = 0; i < tr->plan_n_splits; i++) { - if (sched->splits[i].backend_id != backend_id) { + struct ggml_backend_sched_split * split = &sched->splits[i]; + if (split->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; + for (int j = 0; j < split->n_inputs; j++) { + unsigned char * staged = &tr->input_staged[tr->split_input_ofs[i] + j]; + if (*staged) { + struct ggml_tensor * input_cpy = tensor_copy(split->inputs[j], backend_id, sched->cur_copy); + input_cpy->nb[3] = split->inputs[j]->nb[3]; + } + *staged = 0; } } } @@ -2169,6 +2200,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { 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__); + memset(tr->input_staged, 0, n_inputs_total); for (int i = 0; i < sched->n_splits; i++) { tr->split_order[i] = -1; } @@ -2261,7 +2293,8 @@ 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); + struct ggml_tensor * entry = tensor_copy(split->inputs[j], bid, sched->cur_copy); + ggml_backend_sched_transport_set_layout(entry, input, tr->rings[bid].alignment); size_t input_size; size_t input_size_max; if (!ggml_backend_sched_transport_entry_size(sched->bufts[bid], entry, tr->rings[bid].alignment, &input_size) || @@ -2400,8 +2433,9 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { } } - r->idle_graphs = 0; - tr->n_staged += r->n_staged; + r->idle_graphs = 0; + r->reported_no_room = false; + tr->n_staged += r->n_staged; } if (tr->n_staged == 0) { @@ -2428,7 +2462,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { ggml_backend_buffer_t buf = input->view_src ? input->view_src->buffer : input->buffer; src_buft = ggml_backend_buft_name(buf->buft); struct ggml_backend_sched_ranges rg; - ggml_backend_sched_input_ranges(input, &rg); + ggml_backend_sched_input_ranges(input, NULL, &rg); total += rg.used*rg.n; early += rg.early*rg.n; } @@ -2480,16 +2514,16 @@ 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 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_tensor * input_cpy = tensor_copy(input, split->backend_id, sched->cur_copy); struct ggml_backend_sched_ranges rg; - ggml_backend_sched_input_ranges(input, &rg); + ggml_backend_sched_input_ranges(input, input_cpy, &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_2d_async(r->transfer, input_cpy, input->data, 0, rg.early, rg.n, rg.stride, rg.stride); + ggml_backend_tensor_set_2d_async(r->transfer, input_cpy, input->data, 0, rg.early, rg.n, rg.stride_cpy, rg.stride); if (tr->debug >= 2) { tr->t_issue_us += ggml_time_us() - t0; } @@ -2675,10 +2709,10 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s // 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); + ggml_backend_sched_input_ranges(input, input_cpy, &rg); if (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); + rg.early, rg.used - rg.early, rg.n, rg.stride_cpy, rg.stride); tr->n_bytes_late += (rg.used - rg.early)*rg.n; } continue; @@ -2849,11 +2883,9 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s ggml_backend_event_wait(split_backend, slot->ready); } + enum ggml_status ec = GGML_STATUS_SUCCESS; if (!sched->callback_eval) { - enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &split->graph); - if (ec != GGML_STATUS_SUCCESS) { - return ec; - } + ec = ggml_backend_graph_compute_async(split_backend, &split->graph); } else { // similar to ggml_backend_compare_graph_backend for (int j0 = 0; j0 < split->graph.n_nodes; j0++) { @@ -2872,9 +2904,9 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s struct ggml_cgraph gv = ggml_graph_view(&split->graph, j0, j1 + 1); - enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &gv); + ec = ggml_backend_graph_compute_async(split_backend, &gv); if (ec != GGML_STATUS_SUCCESS) { - return ec; + break; } // TODO: pass backend to the callback, then the user can decide if they want to synchronize @@ -2889,12 +2921,19 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } // every kernel that reads this slot is enqueued, so it may be refilled once the consumer stream reaches this point + // armed even when the split failed: the late copy above is already on this stream, and freeing the ring waits on armed slots only if (slot != NULL) { ggml_backend_event_record(slot->release, split_backend); slot->release_armed = true; tr->rings[split_backend_id].consumed++; tr->n_deliveries++; + } + if (ec != GGML_STATUS_SUCCESS) { + return ec; + } + + if (slot != NULL) { // 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); } @@ -3117,7 +3156,11 @@ bool ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, 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) { + if (dev == NULL) { + continue; + } + const enum ggml_backend_dev_type type = ggml_backend_dev_type(dev); + if (type == GGML_BACKEND_DEVICE_TYPE_META || type == GGML_BACKEND_DEVICE_TYPE_CPU) { continue; } @@ -3135,9 +3178,6 @@ bool ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, if (dev->iface.event_new == NULL) { continue; } - if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) { - continue; - } // 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)) { diff --git a/include/llama.h b/include/llama.h index 3a10d2f3369f..9452bc6c2a52 100644 --- a/include/llama.h +++ b/include/llama.h @@ -427,8 +427,8 @@ extern "C" { struct llama_context * ctx_other; 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 + // 0 is the default and keeps the ordered path, where a decode token pays the transfer and the attention kernels in series + // any other value turns the pipeline on; 1 measures best, and 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 }; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 20680d89ee4e..13f950912d81 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -4117,7 +4117,7 @@ llama_context_params llama_context_default_params() { /*.sampler =*/ nullptr, /*.n_sampler =*/ 0, /*.ctx_other =*/ nullptr, - /*.kv_pipeline_depth =*/ 1, + /*.kv_pipeline_depth =*/ 0, /*.kv_pipeline_budget_mib =*/ 128, }; diff --git a/tests/test-alloc.cpp b/tests/test-alloc.cpp index 6db699fe35bc..24d795f26df2 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -90,7 +90,8 @@ 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 || size > ctx->capacity - ctx->allocated_total()) { + const size_t allocated = ctx->allocated_total(); + if (ctx->fail_alloc || size > ctx->capacity || allocated > ctx->capacity - size) { return nullptr; } ggml_backend_buffer_t & buffer = ctx->buffers.emplace_back(); @@ -1717,12 +1718,18 @@ static void test_transport_multi_stream_ranges() { }); // 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]; + // the copy packs the streams, so it steps by one window padded to the ring's alignment rather than by a whole cache + const size_t stride = (size_t) window->nb[3]; + const size_t align = 128; + const size_t stride_cpy = (used_bytes + align - 1)/align*align; + GGML_ASSERT(stride_cpy < stride); + 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); + const size_t src_offset = (size_t) (d.src - (const char *) window->data); + GGML_ASSERT(src_offset/stride < (size_t) n_stream); + GGML_ASSERT(src_offset%stride + d.size <= used_bytes); + GGML_ASSERT(d.offset == (src_offset/stride)*stride_cpy + src_offset%stride); total += d.size; } GGML_ASSERT(total == (size_t) n_stream*used_bytes); @@ -1730,8 +1737,9 @@ static void test_transport_multi_stream_ranges() { 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); + const size_t src_offset = (size_t) (d.src - (const char *) window->data); + if (src_offset/stride == (size_t) st) { + GGML_ASSERT(src_offset%stride == covered); covered += d.size; } } diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 19c34b78fb57..0cdbb9657595 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -407,7 +407,7 @@ static const cmd_params cmd_params_defaults = { /* main_gpu */ { 0 }, /* no_kv_offload */ { false }, /* kv_cpu_pinned */ { false }, - /* kv_pipeline_depth */ { 1 }, + /* kv_pipeline_depth */ { 0 }, /* kv_pipeline_budget_mib */ { 128 }, /* recurrent_state_offload */ { false }, /* flash_attn */ { LLAMA_FLASH_ATTN_TYPE_AUTO }, From d359a23e3fc7cc0aaeec825904993e54d727f8ca Mon Sep 17 00:00:00 2001 From: piggidragon Date: Wed, 9 Sep 2026 00:59:10 +0200 Subject: [PATCH 34/50] sched, tests : reconcile the ring with the merged ordered range copy The ordered copy of a multi-stream window landed in llama/dev on its own. It brings its own copy of the range geometry, which this branch already had in a form that also carries the copy's stride and the stable prefix, so drop the duplicate and keep the one the ring uses. The ordered copy now steps by the copy's own stride. A copy is laid out for the ring only while it is staged, so the two agree today, but reading it from the copy is what keeps them agreeing. The dummy backend's set_tensor_async records a delivery and now performs it as well when the backend is asked for real memory, so the test that checks the bytes of the ordered copy still sees them. ggml_new_tensor_impl needs a brace for the array inside the union. Assisted-by: Claude Opus 5 --- ggml/src/ggml-backend.cpp | 32 ++------------------------------ ggml/src/ggml.c | 2 +- tests/test-alloc.cpp | 3 +++ 3 files changed, 6 insertions(+), 31 deletions(-) diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index e976cd9461a5..77c5580ae9a0 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -2609,34 +2609,6 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { return true; } -// How a split input breaks into the ranges a graph reads. -// A window over a cache split into streams is one range per stream, keyed on the last dimension: the ranges sit a fixed stride apart and the bytes between them are never read. -// Anything else is one flat range of ggml_nbytes(). -struct ggml_backend_sched_ranges { - int64_t n; // ranges to deliver - size_t stride; // bytes from one range to the next - size_t used; // bytes of a range this graph reads -}; - -static void ggml_backend_sched_input_ranges(const struct ggml_tensor * input, struct ggml_backend_sched_ranges * out) { - out->n = 1; - out->stride = 0; - out->used = ggml_nbytes(input); - - // a range is one stream's byte span, which is what the tensor covers below dimension 3 - 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; - } - - if (input->ne[3] > 1) { - out->n = input->ne[3]; - out->stride = input->nb[3]; - out->used = rows; - } -} - static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t sched) { GGML_ASSERT(sched); struct ggml_backend_sched_split * splits = sched->splits; @@ -2837,7 +2809,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s // ggml_backend_tensor_copy moves ggml_nbytes(), which for a window over several streams is the span the ranges are cut from, gaps and all ggml_backend_buffer_t src_buf = input->view_src ? input->view_src->buffer : input->buffer; struct ggml_backend_sched_ranges rg; - ggml_backend_sched_input_ranges(input, &rg); + ggml_backend_sched_input_ranges(input, input_cpy, &rg); const bool ranged = rg.n > 1 && src_buf != NULL && ggml_backend_buffer_is_host(src_buf); // 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 @@ -2857,7 +2829,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s const int64_t t1 = tr->debug >= 2 ? ggml_time_us() : 0; if (ranged) { // blocking like the copy it replaces: the split backend is idle here, so the ranges go on its own stream and the host waits for them - ggml_backend_tensor_set_2d_async(split_backend, input_cpy, input->data, 0, rg.used, rg.n, rg.stride, rg.stride); + ggml_backend_tensor_set_2d_async(split_backend, input_cpy, input->data, 0, rg.used, rg.n, rg.stride_cpy, rg.stride); ggml_backend_synchronize(split_backend); } else { ggml_backend_tensor_copy(input, input_cpy); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index f21f8cd67379..573a70b3274b 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1836,7 +1836,7 @@ static struct ggml_tensor * ggml_new_tensor_impl( /*.data =*/ obj_alloc_size > 0 ? (void *)(result + 1) : data, /*.name =*/ { 0 }, /*.extra =*/ NULL, - /*.padding =*/ { 0 }, + /*.padding =*/ { { 0 } }, }; // TODO: this should not be needed as long as we don't rely on aligned SIMD loads diff --git a/tests/test-alloc.cpp b/tests/test-alloc.cpp index 24d795f26df2..e8783de19c8f 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -214,6 +214,9 @@ static void dummy_backend_set_tensor_async(ggml_backend_t backend, ggml_tensor * ctx->set_tensor_async_count++; ctx->set_tensor_async_bytes += size; ctx->deliveries.push_back({ tensor, (const char *) data, offset, size }); + if (ctx->real_memory) { + memcpy((char *) tensor->data + offset, data, size); + } } static void dummy_backend_synchronize(ggml_backend_t) {} From da23680de54eb4335bcf36edd79dc7d31f8a0d9b Mon Sep 17 00:00:00 2001 From: piggidragon Date: Wed, 9 Sep 2026 11:52:40 +0200 Subject: [PATCH 35/50] docs : re-measure the parallel and packing tables on the rebased base The ordered copy of a multi-stream window landed in llama/dev, so the baseline the parallel table compares against moves. Only the ordered column of a cache split into streams changes; the rest is within noise. The pipeline is worth +8.5%, +13.8%, +19.8% and +26.0% at 1, 2, 4 and 8 slots over a non-unified cache, against +72% and +73% read off the old baseline, and packing the ring is worth +25.7% at the default budget rather than +49.6%. The single-sequence table does not move at all: a unified window is one range and has no gaps to have been wasting. Gate 5 is re-stated against a build of llama/dev rather than an earlier head of this branch, which is the stronger comparison. Assisted-by: Claude Opus 5 --- docs/kv-transport-pipelining.md | 26 ++++++++++++++------------ 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index cd70c07e6c6e..c274fa6580af 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -186,26 +186,28 @@ Four device-resident layers are worth +14.3% on the ordered path and +1.4% on th | `-npl` | unified ordered | unified pipelined | streams ordered | streams pipelined | |---:|---:|---:|---:|---:| -| 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 | +| 1 | 32.62 | 35.38 | 32.66 | 35.44 | +| 2 | 51.59 | 58.76 | 51.48 | 58.61 | +| 4 | 71.89 | 86.29 | 71.45 | 85.60 | +| 8 | 83.44 | 105.76 | 84.16 | 106.02 | -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. The two caches now measure the same in both arms, so 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. +**The streams-ordered column used to be far lower**, and that was a defect in the ordered copy rather than a property of a cache split into streams. `ggml_backend_tensor_copy` moves `ggml_nbytes()`, which for a window over several streams is the span the ranges are cut from, so the ordered arm copied every gap between them: 34.09, 49.44 and 70.92 at 2, 4 and 8 slots against the numbers above. Fixed separately, before this branch, so the gain the pipeline is credited with here is what it is worth against a baseline that moves the right bytes. + +**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.61, 85.60 and 106.02 here; the difference is the cost of copying the right amount. **A slot holds the window, not the cache it is cut from.** A staged copy keeps its source's layout, and in a cache split into streams that layout steps a whole `kv_size` from one stream to the next while the graph reads only `n_kv` of it. Sizing the slot from that stride reserves every gap the delivery skips, so the ring grew with `n_stream` instead of with the window. The copy now packs the streams: one window padded to the ring's alignment, with `cudaMemcpy2DAsync` given the source stride and the packed stride separately. Nothing about the bytes delivered changes, only where they land. The ring is what the budget is applied to, so this decides whether there is a ring at all. `llama-batched-bench -npp 2048 -ntg 128 -npl 8` at `-c 32768 -no-kvu`, one process per cell, generation t/s: -| budget | depth | ring slot | before | after | +| budget | depth | ring slot | unpacked | packed | |---:|---:|---:|---:|---:| -| 128 (default) | 0 | - | 71.09 | 70.99 | -| 128 (default) | 1 | 8 MiB, then declined / 42.7 MiB | 70.90 | **106.19** | -| 512 | 1 | 68 MiB / 64 MiB | 106.20 | 106.11 | +| 128 (default) | 0 | - | 84.23 | 84.23 | +| 128 (default) | 1 | 8 MiB, then declined / 42.7 MiB | 84.45 | **106.19** | +| 512 | 1 | 68 MiB / 64 MiB | 106.28 | 106.11 | -At the default budget the unpacked ring needs 170 MiB for a window worth 42.7 MiB, so it is declined and the arm reads the ordered path's own throughput. Packed, the same run keeps the ring and gains **+49.6%** over the ordered path. Where the budget was already raised past what the gaps cost, both are the same speed and the packed ring is slightly smaller. A unified cache is one range per delivery and is not affected either way, which the single-sequence A/B confirms: 18.976 -> 29.779 packed against 18.988 -> 29.784 before, at 16,384. +At the default budget the unpacked ring needs 170 MiB for a window worth 42.7 MiB, so it is declined and that arm reads the ordered path's own throughput. Packed, the same run keeps the ring and gains **+25.7%** over the ordered path. Where the budget was already raised past what the gaps cost, both are the same speed and the packed ring is slightly smaller. A unified cache is one range per delivery and is not affected either way, which the single-sequence A/B confirms: 18.96 -> 29.70 packed against 18.99 -> 29.78 before, at 16,384. **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. @@ -285,7 +287,7 @@ The gates, and what was run for them: 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. -Gate 5 and the A/B were re-run for the packed multi-stream ring, on both this head and the commit before it, and the two agree byte for byte: `9c13743e07b55934` at `N = 0`, `1` and `4` with `-sm none`, `7fc64d5ed9709861` at the same depths with `-sm layer`. Packing changes where a range lands in the slot and nothing about the bytes, so an unchanged hash is the result to expect; the gate is there because a stride mistake would not look like one. +Gate 5 was re-run after the ordered range copy landed in `llama/dev` and this branch was rebased onto it, against a build of `llama/dev` itself rather than against an earlier head of this branch: `17f946c340db110b` with `-sm none` and `db661b7a08686b97` with `-sm layer` over both devices, the same on `llama/dev` and at `N = 0`, `1` and `4` here. Packing changes where a range lands in the slot and nothing about the bytes, so an unchanged hash is the result to expect; the gate is there because a stride mistake would not look like one. 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. From 47347c0109379b94dfa23a61c30d11eac0fd7eb1 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 10 Sep 2026 18:11:07 +0200 Subject: [PATCH 36/50] sched, llama : fix the fourth review of the pipelined transport Also gives the stable prefix one value per stream, so a slot that was just reset no longer caps the early region of every other stream. Assisted-by: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JycLdWs6KRgizbdnNAfZqM --- docs/kv-transport-pipelining.md | 12 +- ggml/include/ggml-backend.h | 16 +-- ggml/include/ggml.h | 22 ++-- ggml/src/ggml-backend.cpp | 175 ++++++++++++++++++++++-------- ggml/src/ggml.c | 9 +- include/llama.h | 4 +- src/llama-context.cpp | 3 + src/llama-kv-cache.cpp | 66 ++++++----- src/llama-kv-cache.h | 6 +- src/llama-memory-hybrid-idx.cpp | 5 - tests/test-alloc.cpp | 73 ++++++++----- tools/llama-bench/llama-bench.cpp | 2 +- 12 files changed, 259 insertions(+), 134 deletions(-) diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index c274fa6580af..dae92d143c2c 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -12,7 +12,7 @@ The transfer and the attention arithmetic are the same as before, but the transf --no-kv-offload --kv-cpu-pinned --kv-pipeline-depth 1 ``` -`--kv-pipeline-depth` defaults to 0, so nothing about an existing run changes until it is set. The same value is `llama_context_params::kv_pipeline_depth`, `-kvpd` in `llama-bench`, and `LLAMA_ARG_KV_PIPELINE_DEPTH` in the environment; `GGML_KV_PIPELINE_DEPTH` sets the scheduler default under all of them. +`--kv-pipeline-depth` defaults to 0, so nothing about an existing run changes until it is set. The same value is `llama_context_params::kv_pipeline_depth`, `-kvpd` in `llama-bench`, and `LLAMA_ARG_KV_PIPELINE_DEPTH` in the environment. `N = 1` is the recommendation, not just the smallest value that works: every deeper look-ahead measured is slower, and the ring costs `(N + 2)` slots of device memory. `N = 0` is the default because the feature spends device memory on a cache that is on the host to save it, and because the measurements below are one model on one link. Turn it on, look at `GGML_SCHED_TRANSPORT_DEBUG=1`, and keep it if the deliveries convert. @@ -28,7 +28,9 @@ Three pieces make it work. 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. +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: one value per stream, held in an array the caller keeps. A view inherits the values of the streams its own byte window covers. `llama_kv_cache::update_stable_prefixes()` fills 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. + +Per stream, because the streams of one ubatch sit at very different depths: a slot that was just reset writes at row 12 while the others are at 20,000, and a single value for the body would hold every stream down to that 12. A stream the ubatch does not write at all keeps everything it holds. With seven deep sequences against one freshly reset one, the per-stream prefix moves the split from 42 MiB early / 490 MiB late per graph to 430 MiB early / 92 MiB late, and generation over the eight goes from 2.44 to 2.87 tokens/s. That one is measured on the second card of the same host (RTX 3060, sm_86), eight server slots over a non-unified cache at `-c 32768`, seven prompts of about 3k tokens against one of eight. 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. @@ -305,15 +307,15 @@ A device-resident KV run is unaffected, and was measured to confirm it: 38.5612 - 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 window the graph reads. 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, and the copy packs those ranges so a slot holds the window rather than the whole cache. A window whose streams are not on that dimension keeps the single flat range, which is correct but not accelerated. +- **A multi-stream window is delivered one range per stream**, keyed on the last dimension, and the copy packs those ranges so a slot holds the window rather than the whole cache. A window whose streams are not on that dimension keeps the single flat range, which is correct but not accelerated. Each range carries its own stream's prefix; ranges that agree on it are issued as one strided copy, so streams at the same depth still cost a single call. - **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 exception is a graph that cannot be allocated next to the rings: there every device that was holding one gives it back for good, because the allocator does not say which of them it competed with. - **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. +- **A host write to the cache waits for the delivery.** `llama_memory_clear(mem, true)` waits for the scheduler 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. +- `GGML_KV_PIPELINE_DEPTH` and `GGML_KV_PIPELINE_BUDGET_MIB` set the defaults of a scheduler that nothing else configures. `llama_context` always configures its own from the context parameters, so under `llama-server` and `llama-bench` use `--kv-pipeline-depth` / `LLAMA_ARG_KV_PIPELINE_DEPTH` and `--kv-pipeline-budget` / `LLAMA_ARG_KV_PIPELINE_BUDGET` instead. ## Tensor parallelism diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index 96fbfd02f144..1b836a78678a 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -325,22 +325,22 @@ 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 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. + // Without it a split that reads a host input pays copy + compute in series: the copy goes on the consumer's stream right before the kernels. + // With it the scheduler keeps a ring of staging slots out of ggml-alloc's reach, and sends a later split's stable prefix on a transfer stream. // - // Only persistent host inputs marked with GGML_TENSOR_FLAG_TRANSPORT are eligible, and their stable prefix must be current before each evaluation. + // Only persistent host inputs marked 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, 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. + // `depth` is how many splits ahead deliveries run, 0 disables it, at most GGML_SCHED_MAX_TRANSPORT_SLOTS - GGML_SCHED_TRANSPORT_MARGIN (14). + // The ring holds a couple of slots more, so recycling a slot never waits for a reader that still runs. + // Only CUDA is accepted as the destination: the ring needs a second context that transfers async and orders with events. Others stay ordered. // Costs roughly (depth + 2) * (largest staged split) of device memory. // 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. + // The ring is optional: if a graph does not fit next to it, the scheduler releases the rings that held 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. - // 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. + // A host cache exists to keep device memory free, so the ring has a hard cap: past it 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); diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index fe31b61f153c..c01f08e86445 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -703,22 +703,24 @@ 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 - // padding comes first so that ggml_new_tensor_impl zeroes the whole storage, whatever size_t is + // bytes at the start of each stream that stay unchanged for the current graph evaluation, one entry per stream, NULL for none + // padding comes first so that ggml_new_tensor_impl zeroes the whole storage, whatever a pointer is union { - char padding[8]; - size_t stable_prefix; + char padding[8]; + const size_t * stable_prefix; }; }; 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 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); + // declare that the first nbytes[s] of stream s 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 a value counts from the start of its stream + // the caller keeps the array: one entry per stream of the storage, valid and current for every graph that runs, including a reused one + // a reader whose streams do not tile the storage, or that holds less than one entry says, declares nothing for it rather than guessing + // set it on the tensor that owns the storage, not on a view of it + // a backend may deliver a declared region before the point in the graph that reads it, so NULL and zeroes are always correct + GGML_API void ggml_set_stable_prefix(struct ggml_tensor * tensor, const size_t * nbytes); + GGML_API const size_t * ggml_get_stable_prefix(const struct ggml_tensor * tensor); // Abort callback // If not NULL, called before ggml computation diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 77c5580ae9a0..05d9843ea471 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -778,7 +778,7 @@ static bool ggml_is_view_op(enum ggml_op op) { #endif // How many slots the transport ring keeps behind the look-ahead. -// 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. +// With no margin a delivery recycles the slot of the split that is still running, which is the ordered path with extra steps. #ifndef GGML_SCHED_TRANSPORT_MARGIN #define GGML_SCHED_TRANSPORT_MARGIN 2 #endif @@ -789,7 +789,7 @@ static bool ggml_is_view_op(enum ggml_op op) { #endif // Default cap on the ring itself. -// 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. +// A slot holds one layer's K or V over the whole context, so an uncapped ring spends back the device memory a host cache saves. #ifndef GGML_SCHED_TRANSPORT_BUDGET #define GGML_SCHED_TRANSPORT_BUDGET (128u*1024*1024) #endif @@ -809,7 +809,7 @@ struct ggml_backend_sched_transport_slot { }; // One ring per accelerator the scheduler drives. -// 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. +// A layer-split model gives every device its own splits, so one device must not take another's slots or disable it. struct ggml_backend_sched_transport_ring { bool eligible; // this backend can transfer asynchronously and order with events @@ -831,8 +831,9 @@ struct ggml_backend_sched_transport_ring { }; // Pipelined delivery of host-resident split inputs. -// 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. +// The ordered path copies a split on the consumer's stream right before the kernels read it, so a token pays copy + compute in series. +// This ring sends the stable part of a later split on a separate transfer stream instead. +// The scheduler owns the ring, not ggml-alloc: ggml-alloc frees a copy after its last consumer, while a look-ahead transfer still reads it. 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 @@ -841,7 +842,7 @@ 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 + // delivery order of a split in its own backend's ring, indexed by split id, -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 @@ -1729,25 +1730,45 @@ 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, 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 the cells between one stream's window and the next are never read by this graph, so the copy packs the ranges. +// A window over several streams is one range per stream. The source holds them a stride apart; the copy packs them, as the gaps are never read. struct ggml_backend_sched_ranges { int64_t n; // ranges to deliver size_t stride; // bytes from one range to the next in the host source size_t stride_cpy; // the same in the copy, which packs the ranges the source holds apart 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 + size_t stream; // bytes from one stream of the storage to the next, 0 when the prefix does not apply + int64_t stream0; // stream of the storage the first range sits on + const size_t * prefix; // per-stream stable prefix of the storage, NULL for none }; +// Leading bytes of range i that may go before the split that reads it: its own stream's prefix, bounded by what this graph reads. +static size_t ggml_backend_sched_range_early(const struct ggml_backend_sched_ranges * rg, int64_t i) { + if (rg->prefix == NULL) { + return 0; + } + + // a value past one stream was never per-stream, so it says nothing about this window + const size_t prefix = rg->prefix[rg->stream0 + i]; + if (prefix > rg->stream) { + return 0; + } + + return prefix < rg->used ? prefix : rg->used; +} + // 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. +// The prefix counts from the start of a stream, so it applies only when the view says where the streams are; else the whole window goes late. static void ggml_backend_sched_input_ranges(const struct ggml_tensor * input, const struct ggml_tensor * input_cpy, 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; + out->n = 1; + out->stride = 0; + out->stride_cpy = 0; + out->used = ggml_nbytes(input); + out->stream = 0; + out->stream0 = 0; + out->prefix = NULL; // a range is one stream's byte span, which is what the tensor covers below dimension 3 const size_t rows = ggml_nbytes(input) - (size_t) (input->ne[3] - 1)*input->nb[3]; @@ -1762,12 +1783,22 @@ static void ggml_backend_sched_input_ranges(const struct ggml_tensor * input, co out->used = rows; } - out->early = base->stable_prefix < out->used ? base->stable_prefix : out->used; out->stride_cpy = out->n > 1 && input_cpy ? input_cpy->nb[3] : out->stride; + + // A stream is what the last dimension of this view steps over, and the guard above already put every range on one of them. + // Take the prefixes only when those streams tile the storage, because the array is indexed by the storage's own streams. + const size_t stream = input->nb[3]; + if (stream == 0 || ggml_nbytes(base) % stream != 0) { + return; + } + + out->stream = stream; + out->stream0 = (int64_t) (offs/stream); + out->prefix = ggml_get_stable_prefix(base); } // Whether a split input belongs in its backend's ring. -// 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. +// Independent of the stable prefix: this decides where an input copy lives, which the allocator must know before there is any ubatch. 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)) { @@ -1818,7 +1849,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 ggml_nbytes(): a buffer type may ask for padding that its kernels write into. +// A ring entry costs what the backend allocates for it, which can be more than ggml_nbytes(): a buffer type may add padding its kernels write. 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); @@ -1894,7 +1925,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 they can already be gone on the teardown path. +// Wait through the slots' own release events, never through sched->backends[]: the scheduler does not own those, and they can already be gone. 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]; @@ -2022,6 +2053,25 @@ static bool ggml_backend_sched_transport_decline_all(ggml_backend_sched_t sched) return released; } +// The same for the paths that have no plan to decline: the split list belongs to a graph the rings were never laid out over. +// Returns whether any ring was holding memory. +static bool ggml_backend_sched_transport_disable_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++) { + if (tr->rings[i].buffer == NULL) { + continue; + } + released = true; + ggml_backend_sched_transport_release_ring(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; @@ -2049,7 +2099,7 @@ 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 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. +// The window widens on nearly every prefill ubatch, so grow in powers of two: a few allocations per prompt instead of one 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); @@ -2275,7 +2325,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 at the full context; it is reported rather than enforced, or every large -c would be refused a ring it never grows into + // slot_size_max is what this ring costs at the full context; reported, not enforced, or a large -c is 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 }; @@ -2413,7 +2463,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 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 + // the headroom check passed, so the device is out of memory for reasons this cannot see; a retry per graph costs a 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])); @@ -2464,7 +2514,9 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { struct ggml_backend_sched_ranges rg; ggml_backend_sched_input_ranges(input, NULL, &rg); total += rg.used*rg.n; - early += rg.early*rg.n; + for (int64_t k = 0; k < rg.n; k++) { + early += ggml_backend_sched_range_early(&rg, k); + } } } GGML_LOG_INFO("%s: %s: %d/%d splits staged, %zu KiB per graph, %zu KiB of it early, source %s\n", @@ -2476,8 +2528,40 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { ggml_backend_sched_transport_assign_addresses(sched); } +// Deliver the early or the late part of every range, one call per group of ranges that carry the same number of bytes. +// Streams at the same depth stay a single strided copy; only streams that differ cost a call of their own. +// Returns the bytes issued. +static size_t ggml_backend_sched_transport_deliver( + ggml_backend_t backend, struct ggml_tensor * input_cpy, const struct ggml_tensor * input, + const struct ggml_backend_sched_ranges * rg, bool late) { + size_t n_bytes = 0; + + for (int64_t i = 0; i < rg->n; ) { + const size_t early = ggml_backend_sched_range_early(rg, i); + + int64_t n = 1; + while (i + n < rg->n && ggml_backend_sched_range_early(rg, i + n) == early) { + n++; + } + + const size_t offset = late ? early : 0; + const size_t size = late ? rg->used - early : early; + + if (size > 0) { + ggml_backend_tensor_set_2d_async(backend, input_cpy, + (const char *) input->data + i*rg->stride + offset, + i*rg->stride_cpy + offset, size, n, rg->stride_cpy, rg->stride); + n_bytes += size*n; + } + + i += n; + } + + return n_bytes; +} + // 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 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. +// The margin keeps the recycled slot several readers behind the split just enqueued, so refilling it does not 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; @@ -2513,24 +2597,24 @@ 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 belongs to the ubatch about to run, not to the plan, so it can be less than when the ring was laid out + // the stable part 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_tensor * input_cpy = tensor_copy(input, split->backend_id, sched->cur_copy); struct ggml_backend_sched_ranges rg; ggml_backend_sched_input_ranges(input, input_cpy, &rg); - if (rg.early == 0) { + if (rg.prefix == NULL) { continue; } GGML_ASSERT(input->data != NULL && input_cpy->data != NULL); const int64_t t0 = tr->debug >= 2 ? ggml_time_us() : 0; - ggml_backend_tensor_set_2d_async(r->transfer, input_cpy, input->data, 0, rg.early, rg.n, rg.stride_cpy, rg.stride); + const size_t n_early = ggml_backend_sched_transport_deliver(r->transfer, input_cpy, input, &rg, false); if (tr->debug >= 2) { tr->t_issue_us += ggml_time_us() - t0; } - tr->n_bytes_early += rg.early*rg.n; + tr->n_bytes_early += n_early; } - // 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 + // record the handover now, not when the split runs: the stream is FIFO, so a later event makes the consumer wait for the look-ahead behind it ggml_backend_event_record(slot->ready, r->transfer); r->delivered = true; @@ -2624,21 +2708,27 @@ 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, 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. + const int64_t t_graph_0 = tr->debug >= 2 ? ggml_time_us() : 0; + + // A staged delivery reads its host source long after the call returned, so the previous graph can still read where this ubatch writes. + // Waiting here is what the ordered path gets from its blocking copy, once per graph. It is also why depth and n_copies > 1 do not mix. + // A ring about to stage waits even if it delivered nothing last graph: the prefetch below reads a source the last graph may still write. for (int i = 0; i < sched->n_backends; i++) { if (!tr->rings[i].delivered && !(staged && tr->rings[i].n_staged > 0)) { continue; } + const int64_t t0 = tr->debug >= 2 ? ggml_time_us() : 0; if (tr->rings[i].transfer) { ggml_backend_synchronize(tr->rings[i].transfer); } ggml_backend_synchronize(sched->backends[i]); + if (tr->debug >= 2) { + tr->t_sync_us += ggml_time_us() - t0; + } tr->rings[i].delivered = false; } - // 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. + // Prime every ring before the first consumer runs. After this a delivery follows an enqueued split, so recycling a slot holds back nothing. // 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++) { @@ -2650,8 +2740,6 @@ 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; @@ -2678,15 +2766,11 @@ 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)) { - // 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 + // the stable prefix went out earlier; an earlier split of this graph may still write the rest, so it is safe to read only now + // it goes on the consumer's stream, ordered ahead of the kernels and behind the reader of the slot's last occupant struct ggml_backend_sched_ranges rg; ggml_backend_sched_input_ranges(input, input_cpy, &rg); - if (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_cpy, rg.stride); - tr->n_bytes_late += (rg.used - rg.early)*rg.n; - } + tr->n_bytes_late += ggml_backend_sched_transport_deliver(split_backend, input_cpy, input, &rg, true); continue; } @@ -2806,7 +2890,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } copy_experts(first_id, last_id); } else { - // ggml_backend_tensor_copy moves ggml_nbytes(), which for a window over several streams is the span the ranges are cut from, gaps and all + // ggml_backend_tensor_copy moves ggml_nbytes(), which for a multi-stream window is the whole span, gaps and all ggml_backend_buffer_t src_buf = input->view_src ? input->view_src->buffer : input->buffer; struct ggml_backend_sched_ranges rg; ggml_backend_sched_input_ranges(input, input_cpy, &rg); @@ -2828,7 +2912,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } const int64_t t1 = tr->debug >= 2 ? ggml_time_us() : 0; if (ranged) { - // blocking like the copy it replaces: the split backend is idle here, so the ranges go on its own stream and the host waits for them + // blocking like the copy it replaces: the backend is idle here, so the ranges go on its stream and the host waits ggml_backend_tensor_set_2d_async(split_backend, input_cpy, input->data, 0, rg.used, rg.n, rg.stride_cpy, rg.stride); ggml_backend_synchronize(split_backend); } else { @@ -3287,7 +3371,12 @@ bool ggml_backend_sched_reserve(ggml_backend_sched_t sched, struct ggml_cgraph * ggml_backend_sched_split_graph(sched, measure_graph); if (!ggml_gallocr_reserve_n(sched->galloc, &sched->graph, sched->node_backend_ids, sched->leaf_backend_ids)) { - return false; + // the rings hold device memory this graph needs, and the caller can no longer turn them off + if (!ggml_backend_sched_transport_disable_all(sched) || + !ggml_gallocr_reserve_n(sched->galloc, &sched->graph, sched->node_backend_ids, sched->leaf_backend_ids)) { + return false; + } + 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_reset(sched); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 573a70b3274b..77e10a3cb8ba 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1327,15 +1327,12 @@ 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) { +void ggml_set_stable_prefix(struct ggml_tensor * tensor, const size_t * nbytes) { GGML_ASSERT(tensor); - // 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; + tensor->stable_prefix = nbytes; } -size_t ggml_get_stable_prefix(const struct ggml_tensor * tensor) { +const size_t * ggml_get_stable_prefix(const struct ggml_tensor * tensor) { GGML_ASSERT(tensor); return tensor->stable_prefix; } diff --git a/include/llama.h b/include/llama.h index 9452bc6c2a52..c3533d7dc922 100644 --- a/include/llama.h +++ b/include/llama.h @@ -427,8 +427,8 @@ extern "C" { struct llama_context * ctx_other; 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 is the default and keeps the ordered path, where a decode token pays the transfer and the attention kernels in series - // any other value turns the pipeline on; 1 measures best, and costs (kv_pipeline_depth + 2) * (largest staged split) of device memory + // 0 is the default and keeps the ordered path, where a token pays the transfer and the kernels in series + // any other value turns it on; 1 measures best, and costs (kv_pipeline_depth + 2) * (largest staged split) 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 }; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 13f950912d81..19b97b017a0c 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -153,6 +153,9 @@ llama_context::llama_context( 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)); } + if (cparams.kv_pipeline_depth > 0 && !cparams.kv_cpu_pinned && cparams.offload_kqv) { + LLAMA_LOG_WARN("%s: kv_pipeline_depth needs a host-resident KV cache, staying on the ordered path (see --no-kv-offload, --kv-cpu-pinned)\n", __func__); + } cparams.kv_gpu_layers = params.kv_gpu_layers; cparams.phase_aware_workspace = params.phase_aware_workspace; cparams.live_context_workspace = params.live_context_workspace; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 78813733b4a3..183a96255b6a 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -353,7 +353,18 @@ llama_kv_cache::llama_kv_cache( map_layer_ids[il] = layers.size(); - layers.push_back({ il, k, v, k_store_quantize, v_store_quantize, k_stream, v_stream, }); + layers.push_back({ il, k, v, k_store_quantize, v_store_quantize, k_stream, v_stream, + std::vector(n_stream, 0), std::vector(n_stream, 0), }); + } + + // the layer list is final here, so the arrays keep the address ggml is given + for (auto & layer : layers) { + if (layer.k && (layer.k->flags & GGML_TENSOR_FLAG_TRANSPORT)) { + ggml_set_stable_prefix(layer.k, layer.k_stable.data()); + } + if (layer.v && (layer.v->flags & GGML_TENSOR_FLAG_TRANSPORT)) { + ggml_set_stable_prefix(layer.v, layer.v_stable.data()); + } } if (!offload && placement.gpu_resident_layers > 0) { @@ -486,8 +497,9 @@ 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); + // the scheduler alone, because llama_synchronize would also fold the running decode into the perf counters + if (lctx && lctx->get_sched()) { + ggml_backend_sched_synchronize(lctx->get_sched()); } for (auto & [_, buf] : ctxs_bufs) { @@ -1228,12 +1240,12 @@ 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 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 + // a cache that shares cells sets no stable prefix: the layers it aliases lost the transport flag, and the rest are the owner's to describe if (other) { return; } - // 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 + // before the graph is built, so the plan and the deliveries it issues see the same write position update_stable_prefixes(sinfo); // keep track of the max sequence position that we would overwrite with this ubatch @@ -1669,39 +1681,36 @@ 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 { - // 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; + // the lowest row this ubatch writes in a stream: everything below it keeps what the previous ubatch left there for the whole graph + // per stream: a stream this ubatch does not write keeps all of it, and one row for the whole body would cap every stream at the lowest of them + uint64_t min_row[LLAMA_MAX_SEQ]; + for (uint32_t s = 0; s < n_stream; ++s) { + min_row[s] = get_size(); + } + for (uint32_t s = 0; s < sinfo.n_stream(); ++s) { for (const uint32_t idx : sinfo.idxs[s]) { - min_row = std::min(min_row, (uint64_t) idx); + min_row[sinfo.strm[s]] = std::min(min_row[sinfo.strm[s]], (uint64_t) 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]); + for (uint32_t s = 0; s < n_stream; ++s) { + if (layer.k) { + layer.k_stable[s] = min_row[s]*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 + layer.v_stable[s] = v_trans ? 0 : min_row[s]*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); - } + std::fill(layer.k_stable.begin(), layer.k_stable.end(), 0); + std::fill(layer.v_stable.begin(), layer.v_stable.end(), 0); } } @@ -2482,6 +2491,11 @@ const slot_info_vec_t * sinfos_in) { GGML_UNUSED(flags); + // the writes below go to the same buffers a decode can still be delivering to the device, like clear() above + if (lctx && lctx->get_sched()) { + ggml_backend_sched_synchronize(lctx->get_sched()); + } + // TODO: fix incosistent handling of `seq_id < 0` and `seq_id == -1` in the codebase [TAG_LLAMA_SEQ_ID_NEG] GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size())); diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 39b6eab4c1eb..3fb654820645 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -241,7 +241,7 @@ class llama_kv_cache : public llama_memory_i { bool is_reserve, const slot_info * sinfo = nullptr) const; - // 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 + // tell the scheduler what this ubatch does not write, so a host cache can go out 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; @@ -285,6 +285,10 @@ class llama_kv_cache : public llama_memory_i { std::vector k_stream; std::vector v_stream; + + // stable prefix of each stream, in bytes; ggml keeps the address of these, so they live as long as the cache + mutable std::vector k_stable; + mutable std::vector v_stable; }; bool v_trans = true; // the value tensor is transposed diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index e4b37ab357d7..c354b2213e4c 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -134,11 +134,6 @@ 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); } diff --git a/tests/test-alloc.cpp b/tests/test-alloc.cpp index e8783de19c8f..4b8e5e69556b 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -431,6 +432,15 @@ static transport_graph_pair make_transport_graph_pair(dummy_backend & cpu, size_ return { std::move(result), std::move(buffer), { s0, s1 }, output }; } +// ggml keeps the address of a prefix array, so the test owns one per tensor for as long as it runs +static std::map> g_stable_prefix; + +static void set_stable_prefix(ggml_tensor * tensor, size_t nbytes, int64_t n_stream = 1) { + std::vector & prefix = g_stable_prefix[tensor]; + prefix.assign((size_t) n_stream, nbytes); + ggml_set_stable_prefix(tensor, prefix.data()); +} + static void transport_stats( ggml_backend_sched_t sched, int64_t * deliveries, @@ -1506,7 +1516,7 @@ static void test_transport_prefix_and_configuration() { 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); + 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); @@ -1518,12 +1528,12 @@ static void test_transport_prefix_and_configuration() { 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); + 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); + 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); @@ -1551,8 +1561,8 @@ static void test_transport_entry_allocation() { 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); + set_stable_prefix(graph.sources[0], nbytes/2); + 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); @@ -1622,7 +1632,7 @@ 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. +// Slot k starts at k*slot_size, so a capped slot size must stay a multiple of the alignment, or every slot after the first is misaligned. 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); @@ -1641,7 +1651,7 @@ static void test_transport_slot_alignment() { 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); + 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 @@ -1671,7 +1681,7 @@ static void test_transport_slot_alignment() { 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. +// A window over several streams sits a stride apart in one tensor, with cells between the streams 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); @@ -1708,7 +1718,7 @@ static void test_transport_multi_stream_ranges() { // 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); + set_stable_prefix(store, used_bytes/2, n_stream); 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); @@ -1753,6 +1763,15 @@ static void test_transport_multi_stream_ranges() { 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)); + + // a prefix larger than one stream was never a per-stream value, so the whole window goes late instead of coming from the wrong rows + set_stable_prefix(store, ggml_nbytes(store), n_stream); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), ctx.graph) == GGML_STATUS_SUCCESS); + + int64_t early_bad = 0, late_bad = 0; + transport_stats(sched.get(), NULL, &early_bad, &late_bad); + GGML_ASSERT(early_bad == early); + GGML_ASSERT(late_bad == late + (int64_t) ((size_t) n_stream*used_bytes)); } static void test_transport_empty_graph() { @@ -1777,7 +1796,7 @@ static size_t transport_fallback_buffer_size(int depth) { 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)); + 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))); @@ -1808,7 +1827,7 @@ static size_t transport_scale_buffer_size(int depth, size_t nbytes, size_t capac cuda.context->capacity = capacity; auto graph = make_transport_graph(cpu, nbytes); - ggml_set_stable_prefix(graph.source, nbytes); + 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 }; @@ -1842,7 +1861,7 @@ static void test_transport_releases_ring_for_graph() { GGML_ASSERT(transfers == 0); } -// 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 +// a graph that stages nothing keeps the ring for a few graphs and the transfer context for good: a context shift must not rebuild either 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); @@ -1850,8 +1869,8 @@ static void test_transport_keeps_ring_over_idle_graph() { 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); + set_stable_prefix(first.source, 64); + 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 }; @@ -1872,7 +1891,7 @@ static void test_transport_keeps_ring_over_idle_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 + // one idle graph keeps the ring: freeing and allocating it again blocks the host, which powers-of-two growth exists to avoid run_idle(); GGML_ASSERT(cuda.context->buffers.size() == n_staged_buffers); @@ -1937,7 +1956,7 @@ static void test_transport_environment_is_fallback() { 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); + 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)); @@ -1948,7 +1967,7 @@ static void test_transport_environment_is_fallback() { } { auto graph = make_transport_graph(cpu, 64); - ggml_set_stable_prefix(graph.source, 64); + 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)); @@ -1977,7 +1996,7 @@ 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); + 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 }; @@ -2000,8 +2019,8 @@ static void test_transport_budget_recovers() { 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); + set_stable_prefix(large.source, 256); + 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 }; @@ -2038,8 +2057,8 @@ static void test_transport_partial_backend_failure() { 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); + set_stable_prefix(source_fail, 64); + 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)); @@ -2077,8 +2096,8 @@ static void test_transport_stops_after_backend_failure() { 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); + set_stable_prefix(first.source, 64); + 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 }; @@ -2105,7 +2124,7 @@ 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); + 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 }; @@ -2126,7 +2145,7 @@ static void test_transport_requires_annotation() { 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); + 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 }; @@ -2146,7 +2165,7 @@ 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); + 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 }; diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 0cdbb9657595..3390a8f5b625 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...%d> (default: %s)\n", LLAMA_KV_PIPELINE_DEPTH_MAX, 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()); From 205cb8d0bc51d3904b4b8daedc2d083ea53744b9 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 3 Sep 2026 11:21:53 +0200 Subject: [PATCH 37/50] ggml : keep the source name of a scheduler copy A copy of an input is named "##" in a name field of fixed size. With many devices the backend label of the meta backend lists all of them and the source name is what gets cut, so a consumer can no longer tell which tensor the copy was made from. Cut the label instead. Assisted-by: Claude Opus 5 --- ggml/src/ggml-backend.cpp | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 8e7c9b08b329..fd9788ef025f 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1034,6 +1034,19 @@ static void ggml_backend_sched_print_assignments(ggml_backend_sched_t sched, str } } +// Name a copy of an input "##". The name field has a fixed size, so cut the +// backend label rather than the source name - the source name is what identifies the copy. +static void ggml_backend_sched_name_copy( + struct ggml_tensor * copy, const char * backend_name, const struct ggml_tensor * src, int c) { + const int n_tail = snprintf(NULL, 0, "#%s#%d", src->name, c); + const int n_max = GGML_MAX_NAME - 1 - n_tail; + int n_head = (int) strlen(backend_name); + if (n_head > n_max) { + n_head = n_max > 0 ? n_max : 0; + } + ggml_format_name(copy, "%.*s#%s#%d", n_head, backend_name, src->name, c); +} + static bool ggml_backend_sched_buffer_supported(ggml_backend_sched_t sched, struct ggml_tensor * t, int backend_id) { ggml_backend_buffer_t buf = t->view_src ? t->view_src->buffer : t->buffer; ggml_backend_buffer_type_t buft = NULL; @@ -1381,7 +1394,7 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra tensor_copy = src; // use the original tensor as the current copy } else { tensor_copy = ggml_dup_tensor_layout(sched->ctx, src); - ggml_format_name(tensor_copy, "%s#%s#%d", ggml_backend_name(backend), src->name, c); + ggml_backend_sched_name_copy(tensor_copy, ggml_backend_name(backend), src, c); } ggml_set_input(tensor_copy); ggml_set_output(tensor_copy); // prevent ggml-alloc from overwriting the tensor @@ -1402,7 +1415,7 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra ggml_backend_t backend = sched->backends[cur_backend_id]; for (int c = 0; c < sched->n_copies; c++) { struct ggml_tensor * tensor_copy = ggml_dup_tensor_layout(sched->ctx, src); - ggml_format_name(tensor_copy, "%s#%s#%d", ggml_backend_name(backend), src->name, c); + ggml_backend_sched_name_copy(tensor_copy, ggml_backend_name(backend), src, c); if (sched->n_copies > 1) { ggml_set_input(tensor_copy); ggml_set_output(tensor_copy); // prevent ggml-alloc from overwriting the tensor From e1c5abef8b3586af189f2b5560f75657406c26b8 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 3 Sep 2026 11:21:53 +0200 Subject: [PATCH 38/50] ggml-meta : split a host-resident KV cache by head A KV cache in host memory reaches attention as a scheduler copy, which is a leaf in the compute buffer. Such a leaf never reached the device split-state callback and fell through to MIRRORED, while the queries stayed split by head: each device then attended heads whose keys live on the other device. With more than one KV head that aborts in the FlashAttention kernel, or returns wrong output where the query split happens to stay a multiple of the KV head count. Offer a copied-in leaf to the callback under the name the graph gave it, and let the callback recognise the cache there. The cache folds its heads into one flat axis but the copy arrives permuted, with the heads on an axis of their own, so its segments and granularity are rescaled to whole heads. Assisted-by: Claude Opus 5 --- ggml/src/ggml-backend-meta.cpp | 69 +++++++++++++++++++++++++++++++++- src/llama-model.cpp | 63 +++++++++++++++++++++++++------ 2 files changed, 118 insertions(+), 14 deletions(-) diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 3ec40fb1af7f..ec0bde1c9275 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -485,6 +485,28 @@ static struct ggml_tensor * ggml_backend_meta_buffer_simple_tensor(const struct return it->second[index]; } +// A scheduler copy is named "##", where also carries the suffixes +// ggml appends for views. Recover the graph name the copy was made from. +static std::string ggml_backend_meta_copy_source_name(const char * name) { + std::string ret = name; + const size_t first = ret.find('#'); + if (first == std::string::npos) { + return ret; + } + ret.erase(0, first + 1); + // ggml writes a view suffix as " (...)", a graph name has no spaces + const size_t suffix = ret.find(" ("); + if (suffix != std::string::npos) { + ret.erase(suffix); + return ret; + } + const size_t copy = ret.rfind('#'); + if (copy != std::string::npos && ret.find_first_not_of("0123456789", copy + 1) == std::string::npos) { + ret.erase(copy); + } + return ret; +} + static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(const struct ggml_tensor * tensor, bool assume_sync); static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( @@ -831,10 +853,26 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( if (ggml_nelements(tensor) == 0) { return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; } - if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_COMPUTE && tensor->view_src == nullptr) { + // A host-resident KV cache reaches the graph as a copied-in leaf of the compute buffer. + // Mirroring it while the queries stay split by head makes each device attend the wrong + // heads, so ask the callback; it still answers MIRRORED for names it does not know. + const bool copied_in_leaf = + ggml_backend_buffer_get_usage(tensor->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && + tensor->op == GGML_OP_NONE && (tensor->flags & GGML_TENSOR_FLAG_INPUT) == 0; + + if ((ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_COMPUTE || copied_in_leaf) && + tensor->view_src == nullptr) { ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer)); const ggml_backend_meta_device_context * dev_ctx = (const ggml_backend_meta_device_context *) dev->context; - ggml_backend_meta_split_state ret = dev_ctx->get_split_state(tensor, dev_ctx->get_split_state_ud); + // the callback classifies by name, so offer a copy under the name it was copied from + const ggml_tensor * tensor_query = tensor; + ggml_tensor tensor_named; + if (copied_in_leaf) { + tensor_named = *tensor; + ggml_set_name(&tensor_named, ggml_backend_meta_copy_source_name(tensor->name).c_str()); + tensor_query = &tensor_named; + } + ggml_backend_meta_split_state ret = dev_ctx->get_split_state(tensor_query, dev_ctx->get_split_state_ud); if (ret.axis >= 0 && ret.axis < GGML_MAX_DIMS) { const int64_t granularity = ret.axis == GGML_BACKEND_SPLIT_AXIS_0 ? ggml_blck_size(tensor->type) : 1; int64_t ne_sum = 0; @@ -1416,6 +1454,33 @@ static void ggml_backend_meta_buffer_memset_tensor( static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer); const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); + + // A host-resident attention cache reaches this permuted, as [head_dim, n_kv, n_head_kv, 1]. The + // heads split, but interleaved per cell, so the chunk splice below cannot express the write. + // Each device's heads are one contiguous run per cell: ne[1] cells, from one stride to another. + const bool strided_head_split = + !ggml_is_contiguous(tensor) && + split_state.axis == GGML_BACKEND_SPLIT_AXIS_2 && + split_state.n_segments == 1 && split_state.nr[0] == 1 && + tensor->ne[3] == 1 && tensor->nb[1] > tensor->nb[2] && + offset == 0 && size == ggml_nbytes(tensor); + + if (strided_head_split) { + size_t offset_data = 0; + for (size_t j = 0; j < n_bufs; j++) { + ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + const size_t nbytes = split_state.ne[j] * tensor->nb[2]; + if (nbytes == 0) { + continue; + } + ggml_backend_tensor_set_2d(simple_tensor, (const char *) data + offset_data, 0, nbytes, + tensor->ne[1], simple_tensor->nb[1], tensor->nb[1]); + offset_data += nbytes; + } + GGML_ASSERT(offset_data == (size_t) tensor->ne[2] * tensor->nb[2]); + return; + } + GGML_ASSERT(ggml_is_contiguous(tensor) || split_state.axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); if (split_state.n_segments != 1 || split_state.nr[0] != 1) { diff --git a/src/llama-model.cpp b/src/llama-model.cpp index a65abfc9ca64..1f2d7c5b4459 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -382,8 +382,8 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str static const std::regex pattern_kv_bias ("blk\\.\\d*\\.attn_(k|v)\\.bias"); static const std::regex pattern_qkv_bias ("blk\\.\\d*\\.attn_qkv.bias"); static const std::regex pattern_qk_norm ("blk\\.\\d*\\.attn_(q|k)_norm\\.weight"); - static const std::regex pattern_kv_cache ("cache_(k|v)_l\\d*"); - static const std::regex pattern_idx_cache ("cache_idx_(k|v)_l\\d*"); + static const std::regex pattern_kv_cache ("cache_(k|v)_l\\d+"); + static const std::regex pattern_idx_cache ("cache_idx_(k|v)_l\\d+"); static const std::regex pattern_dsv4_state ("dsv4_(csa|hca|lid)_state_(kv|score)_l\\d*"); static const std::regex pattern_attn_sinks ("blk\\.\\d*\\.attn_sinks.weight"); static const std::regex pattern_attn_out_weight ("blk\\.\\d*\\.attn_output.weight"); @@ -398,9 +398,9 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str static const std::regex pattern_ssm_alpha ("blk\\.\\d*\\.ssm_alpha.weight"); static const std::regex pattern_ssm_beta ("blk\\.\\d*\\.ssm_beta.weight"); static const std::regex pattern_ssm_beta_alpha ("blk\\.\\d*\\.ssm_ba.weight"); - static const std::regex pattern_r_cache ("cache_r_l\\d*"); - static const std::regex pattern_ple_r_cache ("cache_ple_r_l\\d*"); - static const std::regex pattern_s_cache ("cache_s_l\\d*"); + static const std::regex pattern_r_cache ("cache_r_l\\d+"); + static const std::regex pattern_ple_r_cache ("cache_ple_r_l\\d+"); + static const std::regex pattern_s_cache ("cache_s_l\\d+"); static const std::regex pattern_ssm_conv1d ("blk\\.\\d*\\.ssm_conv1d.weight"); static const std::regex pattern_ssm_out_weight ("blk\\.\\d*\\.ssm_out.weight"); @@ -470,7 +470,34 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str return {axis, tensor_axis_0, il, rotation}; }; + // A host-resident cache reaches attention as a scheduler copy, permuted, with the heads on an + // axis of their own. The split follows the cache, so record that axis and its unit - one head. + struct host_cache_copy { + bool valid = false; + ggml_backend_meta_split_axis axis = GGML_BACKEND_SPLIT_AXIS_0; + int64_t unit = 1; // cache elements per element of axis + }; + + const host_cache_copy cache_copy = [&]() -> host_cache_copy { + if (is_dsv4 || !std::regex_match(tensor_name, pattern_kv_cache)) { + return {}; + } + // [head_dim, n_kv, n_head_kv, n_stream] + const uint32_t il = std::stoul(tensor_name.substr(tensor_name.find("_l", 6) + 2)); + const int64_t head_dim = tensor_name[6] == 'k' ? hparams.n_embd_head_k(il) : hparams.n_embd_head_v(il); + if (hparams.n_head_kv(il) > 1 && tensor->ne[0] == head_dim && + tensor->ne[2] == (int64_t) hparams.n_head_kv(il)) { + return {true, GGML_BACKEND_SPLIT_AXIS_2, head_dim}; + } + return {}; + }(); + auto get_tensor_config = [&]() -> tensor_config { + // a graph tensor only reaches here as a scheduler copy, and only a copy of a host-resident + // cache is split - everything else the graph copies in keeps the mirrored fallback + if (!cache_copy.valid && ggml_backend_buffer_get_usage(tensor->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); + } if (is_dsv4) { if (std::regex_match(tensor_name, pattern_kv_cache) || std::regex_match(tensor_name, pattern_dsv4_state)) { @@ -594,7 +621,9 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); }; - auto get_split_segments = [&](int axis, uint32_t il) -> std::vector> { + // ne_axis is the extent of the split axis as the cache sees it, which for a host-resident + // cache copy is not the extent of the axis the copy is split on + auto get_split_segments = [&](int axis, uint32_t il, int64_t ne_axis) -> std::vector> { if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE || ud->model->arch == LLM_ARCH_QWEN4EXP) { const int64_t head_k_dim = hparams.ssm_d_state; @@ -642,7 +671,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str GGML_ASSERT(tensor->ne[axis] == 2*n_ff_exp); return {{n_ff_exp, 2}}; } - return {{tensor->ne[axis], 1}}; + return {{ne_axis, 1}}; } if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_qkv_bias)) { @@ -658,14 +687,14 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str if (tensor->ne[axis] == 2*n_ff) { return {{n_ff, 2}}; } - return {{tensor->ne[axis], 1}}; + return {{ne_axis, 1}}; } if (std::regex_match(tensor_name, pattern_ffn_gate_up_weight)) { const int64_t n_ff_exp = hparams.n_ff_exp(il); GGML_ASSERT(tensor->ne[axis] == 2*n_ff_exp); return {{n_ff_exp, 2}}; } - return {{tensor->ne[axis], 1}}; + return {{ne_axis, 1}}; }; auto get_split_granularity = [&](int64_t blck_size, uint32_t il, const std::vector> & segments) -> std::vector { @@ -788,6 +817,11 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str tensor_config tc = get_tensor_config(); split_state.axis = tc.axis; if (split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS) { + // a host-resident cache copy is split on the axis its heads arrive on, in whole heads, + // so the segments and the granularity of the cache are rescaled to that unit + const bool unfolded = cache_copy.valid && split_state.axis == GGML_BACKEND_SPLIT_AXIS_0; + const int64_t unit = unfolded ? cache_copy.unit : 1; + const int64_t ne_axis = unfolded ? tensor->ne[cache_copy.axis]*unit : tensor->ne[split_state.axis]; const int64_t blck_size = ggml_blck_size(tc.tensor_axis_0->type); const float * tensor_split = ud->model->tensor_split(); std::vector tensor_split_scan; @@ -798,12 +832,17 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str tensor_split_scan[j] += tensor_split_scan[j - 1]; } } - const std::vector> segments = get_split_segments(split_state.axis, tc.il); + const std::vector> segments = get_split_segments(split_state.axis, tc.il, ne_axis); const std::vector granularity = get_split_granularity(blck_size, tc.il, segments); + if (unfolded) { + split_state.axis = cache_copy.axis; + } for (size_t is = 0; is < segments.size(); is++) { - const int64_t ne_s = segments[is].first; + GGML_ASSERT(segments[is].first % unit == 0); + GGML_ASSERT(granularity[is] % unit == 0); + const int64_t ne_s = segments[is].first / unit; const uint32_t nr_s = segments[is].second; - const int64_t g_s = granularity[is]; + const int64_t g_s = granularity[is] / unit; int64_t low = 0; size_t j = 0; for (; j < ud->n_devices - 1; j++) { From fa56aca076dba1ae1267ebcfc39c363a75109452 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 3 Sep 2026 11:21:53 +0200 Subject: [PATCH 39/50] llama : split a V tensor on its own head size The KV granularity was derived from the query granularity through n_gqa, which assumes the V side has the head size of the K side. Count whole KV heads and scale each side by its own head size. Assisted-by: Claude Opus 5 --- src/llama-model.cpp | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 1f2d7c5b4459..66ae3b617fb6 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -781,12 +781,16 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str return {granularity_q}; } - const int64_t granularity_kv = granularity_q / n_gqa; + const int64_t n_head_gran = granularity_q / n_embd_q; // KV heads per granule + const int64_t granularity_kv = n_head_gran * hparams.n_embd_head_k(il); if (std::regex_match(tensor_name, pattern_kv_weight) || std::regex_match(tensor_name, pattern_kv_bias) || std::regex_match(tensor_name, pattern_kv_cache)) { GGML_ASSERT(segments.size() == 1); - return {granularity_kv}; + // the V side can have a head size of its own, the split still lands on head boundaries + const bool is_v = tensor_name.compare(0, 8, "cache_v_") == 0 || + tensor_name.find(".attn_v.") != std::string::npos; + return {is_v ? n_head_gran * hparams.n_embd_head_v(il) : granularity_kv}; } if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_qkv_bias)) { GGML_ASSERT(segments.size() == 2); From 5719cbfa0f50332206e3b2ca50a8fffa5ca05aa1 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 3 Sep 2026 11:21:53 +0200 Subject: [PATCH 40/50] ggml : read back a strided row split from a meta buffer A fused QKV puts Kcur and Vcur in a strided view, which a host-resident cache reads back through the meta buffer. The rows split, so the chunk splice could not express the read. Assisted-by: Claude Opus 5 --- ggml/src/ggml-backend-meta.cpp | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index ec0bde1c9275..30e6070fd4b4 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -1609,6 +1609,33 @@ static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, gg static void ggml_backend_meta_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer); const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); + + // A fused QKV puts Kcur and Vcur in a strided view, which a host-resident cache reads back here. + // The rows split, so the chunk splice below cannot express the read. Each device's part of a row + // is contiguous: ne[1] rows, from the device's own stride to the view's. + const bool strided_rows = + !ggml_is_contiguous(tensor) && + split_state.axis == GGML_BACKEND_SPLIT_AXIS_0 && + split_state.n_segments == 1 && split_state.nr[0] == 1 && + tensor->ne[2] == 1 && tensor->ne[3] == 1 && + offset == 0 && size == ggml_nbytes(tensor); + + if (strided_rows) { + size_t offset_data = 0; + for (size_t j = 0; j < n_bufs; j++) { + const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + const size_t nbytes = ggml_row_size(tensor->type, split_state.ne[j]); + if (nbytes == 0) { + continue; + } + ggml_backend_tensor_get_2d(simple_tensor, (char *) data + offset_data, 0, nbytes, + tensor->ne[1], simple_tensor->nb[1], tensor->nb[1]); + offset_data += nbytes; + } + GGML_ASSERT(offset_data == ggml_row_size(tensor->type, tensor->ne[0])); + return; + } + GGML_ASSERT(ggml_is_contiguous(tensor) || split_state.axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); if (split_state.n_segments != 1 || split_state.nr[0] != 1) { From 5cdde605b04677fe6b26cd6b650cef88b83c8d98 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 3 Sep 2026 11:23:39 +0200 Subject: [PATCH 41/50] llama : keep the recurrent state on device under split mode tensor A linear-attention op packs the state it writes back together with its output, so that split does not line up with the one a host-resident state expects. On device the two orders agree; in host memory they disagree, and the split state of the fused op no longer resolves. The state is small next to the attention cache that -nkvo exists to move, so keep it device-resident. Assisted-by: Claude Opus 5 --- src/llama-context.cpp | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 9aaa0f2a5cf6..f085569d1054 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -144,6 +144,14 @@ llama_context::llama_context( cparams.offload_kqv = params.offload_kqv; cparams.kv_cpu_pinned = params.kv_cpu_pinned; cparams.recurrent_state_offload = params.recurrent_state_offload; + + // A linear-attention op packs the state it writes back together with its output, so that split + // does not line up with the one a host-resident state expects. Keep the state on its device. + if (!cparams.recurrent_state_offload && model.split_mode() == LLAMA_SPLIT_MODE_TENSOR && + (llm_arch_is_recurrent(model.arch) || llm_arch_is_hybrid(model.arch))) { + LLAMA_LOG_INFO("%s: split mode tensor: keeping the recurrent state device-resident\n", __func__); + cparams.recurrent_state_offload = true; + } cparams.offload_attn_compute = params.offload_kqv || (params.op_offload && params.kv_cpu_pinned); cparams.kv_gpu_layers = params.kv_gpu_layers; cparams.phase_aware_workspace = params.phase_aware_workspace; From 346b38b7529f5fcf4d8b39351fcf859fb3eefef0 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 3 Sep 2026 11:22:08 +0200 Subject: [PATCH 42/50] tests : cover a host-resident KV cache split by tensor Run the tensor-split architecture matrix a second time with the cache in host memory, and add an 8-device CI run, where the scheduler copy name is long enough to be truncated. Assisted-by: Claude Opus 5 --- ci/run.sh | 4 +++- tests/test-llama-archs.cpp | 26 +++++++++++++++++++++----- 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/ci/run.sh b/ci/run.sh index 1ceb19fd50aa..110335a4efa5 100755 --- a/ci/run.sh +++ b/ci/run.sh @@ -312,6 +312,8 @@ function gg_run_test_llama_archs_tensor_split { GGML_CUDA_DEVICES=2 ./build-ci-release/bin/test-llama-archs -s 1 2>&1 GGML_CUDA_DEVICES=3 ./build-ci-release/bin/test-llama-archs -s 1 2>&1 GGML_CUDA_DEVICES=4 ./build-ci-release/bin/test-llama-archs -s 1 2>&1 + # the scheduler names a copy after its source, and 8 device names fill the name field + GGML_CUDA_DEVICES=8 ./build-ci-release/bin/test-llama-archs -s 1 -a llama 2>&1 fi if [ ! -z ${GG_BUILD_METAL} ]; then @@ -327,7 +329,7 @@ function gg_run_test_llama_archs_tensor_split { function gg_sum_test_llama_archs_tensor_split { gg_printf '### %s\n\n' "${ci}" - gg_printf 'Runs test-llama-archs with 1 to 4 devices\n' + gg_printf 'Runs test-llama-archs with 1 to 4 and 8 devices\n' gg_printf '- status: %s\n' "$(cat $OUT/${ci}.exit)" gg_printf '```\n' gg_printf '%s\n' "$(cat $OUT/${ci}.log)" diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 3951dbbe8526..b25d932c150a 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -398,9 +398,14 @@ static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/) return true; } +// with offload_kqv=false the cache lives in host memory +struct kv_config { + bool offload_kqv = true; +}; + static std::pair get_model_and_ctx( struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector & devs, - const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false) { + const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false, const kv_config & kvc = {}) { GGML_ASSERT((gguf_ctx == nullptr) != (file == nullptr)); llama_model_params model_params = llama_model_default_params(); model_params.progress_callback = silent_model_load_progress; @@ -413,6 +418,7 @@ static std::pair get_model_and_ctx( ctx_params.n_ctx = 0; ctx_params.n_threads = 4; ctx_params.n_threads_batch = 4; + ctx_params.offload_kqv = kvc.offload_kqv; if (!encode) { ctx_params.n_ubatch = 64; } @@ -1377,9 +1383,10 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const in std::vector devs; std::string label; llama_split_mode split_mode; + kv_config kvc; - device_config(std::vector devs, std::string name, llama_split_mode split_mode) - : devs(std::move(devs)), label(std::move(name)), split_mode(split_mode) {} + device_config(std::vector devs, std::string name, llama_split_mode split_mode, kv_config kvc = {}) + : devs(std::move(devs)), label(std::move(name)), split_mode(split_mode), kvc(std::move(kvc)) {} }; std::vector dev_configs; @@ -1401,6 +1408,15 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const in } dev_configs.emplace_back(devices_meta, "Meta", LLAMA_SPLIT_MODE_TENSOR); + + // a host-resident cache reaches attention as a scheduler copy that is split by head + kv_config kvc_host; + kvc_host.offload_kqv = false; + dev_configs.emplace_back(devices_meta, "Meta -nkvo", LLAMA_SPLIT_MODE_TENSOR, kvc_host); + + for (const device_config & dc : dev_configs) { + max_device_label_length = std::max(max_device_label_length, dc.label.length()); + } } size_t max_arch_name_length = 0; @@ -1473,7 +1489,7 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const in logits_cpu = get_logits(model_and_ctx_cpu.first.get(), model_and_ctx_cpu.second.get(), tokens, encode); } if (dc.split_mode != LLAMA_SPLIT_MODE_TENSOR || llm_arch_supports_sm_tensor(arch)) { - model_and_ctx_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, dc.devs, dc.split_mode, encode); + model_and_ctx_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, dc.devs, dc.split_mode, encode, dc.kvc); logits_dev = get_logits(model_and_ctx_dev.first.get(), model_and_ctx_dev.second.get(), tokens, encode); const double nmse_val = nmse(logits_cpu, logits_dev); snprintf(nmse_str, sizeof(nmse_str), "(%.2e)", nmse_val); @@ -1504,7 +1520,7 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const in ms.save(file); rewind(file); - auto model_and_ctx_roundtrip = get_model_and_ctx(nullptr, file, seed, dc.devs, dc.split_mode, encode); + auto model_and_ctx_roundtrip = get_model_and_ctx(nullptr, file, seed, dc.devs, dc.split_mode, encode, dc.kvc); const std::vector logits_roundtrip = get_logits( model_and_ctx_roundtrip.first.get(), model_and_ctx_roundtrip.second.get(), tokens, encode); status_roundtrip = "\033[1;32mOK\033[0m"; From 1f90680dcaa0660e1b10ebc6925f70bde5753120 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Fri, 4 Sep 2026 23:57:21 +0200 Subject: [PATCH 43/50] ggml-meta : write every stream of a head-split host KV cache The strided head split path ran only for a single stream, so a host-resident cache built with --parallel N fell through to the chunk splice, which cannot express that write. Loop over ne[3] and offset each side by its own stride. Assisted-by: Claude Opus 5 --- ggml/src/ggml-backend-meta.cpp | 29 ++++++++++++++++------------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 30e6070fd4b4..8593ad02ef57 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -1455,29 +1455,32 @@ static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, gg const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer); const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); - // A host-resident attention cache reaches this permuted, as [head_dim, n_kv, n_head_kv, 1]. The - // heads split, but interleaved per cell, so the chunk splice below cannot express the write. + // A host-resident attention cache reaches this permuted, as [head_dim, n_kv, n_head_kv, n_stream]. + // The heads split, but interleaved per cell, so the chunk splice below cannot express the write. // Each device's heads are one contiguous run per cell: ne[1] cells, from one stride to another. const bool strided_head_split = !ggml_is_contiguous(tensor) && split_state.axis == GGML_BACKEND_SPLIT_AXIS_2 && split_state.n_segments == 1 && split_state.nr[0] == 1 && - tensor->ne[3] == 1 && tensor->nb[1] > tensor->nb[2] && + tensor->nb[1] > tensor->nb[2] && offset == 0 && size == ggml_nbytes(tensor); if (strided_head_split) { - size_t offset_data = 0; - for (size_t j = 0; j < n_bufs; j++) { - ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); - const size_t nbytes = split_state.ne[j] * tensor->nb[2]; - if (nbytes == 0) { - continue; + for (int64_t i3 = 0; i3 < tensor->ne[3]; i3++) { + size_t offset_data = i3 * tensor->nb[3]; + for (size_t j = 0; j < n_bufs; j++) { + ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + const size_t nbytes = split_state.ne[j] * tensor->nb[2]; + if (nbytes == 0) { + continue; + } + ggml_backend_tensor_set_2d(simple_tensor, (const char *) data + offset_data, + i3 * simple_tensor->nb[3], nbytes, + tensor->ne[1], simple_tensor->nb[1], tensor->nb[1]); + offset_data += nbytes; } - ggml_backend_tensor_set_2d(simple_tensor, (const char *) data + offset_data, 0, nbytes, - tensor->ne[1], simple_tensor->nb[1], tensor->nb[1]); - offset_data += nbytes; + GGML_ASSERT(offset_data == i3 * tensor->nb[3] + (size_t) tensor->ne[2] * tensor->nb[2]); } - GGML_ASSERT(offset_data == (size_t) tensor->ne[2] * tensor->nb[2]); return; } From 35c62439569352b54ae77a0a3a4248cd6d8e3df2 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sun, 6 Sep 2026 01:19:07 +0200 Subject: [PATCH 44/50] ggml-meta : close the last subgraph on a host-resident cache A split that ends with a view of a host tensor left the subgraph bookkeeping short of the node count and aborted. That happens with more than one cache stream, which the test matrix now covers. Also make the scheduler copy name a stated contract instead of a grammar that two files reconstruct on their own, note the 2d transfer fallback at both strided cache paths, say out loud that split mode tensor overrides the recurrent state placement, and record the Gemma 4 host cache accuracy gap where it is skipped. Assisted-by: Claude Opus 5 --- docs/multi-gpu.md | 2 + ggml/src/ggml-backend-impl.h | 13 ++++++ ggml/src/ggml-backend-meta.cpp | 46 +++++++------------ ggml/src/ggml-backend.cpp | 37 +++++++++++++-- src/llama-context.cpp | 3 +- tests/CMakeLists.txt | 7 +++ tests/test-llama-archs.cpp | 84 ++++++++++++++++++++++++++++------ 7 files changed, 144 insertions(+), 48 deletions(-) diff --git a/docs/multi-gpu.md b/docs/multi-gpu.md index 0d9eea7c2fb8..3d57bfbbd461 100644 --- a/docs/multi-gpu.md +++ b/docs/multi-gpu.md @@ -86,6 +86,8 @@ llama-cli -m model.gguf -sm tensor -ctk f16 -ctv f16 - `--flash-attn off` or (`--flash-attn auto` resolving to `off` when it isn't supported) is a hard error. - KV cache types must be non-quantized: `f32`, `f16`, or `bf16`. Support for quantized KV cache is not implemented and trying to use it will result in an error. - Mark this configuration as experimental in your tooling: validate output quality before deploying. +- `--no-kv-offload` works in this mode: the host-resident cache is split by attention head like the rest. Two limits: a backend without a native 2d copy (CPU, Metal) pays one transfer per cache cell, and Gemma 4 is less accurate this way (perplexity 235.03 with a host cache versus 227.23 with a device cache), so keep the cache on the devices for that architecture. +- A recurrent or hybrid model always keeps its recurrent state on the devices in this mode, even with `--no-recurrent-state-offload`. The linear-attention op writes the state back together with its output, so the two do not agree on a host-resident state. - `--split-mode tensor`is not implemented for all architectures. The following will fail with *"LLAMA_SPLIT_MODE_TENSOR not implemented for architecture '...'"*: - **MoE / hybrid:** Grok, MPT, OLMoE, DeepSeek2, GLM-DSA, Nemotron-H, Nemotron-H-MoE, Granite-Hybrid, LFM2-MoE, Minimax-M2, Mistral4, Kimi-Linear, Jamba, Falcon-H1 diff --git a/ggml/src/ggml-backend-impl.h b/ggml/src/ggml-backend-impl.h index ef05905cf9ab..f60f44997051 100644 --- a/ggml/src/ggml-backend-impl.h +++ b/ggml/src/ggml-backend-impl.h @@ -104,6 +104,19 @@ extern "C" { // temporary workaround to statically allocate tensors from a context in a deduplicated way: GGML_API struct ggml_backend_buffer * ggml_backend_meta_alloc_ctx_tensors_from_buft(struct ggml_context * ctx, ggml_backend_buffer_type_t buft); + // + // Backend (sched) + // + + // The scheduler names a copy of a graph input "##". is the name of the + // tensor the copy was made from and carries any suffix that ggml appends for a view. Only + // identifies the copy, so the backend label is cut when the name does not fit. + GGML_API void ggml_backend_sched_name_copy( + struct ggml_tensor * copy, const char * backend_name, const struct ggml_tensor * src, int c); + + // write the part of a name written by ggml_backend_sched_name_copy into buf + GGML_API void ggml_backend_sched_copy_source_name(const char * name, char * buf, size_t buf_size); + // // Backend (stream) // diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 8593ad02ef57..f86712aff013 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -485,28 +485,6 @@ static struct ggml_tensor * ggml_backend_meta_buffer_simple_tensor(const struct return it->second[index]; } -// A scheduler copy is named "##", where also carries the suffixes -// ggml appends for views. Recover the graph name the copy was made from. -static std::string ggml_backend_meta_copy_source_name(const char * name) { - std::string ret = name; - const size_t first = ret.find('#'); - if (first == std::string::npos) { - return ret; - } - ret.erase(0, first + 1); - // ggml writes a view suffix as " (...)", a graph name has no spaces - const size_t suffix = ret.find(" ("); - if (suffix != std::string::npos) { - ret.erase(suffix); - return ret; - } - const size_t copy = ret.rfind('#'); - if (copy != std::string::npos && ret.find_first_not_of("0123456789", copy + 1) == std::string::npos) { - ret.erase(copy); - } - return ret; -} - static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(const struct ggml_tensor * tensor, bool assume_sync); static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( @@ -869,7 +847,9 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( ggml_tensor tensor_named; if (copied_in_leaf) { tensor_named = *tensor; - ggml_set_name(&tensor_named, ggml_backend_meta_copy_source_name(tensor->name).c_str()); + char source_name[GGML_MAX_NAME]; + ggml_backend_sched_copy_source_name(tensor->name, source_name, sizeof(source_name)); + ggml_set_name(&tensor_named, source_name); tensor_query = &tensor_named; } ggml_backend_meta_split_state ret = dev_ctx->get_split_state(tensor_query, dev_ctx->get_split_state_ud); @@ -1458,6 +1438,7 @@ static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, gg // A host-resident attention cache reaches this permuted, as [head_dim, n_kv, n_head_kv, n_stream]. // The heads split, but interleaved per cell, so the chunk splice below cannot express the write. // Each device's heads are one contiguous run per cell: ne[1] cells, from one stride to another. + // A backend without a native 2d copy (CPU, Metal) pays one transfer per cell here, CUDA does not. const bool strided_head_split = !ggml_is_contiguous(tensor) && split_state.axis == GGML_BACKEND_SPLIT_AXIS_2 && @@ -1616,6 +1597,7 @@ static void ggml_backend_meta_buffer_get_tensor(ggml_backend_buffer_t buffer, co // A fused QKV puts Kcur and Vcur in a strided view, which a host-resident cache reads back here. // The rows split, so the chunk splice below cannot express the read. Each device's part of a row // is contiguous: ne[1] rows, from the device's own stride to the view's. + // A backend without a native 2d copy (CPU, Metal) pays one transfer per row here, CUDA does not. const bool strided_rows = !ggml_is_contiguous(tensor) && split_state.axis == GGML_BACKEND_SPLIT_AXIS_0 && @@ -2277,14 +2259,18 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend, int i_start = 0; for (int i = 0; i < cgraph->n_nodes; i++) { ggml_tensor * node = cgraph->nodes[i]; - if (node->view_src != nullptr && node->view_src->op == GGML_OP_NONE && ggml_backend_buffer_is_host(node->view_src->buffer)) { - continue; - } - const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(node, /*assume_sync =*/ false); - if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) { - max_tmp_size = std::max(max_tmp_size, ggml_nbytes(node)); + // a host-resident KV cache ends a split with a view of itself, that view needs no split state + // but it is still the last node and must close the last subgraph + const bool host_view = node->view_src != nullptr && node->view_src->op == GGML_OP_NONE && + ggml_backend_buffer_is_host(node->view_src->buffer); + bool new_subgraph = i + 1 == cgraph->n_nodes; + if (!host_view) { + const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(node, /*assume_sync =*/ false); + if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) { + max_tmp_size = std::max(max_tmp_size, ggml_nbytes(node)); + new_subgraph = true; + } } - const bool new_subgraph = i + 1 == cgraph->n_nodes || split_state.axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL; if (!new_subgraph) { continue; } diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index fd9788ef025f..94080c86fc73 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1034,9 +1034,9 @@ static void ggml_backend_sched_print_assignments(ggml_backend_sched_t sched, str } } -// Name a copy of an input "##". The name field has a fixed size, so cut the -// backend label rather than the source name - the source name is what identifies the copy. -static void ggml_backend_sched_name_copy( +// The name field has a fixed size, so cut the backend label rather than the source name. +// See the contract at the declaration in ggml-backend-impl.h. +void ggml_backend_sched_name_copy( struct ggml_tensor * copy, const char * backend_name, const struct ggml_tensor * src, int c) { const int n_tail = snprintf(NULL, 0, "#%s#%d", src->name, c); const int n_max = GGML_MAX_NAME - 1 - n_tail; @@ -1047,6 +1047,37 @@ static void ggml_backend_sched_name_copy( ggml_format_name(copy, "%.*s#%s#%d", n_head, backend_name, src->name, c); } +void ggml_backend_sched_copy_source_name(const char * name, char * buf, size_t buf_size) { + GGML_ASSERT(buf_size > 0); + + const char * first = strchr(name, '#'); + const char * src = first != NULL ? first + 1 : name; + size_t len = strlen(src); + + // ggml writes a view suffix as " (...)", a graph name has no spaces + const char * suffix = strstr(src, " ("); + if (suffix != NULL) { + len = suffix - src; + } else { + const char * copy = strrchr(src, '#'); + if (copy != NULL) { + const char * digits = copy + 1; + while (*digits >= '0' && *digits <= '9') { + digits++; + } + if (*digits == '\0') { + len = copy - src; + } + } + } + + if (len > buf_size - 1) { + len = buf_size - 1; + } + memcpy(buf, src, len); + buf[len] = '\0'; +} + static bool ggml_backend_sched_buffer_supported(ggml_backend_sched_t sched, struct ggml_tensor * t, int backend_id) { ggml_backend_buffer_t buf = t->view_src ? t->view_src->buffer : t->buffer; ggml_backend_buffer_type_t buft = NULL; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index f085569d1054..c7ba2f400c08 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -149,7 +149,8 @@ llama_context::llama_context( // does not line up with the one a host-resident state expects. Keep the state on its device. if (!cparams.recurrent_state_offload && model.split_mode() == LLAMA_SPLIT_MODE_TENSOR && (llm_arch_is_recurrent(model.arch) || llm_arch_is_hybrid(model.arch))) { - LLAMA_LOG_INFO("%s: split mode tensor: keeping the recurrent state device-resident\n", __func__); + LLAMA_LOG_WARN("%s: split mode tensor cannot keep the recurrent state in host memory - " + "overriding --no-recurrent-state-offload, which needs more VRAM\n", __func__); cparams.recurrent_state_offload = true; } cparams.offload_attn_compute = params.offload_kqv || (params.op_offload && params.kv_cpu_pinned); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 04693bb0e6ed..aa42167bf30e 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -214,6 +214,13 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) ARGS --test-live-context-workspace ) + llama_test( + test-llama-archs + NAME test-sched-copy-name + LABEL main + ARGS --test-sched-copy-name + ) + set(MODEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-models/") file(MAKE_DIRECTORY "${MODEL_DIR}") diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index b25d932c150a..329fd951fa3c 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -69,7 +70,7 @@ static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) { } static void usage(char ** argv) { - printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-o/--out dir] [-v N] [-h/--help] [--test-phase-workspace] [--test-live-context-workspace]\n", argv[0]); + printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-o/--out dir] [-v N] [-h/--help] [--test-phase-workspace] [--test-live-context-workspace] [--test-sched-copy-name]\n", argv[0]); } static std::vector get_tokens(const uint32_t n_tokens, const uint32_t n_vocab, const size_t seed){ @@ -399,8 +400,10 @@ static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/) } // with offload_kqv=false the cache lives in host memory +// n_seq_max > 1 gives the cache one stream per sequence struct kv_config { - bool offload_kqv = true; + bool offload_kqv = true; + uint32_t n_seq_max = 1; }; static std::pair get_model_and_ctx( @@ -419,6 +422,7 @@ static std::pair get_model_and_ctx( ctx_params.n_threads = 4; ctx_params.n_threads_batch = 4; ctx_params.offload_kqv = kvc.offload_kqv; + ctx_params.n_seq_max = kvc.n_seq_max; if (!encode) { ctx_params.n_ubatch = 64; } @@ -1058,15 +1062,44 @@ static void test_phase_workspace_mismatched_placement(size_t seed) { GGML_ASSERT(llama_contexts_share_workspace(target.get(), draft.get()) == (status == 1)); } +// the meta backend recovers the source of a scheduler copy from its name, so both sides must agree +static void test_sched_copy_name() { + const ggml_init_params params = { + /*.mem_size =*/ 2*ggml_tensor_overhead(), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context_ptr ctx(ggml_init(params)); + GGML_ASSERT(ctx); + ggml_tensor * src = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, 1); + ggml_tensor * copy = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, 1); + + auto check = [&](const char * name, const char * backend_name, const char * expected) { + char buf[GGML_MAX_NAME]; + ggml_set_name(src, name); + ggml_backend_sched_name_copy(copy, backend_name, src, 0); + ggml_backend_sched_copy_source_name(copy->name, buf, sizeof(buf)); + GGML_ASSERT(strcmp(buf, expected) == 0); + }; + + check("cache_k_l0", "CUDA0", "cache_k_l0"); + check("cache_k_l0 (view)", "CUDA0", "cache_k_l0"); + // the backend label is cut when the name does not fit, the source name survives + check("cache_k_l31", "Meta(CUDA0,CUDA1,CUDA2,CUDA3,CUDA4,CUDA5,CUDA6,CUDA7)", "cache_k_l31"); +} + +// the tokens are spread over n_seq sequences, each of which starts at position 0 static std::vector get_logits( - llama_model * model, llama_context * lctx, const std::vector & tokens, bool encode = false) { + llama_model * model, llama_context * lctx, const std::vector & tokens, bool encode = false, uint32_t n_seq = 1) { const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model)); const uint32_t n_ctx = llama_n_ctx(lctx); const uint32_t n_tokens = tokens.size(); - llama_batch batch = llama_batch_init(n_ctx, 0, 1); + GGML_ASSERT(n_seq >= 1 && n_tokens % n_seq == 0); + const uint32_t n_tokens_seq = n_tokens / n_seq; + llama_batch batch = llama_batch_init(n_ctx, 0, n_seq); GGML_ASSERT(n_tokens <= n_ctx); for (uint32_t pos = 0; pos < n_tokens; pos++) { - common_batch_add(batch, tokens[pos], pos, {0}, true); + common_batch_add(batch, tokens[pos], pos % n_tokens_seq, {llama_seq_id(pos / n_tokens_seq)}, true); } batch.n_tokens = n_tokens; if (encode) { @@ -1414,6 +1447,11 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const in kvc_host.offload_kqv = false; dev_configs.emplace_back(devices_meta, "Meta -nkvo", LLAMA_SPLIT_MODE_TENSOR, kvc_host); + // with more than one stream that copy is 4d, one stride per stream + kv_config kvc_host_streams = kvc_host; + kvc_host_streams.n_seq_max = 2; + dev_configs.emplace_back(devices_meta, "Meta -nkvo -np 2", LLAMA_SPLIT_MODE_TENSOR, kvc_host_streams); + for (const device_config & dc : dev_configs) { max_device_label_length = std::max(max_device_label_length, dc.label.length()); } @@ -1448,7 +1486,10 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const in continue; } if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) { - continue; // FIXME: ISWA KV cache initialization needs more fixture params + // FIXME: ISWA KV cache initialization needs more fixture params + // this arch also loses accuracy with a host-resident cache under split mode tensor + // (perplexity 235.03 versus 227.23), the other ISWA archs are clean + continue; } if (arch == LLM_ARCH_EAGLE3 || arch == LLM_ARCH_DFLASH) { continue; @@ -1468,7 +1509,7 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const in GGML_ASSERT(gguf_remove_key(gguf_ctx.get(), "bailingmoe3.kda.safe_gate") >= 0); } std::pair model_and_ctx_cpu; - std::vector logits_cpu; + std::map> logits_cpu_per_n_seq; for (device_config & dc : dev_configs) { // print test config first; should anything fail during model loading or inference, at least we know which test case caused it printf(template_row_cfg.c_str(), @@ -1482,15 +1523,21 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const in std::string status_parallel = "\033[1;33mSKIP\033[0m"; char nmse_str[12] = {0}; - bool skip = !arch_supported(arch) || (dc.split_mode == LLAMA_SPLIT_MODE_TENSOR && dc.devs.empty()); + // an encoder-decoder model needs its own batch layout, so it stays on one sequence + bool skip = !arch_supported(arch) || (dc.split_mode == LLAMA_SPLIT_MODE_TENSOR && dc.devs.empty()) || + (encode && dc.kvc.n_seq_max > 1); if (!skip) { - if (logits_cpu.empty()) { - model_and_ctx_cpu = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}, LLAMA_SPLIT_MODE_LAYER, encode); - logits_cpu = get_logits(model_and_ctx_cpu.first.get(), model_and_ctx_cpu.second.get(), tokens, encode); - } if (dc.split_mode != LLAMA_SPLIT_MODE_TENSOR || llm_arch_supports_sm_tensor(arch)) { + // the reference runs the same batch layout, one sequence per stream + std::vector & logits_cpu = logits_cpu_per_n_seq[dc.kvc.n_seq_max]; + if (logits_cpu.empty()) { + kv_config kvc_cpu; + kvc_cpu.n_seq_max = dc.kvc.n_seq_max; + model_and_ctx_cpu = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}, LLAMA_SPLIT_MODE_LAYER, encode, kvc_cpu); + logits_cpu = get_logits(model_and_ctx_cpu.first.get(), model_and_ctx_cpu.second.get(), tokens, encode, dc.kvc.n_seq_max); + } model_and_ctx_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, dc.devs, dc.split_mode, encode, dc.kvc); - logits_dev = get_logits(model_and_ctx_dev.first.get(), model_and_ctx_dev.second.get(), tokens, encode); + logits_dev = get_logits(model_and_ctx_dev.first.get(), model_and_ctx_dev.second.get(), tokens, encode, dc.kvc.n_seq_max); const double nmse_val = nmse(logits_cpu, logits_dev); snprintf(nmse_str, sizeof(nmse_str), "(%.2e)", nmse_val); status_nmse = "\033[1;32mOK\033[0m"; @@ -1522,7 +1569,7 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const in auto model_and_ctx_roundtrip = get_model_and_ctx(nullptr, file, seed, dc.devs, dc.split_mode, encode, dc.kvc); const std::vector logits_roundtrip = get_logits( - model_and_ctx_roundtrip.first.get(), model_and_ctx_roundtrip.second.get(), tokens, encode); + model_and_ctx_roundtrip.first.get(), model_and_ctx_roundtrip.second.get(), tokens, encode, dc.kvc.n_seq_max); status_roundtrip = "\033[1;32mOK\033[0m"; GGML_ASSERT(logits_roundtrip.size() == logits_dev.size()); for (size_t i = 0; i < logits_roundtrip.size(); i++) { @@ -1557,6 +1604,7 @@ int main(int argc, char ** argv) { std::string out; bool test_phase_workspace = false; bool test_live_context_workspace = false; + bool test_copy_name = false; int verbosity = LOG_LEVEL_ERROR; @@ -1610,6 +1658,10 @@ int main(int argc, char ** argv) { test_live_context_workspace = true; continue; } + if (strcmp(argv[i], "--test-sched-copy-name") == 0) { + test_copy_name = true; + continue; + } } printf("%s: using seed %zu\n", __func__, seed); @@ -1627,6 +1679,10 @@ int main(int argc, char ** argv) { test_live_context_workspace_unsupported(seed); return 0; } + if (test_copy_name) { + test_sched_copy_name(); + return 0; + } if (!out.empty()) { return save_models(arch, seed, verbosity, out); } From 20b0891d143c54e7993d0c732f3f345953533c13 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Mon, 7 Sep 2026 21:40:53 +0200 Subject: [PATCH 45/50] ggml-meta : identify a scheduler copy by its name, not by its flags The scheduler flags a copy as an input when it keeps more than one, which hid a host-resident cache from the split state callback. A cut source name now asserts instead of naming another tensor, and memset writes a head split like set_tensor does. Assisted-by: Claude Opus 5 --- ggml/src/ggml-backend-impl.h | 8 +++++--- ggml/src/ggml-backend-meta.cpp | 33 ++++++++++++++++++++++++++++++--- ggml/src/ggml-backend.cpp | 31 +++++++++++++++++++------------ tests/test-llama-archs.cpp | 10 ++++++++-- 4 files changed, 62 insertions(+), 20 deletions(-) diff --git a/ggml/src/ggml-backend-impl.h b/ggml/src/ggml-backend-impl.h index f60f44997051..36abb660535b 100644 --- a/ggml/src/ggml-backend-impl.h +++ b/ggml/src/ggml-backend-impl.h @@ -110,12 +110,14 @@ extern "C" { // The scheduler names a copy of a graph input "##". is the name of the // tensor the copy was made from and carries any suffix that ggml appends for a view. Only - // identifies the copy, so the backend label is cut when the name does not fit. + // identifies the copy, so the backend label is cut when the name does not fit. A source name that + // does not fit on its own asserts, a cut one would name a different tensor. GGML_API void ggml_backend_sched_name_copy( struct ggml_tensor * copy, const char * backend_name, const struct ggml_tensor * src, int c); - // write the part of a name written by ggml_backend_sched_name_copy into buf - GGML_API void ggml_backend_sched_copy_source_name(const char * name, char * buf, size_t buf_size); + // write the part of a name written by ggml_backend_sched_name_copy into buf, + // returns false and leaves buf alone if the name is not one + GGML_API bool ggml_backend_sched_copy_source_name(const char * name, char * buf, size_t buf_size); // // Backend (stream) diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index f86712aff013..23dc91dee3d4 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -834,9 +834,12 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( // A host-resident KV cache reaches the graph as a copied-in leaf of the compute buffer. // Mirroring it while the queries stay split by head makes each device attend the wrong // heads, so ask the callback; it still answers MIRRORED for names it does not know. + // The name tells a copy apart, the scheduler also flags one as an input when it keeps several. + char source_name[GGML_MAX_NAME]; const bool copied_in_leaf = ggml_backend_buffer_get_usage(tensor->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && - tensor->op == GGML_OP_NONE && (tensor->flags & GGML_TENSOR_FLAG_INPUT) == 0; + tensor->op == GGML_OP_NONE && + ggml_backend_sched_copy_source_name(tensor->name, source_name, sizeof(source_name)); if ((ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_COMPUTE || copied_in_leaf) && tensor->view_src == nullptr) { @@ -847,8 +850,6 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( ggml_tensor tensor_named; if (copied_in_leaf) { tensor_named = *tensor; - char source_name[GGML_MAX_NAME]; - ggml_backend_sched_copy_source_name(tensor->name, source_name, sizeof(source_name)); ggml_set_name(&tensor_named, source_name); tensor_query = &tensor_named; } @@ -1334,6 +1335,32 @@ static void ggml_backend_meta_buffer_memset_tensor( const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer); const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); + + // a host-resident attention cache is permuted, its heads are one run per cell, see set_tensor + const bool strided_head_split = + !ggml_is_contiguous(tensor) && + split_state.axis == GGML_BACKEND_SPLIT_AXIS_2 && + split_state.n_segments == 1 && split_state.nr[0] == 1 && + tensor->nb[1] > tensor->nb[2] && + offset == 0 && size == ggml_nbytes(tensor); + + if (strided_head_split) { + for (int64_t i3 = 0; i3 < tensor->ne[3]; i3++) { + for (size_t j = 0; j < n_bufs; j++) { + ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + const size_t nbytes = split_state.ne[j] * tensor->nb[2]; + if (nbytes == 0) { + continue; + } + for (int64_t i1 = 0; i1 < tensor->ne[1]; i1++) { + ggml_backend_tensor_memset(simple_tensor, value, + i3*simple_tensor->nb[3] + i1*simple_tensor->nb[1], nbytes); + } + } + } + return; + } + GGML_ASSERT(ggml_is_contiguous(tensor) || split_state.axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); if (split_state.n_segments != 1 || split_state.nr[0] != 1) { diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 94080c86fc73..66f4d2dc28b9 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1040,18 +1040,22 @@ void ggml_backend_sched_name_copy( struct ggml_tensor * copy, const char * backend_name, const struct ggml_tensor * src, int c) { const int n_tail = snprintf(NULL, 0, "#%s#%d", src->name, c); const int n_max = GGML_MAX_NAME - 1 - n_tail; + GGML_ASSERT(n_max >= 0 && "source name too long to name a scheduler copy"); int n_head = (int) strlen(backend_name); if (n_head > n_max) { - n_head = n_max > 0 ? n_max : 0; + n_head = n_max; } ggml_format_name(copy, "%.*s#%s#%d", n_head, backend_name, src->name, c); } -void ggml_backend_sched_copy_source_name(const char * name, char * buf, size_t buf_size) { +bool ggml_backend_sched_copy_source_name(const char * name, char * buf, size_t buf_size) { GGML_ASSERT(buf_size > 0); const char * first = strchr(name, '#'); - const char * src = first != NULL ? first + 1 : name; + if (first == NULL) { + return false; + } + const char * src = first + 1; size_t len = strlen(src); // ggml writes a view suffix as " (...)", a graph name has no spaces @@ -1060,22 +1064,25 @@ void ggml_backend_sched_copy_source_name(const char * name, char * buf, size_t b len = suffix - src; } else { const char * copy = strrchr(src, '#'); - if (copy != NULL) { - const char * digits = copy + 1; - while (*digits >= '0' && *digits <= '9') { - digits++; - } - if (*digits == '\0') { - len = copy - src; - } + if (copy == NULL) { + return false; } + const char * digits = copy + 1; + while (*digits >= '0' && *digits <= '9') { + digits++; + } + if (digits == copy + 1 || *digits != '\0') { + return false; + } + len = copy - src; } if (len > buf_size - 1) { - len = buf_size - 1; + return false; } memcpy(buf, src, len); buf[len] = '\0'; + return true; } static bool ggml_backend_sched_buffer_supported(ggml_backend_sched_t sched, struct ggml_tensor * t, int backend_id) { diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 329fd951fa3c..b4411522b9b9 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -1078,7 +1078,7 @@ static void test_sched_copy_name() { char buf[GGML_MAX_NAME]; ggml_set_name(src, name); ggml_backend_sched_name_copy(copy, backend_name, src, 0); - ggml_backend_sched_copy_source_name(copy->name, buf, sizeof(buf)); + GGML_ASSERT(ggml_backend_sched_copy_source_name(copy->name, buf, sizeof(buf))); GGML_ASSERT(strcmp(buf, expected) == 0); }; @@ -1086,6 +1086,13 @@ static void test_sched_copy_name() { check("cache_k_l0 (view)", "CUDA0", "cache_k_l0"); // the backend label is cut when the name does not fit, the source name survives check("cache_k_l31", "Meta(CUDA0,CUDA1,CUDA2,CUDA3,CUDA4,CUDA5,CUDA6,CUDA7)", "cache_k_l31"); + + // a name that was not written by the scheduler is not a copy + char buf[GGML_MAX_NAME]; + GGML_ASSERT(!ggml_backend_sched_copy_source_name("cache_k_l0", buf, sizeof(buf))); + GGML_ASSERT(!ggml_backend_sched_copy_source_name("CUDA0#cache_k_l0", buf, sizeof(buf))); + // a cut name lost its copy index, what is left of the source names another tensor + GGML_ASSERT(!ggml_backend_sched_copy_source_name("#cache_k_l3", buf, sizeof(buf))); } // the tokens are spread over n_seq sequences, each of which starts at position 0 @@ -1431,7 +1438,6 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const in for (size_t i = 0; i < device_count; i++) { ggml_backend_dev_t dev = ggml_backend_dev_get(i); dev_configs.emplace_back(std::vector{dev}, ggml_backend_dev_description(dev), LLAMA_SPLIT_MODE_LAYER); - max_device_label_length = std::max(max_device_label_length, dev_configs.back().label.length()); // cpu-based devices cannot be used in tensor split mode if (ggml_backend_dev_buffer_type(dev) != ggml_backend_cpu_buffer_type()) { From 50b379a5f5d01d733a03d64b4b01df8e9505e3b6 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Mon, 14 Sep 2026 02:04:53 +0200 Subject: [PATCH 46/50] sched : write a multi-stream window to a meta backend in one copy The ranged copy of a host-resident window calls set_tensor_async once per stream. A meta backend only takes a whole contiguous tensor, so with two or more streams it aborts. Keep the whole-span copy for a meta destination. Assisted-by: Claude Opus 5 --- ggml/src/ggml-backend.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 66f4d2dc28b9..bea33db60bc2 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1863,7 +1863,8 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s ggml_backend_buffer_t src_buf = input->view_src ? input->view_src->buffer : input->buffer; struct ggml_backend_sched_ranges rg; ggml_backend_sched_input_ranges(input, &rg); - const bool ranged = rg.n > 1 && src_buf != NULL && ggml_backend_buffer_is_host(src_buf); + // a meta backend writes a whole contiguous tensor, it cannot take one range per stream + const bool ranged = rg.n > 1 && src_buf != NULL && ggml_backend_buffer_is_host(src_buf) && !ggml_backend_is_meta(split_backend); // 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 From 9faffa99a4ba76fbaa4c6c23b874fcd4c74e9ed6 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Thu, 17 Sep 2026 21:51:49 +0200 Subject: [PATCH 47/50] tests : align the parallel check with the host KV rows A multi-stream row decodes its reference as several sequences, so it cannot serve as a one-sequence reference - skip the check there. Give the check the row's offload_kqv, so it covers a host-resident cache with two streams. Assisted-by: Claude Opus 5 --- tests/test-llama-archs.cpp | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index b4411522b9b9..cf193bf168e6 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -1134,7 +1134,7 @@ static std::vector get_logits( // decode two sequences in one batch, compare each with a decode of it alone // logits_a: logits of tokens decoded alone with the same device config -static bool test_parallel_seqs(llama_model * model, const std::vector & tokens, const std::vector & logits_a, bool encode) { +static bool test_parallel_seqs(llama_model * model, const std::vector & tokens, const std::vector & logits_a, bool encode, bool offload_kqv) { const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model)); const uint32_t n_tokens = tokens.size(); @@ -1145,6 +1145,7 @@ static bool test_parallel_seqs(llama_model * model, const std::vector 1 aborts - if (arch != LLM_ARCH_T5) { + // a multi-stream row spreads the reference over several sequences, so it does not match one sequence at pos 0..n-1 + if (arch != LLM_ARCH_T5 && dc.kvc.n_seq_max == 1) { status_parallel = "\033[1;32mOK\033[0m"; - if (!test_parallel_seqs(model_and_ctx_dev.first.get(), tokens, logits_dev, encode)) { + if (!test_parallel_seqs(model_and_ctx_dev.first.get(), tokens, logits_dev, encode, dc.kvc.offload_kqv)) { all_ok = false; status_parallel = "\033[1;31mFAIL\033[0m"; } From a22c586eaadd8f924fcdfe55548c76b623ca4d26 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Mon, 14 Sep 2026 18:41:38 +0200 Subject: [PATCH 48/50] ggml-meta : add events and a ranged write for a head-split host cache A meta event is one event per simple device, recorded and waited for on each device's own stream. set_tensor_async and set_tensor_2d_async take part of the window of a host cache split by head: whole cells from an offset, once per stream, as one 2d copy per stream per device. A mirrored tensor passes the offset through. Add ggml_backend_meta_init_transfer, a meta backend that does not start a communicator, for streams that only move data. Rotate the compute containers of buffers that are only sources of a graph, so tensors bound into a buffer the graph allocator does not own do not pile up. Assisted-by: Claude Opus 5 --- ggml/src/ggml-backend-impl.h | 3 + ggml/src/ggml-backend-meta.cpp | 179 ++++++++++++++++++++++++++++++--- 2 files changed, 170 insertions(+), 12 deletions(-) diff --git a/ggml/src/ggml-backend-impl.h b/ggml/src/ggml-backend-impl.h index 36abb660535b..27c7340a605a 100644 --- a/ggml/src/ggml-backend-impl.h +++ b/ggml/src/ggml-backend-impl.h @@ -101,6 +101,9 @@ extern "C" { GGML_API size_t ggml_backend_meta_n_backends (ggml_backend_t meta_backend); GGML_API ggml_backend_t ggml_backend_meta_simple_backend(ggml_backend_t meta_backend, size_t index); + // a meta backend without a communicator, for moving data on streams of its own: a graph it computes reduces through copies + GGML_API ggml_backend_t ggml_backend_meta_init_transfer(ggml_backend_dev_t meta_dev); + // temporary workaround to statically allocate tensors from a context in a deduplicated way: GGML_API struct ggml_backend_buffer * ggml_backend_meta_alloc_ctx_tensors_from_buft(struct ggml_context * ctx, ggml_backend_buffer_type_t buft); diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 23dc91dee3d4..c192c8a966cb 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -151,6 +151,12 @@ static ggml_backend_buffer_type_t ggml_backend_meta_device_get_buffer_type(ggml_ static ggml_backend_buffer_type_t ggml_backend_meta_device_get_host_buffer_type(ggml_backend_dev_t dev); +static ggml_backend_event_t ggml_backend_meta_device_event_new(ggml_backend_dev_t dev); + +static void ggml_backend_meta_device_event_free(ggml_backend_dev_t dev, ggml_backend_event_t event); + +static void ggml_backend_meta_device_event_synchronize(ggml_backend_dev_t dev, ggml_backend_event_t event); + static bool ggml_backend_meta_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { GGML_ASSERT(ggml_backend_dev_is_meta(dev)); const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; @@ -190,9 +196,9 @@ static const ggml_backend_device_i ggml_backend_meta_device_iface = { /* .supports_op = */ ggml_backend_meta_device_supports_op, /* .supports_buft = */ ggml_backend_meta_device_supports_buft, /* .offload_op = */ nullptr, - /* .event_new = */ nullptr, - /* .event_free = */ nullptr, - /* .event_synchronize = */ nullptr, + /* .event_new = */ ggml_backend_meta_device_event_new, + /* .event_free = */ ggml_backend_meta_device_event_free, + /* .event_synchronize = */ ggml_backend_meta_device_event_synchronize, }; static bool ggml_backend_dev_is_meta(ggml_backend_dev_t dev) { @@ -212,6 +218,45 @@ static ggml_backend_dev_t ggml_backend_meta_dev_simple_dev(ggml_backend_dev_t me return meta_dev_ctx->simple_devs[index]; } +// a meta event is one event per simple device, in the order of the simple devices +static ggml_backend_event_t ggml_backend_meta_device_event_new(ggml_backend_dev_t dev) { + const size_t n_devs = ggml_backend_meta_dev_n_devs(dev); + auto * events = new std::vector(); + events->reserve(n_devs); + for (size_t i = 0; i < n_devs; i++) { + ggml_backend_event_t event = ggml_backend_event_new(ggml_backend_meta_dev_simple_dev(dev, i)); + if (event == nullptr) { + for (ggml_backend_event_t e : *events) { + ggml_backend_event_free(e); + } + delete events; + return nullptr; + } + events->push_back(event); + } + return new ggml_backend_event { + /* .device = */ dev, + /* .context = */ events, + }; +} + +static void ggml_backend_meta_device_event_free(ggml_backend_dev_t dev, ggml_backend_event_t event) { + GGML_UNUSED(dev); + auto * events = (std::vector *) event->context; + for (ggml_backend_event_t e : *events) { + ggml_backend_event_free(e); + } + delete events; + delete event; +} + +static void ggml_backend_meta_device_event_synchronize(ggml_backend_dev_t dev, ggml_backend_event_t event) { + GGML_UNUSED(dev); + for (ggml_backend_event_t e : *(std::vector *) event->context) { + ggml_backend_event_synchronize(e); + } +} + ggml_backend_dev_t ggml_backend_meta_device( ggml_backend_dev_t * devs, size_t n_devs, ggml_backend_meta_get_split_state_t get_split_state, void * get_split_state_ud) { GGML_ASSERT(n_devs <= GGML_BACKEND_META_MAX_DEVICES); @@ -1909,7 +1954,7 @@ struct ggml_backend_meta_context { void * comm_ctx = nullptr; ggml_backend_comm_allreduce_tensor_t comm_allreduce = nullptr; - ggml_backend_meta_context(ggml_backend_dev_t meta_dev, const char * params) { + ggml_backend_meta_context(ggml_backend_dev_t meta_dev, const char * params, bool init_comm) { const size_t n_devs = ggml_backend_meta_dev_n_devs(meta_dev); n_reduce_steps = std::ceil(std::log2(n_devs)); name = "Meta("; @@ -1927,7 +1972,7 @@ struct ggml_backend_meta_context { } name += ")"; - if (n_devs > 1) { + if (n_devs > 1 && init_comm) { ggml_backend_comm_init_t comm_init = (ggml_backend_comm_init_t) ggml_backend_reg_get_proc_address( ggml_backend_dev_backend_reg(ggml_backend_get_device(simple_backends[0])), "ggml_backend_comm_init"); if (comm_init != nullptr) { @@ -1968,12 +2013,68 @@ static void ggml_backend_meta_free(ggml_backend_t backend) { delete backend; } +// A host-resident attention cache copy, [head_dim, n_kv, n_head_kv, n_stream], split on its heads, which are interleaved per cell. +static bool ggml_backend_meta_is_strided_head_split(const ggml_tensor * tensor, const ggml_backend_meta_split_state & split_state) { + return !ggml_is_contiguous(tensor) && + split_state.axis == GGML_BACKEND_SPLIT_AXIS_2 && + split_state.n_segments == 1 && split_state.nr[0] == 1 && + tensor->nb[1] > tensor->nb[2]; +} + +// Write whole cells of a strided head split, n_copies times, from one stream of the copy to the next. +// Each device takes its own run of heads from every cell: one 2d copy per stream per device. +static void ggml_backend_meta_set_head_split_async(ggml_backend_t backend, ggml_tensor * tensor, + const ggml_backend_meta_split_state & split_state, const void * data, size_t offset, size_t size, + size_t n_copies, size_t stride_tensor, size_t stride_data) { + const size_t n_backends = ggml_backend_meta_n_backends(backend); + const size_t nb1 = tensor->nb[1]; + const size_t nb3 = tensor->nb[3]; + + GGML_ASSERT(n_copies <= 1 || stride_tensor == nb3); + GGML_ASSERT(nb3 > 0 && (offset % nb3) % nb1 == 0 && size % nb1 == 0); + + const size_t i3 = offset / nb3; + const size_t i1 = (offset % nb3) / nb1; + const size_t n_rows = size / nb1; + + size_t head_offset = 0; + for (size_t j = 0; j < n_backends; j++) { + const size_t nbytes = split_state.ne[j] * tensor->nb[2]; + if (nbytes == 0) { + continue; + } + ggml_backend_t simple_backend = ggml_backend_meta_simple_backend(backend, j); + ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + for (size_t k = 0; k < std::max(n_copies, 1); k++) { + ggml_backend_tensor_set_2d_async(simple_backend, simple_tensor, + (const char *) data + k*stride_data + head_offset, + (i3 + k)*simple_tensor->nb[3] + i1*simple_tensor->nb[1], nbytes, + n_rows, simple_tensor->nb[1], nb1); + } + head_offset += nbytes; + } + GGML_ASSERT(head_offset == (size_t) tensor->ne[2] * tensor->nb[2]); +} + static void ggml_backend_meta_set_tensor_async(ggml_backend_t backend, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { const size_t n_backends = ggml_backend_meta_n_backends(backend); + const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); + + // a pipelined delivery of a host-resident cache writes part of the window + if (ggml_backend_meta_is_strided_head_split(tensor, split_state)) { + ggml_backend_meta_set_head_split_async(backend, tensor, split_state, data, offset, size, 1, 0, 0); + return; + } + if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + for (size_t j = 0; j < n_backends; j++) { + ggml_backend_tensor_set_async( + ggml_backend_meta_simple_backend(backend, j), ggml_backend_meta_buffer_simple_tensor(tensor, j), data, offset, size); + } + return; + } + GGML_ASSERT(offset == 0); GGML_ASSERT(ggml_is_contiguous(tensor)); - - const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); GGML_ASSERT(split_state.n_segments == 1); GGML_ASSERT(split_state.nr[0] == 1); @@ -2013,6 +2114,27 @@ static void ggml_backend_meta_set_tensor_async(ggml_backend_t backend, ggml_tens } } +static void ggml_backend_meta_set_tensor_2d_async(ggml_backend_t backend, ggml_tensor * tensor, const void * data, + size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data) { + const size_t n_backends = ggml_backend_meta_n_backends(backend); + const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); + + if (ggml_backend_meta_is_strided_head_split(tensor, split_state)) { + ggml_backend_meta_set_head_split_async(backend, tensor, split_state, data, offset, size, n_copies, stride_tensor, stride_data); + return; + } + if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + for (size_t j = 0; j < n_backends; j++) { + ggml_backend_tensor_set_2d_async(ggml_backend_meta_simple_backend(backend, j), ggml_backend_meta_buffer_simple_tensor(tensor, j), + data, offset, size, n_copies, stride_tensor, stride_data); + } + return; + } + for (size_t i = 0; i < n_copies; i++) { + ggml_backend_meta_set_tensor_async(backend, tensor, (const char *) data + i*stride_data, offset + i*stride_tensor, size); + } +} + static void ggml_backend_meta_get_tensor_async(ggml_backend_t backend, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { const size_t n_backends = ggml_backend_meta_n_backends(backend); GGML_ASSERT(offset == 0); @@ -2065,6 +2187,23 @@ static void ggml_backend_meta_synchronize(ggml_backend_t backend) { } } +static void ggml_backend_meta_event_record(ggml_backend_t backend, ggml_backend_event_t event) { + const auto & events = *(std::vector *) event->context; + GGML_ASSERT(events.size() == ggml_backend_meta_n_backends(backend)); + for (size_t j = 0; j < events.size(); j++) { + ggml_backend_event_record(events[j], ggml_backend_meta_simple_backend(backend, j)); + } +} + +// each simple backend waits for the event of its own device +static void ggml_backend_meta_event_wait(ggml_backend_t backend, ggml_backend_event_t event) { + const auto & events = *(std::vector *) event->context; + GGML_ASSERT(events.size() == ggml_backend_meta_n_backends(backend)); + for (size_t j = 0; j < events.size(); j++) { + ggml_backend_event_wait(ggml_backend_meta_simple_backend(backend, j), events[j]); + } +} + static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend, struct ggml_cgraph * cgraph) { GGML_ASSERT(cgraph->grads == nullptr); const size_t n_backends = ggml_backend_meta_n_backends(backend); @@ -2096,6 +2235,13 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend, if (ggml_backend_buffer_is_meta(cgraph->nodes[i]->buffer)) { used_buffers.emplace(cgraph->nodes[i]->buffer); } + // the copies in a transport ring are only sources, but their buffer holds compute tensors that must rotate too + for (int k = 0; k < GGML_MAX_SRC; k++) { + const ggml_tensor * src = cgraph->nodes[i]->src[k]; + if (src != nullptr && ggml_backend_buffer_is_meta(src->buffer)) { + used_buffers.emplace(src->buffer); + } + } } for (ggml_backend_buffer_t buf : used_buffers) { ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) buf->context; @@ -2577,7 +2723,7 @@ static const ggml_backend_i ggml_backend_meta_i = { /* .free = */ ggml_backend_meta_free, /* .set_tensor_async = */ ggml_backend_meta_set_tensor_async, /* .get_tensor_async = */ ggml_backend_meta_get_tensor_async, - /* .set_tensor_2d_async = */ nullptr, + /* .set_tensor_2d_async = */ ggml_backend_meta_set_tensor_2d_async, /* .get_tensor_2d_async = */ nullptr, /* .cpy_tensor_async = */ nullptr, /* .synchronize = */ ggml_backend_meta_synchronize, @@ -2586,8 +2732,8 @@ static const ggml_backend_i ggml_backend_meta_i = { /* .graph_plan_update = */ nullptr, /* .graph_plan_compute = */ nullptr, /* .graph_compute = */ ggml_backend_meta_graph_compute, - /* .event_record = */ nullptr, - /* .event_wait = */ nullptr, + /* .event_record = */ ggml_backend_meta_event_record, + /* .event_wait = */ ggml_backend_meta_event_wait, /* .graph_optimize = */ nullptr, }; @@ -2595,8 +2741,8 @@ bool ggml_backend_is_meta(ggml_backend_t backend) { return backend != nullptr && backend->iface.get_name == ggml_backend_meta_i.get_name; } -static ggml_backend_t ggml_backend_meta_device_init_backend(ggml_backend_dev_t dev, const char * params) { - ggml_backend_meta_context * backend_ctx = new ggml_backend_meta_context(dev, params); +static ggml_backend_t ggml_backend_meta_init_impl(ggml_backend_dev_t dev, const char * params, bool init_comm) { + ggml_backend_meta_context * backend_ctx = new ggml_backend_meta_context(dev, params, init_comm); ggml_backend_t backend = new struct ggml_backend; backend->guid = ggml_backend_meta_guid(); @@ -2606,6 +2752,15 @@ static ggml_backend_t ggml_backend_meta_device_init_backend(ggml_backend_dev_t d return backend; } +static ggml_backend_t ggml_backend_meta_device_init_backend(ggml_backend_dev_t dev, const char * params) { + return ggml_backend_meta_init_impl(dev, params, /*init_comm =*/ true); +} + +ggml_backend_t ggml_backend_meta_init_transfer(ggml_backend_dev_t dev) { + GGML_ASSERT(ggml_backend_dev_is_meta(dev)); + return ggml_backend_meta_init_impl(dev, nullptr, /*init_comm =*/ false); +} + size_t ggml_backend_meta_n_backends(ggml_backend_t meta_backend) { GGML_ASSERT(ggml_backend_is_meta(meta_backend)); const ggml_backend_meta_context * backend_ctx = (const ggml_backend_meta_context *) meta_backend->context; From 734bdd8ed15663cd237573006c1fb3a27d293537 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Mon, 14 Sep 2026 18:41:38 +0200 Subject: [PATCH 49/50] sched : pipeline the host KV delivery under split mode tensor A meta backend is eligible when each of its simple backends is. Its transfer backend has no communicator, and the headroom check uses the simple device with the least free memory, since a meta buffer allocates the whole ring on every device. A device of the meta type that is not the ggml meta backend stays ordered. test-llama-archs runs the host-resident tensor split with two streams a second time at pipeline depth 1. Assisted-by: Claude Opus 5 --- ggml/src/ggml-backend.cpp | 73 +++++++++++++++++++++++++++----------- tests/test-alloc.cpp | 1 + tests/test-llama-archs.cpp | 11 ++++-- 3 files changed, 63 insertions(+), 22 deletions(-) diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 458c7e150f7b..acd1ab108dbe 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -2165,6 +2165,26 @@ static size_t ggml_backend_sched_transport_slot_alloc(size_t need, size_t limit, return std::max(size, need); } +// Free memory of the device a ring lands on, 0 when unknown. +// A meta buffer allocates the whole ring on every simple device, so the device with the least free memory decides; the meta device reports the sum. +static size_t ggml_backend_sched_transport_dev_free(ggml_backend_t backend) { + size_t dev_free = 0, dev_total = 0; + if (ggml_backend_is_meta(backend)) { + for (size_t j = 0; j < ggml_backend_meta_n_backends(backend); j++) { + ggml_backend_dev_t dev = ggml_backend_get_device(ggml_backend_meta_simple_backend(backend, j)); + size_t simple_free = 0, simple_total = 0; + ggml_backend_dev_memory(dev, &simple_free, &simple_total); + dev_free = j == 0 ? simple_free : std::min(dev_free, simple_free); + } + return dev_free; + } + ggml_backend_dev_t dev = ggml_backend_get_device(backend); + if (dev != NULL) { + ggml_backend_dev_memory(dev, &dev_free, &dev_total); + } + return dev_free; +} + // 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; @@ -2179,7 +2199,8 @@ static bool ggml_backend_sched_transport_ensure_backend(ggml_backend_sched_t sch return false; } - ggml_backend_t transfer = ggml_backend_dev_init(dev, NULL); + // a transfer backend never computes, so a meta one goes without the communicator a second set of streams would otherwise start + ggml_backend_t transfer = ggml_backend_is_meta(sched->backends[backend_id]) ? ggml_backend_meta_init_transfer(dev) : ggml_backend_dev_init(dev, NULL); if (transfer == NULL) { return false; } @@ -2484,11 +2505,7 @@ static void ggml_backend_sched_transport_plan(ggml_backend_sched_t sched) { ggml_backend_buffer_type_t buft = sched->bufts[bid]; // 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) { - ggml_backend_dev_memory(dev, &dev_free, &dev_total); - } + const size_t dev_free = ggml_backend_sched_transport_dev_free(sched->backends[bid]); 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 " @@ -2949,8 +2966,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s ggml_backend_buffer_t src_buf = input->view_src ? input->view_src->buffer : input->buffer; struct ggml_backend_sched_ranges rg; ggml_backend_sched_input_ranges(input, input_cpy, &rg); - // a meta backend writes a whole contiguous tensor, it cannot take one range per stream - const bool ranged = rg.n > 1 && src_buf != NULL && ggml_backend_buffer_is_host(src_buf) && !ggml_backend_is_meta(split_backend); + const bool ranged = rg.n > 1 && src_buf != NULL && ggml_backend_buffer_is_host(src_buf); // 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 @@ -3227,6 +3243,23 @@ static void ggml_backend_sched_transport_teardown(ggml_backend_sched_t sched) { } +static bool ggml_backend_sched_transport_backend_supported(ggml_backend_t backend) { + ggml_backend_dev_t dev = ggml_backend_get_device(backend); + if (dev == NULL) { + return false; + } + + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev); + if (reg == NULL || strcmp(ggml_backend_reg_name(reg), "CUDA") != 0) { + return false; + } + + return backend->iface.set_tensor_async != NULL && + backend->iface.event_record != NULL && + backend->iface.event_wait != NULL && + dev->iface.event_new != NULL; +} + bool ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, int depth) { GGML_ASSERT(sched); @@ -3272,22 +3305,22 @@ bool ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, continue; } const enum ggml_backend_dev_type type = ggml_backend_dev_type(dev); - if (type == GGML_BACKEND_DEVICE_TYPE_META || type == GGML_BACKEND_DEVICE_TYPE_CPU) { + if (type == GGML_BACKEND_DEVICE_TYPE_CPU) { 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; + // a meta backend delivers through its simple backends, so each of them must qualify on its own + bool can_transport = true; + if (ggml_backend_is_meta(backend)) { + for (size_t j = 0; j < ggml_backend_meta_n_backends(backend) && can_transport; j++) { + can_transport = ggml_backend_sched_transport_backend_supported(ggml_backend_meta_simple_backend(backend, j)); + } + } else if (type == GGML_BACKEND_DEVICE_TYPE_META) { + can_transport = false; + } else { + can_transport = ggml_backend_sched_transport_backend_supported(backend); } - - if (dev->iface.event_new == NULL) { + if (!can_transport) { continue; } diff --git a/tests/test-alloc.cpp b/tests/test-alloc.cpp index 4b8e5e69556b..b49b25357322 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -2120,6 +2120,7 @@ static void test_transport_stops_after_backend_failure() { GGML_ASSERT(deliveries == 0); } +// only the ggml meta backend is reached through its simple backends, another device of the meta type stays ordered 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); diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index cf193bf168e6..20f41f005524 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -402,8 +402,9 @@ static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/) // with offload_kqv=false the cache lives in host memory // n_seq_max > 1 gives the cache one stream per sequence struct kv_config { - bool offload_kqv = true; - uint32_t n_seq_max = 1; + bool offload_kqv = true; + uint32_t n_seq_max = 1; + uint32_t kv_pipeline_depth = 0; }; static std::pair get_model_and_ctx( @@ -423,6 +424,7 @@ static std::pair get_model_and_ctx( ctx_params.n_threads_batch = 4; ctx_params.offload_kqv = kvc.offload_kqv; ctx_params.n_seq_max = kvc.n_seq_max; + ctx_params.kv_pipeline_depth = kvc.kv_pipeline_depth; if (!encode) { ctx_params.n_ubatch = 64; } @@ -1459,6 +1461,11 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const in kvc_host_streams.n_seq_max = 2; dev_configs.emplace_back(devices_meta, "Meta -nkvo -np 2", LLAMA_SPLIT_MODE_TENSOR, kvc_host_streams); + // the same copy, delivered ahead of the split that reads it + kv_config kvc_host_pipelined = kvc_host_streams; + kvc_host_pipelined.kv_pipeline_depth = 1; + dev_configs.emplace_back(devices_meta, "Meta -nkvo -np 2 -kvpd 1", LLAMA_SPLIT_MODE_TENSOR, kvc_host_pipelined); + for (const device_config & dc : dev_configs) { max_device_label_length = std::max(max_device_label_length, dc.label.length()); } From 0f0f11e133b4ed4daf34e14218e05678285f5e8e Mon Sep 17 00:00:00 2001 From: piggidragon Date: Mon, 14 Sep 2026 18:41:38 +0200 Subject: [PATCH 50/50] docs : measure the pipelined transport under split mode tensor The A/B and server gate scripts take LLAMA_KV_SM, as the parallel gate already did. Assisted-by: Claude Opus 5 --- docs/kv-transport-pipelining.md | 65 ++++++++++++++++++++++++++---- docs/multi-gpu.md | 1 + docs/repro/r4-kv-pipeline-ab.sh | 6 ++- docs/repro/r4-kv-pipeline-exact.sh | 4 +- 4 files changed, 65 insertions(+), 11 deletions(-) diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md index dae92d143c2c..c39df7df6a66 100644 --- a/docs/kv-transport-pipelining.md +++ b/docs/kv-transport-pipelining.md @@ -304,7 +304,7 @@ A device-resident KV run is unaffected, and was measured to confirm it: 38.5612 ## 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. +- CUDA is the only enabled backend, on its own or as every simple device of a meta backend. 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 window the graph reads. 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, and the copy packs those ranges so a slot holds the window rather than the whole cache. A window whose streams are not on that dimension keeps the single flat range, which is correct but not accelerated. Each range carries its own stream's prefix; ranges that agree on it are issued as one strided copy, so streams at the same depth still cost a single call. @@ -312,22 +312,71 @@ A device-resident KV run is unaffected, and was measured to confirm it: 38.5612 - **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 exception is a graph that cannot be allocated next to the rings: there every device that was holding one gives it back for good, because the allocator does not say which of them it competed with. - **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). +- **Tensor parallelism pipelines through the meta backend**, and each device allocates the whole ring. See [Tensor parallelism](#tensor-parallelism). - **A host write to the cache waits for the delivery.** `llama_memory_clear(mem, true)` waits for the scheduler 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` set the defaults of a scheduler that nothing else configures. `llama_context` always configures its own from the context parameters, so under `llama-server` and `llama-bench` use `--kv-pipeline-depth` / `LLAMA_ARG_KV_PIPELINE_DEPTH` and `--kv-pipeline-budget` / `LLAMA_ARG_KV_PIPELINE_BUDGET` instead. ## 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. +`-sm tensor` pipelines the same way. The consumer is a meta backend, and the scheduler sees one ring, one transfer backend and one pair of events per slot; each of them fans out to the simple devices underneath. -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. +This depends on the host-resident cache being split by head, which is #66: before it, `-sm tensor --no-kv-offload` mirrored the cache to every device and produced wrong output. The copy arrives permuted as `[head_dim, n_kv, n_head_kv, n_stream]`, so each device's heads are one run inside every cell rather than a contiguous block. -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. +What the meta backend adds: -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. +- **Events.** A meta event is one event per simple device. Recording it records each part on that device's stream, and waiting on it makes each simple backend wait for its own device's part. `caps.events` still reports false, so nothing else starts using them. +- **A ranged head-split write.** `set_tensor_async` and `set_tensor_2d_async` accept a part of the window: whole cells from an offset, once per stream. Each device takes its run of heads from every cell with one 2d copy per stream, so the early and the late delivery are the same calls as on a single device, one per device. +- **A transfer backend without a communicator.** The transfer backend never computes, so `ggml_backend_meta_init_transfer` builds its streams without starting a second NCCL context. +- **The ring is a meta buffer.** A copy in it gets its per-device tensors from the same split-state callback as the copy the graph allocator would have made. The meta graph compute also rotates the compute containers of buffers that appear only as sources, or the ring's per-device tensors would accumulate one set per plan. + +**Each device allocates the whole slot.** A meta buffer places a tensor at the same offset in every device's buffer and sizes each of those buffers for the whole tensor, so a device that holds half the heads still allocates the full slot. The meta compute buffers already work this way. The budget and the headroom check are applied per device, against the device with the least free memory, so a ring that fits the budget costs that much on every device. + +### Measurements + +RTX 4070 (gen4 x16) + RTX 3060 (gen3 x4), CUDA with NCCL, `Qwen3.8-27B-UD-IQ2_M.gguf`, `-ngl 99 -sm tensor -t 3 -fa on -ctk q8_0 -ctv q8_0 -b 512 -ub 512 -nkvo --kv-cpu-pinned`, under `taskset -c 0,2,4`. The recurrent state stays on the devices, which `-sm tensor` forces. + +`LLAMA_KV_SM=tensor docs/repro/r4-kv-pipeline-ab.sh`, `--kv-pipeline-budget 512`, both passes shown: + +| depth | ordered | pipelined | gain | +|---:|---|---|---:| +| 4,096 | 15.8398, 15.8351 | 21.9574, 21.9368 | **+38.6%** | +| 16,384 | 7.0819, 7.0839 | 8.7811, 8.7829 | **+24.0%** | +| 32,768 | 4.0809, 4.0826 | 4.8851, 4.8852 | **+19.7%** | + +`llama-server`, the tasks of the exactness gate at `-c 32768`: + +| task | prompt | ordered | pipelined | gain | +|---|---:|---:|---:|---:| +| prose | 1,709 | 20.675 | 25.402 | **+22.9%** | +| code | 3,270 | 17.036 | 22.589 | **+32.6%** | +| prose | 14,821 | 7.582 | 9.451 | **+24.7%** | +| code | 29,670 | 4.428 | 5.317 | **+20.1%** | + +Per decode graph at 16,384, `GGML_SCHED_TRANSPORT_DEBUG=2`: + +| | ordered | pipelined | +|---|---:|---:| +| total | 138.85 ms | 111.78 ms | +| blocked in the ordered copy | 103.81 ms | 12.16 ms | +| blocked waiting for the consumer | 31.06 ms | 95.43 ms | +| bytes delivered early / late | 0 / 0 MiB | 549.3 / 3.2 MiB | +| ring | - | 3 slots x 35 MiB, per device | + +The copy is three times the compute here, where on the single RTX 4070 above the two were about equal. The 553 MiB cross in 104 ms, about 5.3 GB/s, and the 3060's half of them crosses a gen3 x4 link. The pipeline hides the compute behind the copy, and the consumer wait now contains the rest of the transfer, so the token is bounded by the slower link rather than by the order of the work. The ceiling is `max(copy, compute)` plus the work outside the split loop, the same as on one device, and at 16,384 the pipeline is within a few milliseconds of it. + +So the gain is a property of the link, as it is on one device, with one addition: the devices compute in lock step, so the slowest link sets the pace for all of them. Devices on equal links carry an equal share of the bytes each; that has not been measured here. + +### Validation under `-sm tensor` + +On the same two devices: + +- `docs/repro/r4-kv-pipeline-exact.sh` with `LLAMA_KV_SM=tensor`: all eight tasks identical at `N = 0`, `1` and `4`. +- `docs/repro/r4-kv-pipeline-parallel-exact.sh` with `LLAMA_KV_SM=tensor`: `125cb9c2082d36cf` at `N = 0`, `1` and `4`, 8 concurrent sequences over a cache split into streams. With `-sm none` and `-sm layer` the same gate still gives `17f946c340db110b` and `db661b7a08686b97` at `N = 0` and `1`. +- Greedy `llama-completion`, 64 tokens behind a 3k prompt: Qwen3.8-27B-UD-IQ2_M gives `64e86551f7ef1638` with a device-resident cache and at `N = 0`, `1` and `4` with a host one. gemma-4-26B-A4B gives `5525e3f5ac7337d7` at `N = 0` and `1` with `-ts 50,50`, and `4f8986fb3655a567` at both with `-ts 55,45`. +- `test-llama-archs` adds a `Meta -nkvo -np 2 -kvpd 1` configuration, which stages the cache on the meta ring and delivers it through the ranged head-split write. It passes on 2, 3 and 4 CUDA devices. It evaluates one ubatch, so it delivers only the late part. +- `test-alloc` passes. Its meta test now covers a device of the meta type that is not the ggml meta backend, which stays ordered. ## Future work -- 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. +- Allocate each device's part of a meta ring at its own share of the heads, instead of the whole slot on every device. diff --git a/docs/multi-gpu.md b/docs/multi-gpu.md index 3d57bfbbd461..fcab4546adc6 100644 --- a/docs/multi-gpu.md +++ b/docs/multi-gpu.md @@ -87,6 +87,7 @@ llama-cli -m model.gguf -sm tensor -ctk f16 -ctv f16 - KV cache types must be non-quantized: `f32`, `f16`, or `bf16`. Support for quantized KV cache is not implemented and trying to use it will result in an error. - Mark this configuration as experimental in your tooling: validate output quality before deploying. - `--no-kv-offload` works in this mode: the host-resident cache is split by attention head like the rest. Two limits: a backend without a native 2d copy (CPU, Metal) pays one transfer per cache cell, and Gemma 4 is less accurate this way (perplexity 235.03 with a host cache versus 227.23 with a device cache), so keep the cache on the devices for that architecture. +- With `--no-kv-offload --kv-cpu-pinned`, `--kv-pipeline-depth 1` delivers the cache while the previous split computes, as with one device. See [kv-transport-pipelining.md](kv-transport-pipelining.md#tensor-parallelism). - A recurrent or hybrid model always keeps its recurrent state on the devices in this mode, even with `--no-recurrent-state-offload`. The linear-attention op writes the state back together with its output, so the two do not agree on a host-resident state. - `--split-mode tensor`is not implemented for all architectures. The following will fail with *"LLAMA_SPLIT_MODE_TENSOR not implemented for architecture '...'"*: diff --git a/docs/repro/r4-kv-pipeline-ab.sh b/docs/repro/r4-kv-pipeline-ab.sh index d0e50427ab74..04e039e90f4c 100755 --- a/docs/repro/r4-kv-pipeline-ab.sh +++ b/docs/repro/r4-kv-pipeline-ab.sh @@ -3,11 +3,13 @@ # 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 ...] +# LLAMA_KV_SM=tensor splits the model and its cache by head over every device. 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}" +SM="${LLAMA_KV_SM:-none}" 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. @@ -30,7 +32,7 @@ run () { # $1 label, $2 pipeline depth, $3 context depth, $4 reps err="$(mktemp)" rc=0 taskset -c "$PIN" "$BUILD/bin/llama-bench" -m "$MODEL" --kv-pipeline-depth "$2" \ - --kv-pipeline-budget "$BUDGET" -ngl 99 -sm none -mg 0 -t 3 -nkvo 1 -kvcp 1 -rso 1 \ + --kv-pipeline-budget "$BUDGET" -ngl 99 -sm "$SM" -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 \ > "$out" 2> "$err" || rc=$? if [ "$rc" -eq 0 ]; then @@ -55,7 +57,7 @@ for D in "${DEPTHS[@]}"; do R=5 fi echo "== context depth=$D reps=$R" - flock "$LOCK" bash -c "set -euo pipefail; $(declare -f run); BUILD='$BUILD'; MODEL='$MODEL'; PIN='$PIN'; BUDGET='$BUDGET' + flock "$LOCK" bash -c "set -euo pipefail; $(declare -f run); BUILD='$BUILD'; MODEL='$MODEL'; PIN='$PIN'; BUDGET='$BUDGET'; SM='$SM' run ordered 0 $D $R run pipelined 1 $D $R run ordered2 0 $D $R diff --git a/docs/repro/r4-kv-pipeline-exact.sh b/docs/repro/r4-kv-pipeline-exact.sh index 017926d29a0c..ab929e746ce1 100755 --- a/docs/repro/r4-kv-pipeline-exact.sh +++ b/docs/repro/r4-kv-pipeline-exact.sh @@ -4,6 +4,7 @@ # # 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). +# LLAMA_KV_SM=tensor splits the model and its cache by head over every device. set -u MODEL="${LLAMA_KV_MODEL:?set LLAMA_KV_MODEL to a .gguf path}" BUILD="${LLAMA_KV_BUILD:-build}" @@ -12,6 +13,7 @@ PORT="${LLAMA_KV_PORT:-18099}" LENGTHS="${LLAMA_KV_LENGTHS:-2048,18432}" CTX="${LLAMA_KV_CTX:-32768}" BUDGET="${LLAMA_KV_BUDGET:-512}" +SM="${LLAMA_KV_SM:-none}" HERE="$(cd "$(dirname "$0")" && pwd)" DEPTHS=(0 1 4); [ $# -gt 0 ] && DEPTHS=("$@") @@ -22,7 +24,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" \ - --kv-pipeline-budget "$BUDGET" -ngl 99 -sm none -mg 0 -t 3 -nkvo --kv-cpu-pinned --recurrent-state-offload \ + --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" --parallel 1 \ --host 127.0.0.1 --port "$PORT" --no-warmup > "$LOG" 2>&1 & SRV=$!