diff --git a/AGENTS.md b/AGENTS.md index 5a329c4d205..8fb894636bf 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -18,10 +18,24 @@ small: `GGML_TYPE_Q2_0S` so it cannot collide with upstream's serialized Q2_0 weight format. - A profit-only adaptive draft-max controller for DFlash. +- Pipelined delivery of a host-resident KV cache, selected with + `--kv-pipeline-depth` and bounded by `--kv-pipeline-budget`. - Reasoning-loop detection and the opted-in realtime `/v1/chat/completions/control` endpoint. - INI presets and KLD measurement support in `llama-perplexity`. +`--kv-pipeline-depth N` issues the host-to-device delivery of a host-resident KV +cache N splits ahead of the split that reads it, on a transfer stream of its own, +so a decode token stops paying transfer and attention in series. It is on by +default at `N = 1`, `0` restores the ordered path exactly, and it only engages +where a host-resident cache produces the deliveries (`--no-kv-offload`, with or +without `--kv-cpu-pinned`). Output is byte-identical either way. The staging it +needs grows with the context, so `--kv-pipeline-budget` (default 128 MiB) caps it +outright and the scheduler declines past the cap rather than spending the device +memory a host-resident cache exists to save. Measured single-GPU only. See +[`docs/kv-transport-pipelining.md`](docs/kv-transport-pipelining.md) for the +design, the numbers, the limits, and the reproduction scripts. + DFlash GGUFs must use upstream's `dflash` architecture, metadata, and tensor names. diff --git a/common/arg.cpp b/common/arg.cpp index f4016aa1ef2..5f62cef069a 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2600,6 +2600,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 0ce5d9bd5e8..cab11aa47b4 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1890,6 +1890,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.phase_aware_workspace = params.phase_aware_workspace; cparams.live_context_workspace = params.live_context_workspace; diff --git a/common/common.h b/common/common.h index a45faa93c5b..1b487bd8d50 100644 --- a/common/common.h +++ b/common/common.h @@ -642,6 +642,8 @@ struct common_params { bool recurrent_state_offload = false; // for hybrid models, keep recurrent state GPU-resident even with no_kv_offload bool phase_aware_workspace = false; // resize compute schedulers between prompt and generation phases bool live_context_workspace = false; // size supported attention workspaces from the padded live KV extent + int32_t kv_pipeline_depth = 1; // splits of look-ahead for pipelined delivery of a host-resident KV cache (0 = off) + int32_t kv_pipeline_budget_mib = 128; // hard cap on the device memory that delivery may use (0 = uncapped) bool warmup = true; // warmup run bool check_tensors = false; // validate tensor data bool no_op_offload = false; // globally disable offload host tensor operations to device diff --git a/docs/beellama-args.md b/docs/beellama-args.md index 2ba91e02373..feb7c364a7c 100644 --- a/docs/beellama-args.md +++ b/docs/beellama-args.md @@ -202,6 +202,8 @@ context-dependent H2D traffic; it is not equivalent to CPU attention. | `--recurrent-state-offload`, `--no-recurrent-state-offload` | `LLAMA_ARG_RECURRENT_STATE_OFFLOAD` | Disabled | For hybrid models under `--no-kv-offload`, keeps the fixed recurrent R/S state on the GPU while long attention KV remains in host memory. With ordinary KV offload enabled, recurrent state is already device-resident. | | `--kv-gpu-layers N` | `LLAMA_ARG_KV_GPU_LAYERS` | `0` | Under `--no-kv-offload`, retains the first `N` attention-KV layers on the device. This trades VRAM for lower repeated host-to-device traffic; a value at or above the attention-layer count matches full KV offload for those layers. | | `--spec-draft-kv-gpu-layers N`, `--kv-gpu-layers-draft N` | `LLAMA_ARG_SPEC_DRAFT_KV_GPU_LAYERS` | Inherit target | Overrides the target KV policy for a separate draft context and keeps the first `N` independently owned draft attention-KV layers device-resident. Shared KV layers follow their owner. The INI key is `spec-draft-kv-gpu-layers`. | +| `--kv-pipeline-depth N` | `LLAMA_ARG_KV_PIPELINE_DEPTH` | `1` | How many splits ahead the scheduler delivers a host-resident KV cache to the accelerator, so that the transfer runs while the previous split computes instead of in series with it. `0` restores the ordered path. Only engages where a host-resident cache produces the deliveries, and costs `(N + 2) x (largest staged split)` of device memory. See [`kv-transport-pipelining.md`](kv-transport-pipelining.md). | +| `--kv-pipeline-budget N` | `LLAMA_ARG_KV_PIPELINE_BUDGET` | `128` | Hard cap, in MiB, on the device memory pipelined delivery may use. A staging slot holds one attention layer's K or V over the whole context, so the requirement grows with the context; past this cap the scheduler declines and keeps the ordered path, at no steady-state cost, so a host-resident cache never quietly trades away the device memory it exists to save. `0` removes the cap. | | `--op-offload`, `--no-op-offload` | — | Enabled | Controls whether operations over host tensors may execute on the accelerator. The negated form selects the separately tested CPU execution route; it is not the serving-performance path measured below. | Target and draft attention-KV placement is independent of model-weight diff --git a/docs/kv-transport-pipelining.md b/docs/kv-transport-pipelining.md new file mode 100644 index 00000000000..c1f5188b4ef --- /dev/null +++ b/docs/kv-transport-pipelining.md @@ -0,0 +1,281 @@ +# Pipelined delivery of a host-resident KV cache + +With `--no-kv-offload` (optionally with `--kv-cpu-pinned`), the attention history +lives in host RAM and has to reach the accelerator on every decode token. The +backend scheduler used to issue that transfer on the consumer's own stream, right +before the kernels that read it, so a token cost `copy + compute` in series. + +This is the R4 item of the KV-offload roadmap +([PR #31](https://github.com/GenerelSchwerz/llama.cpp/pull/31)): the transfer and +the attention arithmetic are the same as before, but the transfer is issued one +split ahead, on a stream of its own, so the copy engine retires it underneath the +kernels of the split before it. + +`--kv-pipeline-depth N` controls it. It is on by default at `N = 1` and only has +an effect where a host-resident cache produces the deliveries; `0` restores the +ordered path exactly. + +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 (measurements §5 of PR #31: `+1.38%`, and not +exact). + +So the scheduler allocates its own ring and points the staged input copies at it +before the graph is allocated. A tensor that already has `data` is left alone by +`ggml_gallocr_init_tensor`, so the ring sits outside the allocator's reuse +analysis rather than competing with it. + +Each slot has one ownership cycle: + +1. the transfer stream owns an idle slot and writes one future split's prefix into it; +2. it records the slot's `ready` event, which the consumer stream waits for before launching the split that reads the slot; +3. the consumer records `release` once every kernel that reads the slot has been enqueued, and the transfer stream waits for that before overwriting the slot for a later split. + +Membership in the ring is decided once, when the ring is laid out, and execution +goes by the recorded answer. How much of a staged input can go early moves with +every ubatch; *which* input copies live in the ring must not, because their +addresses were handed out at allocation time. + +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 from PR #31 §3 | share | +|---|---:|---|---|---:|---:|---:| +| 4,096 | 5 | 31.7324, 31.7363 | 37.0889, 37.0741 | **+16.9%** | 38.49 | 96.4% | +| 16,384 | 3 | 19.6765, 19.6854 | 31.5352, 31.5807 | **+60.4%** | 34.88 | 90.4% | +| 32,768 | 3 | 13.0264, 13.0254 | 15.5325, 15.5329 | **+19.3%** | 20.83 | 74.6% | + +These are the uncapped numbers, measured before `--kv-pipeline-budget` existed; +they are what the ring can buy, and the 32,768 row needs +`--kv-pipeline-budget 512` to reproduce, because 213 MiB is over the 128 MiB +default. At the default the 4,096 and 16,384 rows stand and 32,768 declines to +the ordered path. See [The budget](#the-budget). + +Server decode behind an 18,422-token prompt +(`docs/repro/r4-kv-pipeline-exact.sh`): **18.468 -> 30.685 t/s, +66.2%**. + +### Across context depth, with device memory + +`docs/repro/r4-kv-pipeline-context-sweep.sh`, A/B/A/B, peak device memory sampled +with `nvidia-smi` across each arm. Both passes agreed to the digits shown. + +| Context | ordered | pipelined | gain | peak device memory | delta | ring | +|---|---:|---:|---:|---|---:|---:| +| 4,096 | 31.66 | 37.01 | **+16.9%** | 10,169 -> 10,197 MiB | +28 MiB | 27 MiB | +| 16,384 | 19.64 | 31.43 | **+60.1%** | 10,159 -> 10,263 MiB | +104 MiB | 107 MiB | +| 32,768 | 12.99 | 15.49 | **+19.2%** | 10,161 -> 10,367 MiB | +206 MiB | 213 MiB | +| 65,536 | 7.74 | 8.74 | **+12.9%** | 10,163 -> 10,573 MiB | +410 MiB | 428 MiB | +| 131,072 | 4.29 | 4.68 | **+9.1%** | 10,537 -> 11,355 MiB | +818 MiB | 855 MiB | +| 262,144 | 2.25 | 2.24 | **declined** | 11,329 -> 11,391 MiB | +62 MiB | not allocated | + +Two curves run in opposite directions here, and both matter. + +**The gain narrows with depth.** A token is copy plus compute; as the context +grows the copy grows with it while the compute per staged split does not, so the +share of the token that can hide a transfer shrinks. At 16,384 compute still +covers most of the copy; by 131,072 it covers a tenth of it. That is arithmetic, +not an implementation limit, and no amount of look-ahead changes it. + +**The ring's cost does not narrow.** It is `(depth + 2)` slots of one staged +split, and a staged split is K and V of one attention layer over the whole +context: it doubles every time the context doubles. At 131,072 it claims 818 MiB +of an 11,902 MiB card to buy 9.1%. + +At 262,144 the ring would need 1.7 GiB against 573 MiB free, so it declines and +the run stays on the ordered path -- 2.25 against 2.24 t/s, inside the spread of +the ordered arm's own two passes, and 62 MiB of device memory for the transfer +backend's context. Declining is the intended outcome, not a failure: the +62 MiB +and the unchanged throughput are what "the guard did its job" looks like. + +**On a memory-constrained card, past roughly 64k the same device memory is +probably better spent on `--kv-gpu-layers`.** At 131,072 a staged split is +285 MiB, so the 818 MiB the ring takes is about three attention layers' worth of +K and V; making three of sixteen layers device-resident removes about 19% of the +host-to-device traffic against the 9.1% the ring buys. That comparison has not +been measured here and it will move with the model's layer count and the card, so +it is a pointer for whoever tunes a deployment, not a recommendation. + +### The budget + +The table above is what the feature costs uncapped, and it is the reason it is +capped. A host-resident KV cache exists to keep device memory free; a transport +that speeds it up by spending hundreds of MiB of that memory is working against +the thing it is accelerating. `--kv-pipeline-budget` (default 128 MiB) is an +absolute cap on the ring, not a fraction of what happens to be free: + +- Under the cap the ring is allocated and the deliveries pipeline: 4,096 and + 16,384 in the table, at 28 MiB and 104 MiB. +- Over it the scheduler declines and keeps the ordered path, and the decision is + latched, because a context only grows and a ring allocated for the small + windows of early prefill would only have to be given back later. +- Declining costs nothing in steady state. Both the ring and the transfer + backend's device context are released: at 32,768 with the default budget, + device memory settles at 10,161 MiB, the same as the ordered path, and + throughput matches it (12.965 against 12.984 t/s). + +Raising the budget trades that memory back for speed where it is worth it: +`--kv-pipeline-budget 512` at 32,768 gives 15.487 t/s for 206 MiB. + +**Known limitation.** The cap is applied per graph, so a run whose context grows +past it still allocates a ring for the early prefill graphs and releases it once +the window outgrows the budget -- at 32,768 that shows up as a transient peak of ++112 MiB even though the steady state is +0. Deciding against the context's final +size rather than the current graph's would remove it, and needs the KV geometry +the scheduler does not have. + +Where the split-loop host time goes, per decode graph at 18.5k +(`GGML_SCHED_TRANSPORT_DEBUG=2`): + +| | ordered | pipelined | +|---|---:|---:| +| total | 52.11 ms | 30.37 ms | +| blocked in the ordered `ggml_backend_tensor_copy` | 26.80 ms | 0.15 ms | +| blocked waiting for the consumer backend | 25.14 ms | 26.69 ms | +| issuing early deliveries | 0.00 ms | 0.04 ms | +| bytes delivered early / late | 0 / 0 MiB | 619.2 / 1.3 MiB | + +The blocking host-to-device copy is gone and the consumer wait is unchanged, +which is the shape a working overlap has: the transfer left the host's critical +path without being added to the consumer's. + +**32,768 is the weak point and is reported as such.** The gain there is 74.6% of +the probe ceiling, against 90-96% at shallower depths. At that depth the copy per +staged split (about 2.9 ms) exceeds the compute between staged splits (about +1.9 ms), so one split of look-ahead cannot cover it. Raising the look-ahead does +not help: at 32,768 `N = 2` measured 15.5274 and `N = 3` measured 15.1114 against +15.5273 for `N = 1`, and at 16,384 the same sweep gave 30.34 and 29.15 against +31.56. `N = 1` is the best setting at every depth measured, which is why it is +the default. Closing the 32,768 gap is a separate piece of work, not a knob. + +Do not compare these numbers against runs on other models, prompts, cache +settings, hardware, or commits. + +## Validation + +The gates, and what was run for them: + +1. **Byte-identical greedy server output against the control.** Four fixed tasks + at `temperature 0, top_k 1, seed 1234`, plus two tasks behind an 18,422-token + prompt, hashed and compared against a build of the parent commit. Identical at + `N = 0`, `N = 1` and `N = 4`. `docs/repro/r4-kv-pipeline-exact.sh`. +2. **A/B/A/B at 4,096 / 16,384 / 32,768 with reversed arm order.** + `docs/repro/r4-kv-pipeline-ab.sh`; the table above is its output. +3. **Device allocation high-water reported.** Above. +4. **Telemetry showing the deliveries actually converted.** + `GGML_SCHED_TRANSPORT_DEBUG=1` reports the plan (staged splits, bytes per + graph, how much of it goes early, and the source buffer type); + `=2` adds the per-graph host-time breakdown above. + `ggml_backend_sched_get_transport_pipeline_stats()` exposes the same counters + to callers. + +A device-resident KV run is unaffected, and was measured to confirm it: 39.13 t/s +on the parent commit against 39.10 t/s here at `tg128 @ d4096`, with the +transport never enabled because the scheduler is given a depth of 0. + +## Scope and limits + +- Only inputs carrying a stable prefix are eligible. Everything else -- weights, + user inputs, a transposed V cache, an input copy with a reader in a later + split, any backend that cannot transfer asynchronously or record events -- + keeps the ordered path untouched. +- The ring costs `(depth + 2) x (largest staged split)` of device memory, and a + staged split is both K and V of one attention layer over the whole context. That + is linear in context length, and it is what bounds the feature at depth rather + than anything about the transfer itself. +- **One ring, on the first backend that qualifies. Additional accelerators keep + the ordered path, and nothing here has been measured on more than one GPU** -- + every number in this document is `-sm none -mg 0` on a single RTX 4070. A + multi-GPU host-resident cache needs its own validation before any of this is + claimed for it. +- The scheduler must be configured with the device's own default buffer type. A + scheduler built on a split or host buffer type keeps the ordered path. +- `GGML_KV_PIPELINE_DEPTH` overrides the depth for tools that do not expose the + command-line option, such as `llama-bench`. diff --git a/docs/repro/r4-kv-pipeline-ab.sh b/docs/repro/r4-kv-pipeline-ab.sh new file mode 100755 index 00000000000..27e8b0ed302 --- /dev/null +++ b/docs/repro/r4-kv-pipeline-ab.sh @@ -0,0 +1,31 @@ +#!/bin/bash +# R4: pipelined delivery of a host-resident KV cache, A/B/A/B with reversed arm order. +# The two arms are the same binary: GGML_KV_PIPELINE_DEPTH=0 is the ordered path. +# +# LLAMA_KV_MODEL=/path/model.gguf docs/repro/r4-kv-pipeline-ab.sh [depth ...] +set -u +MODEL="${LLAMA_KV_MODEL:?set LLAMA_KV_MODEL to a .gguf path}" +BUILD="${LLAMA_KV_BUILD:-build}" +PIN="${LLAMA_KV_TASKSET:-0,2,4}" +LOCK=/tmp/beellama-single-gpu.lock + +run () { # $1 label, $2 pipeline depth, $3 context depth, $4 reps + GGML_KV_PIPELINE_DEPTH=$2 taskset -c "$PIN" "$BUILD/bin/llama-bench" -m "$MODEL" \ + -ngl 99 -sm none -mg 0 -t 3 -nkvo 1 --kv-cpu-pinned --recurrent-state-offload \ + -fa on -ctk q8_0 -ctv q8_0 -b 512 -ub 512 --no-warmup -p 0 -n 128 -d "$3" -r "$4" -o json \ + 2>/dev/null \ + | python3 -c "import json,sys;d=json.load(sys.stdin);print(' %-12s %.4f +- %.4f'%('$1',d[0]['avg_ts'],d[0]['stddev_ts']))" +} + +DEPTHS=(4096 16384 32768); [ $# -gt 0 ] && DEPTHS=("$@") +rc=0 +for D in "${DEPTHS[@]}"; do + R=3; [ "$D" -le 4096 ] && R=5 + echo "== context depth=$D reps=$R" + flock "$LOCK" bash -c "$(declare -f run); BUILD='$BUILD'; MODEL='$MODEL'; PIN='$PIN' + run ordered 0 $D $R + run pipelined 1 $D $R + run ordered2 0 $D $R + run pipelined2 1 $D $R" || rc=$? +done +exit $rc diff --git a/docs/repro/r4-kv-pipeline-context-sweep.sh b/docs/repro/r4-kv-pipeline-context-sweep.sh new file mode 100755 index 00000000000..b8a2b91bf72 --- /dev/null +++ b/docs/repro/r4-kv-pipeline-context-sweep.sh @@ -0,0 +1,42 @@ +#!/bin/bash +# R4 across context depth: throughput and device allocation high-water, ordered against +# pipelined, on the same binary. The ring holds one split's whole delivery per slot, so its +# cost grows with the context; this is what measures where that stops being affordable. +# +# LLAMA_KV_MODEL=/path/model.gguf docs/repro/r4-kv-pipeline-context-sweep.sh [depth ...] +set -u +MODEL="${LLAMA_KV_MODEL:?set LLAMA_KV_MODEL to a .gguf path}" +BUILD="${LLAMA_KV_BUILD:-build}" +PIN="${LLAMA_KV_TASKSET:-0,2,4}" +NGEN="${LLAMA_KV_NGEN:-64}" +LOCK=/tmp/beellama-single-gpu.lock + +arm () { # $1 pipeline depth, $2 context depth, $3 reps + local vram; vram=$(mktemp) + ( while true; do nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits; sleep 0.25; done ) > "$vram" 2>/dev/null & + local sampler=$! + local ts + ts=$(GGML_KV_PIPELINE_DEPTH=$1 taskset -c "$PIN" "$BUILD/bin/llama-bench" -m "$MODEL" \ + -ngl 99 -sm none -mg 0 -t 3 -nkvo 1 --kv-cpu-pinned --recurrent-state-offload \ + -fa on -ctk q8_0 -ctv q8_0 -b 512 -ub 512 --no-warmup -p 0 -n "$NGEN" -d "$2" -r "$3" -o json \ + 2>/dev/null \ + | python3 -c "import json,sys +try: + d=json.load(sys.stdin); print('%.4f'%d[0]['avg_ts']) +except Exception: + print('FAILED')") + kill $sampler 2>/dev/null; wait $sampler 2>/dev/null + printf ' %-10s %-10s %s MiB\n' "depth=$1" "$ts" "$(sort -n "$vram" | tail -1)" + rm -f "$vram" +} + +DEPTHS=(4096 16384 32768 65536 131072 262144); [ $# -gt 0 ] && DEPTHS=("$@") +for D in "${DEPTHS[@]}"; do + R=3; [ "$D" -gt 32768 ] && R=1 + echo "== context depth=$D reps=$R (t/s, peak device memory)" + flock "$LOCK" bash -c "$(declare -f arm); BUILD='$BUILD'; MODEL='$MODEL'; PIN='$PIN'; NGEN='$NGEN' + arm 0 $D $R + arm 1 $D $R + arm 0 $D $R + arm 1 $D $R" +done diff --git a/docs/repro/r4-kv-pipeline-exact.py b/docs/repro/r4-kv-pipeline-exact.py new file mode 100644 index 00000000000..0a0800f3f79 --- /dev/null +++ b/docs/repro/r4-kv-pipeline-exact.py @@ -0,0 +1,70 @@ +# Greedy server output, hashed, over several prefill corpora and prefill lengths. +# Run through r4-kv-pipeline-exact.sh. Compare the hashes across pipeline depths and against a +# build of the parent commit: the pipelined path must reproduce the ordered path exactly. +import hashlib, json, sys, urllib.request + +PORT = sys.argv[1] +LENGTHS = [int(x) for x in sys.argv[2].split(",")] # approximate prefill tokens + +# Four corpora with different token statistics, so that the deliveries being pipelined are not +# always the same shape of content: prose, source code, structured records, and dialogue. +CORPORA = { + "prose": ("A B-tree index stores keys in sorted order across a shallow, balanced tree. " + "Range queries descend once to the first qualifying leaf and then walk the leaf " + "chain sequentially, so the cost is one descent plus the size of the range. "), + "code": ("static int walk_leaf_chain(struct btree *t, uint64_t lo, uint64_t hi, " + "int (*cb)(void *, uint64_t), void *ctx) {\n" + " struct leaf *l = btree_descend(t, lo);\n" + " while (l && l->keys[0] <= hi) {\n" + " for (int i = 0; i < l->n; i++) { if (l->keys[i] > hi) return 0; " + "cb(ctx, l->keys[i]); }\n" + " l = l->next;\n }\n return 0;\n}\n"), + "records": ('{"id":%d,"region":"eu-central","bytes":918273,"status":"ok",' + '"latency_ms":12.75,"tags":["index","range","btree"]}\n'), + "dialogue": ("Q: Why does the planner prefer a sequential scan here?\n" + "A: Because the predicate matches most of the table, and random leaf access " + "would cost more than reading every page once.\n"), +} + +QUESTIONS = { + "prose": "Summarise the text above in exactly five sentences.", + "code": "Describe what the function above does, then name one bug it could hide.", + "records": "How many distinct fields does each record above have, and what are they?", + "dialogue": "State the single claim the answers above keep returning to.", +} + +def filler(name, target_tokens): + unit = CORPORA[name] + # roughly four characters to the token; the exact prefill length is reported per task + reps = max(1, (target_tokens * 4) // len(unit % 0 if "%d" in unit else unit)) + if "%d" in unit: + return "".join(unit % i for i in range(reps)) + return unit * reps + +def ask(label, prompt, ntok): + body = json.dumps({"model": "m", "messages": [{"role": "user", "content": prompt}], + "max_tokens": ntok, "temperature": 0, "top_k": 1, "seed": 1234}).encode() + req = urllib.request.Request(f"http://127.0.0.1:{PORT}/v1/chat/completions", body, + {"Content-Type": "application/json"}) + try: + with urllib.request.urlopen(req, timeout=14400) as r: + d = json.load(r) + except Exception as e: + print(f"{label} REQUEST_FAILED {type(e).__name__}", flush=True) + return False + m = d["choices"][0]["message"] + # reasoning models put most of the generation in reasoning_content; hash both + text = (m.get("reasoning_content") or "") + "\x00" + (m.get("content") or "") + t = d.get("timings", {}) + print(f"{label:<18} {hashlib.sha256(text.encode()).hexdigest()[:16]} " + f"prompt_n={t.get('prompt_n'):<7} n={t.get('predicted_n'):<4} " + f"pp={t.get('prompt_per_second'):8.2f} tg={t.get('predicted_per_second'):7.3f}", flush=True) + return True + +ok = True +for length in LENGTHS: + ntok = 256 if length <= 4096 else 128 + for name in CORPORA: + prompt = filler(name, length) + "\n\n" + QUESTIONS[name] + ok &= ask(f"{name}@{length}", prompt, ntok) +sys.exit(0 if ok else 1) diff --git a/docs/repro/r4-kv-pipeline-exact.sh b/docs/repro/r4-kv-pipeline-exact.sh new file mode 100755 index 00000000000..6ab8d7b6b61 --- /dev/null +++ b/docs/repro/r4-kv-pipeline-exact.sh @@ -0,0 +1,35 @@ +#!/bin/bash +# R4 gate 1: greedy server output must be byte-identical to the ordered path, across several +# prefill corpora and prefill lengths. Compare the hashes across pipeline depths, and against a +# build of the parent commit. +# +# LLAMA_KV_MODEL=/path/model.gguf docs/repro/r4-kv-pipeline-exact.sh [pipeline-depth ...] +# LLAMA_KV_LENGTHS=2048,18432,65536 selects the prefill lengths (default 2048,18432). +set -u +MODEL="${LLAMA_KV_MODEL:?set LLAMA_KV_MODEL to a .gguf path}" +BUILD="${LLAMA_KV_BUILD:-build}" +PIN="${LLAMA_KV_TASKSET:-0,2,4}" +PORT="${LLAMA_KV_PORT:-18099}" +LENGTHS="${LLAMA_KV_LENGTHS:-2048,18432}" +CTX="${LLAMA_KV_CTX:-32768}" +HERE="$(cd "$(dirname "$0")" && pwd)" + +DEPTHS=(0 1 4); [ $# -gt 0 ] && DEPTHS=("$@") +rc=0 +for D in "${DEPTHS[@]}"; do + echo "== pipeline depth=$D ctx=$CTX prefill lengths=$LENGTHS" + LOG=$(mktemp /tmp/r4-kv-pipeline.XXXX.log) + GGML_KV_PIPELINE_DEPTH=$D taskset -c "$PIN" "$BUILD/bin/llama-server" -m "$MODEL" \ + -ngl 99 -sm none -mg 0 -t 3 -nkvo --kv-cpu-pinned --recurrent-state-offload \ + -fa on -ctk q8_0 -ctv q8_0 -b 512 -ub 512 -c "$CTX" --parallel 1 \ + --host 127.0.0.1 --port "$PORT" --no-warmup > "$LOG" 2>&1 & + SRV=$! + for _ in $(seq 1 600); do + curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1 && break + sleep 1 + done + python3 "$HERE/r4-kv-pipeline-exact.py" "$PORT" "$LENGTHS" || rc=$? + kill $SRV 2>/dev/null; wait $SRV 2>/dev/null + rm -f "$LOG" +done +exit $rc diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index 7ff367079cf..43877fd0794 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -371,6 +371,35 @@ extern "C" { GGML_API void ggml_backend_sched_request_shared_buffer_shrink( ggml_backend_sched_t sched); + // Pipelined delivery of host-resident split inputs. + // + // Without it, a split that reads a host-resident input pays copy + compute in series: + // the transfer is issued on the consumer's own stream immediately before the kernels + // that read it. With it, the scheduler keeps a ring of `depth` staging slots outside + // the graph allocator's reach and issues the stable prefix of a later split's inputs on + // a separate transfer stream while the current split computes, so the transfer retires + // underneath the kernels. + // + // Only inputs that carry a stable prefix (ggml_set_stable_prefix) are eligible: without + // one, the scheduler cannot know that an earlier split of the same graph will not still + // write the bytes it would deliver ahead of time. Everything else keeps the ordered path. + // + // `depth` is how many splits ahead deliveries run; 0 disables pipelining. The ring holds a + // couple of slots more than that, so that recycling a slot never has to wait for a reader + // that is still running. Requires a destination backend with asynchronous transfers and + // events; where that is missing the setting is ignored. Costs roughly (depth + 2) * + // (largest staged split) of device memory. Must be called before the first graph is + // allocated. + GGML_API void ggml_backend_sched_set_transport_pipeline_depth(ggml_backend_sched_t sched, int depth); + + // 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 094852cd4b4..d7924c9bc0f 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -734,11 +734,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 7edce750b58..c9573df8e5b 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -434,7 +434,10 @@ struct ggml_backend_meta_buffer_context { // FIXME // The size of the split state cache is unbounded and can theoretically grow infinitely large. // However, it is also expensive to build and clearing it on every rebuild in ggml_backend_meta_graph_compute is too expensive. - static constexpr size_t nbtc = GGML_TENSOR_SIZE - sizeof(ggml_tensor::padding); + // ggml_tensor::stable_prefix is a hint for the backend scheduler that this backend does + // not consume, and it changes from graph to graph, so keep it out of the compared image + // together with the trailing padding. + static constexpr size_t nbtc = offsetof(ggml_tensor, stable_prefix); std::map, std::pair> split_state_cache; int debug; diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 6b114b403a3..9b500afdd1b 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -14,6 +14,7 @@ #include "ggml-impl.h" #include +#include #include #include #include @@ -786,6 +787,112 @@ static bool ggml_backend_sched_allows_bufferless_kvarn_src( #define GGML_SCHED_MAX_COPIES 4 #endif +#ifndef GGML_SCHED_MAX_TRANSPORT_SLOTS +#define GGML_SCHED_MAX_TRANSPORT_SLOTS 16 +#endif + +// How many slots the transport ring keeps behind the look-ahead. A delivery that runs L splits +// ahead recycles the slot of the split L - n_slots back, so with n_slots == L + 1 it would +// recycle the split that was enqueued a moment ago and is still running, and every delivery +// would have to wait for the consumer to catch up -- the ordered path with extra steps. Two +// slots of margin put the recycled reader far enough behind to have finished. +#ifndef GGML_SCHED_TRANSPORT_MARGIN +#define GGML_SCHED_TRANSPORT_MARGIN 2 +#endif + +// 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; @@ -846,6 +953,9 @@ struct ggml_backend_sched { bool op_offload; + // pipelined delivery of host-resident split inputs + struct ggml_backend_sched_transport transport; + int debug; // used for debugging graph reallocations [GGML_SCHED_DEBUG_REALLOC] @@ -1633,6 +1743,542 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra } } +// +// pipelined transport of host-resident split inputs +// + +static void ggml_backend_sched_transport_teardown(ggml_backend_sched_t sched); + +static bool ggml_backend_sched_transport_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++) { @@ -1652,6 +2298,10 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { } } + // lay out the transport ring and point the staged input copies at it before the graph is + // allocated, so ggml-alloc sees those copies as already allocated and leaves them alone + ggml_backend_sched_transport_plan(sched); + // allocate graph if (backend_ids_changed || !ggml_gallocr_alloc_graph(sched->galloc, &sched->graph)) { #ifndef NDEBUG @@ -1671,6 +2321,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]); } @@ -1693,6 +2348,34 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s std::vector ids; std::vector used_ids; + struct ggml_backend_sched_transport * tr = &sched->transport; + // a reused graph keeps the plan that was made for it, so the split list it describes must be + // the one about to run + int n_inputs_now = 0; + for (int i = 0; i < sched->n_splits; i++) { + n_inputs_now += splits[i].n_inputs; + } + const bool staged = tr->n_staged > 0 && tr->plan_n_splits == sched->n_splits && + tr->plan_n_inputs == n_inputs_now; + + // the plan outlives the graph it was made for -- a reused graph keeps it -- so the cursor + // that walks it has to start over on every evaluation, not only when the plan is rebuilt + // prime every ring before the first consumer runs; from here on deliveries are issued only + // after a split has been enqueued, never before, so that recycling a slot can never hold + // back work the consumer could already be running + if (staged) { + for (int i = 0; i < sched->n_backends; i++) { + tr->rings[i].consumed = 0; + tr->rings[i].scan_cursor = 0; + } + for (int i = 0; i < sched->n_backends; i++) { + ggml_backend_sched_transport_prefetch(sched, i); + } + } + + const bool timing = tr->debug >= 2; + const int64_t t_graph_start = timing ? ggml_time_us() : 0; + for (int split_id = 0; split_id < sched->n_splits; split_id++) { struct ggml_backend_sched_split * split = &splits[split_id]; int split_backend_id = split->backend_id; @@ -1704,6 +2387,23 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s struct ggml_tensor * input = split->inputs[input_id]; struct ggml_tensor * input_cpy = tensor_copy(input, split_backend_id, sched->cur_copy); + if (staged && ggml_backend_sched_input_is_staged(sched, split_id, input_id)) { + // whatever prefix was stable went out on the transfer stream earlier; the rest is + // what an earlier split of this graph may still have written, and it is only safe + // to read now that every earlier split has run. It goes on the consumer's own + // stream, where it is already ordered ahead of the kernels and behind the reader + // of whatever occupied this slot before, and where it cannot push the transfer + // stream's look-ahead behind it. + const size_t prefix = ggml_backend_sched_input_stable_prefix(input); + const size_t nbytes = ggml_nbytes(input); + if (nbytes > prefix) { + ggml_backend_tensor_set_async(split_backend, input_cpy, + (const char *) input->data + prefix, prefix, nbytes - prefix); + tr->n_bytes_late += nbytes - prefix; + } + continue; + } + if (input->flags & GGML_TENSOR_FLAG_INPUT) { // inputs from the user must be copied immediately to prevent the user overwriting the data before the copy is done if (sched->events[split_backend_id][sched->cur_copy] != NULL) { @@ -1714,11 +2414,15 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s ggml_backend_tensor_copy(input, input_cpy); } else { // wait for the split backend to finish using the input before overwriting it + const int64_t t_sync0 = timing ? ggml_time_us() : 0; if (sched->events[split_backend_id][sched->cur_copy] != NULL) { ggml_backend_event_wait(split_backend, sched->events[split_backend_id][sched->cur_copy]); } else { ggml_backend_synchronize(split_backend); } + if (timing) { + tr->t_sync_us += ggml_time_us() - t_sync0; + } // when offloading MoE weights, we can reduce the amount of data copied by copying only the experts that are used ggml_tensor * node = split->graph.nodes[0]; @@ -1810,18 +2514,31 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s // TODO: add public function to facilitate this, since applications do not have direct access to the backend interface const bool async_ok = split_backend->iface.cpy_tensor_async && split_backend->iface.cpy_tensor_async(input_backend, split_backend, input, input_cpy); if (!async_ok) { + const int64_t t_sync1 = timing ? ggml_time_us() : 0; ggml_backend_synchronize(input_backend); if (sched->events[split_backend_id][sched->cur_copy] != NULL) { ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]); } else { ggml_backend_synchronize(split_backend); } + const int64_t t_copy0 = timing ? ggml_time_us() : 0; ggml_backend_tensor_copy(input, input_cpy); + if (timing) { + tr->t_sync_us += t_copy0 - t_sync1; + tr->t_copy_us += ggml_time_us() - t_copy0; + } } } } } + // order the consumer behind this split's early deliveries + struct ggml_backend_sched_transport_slot * slot = NULL; + if (staged && tr->split_order[split_id] >= 0) { + slot = &tr->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) { @@ -1863,6 +2580,20 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } } + // every kernel that reads this split's slot is enqueued, so the slot may be refilled + // once the consumer stream reaches this point + if (slot != NULL) { + ggml_backend_event_record(slot->release, split_backend); + slot->release_armed = true; + tr->rings[split_backend_id].consumed++; + tr->n_deliveries++; + + // run the transfer stream ahead of the consumer: with this split's kernels already + // enqueued, the deliveries for the next few staged splits can go out even if + // recycling their slot has to wait for a reader that is now running + ggml_backend_sched_transport_prefetch(sched, split_backend_id); + } + // record the event of this copy if (split->n_inputs > 0) { if (sched->events[split_backend_id][sched->cur_copy] != NULL) { @@ -1871,6 +2602,28 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } } + if (timing) { + tr->t_graph_us += ggml_time_us() - t_graph_start; + if (++tr->n_graphs % 64 == 0) { + const double n = 64.0; + const double d_graph = (tr->t_graph_us - tr->p_graph_us)/1e3/n; + const double d_sync = (tr->t_sync_us - tr->p_sync_us) /1e3/n; + const double d_copy = (tr->t_copy_us - tr->p_copy_us) /1e3/n; + const double d_issue = (tr->t_issue_us - tr->p_issue_us)/1e3/n; + GGML_LOG_INFO("%s: %" PRId64 " graphs, last 64: %.2f ms/graph = sync %.2f + ordered copy %.2f + " + "issue %.2f + rest %.2f; %.1f MiB early, %.1f MiB late\n", __func__, tr->n_graphs, + d_graph, d_sync, d_copy, d_issue, d_graph - d_sync - d_copy - d_issue, + (tr->n_bytes_early - tr->p_bytes_early)/1048576.0/n, + (tr->n_bytes_late - tr->p_bytes_late) /1048576.0/n); + tr->p_graph_us = tr->t_graph_us; + tr->p_sync_us = tr->t_sync_us; + tr->p_copy_us = tr->t_copy_us; + tr->p_issue_us = tr->t_issue_us; + tr->p_bytes_early = tr->n_bytes_early; + tr->p_bytes_late = tr->n_bytes_late; + } + } + return GGML_STATUS_SUCCESS; } @@ -1941,15 +2694,130 @@ ggml_backend_sched_t ggml_backend_sched_new( sched->galloc = ggml_gallocr_new_n(sched->bufts, n_backends); sched->op_offload = op_offload; + sched->transport.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; + } + ggml_backend_sched_reset(sched); 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]); @@ -2086,6 +2954,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 81f2a8ae0e4..deff8604443 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1379,6 +1379,17 @@ size_t ggml_nbytes_pad(const struct ggml_tensor * tensor) { return GGML_PAD(ggml_nbytes(tensor), GGML_MEM_ALIGN); } +void ggml_set_stable_prefix(struct ggml_tensor * tensor, size_t nbytes) { + GGML_ASSERT(tensor); + const size_t total = ggml_nbytes(tensor); + tensor->stable_prefix = nbytes < total ? nbytes : total; +} + +size_t ggml_get_stable_prefix(const struct ggml_tensor * tensor) { + GGML_ASSERT(tensor); + return tensor->stable_prefix; +} + int64_t ggml_blck_size(enum ggml_type type) { assert(type >= 0); assert(type < GGML_TYPE_COUNT); @@ -1875,7 +1886,7 @@ static struct ggml_tensor * ggml_new_tensor_impl( /*.data =*/ obj_alloc_size > 0 ? (void *)(result + 1) : data, /*.name =*/ { 0 }, /*.extra =*/ NULL, - /*.padding =*/ { 0 }, + /*.stable_prefix=*/ 0, }; // TODO: this should not be needed as long as we don't rely on aligned SIMD loads diff --git a/include/llama.h b/include/llama.h index c2bb6f9bd03..e294a74919e 100644 --- a/include/llama.h +++ b/include/llama.h @@ -516,6 +516,14 @@ extern "C" { // Fit/no-alloc contexts still measure the full prompt reservation. bool live_context_workspace; // grow supported attention workspace plans with the padded live physical KV extent. // false preserves the full-context reservation and initialization order. + uint32_t kv_pipeline_depth; // how many splits ahead the scheduler delivers a host-resident KV cache to the + // accelerator, so that the transfer runs while the previous split computes. + // 0 keeps the ordered path, where a decode token pays the transfer and the + // attention kernels in series. Costs (kv_pipeline_depth + 2) * (largest staged + // split) of device memory. + 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 d97c25ec39f..3b0b16dff2b 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -302,6 +302,8 @@ llama_context::llama_context( cparams.recurrent_state_offload = params.recurrent_state_offload; cparams.phase_aware_workspace = params.phase_aware_workspace; cparams.live_context_workspace = params.live_context_workspace; + cparams.kv_pipeline_depth = params.kv_pipeline_depth; + cparams.kv_pipeline_budget_mib = params.kv_pipeline_budget_mib; cparams.kv_gpu_layers = params.kv_gpu_layers; cparams.offload_attn_compute = params.offload_kqv || (params.op_offload && (params.kv_cpu_pinned || params.kv_gpu_layers > 0)); @@ -1120,6 +1122,11 @@ void llama_context::sched_reserve(uint32_t n_tokens_req, uint32_t n_kv_req) { sched.reset(ggml_backend_sched_new( backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, pipeline_parallel, cparams.op_offload)); + // only a host-resident KV cache produces the deliveries this pipelines + ggml_backend_sched_set_transport_pipeline_budget(sched.get(), + (size_t) cparams.kv_pipeline_budget_mib * 1024 * 1024); + ggml_backend_sched_set_transport_pipeline_depth(sched.get(), + cparams.kv_cpu_pinned || !cparams.offload_kqv ? (int) cparams.kv_pipeline_depth : 0); if (sched_shared_buffers) { ggml_backend_sched_set_shared_buffers(sched.get(), sched_shared_buffers.get()); sched_shared_generation = @@ -4564,6 +4571,8 @@ llama_context_params llama_context_default_params() { /*.kv_gpu_layers =*/ 0, /*.phase_aware_workspace =*/ false, /*.live_context_workspace =*/ false, + /*.kv_pipeline_depth =*/ 1, + /*.kv_pipeline_budget_mib =*/ 128, }; return result; diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 37438bd6448..34ab37519d5 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -65,6 +65,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 4b8a2bbc1f5..fb53112034a 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -2803,6 +2803,10 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & return; } + // before the graph is built and allocated, so that the scheduler's delivery plan and the + // deliveries it then issues are decided against the same write position + update_stable_prefixes(sinfo); + // keep track of the max sequence position that we would overwrite with this ubatch // for non-SWA cache, this would be always empty llama_seq_id seq_pos_max_rm[LLAMA_MAX_SEQ]; @@ -3552,7 +3556,48 @@ ggml_tensor * llama_kv_cache::build_input_v_rot(ggml_context * ctx) const { return res; } +void llama_kv_cache::update_stable_prefixes(const slot_info & sinfo) const { + // Rows written by this ubatch, in the flattened [n_embd_gqa, kv_size*n_stream] body: every + // byte below the lowest of them keeps whatever the previous ubatch left there for the whole + // graph, so a delivery of that region may be issued before the split that reads it. + uint64_t min_row = UINT64_MAX; + for (uint32_t s = 0; s < sinfo.n_stream(); ++s) { + const uint64_t offs = (uint64_t) sinfo.strm[s]*get_size(); + for (const uint32_t idx : sinfo.idxs[s]) { + min_row = std::min(min_row, offs + idx); + } + } + + if (min_row == UINT64_MAX) { + clear_stable_prefixes(); + return; + } + + for (const auto & layer : layers) { + if (layer.k) { + ggml_set_stable_prefix(layer.k, min_row*layer.k->nb[1]); + } + if (layer.v) { + // the transposed V cache scatters each ubatch across the whole tensor, so there is + // no leading region that this ubatch leaves alone + ggml_set_stable_prefix(layer.v, v_trans ? 0 : min_row*layer.v->nb[1]); + } + } +} + +void llama_kv_cache::clear_stable_prefixes() const { + for (const auto & layer : layers) { + if (layer.k) { + ggml_set_stable_prefix(layer.k, 0); + } + if (layer.v) { + ggml_set_stable_prefix(layer.v, 0); + } + } +} + void llama_kv_cache::set_input_k_idxs(ggml_tensor * dst, const llama_ubatch * ubatch, const slot_info & sinfo) const { + const uint32_t n_tokens = ubatch->n_tokens; GGML_ASSERT(n_tokens == (int64_t) sinfo.size()*sinfo.n_stream()); @@ -4237,6 +4282,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 78e2ced63ba..0e9c567c9ea 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -328,6 +328,13 @@ class llama_kv_cache : public llama_memory_i { ggml_tensor * build_input_k_rot(ggml_context * ctx) const; ggml_tensor * build_input_v_rot(ggml_context * ctx) const; + // Tell the backend scheduler which part of each layer's persistent K/V storage this ubatch + // will not write, so a host-resident cache can be delivered to the accelerator ahead of the + // attention that reads it. Must be refreshed for every ubatch, including when the graph is + // reused, because the write position moves while the graph does not. + void update_stable_prefixes(const slot_info & sinfo) const; + void clear_stable_prefixes() const; + bool can_use_compact_causal_mask( const llama_ubatch & ubatch, bool causal_attn, uint32_t n_kv, bool is_reserve,