sched: pipeline the delivery of a host-resident KV cache (R4)

AGENTS.md

@@ -18,10 +18,24 @@ small:
   `GGML_TYPE_Q2_0S` so it cannot collide with upstream's serialized Q2_0 weight
   format.
 - A profit-only adaptive draft-max controller for DFlash.
+- Pipelined delivery of a host-resident KV cache, selected with
+  `--kv-pipeline-depth` and bounded by `--kv-pipeline-budget`.
 - Reasoning-loop detection and the opted-in realtime
   `/v1/chat/completions/control` endpoint.
 - INI presets and KLD measurement support in `llama-perplexity`.
 
+`--kv-pipeline-depth N` issues the host-to-device delivery of a host-resident KV
+cache N splits ahead of the split that reads it, on a transfer stream of its own,
+so a decode token stops paying transfer and attention in series. It is on by
+default at `N = 1`, `0` restores the ordered path exactly, and it only engages
+where a host-resident cache produces the deliveries (`--no-kv-offload`, with or
+without `--kv-cpu-pinned`). Output is byte-identical either way. The staging it
+needs grows with the context, so `--kv-pipeline-budget` (default 128 MiB) caps it
+outright and the scheduler declines past the cap rather than spending the device
+memory a host-resident cache exists to save. Measured single-GPU only. See
+[`docs/kv-transport-pipelining.md`](docs/kv-transport-pipelining.md) for the
+design, the numbers, the limits, and the reproduction scripts.
+
 DFlash GGUFs must use upstream's `dflash` architecture, metadata, and tensor
 names.
 

common/arg.cpp

@@ -2600,6 +2600,34 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.kv_cpu_pinned = value;
         }
     ).set_env("LLAMA_ARG_KV_CPU_PINNED"));
+    add_opt(common_arg(
+        {"--kv-pipeline-depth"}, "N",
+        string_format("how many splits ahead the scheduler delivers a host-resident KV cache to the accelerator, so "
+                      "that the transfer runs while the previous split computes; 0 keeps the ordered path, where a "
+                      "decode token pays the transfer and the attention kernels in series. Only takes effect with a "
+                      "host-resident cache, e.g. --no-kv-offload or --kv-cpu-pinned, and costs (N + 2) * (largest "
+                      "staged split) of device memory (default: %d)", params.kv_pipeline_depth),
+        [](common_params & params, int value) {
+            if (value < 0 || value > 14) {
+                throw std::invalid_argument("--kv-pipeline-depth must be between 0 and 14");
+            }
+            params.kv_pipeline_depth = value;
+        }
+    ).set_env("LLAMA_ARG_KV_PIPELINE_DEPTH"));
+    add_opt(common_arg(
+        {"--kv-pipeline-budget"}, "N",
+        string_format("hard cap, in MiB, on the device memory that pipelined delivery of a host-resident KV cache "
+                      "may use. A staging slot holds one attention layer's K or V over the whole context, so the "
+                      "requirement grows with the context; past this cap the scheduler declines and keeps the "
+                      "ordered path, so a host-resident cache never quietly trades away the device memory it exists "
+                      "to save. 0 removes the cap (default: %d)", params.kv_pipeline_budget_mib),
+        [](common_params & params, int value) {
+            if (value < 0) {
+                throw std::invalid_argument("--kv-pipeline-budget must not be negative");
+            }
+            params.kv_pipeline_budget_mib = value;
+        }
+    ).set_env("LLAMA_ARG_KV_PIPELINE_BUDGET"));
     add_opt(common_arg(
         {"--recurrent-state-offload"},
         {"--no-recurrent-state-offload"},

common/common.cpp

@@ -1890,6 +1890,8 @@ struct llama_context_params common_context_params_to_llama(const common_params &
     cparams.cb_eval_user_data = params.cb_eval_user_data;
     cparams.offload_kqv       = !params.no_kv_offload;
     cparams.kv_cpu_pinned     = params.kv_cpu_pinned;
+    cparams.kv_pipeline_depth = params.kv_pipeline_depth < 0 ? 0 : (uint32_t) params.kv_pipeline_depth;
+    cparams.kv_pipeline_budget_mib = params.kv_pipeline_budget_mib < 0 ? 0 : (uint32_t) params.kv_pipeline_budget_mib;
     cparams.recurrent_state_offload = params.recurrent_state_offload;
     cparams.phase_aware_workspace = params.phase_aware_workspace;
     cparams.live_context_workspace = params.live_context_workspace;

common/common.h

@@ -642,6 +642,8 @@ struct common_params {
     bool recurrent_state_offload = false; // for hybrid models, keep recurrent state GPU-resident even with no_kv_offload
     bool phase_aware_workspace = false; // resize compute schedulers between prompt and generation phases
     bool live_context_workspace = false; // size supported attention workspaces from the padded live KV extent
+    int32_t kv_pipeline_depth = 1;  // splits of look-ahead for pipelined delivery of a host-resident KV cache (0 = off)
+    int32_t kv_pipeline_budget_mib = 128; // hard cap on the device memory that delivery may use (0 = uncapped)
     bool warmup            = true;  // warmup run
     bool check_tensors     = false; // validate tensor data
     bool no_op_offload     = false; // globally disable offload host tensor operations to device

docs/beellama-args.md

@@ -202,6 +202,8 @@ context-dependent H2D traffic; it is not equivalent to CPU attention.
 | `--recurrent-state-offload`, `--no-recurrent-state-offload` | `LLAMA_ARG_RECURRENT_STATE_OFFLOAD` | Disabled | For hybrid models under `--no-kv-offload`, keeps the fixed recurrent R/S state on the GPU while long attention KV remains in host memory. With ordinary KV offload enabled, recurrent state is already device-resident. |
 | `--kv-gpu-layers N` | `LLAMA_ARG_KV_GPU_LAYERS` | `0` | Under `--no-kv-offload`, retains the first `N` attention-KV layers on the device. This trades VRAM for lower repeated host-to-device traffic; a value at or above the attention-layer count matches full KV offload for those layers. |
 | `--spec-draft-kv-gpu-layers N`, `--kv-gpu-layers-draft N` | `LLAMA_ARG_SPEC_DRAFT_KV_GPU_LAYERS` | Inherit target | Overrides the target KV policy for a separate draft context and keeps the first `N` independently owned draft attention-KV layers device-resident. Shared KV layers follow their owner. The INI key is `spec-draft-kv-gpu-layers`. |
+| `--kv-pipeline-depth N` | `LLAMA_ARG_KV_PIPELINE_DEPTH` | `1` | How many splits ahead the scheduler delivers a host-resident KV cache to the accelerator, so that the transfer runs while the previous split computes instead of in series with it. `0` restores the ordered path. Only engages where a host-resident cache produces the deliveries, and costs `(N + 2) x (largest staged split)` of device memory. See [`kv-transport-pipelining.md`](kv-transport-pipelining.md). |
+| `--kv-pipeline-budget N` | `LLAMA_ARG_KV_PIPELINE_BUDGET` | `128` | Hard cap, in MiB, on the device memory pipelined delivery may use. A staging slot holds one attention layer's K or V over the whole context, so the requirement grows with the context; past this cap the scheduler declines and keeps the ordered path, at no steady-state cost, so a host-resident cache never quietly trades away the device memory it exists to save. `0` removes the cap. |
 | `--op-offload`, `--no-op-offload` | — | Enabled | Controls whether operations over host tensors may execute on the accelerator. The negated form selects the separately tested CPU execution route; it is not the serving-performance path measured below. |
 
 Target and draft attention-KV placement is independent of model-weight