diff --git a/arc-tools/arcspec_perseq_ladder.sh b/arc-tools/arcspec_perseq_ladder.sh new file mode 100755 index 000000000..062eb2217 --- /dev/null +++ b/arc-tools/arcspec_perseq_ladder.sh @@ -0,0 +1,544 @@ +#!/bin/bash +# ArcSpec — per-sequence KV advance: THROUGHPUT on the ladder B = 1, 8, 32, 128. +# +# ───────────────────────────────────────────────────────────────────────────── +# THE QUESTION +# +# #116 proved per-sequence advance *works* at B=8 (it returned 1 token with no +# `finish_reason` before; it returns the full 48 after). It did not prove the +# mechanism is worth anything. MTP was measured at 1.93 tok/step at one user +# collapsing to 1.06 at 128, and this stack targets exactly that collapse. +# +# ⚠️ REPORT BOTH NUMBERS OR THE RESULT IS UNREADABLE. tok/step rising while +# aggregate throughput falls has already happened on this chain (+9%/+13% +# tok/step, −5×/−20% aggregate). A tok/step number alone is not a throughput +# result, and an aggregate number alone cannot say what moved. So every cell +# reports: +# +# aggregate tok/s (wall clock, the number that matters) +# tok_per_step (per-sequence multiplier) +# tok_per_batch_step (tokens out per target forward, whole batch) +# batch_steps/s (the decomposition: aggregate ≈ tok_per_batch_step +# × batch_steps/s — when aggregate does not move, +# this says which factor ate it, in ms) +# +# ───────────────────────────────────────────────────────────────────────────── +# ⚠️ TWO PROMPT REGIMES, MEAN-MATCHED. Both are needed, for opposite reasons. +# +# spread 144 24 320 64 260 40 200 96 words, cycled across workers +# uniform 144 words for every worker (= the spread's mean, and its first +# element, so B=1 is byte-identical in both regimes) +# +# **spread** is the regime this stack exists for. Uniform prompts hide the +# failure mode: the `xs` window defect #116 fixed could not even be reached with +# equal-length prompts, because the retained window is sized for the greediest +# row and a short row only becomes shorter-than-the-window when arrivals differ. +# A benchmark that sends one prompt to every worker measures the case that +# already worked. +# +# **uniform** is the only regime where the batch is WHOLE, and that is why it is +# not optional. The scheduler A/B has now measured the bucketing law directly: +# +# running bucket size = B / (distinct cache lengths) +# +# holding at 8/8=1 and 32/8=4, with `1 running, 7 waiting` sustained. Under a +# spread of 8 distinct lengths, B=8 therefore runs ONE sequence at a time — that +# chain measured B=8 spread at 7.91 tok/s against B=1's 15.36, i.e. batching is +# *negative* on realistic traffic. So a flat aggregate in the spread arm is +# unreadable on its own: it could be this fix not paying, or it could be that +# the batch never existed. Uniform is where a whole batch actually forms, so it +# is the only place this fix's effect on aggregate is visible without the +# serialisation swamping it. +# +# Mean-matched on purpose: same total prompt tokens in both regimes, so the only +# thing that differs is the *spread*, and a difference is attributable to +# raggedness rather than to prompt length. +# +# Session 8 measured V4 serving dying outright at ~1,055-word prompts and this +# branch does NOT carry #97 (`pin V4 against length`), so 320 keeps a >3x margin +# to that cliff while still being a real spread. +# +# ───────────────────────────────────────────────────────────────────────────── +# ⚠️ THE SCHEDULER CONFOUND, MEASURED RATHER THAN ASSUMED. +# +# Both schedulers bucket the running set by EXACT cache length and run ONE +# bucket per step, preempting the rest (`default_scheduler.rs`, +# `paged_attention/scheduler.rs`). So "B=128" in this harness can be a 3-wide +# step in the engine, and an aggregate number measured inside that is a number +# about the scheduler, not about KV advance. +# +# `SCHED[agg]` (added with this harness) is emitted on the SAME log fence as +# `MTP[agg]`, so each cell reports, over exactly its own window: +# +# buckets_per_step 1.0 = the batch was never split +# running_bucket_size the width the engine ACTUALLY ran +# offered_per_step the width this harness thinks it asked for +# +# If `running_bucket_size` << `offered_per_step`, the throughput cell is a +# scheduler measurement and says so. +# +# ───────────────────────────────────────────────────────────────────────────── +# THE DESIGN +# +# Two server processes, one per arm; four concurrency cells inside each. +# OFF no flags (control) +# ON ARC_V4_XS_PER_SEQ=1 ARC_MTP_PER_SEQ_KV=1 (treatment) +# Same binary, same model, same prompts, same MTP depth, same --max-seqs. One +# server per arm rather than one per cell because the engine counters are +# cumulative and monotone: differencing them across a wall-clock fence gives a +# cell's own numbers without paying the ~2m30s V4 load eight times. +# +# ⚠️ ENGAGEMENT IS ASSERTED, NOT ASSUMED (D18). "per-sequence KV advance is ON" +# and "Ragged batch admission is ON" must be PRESENT in the ON log and ABSENT +# in OFF, "cannot honour it" must be absent from ON, and `per_seq_steps` must be +# non-zero — granted-but-inert has to show as a visible zero, not as silence. +# +# Environment failure exits 2. A gate that ran and failed exits 1. +set -u + +ARCENV="${ARCENV:-/root/arcenv.sh}" +WT="${WT:-/root/wt-ladder}" +BIN="${BIN:-$WT/target/release/mistralrs}" +MODEL="${MODEL:-/root/models/v4-src}" +UQFF="${UQFF:-/root/models/v4-uqff/qtip2b-0.uqff}" +ARCH="${ARCH:-deepseekv4}" +PREFLIGHT="${PREFLIGHT:-/root/arc-tools/gpu_box_preflight.sh}" +OUT="${OUT:-/root/logs/ladder}" +STATUS="${STATUS:-/root/logs/ladder.status}" +PORT="${PORT:-1243}" +MTP_DEPTH="${MTP_DEPTH:-3}" +MAX_SEQS="${MAX_SEQS:-256}" +SEED="${SEED:-1234}" +BATCHES="${BATCHES:-1 8 32 128}" +WARMUP="${WARMUP:-20}" +STEADY="${STEADY:-45}" +MAX_TOKENS="${MAX_TOKENS:-4096}" +# 144 first so B=1 is byte-identical between the two regimes; mean is 144. +WORD_SPREAD="${WORD_SPREAD:-144 24 320 64 260 40 200 96}" +UNIFORM_WORDS="${UNIFORM_WORDS:-144}" +REGIMES="${REGIMES:-uniform spread}" + +mkdir -p "$OUT" "$(dirname "$STATUS")" +: > "$STATUS" +say() { echo "[$(date -u +%H:%M:%SZ)] $*" | tee -a "$STATUS"; } +die() { say "ABORT: $*"; echo "ABORT: $*" >> "$OUT/summary.txt"; exit 2; } + +say "ArcSpec per-seq throughput ladder — start" +[ -f "$ARCENV" ] || die "$ARCENV not found" +# shellcheck disable=SC1090 +. "$ARCENV" +say "env: CUDA_HOME=${CUDA_HOME:-unset} LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-unset}" +[ -x "$PREFLIGHT" ] || die "$PREFLIGHT missing — refusing to measure" +"$PREFLIGHT" > "$OUT/preflight.log" 2>&1 || die "preflight refused the box" +say "preflight OK" +[ -x "$BIN" ] || die "$BIN not found" +SHA=$(git -C "$WT" rev-parse HEAD 2>/dev/null) || die "no git ref at $WT" +say "binary ref: $SHA" +nvidia-smi --query-gpu=name,memory.total,memory.used --format=csv > "$OUT/gpu.txt" 2>&1 + +# ── The load driver. Streaming, so a cell yields tokens even when no request +# finishes inside it (at B=128 V4 needs minutes per request). +cat > "$OUT/drive.py" <<'PYEOF' +import argparse, json, threading, time, urllib.request + +FILLER = ("alpha bravo charlie delta echo foxtrot golf hotel india juliet " + "kilo lima mike november oscar papa quebec romeo sierra tango").split() + +def prompt_of(words): + return ("Continue this list, one item per line. " + + " ".join(FILLER[i % len(FILLER)] for i in range(words))) + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--port", type=int, required=True) + ap.add_argument("--k", type=int, required=True) + ap.add_argument("--warmup", type=float, required=True) + ap.add_argument("--steady", type=float, required=True) + ap.add_argument("--spread", required=True, help="comma-separated word counts") + ap.add_argument("--max-tokens", type=int, default=4096) + ap.add_argument("--seed", type=int, default=1234) + a = ap.parse_args() + + # 🔑 RAGGED BY CONSTRUCTION: worker i takes spread[i % len(spread)], so the + # arrivals differ in length the way real ones do. Deterministic (no RNG), so + # the two arms see byte-identical prompt sets. + spread = [int(x) for x in a.spread.split(",") if x] + prompts = [prompt_of(spread[i % len(spread)]) for i in range(a.k)] + url = f"http://127.0.0.1:{a.port}/v1/completions" + stop = threading.Event() + lock = threading.Lock() + n = {"chunks": 0, "errors": 0, "started": 0, "finished": 0} + counting = [False] + + def worker(wid): + body_prompt = prompts[wid] + while not stop.is_set(): + body = json.dumps({ + "model": "default", "prompt": body_prompt, + "max_tokens": a.max_tokens, "temperature": 0, + "seed": a.seed, "stream": True, + }).encode() + req = urllib.request.Request( + url, data=body, headers={"Content-Type": "application/json"}) + try: + with lock: + n["started"] += 1 + with urllib.request.urlopen(req, timeout=1800) as r: + for raw in r: + if stop.is_set(): + break + line = raw.decode("utf-8", "ignore").strip() + if not line.startswith("data:"): + continue + payload = line[5:].strip() + if payload == "[DONE]": + with lock: + n["finished"] += 1 + break + try: + ch = json.loads(payload) + except Exception: + continue + txt = (ch.get("choices") or [{}])[0].get("text", "") + if txt: + with lock: + if counting[0]: + n["chunks"] += 1 + except Exception: + with lock: + n["errors"] += 1 + time.sleep(0.5) + + threads = [threading.Thread(target=worker, args=(i,), daemon=True) + for i in range(a.k)] + for t in threads: + t.start() + time.sleep(a.warmup) + with lock: + n["chunks"] = 0 + counting[0] = True + t0 = time.time() + fence = {"t_start": time.strftime("%Y-%m-%dT%H:%M:%S", time.gmtime())} + time.sleep(a.steady) + with lock: + counting[0] = False + snap = dict(n) + fence["t_end"] = time.strftime("%Y-%m-%dT%H:%M:%S", time.gmtime()) + stop.set() + elapsed = time.time() - t0 + snap.update(fence) + snap["k"] = a.k + snap["prompt_words"] = [spread[i % len(spread)] for i in range(a.k)] + snap["elapsed_s"] = round(elapsed, 3) + snap["driver_tok_s"] = round(snap["chunks"] / elapsed, 3) if elapsed else 0.0 + print(json.dumps(snap)) + for t in threads: + t.join(timeout=3) + +main() +PYEOF + +# ── The extractor. Differences cumulative MTP[agg] AND SCHED[agg] across the +# cell's wall-clock fence. Everything reported is a delta of raw counts. +cat > "$OUT/fence.py" <<'PYEOF' +"""Delta of the engine's cumulative counters across one cell's fence. + +`MTP[agg] ... steps=S committed=C batch_steps=BS` and `SCHED[agg] ... calls=N +buckets=B offered=O chosen=C` are cumulative for the process and emitted on the +same log fence, so a cell's own numbers are the difference between the last +line at or before its end and the last line at or before its start. + +Reports nothing it did not measure: a cell with <2 usable lines prints NA with +a reason rather than 0.0, because a zero would read as "the engine produced +nothing" when it means "the reporter never fired twice" (D18). +""" +import json, re, sys + +TS = re.compile(r"(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2})") +FIELDS = { + "MTP[agg]": ("accepted", "proposed", "steps", "drafted_steps", "committed", + "batch_steps"), + "SCHED[agg]": ("calls", "buckets", "offered", "chosen", "shattered"), +} + +def parse(path, marker): + keys = FIELDS[marker] + out = [] + for line in open(path, errors="ignore"): + if marker not in line: + continue + mt = TS.search(line) + if not mt: + continue + rec, ok = {"ts": mt.group(1)}, True + for k in keys: + m = re.search(rf"\b{k}=(\d+)\b", line) + if not m: + ok = False + break + rec[k] = int(m.group(1)) + if ok: + out.append(rec) + return out + +def bracket(recs, t0, t1): + """(last record at or before t0, last record at or before t1).""" + before = [r for r in recs if r["ts"] <= t0] + within = [r for r in recs if r["ts"] <= t1] + if not before or not within: + return None, None + return before[-1], within[-1] + +def delta(path, marker, t0, t1): + recs = parse(path, marker) + if len(recs) < 2: + return {"NA": f"{marker}: only {len(recs)} line(s) in the whole log"} + a, b = bracket(recs, t0, t1) + if a is None or b is None: + return {"NA": f"{marker}: no line brackets [{t0}, {t1}]"} + d = {k: b[k] - a[k] for k in FIELDS[marker]} + if all(v == 0 for v in d.values()): + return {"NA": f"{marker}: counters did not advance inside the fence"} + return d + +def main(): + log, t0, t1, elapsed, driver_tok_s = ( + sys.argv[1], sys.argv[2], sys.argv[3], float(sys.argv[4]), float(sys.argv[5])) + out = {"t_start": t0, "t_end": t1, "elapsed_s": elapsed, + "aggregate_tok_s": driver_tok_s} + m = delta(log, "MTP[agg]", t0, t1) + s = delta(log, "SCHED[agg]", t0, t1) + out["mtp_raw"], out["sched_raw"] = m, s + r = lambda x: round(x, 4) + if "NA" not in m: + out["tok_per_step"] = r(m["committed"] / m["steps"]) if m["steps"] else None + out["tok_per_batch_step"] = ( + r(m["committed"] / m["batch_steps"]) if m["batch_steps"] else None) + out["mean_batch"] = r(m["steps"] / m["batch_steps"]) if m["batch_steps"] else None + out["batch_steps_per_s"] = r(m["batch_steps"] / elapsed) if elapsed else None + out["engine_tok_s"] = r(m["committed"] / elapsed) if elapsed else None + if m["batch_steps"] and elapsed: + out["ms_per_batch_step"] = r(1000.0 * elapsed / m["batch_steps"]) + if "NA" not in s: + out["buckets_per_step"] = r(s["buckets"] / s["calls"]) if s["calls"] else None + out["running_bucket_size"] = r(s["chosen"] / s["calls"]) if s["calls"] else None + out["offered_per_step"] = r(s["offered"] / s["calls"]) if s["calls"] else None + out["shattered_frac"] = r(s["shattered"] / s["calls"]) if s["calls"] else None + print(json.dumps(out)) + +main() +PYEOF + +run_arm() { + local arm="$1"; shift + local slog="$OUT/server.$arm.log" + say "arm $arm — starting server (${*:-no flags})" + # shellcheck disable=SC2086 + env RUST_LOG=info ARC_MTP_LOG_ACCEPTANCE=1 "$@" \ + "$BIN" serve -p "$PORT" -m "$MODEL" -a "$ARCH" --from-uqff "$UQFF" \ + --max-seqs "$MAX_SEQS" --mtp-depth "$MTP_DEPTH" --prefix-cache-n 0 \ + --paged-attn off --seed "$SEED" \ + < /dev/null > "$slog" 2>&1 & + local spid=$! + local up=0 + for _ in $(seq 1 450); do + curl -s -m 5 "http://127.0.0.1:$PORT/v1/models" >/dev/null 2>&1 && { up=1; break; } + kill -0 "$spid" 2>/dev/null || break + sleep 2 + done + [ "$up" = "1" ] || { say "arm $arm — server never came up"; tail -25 "$slog" | tee -a "$STATUS"; return 1; } + + local logged + logged=$(grep -m1 -o 'git revision: [0-9a-f]*' "$slog" | awk '{print $3}') + if [ -n "$SHA" ] && [ "$logged" != "$SHA" ]; then + say "arm $arm — PROVENANCE FAILED: server $logged != worktree $SHA" + kill -TERM "$spid" 2>/dev/null || true + return 1 + fi + say "arm $arm — server up, revision $logged" + + local regime k spread_csv tag + for regime in $REGIMES; do + if [ "$regime" = "uniform" ]; then + spread_csv="$UNIFORM_WORDS" + else + spread_csv=$(echo "$WORD_SPREAD" | tr ' ' ',') + fi + for k in $BATCHES; do + tag="$arm.$regime.k$k" + say "arm $arm — cell $regime B=$k (warmup ${WARMUP}s, steady ${STEADY}s, words $spread_csv)" + python3 "$OUT/drive.py" --port "$PORT" --k "$k" --warmup "$WARMUP" \ + --steady "$STEADY" --spread "$spread_csv" --max-tokens "$MAX_TOKENS" \ + --seed "$SEED" > "$OUT/$tag.drive.json" 2>"$OUT/$tag.drive.err" \ + || { say "cell $tag — driver failed"; return 1; } + local t0 t1 el tok + t0=$(python3 -c "import json,sys;print(json.load(open(sys.argv[1]))['t_start'])" "$OUT/$tag.drive.json") + t1=$(python3 -c "import json,sys;print(json.load(open(sys.argv[1]))['t_end'])" "$OUT/$tag.drive.json") + el=$(python3 -c "import json,sys;print(json.load(open(sys.argv[1]))['elapsed_s'])" "$OUT/$tag.drive.json") + tok=$(python3 -c "import json,sys;print(json.load(open(sys.argv[1]))['driver_tok_s'])" "$OUT/$tag.drive.json") + python3 "$OUT/fence.py" "$slog" "$t0" "$t1" "$el" "$tok" > "$OUT/$tag.cell.json" \ + || { say "cell $tag — fence extraction failed"; return 1; } + say "$tag -> $(cat "$OUT/$tag.cell.json")" + sleep 5 + done + done + + kill -TERM "$spid" 2>/dev/null || true + for _ in $(seq 1 30); do kill -0 "$spid" 2>/dev/null || break; sleep 2; done + kill -KILL "$spid" 2>/dev/null || true + sleep 5 + say "arm $arm — done" + return 0 +} + +run_arm OFF || die "OFF arm failed" +run_arm ON ARC_V4_XS_PER_SEQ=1 ARC_MTP_PER_SEQ_KV=1 || die "ON arm failed" +# The third arm prices the window pin, which is OPT-IN (default OFF) — so every +# arm above is unpinned and this one is the treatment. Same flags as ON, pin +# enabled, so ON vs ON_PINNED isolates one change: whether the compressor's raw +# window is reallocated every decode step on all 41 compressed layers, or held +# at its documented bound. +# +# This arm IS the experiment named in `xs_pin_window_enabled_from`'s FLIP +# CONDITION. If ON_PINNED wins at uniform B=32 with identical generated tokens, +# `ARC_V4_XS_PIN_WINDOW` defaults to ON and that doc comment is rewritten in the +# same change that records the number. +run_arm ON_PINNED ARC_V4_XS_PER_SEQ=1 ARC_MTP_PER_SEQ_KV=1 ARC_V4_XS_PIN_WINDOW=1 \ + || die "ON_PINNED arm failed" + +cat > "$OUT/report.py" <<'PYEOF' +import json, os, sys +out, batches, regimes = sys.argv[1], sys.argv[2].split(), sys.argv[3].split() +ARMS = ("OFF", "ON", "ON_PINNED") +def cell(arm, regime, k): + p = os.path.join(out, f"{arm}.{regime}.k{k}.cell.json") + try: + return json.load(open(p)) + except Exception as e: + return {"NA": str(e)} +def g(c, k): + v = c.get(k) + return "NA" if v is None else v +print("\n=== LADDER: aggregate throughput and its decomposition ===") +print("aggregate_tok/s is wall-clock and is THE number. tok/batch_step x") +print("batch_steps/s reconstructs it: when aggregate does not move, those two") +print("say which factor ate it. run_bucket is the width the engine ACTUALLY ran") +print("— compare it to offered before reading any cell as a batch result.") +hdr = (f"{'regime':<8} {'B':>4} {'arm':<12} {'aggregate_tok/s':>15} {'tok/step':>9} " + f"{'tok/bstep':>10} {'bsteps/s':>9} {'ms/bstep':>9} " + f"{'buckets/step':>13} {'run_bucket':>11} {'offered':>8}") +print(hdr); print("-" * len(hdr)) +for regime in regimes: + for k in batches: + for arm in ARMS: + c = cell(arm, regime, k) + print(f"{regime:<8} {k:>4} {arm:<12} {str(g(c,'aggregate_tok_s')):>15} " + f"{str(g(c,'tok_per_step')):>9} {str(g(c,'tok_per_batch_step')):>10} " + f"{str(g(c,'batch_steps_per_s')):>9} {str(g(c,'ms_per_batch_step')):>9} " + f"{str(g(c,'buckets_per_step')):>13} {str(g(c,'running_bucket_size')):>11} " + f"{str(g(c,'offered_per_step')):>8}") +print("\n=== ON vs OFF, per regime ===") +for regime in regimes: + print(f"\n-- {regime} --") + for k in batches: + o, n = cell("OFF", regime, k), cell("ON", regime, k) + def ratio(key): + a, b = o.get(key), n.get(key) + if not a or not b: + return "NA" + return f"{b/a:.3f}x ({a} -> {b})" + print(f" B={k:<4} aggregate {ratio('aggregate_tok_s')}") + print(f" tok/step {ratio('tok_per_step')}") + print(f" tok/bstep {ratio('tok_per_batch_step')}") + print(f" bsteps/s {ratio('batch_steps_per_s')}") + print(f" run_bucket {ratio('running_bucket_size')}") +print("\n=== THE WINDOW PIN, isolated (ON = control, ON_PINNED = treatment) ===") +print("Same per-seq flags in both; the only difference is whether the") +print("compressor's raw window is reallocated every decode step on 41 layers.") +print("Ratios are treatment/control, so >1.000x means the pin helped.") +print("This is the number in ARC_V4_XS_PIN_WINDOW's FLIP CONDITION: if the") +print("uniform B=32 row is >1.000x with identical output, the default flips ON.") +for regime in regimes: + for k in batches: + # `u` is the unpinned control, `p_` the pinned treatment; the pin is + # opt-in, so the DEFAULT arm (ON) is the control. This pairing was the + # other way round while the pin defaulted on — the ratio is + # treatment/control either way, and swapping the arms without swapping + # this would have silently inverted every number below. + u, p_ = cell("ON", regime, k), cell("ON_PINNED", regime, k) + def r2(key): + a, b = u.get(key), p_.get(key) + if not a or not b: + return "NA" + return f"{b/a:.3f}x ({a} -> {b})" + print(f" {regime:<8} B={k:<4} aggregate {r2('aggregate_tok_s')} " + f"ms/bstep {r2('ms_per_batch_step')}") + +print("\n=== THE BUCKETING LAW, checked against this run ===") +print("The scheduler A/B measured `running bucket = B / distinct lengths`.") +print("uniform has 1 distinct length, spread has 8, so the prediction is:") +for regime in regimes: + dist = 1 if regime == "uniform" else 8 + for k in batches: + pred = max(1.0, int(k) / dist) + for arm in ARMS: + c = cell(arm, regime, k) + got = c.get("running_bucket_size") + if got is None: + continue + print(f" {regime:<8} B={k:<4} {arm:<12} predicted {pred:>6.2f} measured {got:>7.3f}") +print("\n⚠️ Read every spread cell against its run_bucket. If run_bucket is ~1,") +print("that cell is a measurement of the scheduler serialising the batch, not") +print("of per-sequence KV advance, and its aggregate cannot be attributed here.") +PYEOF + +{ + echo "=== ref ==="; echo "$SHA" + echo "=== prompts ===" + echo " spread : $WORD_SPREAD (8 distinct lengths)" + echo " uniform: $UNIFORM_WORDS (the spread's mean AND its first element, so B=1 matches)" + echo + echo "=== ENGAGEMENT ===" + for m in "per-sequence KV advance is ON" "Ragged batch admission is ON"; do + on=$(grep -c "$m" "$OUT/server.ON.log" 2>/dev/null || echo 0) + off=$(grep -c "$m" "$OUT/server.OFF.log" 2>/dev/null || echo 0) + echo " '$m': ON=$on (want >=1) OFF=$off (want 0)" + done + echo " 'cannot honour it' in ON: $(grep -c 'cannot honour it' "$OUT/server.ON.log" 2>/dev/null || echo 0) (want 0)" + echo " per_seq_steps final, ON: $(grep -o 'per_seq_steps=[0-9]*' "$OUT/server.ON.log" 2>/dev/null | tail -1) (want >0)" + echo " per_seq_steps final, OFF: $(grep -o 'per_seq_steps=[0-9]*' "$OUT/server.OFF.log" 2>/dev/null | tail -1) (want =0)" + echo " 'xs rolling cache' errors ON=$(grep -c 'xs rolling cache' "$OUT/server.ON.log" 2>/dev/null || echo 0) OFF=$(grep -c 'xs rolling cache' "$OUT/server.OFF.log" 2>/dev/null || echo 0) (want 0/0)" + echo + echo " -- window pin engagement (a wrong flag name would make ON and ON_PINNED" + echo " the same build and report a free pin from two identical arms) --" + for a in OFF ON ON_PINNED; do + echo " $a: $(grep -m1 -o 'xs rolling window is [A-Z]*' "$OUT/server.$a.log" 2>/dev/null || echo 'NOT LOGGED')" + done + # ⚠️ Read from the LOG, never from the binary. BOTH mode strings are compiled + # in unconditionally, so `strings | grep RESIZING` would succeed in + # every arm and prove nothing. Only the runtime line says which branch a + # process actually took. `RUST_LOG=info` is set once in `run_arm`, shared by + # every arm, so no arm can have this line filtered out while another keeps it. + # + # ⚠️ The default direction matters here, and it INVERTED when the pin became + # opt-in: unset now means RESIZING (`xs_rolling.rs`, + # `xs_pin_window_enabled_from` is `== Some("1")`). So the CONTROL is the + # default and the TREATMENT is the one carrying the flag. A typo in the flag + # name now yields a treatment that silently ran the control — still a + # perfectly clean-looking 1.000x, which is why this check tests BOTH + # directions rather than inferring one from the other. + pin_off=$(grep -c 'xs rolling window is RESIZING' "$OUT/server.ON.log" 2>/dev/null || echo 0) + pin_on=$(grep -c 'xs rolling window is PINNED' "$OUT/server.ON_PINNED.log" 2>/dev/null || echo 0) + treat_leaked=$(grep -c 'xs rolling window is RESIZING' "$OUT/server.ON_PINNED.log" 2>/dev/null || echo 0) + if [ "$pin_off" -ge 1 ] && [ "$pin_on" -ge 1 ] && [ "$treat_leaked" -eq 0 ]; then + echo " => VALID: ON took the resizing path, ON_PINNED took the pinned path." + echo " The treatment reached the pinned branch; it did not silently run the control." + else + echo " => ⚠️ PIN A/B IS VOID: ON resizing=$pin_off, ON_PINNED pinned=$pin_on," + echo " treatment-leaked-to-resizing=$treat_leaked." + echo " Any ON-vs-ON_PINNED ratio below compares one thing with itself. Do not read it." + fi + python3 "$OUT/report.py" "$OUT" "$BATCHES" "$REGIMES" +} | tee "$OUT/summary.txt" + +say "DONE — $OUT/summary.txt" diff --git a/mistralrs-core/src/kv_cache/mod.rs b/mistralrs-core/src/kv_cache/mod.rs index b3e21b693..f21b01612 100644 --- a/mistralrs-core/src/kv_cache/mod.rs +++ b/mistralrs-core/src/kv_cache/mod.rs @@ -13,7 +13,7 @@ mod hybrid_cache; mod rotating_cache; mod single_cache; pub mod turboquant_cache; -mod xs_rolling; +pub(crate) mod xs_rolling; /// Thread-local `ARC_V4_XS_PER_SEQ` override, so a test outside this module can /// exercise both sides of the flag (the production read is a `OnceLock`). #[cfg(test)] @@ -2874,6 +2874,35 @@ mod clone_in_cache_invariant_tests { use crate::sequence::{SeqStepType, SequenceGroup, SequenceRecognizer}; use candle_core::Device; + /// Run `f` with the xs window pin OFF. + /// + /// Three tests below assert the geometry of the *unpinned* retained window + /// — exact tail widths, and that column 0 holds token `base`. Both are + /// properties of the resizing policy specifically: #121's pin holds the + /// buffer at `span_groups * ratio + margin` and leaves real tokens sitting + /// ahead of `base`, so column 0 is no longer `base` and the widths are the + /// pinned constants instead. + /// + /// The pin is opt-in (`ARC_V4_XS_PIN_WINDOW=1`), so as things stand these + /// three would pass without asking. They ask anyway, deliberately: the + /// default is expected to flip to ON once the pin is measured on this tree + /// (see `xs_rolling::xs_pin_window_enabled_from`'s FLIP CONDITION), and a + /// test that depends on a default it does not name is a test that breaks on + /// the day someone flips it. Naming the mode makes these three immune to + /// that, in either direction. + /// + /// They are not testing a dead path either way: the resizing window is a + /// supported serving mode and is one arm of the pin A/B. The pinned + /// geometry is covered separately by `xs_rolling`'s own tests and by + /// `deepseek4`'s pin A/B. + /// + /// Uses #121's `pin_test_override` — thread-local, not an env mutation, + /// because `cargo test` is multi-threaded and `xs_pin_window_enabled` + /// latches its env answer in a `OnceLock`. + fn unpinned(f: impl FnOnce() -> T) -> T { + crate::kv_cache::xs_rolling::pin_test_override::with(false, f) + } + /// A cache slot whose only interesting property is its `current_seq_len`. /// `all_data` stays `None` — `first_mismatched_cache_len` reads lengths, /// never tensors, which is exactly the point: `CACHE_GROW_SIZE = 512` means @@ -3958,8 +3987,23 @@ mod clone_in_cache_invariant_tests { /// compressor history and cannot share one dense buffer. It must be /// refused, by name — never papered over, and never a panic on the engine /// task. + /// + /// ⚠️ Run with the window pin OFF, deliberately. The widths 18 and 22 are + /// what the *resizing* buffer produced, and they are the whole + /// discriminator: pinning the window makes both 144 + /// (`span_groups * ratio + margin`), so the shape mismatch at :499 cannot + /// arise and this fixture would silently stop testing anything. The + /// historical defect is still worth pinning — the pin is a flag, and the + /// refusal is what has to hold when it is off. + /// `the_pin_removes_this_fixtures_discriminator` records the other side. #[test] fn ragged_xs_tail_is_refused_by_name_not_panicked() { + crate::kv_cache::xs_rolling::pin_test_override::with(false, || { + ragged_xs_tail_is_refused_by_name_not_panicked_inner(); + }); + } + + fn ragged_xs_tail_is_refused_by_name_not_panicked_inner() { let mut short = xs_state(128, 1); feed_xs(&mut short, 274); let mut long = xs_state(128, 1); @@ -4015,6 +4059,38 @@ mod clone_in_cache_invariant_tests { ); } + /// The other side of the fixture above: with the window pinned, the two + /// sequences' tails are the SAME width, so wave51-CB's shape mismatch is + /// not merely refused — it cannot be constructed. + /// + /// That is a consequence of the pin worth recording, not a reason to stop + /// refusing: the two caches still hold different history, and the batch is + /// still refused, just one layer earlier and by the token-count invariant + /// (`ensure_uniform_batch_cache_lens`) rather than by a tensor shape. A + /// refusal that depends on two buffers happening to differ in size is a + /// weaker guarantee than one that reads the lengths, and this shows the + /// weaker one is not what is holding. + #[test] + fn the_pin_removes_this_fixtures_discriminator() { + crate::kv_cache::xs_rolling::pin_test_override::with(true, || { + let mut short = xs_state(128, 1); + feed_xs(&mut short, 274); + let mut long = xs_state(128, 1); + feed_xs(&mut long, 278); + let (sw, lw) = ( + short.tail.as_ref().unwrap().dims()[1], + long.tail.as_ref().unwrap().dims()[1], + ); + assert_eq!( + (sw, lw), + (144, 144), + "pinned HCA windows must both be span_groups*ratio+margin = 144" + ); + // …and the lengths still disagree, which is what the refusal reads. + assert_ne!(short.current_seq_len(), long.current_seq_len()); + }); + } + /// The ragged-tail refusal must not depend on the K/V slots noticing /// first: on V4 the K/V halves are `[1, 1, capacity, head_dim]`, whose /// dim-1 is the head count, so they batch happily at *any* pair of lengths. @@ -4073,6 +4149,10 @@ mod clone_in_cache_invariant_tests { /// `shape mismatch on dim 1, 4 <> 132`. #[test] fn xs_base_divergence_at_equal_lengths_is_reconciled_not_refused() { + unpinned(xs_base_divergence_at_equal_lengths_is_reconciled_not_refused_inner); + } + + fn xs_base_divergence_at_equal_lengths_is_reconciled_not_refused_inner() { // Restored from a prefix-cache entry stored at 300 tokens, truncated to // 260 — `base` stays at canonical(300), which is past canonical(260). let mut restored = xs_state(128, 1); @@ -4189,6 +4269,10 @@ mod clone_in_cache_invariant_tests { /// looks at shapes. #[test] fn per_row_xs_bases_survive_a_batch_round_trip_and_are_not_flattened() { + unpinned(per_row_xs_bases_survive_a_batch_round_trip_and_are_not_flattened_inner); + } + + fn per_row_xs_bases_survive_a_batch_round_trip_and_are_not_flattened_inner() { xs_rolling::test_override::with(true, || { let mut restored = xs_state(128, 1); feed_xs(&mut restored, 300); @@ -4356,6 +4440,10 @@ mod clone_in_cache_invariant_tests { /// (token `t` carries the value `t`) and reads the actual numbers back. #[test] fn trimming_the_retained_window_drops_the_oldest_rows_not_the_newest() { + unpinned(trimming_the_retained_window_drops_the_oldest_rows_not_the_newest_inner); + } + + fn trimming_the_retained_window_drops_the_oldest_rows_not_the_newest_inner() { use candle_core::IndexOp; let dev = candle_core::Device::Cpu; let mut state = xs_state(128, 1); diff --git a/mistralrs-core/src/kv_cache/xs_rolling.rs b/mistralrs-core/src/kv_cache/xs_rolling.rs index fa30bed65..5e580547e 100644 --- a/mistralrs-core/src/kv_cache/xs_rolling.rs +++ b/mistralrs-core/src/kv_cache/xs_rolling.rs @@ -100,34 +100,42 @@ //! [`XsRollingCache::advance`] keeps a verbatim uniform fast path and every //! B=1 and every uniform batch runs byte-identical code. //! -//! # What `W` is now allowed to be +//! # The window is a fixed-size buffer, and `advance` was the one thing +//! # breaking that //! -//! Nothing below still requires `W == max_i (tokens_i - base_i)`. -//! [`plan_xs_advance`] and [`XsRollingCache::set_row_lens`] take `W` as *given* -//! and ask only that every row fit inside it; every read offset is -//! `need_start + W - tokens_i`, which tracks `W` exactly. So `W` may be any -//! width at or above the widest row's run — including a **constant** one. -//! `an_oversized_window_reads_exactly_the_same_absolute_tokens` pins that. +//! Read the description of `tail` at the top of this file again: *"the raw rows +//! behind `tokens`, **bounded by `span_groups * ratio + margin` and independent +//! of context length**"*. That is an invariant this type has always claimed — +//! a buffer whose size is a function of the layer's geometry and nothing else. //! -//! It is worth writing down because this buffer does **not** have a constant -//! size today. `tail` is rebuilt by `Tensor::cat` + `narrow` every step at -//! width `tokens - base`, and that width is not stable: `tokens` climbs by one -//! per token while `base` jumps a whole `ratio` at a group boundary, so the -//! allocation cycles through `ratio` consecutive sizes. The ArcGraph chain -//! measured exactly that from the outside — `4096 × {18, 19, 20, 21}` at -//! `hidden = 4096`, `ratio = 4`. (Their separate hypothesis, that this is what -//! stops a CUDA graph replaying, they have since retracted as unproven; the -//! size measurement is a fact about this buffer either way.) +//! It was not true of the allocation. `advance` rebuilt `tail` with +//! `Tensor::cat` + `narrow` on every step at width exactly `tokens - base`, and +//! that width is not stable: `tokens` climbs by one per token while `base` jumps +//! a whole `ratio` at a group boundary, so the buffer cycled through `ratio` +//! consecutive sizes and was reallocated every single decode step, on every one +//! of the 41 compressed layers. The ArcGraph chain measured exactly that shape +//! from outside the engine — `4096 × {18, 19, 20, 21}` at `hidden = 4096`, +//! `ratio = 4`. (Their further hypothesis, that this is what stops a CUDA graph +//! replaying, they have retracted as unproven. The reallocation is a per-token +//! cost on the hot path whether or not anything is capturing, which is the +//! reason it is fixed here.) //! -//! Pinning `W` to the bound the retention rule already guarantees it never -//! exceeds — `span_groups * ratio + margin` — would make the allocation -//! constant per step **without touching any of the semantics above**: a row -//! would simply carry more dead lead columns, exactly as a short row in a -//! ragged batch already does. The places that decide `W` are the retention -//! narrow at the end of [`XsRollingCache::advance_ragged`] and its counterpart -//! in [`XsRollingCache::advance_uniform`], plus the fast-path predicate -//! [`XsRollingCache::can_advance_uniform`], which reads the exact equality as -//! its dispatch condition. Nothing here depends on that change being made. +//! [`xs_pin_window_enabled`] holds the buffer at the documented bound instead. +//! It required no new semantics: [`plan_xs_advance`] and +//! [`XsRollingCache::set_row_lens`] already take `W` as *given* and ask only +//! that every row fit inside it, and every read offset is +//! `need_start + W - tokens_i`, which tracks `W` exactly. A pinned window is +//! therefore just the widest case of the ragged one that already had to work — +//! a row carrying dead lead columns it never promised. +//! `an_oversized_window_reads_exactly_the_same_absolute_tokens` pins the +//! arithmetic and `pinning_the_window_is_numerically_inert` pins the outcome. +//! +//! The two widths, reconciled: the steady-state band is +//! `[capacity - ratio, capacity - 1]`, i.e. `{20..23}` at CSA geometry, and the +//! `{18..21}` measured outside was the pre-saturation **ramp** — `base` is still +//! at 0 while `tokens` climbs, so `W` has not finished growing. +//! `the_window_ramps_then_settles_to_ratio_consecutive_sizes` checks both +//! halves of that, so the bound is right for a reason rather than by luck. use std::sync::OnceLock; @@ -236,6 +244,161 @@ pub fn xs_per_sequence_enabled() -> bool { }) } +/// `ARC_V4_XS_PIN_WINDOW=1` — hold the retained raw window at a **constant** +/// width instead of re-narrowing it to `tokens - base` every step. +/// **Default OFF.** See [`xs_pin_window_enabled_from`] for the polarity, the +/// reason, and the named measurement that flips this default back on. +/// +/// # The argument that once made this default-ON, kept because it is still true +/// +/// This shipped default-ON on #121's branch, on the reasoning below. The +/// reasoning is sound as far as it goes — it is why the pin is *safe* — but it +/// is an argument about correctness, and the default was also a claim about +/// throughput, which is what was never measured on this tree. Correctness +/// argument retained; default retracted. +/// +/// The usual rule on this project is measure-then-default-on, and it exists +/// because FP8 KV shipped defaulted-on and broke every forward for a day. The +/// difference is what the two changes touch. FP8 KV changed the *values*. This +/// changes only the *size of an allocation*: the compressor is handed a slice +/// covering the same absolute tokens either way, because every offset is +/// derived from the row's own token count rather than from the buffer's width +/// (see [`plan_xs_advance`]). `pinning_the_window_is_numerically_inert` pins +/// that as an exact-equality test rather than as an argument. +/// +/// # The bound, which is the reviewable part +/// +/// The retained width is `W = tokens - base`, and `base` is set from the +/// retention rule at `tokens - margin`, then only ever raised: +/// +/// ```text +/// keep_from = ((tokens - margin) / ratio + 1 - span_groups) * ratio +/// base = max(keep_from, previous base), capped at tokens +/// ``` +/// +/// Write `m = tokens - margin` and `q = floor(m / ratio)`. Then +/// `q * ratio > m - ratio`, so +/// +/// ```text +/// keep_from = (q + 1 - span_groups) * ratio > m - span_groups * ratio +/// W = tokens - base <= tokens - keep_from < margin + span_groups * ratio +/// ``` +/// +/// so **`W <= span_groups * ratio + margin - 1`, always**. Pinning to +/// `span_groups * ratio + margin` is therefore provably sufficient and provably +/// never truncates a row's promised history. At V4's CSA geometry +/// (`ratio 4, span_groups 2, margin 16`) that is 24, against a measured width +/// cycling through `{20, 21, 22, 23}` — `ratio` consecutive sizes, which is +/// exactly the `4096 x {18..21}`-shaped reallocation the ArcGraph chain +/// observed from outside the engine. +/// +/// This is a bound with a proof, not a chosen constant. The capacity is a +/// function of the layer's own `ratio`/`span_groups`/`margin`, so it does not +/// grow with context length. +pub fn xs_pin_window_enabled() -> bool { + #[cfg(test)] + match pin_test_override::current() { + Some(on) => return on, + None => {} + } + static ON: OnceLock = OnceLock::new(); + *ON.get_or_init(|| { + let on = xs_pin_window_enabled_from(std::env::var("ARC_V4_XS_PIN_WINDOW").ok().as_deref()); + // Named once per process, so an A/B that means to compare pinned + // against unpinned can ASSERT it got two different builds of behaviour + // rather than two identical arms reporting a 1.000x ratio. A flag whose + // name is wrong is indistinguishable from a flag that costs nothing. + tracing::info!( + target: "xs_rolling", + "xs rolling window is {} (ARC_V4_XS_PIN_WINDOW)", + if on { "PINNED" } else { "RESIZING" } + ); + on + }) +} + +/// Pure half of [`xs_pin_window_enabled`], so the polarity is testable without +/// mutating the process environment. +/// +/// # Why the pin is OPT-IN +/// +/// This shipped default-ON, with `=0` as the control arm, because that made the +/// A/B guard meaningful. It is now default-OFF, `== Some("1")`, and the A/B +/// carries the flag on its TREATMENT arm instead. +/// +/// The width bound above is proved, and the pin is very probably a win. But +/// "probably a win" is not the bar. It changes what the serving path retains — +/// at V4's HCA geometry the window is held at 144 columns where the resizing +/// policy keeps as few as 4 at some residues — and the +23.2% at uniform B=32 +/// that motivates it was measured on this change's own branch, against a tree +/// that no longer exists. That is an unmeasured default, and unmeasured +/// defaults are how this repo lost a week: TCFRAG was default-ON with +/// "UNVERIFIED ON HARDWARE — NEVER RUN" in its own header and held 63 GB +/// (#209), and the fused-512 attention path silently dropped the attention mask +/// for four days at 12% agreement. +/// +/// So: **unverified means default-off, and "unverified" means unmeasured, not +/// new.** +/// +/// # FLIP CONDITION — this default is meant to change +/// +/// Off is a **temporary state with an owner**, not the finish line. Arc's +/// larger problem is not unbuilt work, it is finished, correct, tested work +/// left switched off — a complete GPU sampler with no callers, a fused decode +/// path nobody enabled, 34 tok/s of landed work while master served 17. A flag +/// that ships off with no named experiment behind it is how that happens, so +/// here is the experiment. +/// +/// * **Gate:** `ARC_V4_XS_PIN_WINDOW`, read at +/// `mistralrs-core/src/kv_cache/xs_rolling.rs` in +/// [`xs_pin_window_enabled_from`] (this function). +/// * **The measurement that flips it:** one binary, uniform B=32 decode, same +/// prompt and same seed, two arms — `ARC_V4_XS_PIN_WINDOW=1` against the flag +/// unset. `arc-tools/arcspec_perseq_ladder.sh` already runs exactly this +/// shape and asserts, from the per-process `xs rolling window is …` log line, +/// that the two arms really differed. Pass = the pinned arm is faster on +/// aggregate tok/s, with no change in generated tokens. +/// * **What it was, on its own branch and not on this tree:** +23.2% at uniform +/// B=32. That number was taken on #121 before rebase, against a tree that no +/// longer exists. It is **not a current fact** and must not be quoted as one. +/// * **On pass, the default flips to ON in the same change that records the +/// number.** Leaving it off after the measurement succeeds is a failure +/// state, not a safe one. +/// + +/// Nothing about CUDA-graph capture depends on this default: capture asks for +/// the pin per-cache through [`XsRollingCache::pin_tail_width`], and +/// `pin_is_on` honours that trigger independently of this flag. +/// +/// Read by VALUE and not by presence — #212 converted 23 `ARC_*` flags because +/// `var_os(..).is_some()` made `ARC_FOO=0` mean ON, which silently turns an A/B +/// control arm into a second treatment arm. +pub fn xs_pin_window_enabled_from(value: Option<&str>) -> bool { + value == Some("1") +} + +/// Test-only override for [`xs_pin_window_enabled`], thread-local for the same +/// reason [`test_override`] is. +#[cfg(test)] +pub(crate) mod pin_test_override { + use std::cell::Cell; + + thread_local! { + static STATE: Cell> = const { Cell::new(None) }; + } + + pub(crate) fn current() -> Option { + STATE.with(|s| s.get()) + } + + pub(crate) fn with(on: bool, f: impl FnOnce() -> R) -> R { + let prev = STATE.with(|s| s.replace(Some(on))); + let out = f(); + STATE.with(|s| s.set(prev)); + out + } +} + /// Test-only override for [`xs_per_sequence_enabled`]. /// /// The production read is latched in a `OnceLock` on purpose — a mid-run change @@ -537,8 +700,14 @@ impl XsRollingCache { /// `r + margin + (span_groups - 1) * ratio`, maximised at `r = ratio - 1` /// as `span_groups * ratio + margin - 1`. So pinning here only ever /// retains MORE than the compressor needs, never less. + /// + /// The same quantity as [`Self::window_capacity`], which #121 derived + /// independently for the serving pin. Delegating rather than repeating the + /// expression: two pins that computed "the pinned width" separately could + /// drift, and then capture and serving would disagree about a number whose + /// whole value is that it does not move. pub fn graph_tail_width(&self) -> usize { - self.span_groups * self.ratio + self.margin + self.window_capacity() } /// Hold the retained raw tail at [`Self::graph_tail_width`] columns instead @@ -662,24 +831,103 @@ impl XsRollingCache { )) })?; let have = tail.dim(1)?; - // `want > have` cannot happen (`set_row_lens` refuses such a batch), - // but if it ever did, keeping the whole window and RAISING `base` - // degrades to "refuses a rollback it could have served" — never to - // resuming the compressor from a gap. - let want = (tok - base[i]).min(have); + // What this row actually promised. `promised > have` cannot happen + // (`set_row_lens` refuses such a batch), but if it ever did, keeping + // the whole window and RAISING `base` degrades to "refuses a rollback + // it could have served" — never to resuming the compressor from a gap. + let promised = (tok - base[i]).min(have); + // 🔑 With the window pinned, do NOT re-narrow to `promised` here. + // `clone_out_cache` calls this once per layer per sequence on EVERY + // engine step, so narrowing would reallocate the buffer every step and + // undo the pin exactly on the hot path it exists for. Keeping the + // pinned width costs nothing: the extra columns are older real tokens + // ahead of this row's `base`, and `plan_xs_advance` refuses to read + // below `base` regardless of how wide the buffer is. + // + // Clamped by `tok` so a per-sequence cache never ends up wider than the + // row is long — that state is legal but it would send every subsequent + // B=1 step down `advance_ragged`, which is the slower path. + let width = if self.pin_is_on() { + self.window_capacity().min(tok).min(have).max(promised) + } else { + promised + }; let mut out = self.clone(); out.comp.all_data = Some(comp); out.comp.current_seq_len = tok / self.ratio; - out.tail = Some(if want == have { + out.tail = Some(if width == have { tail } else { - tail.narrow(1, have - want, want)?.contiguous()? + tail.narrow(1, have - width, width)?.contiguous()? }); out.tokens = vec![tok]; - out.base = vec![tok - want]; + // The resume point is this row's own, unchanged by how wide the buffer + // it landed in happens to be. + out.base = vec![tok - promised]; Ok(out) } + /// The width the retained raw window is pinned to — `span_groups * ratio + + /// margin`, the bound derived in `xs_pin_window_enabled` that + /// `tokens - base` provably never reaches. + /// + /// (Not an intra-doc link: `XsRollingCache` is re-exported from the crate + /// root but `xs_pin_window_enabled` is not, so rustdoc reads a public item + /// linking to a private one and warns. Exporting it alongside its exact + /// analogue `xs_per_sequence_enabled` would also close this, but that is an + /// API-surface change and not this PR's business.) + /// + /// A function of the layer's own geometry, so it does not grow with context + /// length: 24 for V4's CSA layers, 144 for HCA. That is the whole point — + /// "the capacity it can never exceed" has to be a small number with a + /// proof, not a context-length ceiling. + pub fn window_capacity(&self) -> usize { + self.span_groups * self.ratio + self.margin + } + + /// The width to retain this step: `needed` when the pin is off (which is + /// bit-identical to what this always did), otherwise the pinned capacity — + /// clamped so it can never exceed what the window physically holds, and + /// never fall below what the rows actually promised. + /// + /// `w_win` matters after a rollback: `set_len` narrows the buffer, so the + /// step after one may have fewer than `capacity` columns to give. Clamping + /// there degrades to "this one step reallocates", never to a truncated + /// history — `needed <= w_win` holds because `tokens - base <= w_phys` is + /// the type's invariant and `w_win = w_phys + t_new`. + fn retained_width(&self, needed: usize, tokens_new: usize, w_win: usize) -> usize { + if !self.pin_is_on() { + return needed; + } + self.window_capacity() + .min(tokens_new) + .min(w_win) + .max(needed) + } + + /// Is the width pin on, by either of its two triggers? + /// + /// There are two, because the pin was arrived at twice for two different + /// reasons and both callers are real: + /// + /// * [`Self::pin_tail_width`] — a per-cache request, made by CUDA-graph + /// capture, for which a moving allocation size is not slow but *invalid* + /// (an allocation that misses the warm pool during capture becomes an + /// unstable graph memory node). See that method for the measured + /// H200 evidence. + /// * [`xs_pin_window_enabled`] — a process-wide serving switch, for which + /// the point is that the reallocation per decode step is simply waste. + /// + /// One width policy answers both: the width each asks for is + /// `span_groups * ratio + margin`, the same expression, so they cannot + /// disagree about what "pinned" means. Keeping both triggers rather than + /// collapsing them keeps capture working when the env switch is off, which + /// is its default. + #[inline] + fn pin_is_on(&self) -> bool { + self.pin_tail || xs_pin_window_enabled() + } + /// Physical width of the retained raw window (0 when there is none). pub fn tail_width(&self) -> Result { match self.tail.as_ref() { @@ -717,11 +965,22 @@ impl XsRollingCache { /// batch with `ARC_V4_XS_PER_SEQ` off satisfies it (one length, one base, /// and `clone_in_cache` refuses any tail-width disagreement outright), so /// the "flag off is byte-identical" claim is unchanged. + /// + /// The width is deliberately **not** part of the test. It was, while the + /// buffer was re-narrowed to `tokens - base` every step; pinning the window + /// ([`xs_pin_window_enabled`]) makes it wider than that on purpose, and the + /// scalar path now measures its offsets from the buffer's physical start, + /// so a wider buffer is no longer an assumption it makes. What it still + /// needs is that the window begins at or after token 0 — `w_phys <= tokens` + /// — because otherwise its `tokens - w_phys` would underflow. That case is + /// a batch assembled around a longer neighbour, and `advance_ragged` is the + /// path written for it. fn can_advance_uniform(&self) -> Result { let (t0, b0) = (self.tokens[0], self.base[0]); Ok(self.tokens.iter().all(|&t| t == t0) && self.base.iter().all(|&b| b == b0) - && self.tail_width()? == t0 - b0) + && b0 <= t0 + && self.tail_width()? <= t0) } pub fn reset(&mut self) { @@ -1149,10 +1408,24 @@ impl XsRollingCache { { let tokens = self.tokens[0]; let base = self.base[0]; + let w_phys = self.tail_width()?; - // `window` covers tokens [win_start, tokens + t_new). + // `window` covers tokens [win_start, tokens + t_new), where `win_start` + // is the buffer's PHYSICAL start, not `base`. + // + // Those were the same number while the buffer was re-narrowed to + // exactly `tokens - base` every step. They are not once it is pinned: a + // pinned window keeps older columns this row no longer promises, so its + // physical start sits at or below `base`. Measuring offsets from the + // physical start is what makes the two cases one code path — and it is + // the same thing the ragged path already does, where a neighbour's + // width has always been able to widen the window. + // + // `tokens - w_phys` cannot underflow here: `can_advance_uniform` + // refuses this path when `w_phys > tokens`, sending that case to + // `advance_ragged`, which handles a window wider than the row is long. let (window, win_start) = match self.tail.as_ref() { - Some(tail) if tokens > base => (Tensor::cat(&[tail, xs_new], 1)?, base), + Some(tail) if w_phys > 0 => (Tensor::cat(&[tail, xs_new], 1)?, tokens - w_phys), _ => (xs_new.clone(), tokens), }; let tokens_new = tokens + t_new; @@ -1167,11 +1440,16 @@ impl XsRollingCache { need_start, (g_done + 1).saturating_sub(self.span_groups) * self.ratio ); - if need_start < win_start { + // Bounded by `base`, not by the buffer's physical start — the same + // tightening `plan_xs_advance` carries. A pinned window physically + // holds tokens below `base`, and they are real, but they are not + // history this row promised to keep; reading them because the + // buffer happens to be wider would make the pin change answers. + if need_start < base { candle_core::bail!( "xs rolling cache: compressor history gap — row {g_done} needs tokens from \ - {need_start} but the retained window starts at {win_start}. This means a \ - truncation was applied without going through `set_len`." + {need_start} but that row only retained from {base}. This means a truncation \ + was applied without going through `set_len`." ); } let off = need_start - win_start; @@ -1209,22 +1487,32 @@ impl XsRollingCache { let rollback_floor = tokens_new.saturating_sub(self.margin); let keep_from = ((rollback_floor / self.ratio + 1).saturating_sub(self.span_groups) * self.ratio) - .max(win_start); - // Under `pin_tail_width` the retention point is the FIXED-width one - // instead, so the tail (and the cat, narrow and contiguous that build - // it) keeps one shape for the whole run — a capture precondition, see - // `pin_tail_width`. `graph_tail_width()` is an upper bound on what the - // unpinned policy would keep, so this only ever retains more, never - // less; `.max(win_start)` covers the transition step, where the window - // in hand may still be narrower than the pinned width. - let pinned_base = (self.pin_tail && tokens_new >= self.graph_tail_width()) - .then(|| (tokens_new - self.graph_tail_width()).max(win_start)); - let new_base = pinned_base.unwrap_or(keep_from).min(tokens_new); - self.tail = Some( - window - .narrow(1, new_base - win_start, tokens_new - new_base)? - .contiguous()?, - ); + .max(base); + let new_base = keep_from.min(tokens_new); + // End-anchored, which is what it always was: with the pin off, + // `w_win - needed == new_base - win_start` identically, so this is the + // same narrow the scalar path has always taken. With it on, the extra + // columns are older real tokens ahead of `new_base` that nothing reads. + // + // This supersedes the `pinned_base` form that reached the integration + // branch first (#181's `pin_tail_width`). Both hold the retained tail + // at the SAME constant width -- `graph_tail_width()` and + // `window_capacity()` are the same expression, `span_groups * ratio + + // margin` -- which is the property capture needs. They differed only in + // where the constancy came from: that one moved `base` earlier so the + // logical span `tokens - base` became constant, this one leaves `base` + // at the retention point the type documents and widens the physical + // buffer instead. This form is the one that composes, because + // `win_start` here is the buffer's PHYSICAL start (`tokens - w_phys`) + // rather than `base`, so a buffer wider than the row promised is + // already a representable state -- and `plan_xs_advance` still refuses + // to read below `base`, so the pin cannot change any answer. Both + // triggers are honoured; see `retained_width`. + let w_win = window.dim(1)?; + let needed = tokens_new - new_base; + let width = self.retained_width(needed, tokens_new, w_win); + debug_assert!(width >= needed && width <= w_win); + self.tail = Some(window.narrow(1, w_win - width, width)?.contiguous()?); for (t, ba) in self.tokens.iter_mut().zip(self.base.iter_mut()) { *t = tokens_new; *ba = new_base; @@ -1320,11 +1608,10 @@ impl XsRollingCache { // row's own token count and `plan_xs_advance` refuses any that would // reach below the row's `base`. let w_win = window.dim(1)?; - self.tail = Some( - window - .narrow(1, w_win - plan.tail_width, plan.tail_width)? - .contiguous()?, - ); + let widest_row = plan.tokens_new.iter().copied().max().unwrap_or(0); + let width = self.retained_width(plan.tail_width, widest_row, w_win); + debug_assert!(width >= plan.tail_width && width <= w_win); + self.tail = Some(window.narrow(1, w_win - width, width)?.contiguous()?); self.tokens = plan.tokens_new; self.base = plan.base_new; self.compressed_rows() @@ -1698,3 +1985,60 @@ mod tests { ); } } + +#[cfg(test)] +mod pin_gate_polarity_tests { + use super::*; + + /// The pin is OFF unless explicitly asked for. + /// + /// It shipped default-ON in #121, which made the A/B's control arm the one + /// carrying `=0`. That is the wrong way round for something unmeasured on + /// this tree: the pin changes what serving retains (144 columns at V4's HCA + /// geometry, against as few as 4 at some residues under the resizing + /// policy), and its +23.2% at uniform B=32 was measured on #121's own + /// branch. Same lesson as #209's TCFRAG retirement. + /// + /// This is NOT a parking assertion. See `xs_pin_window_enabled_from`'s doc + /// comment for the single measurement that flips this default back on. + #[test] + fn the_pin_is_off_unless_explicitly_enabled() { + assert!( + !xs_pin_window_enabled_from(None), + "unset must mean OFF — the pin is unmeasured on the current tree" + ); + assert!(xs_pin_window_enabled_from(Some("1"))); + } + + /// Read by VALUE. #212 converted 23 `ARC_*` flags because + /// `var_os(..).is_some()` made `ARC_FOO=0` mean ON — which turns an A/B's + /// control arm into a second treatment arm and reports a clean 1.000x. + #[test] + fn zero_means_off_and_so_does_every_other_value() { + for v in ["0", "", "false", "no", "off", "true", "yes", "on", "2", " 1"] { + assert!( + !xs_pin_window_enabled_from(Some(v)), + "{v:?} must not enable the pin; only the exact string \"1\" does" + ); + } + } + + /// Capture must not depend on the env default: it asks per-cache through + /// `pin_tail_width()`, and `pin_is_on` honours that trigger on its own. + /// If this regressed, CUDA-graph capture would silently go back to a + /// per-step-varying allocation size — which is not a slow graph but an + /// invalid one. + #[test] + fn capture_still_gets_a_constant_width_with_the_env_flag_off() { + pin_test_override::with(false, || { + let mut c = XsRollingCache::new(4, 2, 16, 64); + assert!(!c.pin_is_on(), "env off and pin_tail unset => not pinned"); + c.pin_tail_width(); + assert!( + c.pin_is_on(), + "capture's per-cache trigger must work with the env flag off" + ); + assert_eq!(c.graph_tail_width(), c.window_capacity()); + }); + } +} diff --git a/mistralrs-core/src/lib.rs b/mistralrs-core/src/lib.rs index 26b8b14d9..0c9b1f2d9 100644 --- a/mistralrs-core/src/lib.rs +++ b/mistralrs-core/src/lib.rs @@ -174,7 +174,9 @@ pub use response::*; pub use sampler::{ CustomLogitsProcessor, DrySamplingParams, SamplingParams, StopTokens, TopLogprob, }; -pub use scheduler::{DefaultSchedulerMethod, SchedulerConfig}; +pub use scheduler::{ + sched_bucket_marker, sched_buckets, DefaultSchedulerMethod, SchedBuckets, SchedulerConfig, +}; pub use search::{SearchCallback, SearchFunctionParameters, SearchResult}; use serde::Serialize; pub use speech_models::{utils as speech_utils, SpeechGenerationConfig, SpeechLoaderType}; diff --git a/mistralrs-core/src/models/deepseek4.rs b/mistralrs-core/src/models/deepseek4.rs index 0aa8909c3..766a5f59c 100644 --- a/mistralrs-core/src/models/deepseek4.rs +++ b/mistralrs-core/src/models/deepseek4.rs @@ -5957,6 +5957,9 @@ mod kv_footprint_tests { #[cfg(test)] mod tests { use super::*; + // Only the tests drive the pin override; importing it at module scope made + // it an unused import on every non-test build. + use crate::kv_cache::xs_rolling; /// The **verbatim** `config.json` published at /// @@ -8149,6 +8152,156 @@ mod tests { ragged_batch_matches_b1(128, &[129, 143, 1024, 1150], 3) } + /// 🔑 Settles the one discrepancy between two independent models of this + /// buffer, and does it without a GPU. + /// + /// The ArcGraph chain measured the reallocation cycling through + /// `4096 x {18, 19, 20, 21}`. Running the retention rule forward predicts + /// `{20, 21, 22, 23}`. Both are `ratio`-consecutive, both contain 21, and + /// they are offset by exactly 2 — so one of the two models is wrong about + /// *why*, even though the bound holds either way. + /// + /// The pre-committed reconciliation was: theirs is the pre-saturation RAMP + /// (their run generated ~21 tokens, so `base` had not finished advancing + /// and `W = tokens - base` was still climbing), mine is the steady state. + /// This checks that claim directly by recording the width from token 1, + /// and it is the difference between a bound that is right for the right + /// reason and one that happens to be right — which is the trap the next + /// context length springs. + #[test] + fn the_window_ramps_then_settles_to_ratio_consecutive_sizes() -> Result<()> { + xs_rolling::pin_test_override::with(false, || { + let device = Device::Cpu; + let (hidden, head_dim, ratio) = (32usize, 16usize, 4usize); + let compressor = rolling_test_compressor(ratio, hidden, head_dim, &device)?; + let stream = per_row_streams(1, 130, hidden, &device)?.remove(0); + let mut s = XsRollingCache::new(ratio, compressor.coff, head_dim, 4096); + + let mut widths = Vec::new(); + for i in 0..128 { + s.advance(&stream.narrow(1, i, 1)?, |w| compressor.forward_from_xs(w))?; + widths.push(s.tail_width()?); + } + let cap = ratio * compressor.coff + crate::kv_cache::XS_TAIL_MARGIN_TOKENS; + + // The ramp: early widths are BELOW the steady-state band, because + // `base` is still pinned at 0 while `tokens` climbs. + let ramp: std::collections::BTreeSet<_> = widths[..20].iter().copied().collect(); + let steady: std::collections::BTreeSet<_> = widths[64..].iter().copied().collect(); + assert!( + ramp.iter().min() < steady.iter().min(), + "the early widths must be a ramp, not the steady band: ramp {ramp:?} \ + steady {steady:?}" + ); + // The steady state is exactly `ratio` consecutive sizes… + assert_eq!( + steady.len(), + ratio, + "steady-state width must cycle through exactly `ratio` sizes, got {steady:?}" + ); + let lo = *steady.iter().next().unwrap(); + assert_eq!( + steady.iter().copied().collect::>(), + (lo..lo + ratio).collect::>(), + "the steady band must be consecutive, got {steady:?}" + ); + // …and it is the top of the range the bound allows: [cap-ratio, cap). + assert_eq!( + (lo, lo + ratio - 1), + (cap - ratio, cap - 1), + "steady band should be [cap-ratio, cap-1] = [{}, {}]", + cap - ratio, + cap - 1 + ); + // The bound, over every step including the ramp. + assert!( + widths.iter().all(|&w| w < cap), + "some width reached the pinned capacity {cap}, so pinning there could truncate: \ + max was {:?}", + widths.iter().max() + ); + Ok(()) + }) + } + + /// 🔑 Pinning the retained window must change the ALLOCATION and nothing + /// else. This is the evidence that the pin is SAFE — the two settings are + /// run against the same stream for 40 steps and required to agree EXACTLY, + /// on the compressed rows and on both time bases. + /// + /// It was once also offered as the evidence for defaulting + /// `ARC_V4_XS_PIN_WINDOW` on without a throughput number. It is not that: + /// it proves the pin changes no answer, which is necessary and not + /// sufficient. The flag is opt-in until the throughput number exists on + /// this tree — see `xs_rolling::xs_pin_window_enabled_from`'s FLIP + /// CONDITION for the experiment that settles it. + /// + /// The argument for why they must — every offset is derived from the row's + /// own token count, so a wider buffer shifts `off` by exactly the widening + /// and the compressor sees the same absolute tokens — is the kind of + /// argument that has been wrong on this chain before. So it is checked. + #[test] + fn pinning_the_window_is_numerically_inert() -> Result<()> { + let device = Device::Cpu; + let (hidden, head_dim, ratio) = (32usize, 16usize, 4usize); + let compressor = rolling_test_compressor(ratio, hidden, head_dim, &device)?; + let stream = per_row_streams(1, 220, hidden, &device)?.remove(0); + + let run = |pinned: bool| -> Result<(Tensor, Vec, Vec, Vec)> { + xs_rolling::pin_test_override::with(pinned, || { + let mut s = XsRollingCache::new(ratio, compressor.coff, head_dim, 4096); + s.advance(&stream.narrow(1, 0, 37)?, |w| compressor.forward_from_xs(w))?; + let mut widths = Vec::new(); + for i in 0..40 { + s.advance(&stream.narrow(1, 37 + i, 1)?, |w| { + compressor.forward_from_xs(w) + })?; + widths.push(s.tail_width()?); + } + let comp = s.compressed_rows()?.expect("past one full group"); + let (tok, base) = s.row_lens(); + Ok((comp, tok.to_vec(), base.to_vec(), widths)) + }) + }; + + let (comp_off, tok_off, base_off, widths_off) = run(false)?; + let (comp_on, tok_on, base_on, widths_on) = run(true)?; + + assert_eq!( + max_abs_diff(&comp_off, &comp_on)?, + 0.0, + "pinning the window changed the compressed rows — it must only change the allocation" + ); + assert_eq!(tok_off, tok_on, "token counts diverged"); + assert_eq!(base_off, base_on, "resume points diverged"); + + // …and the fixture must actually exercise the difference, or the + // equality above proves nothing (five fixtures on this chain have + // survived a mutation by not distinguishing the two answers). + let distinct_off: std::collections::BTreeSet<_> = widths_off.iter().collect(); + assert!( + distinct_off.len() > 1, + "unpinned widths did not vary ({widths_off:?}), so this proves nothing about pinning" + ); + assert_eq!( + distinct_off.len(), + ratio, + "unpinned width should cycle through `ratio` consecutive sizes, got {distinct_off:?}" + ); + let cap = ratio * compressor.coff + crate::kv_cache::XS_TAIL_MARGIN_TOKENS; + assert!( + widths_on.iter().all(|&w| w == cap), + "pinned widths must be the constant {cap}, got {widths_on:?}" + ); + // The bound the pin rests on: the unpinned width never reaches it. + assert!( + widths_off.iter().all(|&w| w < cap), + "the pinned capacity {cap} must exceed every width the retention rule produces, \ + got {widths_off:?}" + ); + Ok(()) + } + /// The control the ragged tests need: a UNIFORM batch takes the untouched /// scalar path and is also token-identical. If this ever failed, the /// "flag off is byte-identical" claim would be false and the ragged @@ -8297,8 +8450,73 @@ mod tests { /// the per-sequence invariant `window width == tokens - base`, taking that /// row's share from the END of the shared window. Taking it from the front /// would hand the row somebody else's older tokens under its own `base`. + /// ⚠️ Run with the window pin OFF, deliberately. Its premise is a row + /// *narrower* than the shared window, which is what a resizing buffer + /// produces; pinned, every row is the same width and the premise — and the + /// re-anchoring it checks — cannot be constructed. The re-anchoring still + /// has to be right when the pin is off, and + /// `splitting_a_pinned_row_keeps_the_buffer_and_the_resume_point` covers + /// the pinned side, where the requirement is the opposite one. #[test] fn splitting_a_batched_row_restores_the_per_sequence_window() -> Result<()> { + xs_rolling::pin_test_override::with(false, splitting_a_batched_row_inner) + } + + /// 🔑 The pinned counterpart, and the reason it matters is cost, not + /// correctness: `clone_out_cache` calls `split_row` once per layer per + /// sequence on EVERY engine step. A split that re-narrowed to + /// `tokens - base` would reallocate the buffer every step and undo the pin + /// precisely on the hot path it exists for. So the split must hand back the + /// SAME tensor, and must not move the row's resume point while doing it. + #[test] + fn splitting_a_pinned_row_keeps_the_buffer_and_the_resume_point() -> Result<()> { + xs_rolling::pin_test_override::with(true, || { + let device = Device::Cpu; + let (hidden, head_dim, ratio) = (32usize, 16usize, 4usize); + let compressor = rolling_test_compressor(ratio, hidden, head_dim, &device)?; + let lens = [37usize, 40]; + let streams = per_row_streams(2, 64, hidden, &device)?; + let mut refs: Vec = Vec::new(); + for (i, &l) in lens.iter().enumerate() { + let mut s = XsRollingCache::new(ratio, compressor.coff, head_dim, 4096); + s.advance(&streams[i].narrow(1, 0, l)?, |w| { + compressor.forward_from_xs(w) + })?; + refs.push(s); + } + let cap = ratio * compressor.coff + crate::kv_cache::XS_TAIL_MARGIN_TOKENS; + for r in &refs { + assert_eq!( + r.tail.as_ref().unwrap().dim(1)?, + cap, + "a pinned per-sequence window must already be the capacity" + ); + } + let batched = batch_xs(&refs)?; + let shared = batched.tail.as_ref().unwrap().dim(1)?; + assert_eq!(shared, cap, "pinned rows batch without any padding at all"); + + let comps = batched.comp.all_data.as_ref().unwrap().chunk(2, 0)?; + let tails = batched.tail.as_ref().unwrap().chunk(2, 0)?; + for i in 0..2 { + let out = batched.split_row(i, comps[i].clone(), tails[i].clone())?; + assert_eq!( + out.tail.as_ref().unwrap().dim(1)?, + cap, + "row {i}: the split re-narrowed a pinned buffer, which would reallocate it \ + on every engine step" + ); + assert_eq!( + (out.row_lens().0[0], out.row_lens().1[0]), + (refs[i].row_lens().0[0], refs[i].row_lens().1[0]), + "row {i}: keeping the buffer must not move the resume point" + ); + } + Ok(()) + }) + } + + fn splitting_a_batched_row_inner() -> Result<()> { let device = Device::Cpu; let (hidden, head_dim, ratio) = (32usize, 16usize, 4usize); let compressor = rolling_test_compressor(ratio, hidden, head_dim, &device)?; diff --git a/mistralrs-core/src/pipeline/mtp_pipeline.rs b/mistralrs-core/src/pipeline/mtp_pipeline.rs index b026a8df6..169331801 100644 --- a/mistralrs-core/src/pipeline/mtp_pipeline.rs +++ b/mistralrs-core/src/pipeline/mtp_pipeline.rs @@ -1522,6 +1522,16 @@ impl AcceptanceTelemetry { for (b, per_b) in self.snapshot_by_batch() { tracing::info!(target: "mtp_speculative", "{}", per_b.marker(&format!("b={b}"))); } + // 🔑 On the SAME fence, deliberately. `tok_per_step` and aggregate + // throughput are both read against how wide the steps actually were, + // and both schedulers run one cache-length bucket per step — so a batch + // can be serialised for reasons that have nothing to do with MTP. A + // harness that differences these counters across a wall-clock boundary + // gets the scheduler's window and MTP's window identical only if they + // are emitted together. See `scheduler::bucket_telemetry`. + if let Some(sched) = crate::scheduler::sched_bucket_marker() { + tracing::info!(target: "mtp_speculative", "{sched}"); + } tracing::info!(target: "mtp_speculative", "{}", snap.report_line()); // The per-position breakdown, in prose, right where the aggregate rate // is. This is the line the acceptance-gap diagnosis needs: the scalar diff --git a/mistralrs-core/src/scheduler/bucket_telemetry.rs b/mistralrs-core/src/scheduler/bucket_telemetry.rs new file mode 100644 index 000000000..2b0985a4f --- /dev/null +++ b/mistralrs-core/src/scheduler/bucket_telemetry.rs @@ -0,0 +1,241 @@ +//! Parent system: ArcInfer / ArcSched +//! +//! Process-wide counters for the one thing that decides how wide a decode step +//! actually is: **bucketing**. +//! +//! Both schedulers partition the running set by an exact cache length and run +//! **one bucket per step**, preempting the rest +//! (`default_scheduler.rs:bucket_and_waitlist_seqs_waiting`, +//! `paged_attention/scheduler.rs`). So a throughput number measured at "B=128" +//! is not necessarily a number about 128-wide steps — if the batch shattered +//! into 40 buckets, the engine ran 3-wide steps and the label lies. +//! +//! That is not a hypothetical confound for the per-sequence-advance A/B: the +//! whole point of ragged admission ([`super::RaggedAdmission`]) is to collapse +//! those buckets into one, so an aggregate throughput delta between the ON and +//! OFF arms could come from the KV mechanism, from the bucketing, or from both. +//! Without these counters the two are indistinguishable, and "aggregate did not +//! move" would be unattributable. +//! +//! The marker is emitted on the same log fence as `MTP[agg]`, so a harness that +//! differences cumulative counters across a wall-clock boundary gets the +//! scheduler's numbers for exactly the window it got MTP's. + +use std::sync::atomic::{AtomicUsize, Ordering}; + +/// Cumulative bucketing counters. Monotone; difference two snapshots to get a +/// window. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct SchedBuckets { + /// Bucketing calls with at least one running sequence. + pub calls: usize, + /// Sum of the number of buckets formed. + pub buckets: usize, + /// Sum of the running-set size *before* selection — the batch the engine + /// was offered. + pub offered: usize, + /// Sum of the winning bucket's size — the batch the engine actually ran. + pub chosen: usize, + /// Calls that formed more than one bucket, i.e. that preempted somebody. + pub shattered: usize, +} + +impl SchedBuckets { + /// `self - earlier`, for differencing across a fence. Saturating, so a + /// counter reset between snapshots degrades to zero rather than to a + /// gigantic wrapped number that would read as a spectacular result. + #[must_use] + pub fn since(&self, earlier: &Self) -> Self { + Self { + calls: self.calls.saturating_sub(earlier.calls), + buckets: self.buckets.saturating_sub(earlier.buckets), + offered: self.offered.saturating_sub(earlier.offered), + chosen: self.chosen.saturating_sub(earlier.chosen), + shattered: self.shattered.saturating_sub(earlier.shattered), + } + } + + #[allow(clippy::cast_precision_loss)] + fn per_call(&self, total: usize) -> Option { + (self.calls > 0).then(|| total as f64 / self.calls as f64) + } + + /// Mean buckets formed per scheduling step. 1.0 means the batch was never + /// split; 40.0 at a nominal B=128 means the engine ran ~3-wide steps. + #[must_use] + pub fn buckets_per_step(&self) -> Option { + self.per_call(self.buckets) + } + + /// Mean size of the bucket that actually ran — the **real** batch width, + /// as opposed to the number of sequences in flight. + #[must_use] + pub fn running_bucket_size(&self) -> Option { + self.per_call(self.chosen) + } + + /// Mean number of running sequences offered to the scheduler per step. + #[must_use] + pub fn offered_per_step(&self) -> Option { + self.per_call(self.offered) + } + + /// Mean sequences preempted per step (`offered - chosen`). + #[must_use] + pub fn preempted_per_step(&self) -> Option { + self.per_call(self.offered.saturating_sub(self.chosen)) + } + + /// Fraction of steps that split the batch at all. + #[allow(clippy::cast_precision_loss)] + #[must_use] + pub fn shattered_frac(&self) -> Option { + (self.calls > 0).then(|| self.shattered as f64 / self.calls as f64) + } + + /// The machine-greppable one-liner, in the project's marker convention + /// (`MTP[...]`, `SPEED[...]`). Every raw count is on the line so the ratios + /// are auditable without trusting the formatter. + #[must_use] + pub fn marker(&self, scope: &str) -> String { + let fmt = |v: Option| v.map_or_else(|| "n/a".to_string(), |x| format!("{x:.4}")); + format!( + "SCHED[{scope}] calls={} buckets={} offered={} chosen={} shattered={} \ + buckets_per_step={} running_bucket_size={} offered_per_step={} \ + preempted_per_step={} shattered_frac={}", + self.calls, + self.buckets, + self.offered, + self.chosen, + self.shattered, + fmt(self.buckets_per_step()), + fmt(self.running_bucket_size()), + fmt(self.offered_per_step()), + fmt(self.preempted_per_step()), + fmt(self.shattered_frac()), + ) + } +} + +static CALLS: AtomicUsize = AtomicUsize::new(0); +static BUCKETS: AtomicUsize = AtomicUsize::new(0); +static OFFERED: AtomicUsize = AtomicUsize::new(0); +static CHOSEN: AtomicUsize = AtomicUsize::new(0); +static SHATTERED: AtomicUsize = AtomicUsize::new(0); + +/// Record one bucketing decision. Called once per scheduling step that had +/// anything to schedule. +pub fn record_bucketing(buckets: usize, offered: usize, chosen: usize) { + if offered == 0 { + return; + } + CALLS.fetch_add(1, Ordering::Relaxed); + BUCKETS.fetch_add(buckets, Ordering::Relaxed); + OFFERED.fetch_add(offered, Ordering::Relaxed); + CHOSEN.fetch_add(chosen, Ordering::Relaxed); + if buckets > 1 { + SHATTERED.fetch_add(1, Ordering::Relaxed); + } +} + +/// Snapshot of the process-wide counters. +#[must_use] +pub fn sched_buckets() -> SchedBuckets { + SchedBuckets { + calls: CALLS.load(Ordering::Relaxed), + buckets: BUCKETS.load(Ordering::Relaxed), + offered: OFFERED.load(Ordering::Relaxed), + chosen: CHOSEN.load(Ordering::Relaxed), + shattered: SHATTERED.load(Ordering::Relaxed), + } +} + +/// The aggregate `SCHED[agg] …` line, or `None` when nothing was ever +/// scheduled — the honest answer to "how wide were the steps" when none ran is +/// *nothing*, not `0`. +#[must_use] +pub fn sched_bucket_marker() -> Option { + let snap = sched_buckets(); + (snap.calls > 0).then(|| snap.marker("agg")) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The ratios are what get quoted, so they are what must be pinned. A batch + /// that shatters is the case the counters exist to make visible. + #[test] + fn a_shattered_batch_reports_the_width_that_actually_ran() { + // Three steps: 128 offered, split 40 ways, 3 of them ran. + let mut snap = SchedBuckets::default(); + for _ in 0..3 { + snap.calls += 1; + snap.buckets += 40; + snap.offered += 128; + snap.chosen += 3; + snap.shattered += 1; + } + assert_eq!(snap.buckets_per_step(), Some(40.0)); + assert_eq!(snap.offered_per_step(), Some(128.0)); + // 🔑 The number a throughput result has to be read against: the engine + // ran 3-wide steps while the harness called it B=128. + assert_eq!(snap.running_bucket_size(), Some(3.0)); + assert_eq!(snap.preempted_per_step(), Some(125.0)); + assert_eq!(snap.shattered_frac(), Some(1.0)); + let m = snap.marker("agg"); + assert!(m.contains("running_bucket_size=3.0000"), "{m}"); + assert!(m.contains("buckets_per_step=40.0000"), "{m}"); + } + + /// An unshattered batch must be visibly different, or the counter cannot + /// distinguish "the scheduler serialised us" from "it did not". + #[test] + fn an_intact_batch_reports_one_bucket_and_no_preemption() { + let snap = SchedBuckets { + calls: 10, + buckets: 10, + offered: 1280, + chosen: 1280, + shattered: 0, + }; + assert_eq!(snap.buckets_per_step(), Some(1.0)); + assert_eq!(snap.running_bucket_size(), Some(128.0)); + assert_eq!(snap.preempted_per_step(), Some(0.0)); + assert_eq!(snap.shattered_frac(), Some(0.0)); + } + + /// Differencing across a fence is how a cell's numbers are taken; a reset + /// in between must not read as a huge result. + #[test] + fn differencing_saturates_rather_than_wrapping() { + let early = SchedBuckets { + calls: 100, + buckets: 200, + offered: 1000, + chosen: 500, + shattered: 50, + }; + let late = SchedBuckets { + calls: 150, + buckets: 260, + offered: 1600, + chosen: 900, + shattered: 60, + }; + let d = late.since(&early); + assert_eq!(d.calls, 50); + assert_eq!(d.buckets, 60); + assert_eq!(d.running_bucket_size(), Some(8.0)); + // A counter that went backwards (process restart) reads as zero work, + // never as a wrapped `usize`. + assert_eq!(early.since(&late), SchedBuckets::default()); + } + + /// Nothing scheduled means no marker, not a marker full of zeros. + #[test] + fn no_scheduling_reports_nothing_rather_than_zero() { + assert!(SchedBuckets::default().buckets_per_step().is_none()); + assert!(SchedBuckets::default().running_bucket_size().is_none()); + } +} diff --git a/mistralrs-core/src/scheduler/default_scheduler.rs b/mistralrs-core/src/scheduler/default_scheduler.rs index 16d3466da..e3f550839 100644 --- a/mistralrs-core/src/scheduler/default_scheduler.rs +++ b/mistralrs-core/src/scheduler/default_scheduler.rs @@ -268,6 +268,13 @@ impl BucketingManager for FixedBucketingManager { discrete: bool, ) -> BucketedSeqs { // Now, get the sequences with the smallest sequence lengths, and allow them to catch up. + // + // How wide the step the engine is about to run actually is. A batch + // that shatters into N buckets runs one of them and preempts the rest, + // so "B=128" in a harness can be 3-wide in the engine — and a + // throughput number is unreadable without knowing which. See + // [`super::bucket_telemetry`]. + let offered = running.len(); let mut seq_buckets: HashMap> = HashMap::new(); let mut seq_priorities: HashMap = HashMap::new(); let ragged_decode = crate::kv_cache::ragged_decode_supported(); @@ -355,6 +362,7 @@ impl BucketingManager for FixedBucketingManager { ); } } + let n_buckets = seq_buckets.len(); let running = if seq_buckets.len() <= 1 { // Full steam ahead or have everything seq_buckets @@ -430,6 +438,7 @@ impl BucketingManager for FixedBucketingManager { // Know min_seqs.len < running.len() <= max highest_priority_seqs }; + super::bucket_telemetry::record_bucketing(n_buckets, offered, running.len()); BucketedSeqs { running, waiting } } } diff --git a/mistralrs-core/src/scheduler/mod.rs b/mistralrs-core/src/scheduler/mod.rs index 34c8d56fb..f3606ac3b 100644 --- a/mistralrs-core/src/scheduler/mod.rs +++ b/mistralrs-core/src/scheduler/mod.rs @@ -1,7 +1,9 @@ +pub mod bucket_telemetry; mod default_scheduler; use std::sync::Arc; +pub use bucket_telemetry::{sched_bucket_marker, sched_buckets, SchedBuckets}; pub use default_scheduler::{DefaultScheduler, DefaultSchedulerMethod, DefaultSchedulerOutput}; use tokio::sync::Mutex; diff --git a/mistralrs-core/tests/capability_reachability.rs b/mistralrs-core/tests/capability_reachability.rs index bffe418f7..0ce630f50 100644 --- a/mistralrs-core/tests/capability_reachability.rs +++ b/mistralrs-core/tests/capability_reachability.rs @@ -119,6 +119,39 @@ static REGISTRY: &[Capability] = &[ }, status: Status::Live, }, + // ── SHIPPED OFF, WITH A NAMED EXPERIMENT ──────────────────────────────── + // + // The entry below is `Live` because the GATE is reached on every step — + // that is what this file checks, and it must not go dark. What is switched + // off is the gate's DEFAULT ANSWER, which no reachability check can see. + // + // Recording it here anyway, because "built, correct, tested, and switched + // off" is this repo's most expensive failure mode and this file is where + // someone goes looking for it. 26 pieces of finished work are parked that + // way right now, including a complete GPU sampler with zero callers; we + // shipped 34 tok/s of landed work while master served 17, for exactly this + // reason. A flag that ships off with no named experiment behind it is how + // that happens, so the experiment is named: + // + // flag ARC_V4_XS_PIN_WINDOW=1 + // gate kv_cache/xs_rolling.rs :: xs_pin_window_enabled_from + // test arc-tools/arcspec_perseq_ladder.sh — ON vs ON_PINNED, uniform + // B=32, one binary, same prompt and seed + // pass ON_PINNED faster on aggregate tok/s, generated tokens identical + // was +23.2% at uniform B=32 on #121's own pre-rebase branch — NOT a + // current fact, do not quote it as one + // then the default flips to ON in the same change that records the + // number. Leaving it off after a passing measurement is a failure + // state, not the finish line. + Capability { + name: "xs window pin: built and tested, default OFF pending one measurement", + parent: "ArcInfer / ArcKV", + check: Check::Symbol { + symbol: "xs_pin_window_enabled", + defined_in: "mistralrs-core/src/kv_cache/xs_rolling.rs", + }, + status: Status::Live, + }, Capability { name: "ragged decode: per-row lengths on the xs cache", parent: "ArcInfer / ArcKV",