diff --git a/common/arg.cpp b/common/arg.cpp index ae6db10b6e9..65022db235a 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -30,6 +30,7 @@ #include #include #include +#include #include #include #include @@ -2429,6 +2430,39 @@ 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 > 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; + } + ).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, %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; + 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"); + } + 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 05af5273742..b171ad9b29d 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 = (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/common/common.h b/common/common.h index 74f0f1281c7..ea6d053eab4 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 00000000000..c27e8e9763c --- /dev/null +++ b/docs/kv-transport-pipelining.md @@ -0,0 +1,303 @@ +# 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. + +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). + +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`), at `--kv-pipeline-budget 512`, both passes shown: + +| depth | ordered | pipelined | gain | +|---:|---|---|---:| +| 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. + +> 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`, the four 18,432-prefill tasks of the exactness gate at `-c 32768`: + +| task | prompt | ordered | pipelined | gain | +|---|---:|---:|---:|---:| +| 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`: + +| 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 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 + +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 + +`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 | 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. + +**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. 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. + +**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. + +**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. + +### Parallel sequences + +`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.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 | + +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. + +**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 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 + +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, 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; 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. + +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 + +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**. + +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 + +Per decode graph, `GGML_SCHED_TRANSPORT_DEBUG=2`, behind a 19,246-token prompt: + +| | ordered | pipelined | +|---|---:|---:| +| 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 | 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 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. + +## Validation + +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. + +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, 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. + +## 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. +- **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. +- **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. + +## 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. + +## 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. diff --git a/docs/repro/r4-kv-pipeline-ab.sh b/docs/repro/r4-kv-pipeline-ab.sh new file mode 100755 index 00000000000..d0e50427ab7 --- /dev/null +++ b/docs/repro/r4-kv-pipeline-ab.sh @@ -0,0 +1,63 @@ +#!/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: --kv-pipeline-depth 0 is the ordered path. +# +# LLAMA_KV_MODEL=/path/model.gguf docs/repro/r4-kv-pipeline-ab.sh [depth ...] +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, 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" \ + --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 \ + > "$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) +if [ $# -gt 0 ]; then + DEPTHS=("$@") +fi +for D in "${DEPTHS[@]}"; do + R=3 + if [ "$D" -le 4096 ]; then + 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' + run ordered 0 $D $R + run pipelined 1 $D $R + run ordered2 0 $D $R + run pipelined2 1 $D $R" +done 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..64bac9feba3 --- /dev/null +++ b/docs/repro/r4-kv-pipeline-context-sweep.sh @@ -0,0 +1,71 @@ +#!/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 -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, 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 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=$! + 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" "$out" "$err" +} + +DEPTHS=(4096 16384 32768 65536 131072 262144) +if [ $# -gt 0 ]; then + DEPTHS=("$@") +fi +for D in "${DEPTHS[@]}"; do + 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 "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 + 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..3330893ddfc --- /dev/null +++ b/docs/repro/r4-kv-pipeline-exact.py @@ -0,0 +1,87 @@ +# Greedy server output, hashed, over several prefill corpora and prefill lengths. +# 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. +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 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): + # 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() + 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", {}) + # 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] + 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) + return not reused + +ok = True +for length in LENGTHS: + ntok = 256 if length <= 4096 else 128 + 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 new file mode 100755 index 00000000000..017926d29a0 --- /dev/null +++ b/docs/repro/r4-kv-pipeline-exact.sh @@ -0,0 +1,52 @@ +#!/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. +# +# 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}" +BUDGET="${LLAMA_KV_BUDGET:-512}" +HERE="$(cd "$(dirname "$0")" && pwd)" + +DEPTHS=(0 1 4); [ $# -gt 0 ] && DEPTHS=("$@") +rc=0 +BASE="" +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) + 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 \ + -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 + OUT=$(mktemp /tmp/r4-kv-pipeline.XXXX.hashes) + if python3 "$HERE/r4-kv-pipeline-exact.py" "$PORT" "$LENGTHS" "$OUT"; then + if [ "$I" -eq 0 ]; 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" + if [ "$I" -eq 0 ] && [ -z "$BASE" ]; then + break + fi +done +[ -z "$BASE" ] || rm -f "$BASE" +exit $rc diff --git a/docs/repro/r4-kv-pipeline-parallel-exact.sh b/docs/repro/r4-kv-pipeline-parallel-exact.sh new file mode 100755 index 00000000000..779b0dae8b9 --- /dev/null +++ b/docs/repro/r4-kv-pipeline-parallel-exact.sh @@ -0,0 +1,79 @@ +#!/bin/bash +# R4 gate 2: greedy output of concurrent sequences over a cache split into streams must be byte-identical to the ordered path. +# Gate 1 covers one sequence per ubatch, where a delivery is a single range. This covers a ubatch that spans several streams, where a delivery is one range per stream and the cells between them are never read. +# llama-parallel is used rather than the server because it seeds its client schedule, so the batches are the same run to run and the outputs can be compared directly. +# +# LLAMA_KV_MODEL=/path/model.gguf docs/repro/r4-kv-pipeline-parallel-exact.sh [pipeline-depth ...] +# LLAMA_KV_NP=8 sets the concurrent sequences, LLAMA_KV_NS=16 the total. +# LLAMA_KV_SM=layer spreads the model over every device, which gives each of them its own ring. +set -euo pipefail +MODEL="${LLAMA_KV_MODEL:?set LLAMA_KV_MODEL to a .gguf path}" +BUILD="${LLAMA_KV_BUILD:-build}" +PIN="${LLAMA_KV_TASKSET:-0,2,4}" +CTX="${LLAMA_KV_CTX:-16384}" +BUDGET="${LLAMA_KV_BUDGET:-512}" +NP="${LLAMA_KV_NP:-8}" +NS="${LLAMA_KV_NS:-16}" +SM="${LLAMA_KV_SM:-none}" + +DEPTHS=(0 1 4); [ $# -gt 0 ] && DEPTHS=("$@") + +# Without a pinned host cache, a host-resident recurrent state and a budget the run does not exercise the path the doc reports on. +if ! HELP="$("$BUILD/bin/llama-parallel" --help 2>&1)"; then + echo "cannot run $BUILD/bin/llama-parallel:" >&2 + echo "$HELP" >&2 + exit 1 +fi +for opt in --kv-cpu-pinned --recurrent-state-offload --kv-pipeline-depth --kv-pipeline-budget --no-kv-unified; do + if ! grep -q -- "$opt" <<< "$HELP"; then + echo "$BUILD/bin/llama-parallel has no $opt option" >&2 + exit 1 + fi +done + +# Keep only what the clients produced: drop the log timestamps, the colour codes and the timing summary, all of which differ between runs by design. +transcript () { + sed -e 's/\x1b\[[0-9;]*m//g' \ + -e '/^[0-9][0-9.]* [A-Z] /d' \ + -e '/speed:/d' -e '/^Cache misses/d' "$1" +} + +rc=0 +BASE="" +for I in "${!DEPTHS[@]}"; do + D="${DEPTHS[$I]}" + echo "== pipeline depth=$D ctx=$CTX np=$NP ns=$NS sm=$SM streams (non-unified)" + RAW="$(mktemp /tmp/r4-kv-parallel.XXXX.log)" + OUT="$(mktemp /tmp/r4-kv-parallel.XXXX.txt)" + arm_rc=0 + taskset -c "$PIN" "$BUILD/bin/llama-parallel" -m "$MODEL" \ + --kv-pipeline-depth "$D" --kv-pipeline-budget "$BUDGET" \ + -ngl 99 -sm "$SM" -mg 0 -t 3 -nkvo --kv-cpu-pinned --recurrent-state-offload \ + -fa on -ctk q8_0 -ctv q8_0 -b 512 -ub 512 -c "$CTX" -no-kvu \ + -np "$NP" -ns "$NS" --temp 0 > "$RAW" 2>&1 || arm_rc=$? + if [ "$arm_rc" -ne 0 ]; then + echo " depth $D: FAILED (exit $arm_rc)" >&2 + tail -20 "$RAW" >&2 + rm -f "$RAW" "$OUT" + exit "$arm_rc" + fi + transcript "$RAW" > "$OUT" + if [ ! -s "$OUT" ]; then + echo " depth $D: FAILED (no client output)" >&2 + rm -f "$RAW" "$OUT" + exit 1 + fi + echo " $(sha256sum < "$OUT" | cut -c1-16) $(wc -l < "$OUT") lines" + if [ "$I" -eq 0 ]; then + BASE="$OUT" + else + if ! cmp -s "$BASE" "$OUT"; then + diff -u "$BASE" "$OUT" | head -40 + rc=1 + fi + rm -f "$OUT" + fi + rm -f "$RAW" +done +rm -f "$BASE" +exit $rc diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index 7504779f278..505e1f31df7 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -323,6 +323,30 @@ 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 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, 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. + // 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. + 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. + // 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. + GGML_API void ggml_backend_sched_get_transport_pipeline_stats(ggml_backend_sched_t sched, int64_t * n_deliveries, int64_t * n_bytes_early, int64_t * n_bytes_late); + // Initialize backend buffers from a measure graph GGML_API void ggml_backend_sched_reserve_size(ggml_backend_sched_t sched, struct ggml_cgraph * measure_graph, size_t * sizes); GGML_API bool ggml_backend_sched_reserve(ggml_backend_sched_t sched, struct ggml_cgraph * measure_graph); // returns success diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 56f9740cd60..fdb70ce004f 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -648,11 +648,12 @@ extern "C" { // this tensor... enum ggml_tensor_flag { - GGML_TENSOR_FLAG_INPUT = 1, // ...is an input for the GGML compute graph - GGML_TENSOR_FLAG_OUTPUT = 2, // ...is an output for the GGML compute graph - GGML_TENSOR_FLAG_PARAM = 4, // ...contains trainable parameters - GGML_TENSOR_FLAG_LOSS = 8, // ...defines loss for numerical optimization (multiple loss tensors add up) - GGML_TENSOR_FLAG_COMPUTE = 16, // ...must be computed + GGML_TENSOR_FLAG_INPUT = 1, // ...is an input for the GGML compute graph + GGML_TENSOR_FLAG_OUTPUT = 2, // ...is an output for the GGML compute graph + GGML_TENSOR_FLAG_PARAM = 4, // ...contains trainable parameters + GGML_TENSOR_FLAG_LOSS = 8, // ...defines loss for numerical optimization (multiple loss tensors add up) + GGML_TENSOR_FLAG_COMPUTE = 16, // ...must be computed + GGML_TENSOR_FLAG_TRANSPORT = 32, // ...is persistent host storage that can use split-input transport }; enum ggml_tri_type { @@ -701,11 +702,22 @@ extern "C" { void * extra; // extra things e.g. for ggml-cuda.cu - char padding[8]; + // bytes at the start of every stream that stay unchanged for the current graph evaluation, 0 for none + union { + size_t stable_prefix; + char padding[8]; + }; }; 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); + // Abort callback // If not NULL, called before ggml computation // If it returns true, the computation is aborted diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 8138028ee8c..9acc68208e4 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 @@ -771,6 +772,119 @@ 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. +// 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, 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 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 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. +// 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 + 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, 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 + + 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; // 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 + + bool reported_no_room; +}; + +// 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. +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]; + + // 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 + 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 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; + + 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; + + // why the look-ahead stopped: the depth it was given, or a slot whose reader has not run yet + int64_t n_stop_depth; + 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; + bool named_ordered; // debug >= 3 names them once, they are the same every graph + + int debug; +}; + struct ggml_backend_sched_split { int backend_id; int i_start; @@ -810,6 +924,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; @@ -830,6 +945,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] @@ -1066,6 +1184,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 = { @@ -1598,6 +1717,803 @@ 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; +} + +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; +} + +// 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. +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 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; + + out->n = 1; + out->stride = 0; + out->used = ggml_nbytes(input); + out->early = 0; + + // 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; + } + + out->early = base->stable_prefix < out->used ? base->stable_prefix : out->used; +} + +// 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. +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]; + 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) { + 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; + } + + // 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); + 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; +} + +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; +} + +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. +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; + + 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); + // 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); + 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)); + } + GGML_ASSERT(offset <= r->slot_size); + } +} + +// 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]; + + if (r->buffer == NULL) { + return; + } + + // nothing may still be reading from or writing into the ring + if (r->transfer) { + ggml_backend_synchronize(r->transfer); + } + // 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; + 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; + } +} + +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); + + 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; +} + +// Count one graph that staged nothing on this ring, and give the staging back once there have been a few in a row. +// 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]; + + 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_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: 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; + + 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; + } + + 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; + } + } +} + +// Stop asking this backend for a ring, and give back the device context with it. +// 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 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++) { + 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; + + 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; + } + *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); +} + +// 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. +// 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); + + size_t size = 1; + while (size < need && size <= SIZE_MAX/2) { + size *= 2; + } + size = std::max(std::min(size, limit), need); + size -= size % alignment; + + return std::max(size, need); +} + +// 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]; + + 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 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; + + tr->n_staged = 0; + tr->plan_n_splits = 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; + tr->rings[i].scan_cursor = 0; + } + + if (!ggml_backend_sched_transport_enabled(sched) || sched->n_splits == 0) { + ggml_backend_sched_transport_release_idle(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)); + 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; + } + + 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; + + // 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; + } + ggml_backend_sched_transport_release_idle(sched); + return; + } + + 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__); + ggml_backend_sched_transport_release_idle(sched); + 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_gen = sched->splits_gen; + + 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++; + } + } + } + + if (n_candidates == 0) { + for (int i = 0; i < sched->n_splits; i++) { + tr->split_order[i] = -1; + } + ggml_backend_sched_transport_release_idle(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 + size_t staged_hash_size = n_candidates; + staged_hash_size += staged_hash_size/4 + 1; + 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; + } + + 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_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; + } + } + } + } + + 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; + } + } + } + + // 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 + 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; + + tr->split_order[i] = -1; + + 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; + } + 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_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; + break; + } + } + + if (need == 0 || size_overflow[bid]) { + continue; + } + + tr->split_order[i] = tr->rings[bid].n_staged++; + slot_size[bid] = std::max(slot_size[bid], need); + 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]) { + 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; + } + // 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_ring_idle(sched, bid); + continue; + } + + 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; + } + + 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 " + "--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_decline_backend(sched, bid); + continue; + } + + // 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])); + ggml_backend_sched_transport_disable_backend(sched, bid); + continue; + } + + if (r->buffer == NULL || r->slot_size < slot_size[bid]) { + ggml_backend_sched_transport_free_ring(sched, bid); + + 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); + } + 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__, + ggml_backend_name(sched->backends[bid]), ring_size >> 20, + GGML_SCHED_TRANSPORT_HEADROOM >> 20, dev_free >> 20); + r->reported_no_room = true; + } + ggml_backend_sched_transport_decline_backend(sched, bid); + continue; + } + + // 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); + } + 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, 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); + 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 + 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_disable_backend(sched, bid); + 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; + + 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_alloc >> 10); + } + } + + r->idle_graphs = 0; + tr->n_staged += r->n_staged; + } + + 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); + 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", + __func__, ggml_backend_name(sched->backends[bid]), tr->rings[bid].n_staged, + sched->n_splits, total >> 10, early >> 10, src_buft); + } + } + + 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 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]; + + if (r->n_staged == 0) { + return; + } + + 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[o % tr->n_slots]; + + // the previous occupant of this slot must be read before the slot is overwritten + // 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; + tr->n_wait_recycle++; + } + + 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 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) { + 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); + if (tr->debug >= 2) { + tr->t_issue_us += ggml_time_us() - t0; + } + tr->n_bytes_early += rg.early*rg.n; + } + + // 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 = o + 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++) { @@ -1617,6 +2533,9 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { } } + // 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 if (backend_ids_changed || !ggml_gallocr_alloc_graph(sched->galloc, &sched->graph)) { #ifndef NDEBUG @@ -1636,14 +2555,27 @@ 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]); + sched->transport.rings[i].delivered = false; } + 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 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, 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)) { GGML_LOG_ERROR("%s: failed to allocate graph\n", __func__); return false; @@ -1663,6 +2595,39 @@ 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 + 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. + for (int i = 0; i < sched->n_backends; i++) { + if (!tr->rings[i].delivered && !(staged && tr->rings[i].n_staged > 0)) { + continue; + } + if (tr->rings[i].transfer) { + 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; 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++) { + 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 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; @@ -1671,11 +2636,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 @@ -1684,14 +2653,41 @@ 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)) { + // 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) { + 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; + } + 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)); + named_ordered_now = true; + } } else { // wait for the split backend to finish using the input before overwriting it if (sched->events[split_backend_id][sched->cur_copy] != NULL) { @@ -1789,18 +2785,39 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s // try async copy, but if not possible, we can still use a sync copy without synchronizing the dst backend, since we handle the synchronization here with multiple copies and events // TODO: add public function to facilitate this, since applications do not have direct access to the backend interface if (!split_backend->iface.cpy_tensor_async || !split_backend->iface.cpy_tensor_async(input_backend, split_backend, input, input_cpy)) { + const int64_t t0 = tr->debug >= 2 ? ggml_time_us() : 0; ggml_backend_synchronize(input_backend); if (sched->events[split_backend_id][sched->cur_copy] != NULL) { ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]); } else { ggml_backend_synchronize(split_backend); } + if (tr->debug >= 2) { + tr->t_sync_us += ggml_time_us() - t0; + } + const int64_t t1 = tr->debug >= 2 ? ggml_time_us() : 0; ggml_backend_tensor_copy(input, input_cpy); + if (tr->debug >= 2) { + tr->t_copy_us += ggml_time_us() - t1; + tr->n_bytes_ordered += ggml_nbytes(input); + } + if (tr->debug >= 3 && !tr->named_ordered) { + GGML_LOG_INFO("%s: ordered copy %s %zu KiB from %s\n", __func__, input->name, + ggml_nbytes(input) >> 10, ggml_backend_buft_name(input->buffer->buft)); + named_ordered_now = true; + } } } } } + // 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) { @@ -1840,6 +2857,17 @@ 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 + if (slot != NULL) { + ggml_backend_event_record(slot->release, split_backend); + slot->release_armed = true; + tr->rings[split_backend_id].consumed++; + tr->n_deliveries++; + + // 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); + } + // 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); @@ -1848,9 +2876,82 @@ 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++; + + // 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, " + "issue %.2f ms, early %.1f MiB, late %.1f MiB, ordered %.1f MiB, " + "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, + (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_wait_recycle - tr->p_wait_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->p_stop_depth = tr->n_stop_depth; + tr->p_wait_recycle = tr->n_wait_recycle; + } + } + 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, @@ -1916,17 +3017,155 @@ 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; + } sched->op_offload = op_offload; ggml_backend_sched_reset(sched); + int transport_depth; + if (ggml_backend_sched_transport_depth_from_env(&transport_depth)) { + const bool ok = ggml_backend_sched_set_transport_pipeline_depth(sched, transport_depth); + GGML_ASSERT(ok); + } + 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); + } + sched->transport.n_staged = 0; +} + + +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; + } + + if (depth < 0 || depth > GGML_SCHED_MAX_TRANSPORT_SLOTS - GGML_SCHED_TRANSPORT_MARGIN) { + return false; + } + + ggml_backend_sched_transport_teardown(sched); + for (int i = 0; i < sched->n_backends; i++) { + tr->rings[i].eligible = false; + tr->rings[i].reported_no_room = false; + } + + tr->depth = depth; + tr->n_slots = depth + GGML_SCHED_TRANSPORT_MARGIN; + + if (depth < 1) { + return true; + } + + // 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; + + 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]; + 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; + } + + 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)) { + 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 CUDA backend supports pipelined host transport, staying on the ordered path\n", __func__); + } + + return true; +} + +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; + } + + 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; + } + } + + return 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.ring_split); + 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]); @@ -2018,6 +3257,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; @@ -2055,8 +3296,14 @@ 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]); + sched->transport.rings[i].delivered = false; } if (!sched->is_alloc) { // if the graph is not already allocated, always use copy 0 after a synchronization diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 708b16edd47..f91e2e5598b 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1260,6 +1260,9 @@ static_assert(GGML_GLU_OP_COUNT == 6, "GGML_GLU_OP_COUNT != 6"); static_assert(sizeof(struct ggml_object)%GGML_MEM_ALIGN == 0, "ggml_object size must be a multiple of GGML_MEM_ALIGN"); static_assert(sizeof(struct ggml_tensor)%GGML_MEM_ALIGN == 0, "ggml_tensor size must be a multiple of GGML_MEM_ALIGN"); +static_assert(sizeof(((struct ggml_tensor *) 0)->padding) == 8, "ggml_tensor trailing storage must be 8 bytes"); +static_assert(sizeof(((struct ggml_tensor *) 0)->stable_prefix) <= sizeof(((struct ggml_tensor *) 0)->padding), "stable_prefix must fit in trailing storage"); +static_assert(offsetof(struct ggml_tensor, stable_prefix) + sizeof(((struct ggml_tensor *) 0)->padding) == sizeof(struct ggml_tensor), "ggml_tensor trailing storage must remain last"); //////////////////////////////////////////////////////////////////////////////// @@ -1323,6 +1326,19 @@ 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); + // 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) { + GGML_ASSERT(tensor); + return tensor->stable_prefix; +} + int64_t ggml_blck_size(enum ggml_type type) { assert(type >= 0); assert(type < GGML_TYPE_COUNT); diff --git a/include/llama.h b/include/llama.h index 27c9a74846e..3a10d2f3369 100644 --- a/include/llama.h +++ b/include/llama.h @@ -38,6 +38,11 @@ #define LLAMA_TOKEN_NULL -1 +// bounds of llama_context_params::kv_pipeline_depth and ::kv_pipeline_budget_mib +// a staging slot holds one attention layer's K or V, so a budget of 64 GiB is already uncapped: past it the value is a typo, not a budget +#define LLAMA_KV_PIPELINE_DEPTH_MAX 14 +#define LLAMA_KV_PIPELINE_BUDGET_MIB_MAX 65536 + #define LLAMA_FILE_MAGIC_GGLA 0x67676c61u // 'ggla' #define LLAMA_FILE_MAGIC_GGSN 0x6767736eu // 'ggsn' #define LLAMA_FILE_MAGIC_GGSQ 0x67677371u // 'ggsq' @@ -420,6 +425,12 @@ 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, 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 { @@ -740,6 +751,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 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 ca2218d8100..bfb14cf6146 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -145,6 +145,14 @@ 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; + 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; @@ -865,10 +873,16 @@ 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; 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()); + // 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"); + } 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()); @@ -4071,6 +4085,8 @@ llama_context_params llama_context_default_params() { /*.sampler =*/ nullptr, /*.n_sampler =*/ 0, /*.ctx_other =*/ nullptr, + /*.kv_pipeline_depth =*/ 1, + /*.kv_pipeline_budget_mib =*/ 128, }; return result; diff --git a/src/llama-cparams.h b/src/llama-cparams.h index a51edaa6acf..0555d549e6d 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -60,6 +60,8 @@ struct llama_cparams { bool recurrent_state_offload; bool phase_aware_workspace; bool live_context_workspace; + uint32_t kv_pipeline_depth; + uint32_t kv_pipeline_budget_mib; std::vector embeddings_layer_inp; // [n_layer()] extract input embeddings for layer diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index ac7908b9946..d81b1baeec3 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -264,6 +264,14 @@ llama_kv_cache::llama_kv_cache( layers.push_back(layer_share); layers.back().il = il; + // this cache writes the shared tensors at its own slot, so their stable prefix would have two writers: keep them on the ordered path + if (layers.back().k) { + layers.back().k->flags &= ~GGML_TENSOR_FLAG_TRANSPORT; + } + if (layers.back().v) { + layers.back().v->flags &= ~GGML_TENSOR_FLAG_TRANSPORT; + } + continue; } } @@ -315,6 +323,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)) { @@ -469,6 +486,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); } @@ -844,9 +866,16 @@ 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, so this is set before a delivery can be in flight + set_lctx(lctx); + bool do_shift = get_has_shift(); return std::make_unique(this, lctx, do_shift, std::move(sc_info)); @@ -1200,10 +1229,14 @@ 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 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 + 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]; @@ -1636,6 +1669,43 @@ 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 { + // 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]) { + min_row = std::min(min_row, (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]); + } + } +} + +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()); @@ -2315,6 +2385,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 af765830d01..39b6eab4c1e 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; @@ -238,6 +241,11 @@ 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 + // 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; @@ -305,6 +313,11 @@ 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 + // 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 const llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE; diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index d65edf9dc85..fb122cf4a3a 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -129,6 +129,11 @@ llama_memory_context_ptr llama_memory_hybrid_idx::init_full() { } llama_memory_context_ptr llama_memory_hybrid_idx::init_update(llama_context * lctx, bool optimize) { + // the indexer builds no update graph, so it gets no context of its own, but clear() still has to wait for a decode that reads its buffers + if (mem_idx) { + mem_idx->set_lctx(lctx); + } + return std::make_unique(this, lctx, optimize); } diff --git a/tests/test-alloc.cpp b/tests/test-alloc.cpp index 879cb953f24..78f2aadcb4d 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -5,8 +5,11 @@ #include "ggml.h" #include +#include +#include #include #include +#include #include // @@ -16,10 +19,45 @@ 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; 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 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; + size_t alloc_size_pad = 0; + + // 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; + 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 = {}; ggml_backend_buffer_i buffer_interface; std::vector buffers; @@ -43,7 +81,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(); @@ -63,8 +101,14 @@ 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 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; } // ggml_backend_buffer interface @@ -88,7 +132,17 @@ static void * dummy_backend_buffer_get_base(ggml_backend_buffer_t buffer) { return ctx->buffer_bases[i - ctx->buffers.begin()]; } -static ggml_status dummy_backend_buffer_init_tensor(ggml_backend_buffer_t, ggml_tensor *) { +static ggml_status dummy_backend_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { + dummy_backend_context * ctx = (dummy_backend_context *) buffer->context; + const size_t size = ggml_backend_buffer_get_alloc_size(buffer, tensor); + + for (auto & b : ctx->bindings) { + if (b.tensor == tensor) { + b = { tensor, buffer, (const char *) tensor->data, size }; + return GGML_STATUS_SUCCESS; + } + } + ctx->bindings.push_back({ tensor, buffer, (const char *) tensor->data, size }); return GGML_STATUS_SUCCESS; } @@ -104,13 +158,39 @@ static void dummy_backend_buffer_clear(ggml_backend_buffer_t, uint8_t) {} struct dummy_backend { std::unique_ptr context; + std::unique_ptr registry; std::unique_ptr device; std::unique_ptr handle; ggml_backend_buffer_type buffer_type; }; -static const char * dummy_backend_get_name(ggml_backend_t) { - return "dummy_backend"; +static const char * dummy_backend_get_name(ggml_backend_t backend) { + return ((dummy_backend_context *) backend->context)->registry_name; +} + +static void dummy_backend_free(ggml_backend_t backend) { + dummy_backend_context * ctx = (dummy_backend_context *) backend->context; + ctx->transfer_backend_count--; + delete backend; +} + +static void dummy_backend_set_tensor_async(ggml_backend_t backend, ggml_tensor * 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 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 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 *) { @@ -119,8 +199,52 @@ 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; + ctx->transfer_backend_inits++; + 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 *) { @@ -131,12 +255,21 @@ static bool dummy_backend_device_supports_buft(ggml_backend_dev_t device, ggml_b return device->context == buft->context; } -static dummy_backend dummy_backend_init(size_t max_buffer_size, size_t alignment = 8, bool unique_alloc_addresses = false) { +static dummy_backend dummy_backend_init( + size_t max_buffer_size, + size_t alignment = 8, + bool unique_alloc_addresses = false, + enum ggml_backend_dev_type device_type = GGML_BACKEND_DEVICE_TYPE_CPU, + const char * registry_name = "dummy", + bool buffer_is_host = true) { dummy_backend b{}; b.context = std::make_unique(); b.context->alignment = alignment; b.context->max_buffer_size = max_buffer_size; b.context->unique_alloc_addresses = unique_alloc_addresses; + b.context->device_type = device_type; + b.context->registry_name = registry_name; + b.context->buffer_is_host = buffer_is_host; b.context->buffer_interface.free_buffer = dummy_backend_buffer_free_buffer; b.context->buffer_interface.get_base = dummy_backend_buffer_get_base; @@ -151,21 +284,40 @@ static dummy_backend dummy_backend_init(size_t max_buffer_size, size_t alignment 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; + + 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; } @@ -189,6 +341,64 @@ 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 }; +} + +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, + 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); @@ -1127,6 +1337,668 @@ static void test_resizable_buffers_owner_borrower_teardown_order() { } } +static void test_transport_prefix_and_configuration() { + dummy_backend cuda = dummy_backend_init(SIZE_MAX, 8, true, GGML_BACKEND_DEVICE_TYPE_GPU, "CUDA", false); + dummy_backend cpu = dummy_backend_init(SIZE_MAX, 8, true); + auto graph = make_transport_graph(cpu, 64); + + ggml_backend_t backends[] = { cuda.handle.get(), cpu.handle.get() }; + ggml_backend_buffer_type_t bufts[] = { &cuda.buffer_type, &cpu.buffer_type }; + { + ggml_backend_sched_ptr sched(ggml_backend_sched_new(backends, bufts, 2, 128, false, false)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_budget(sched.get(), 1024)); + GGML_ASSERT(ggml_backend_sched_set_transport_pipeline_depth(sched.get(), 1)); + ggml_backend_sched_set_tensor_backend(sched.get(), graph.output, cuda.handle.get()); + + ggml_set_stable_prefix(graph.source, 32); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), graph.ctx.graph)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), graph.ctx.graph) == GGML_STATUS_SUCCESS); + + int64_t deliveries = 0; + int64_t early = 0; + int64_t late = 0; + transport_stats(sched.get(), &deliveries, &early, &late); + GGML_ASSERT(deliveries == 1 && early == 32 && late == 32); + GGML_ASSERT(!ggml_backend_sched_set_transport_pipeline_depth(sched.get(), 0)); + GGML_ASSERT(!ggml_backend_sched_set_transport_pipeline_budget(sched.get(), 0)); + + ggml_set_stable_prefix(graph.source, 0); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), graph.ctx.graph) == GGML_STATUS_SUCCESS); + transport_stats(sched.get(), &deliveries, &early, &late); + GGML_ASSERT(deliveries == 2 && early == 32 && late == 96); + + ggml_set_stable_prefix(graph.source, 64); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched.get(), graph.ctx.graph) == GGML_STATUS_SUCCESS); + transport_stats(sched.get(), &deliveries, &early, &late); + GGML_ASSERT(deliveries == 3 && early == 96 && late == 96); + GGML_ASSERT(cuda.context->event_wait_count > 0); + } + GGML_ASSERT(cuda.context->transfer_backend_count == 0); +} + +// 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); + } +} + +// 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() { + 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); + // 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); + + 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); + 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 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); +} + +// 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); + 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 n_staged_buffers = cuda.context->buffers.size(); + GGML_ASSERT(cuda.context->transfer_backend_count == 1); + GGML_ASSERT(cuda.context->transfer_backend_inits == 1); + + 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()); + 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); +#else + GGML_ASSERT(setenv(name, value, 1) == 0); +#endif +} + +// 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 + _putenv_s(name, had_value ? value.c_str() : ""); +#else + if (had_value) { + setenv(name, value.c_str(), 1); + } else { + unsetenv(name); + } +#endif + } +}; + +static void test_transport_environment_is_fallback() { + 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 }; + + 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); + 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)); + + // 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() { + 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_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); + 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 test_backend_graph_optimize(ggml_backend_t, ggml_cgraph * graph, ggml_backend_graph_optimize_params * params) { GGML_ASSERT(graph->n_nodes == 3); params->add_alloc_dep(params->user_data, graph->nodes[0], graph->nodes[2]); @@ -1193,6 +2065,22 @@ int main() { run("test_resizable_buffers_owner_borrower_allocation_failure", test_resizable_buffers_owner_borrower_allocation_failure); run("test_resizable_buffers_owner_borrower_scheduler_failure", test_resizable_buffers_owner_borrower_scheduler_failure); run("test_resizable_buffers_owner_borrower_teardown_order", test_resizable_buffers_owner_borrower_teardown_order); + run("test_transport_prefix_and_configuration", test_transport_prefix_and_configuration); + run("test_transport_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_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); + 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); run("test_graph_optimize_alloc_dep", test_graph_optimize_alloc_dep); return 0; } diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 1fff21f701e..1b2d03c89d1 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -357,6 +357,10 @@ struct cmd_params { std::vector lazy_mode; std::vector main_gpu; 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; std::vector> tensor_split; @@ -402,6 +406,10 @@ static const cmd_params cmd_params_defaults = { /* lazy_mode */ { LLAMA_LAZY_MODE_AUTO }, /* main_gpu */ { 0 }, /* 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 */ { {} }, /* tensor_split */ { std::vector(llama_max_devices(), 0.0f) }, @@ -472,6 +480,10 @@ 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(" -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"); printf(" -lm, --load-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").c_str()); @@ -841,6 +853,52 @@ 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 == "-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 > LLAMA_KV_PIPELINE_DEPTH_MAX) { + invalid_param = true; + 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) { + invalid_param = true; + break; + } + auto p = parse_int_range(argv[i]); + for (int budget : p) { + if (budget < 0 || budget > LLAMA_KV_PIPELINE_BUDGET_MIB_MAX) { + invalid_param = true; + 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) { + 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 +1246,18 @@ 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.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; + } if (params.flash_attn.empty()) { params.flash_attn = cmd_params_defaults.flash_attn; } @@ -1251,6 +1321,10 @@ struct cmd_params_instance { llama_lazy_mode lazy_mode; int main_gpu; 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; std::vector tensor_split; @@ -1332,6 +1406,10 @@ 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.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; cparams.op_offload = !no_op_offload; @@ -1366,6 +1444,10 @@ 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 & 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) for (const auto & cm : params.cpu_mask) @@ -1396,6 +1478,10 @@ static std::vector get_cmd_params_instances(const cmd_param /* .lazy_mode = */ lzm, /* .main_gpu = */ mg, /* .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, /* .tensor_split = */ ts, @@ -1433,6 +1519,10 @@ static std::vector get_cmd_params_instances(const cmd_param /* .lazy_mode = */ lzm, /* .main_gpu = */ mg, /* .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, /* .tensor_split = */ ts, @@ -1470,6 +1560,10 @@ static std::vector get_cmd_params_instances(const cmd_param /* .lazy_mode = */ lzm, /* .main_gpu = */ mg, /* .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, /* .tensor_split = */ ts, @@ -1512,6 +1606,10 @@ struct test { llama_lazy_mode lazy_mode; int main_gpu; 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; std::vector tensor_split; @@ -1552,6 +1650,10 @@ 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; + 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; tensor_split = inst.tensor_split; @@ -1616,9 +1718,10 @@ 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", - "tensor_buft_overrides", "load_mode", "lazy_mode", - "embeddings", + "main_gpu", "no_kv_offload", "kv_cpu_pinned", "kv_pipeline_depth", "kv_pipeline_budget_mib", + "recurrent_state_offload", + "flash_attn", "devices", "tensor_split", + "tensor_buft_overrides", "load_mode", "lazy_mode", "embeddings", "no_op_offload", "no_host", "fit_target", "fit_min_ctx", "n_prompt", "n_gen", "n_depth", "test_time", "avg_ns", "stddev_ns", "avg_ts", "stddev_ts" @@ -1631,12 +1734,13 @@ 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 == "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; } - 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 +1812,10 @@ 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(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), tensor_split_str, @@ -1904,6 +2012,18 @@ struct markdown_printer : public printer { if (field == "test") { return 15; } + if (field == "kv_cpu_pinned") { + return 4; + } + if (field == "kv_pipeline_depth") { + return 4; + } + if (field == "kv_pipeline_budget_mib") { + return 5; + } + if (field == "recurrent_state_offload") { + return 3; + } if (field == "no_op_offload") { return 4; } @@ -1929,6 +2049,18 @@ struct markdown_printer : public printer { if (field == "n_threads") { return "threads"; } + if (field == "kv_cpu_pinned") { + return "kvcp"; + } + if (field == "kv_pipeline_depth") { + return "kvpd"; + } + if (field == "kv_pipeline_budget_mib") { + return "kvpb"; + } + if (field == "recurrent_state_offload") { + return "rso"; + } if (field == "no_kv_offload") { return "nkvo"; } @@ -2010,6 +2142,20 @@ 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.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"); + } if (params.no_kv_offload.size() > 1 || params.no_kv_offload != cmd_params_defaults.no_kv_offload) { fields.emplace_back("no_kv_offload"); }