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

common/arg.cpp

@@ -30,6 +30,7 @@
 #include <filesystem>
 #include <fstream>
 #include <list>
+#include <limits>
 #include <numeric>
 #include <regex>
 #include <set>
@@ -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<size_t>::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"},

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;

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

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

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

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)

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