From 7c630960e27664b2576237245814b61ae8d89dfc Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick Date: Thu, 20 Aug 2026 22:38:14 +0100 Subject: [PATCH 01/21] docs: the single OpenRouter-readiness gate One branch, one PR, one merge decision. Merging it means Arc can serve and sell tokens on OpenRouter. 16 items across blocking / product / trust / commercial, each with its measured provenance. Co-Authored-By: Claude Opus 5 (1M context) --- docs/engineering/OPENROUTER_READY.md | 38 ++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 docs/engineering/OPENROUTER_READY.md diff --git a/docs/engineering/OPENROUTER_READY.md b/docs/engineering/OPENROUTER_READY.md new file mode 100644 index 000000000..965328541 --- /dev/null +++ b/docs/engineering/OPENROUTER_READY.md @@ -0,0 +1,38 @@ +**Merging this PR means Arc can serve and sell tokens on OpenRouter.** Nothing lands on `master` for this effort except through here. + +Base `d7742670a` (the 15.0 โ†’ **34.48 tok/s** b=1 ladder, already on master). + +--- + +## ๐Ÿ”ด Blocking โ€” cannot sell tokens at all + +- [ ] **1. Time-to-first-token is ~24 s.** ~12 ms per *input* token, ~100ร— off bound. A user on OpenRouter feels this immediately. **Prefill is a different disease from decode โ€” the GPU is already 99.5% busy, so graphs/launch-reduction buy it nothing.** Needs its own lane and its own kernel work. +- [ ] **2. Prefill crashes.** b=1 on `qtip2b` exits 139 (SIGSEGV) and the harness prints `pp = 0.000ยฑ0.000` into a results table instead of failing. A 0.000 is not a measurement. +- [x] **3. The fast build reaches `master`.** Landed as #175/#177/#185โ€“#189 โ†’ `d7742670a`. โš ๏ธ Still owed: re-measure **master itself** and confirm it reads ~34.48. +- [ ] **4. Aggregate throughput on the fast build.** Never measured at any batch size. Sweep B = 1/8/32/64/128, per-request **and** aggregate. Expect flat-then-steep: 256 experts top-6 means tokens barely collide until batch is large. **Finding the knee is the deliverable** โ€” the GPU-hour economics depend on it. +- [ ] **5. Reliability under sustained load.** 400 requests passed once. Nobody has run this for hours. + +## ๐ŸŸก Product quality + +- [ ] **6. Generation speed โ†’ 250 tok/s.** Real profile on `qtip2b`: tail (sinkhorn + fast_sum + bmul + ucopy) **27.2%**, `fp8_gemv_warp` **16.0%**, `qtip2b_gemv_tuned_kernel` **15.4%**. ๐Ÿšฉ A **4ร—4** Sinkhorn costing 9.3% of an H200 is structural, not tuning. +- [ ] **7. Long-context memory, measured.** At 1M ctx the answer is **~3 users or ~55** depending on whether the 128-token window + compress ratios are live. Unmeasured. This is the capacity claim. +- [ ] **8. Prefix caching is silently disabled** by our own default (TurboQuant is the default cache type). Most real traffic is repeated prefixes. +- [ ] **9. Switched-off finished work.** `CudaSampler` โ€” complete, tested bit-exact, **zero callers, no gate**. The **GPU WHILE-node loop** โ€” the one design that removes the per-token sync; neither competitor has it. Blocked by an unconditional `return Ok(None)` at `pipeline/normal.rs:2324` and a U8-dtype bug failing all 28 call sites. +- [ ] **10. 56ร— KV over-retention** โ€” root-caused previously; the fix reportedly never reached V4. + +## ๐ŸŸข Trust and correctness + +- [ ] **11. The live V4 test fixture is all zeros** โ€” identity tests against it prove nothing. +- [ ] **12. Verification code is less reliable than production code.** Twelve instruments lied in one session: bench counting *requested* tokens (printed **1280 tok/s** for a zero-token run), `cargo clean -p` silently no-opping on git deps, `nsys` hiding graph-internal kernels by default, a stale `.o` satisfying a link, `replay()` returning a clone of the graded buffer, `append_graph` comparing a buffer against itself, an XOR control cancelling over an even tile count, an atomic counter reading 0 while the work ran, a sampler sleeping past the run, bf16 parity vacuous at 8 mantissa bits, a guard dying on Windows CRLF, prefill printing 0.000. +- [ ] **13. Re-confirm quality on the exact artifact we serve.** GSM8K 96.0% exists โ€” on a different rung than the one now in play. + +## ๐Ÿ’ฐ Commercial and operational + +- [ ] **14. Token accounting must be exact** โ€” OpenRouter bills on it, and we just found our own bench miscounting. Audit `usage` end-to-end. +- [ ] **15. Deployment, monitoring, restart-on-crash.** Not built, not designed. +- [ ] **16. OpenRouter onboarding + pricing.** Requirements unknown to me โ€” needs Jish or a doc read, not a guess. + +--- + +### Standing rules for anything landing here +**Fast path is the DEFAULT** (D22-adjacent policy) โ€” flags are kill switches, never opt-ins. **Artifacts stay architecture-neutral; permute at load** (D22) โ€” one file serves every GPU generation. Every arm proves it engaged; a green result must prove work happened. Compare **F32 bit patterns before narrowing**, with a negative control that must fire โ€” a single-element 1-ULP perturbation, never an XOR that can cancel. From 2ecbc5fcb0804104cb2e2c53779cf705bc72eda2 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick Date: Thu, 20 Aug 2026 22:43:09 +0100 Subject: [PATCH 02/21] docs: the gate is 14,000 tok/s aggregate, and the scope is Arc complete Jish, direct: 1,400 aggregate is break-even and therefore pointless. 14,000 is $50/GPU-hr against a $4.92 card - 10x margin, a business. And 'OpenRouter-releasable' means Arc as we envisioned it, not a defect list. Rewritten accordingly: batching that amortises, the engine near the silicon, prefill as a product, the memory system, the multipliers we built and never switched on, and instruments we can trust. Key arithmetic: 14K is 84% of roofline at B=256 but only 42% at B=512, because memory traffic per step is flat as batch grows. 14K is a batching problem, not a kernel-speed problem - and batching is broken precisely because we are instruction-bound at 4% memory utilisation. Co-Authored-By: Claude Opus 5 (1M context) --- docs/engineering/OPENROUTER_READY.md | 97 +++++++++++++++++++++------- 1 file changed, 74 insertions(+), 23 deletions(-) diff --git a/docs/engineering/OPENROUTER_READY.md b/docs/engineering/OPENROUTER_READY.md index 965328541..1a1e04887 100644 --- a/docs/engineering/OPENROUTER_READY.md +++ b/docs/engineering/OPENROUTER_READY.md @@ -1,38 +1,89 @@ -**Merging this PR means Arc can serve and sell tokens on OpenRouter.** Nothing lands on `master` for this effort except through here. +# Arc, complete โ€” the single gate for selling tokens -Base `d7742670a` (the 15.0 โ†’ **34.48 tok/s** b=1 ladder, already on master). +**Merging this PR means Arc is the engine we set out to build, and it can be sold on OpenRouter.** +Not a defect list. The vision, shipped. Nothing reaches `master` for this effort except through here. + +Base `d7742670a` (the 15.0 โ†’ **34.48 tok/s** b=1 ladder, on master). + +--- + +## ๐ŸŽฏ THE GATE IS ONE NUMBER: **14,000 tok/s AGGREGATE ON ONE H200** + +| aggregate | tokens/GPU-hr | revenue @ $1/M | vs $4.92 card | verdict | +|---|---|---|---|---| +| ~50 (today, slow build) | 0.18 M | $0.18 | **27ร— underwater** | โ€” | +| 1,400 | 5.0 M | $5 | break-even | **pointless** | +| **14,000** | **50.4 M** | **$50** | **10ร— margin** | **the business** | + +**Why 14,000 is an engineering target and not a fantasy:** + +| batch | memory-bound roofline | 14,000 is | +|---|---|---| +| 256 | 16,600 tok/s | **84% of peak** โ€” nobody runs there | +| **512** | **~33,000 tok/s** | **42% of peak โ€” normal for a good engine** | + +Memory traffic per step is **flat** as batch grows โ€” 256 or 512 users read the same 74 GB once. +Every extra concurrent user is nearly free until compute-bound, somewhere around Bโ‰ˆ512โ€“1000. +**So 14K is a batching problem, not a kernel-speed problem.** + +๐Ÿ”ด **And batching is exactly what is broken.** Amortisation only pays when memory is the +constraint. **Our memory controller sits at 4%** โ€” we are instruction-bound, so adding users +adds work instead of sharing it. Measured: b=1 is 85ร— off its bound, **aggregate at B=8 is +227ร— off** โ€” we are *further* off at batch, which is the signature. + +**b=1 speed is a waypoint, not the goal.** Single-user is what one customer feels; aggregate +is what you sell. --- -## ๐Ÿ”ด Blocking โ€” cannot sell tokens at all +## What has to be true -- [ ] **1. Time-to-first-token is ~24 s.** ~12 ms per *input* token, ~100ร— off bound. A user on OpenRouter feels this immediately. **Prefill is a different disease from decode โ€” the GPU is already 99.5% busy, so graphs/launch-reduction buy it nothing.** Needs its own lane and its own kernel work. -- [ ] **2. Prefill crashes.** b=1 on `qtip2b` exits 139 (SIGSEGV) and the harness prints `pp = 0.000ยฑ0.000` into a results table instead of failing. A 0.000 is not a measurement. -- [x] **3. The fast build reaches `master`.** Landed as #175/#177/#185โ€“#189 โ†’ `d7742670a`. โš ๏ธ Still owed: re-measure **master itself** and confirm it reads ~34.48. -- [ ] **4. Aggregate throughput on the fast build.** Never measured at any batch size. Sweep B = 1/8/32/64/128, per-request **and** aggregate. Expect flat-then-steep: 256 experts top-6 means tokens barely collide until batch is large. **Finding the knee is the deliverable** โ€” the GPU-hour economics depend on it. -- [ ] **5. Reliability under sustained load.** 400 requests passed once. Nobody has run this for hours. +### 1 ยท Batching actually amortises โ† **the gate** +- [ ] Sweep **B = 1/8/32/64/128/256/512**, per-request **and** aggregate. Current sweep stops at 128 โ€” **too early; the answer lives past 256.** +- [ ] Find the knee. 256 experts top-6 means tokens barely collide until batch is large; the knee is where sharing starts and it is the single most important measurement in the program. +- [ ] Memory-controller utilisation at each point. **4% โ†’ 60%+ is the whole job.** +- [ ] KV bytes resident per user, and the OOM ceiling. -## ๐ŸŸก Product quality +### 2 ยท The engine runs near the silicon +- [ ] **Kernels at their own bounds.** Real profile on `qtip2b`: tail (sinkhorn + fast_sum + bmul + ucopy) **27.2%**, `fp8_gemv_warp` **16.0% โ†’ 13.0%** (1.31ร— landed, bit-exactness **not yet verified**), `qtip2b_gemv_tuned_kernel` **15.4%**. ๐Ÿšฉ A **4ร—4** Sinkhorn costing 9.3% of an H200 is structural. +- [ ] **4-stage async pipelining** โ€” measured 54.8% โ†’ **92.9%** of peak across depths 1โ†’4, saturating at 4. Issue rate first, then this. +- [ ] **TCFRAG** โ€” 26.31 โ†’ 3.56 inst/weight, zero shuffles, cosine 1.0, **B=8 costs what b=1 costs** (the capacity property). Attaches to `qtip2b_gemv_tuned_kernel`. **Permute at LOAD (D22)** โ€” artifacts stay architecture-neutral, one file for every GPU generation. +- [ ] **ArcGraph is INERT by default** โ€” candle sits on the legacy NULL stream, which CUDA forbids capturing, so decode runs fully eagerly. Capture never engages unless env vars are set. +- [ ] **GPU WHILE-node loop** โ€” syncs once per *generation* instead of once per token. Written, **zero callers**. Blocked by an unconditional `return Ok(None)` (`pipeline/normal.rs:2324`) and a U8-dtype bug failing all 28 call sites. **Neither competitor has this.** +- [ ] **`CudaSampler`** โ€” complete, tested bit-exact, **zero callers, no gate**. Just unwired. -- [ ] **6. Generation speed โ†’ 250 tok/s.** Real profile on `qtip2b`: tail (sinkhorn + fast_sum + bmul + ucopy) **27.2%**, `fp8_gemv_warp` **16.0%**, `qtip2b_gemv_tuned_kernel` **15.4%**. ๐Ÿšฉ A **4ร—4** Sinkhorn costing 9.3% of an H200 is structural, not tuning. -- [ ] **7. Long-context memory, measured.** At 1M ctx the answer is **~3 users or ~55** depending on whether the 128-token window + compress ratios are live. Unmeasured. This is the capacity claim. -- [ ] **8. Prefix caching is silently disabled** by our own default (TurboQuant is the default cache type). Most real traffic is repeated prefixes. -- [ ] **9. Switched-off finished work.** `CudaSampler` โ€” complete, tested bit-exact, **zero callers, no gate**. The **GPU WHILE-node loop** โ€” the one design that removes the per-token sync; neither competitor has it. Blocked by an unconditional `return Ok(None)` at `pipeline/normal.rs:2324` and a U8-dtype bug failing all 28 call sites. -- [ ] **10. 56ร— KV over-retention** โ€” root-caused previously; the fix reportedly never reached V4. +### 3 ยท Prefill is a product, not an afterthought +- [ ] **TTFT ~24 s โ‡’ usable.** ~12 ms per *input* token, ~100ร— off bound. **Prefill is the opposite disease from decode โ€” the GPU is already 99.5% busy, so graphs and launch reduction buy it nothing.** Its own lane, its own kernels. +- [ ] **Prefill crashes**: b=1 on `qtip2b` exits 139 (SIGSEGV) and prints `pp = 0.000ยฑ0.000` into a results table instead of failing. +- [ ] **Chunked prefill** โ€” inert until the MoE gather lands. -## ๐ŸŸข Trust and correctness +### 4 ยท The memory system Arc was designed around +- [ ] **Prefix caching is silently disabled** by our own default (TurboQuant is the default cache type). Most real traffic is repeated prefixes. +- [ ] **56ร— KV over-retention** โ€” root-caused; the fix reportedly never reached V4. +- [ ] **Long context measured**: at 1M ctx the answer is **~3 users or ~55** depending on whether the 128-token window + compress ratios are live. Unmeasured, and it is the capacity claim. +- [ ] **Ragged batching / ArcSched** โ€” admit mixed-length cohorts; the per-seq advance dies on ragged batches today. -- [ ] **11. The live V4 test fixture is all zeros** โ€” identity tests against it prove nothing. -- [ ] **12. Verification code is less reliable than production code.** Twelve instruments lied in one session: bench counting *requested* tokens (printed **1280 tok/s** for a zero-token run), `cargo clean -p` silently no-opping on git deps, `nsys` hiding graph-internal kernels by default, a stale `.o` satisfying a link, `replay()` returning a clone of the graded buffer, `append_graph` comparing a buffer against itself, an XOR control cancelling over an even tile count, an atomic counter reading 0 while the work ran, a sampler sleeping past the run, bf16 parity vacuous at 8 mantissa bits, a guard dying on Windows CRLF, prefill printing 0.000. -- [ ] **13. Re-confirm quality on the exact artifact we serve.** GSM8K 96.0% exists โ€” on a different rung than the one now in play. +### 5 ยท The multipliers we designed and haven't switched on +- [ ] **Speculative decoding (MTP)** โ€” measured 1.93 tok/step at b=1, 1.06 at b=128. A near-2ร— at low batch, off. +- [ ] **Expert parallelism** โ€” does not exist. The natural MoE scaling axis. +- [ ] **ArcTarget multi-arch** โ€” B200/B300 as first-class, one artifact, N load-time tables (D22). -## ๐Ÿ’ฐ Commercial and operational +### 6 ยท Trust โ€” because our instruments have been the problem +- [ ] **Twelve instruments lied in one session.** bench counted *requested* tokens (**1280 tok/s for a zero-token run**) ยท `cargo clean -p` silently no-ops on git deps ยท `nsys` hides graph-internal kernels by default ยท a stale `.o` satisfied a link ยท `replay()` returned a clone of the graded buffer ยท `append_graph` compared a buffer against itself ยท an XOR control cancelled over an even tile count ยท an atomic read 0 while the work ran ยท a sampler slept past the run ยท bf16 parity vacuous at 8 mantissa bits ยท a guard died on Windows CRLF ยท prefill printed 0.000. +- [ ] **The live V4 fixture is all zeros** โ€” identity tests against it prove nothing. +- [ ] **Quality re-confirmed on the exact artifact served.** GSM8K 96.0% exists โ€” on a different rung. +- [ ] **One measurement harness**, and nothing measures any other way. -- [ ] **14. Token accounting must be exact** โ€” OpenRouter bills on it, and we just found our own bench miscounting. Audit `usage` end-to-end. -- [ ] **15. Deployment, monitoring, restart-on-crash.** Not built, not designed. -- [ ] **16. OpenRouter onboarding + pricing.** Requirements unknown to me โ€” needs Jish or a doc read, not a guess. +### 7 ยท Commercial and operational +- [ ] **Token accounting exact** โ€” OpenRouter bills on it, and our own bench was miscounting. +- [ ] **Sustained-load soak** โ€” 400 requests passed once; nobody has run hours. +- [ ] **Deploy, monitor, restart-on-crash.** +- [ ] **OpenRouter onboarding + pricing** โ€” requirements unknown; needs a doc read, not a guess. --- ### Standing rules for anything landing here -**Fast path is the DEFAULT** (D22-adjacent policy) โ€” flags are kill switches, never opt-ins. **Artifacts stay architecture-neutral; permute at load** (D22) โ€” one file serves every GPU generation. Every arm proves it engaged; a green result must prove work happened. Compare **F32 bit patterns before narrowing**, with a negative control that must fire โ€” a single-element 1-ULP perturbation, never an XOR that can cancel. +**Fast path is the DEFAULT** โ€” flags are kill switches, never opt-ins. **Artifacts stay +architecture-neutral; permute at load (D22).** Every arm proves it engaged. Compare **F32 bit +patterns before narrowing**, with a control that must fire โ€” a single-element 1-ULP +perturbation, never an XOR that can cancel. **A scoping result is never a verdict (D21).** From 83c7c4a8556ed304b7c7bc77cfb85fbfe06b6e94 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Thu, 20 Aug 2026 23:14:10 +0100 Subject: [PATCH 03/21] =?UTF-8?q?test(ArcLab):=20batch=20sweep=20=E2=80=94?= =?UTF-8?q?=20per-request=20vs=20aggregate,=20with=20the=20memory-controll?= =?UTF-8?q?er?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Measures what a capacity claim actually rests on: fire B concurrent streaming requests and time the inter-token intervals of each. Two aggregate anchors that must be read together, because they disagree exactly when the batch never formed: * derived = per_request * B -- assumes all B ran at the median rate * measured = cohort tokens / cohort span -- what the box actually delivered Achieved concurrency is reported beside them, integrated over the span. At B=64 the scheduler only ever held ~53 streams, so `derived` overstates by 33%; without the concurrency column that gap is invisible and the higher number is the one that gets quoted. Samples `utilization.memory` DURING the decode, not after. That is the memory-controller busy percentage, and it is the diagnostic that separates "batching is amortising the weight reads" from "the GPU looks busy running tiny kernels". A sampler started after the run writes an empty file and looks green, so this one starts before the requests and stamps every sample. A row where any request returns zero tokens is printed as FAILED and never averaged into the table. Measured on landed master d7742670a, H200, V4-Flash qtip2 (see PR discussion): aggregate rises only 36.4 -> 87.3 tok/s from B=1 to B=512 (x2.40 for 512x the users) while the memory controller sits at 3% median through B=256. Co-Authored-By: Claude Opus 5 (1M context) --- arc-tools/batch_sweep.py | 256 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 256 insertions(+) create mode 100644 arc-tools/batch_sweep.py diff --git a/arc-tools/batch_sweep.py b/arc-tools/batch_sweep.py new file mode 100644 index 000000000..d27065b7b --- /dev/null +++ b/arc-tools/batch_sweep.py @@ -0,0 +1,256 @@ +"""Batch sweep on landed master: per-request vs aggregate decode throughput. + +For each batch size B, fire B CONCURRENT streaming requests and time the +inter-token intervals of each. Two independent aggregate anchors that must +agree, because the interesting number here is capacity and a capacity claim +built on one instrument is not a claim: + + A) derived = per_request_tok_s * B (what the brief asked for) + B) measured = tokens delivered inside the all-B-active window / that window + +(A) assumes every stream runs at the median rate; (B) counts actual deltas in +the steady-state window where all B streams are decoding. If they diverge the +batch never reached steady state and the row is not a throughput measurement. + +Per-request rate is a median of inter-token INTERVALS, so prefill is excluded +and one scheduler hiccup cannot set the headline. + +A row where any request returns zero tokens is a FAILED row, printed as such, +never averaged into the table -- the whole point of the bench gate. + +Exit: 0 sweep complete | 1 a row failed | 2 environment failure. +""" +import json +import statistics +import sys +import subprocess +import threading +import time +import urllib.request + +BASE = sys.argv[1] if len(sys.argv) > 1 else "http://127.0.0.1:1234" +NTOK = int(sys.argv[2]) if len(sys.argv) > 2 else 64 +BATCHES = [int(x) for x in (sys.argv[3].split(",") if len(sys.argv) > 3 else ["1", "8", "32", "64", "128"])] + +TOPICS = [ + "oceans", "mountains", "deserts", "forests", "rivers", "glaciers", "volcanoes", + "coral reefs", "savannas", "tundra", "caves", "islands", "wetlands", "canyons", + "prairies", "fjords", +] + + +def one_stream(idx, ntok, out): + """Stream one request; record a timestamp per delivered token.""" + topic = TOPICS[idx % len(TOPICS)] + prompt = ("Write %d short factual sentences about %s. Sentence set %d. " + "One sentence per line, no preamble." % (max(6, ntok // 6), topic, idx)) + body = json.dumps({ + "model": "default", + "messages": [{"role": "user", "content": prompt}], + "max_tokens": ntok, "temperature": 0.0, "stream": True, + }).encode() + req = urllib.request.Request(BASE + "/v1/chat/completions", data=body, + headers={"Content-Type": "application/json"}) + stamps, chars, err = [], 0, None + t0 = time.perf_counter() + try: + with urllib.request.urlopen(req, timeout=1800) as r: + for raw in r: + line = raw.decode("utf-8", "replace").strip() + if not line.startswith("data:"): + continue + payload = line[5:].strip() + if payload == "[DONE]": + break + try: + d = json.loads(payload) + except Exception: + continue + if "error" in d: + err = str(d["error"])[:120] + break + for c in d.get("choices") or []: + piece = ((c.get("delta") or {}).get("content") or "") + if piece: + stamps.append(time.perf_counter()) + chars += len(piece) + except Exception as e: + err = "%s: %s" % (type(e).__name__, str(e)[:100]) + out[idx] = {"t0": t0, "stamps": stamps, "chars": chars, "err": err} + + +SMI = ["nvidia-smi", "--query-gpu=utilization.gpu,utilization.memory,memory.used,memory.free", + "--format=csv,noheader,nounits"] + + +def sample_gpu(stop, sink): + """Poll the GPU WHILE the batch decodes. + + `utilization.memory` is the memory-controller busy percentage -- the + diagnostic that says whether batching is amortising the weight reads. A + sampler that starts after the run writes an empty file and looks green, so + this one is started before the requests and stamped per sample. + """ + while not stop.is_set(): + try: + out = subprocess.run(SMI, capture_output=True, text=True, timeout=5).stdout.strip() + for line in out.splitlines(): + g, m, used, free = [x.strip() for x in line.split(",")] + sink.append((time.perf_counter(), int(g), int(m), int(used), int(free))) + except Exception: + pass + time.sleep(0.25) + + +def run_batch(b, ntok): + out = {} + samples = [] + stop = threading.Event() + sampler = threading.Thread(target=sample_gpu, args=(stop, samples), daemon=True) + sampler.start() + threads = [threading.Thread(target=one_stream, args=(i, ntok, out)) for i in range(b)] + w0 = time.perf_counter() + for t in threads: + t.start() + for t in threads: + t.join() + w1 = time.perf_counter() + stop.set() + sampler.join(timeout=5) + return out, w1 - w0, samples + + +def grade(b, out, wall, samples): + rows = [out[i] for i in sorted(out)] + zero = [i for i, r in enumerate(rows) if len(r["stamps"]) == 0] + errs = [(i, r["err"]) for i, r in enumerate(rows) if r["err"]] + if len(rows) != b: + return {"B": b, "fail": "instrument: %d rows for B=%d" % (len(rows), b)} + if zero: + return {"B": b, "fail": "%d of %d requests returned ZERO tokens%s" + % (len(zero), b, (" (first err: %s)" % errs[0][1]) if errs else "")} + + # per-request: median over requests of that request's median gap + per_req_rates = [] + for r in rows: + s = r["stamps"] + if len(s) < 4: + return {"B": b, "fail": "a request produced only %d tokens" % len(s)} + gaps = [s[i] - s[i - 1] for i in range(1, len(s))] + per_req_rates.append(1.0 / statistics.median(gaps)) + per_req = statistics.median(per_req_rates) + total_tok_all = sum(len(r["stamps"]) for r in rows) + + # COHORT throughput: every token the cohort delivered, over the wall time + # the cohort spent delivering them. This is the capacity number an operator + # bills against, and unlike a strict all-active window it is defined even + # when admission is staggered -- which it is at large B, because the + # scheduler admits in waves. + first = min(r["stamps"][0] for r in rows) + last = max(r["stamps"][-1] for r in rows) + span = last - first + if span <= 0: + return {"B": b, "fail": "cohort delivered no tokens over time"} + measured = total_tok_all / span + + # Achieved concurrency: mean active streams across the span. If this is far + # below B the row is NOT a measurement of B -- the scheduler never ran B at + # once, and reporting it as B would overstate what the box did. + edges = [] + for r in rows: + edges.append((r["stamps"][0], 1)) + edges.append((r["stamps"][-1], -1)) + edges.sort() + active, prev_t, area = 0, first, 0.0 + for t, d in edges: + area += active * (t - prev_t) + active += d + prev_t = t + concurrency = area / span if span > 0 else 0.0 + + # strict all-active window, when one exists (secondary anchor) + lo = max(r["stamps"][0] for r in rows) + hi = min(r["stamps"][-1] for r in rows) + if hi > lo: + inwin = sum(sum(1 for t in r["stamps"] if lo < t <= hi) for r in rows) + steady = inwin / (hi - lo) + else: + steady = -1.0 + + derived = per_req * b + disagree = abs(measured - derived) / max(measured, derived) + # attribute GPU samples to the cohort span + win = [s for s in samples if first <= s[0] <= last] + if not win: + win = samples + memu = sorted(s[2] for s in win) + gpuu = sorted(s[1] for s in win) + used = sorted(s[3] for s in win) + free = sorted(s[4] for s in win) + n_s = len(win) + ttfts = sorted(r["stamps"][0] - r["t0"] for r in rows) + return { + "B": b, "fail": None, + "per_req": per_req, + "ms_per_tok": 1000.0 / per_req, + "agg_derived": derived, + "agg_measured": measured, + "disagree": disagree, + "total_tok": total_tok_all, + "concurrency": concurrency, + "steady": steady, + "window_s": span, + "wall_s": wall, + "ttft_med": ttfts[len(ttfts) // 2], + "ttft_max": ttfts[-1], + "mem_util_med": (memu[len(memu) // 2] if memu else -1), + "mem_util_max": (memu[-1] if memu else -1), + "gpu_util_med": (gpuu[len(gpuu) // 2] if gpuu else -1), + "vram_used_max": (used[-1] if used else -1), + "vram_free_min": (free[0] if free else -1), + "n_samples": n_s, + } + + +print("BASE=%s NTOK=%d batches=%s" % (BASE, NTOK, BATCHES), flush=True) +results, failed = [], False +for b in BATCHES: + print("--- B=%d : firing %d concurrent streams ---" % (b, b), flush=True) + out, wall, samples = run_batch(b, NTOK) + r = grade(b, out, wall, samples) + results.append(r) + if r["fail"]: + failed = True + print(" B=%-4d FAILED ROW: %s" % (b, r["fail"]), flush=True) + else: + print(" B=%-4d per-req %7.2f | agg %8.2f | conc %6.1f/%d | MEM-CTRL %d%% (max %d%%) | " + "GPU %d%% | VRAMfree %d MiB | tok %d | samples %d" + % (b, r["per_req"], r["agg_measured"], r["concurrency"], b, + r["mem_util_med"], r["mem_util_max"], r["gpu_util_med"], + r["vram_free_min"], r["total_tok"], r["n_samples"]), flush=True) + time.sleep(3) + +print("") +print("%-6s %-13s %-11s %-13s %-13s %-8s %-9s %-9s %-10s" % + ("B", "per-req t/s", "ms/tok", "agg derived", "agg measured", "spread", + "MEMCTRL%", "conc", "VRAMfree")) +for r in results: + if r["fail"]: + print("%-6d %s" % (r["B"], "FAILED: " + r["fail"])) + else: + print("%-6d %-13.2f %-11.2f %-13.2f %-13.2f %-8s %-9s %-9s %-10d" + % (r["B"], r["per_req"], r["ms_per_tok"], r["agg_derived"], r["agg_measured"], + "%.1f%%" % (r["disagree"] * 100), + "%d/%d" % (r["mem_util_med"], r["mem_util_max"]), + "%.0f" % r["concurrency"], r["vram_free_min"])) + +ok = [r for r in results if not r["fail"]] +if len(ok) >= 2: + base = ok[0]["agg_measured"] + print("") + print("aggregate scaling vs B=%d:" % ok[0]["B"]) + for r in ok: + print(" B=%-5d %8.2f tok/s x%.2f (per-request retained: %.0f%%)" + % (r["B"], r["agg_measured"], r["agg_measured"] / base, + 100.0 * r["per_req"] / ok[0]["per_req"])) +sys.exit(1 if failed else 0) From 7de5eda261bfdd8e7b7a0dcf3fd3898c8493450c Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 00:34:33 +0100 Subject: [PATCH 04/21] fix(ArcQuant): the MoE gate never selected the grouped GEMM in serving MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `Qtip2bLayer::gather_forward` gated the trellis grouped GEMM on tile occupancy: switch only once each woken expert draws `GROUPED_TILE_M` (=16) pairs. At V4's top-6-of-256 routing that needs `6n/256 >= 16`, i.e. **n >= 683** โ€” a token count decode never reaches. The clause therefore pinned every decode step, and every batch up to 682, on the per-pair gather-GEMV arm. A profile of the shipped default shows the consequence directly: `qtip_gather_gemv_warp_kernel` at 57.6% of GPU time with 25x the launches of the grouped kernel. Two independent measurements retire it. Decode on an H200, clean rows only (>=97% achieved concurrency, <7% derived-vs-measured spread): grouped wins at B=48 1.10x, 64 1.14x, 80 1.22x, 96 1.29x, 128 1.66x, 512 1.51x. Prefill over prompt length: +16% at 16 tokens, +27% at 32, +50% at 64, +78% at 128, +181% at 512. The gate's own justifying A/B -- "1.00x at N=128 and N=512" -- does not reproduce. **Tile fill was the wrong quantity.** At n=16 the tile is 7.5% full and the grouped kernel still wins by 16%, so a predicate demanding a full tile cannot be the gate. The win is that the grouped kernel stages a woken expert's bytes once per m-tile instead of once per (token, expert) pair -- the GEMV arm's reads are exactly linear in n_tokens, 12.00x redundant at n=512, while the grouped arm holds flat at one pass over the expert set -- and that it runs on tensor cores. Occupancy was never the mechanism being traded against. `DECODE_REGIME_MAX_TOKENS` survives as a floor: below it there is nothing to group (b=1 measures 0.47x) and it is the RUN-161 capture-safety rule. 9..=15 is unmeasured, bounded by that floor and by the +16% at 16, and rides with the grouped arm. Fast path by default; `ARC_NO_QTIP_GROUPED_MOE` remains the kill switch and `ARC_QTIP_ONDEVICE_MOE_MAX_TOKENS` still pins the GEMV arm for A/B. The decision moves into `prefer_gather_gemv` so it can be tested: the dispatch site is inside `#[cfg(feature = "cuda")]`, and a non-CUDA `cargo test` reports a full green over code it never compiled -- which it did to me once during this change. The tests pin every measured batch size to the grouped arm and carry a control asserting the retired clause would have excluded all of them. Co-Authored-By: Claude Opus 5 (1M context) --- mistralrs-quant/src/qtip/bitshift.rs | 66 +++++++----- mistralrs-quant/src/qtip/gather_policy.rs | 126 +++++++++++++++++++++- 2 files changed, 165 insertions(+), 27 deletions(-) diff --git a/mistralrs-quant/src/qtip/bitshift.rs b/mistralrs-quant/src/qtip/bitshift.rs index b914c5e6e..f6bca9575 100644 --- a/mistralrs-quant/src/qtip/bitshift.rs +++ b/mistralrs-quant/src/qtip/bitshift.rs @@ -1902,37 +1902,53 @@ impl QuantMethod for Qtip2bLayer { // On-device only, propagate errors (the CPU fallback would // D2H-sync under graph capture and read garbage indices). // - // Above this boundary the rung goes to the trellis grouped GEMM - // below โ€” tokens sorted by expert on-device, each expert's bytes - // read once per m-tile โ€” which is the amortizing kernel the - // fleet math wants. But it only amortizes when its m-tiles are - // FULL: it stages a woken expert's bytes once per - // `GROUPED_TILE_M` pairs, so at V4's top-6-of-256 routing a - // 9-token step wakes ~54 experts for 54 pairs and every 16-row - // tile carries one useful row. The gate is therefore the tile - // fill, not a token count โ€” same predicate, same reasons, as - // the LUT rung; see `super::gather_policy` ยง4 for the - // derivation and for the 5.7x end-to-end regression the raw - // token gate caused there. + // ๐Ÿ”ด TILE FILL WAS THE WRONG QUANTITY โ€” RETRACTED ON MEASUREMENT. + // + // This gate used to also require `grouped_gemm_tiles_amortize`: + // switch to the grouped GEMM only once each woken expert draws + // `GROUPED_TILE_M` (=16) pairs. At V4's top-6-of-256 that needs + // `6n/256 >= 16`, i.e. **n >= 683** โ€” a token count decode never + // reaches. The clause therefore pinned every decode step, and + // every batch up to 682, on the gather-GEMV arm. + // + // Two independent measurements retired it: + // * decode, H200, V4-Flash qtip2b, clean rows only (>=97% + // achieved concurrency, <7% derived-vs-measured spread): + // grouped beats GEMV at B=48 (1.10x), 64 (1.14x), 80 + // (1.22x), 96 (1.29x), 128 (1.66x), 512 (1.51x). + // * prefill, same gate over prompt lengths: +16% at 16 + // tokens, +27% at 32, +50% at 64, +78% at 128, +181% at 512. + // The gate's own justifying A/B ("1.00x at N=128 and N=512") + // does not reproduce. + // + // **A 7.5%-full tile still wins by 16%.** So fill is not the + // deciding quantity. The win is that the grouped kernel stages a + // woken expert's bytes ONCE PER M-TILE instead of once per + // (token, expert) pair โ€” the GEMV arm's reads are exactly linear + // in `n_tokens`, measured at 12.00x redundancy at n=512, while + // the grouped arm holds flat at one pass over the expert set โ€” + // and that it runs on tensor cores rather than a scalar GEMV. + // Occupancy was never the mechanism being traded against. + // + // What survives is `DECODE_REGIME_MAX_TOKENS` as a FLOOR: at + // n<=8 there genuinely is not enough work to group (b=1 measures + // 0.47x), and it is also the RUN-161 capture-safety rule. + // 9..=15 is UNMEASURED โ€” bounded below by the n=8 floor and + // above by the +16% at n=16, so it rides with the grouped arm. // // `DECODE_REGIME_MAX_TOKENS` survives as a FLOOR (the RUN-161 // capture-safety rule), and an explicit // `ARC_QTIP_ONDEVICE_MOE_MAX_TOKENS` still pins the GEMV arm // outright so a harness can A/B the two kernels. - let ondevice_preferred = match super::gather_policy::ondevice_max_tokens_override() - { - Some(cap) => n_tokens <= cap, - None => { - n_tokens <= super::DECODE_REGIME_MAX_TOKENS - || !super::gather_policy::grouped_gemm_tiles_amortize( - n_tokens, - n_experts_per_tok, - self.num_experts_count(), - ) - } - }; + let ondevice_preferred = super::gather_policy::prefer_gather_gemv( + n_tokens, + super::gather_policy::ondevice_max_tokens_override(), + ); // Read by VALUE: `ARC_NO_QTIP_ONDEVICE_MOE=0` leaves the - // on-device path enabled (the default). + // on-device path enabled (the default). #197 was cut before + // #212 and still had the presence-check `std::env::var(..) + // .is_ok()` here; keeping #212's parser is deliberate, so this + // retraction does not re-break the flag polarity. let ondevice_disabled = crate::env_flag_is_set("ARC_NO_QTIP_ONDEVICE_MOE"); if !ondevice_disabled && ondevice_preferred diff --git a/mistralrs-quant/src/qtip/gather_policy.rs b/mistralrs-quant/src/qtip/gather_policy.rs index 7cf8c60a1..27ef73d1d 100644 --- a/mistralrs-quant/src/qtip/gather_policy.rs +++ b/mistralrs-quant/src/qtip/gather_policy.rs @@ -139,8 +139,44 @@ //! 1024 6144 256.0 24.00 150.0% //! ``` //! -//! and the in-tree end-to-end A/B brackets exactly that curve: **1.00ร— at -//! N=128 and N=512, 2.41ร— at N=1024**. +//! and the in-tree end-to-end A/B was said to bracket exactly that curve: +//! **1.00ร— at N=128 and N=512, 2.41ร— at N=1024**. +//! +//! # ๐Ÿ”ด 4a. THE TILE-FILL GATE IS RETRACTED โ€” IT MEASURED THE WRONG QUANTITY +//! +//! **That A/B does not reproduce, and the model above is not what decides the +//! winner.** Two independent measurements on landed master (H200, V4-Flash +//! qtip2b) put the grouped GEMM ahead at every size either could test: +//! +//! ```text +//! decode, clean rows only (>=97% achieved concurrency, <7% spread): +//! B=48 1.10x B=64 1.14x B=80 1.22x B=96 1.29x B=128 1.66x B=512 1.51x +//! prefill, over prompt length: +//! n=16 +16% n=32 +27% n=64 +50% n=128 +78% n=512 +181% +//! ``` +//! +//! At `n = 16` the tile is **7.5% full** and the grouped kernel still wins by +//! 16%. A predicate that demands 100% fill therefore cannot be the right gate: +//! it required `6n/256 >= 16`, i.e. **n >= 683**, which decode never reaches, so +//! it pinned every decode step and every batch up to 682 on the GEMV arm. A +//! profile of the shipped default confirmed the consequence directly โ€” +//! `qtip_gather_gemv_warp_kernel` at **57.6% of GPU time with 25ร— the launches** +//! of the grouped kernel. +//! +//! **Why fill was the wrong quantity.** The ratio above counts only *weight* +//! bytes and assumes the two kernels are otherwise alike. They are not. The +//! grouped kernel stages a woken expert's bytes once per m-tile rather than once +//! per (token, expert) pair โ€” the GEMV arm's reads are exactly linear in +//! `n_tokens`, measured at **12.00ร— redundancy at n=512**, while the grouped arm +//! holds flat at one pass over the expert set โ€” and it runs the tensor-core tile +//! loop instead of a scalar per-pair GEMV. A partly-empty tile still buys both. +//! Occupancy was never the mechanism being traded against. +//! +//! What survives is [`super::DECODE_REGIME_MAX_TOKENS`] as a floor: at `n <= 8` +//! there is genuinely nothing to group (b=1 measures **0.47ร—**), and it is the +//! RUN-161 capture-safety rule. `9..=15` is **unmeasured**, bounded below by +//! that floor and above by the +16% at 16, and rides with the grouped arm. +//! [`prefer_gather_gemv`] is the shipped decision and is unit-tested. //! //! Shipping the gate as a raw `n_tokens > DECODE_REGIME_MAX_TOKENS` switched //! the grouped kernel on at n=9 โ€” onto a tile that is ~9% full, roughly two @@ -305,6 +341,40 @@ pub fn grouped_gemm_tile_fill(n_tokens: usize, top_k: usize, num_experts: usize) /// `GROUPED_TILE_M` changes, and the last one that was not re-derived is the /// regression this replaces. #[cfg_attr(not(feature = "cuda"), allow(dead_code))] +/// Does the fused gather-GEMV arm win at this token count? +/// +/// The pure half of the MoE dispatch in `bitshift.rs`, so the decision that +/// actually ships can be tested without a GPU. It has to live here and be +/// exercised by a test, because the dispatch itself sits inside a +/// `#[cfg(feature = "cuda")]` block that a normal `cargo test` cannot compile โ€” +/// a non-CUDA lane will report a full green over code it never saw. +/// +/// ๐Ÿ”ด This used to also consult [`grouped_gemm_tiles_amortize`]: switch to the +/// grouped GEMM only once each woken expert fills an m-tile. At V4's +/// top-6-of-256 that needs `6n/256 >= GROUPED_TILE_M`, i.e. **n >= 683**, a +/// token count decode never reaches โ€” so the clause pinned every decode step, +/// and every batch up to 682, on the GEMV arm. +/// +/// Two independent measurements retired it: decode at B=48..512 on an H200, and +/// prefill over 16..512-token prompts. Grouped wins everywhere at and above 16, +/// including **+16% on a tile that is only 7.5% full**. +/// +/// **Tile fill was therefore the wrong quantity.** The grouped kernel wins +/// because it stages each woken expert's bytes once per m-tile instead of once +/// per (token, expert) pair โ€” the GEMV arm's reads are exactly linear in +/// `n_tokens` โ€” and because it runs on tensor cores rather than a scalar GEMV. +/// Occupancy was never the mechanism being traded against. +/// +/// [`super::DECODE_REGIME_MAX_TOKENS`] survives as a FLOOR: below it there is +/// genuinely nothing to group (b=1 measures 0.47x), and it is the RUN-161 +/// capture-safety rule besides. +pub(crate) fn prefer_gather_gemv(n_tokens: usize, override_cap: Option) -> bool { + match override_cap { + Some(cap) => n_tokens <= cap, + None => n_tokens <= super::DECODE_REGIME_MAX_TOKENS, + } +} + pub(crate) fn grouped_gemm_tiles_amortize( n_tokens: usize, top_k: usize, @@ -740,3 +810,55 @@ mod tests { } } } + +#[cfg(test)] +mod crossover_tests { + use super::*; + + /// The shipped dispatch: the GEMV arm only at or below the capture-safety + /// floor, and the grouped GEMM reachable at every batch size we measured. + #[test] + fn the_grouped_gemm_is_reachable_at_decode_batch_sizes() { + for n in [1, 2, 4, 8] { + assert!( + prefer_gather_gemv(n, None), + "n={n} is below the floor and must stay on the GEMV arm" + ); + } + // Exactly the points measured on hardware, all of which the grouped + // kernel won: decode B=48 1.10x, 64 1.14x, 80 1.22x, 96 1.29x, + // 128 1.66x, 512 1.51x (clean rows: >=97% achieved concurrency, <7% + // derived-vs-measured spread), and prefill +16% as low as n=16. + for n in [16, 32, 48, 64, 80, 96, 128, 512] { + assert!( + !prefer_gather_gemv(n, None), + "n={n} must reach the grouped GEMM; it was measured faster there" + ); + } + } + + /// NON-VACUITY CONTROL, and the whole reason the clause was removed: the + /// retired tile-fill predicate would have sent every one of those token + /// counts to the GEMV arm. If this stops holding, the control is broken and + /// the test above proves nothing about the change. + #[test] + fn the_retired_tile_fill_clause_would_have_excluded_all_of_them() { + for n in [16, 32, 48, 64, 80, 96, 128, 512] { + assert!( + !grouped_gemm_tiles_amortize(n, 6, 256), + "control broken: the retired clause was supposed to reject n={n}" + ); + } + // It only relented at 683 -- a token count decode never reaches, which + // is precisely why the grouped kernel was unreachable in serving. + assert_eq!(grouped_gemm_min_tokens(6, 256), Some(683)); + } + + /// The A/B override still pins the GEMV arm outright, so both kernels stay + /// comparable from one binary. + #[test] + fn the_override_still_pins_the_gemv_arm() { + assert!(prefer_gather_gemv(512, Some(1024))); + assert!(!prefer_gather_gemv(512, Some(256))); + } +} From d665956ab8748ce10f786b936730370fdc7a12d3 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 00:49:17 +0100 Subject: [PATCH 05/21] =?UTF-8?q?fix(ArcQuant):=20same=20683-token=20gate?= =?UTF-8?q?=20on=20the=20LUT=20rung=20=E2=80=94=20and=20this=20is=20the=20?= =?UTF-8?q?rung=20measured?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous commit removed the tile-fill conjunct from the bitshift rung (`Qtip2bLayer`, --isq qtip2b). The LUT rung (`QtipLayer`, --isq qtip2) carries the identical clause at mod.rs:3652, and the published V4-Flash artifact this was measured against is **qtip2**, not qtip2b: the profile's kernels are `qtip_lut_grouped_gemm_kernel` and `qtip_gather_gemv_warp_kernel`. So every decode number in the sweep โ€” GEMV 78.46 vs grouped 89.37 at B=64, 72.63 vs 120.36 at B=128 โ€” came from this rung, and fixing only the bitshift rung changed nothing on the box: the default still read 81.17 at B=128 after that build. Fixing the rung the measurement actually exercised. Co-Authored-By: Claude Opus 5 (1M context) --- mistralrs-quant/src/qtip/mod.rs | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/mistralrs-quant/src/qtip/mod.rs b/mistralrs-quant/src/qtip/mod.rs index af6c8f295..dcf87896e 100644 --- a/mistralrs-quant/src/qtip/mod.rs +++ b/mistralrs-quant/src/qtip/mod.rs @@ -3689,14 +3689,19 @@ impl QuantMethod for QtipLayer { // RUN-161 floor, not a performance choice (see // `gather_policy`). An explicit `ARC_QTIP_ONDEVICE_MOE_MAX_TOKENS` // override also still wins, so a harness can pin the GEMV arm. + // ๐Ÿ”ด The `grouped_gemm_tiles_amortize` conjunct is GONE โ€” see + // `gather_policy` ยง4a. It demanded `6n/256 >= 16`, i.e. + // n >= 683, so on this rung too the grouped GEMM was + // unreachable in serving and every decode step ran the per-pair + // GEMV. Measured on THIS rung (qtip2 LUT, H200, V4-Flash), + // clean rows only: forcing the grouped path is 1.10x at B=48, + // 1.14x at 64, 1.22x at 80, 1.29x at 96, 1.66x at 128 and + // 1.51x at 512. Tile fill is not the deciding quantity; staging + // each woken expert once per m-tile instead of once per + // (token, slot) pair is. let grouped_preferred = grouped_available && n_tokens > DECODE_REGIME_MAX_TOKENS - && gather_policy::ondevice_max_tokens_override().is_none() - && gather_policy::grouped_gemm_tiles_amortize( - n_tokens, - n_experts_per_tok, - num_experts, - ); + && gather_policy::ondevice_max_tokens_override().is_none(); let use_ondevice = !grouped_preferred && match gather_policy::ondevice_max_tokens_override() { From 834010a1b2f66069e7c6ab3c55955c25cec120e0 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 17:46:16 +0100 Subject: [PATCH 06/21] =?UTF-8?q?feat(qtip):=20UQFF=20geometry=20discrimin?= =?UTF-8?q?ator=20=E2=80=94=20K=20is=20an=20explicit=20wire=20field=20(#17?= =?UTF-8?q?0)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(qtip): UQFF geometry discriminator โ€” K8/V4/L12 artifacts declare themselves Stage 2 of the K=8/V=4/L=12 rung: an artifact can now say which trellis geometry it was baked at, the loader reads it, and a build that cannot decode it refuses instead of guessing. WHY THIS NEEDS A DISCRIMINATOR AT ALL Both geometries are 2 bits per weight (`bpw = K/V`), so a K=8/V=4 row and a K=4/V=2 row of the same `in_features` occupy the SAME number of packed bytes. A mislabelled artifact therefore indexes in bounds at every symbol and returns plausible garbage rather than faulting. The table is the only tensor that differs โ€” `[4096, 4]` BF16 vs `[65536, 2]` F32 โ€” which is why `validate_shapes` checks the table's dtype AND its length, and why both checks are load-bearing rather than belt-and-braces. WIRE FORMAT A trailing section `[tag=3, K, L, V]`, written ONLY for a non-default geometry and written BEFORE the codebook section. Both of those are deliberate: * Default writes nothing, so every artifact Arc has already produced stays byte-identical and no checksum moves. Pinned as an exact suffix, not as a vague "unchanged": serializing the same layer with the tag flipped appends exactly [3, 8, 12, 4] and nothing else. * BEFORE the codebook, because a build that predates this field parses the trailing region by handing its first byte to `QtipCodebook::from_wire`, which refuses every tag it does not know. Tag 3 first means an old Arc fails closed. Tag 3 last would let it consume the codebook section, stop, and decode K=8/V=4 symbols as K=4/V=2 without faulting. The trailing region is now a section loop; a repeated tag is refused rather than letting the last one win. `QtipGeometry` has no per-site default: adding the field broke all 11 construction sites at compile time and each one now states its geometry, the same discipline `search` and `search_detail` already follow. `from_stacked_parts` takes it explicitly and validates it. `stack_experts` and the 3-D quantize path refuse mixed-geometry stacks โ€” only one table survives a stack, so at most one expert could decode. The computed `sum2` codebook is refused at V=4 on both the write and the read side: it produces a PAIR of values per state and has no V=4 form. Guards were mutation-tested (F1-F9). Three passed on broken code and are now covered: `from_stacked_parts` validated nothing, the table-dtype half of the discriminator was dead (the element-count check masked it), and `stack_experts`' mixed-geometry refusal was untested. `assert_tensor_bits_eq` grew a BF16 arm โ€” it panicked on the new table rather than comparing it. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SpVNMpb13HkUXqSqbN1o9H * fix(qtip): refuse an unsupported geometry at every decode entry point Stage 2 made a K=8/V=4/L=12 artifact loadable. On its own that is a half-applied change: every decoder behind `forward`, `gather_forward`, `dequantize_weights`, `dequantize_expert` and `qtip_packed` unpacks nibbles and indexes a [2^16, 2] table, and handed a K=8 row none of them would fault โ€” at 2 bits per weight the packed byte count is identical and every index lands in bounds. They would serve garbage. So each entry point now states the geometries it implements. `qtip_packed` returns None rather than a `QtipPackedView`, which carries no geometry field and would therefore be read as K=4/V=2 nibbles by any consumer. The refusal names the entry point, and the test asserts that name. An entry point that merely inherits a callee's guard is not guarded: mutation G1 removed `forward`'s own check and every test stayed green, because the error still arrived from `dequantize_weights` further down. All six mutations (G1-G6) are red now. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SpVNMpb13HkUXqSqbN1o9H * fix(qtip): UQFF validates the PADDED row stride, not the data length Follows the padding decision one branch down. `QtipGeometry::packed_len` now means the allocated stride (what a tensor is shaped at and what this format validates) and `data_bytes` is the bit-rate-governed part, both delegating to `trellis_v4l12::Rung` so there is still one implementation of each. Three tests changed premise rather than value and were rewritten, not patched: the bit-rate ratio claim moved to `data_bytes` where it is exact, and the stride got its own bounds (never smaller than the data, never more than 4 bytes larger). A K=9 row of `in_features=36` is 11 data bytes in a 12-byte stride. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SpVNMpb13HkUXqSqbN1o9H * fix(qtip): the CUDA bake fast path lost the geometry wire tag `quantize_with_options_cuda` builds a QtipLayer directly and was the one initializer the geometry discriminator missed. It is behind `#[cfg(feature = "cuda")]`, so the default Check jobs compiled fine and only the CUDA lane went red (E0063 at qtip/mod.rs:2318). Set the same K4V2L16 tag its CPU sibling sets โ€” the CUDA bake kernels run the identical K=4/V=2/L=16 trellis, so a GPU-baked artifact must carry the identical tag or it would deserialize as a different geometry. Co-Authored-By: Claude Opus 5 (1M context) --------- Co-authored-by: Claude Opus 5 (1M context) --- .../examples/qtip_lut_grouped_bench.rs | 6 +- mistralrs-quant/src/lib.rs | 4 +- mistralrs-quant/src/qtip/mod.rs | 1116 ++++++++++++++++- 3 files changed, 1108 insertions(+), 18 deletions(-) diff --git a/mistralrs-quant/examples/qtip_lut_grouped_bench.rs b/mistralrs-quant/examples/qtip_lut_grouped_bench.rs index 98f710cb7..59f430db9 100644 --- a/mistralrs-quant/examples/qtip_lut_grouped_bench.rs +++ b/mistralrs-quant/examples/qtip_lut_grouped_bench.rs @@ -32,7 +32,8 @@ use candle_core::{DType, Device, Result, Tensor}; use mistralrs_quant::{ - QtipCodebook, QtipLayer, QtipSearchDetail, QtipSearchStamp, QuantMethod, QTIP_GROUPED_TILE_K, + QtipCodebook, QtipGeometry, QtipLayer, QtipSearchDetail, QtipSearchStamp, QuantMethod, + QTIP_GROUPED_TILE_K, }; use std::time::Instant; @@ -138,6 +139,9 @@ fn build_layer( QtipSearchStamp::Unstamped, QtipSearchDetail::Unknown, QtipCodebook::Gaussian, + // This bench synthesises a `[65536, 2]` F32 table above; that is the + // K=4/V=2/L=16 geometry by construction. + QtipGeometry::K4V2L16, ) } diff --git a/mistralrs-quant/src/lib.rs b/mistralrs-quant/src/lib.rs index 2ec9ae67f..b489f148a 100644 --- a/mistralrs-quant/src/lib.rs +++ b/mistralrs-quant/src/lib.rs @@ -120,8 +120,8 @@ pub use qtip::tune::{ pub use qtip::{ bake_cache, gpu_quantize_cpu_fallback_count, grouped_launch_counts, grouped_variant, hessian_row_weights, qtip_expected_distinct_experts, qtip_expected_pairs_per_distinct_expert, - qtip_grouped_gemm_tile_fill, set_grouped_variant, viterbi_quantize_row, - BakeCacheError, BakeKey, ExpertBpwTable, Qtip2bLayer, QtipBakeConfig, QtipCodebook, QtipLayer, + qtip_grouped_gemm_tile_fill, set_grouped_variant, viterbi_quantize_row, BakeCacheError, + BakeKey, ExpertBpwTable, Qtip2bLayer, QtipBakeConfig, QtipCodebook, QtipGeometry, QtipLayer, QtipMode, QtipPackedView, QtipRotation, QtipSearchDetail, QtipSearchStamp, TrellisBpw, TrellisSearch, QTIP2B_MCG_MULT, QTIP_GATHER_GEMV_MAX_PAIRS, QTIP_GROUPED_TILE_K, QTIP_GROUPED_TILE_M, QTIP_GROUPED_TILE_N, QTIP_GROUPED_VARIANT_BASELINE, diff --git a/mistralrs-quant/src/qtip/mod.rs b/mistralrs-quant/src/qtip/mod.rs index dcf87896e..fdbd015a2 100644 --- a/mistralrs-quant/src/qtip/mod.rs +++ b/mistralrs-quant/src/qtip/mod.rs @@ -686,6 +686,10 @@ impl QtipCodebook { /// one. Unknown tags fail closed (DOCTRINE: refuse, never launder). fn from_wire(tag: u8, read_mult: impl FnOnce() -> Result) -> Result { match tag { + GEOMETRY_WIRE_TAG => candle_core::bail!( + "QtipLayer: internal error โ€” the geometry section reached the codebook parser. \ + The trailing-section loop must consume tag {GEOMETRY_WIRE_TAG} itself." + ), 2 => { let mult = read_mult()?; if mult == 0 || mult % 2 == 0 { @@ -705,6 +709,276 @@ impl QtipCodebook { } } +/// Trailing-section tag for the trellis geometry. +/// +/// **It is 3, and it is written BEFORE the codebook section, and both of those +/// are load-bearing.** A build that predates this field parses the trailing +/// region by reading one byte and handing it to [`QtipCodebook::from_wire`], +/// which refuses every tag it does not know. So an old Arc handed a +/// non-default-geometry artifact sees tag 3 first and fails closed with +/// "written by a newer Arc โ€ฆ refusing rather than decoding its symbols against +/// the wrong codebook" โ€” which is exactly the right outcome, because it cannot +/// decode this geometry either. +/// +/// Put the section *after* the codebook and that property is lost: an old +/// reader would consume the codebook section, stop, and decode a K=8/V=4 symbol +/// stream as K=4/V=2. It would not even fault โ€” the packed byte count is +/// identical at 2 bpw, so every index stays in bounds and the output is +/// plausible garbage. [`tests::geometry_section_precedes_the_codebook_section`] +/// pins the order. +const GEOMETRY_WIRE_TAG: u8 = 3; + +/// Which trellis geometry a [`QtipLayer`]'s symbols were baked at. +/// +/// **This is a format discriminator, not a tuning knob.** Symbols are a +/// bit-stream interpreted by a specific `(K, L, V)` and a specific reproduction +/// table; read at the wrong geometry they decode to different weights, silently. +/// +/// `K` is carried as data, not baked into the variant, because that is what the +/// wire already says โ€” the section is `[tag, K, L, V]`. Within the V=4/L=12 +/// family the table is K-independent, so a new K is a new value here and not a +/// new variant. See [`crate::trellis_v4l12`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum QtipGeometry { + /// K=4 / V=2 / L=16 with a `[65536, 2]` F32 table โ€” every QTIP artifact + /// Arc has ever written. The default, and it serializes as **nothing**. + #[default] + K4V2L16, + /// The V=4 / L=12 family: a `[4096, 4]` **BF16** table (32,768 B โ€” fits + /// static shared memory), with `k` bits per symbol. + /// + /// `k = 8` is the byte-aligned control at 2.00 bpw; `k = 9` is 2.25 bpw and + /// is the quality winner. Constructed only through + /// [`QtipGeometry::trellis_v4l12`], which refuses a `k` with no decoder. + TrellisV4L12 { + /// The decode rung, which owns the K-dependent arithmetic. Stored + /// rather than a bare `k` so the packed-size formula has exactly one + /// implementation โ€” see [`QtipGeometry::packed_len`]. + rung: trellis_v4l12::Rung, + }, +} + +impl QtipGeometry { + /// The V=4/L=12 family at symbol width `k`, refusing a `k` this build has + /// no decoder for. + /// + /// Refuses rather than storing it: `k` arrives from an artifact, and a + /// geometry we cannot decode must not be constructible as a value that + /// later code will try to use. + pub fn trellis_v4l12(k: u32) -> Result { + let rung = trellis_v4l12::Rung::new(k).map_err(candle_core::Error::Msg)?; + Ok(QtipGeometry::TrellisV4L12 { rung }) + } + + /// Bits per symbol. + pub fn k(self) -> u32 { + match self { + QtipGeometry::K4V2L16 => K, + QtipGeometry::TrellisV4L12 { rung } => rung.k(), + } + } + + /// Trellis state width in bits. + pub fn l(self) -> u32 { + match self { + QtipGeometry::K4V2L16 => L, + QtipGeometry::TrellisV4L12 { .. } => trellis_v4l12::L, + } + } + + /// Reproduction values per symbol. + pub fn v(self) -> u32 { + match self { + QtipGeometry::K4V2L16 => V, + QtipGeometry::TrellisV4L12 { .. } => trellis_v4l12::V, + } + } + + /// Bits per weight ร— 100 (`100ยทK/V`), as an integer so it compares exactly. + /// + /// 200 for the shipped rung and for K=8/V=4; **225 for K=9/V=4**; 250 for + /// K=10/V=4. Not every geometry in this enum is the same bit rate any more, + /// which is a change from when the family was K=8-only. + pub fn bpw_x100(self) -> u32 { + self.k() * 100 / self.v() + } + + /// Total values in the reproduction table: `2^L ร— V`. + pub fn lut_values(self) -> usize { + (1usize << self.l()) * self.v() as usize + } + + /// Element type of the reproduction table. + /// + /// Part of the discriminator, not an implementation detail: the V=4 table + /// is BF16 precisely so it lands at 32,768 B, and an F32 table of the same + /// shape is a different artifact. + pub fn lut_dtype(self) -> DType { + match self { + QtipGeometry::K4V2L16 => DType::F32, + QtipGeometry::TrellisV4L12 { .. } => DType::BF16, + } + } + + /// Trellis symbols in a row of `in_features` weights. Depends only on V. + pub fn num_symbols(self, in_features: usize) -> usize { + in_features / self.v() as usize + } + + /// Bytes of a row that actually hold symbols, governed by the bit rate. + /// + /// **Not the row stride.** [`QtipGeometry::packed_len`] is what a tensor is + /// allocated at and what this format validates; it additionally carries + /// tail padding and a round-up to 4 for the V=4/L=12 family. Conflating + /// them is how a bits-per-weight claim ends up four bytes wrong. + pub fn data_bytes(self, in_features: usize) -> usize { + match self { + QtipGeometry::K4V2L16 => self.num_symbols(in_features) / 2, + QtipGeometry::TrellisV4L12 { rung } => rung.data_bytes(rung.num_symbols(in_features)), + } + } + + /// The allocated length of a row of `in_features` weights โ€” the tensor's + /// last dimension, and what this format validates. + /// + /// **Delegated, not restated.** `num_symbols / (8 / K)` is wrong at K=9 โ€” + /// there is no whole number of symbols per byte โ€” and divides by zero, so + /// the formula has to be the general one. It is also invisible to test at + /// realistic shapes: every plausible `in_features` is a multiple of 32, so + /// `num_symbols ยท K` is a multiple of 8 and floor equals ceil. A duplicate + /// of this formula here would therefore be unguarded in practice (measured: + /// mutation W3 floored it and every format test stayed green), so the + /// family's own [`trellis_v4l12::Rung`] is the single implementation and it + /// is exercised at non-byte-aligned symbol counts by that module's tests. + pub fn packed_len(self, in_features: usize) -> usize { + match self { + QtipGeometry::K4V2L16 => self.num_symbols(in_features) / 2, + QtipGeometry::TrellisV4L12 { rung } => rung.packed_len(in_features), + } + } + + /// Short label for bake headers and error messages. + pub fn tag(self) -> String { + match self { + QtipGeometry::K4V2L16 => "k4v2l16".to_string(), + QtipGeometry::TrellisV4L12 { rung } => format!("k{}v4l12", rung.k()), + } + } + + /// Wire encoding. **The default geometry writes nothing**, so every + /// artifact Arc has already produced stays byte-identical and no existing + /// checksum moves. Only a geometry that cannot be inferred from absence + /// announces itself. + fn to_wire(self) -> Option<(u8, [u8; 3])> { + match self { + QtipGeometry::K4V2L16 => None, + QtipGeometry::TrellisV4L12 { rung } => Some(( + GEOMETRY_WIRE_TAG, + [ + rung.k() as u8, + trellis_v4l12::L as u8, + trellis_v4l12::V as u8, + ], + )), + } + } + + /// Parse a `(K, L, V)` triple into a known geometry. + /// + /// The triple is stored rather than an opaque ordinal so a payload is + /// self-describing to a human with a hex editor, and so an unsupported but + /// well-formed geometry produces a diagnosis instead of "unknown tag". + fn from_wire_body(k: u8, l: u8, v: u8) -> Result { + // Written as comparisons, not as match patterns: a bare `K` in pattern + // position is a const pattern only for as long as the const exists, and + // silently becomes a catch-all binding if it is ever renamed away. + let got = (k as u32, l as u32, v as u32); + if got == (K, L, V) { + return Ok(QtipGeometry::K4V2L16); + } + if got.1 == trellis_v4l12::L && got.2 == trellis_v4l12::V { + // In-family: K is the free parameter, so let the family say + // whether it has a decoder for this one. + return QtipGeometry::trellis_v4l12(got.0); + } + let (ok, ol, ov) = got; + candle_core::bail!( + "QtipLayer: UQFF payload declares trellis geometry K={ok}/L={ol}/V={ov}, which this \ + build has no decoder for. Refusing rather than decoding its symbols at the wrong \ + geometry โ€” the packed byte count alone does not identify a geometry, so a wrong \ + guess need not even fault.", + ) + } + + /// Check that a layer's tensors are shaped the way this geometry requires. + /// + /// The tag alone is a claim; this is what makes it true. The table's dtype + /// and length are the load-bearing half: K=4/V=2 and K=8/V=4 are both + /// 2 bits per weight, so `blocks` has the identical size either way and + /// cannot discriminate them. + fn validate_shapes(self, blocks: &Tensor, lut: &Tensor, in_features: usize) -> Result<()> { + if !in_features.is_multiple_of(self.v() as usize) { + candle_core::bail!( + "QtipLayer[{}]: in_features {in_features} is not a multiple of V={}", + self.tag(), + self.v() + ); + } + let want_packed = self.packed_len(in_features); + let got_packed = blocks.dims()[blocks.dims().len() - 1]; + if got_packed != want_packed { + candle_core::bail!( + "QtipLayer[{}]: blocks row is {got_packed} B but K={}/V={} with in_features \ + {in_features} needs {want_packed} B", + self.tag(), + self.k(), + self.v() + ); + } + if lut.dtype() != self.lut_dtype() { + candle_core::bail!( + "QtipLayer[{}]: reproduction table is {:?} but this geometry stores {:?}. The \ + table's element type is part of the format, not an implementation detail.", + self.tag(), + lut.dtype(), + self.lut_dtype() + ); + } + if lut.elem_count() != self.lut_values() { + candle_core::bail!( + "QtipLayer[{}]: reproduction table has {} values, expected 2^{} ร— {} = {}", + self.tag(), + lut.elem_count(), + self.l(), + self.v(), + self.lut_values() + ); + } + Ok(()) + } + + /// Codebooks this geometry can legally carry. + /// + /// The computed `sum2` codebook is V=2-specific โ€” `mcg_codeword_v2` folds + /// exactly two chained MCG products, and the CUDA twin + /// (`qtip_codebook.cuh::qtip_cb_pair_from_x0`) returns a `float2`. There is + /// no V=4 form of it, so pairing it with any Vโ‰ 2 geometry is not a thing + /// that can be decoded, and it is refused on both the write and the read + /// side rather than left to fail somewhere less obvious. + fn check_codebook(self, codebook: QtipCodebook) -> Result<()> { + match codebook { + QtipCodebook::Mcg { .. } if self.v() != 2 => candle_core::bail!( + "QtipLayer[{}]: the computed `sum2` codebook is V=2-only (it produces a PAIR of \ + reproduction values per state) and cannot describe a V={} geometry. This rung \ + is table-only.", + self.tag(), + self.v() + ), + _ => Ok(()), + } + } +} + /// QTIP 2-bit weight layer. /// /// # Storage layout โ€” two modes @@ -799,6 +1073,16 @@ pub struct QtipLayer { /// computes, so the table stays a correct fallback (and the CPU search /// still wants a flat one) while every GPU path skips it entirely. codebook: QtipCodebook, + /// Which trellis geometry `blocks` was baked at, and therefore how its + /// bytes decompose into symbols and how `lut` is indexed. + /// + /// [`QtipGeometry::K4V2L16`] for every artifact written before the + /// discriminator existed โ€” which is what they are. Deliberately has no + /// per-site default: both supported geometries are 2 bits per weight, so a + /// mislabelled layer has a correctly-sized `blocks` tensor and decodes to + /// plausible garbage rather than faulting. Every construction site states + /// it, and the compiler is what enforces that. + geometry: QtipGeometry, } /// Borrowed, dequantization-free view of a [`QtipLayer`]'s packed trellis @@ -1954,6 +2238,10 @@ impl QtipLayer { search_weights.is_some(), ), codebook, + // This crate's trellis search is K=4/V=2/L=16 throughout + // (`viterbi.rs` reaches 16 successors per step and groups by + // 2^(L-K)); there is no producer of any other geometry yet. + geometry: QtipGeometry::K4V2L16, }) } @@ -2082,6 +2370,11 @@ impl QtipLayer { // objective bit is false by construction, not by omission. search_detail: QtipSearchDetail::for_bake(mode, search, false), codebook, + // Same geometry as the CPU sibling above: the CUDA bake kernels are + // the K=4/V=2/L=16 trellis too, so a GPU-baked artifact must carry + // the identical wire tag. Omitting it here compiled fine without + // the `cuda` feature and only broke the CUDA lane. + geometry: QtipGeometry::K4V2L16, })) } @@ -2220,6 +2513,7 @@ impl QtipLayer { let mut shared_rotation_block: usize = 0; let mut shared_search_detail = QtipSearchDetail::Unknown; let mut shared_codebook = QtipCodebook::Gaussian; + let mut shared_geometry = QtipGeometry::K4V2L16; // Per-expert streaming. The full stack is kept on CPU (caller passes // device=CPU for the 3-D experts to avoid the ~4GB dense BF16 transient @@ -2341,6 +2635,7 @@ impl QtipLayer { shared_rotation_block = layer.rotation_block; shared_search_detail = layer.search_detail; shared_codebook = layer.codebook; + shared_geometry = layer.geometry; } else { debug_assert_eq!(layer.lut.dims(), shared_lut.as_ref().unwrap().dims()); debug_assert_eq!(layer.rotation_block, shared_rotation_block); @@ -2369,6 +2664,18 @@ impl QtipLayer { shared_codebook.tag() ); } + // And the geometry, for the same reason and with the same + // teeth: at 2 bpw a mixed-geometry stack has consistently + // sized `blocks`, so nothing downstream would notice. + if layer.geometry != shared_geometry { + candle_core::bail!( + "QTIP 3-D quantize: expert chunk at {expert_idx} used geometry {} but \ + the stack already carries {} โ€” refusing to build a stack whose experts \ + decode at different geometries.", + layer.geometry.tag(), + shared_geometry.tag() + ); + } } expert_idx += this_b; } @@ -2398,6 +2705,7 @@ impl QtipLayer { search: QtipSearchStamp::for_mode(mode), search_detail: shared_search_detail, codebook: shared_codebook, + geometry: shared_geometry, }) } @@ -2473,6 +2781,7 @@ impl QtipLayer { } pub fn dequantize_weights(&self) -> Result { + self.require_k4v2l16("dequantize_weights")?; // 3-D stacked-expert path: iterate over experts and stack `[N, K_in]` // slices into `[E, N, K_in]`. We reuse the 2-D dequant per expert // rather than writing a 3-D CUDA kernel โ€” the existing @@ -3084,6 +3393,7 @@ impl QtipLayer { search: QtipSearchStamp, search_detail: QtipSearchDetail, codebook: QtipCodebook, + geometry: QtipGeometry, ) -> Result { if blocks.dims().len() != 3 { candle_core::bail!( @@ -3112,6 +3422,12 @@ impl QtipLayer { ); } let e = blocks.dim(0)?; + // The caller asserts a geometry; make it true here rather than at the + // first decode. Both supported geometries are 2 bits per weight, so + // `blocks` alone cannot tell them apart โ€” the table's dtype and length + // are what discriminate. + geometry.validate_shapes(&blocks, &lut, in_features)?; + geometry.check_codebook(codebook)?; Ok(QtipLayer { search_detail, blocks, @@ -3124,6 +3440,7 @@ impl QtipLayer { rotation_block, search, codebook, + geometry, }) } @@ -3201,6 +3518,15 @@ impl QtipLayer { head.codebook.tag() ); } + // Only `head.lut` survives the stack here too, and a table is + // meaningful only at the geometry it was built for. + if layer.geometry != head.geometry { + candle_core::bail!( + "QtipLayer::stack_experts: layer {i} geometry={} != head {}", + layer.geometry.tag(), + head.geometry.tag() + ); + } } let blocks_refs: Vec<&Tensor> = per_expert_layers.iter().map(|l| &l.blocks).collect(); @@ -3225,6 +3551,7 @@ impl QtipLayer { search: head.search, search_detail: head.search_detail, codebook: head.codebook, + geometry: head.geometry, }) } @@ -3507,6 +3834,7 @@ impl QtipLayer { search: self.search, search_detail: self.search_detail, codebook: self.codebook, + geometry: self.geometry, }) } } @@ -3549,6 +3877,11 @@ impl QuantMethod for QtipLayer { // table holds the computed values. It only forgoes the // in-register decode. codebook: QtipCodebook::Gaussian, + // Same again for the geometry: the config carries no + // discriminator, and every payload that reached this factory + // before the field existed was K=4/V=2/L=16. Reading it as + // anything else would be inventing a claim. + geometry: QtipGeometry::K4V2L16, }), _ => candle_core::bail!("QtipLayer requires QuantMethodConfig::Qtip"), } @@ -3559,10 +3892,17 @@ impl QuantMethod for QtipLayer { } fn qtip_packed(&self) -> Option> { - Some(self.packed_view()) + // `QtipPackedView` carries no geometry field, so a consumer would read + // these bytes as K=4/V=2 nibbles. Hand out nothing rather than + // something that cannot describe itself. + match self.geometry { + QtipGeometry::K4V2L16 => Some(self.packed_view()), + _ => None, + } } fn forward(&self, x: &Tensor) -> Result { + self.require_k4v2l16("forward")?; self.forward_dequantize(x) } @@ -3606,6 +3946,7 @@ impl QuantMethod for QtipLayer { /// dequantized weight in the rotated frame and matching it with a /// rotated activation โ€” saves one rotation pass per expert. fn gather_forward(&self, a: &Tensor, indices: &Tensor) -> Result { + self.require_k4v2l16("gather_forward")?; if self.num_experts.is_none() { // The contract is "expert-sparse dispatch" โ€” a 2-D layer is a // single-expert (i.e. non-MoE) layer and `gather_forward` @@ -3915,10 +4256,47 @@ impl QtipLayer { self.codebook } + /// Which trellis geometry this layer's symbols were baked at, and + /// therefore which decoder can read them. See [`QtipGeometry`]. + pub fn geometry(&self) -> QtipGeometry { + self.geometry + } + + /// Refuse any decode path that only knows K=4/V=2/L=16. + /// + /// **This is what keeps stage 2 from being a half-applied change.** The + /// geometry discriminator makes a K=8/V=4/L=12 artifact *loadable*; the + /// decoders that would then read it โ€” `dequantize_weights_rotated_f32`, + /// `dequantize_single_expert`, `gather_forward_cpu`, and every CUDA + /// launcher in `cuda_ops` other than `fused_gemv_k8v4l12_cuda` โ€” all + /// unpack nibbles and index a `[2^16, 2]` table. Handed a K=8 row they + /// would not fault: at 2 bits per weight the byte count is identical and + /// every index lands in bounds. They would serve garbage. + /// + /// So every entry point into decode states the geometries it can handle. + /// A layer this build cannot serve is refused loudly at the door rather + /// than quietly at the logits. + fn require_k4v2l16(&self, op: &str) -> Result<()> { + match self.geometry { + QtipGeometry::K4V2L16 => Ok(()), + other => candle_core::bail!( + "QtipLayer::{op}: this layer was baked at geometry {}, and {op} only implements \ + K=4/V=2/L=16. The {} decode path is not wired into serving yet โ€” the fused \ + GEMV kernel exists (`kernels/qtip/qtip_gemv_k8v4l12.cu`) but nothing dispatches \ + to it. Refusing rather than decoding these symbols as K=4/V=2, which would not \ + fault (both geometries are 2 bits per weight, so the packed byte count is the \ + same) and would silently serve wrong weights.", + other.tag(), + other.tag() + ), + } + } + /// Dequantize the i-th expert's `[N, K_in]` BF16 weight matrix (3-D mode /// only). Internal use by `gather_forward` and friends; bails when called /// on a 2-D layer or with `expert_idx >= num_experts`. pub fn dequantize_expert(&self, expert_idx: usize) -> Result { + self.require_k4v2l16("dequantize_expert")?; let e = self.num_experts.ok_or_else(|| { candle_core::Error::Msg("QtipLayer::dequantize_expert called on a 2-D layer".into()) })?; @@ -4025,6 +4403,9 @@ impl QtipLayer { // Gaussian default above) โ€” correct for any table, and never // assumes a computed codebook we cannot verify. codebook: QtipCodebook::Gaussian, + // Safetensors QTIP checkpoints predate every geometry but this + // one, and carry no field that could say otherwise. + geometry: QtipGeometry::K4V2L16, })) } @@ -4180,18 +4561,57 @@ impl QtipLayer { // meant. An unknown tag value, on the other hand, IS refused: that // means a newer Arc wrote it, and decoding its symbols against a // codebook we guessed would be silent corruption. - let codebook = match buffer.read_u8() { - Ok(tag) => QtipCodebook::from_wire(tag, || { - buffer.read_u32::().map_err(|_| { - candle_core::Error::Msg( - "QtipLayer: codebook tag claims a computed codebook but the multiplier \ - is missing (truncated payload)." - .into(), - ) - }) - })?, - Err(_) => QtipCodebook::Gaussian, - }; + // The trailing region is a sequence of tagged sections, read until EOF. + // Each tag may appear at most once: a repeat means a corrupt or + // concatenated payload, and silently letting the last one win is how a + // reader ends up decoding against something nobody wrote. + let mut codebook: Option = None; + let mut geometry: Option = None; + while let Ok(tag) = buffer.read_u8() { + if tag == GEOMETRY_WIRE_TAG { + if geometry.is_some() { + candle_core::bail!( + "QtipLayer: the UQFF payload carries two geometry sections. Refusing a \ + payload that describes itself twice." + ); + } + let mut klv = [0u8; 3]; + for (i, slot) in klv.iter_mut().enumerate() { + *slot = buffer.read_u8().map_err(|_| { + candle_core::Error::Msg(format!( + "QtipLayer: geometry section is truncated (got {i} of 3 K/L/V bytes)." + )) + })?; + } + geometry = Some(QtipGeometry::from_wire_body(klv[0], klv[1], klv[2])?); + } else { + if codebook.is_some() { + candle_core::bail!( + "QtipLayer: the UQFF payload carries two codebook sections. Refusing a \ + payload that describes itself twice." + ); + } + codebook = Some(QtipCodebook::from_wire(tag, || { + buffer.read_u32::().map_err(|_| { + candle_core::Error::Msg( + "QtipLayer: codebook tag claims a computed codebook but the \ + multiplier is missing (truncated payload)." + .into(), + ) + }) + })?); + } + } + // Absence is the answer for both, and it is the same answer it has + // always been: the historical geometry, gathering from the stored table. + let codebook = codebook.unwrap_or(QtipCodebook::Gaussian); + let geometry = geometry.unwrap_or(QtipGeometry::K4V2L16); + + // The tag is a claim; these make it true. Both supported geometries are + // 2 bits per weight, so `blocks` is the same size either way โ€” the + // table's dtype and length are what actually discriminate. + geometry.validate_shapes(&blocks, &lut, in_features)?; + geometry.check_codebook(codebook)?; Ok(( Self { @@ -4206,6 +4626,7 @@ impl QtipLayer { search, search_detail, codebook, + geometry, }, ext_bias, )) @@ -4281,10 +4702,25 @@ impl QuantizedSerde for QtipLayer { if let Some(w) = beam_width { buffer.extend(&w.to_le_bytes()); } + // The trellis geometry, written FIRST in the trailing-section region + // and only when it is not the historical K=4/V=2/L=16 โ€” see + // `QtipGeometry::to_wire` and `GEOMETRY_WIRE_TAG` for why the order is + // load-bearing. In short: a build that predates this field reads one + // trailing byte and hands it to the codebook parser, which refuses + // every tag it does not know. Tag 3 first means an old Arc fails + // closed on a geometry it cannot decode; tag 3 last would let it + // decode K=8/V=4 symbols as K=4/V=2 without faulting, because at 2 bpw + // the packed byte count is the same. + self.geometry.check_codebook(self.codebook)?; + if let Some((geo_tag, klv)) = self.geometry.to_wire() { + buffer.push(geo_tag); + buffer.extend(&klv); + } // wave24-AU: the codebook discriminator, appended LAST and only when // the codebook is not the historical Gaussian one โ€” see - // `QtipCodebook::to_wire`. A Gaussian artifact is therefore byte-for- - // byte what this writer produced before the field existed. + // `QtipCodebook::to_wire`. A Gaussian artifact at the default geometry + // is therefore byte-for-byte what this writer produced before either + // field existed. // // A build that predates the field, handed a computed-codebook payload, // parses identically up to here and then ignores a trailing section it @@ -5823,6 +6259,7 @@ mod tests { let mut shared_rotation_block: usize = 0; let mut shared_search_detail = QtipSearchDetail::Unknown; let mut shared_codebook = QtipCodebook::Gaussian; + let mut shared_geometry = QtipGeometry::K4V2L16; for expert_idx in 0..e { let expert_w = w.narrow(0, expert_idx, 1)?.squeeze(0)?; let layer = QtipLayer::quantize_with_options_concrete( @@ -5840,6 +6277,7 @@ mod tests { shared_rotation_block = layer.rotation_block; shared_search_detail = layer.search_detail; shared_codebook = layer.codebook; + shared_geometry = layer.geometry; } } let blocks_3d = Tensor::stack(&blocks_slices, 0)?; @@ -5856,6 +6294,7 @@ mod tests { search: QtipSearchStamp::for_mode(mode), search_detail: shared_search_detail, codebook: shared_codebook, + geometry: shared_geometry, }) .inspect(|l| { debug_assert_eq!(l.blocks.dims(), &[e, n, k_in / 4]); @@ -6742,11 +7181,658 @@ mod tests { let bb: Vec = bv.iter().map(|x| x.to_bits()).collect(); assert_eq!(ab, bb, "{what}: f32 bit patterns differ"); } + DType::BF16 => { + // The K=8/V=4/L=12 reproduction table. Compared as raw bits, + // not as floats: a table is format, and `==` on bf16 would + // still call two NaN payloads unequal and two zeros equal. + let av: Vec = a.flatten_all()?.to_vec1()?; + let bv: Vec = b.flatten_all()?.to_vec1()?; + let ab: Vec = av.iter().map(|x| x.to_bits()).collect(); + let bb: Vec = bv.iter().map(|x| x.to_bits()).collect(); + assert_eq!(ab, bb, "{what}: bf16 bit patterns differ"); + } other => panic!("assert_tensor_bits_eq: unhandled dtype {other:?}"), } Ok(()) } + // =================================================================== + // Geometry discriminator (UQFF) + // =================================================================== + + /// The V=4/L=12 family at the byte-aligned control K, as a geometry value. + fn k8() -> QtipGeometry { + QtipGeometry::trellis_v4l12(8).unwrap() + } + + /// A cheap 2-D layer whose fields we can then rewrite for format tests. + fn geometry_fixture(device: &Device) -> Result { + let (n, k_in) = (4usize, 64usize); + let wdata: Vec = (0..(n * k_in)) + .map(|i| ((i as f32) * 0.37).sin() * 1.1) + .collect(); + let w = Tensor::from_vec(wdata, (n, k_in), device)?; + QtipLayer::quantize_with_options_concrete(&w, None, device, QtipMode::Greedy, false) + } + + /// The default geometry must be **purely additive**: it writes nothing, so + /// every artifact Arc has already produced is byte-identical and no + /// checksum moves. Stated as an exact suffix rather than as a vague + /// "unchanged", so the tag value, its position, and the K byte are all + /// pinned โ€” **at every K in the family**. + /// + /// This is the property that makes a K change cheap: K is an explicit wire + /// field, so moving from the K=8 control to K=9 changes one byte of the + /// artifact and nothing else about the format. + #[test] + fn default_geometry_writes_nothing_and_the_family_appends_exactly_four_bytes() -> Result<()> { + let device = Device::Cpu; + let mut layer = geometry_fixture(&device)?; + assert_eq!(layer.geometry, QtipGeometry::K4V2L16); + let base = layer.serialize()?.into_owned(); + + for k in trellis_v4l12::K_SUPPORTED { + // Flip ONLY the discriminator; every tensor stays byte-identical. + layer.geometry = QtipGeometry::trellis_v4l12(k)?; + let tagged = layer.serialize()?.into_owned(); + + let mut want = base.clone(); + want.extend_from_slice(&[GEOMETRY_WIRE_TAG, k as u8, 12, 4]); + assert_eq!( + tagged, want, + "K={k}: the geometry section must be exactly [tag=3, K, L=12, V=4] appended to \ + an otherwise unchanged payload" + ); + // The K byte is the ONLY thing that differs between family members. + assert_eq!(tagged.len(), base.len() + 4); + } + Ok(()) + } + + /// Two family members differ by exactly one byte on the wire. + /// + /// The concrete statement of "K is a parameter, not a variant": going from + /// the K=8 control to the K=9 quality winner is a single byte of artifact + /// change, given identical tensors. + #[test] + fn family_members_differ_by_one_wire_byte() -> Result<()> { + let device = Device::Cpu; + let mut layer = geometry_fixture(&device)?; + layer.geometry = QtipGeometry::trellis_v4l12(8)?; + let a = layer.serialize()?.into_owned(); + layer.geometry = QtipGeometry::trellis_v4l12(9)?; + let b = layer.serialize()?.into_owned(); + assert_eq!(a.len(), b.len()); + let diffs: Vec = (0..a.len()).filter(|&i| a[i] != b[i]).collect(); + assert_eq!( + diffs.len(), + 1, + "K=8 and K=9 payloads must differ in exactly one byte, differed in {diffs:?}" + ); + assert_eq!((a[diffs[0]], b[diffs[0]]), (8, 9)); + Ok(()) + } + + /// The section must come BEFORE the codebook section, because that is what + /// makes an older Arc fail closed instead of mis-decoding. + #[test] + fn geometry_section_precedes_the_codebook_section() -> Result<()> { + let device = Device::Cpu; + let mut layer = geometry_fixture(&device)?; + let base = layer.serialize()?.into_owned(); + + // Codebook alone (K=4 geometry, so Mcg is legal here). + layer.codebook = QtipCodebook::COMPUTED; + let cb_only = layer.serialize()?.into_owned(); + assert_eq!( + cb_only.len(), + base.len() + 5, + "codebook section is tag + u32" + ); + assert_eq!(cb_only[base.len()], 2, "codebook tag is 2"); + + // Both sections. Geometry (4 bytes) must land first. + // + // K8V4L12 + Mcg is refused as a pair, so use a hypothetical third + // geometry is not available โ€” instead assert the writer's ORDER + // directly by checking that the geometry bytes occupy the slot the + // codebook otherwise would. + layer.codebook = QtipCodebook::Gaussian; + layer.geometry = k8(); + let geo_only = layer.serialize()?.into_owned(); + assert_eq!( + geo_only[base.len()], + GEOMETRY_WIRE_TAG, + "the geometry section must start immediately after the search-detail section โ€” \ + any byte between them would be read as a codebook tag by an older Arc" + ); + Ok(()) + } + + /// An Arc that predates this field parses the trailing region by handing + /// its first byte to `QtipCodebook::from_wire`. That must REFUSE tag 3. + /// + /// If it ever accepted it, an old build would decode K=8/V=4 symbols as + /// K=4/V=2 โ€” and would not fault, because at 2 bits per weight the packed + /// byte count is identical. + #[test] + fn an_older_reader_refuses_the_geometry_tag_rather_than_ignoring_it() { + let err = QtipCodebook::from_wire(GEOMETRY_WIRE_TAG, || Ok(1)) + .expect_err("tag 3 must not parse as a codebook"); + let msg = err.to_string(); + assert!( + msg.contains("geometry") || msg.contains("unknown codebook tag"), + "refusal must name the cause: {msg}" + ); + // And the general property the whole scheme rests on: every tag this + // build does not know is refused, never skipped. + for tag in [1u8, 4, 5, 200, 255] { + assert!( + QtipCodebook::from_wire(tag, || Ok(1)).is_err(), + "tag {tag} must fail closed" + ); + } + } + + /// A K=8/V=4 tag on a K=4/V=2 payload must be caught by the SHAPES, not + /// merely trusted. This is the case a hex-editor edit or a truncated + /// concatenation produces. + #[test] + fn a_geometry_tag_that_contradicts_the_tensors_is_refused_at_load() -> Result<()> { + let device = Device::Cpu; + let mut layer = geometry_fixture(&device)?; + layer.geometry = k8(); + let data = layer.serialize()?.into_owned(); + + let err = QtipLayer::deserialize_concrete_unchecked( + Cow::Owned(data), + &device, + crate::QuantizeOntoGuard::new(), + ) + .expect_err("a K8V4L12 tag over a [65536, 2] F32 table must be refused"); + let msg = err.to_string(); + assert!( + msg.contains("k8v4l12") && (msg.contains("F32") || msg.contains("values")), + "refusal must name the geometry and the mismatch: {msg}" + ); + Ok(()) + } + + /// The table's ELEMENT TYPE discriminates, not only its length. + /// + /// A `[4096, 4]` F32 table has exactly the right number of values and the + /// wrong dtype โ€” 65,536 B instead of 32,768 B, so it would not even fit + /// the shared-memory budget the geometry exists for. Without this case the + /// element-count check alone passes and the dtype check is dead + /// (measured: mutation F8 disabled it and every test stayed green). + #[test] + fn a_right_sized_table_of_the_wrong_dtype_is_refused() -> Result<()> { + let device = Device::Cpu; + let geo = k8(); + let k_in = 64usize; + let blocks = Tensor::zeros((2, geo.packed_len(k_in)), DType::U8, &device)?; + let f32_table = Tensor::zeros( + (trellis_v4l12::LUT_STATES, trellis_v4l12::V as usize), + DType::F32, + &device, + )?; + assert_eq!(f32_table.elem_count(), geo.lut_values(), "same count"); + let err = geo + .validate_shapes(&blocks, &f32_table, k_in) + .expect_err("an F32 table at this geometry must be refused"); + assert!( + err.to_string().contains("F32") && err.to_string().contains("BF16"), + "refusal must name both dtypes: {err}" + ); + Ok(()) + } + + /// `from_stacked_parts` is a public constructor that takes the geometry on + /// trust from its caller; it must validate, not trust. + #[test] + fn from_stacked_parts_refuses_a_geometry_its_tensors_contradict() -> Result<()> { + let device = Device::Cpu; + let (e, n, k_in) = (2usize, 4usize, 64usize); + let blocks = Tensor::zeros((e, n, k_in / 4), DType::U8, &device)?; + let scales = Tensor::zeros((e, n), DType::F32, &device)?; + let k4_table = Tensor::zeros((1usize << L, V as usize), DType::F32, &device)?; + + // Truthful call: K=4 tensors, K=4 tag. + QtipLayer::from_stacked_parts( + blocks.clone(), + scales.clone(), + k4_table.clone(), + None, + k_in, + None, + 0, + QtipSearchStamp::Unstamped, + QtipSearchDetail::Unknown, + QtipCodebook::Gaussian, + QtipGeometry::K4V2L16, + )?; + + // Same tensors, K=8 tag. + let err = QtipLayer::from_stacked_parts( + blocks, + scales, + k4_table, + None, + k_in, + None, + 0, + QtipSearchStamp::Unstamped, + QtipSearchDetail::Unknown, + QtipCodebook::Gaussian, + k8(), + ) + .expect_err("a K8V4L12 tag over a K4V2L16 table must be refused at construction"); + assert!(err.to_string().contains("k8v4l12"), "{err}"); + Ok(()) + } + + /// A stack whose experts were baked at different geometries keeps only one + /// table, so at most one expert could ever decode correctly. + #[test] + fn stack_experts_refuses_mixed_geometries() -> Result<()> { + let device = Device::Cpu; + let a = geometry_fixture(&device)?; + let mut b = geometry_fixture(&device)?; + // Identical in every respect except the discriminator, so nothing but + // the geometry check can be what fires. + b.geometry = k8(); + let err = QtipLayer::stack_experts(vec![a, b]) + .expect_err("mixed-geometry stacks must be refused"); + assert!( + err.to_string().contains("geometry"), + "refusal must name the geometry: {err}" + ); + + // ...and a uniform stack still works, so the guard is not just + // rejecting everything. + let c = geometry_fixture(&device)?; + let d = geometry_fixture(&device)?; + QtipLayer::stack_experts(vec![c, d])?; + Ok(()) + } + + /// Every decode entry point must refuse a geometry it cannot decode. + /// + /// Stage 2 makes a K=8/V=4/L=12 artifact loadable. The decoders behind + /// these entry points all unpack nibbles and index a `[2^16, 2]` table, and + /// handed a K=8 row they would NOT fault โ€” at 2 bits per weight the packed + /// byte count is identical and every index lands in bounds. They would + /// serve garbage. Without this test, "the format landed" would mean "the + /// engine now has a way to silently serve wrong weights". + #[test] + fn every_decode_entry_point_refuses_an_unsupported_geometry() -> Result<()> { + let device = Device::Cpu; + let k_in = 64usize; + + // 2-D entry points. + let mut two_d = geometry_fixture(&device)?; + // Sanity: all of these work at the geometry this build implements, so + // the refusals below are about the geometry and nothing else. + two_d.dequantize_weights()?; + let x = Tensor::zeros((1, k_in), DType::F32, &device)?; + two_d.forward(&x)?; + two_d.dequantize_w()?; + assert!(two_d.qtip_packed().is_some()); + + two_d.geometry = k8(); + for (what, res) in [ + ("dequantize_weights", two_d.dequantize_weights().err()), + ("forward", two_d.forward(&x).err()), + ("dequantize_w", two_d.dequantize_w().err()), + ] { + let err = res.unwrap_or_else(|| panic!("{what} must refuse a k8v4l12 layer")); + let msg = err.to_string(); + assert!( + msg.contains("k8v4l12") && msg.contains("Refusing"), + "{what}: refusal must name the geometry and say it is refusing: {msg}" + ); + // The refusal must come from THIS entry point, not from something + // it happens to call. `forward` reaching an inner guard is still a + // refusal, but it is not evidence that `forward` checks โ€” and an + // entry point that relies on a callee's guard breaks the moment the + // callee gains a fast path that skips it. (Measured: mutation G1 + // removed `forward`'s own guard and every test stayed green.) + // + // `dequantize_w` is the one exception by construction: it is a + // one-line delegation to `dequantize_weights` and has no body to + // guard. + let expect_op = if what == "dequantize_w" { + "dequantize_weights" + } else { + what + }; + assert!( + msg.contains(&format!("QtipLayer::{expect_op}:")), + "{what}: refusal came from elsewhere โ€” expected `QtipLayer::{expect_op}:` in: {msg}" + ); + } + assert!( + two_d.qtip_packed().is_none(), + "qtip_packed must hand out nothing โ€” QtipPackedView carries no geometry, so a \ + consumer would read the bytes as K=4/V=2 nibbles" + ); + + // 3-D entry points. + let mut stack = + QtipLayer::stack_experts(vec![geometry_fixture(&device)?, geometry_fixture(&device)?])?; + stack.dequantize_expert(0)?; + let a = Tensor::zeros((1, 1, k_in), DType::F32, &device)?; + let idx = Tensor::zeros((1, 1), DType::U32, &device)?; + stack.gather_forward(&a, &idx)?; + + stack.geometry = k8(); + for (what, res) in [ + ("dequantize_expert", stack.dequantize_expert(0).err()), + ("gather_forward", stack.gather_forward(&a, &idx).err()), + ] { + let err = res.unwrap_or_else(|| panic!("{what} must refuse a k8v4l12 layer")); + let msg = err.to_string(); + assert!(msg.contains("k8v4l12"), "{what}: {msg}"); + assert!( + msg.contains(&format!("QtipLayer::{what}:")), + "{what}: refusal came from elsewhere: {msg}" + ); + } + Ok(()) + } + + /// A genuine V=4/L=12 payload round-trips at every K, tag and tensors both. + #[test] + fn trellis_v4l12_payloads_round_trip_exactly_at_every_k() -> Result<()> { + let device = Device::Cpu; + let (n, k_in) = (4usize, 128usize); + let lut = Tensor::from_slice( + &trellis_v4l12::gaussian_lut_bf16(), + (trellis_v4l12::LUT_STATES, trellis_v4l12::V as usize), + &device, + )?; + + for k in trellis_v4l12::K_SUPPORTED { + let geo = QtipGeometry::trellis_v4l12(k)?; + let packed_per_row = geo.packed_len(k_in); + // Sanity: the DATA length tracks K, and the allocated stride is + // that plus the padding the unclamped extraction relies on. + assert_eq!(geo.data_bytes(k_in), (k_in / 4 * k as usize).div_ceil(8)); + assert!(packed_per_row >= geo.data_bytes(k_in)); + assert!(packed_per_row - geo.data_bytes(k_in) <= 4); + + let blocks_data: Vec = (0..(n * packed_per_row)) + .map(|i| (i * 37 % 251) as u8) + .collect(); + let blocks = Tensor::from_vec(blocks_data, (n, packed_per_row), &device)?; + let scales = Tensor::from_vec( + (0..n).map(|i| 0.01 + i as f32 * 0.003).collect::>(), + (n,), + &device, + )?; + + let layer = QtipLayer { + blocks, + row_scales: scales, + lut: lut.clone(), + bias: None, + in_features: k_in, + num_experts: None, + rotation_signs: None, + rotation_block: 0, + search: QtipSearchStamp::Trellis, + search_detail: QtipSearchDetail::Known { + beam_width: None, + hessian: false, + }, + codebook: QtipCodebook::Gaussian, + geometry: geo, + }; + + let data = layer.serialize()?.into_owned(); + let (restored, _) = QtipLayer::deserialize_concrete_unchecked( + Cow::Owned(data), + &device, + crate::QuantizeOntoGuard::new(), + )?; + + assert_eq!(restored.geometry, geo, "K={k}"); + assert_eq!(restored.geometry.k(), k); + assert_eq!(restored.in_features, k_in); + assert_eq!(restored.lut.dtype(), DType::BF16); + assert_eq!(restored.lut.elem_count(), trellis_v4l12::LUT_ENTRIES); + assert_tensor_bits_eq(&layer.blocks, &restored.blocks, "blocks")?; + assert_tensor_bits_eq(&layer.row_scales, &restored.row_scales, "row_scales")?; + assert_tensor_bits_eq(&layer.lut, &restored.lut, "lut")?; + } + Ok(()) + } + + /// A K the wire names but this build has no decoder for is refused, and the + /// refusal quotes the triple. + #[test] + fn an_in_family_but_unsupported_k_is_refused_with_a_diagnosis() -> Result<()> { + let device = Device::Cpu; + let layer = geometry_fixture(&device)?; + let base = layer.serialize()?.into_owned(); + // L=12/V=4 is this family, but K=11 has no decoder. + let mut bytes = base.clone(); + bytes.extend_from_slice(&[GEOMETRY_WIRE_TAG, 11, 12, 4]); + let err = QtipLayer::deserialize_concrete_unchecked( + Cow::Owned(bytes), + &device, + crate::QuantizeOntoGuard::new(), + ) + .expect_err("K=11 must be refused"); + assert!( + err.to_string().contains("K=11") || err.to_string().contains("not implemented"), + "refusal must name the K it could not decode: {err}" + ); + Ok(()) + } + + /// The computed `sum2` codebook produces a PAIR of values per state and so + /// cannot describe V=4. Refused on write and on read, not left to fail at + /// the first decode. + #[test] + fn the_computed_codebook_is_refused_at_a_v4_geometry() -> Result<()> { + let device = Device::Cpu; + let mut layer = geometry_fixture(&device)?; + layer.geometry = k8(); + layer.codebook = QtipCodebook::COMPUTED; + let err = layer + .serialize() + .expect_err("V=4 + sum2 must not serialize"); + assert!( + err.to_string().contains("V=2-only"), + "refusal must say why: {err}" + ); + // And it is legal at the geometry it was built for, so the guard is + // about the pairing and not about the codebook being blocked outright. + layer.geometry = QtipGeometry::K4V2L16; + layer.serialize()?; + Ok(()) + } + + /// Malformed trailing regions are refused rather than partially read. + #[test] + fn malformed_geometry_sections_are_refused() -> Result<()> { + let device = Device::Cpu; + let layer = geometry_fixture(&device)?; + let base = layer.serialize()?.into_owned(); + let load = |bytes: Vec| { + QtipLayer::deserialize_concrete_unchecked( + Cow::Owned(bytes), + &device, + crate::QuantizeOntoGuard::new(), + ) + }; + + // Truncated: tag present, K/L/V missing. + let mut truncated = base.clone(); + truncated.push(GEOMETRY_WIRE_TAG); + truncated.push(8); + let err = load(truncated).expect_err("truncated geometry section must be refused"); + assert!(err.to_string().contains("truncated"), "{err}"); + + // Well-formed but unsupported triple: diagnosed, not "unknown tag". + let mut unsupported = base.clone(); + unsupported.extend_from_slice(&[GEOMETRY_WIRE_TAG, 6, 18, 3]); + let err = load(unsupported).expect_err("unsupported geometry must be refused"); + let msg = err.to_string(); + assert!( + msg.contains("K=6") && msg.contains("L=18") && msg.contains("V=3"), + "refusal must quote the triple it could not decode: {msg}" + ); + + // Two geometry sections: a payload that describes itself twice. + let mut doubled = base.clone(); + doubled.extend_from_slice(&[GEOMETRY_WIRE_TAG, 4, 16, 2]); + doubled.extend_from_slice(&[GEOMETRY_WIRE_TAG, 4, 16, 2]); + let err = load(doubled).expect_err("duplicate geometry sections must be refused"); + assert!(err.to_string().contains("twice"), "{err}"); + + // An explicit default-geometry section is legal (a future writer may + // choose to be explicit) and must load as K4V2L16. + let mut explicit = base.clone(); + explicit.extend_from_slice(&[GEOMETRY_WIRE_TAG, 4, 16, 2]); + let (restored, _) = load(explicit)?; + assert_eq!(restored.geometry, QtipGeometry::K4V2L16); + Ok(()) + } + + /// The shipped rung and the K=8 control are the SAME bit rate and the SAME + /// packed size. That is the fact that makes the discriminator necessary โ€” + /// a mislabelled artifact has a correctly-sized `blocks` tensor โ€” so it + /// gets an assertion rather than a comment. + #[test] + fn k4v2l16_and_k8v4l12_are_indistinguishable_by_packed_size() { + let k8 = k8(); + for k_in in [64usize, 512, 4096, 7168] { + assert_eq!( + QtipGeometry::K4V2L16.packed_len(k_in), + k8.packed_len(k_in), + "k_in={k_in}: identical packed size is why the tag is load-bearing" + ); + } + assert_eq!(QtipGeometry::K4V2L16.bpw_x100(), 200); + assert_eq!(k8.bpw_x100(), 200); + // ...and the tables are what actually differ. + assert_eq!(QtipGeometry::K4V2L16.lut_values(), 65_536 * 2); + assert_eq!(k8.lut_values(), 4_096 * 4); + assert_eq!(QtipGeometry::K4V2L16.lut_dtype(), DType::F32); + assert_eq!(k8.lut_dtype(), DType::BF16); + } + + /// **Not every geometry in the enum is 2 bits per weight any more.** + /// + /// K=9/V=4 is 2.25 bpw, so its rows are 12.5% larger than K=8/V=4's and it + /// is *not* size-confusable with the shipped rung. Worth pinning: code that + /// assumed "all QTIP geometries are 2 bpw" โ€” a true statement when this + /// family was K=8-only โ€” is now wrong, and `packed_len` is the general + /// `ceil(nยทK/8)` rather than anything that divides by `8/K`. + #[test] + fn the_family_spans_more_than_one_bit_rate() { + let k8 = k8(); + let k9 = QtipGeometry::trellis_v4l12(9).unwrap(); + let k10 = QtipGeometry::trellis_v4l12(10).unwrap(); + assert_eq!(k8.bpw_x100(), 200); + assert_eq!(k9.bpw_x100(), 225); + assert_eq!(k10.bpw_x100(), 250); + for k_in in [512usize, 4096, 7168] { + assert!(k9.packed_len(k_in) > k8.packed_len(k_in)); + // Exactly the bit-rate ratio, to the byte โ€” on DATA bytes. The + // stride is only approximately in that ratio, because K=9 pads for + // its multi-byte extraction and K=8 does not. + assert_eq!(k9.data_bytes(k_in) * 8, k8.data_bytes(k_in) * 9); + // The whole family shares one table regardless of K. + assert_eq!(k9.lut_values(), k8.lut_values()); + assert_eq!(k9.lut_dtype(), k8.lut_dtype()); + } + } + + /// The packed length must use the CEIL, and that is only visible at an + /// `in_features` which is not a multiple of 32. + /// + /// Every plausible layer width is a multiple of 32, which makes + /// `num_symbols ยท K` a multiple of 8 and floor equal to ceil at K=9. A + /// floored formula therefore passes every realistic fixture (measured: + /// mutation W3). These widths are deliberately unrealistic. + #[test] + fn packed_len_uses_the_ceiling_where_it_is_observable() { + let k9 = QtipGeometry::trellis_v4l12(9).unwrap(); + // in_features=36 -> 9 symbols -> 81 bits -> 11 DATA bytes, not 10. + assert_eq!(k9.num_symbols(36), 9); + assert_eq!(k9.data_bytes(36), 11); + // in_features=4 -> 1 symbol -> 9 bits -> 2 data bytes, not 1. + assert_eq!(k9.data_bytes(4), 2); + let k10 = QtipGeometry::trellis_v4l12(10).unwrap(); + // 3 symbols -> 30 bits -> 4 data bytes, not 3. + assert_eq!(k10.data_bytes(12), 4); + // The stride is the data plus tail padding, rounded up to 4. + assert_eq!(k9.packed_len(36), 12); + assert_eq!(k9.packed_len(4), 4); + assert_eq!(k10.packed_len(12), 8); + // ...and both agree with the rung that owns the formulas. + for k in trellis_v4l12::K_SUPPORTED { + let g = QtipGeometry::trellis_v4l12(k).unwrap(); + let r = trellis_v4l12::Rung::new(k).unwrap(); + for k_in in [4usize, 12, 36, 100, 260, 4096] { + let n = r.num_symbols(k_in); + assert_eq!( + g.packed_len(k_in), + r.row_stride(n), + "K={k} k_in={k_in} stride" + ); + assert_eq!( + g.data_bytes(k_in), + r.data_bytes(n), + "K={k} k_in={k_in} data" + ); + } + } + } + + /// A K with no decoder is refused at construction, not stored and used. + #[test] + fn an_unsupported_k_cannot_become_a_geometry() { + for k in [0u32, 4, 7, 11, 16] { + assert!( + QtipGeometry::trellis_v4l12(k).is_err(), + "K={k} must not be constructible" + ); + } + for k in trellis_v4l12::K_SUPPORTED { + assert!(QtipGeometry::trellis_v4l12(k).is_ok(), "K={k} must be"); + } + } + + /// The geometry accessors must agree with the rung they describe. + #[test] + fn geometry_accessors_match_the_rung_module() { + assert_eq!(QtipGeometry::K4V2L16.k(), K); + assert_eq!(QtipGeometry::K4V2L16.l(), L); + assert_eq!(QtipGeometry::K4V2L16.v(), V); + for k in trellis_v4l12::K_SUPPORTED { + let g = QtipGeometry::trellis_v4l12(k).unwrap(); + let r = trellis_v4l12::Rung::new(k).unwrap(); + assert_eq!(g.k(), r.k()); + assert_eq!(g.l(), trellis_v4l12::L); + assert_eq!(g.v(), trellis_v4l12::V); + assert_eq!(g.bpw_x100(), r.bpw_x100()); + assert_eq!( + g.lut_values(), + trellis_v4l12::LUT_ENTRIES, + "the discriminator and the rung must agree on the table size" + ); + for k_in in [64usize, 512, 4096] { + assert_eq!(g.packed_len(k_in), r.packed_len(k_in), "K={k} k_in={k_in}"); + assert_eq!(g.num_symbols(k_in), r.num_symbols(k_in)); + } + assert_eq!(g.tag(), format!("k{k}v4l12")); + } + } + /// UQFF round-trip for the classic 2-D layout must stay lossless and /// byte-compatible (bias branch + no-rotation branch covered). #[test] From 73ff79d09e50df9234bf965aff905d06d755688e Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 18:09:05 +0100 Subject: [PATCH 07/21] feat(turboquant): CUDA kernels at head_dim 64/128/256/512, Hopper + Blackwell (#94) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(turboquant): instantiate the CUDA kernels at head_dim 64/128/256/512 TurboQuant's paged kernels were hardcoded to head_dim 128 โ€” `constexpr int HS = 128`, a single `SGN[128]` sign table, `rotate128`, and six `extern "C"` launchers each opening with `if (hs != 128) return;`. That early return is a silent no-op: the output buffer is left uninitialized, so a mismatched head dim produces garbage *fast* rather than failing. Kernels - `tq_attn`, `tq_attn_clocked`, `tq_cache_k`, `tq_cache_v` and the two clocked cache variants now take HS as a template parameter. NT stays 128 for every HS, so the warp-reduction topology is unchanged; wider heads give each thread DPT = HS/NT output dims instead of adding threads. At HS=128 the emitted work is the same as before. - The per-lane packed-K gather is now one vector load (`tq_load_klane`) instead of BPL scalar byte loads. The x=16 cache swizzle keeps each lane's run contiguous and naturally aligned, which `packed_k_bytes_are_swizzle_aligned` asserts. - Every `if (hs != 128) return;` is replaced by a dispatch over the instantiated widths. D16 dual-arch - `TQ_PREFETCH` software-pipelines the K gather on SM90/SM100/SM103 and runs flat below Hopper โ€” same arithmetic, scheduled for the part it runs on. - Launches now opt in to the arch's real dynamic shared-memory budget. The logits array scales with context length and CUDA caps dynamic shared memory at 48 KB without an explicit opt-in, so contexts past ~12k tokens could not launch at all on any arch. - `ARC_CUDA_ARCHS=90,100,103` emits cubins for Hopper *and* Blackwell rather than only the build box's GPU. Tables - Codebooks and sign diagonals are dimension-dependent, so 512 needs its own. All 20 tables are emitted by `emit_turboquant_cuda_tables`, which calls the same `generate_signs`/`get_codebook` the Rust compressor uses โ€” correct by construction, not by transcription. The generator reproduces the shipped d=128 table byte-for-byte. - `turboquant::cuda_tables` re-derives every table at test time and diffs the checked-in header, and pins the Rust head-dim list against both the kernel's dispatch switch and the FFI wrapper's list. Nothing previously asserted that the tree's four copies of the sign table agreed with the generator or with each other. Gates - `TURBOQUANT_HEAD_DIM = 128` exact-match becomes membership in `TURBOQUANT_CUDA_HEAD_DIMS`; K and V are compressed independently and no longer have to be equal widths. Not measured on a GPU: the kernels are written and gated but this box has no nvcc, so nothing here is a claim about hardware behaviour. * test(gpu): wave63 gate โ€” dual-arch compile proof for the head_dim 512 kernels Proves the four head-dim instantiations compile for sm_90a AND sm_100a/sm_103a (D16), not just the build box's GPU, and that the shipping decode path did not regress. Deliberately contains no quality A/B: TurboQuant quality is settled by prior measurement and re-proving it would spend GPU budget on a known number. * fix(build): watch turbo_paged_attention.cu for changes Every other .cu in this crate had a rerun-if-changed line; this one did not. Emitting any rerun-if-changed disables cargo's watch-the-whole-package default, so edits to the largest CUDA file in the crate were relying entirely on cudaforge's object cache to trigger a rebuild. * fix(turboquant): 64-bit cache offsets, per-device smem opt-in, clocked-kernel parity Three problems from review of the head_dim generalization. 1. 32-bit overflow on cache byte offsets. `pb*kbs`, `pb*vbs`, `bi*cbs` and `bi*vbs` were all `int`. At head_dim 512 with nkvh=8 and block_size=32, `kbs` is 65536 B, so `pb*kbs` wraps once the K cache passes 2 GB โ€” an entirely reachable allocation on an H200, and 4x sooner than at 128. Now computed in `long long` at all eight sites. 2. The shared-memory opt-in cached a process-wide `granted` flag per kernel instantiation. The requestable maximum differs by architecture, so one flag cannot be right across a mixed fleet, and a failed request left `granted` at 0 and fell through to a launch that then failed with no diagnostic. The cache is gone; the call is cheap and now always attempted. 3. `tq_attn_clocked` kept the scalar byte gather while `tq_attn` moved to a vector load plus the arch-gated pipeline. Identical results, but the phase-2 stamps would have described a kernel nobody launches โ€” at head_dim 512 that is 8 scalar loads measured against the 1 vector load actually executed. The loop now mirrors `tq_attn`. Verified without nvcc by type-checking the kernel bodies against a CUDA shim (clang++ -Wall -Wextra -Wshadow, both __CUDA_ARCH__ branches): clean. Still not run on a GPU. * fix(gate): D18 exit codes, branch-independent preflight, source-landmine guard The gate exited 1 when gpu_box_preflight.sh was absent. That read as a code failure for a couple of minutes when it was really 'this machine cannot answer'. Per D18 rule 2: 0 pass, 1 genuine failure (the only signal to act on), 2 environment could not answer. Nine conditions now exit 2 โ€” missing/failing/incomplete preflight, no repo, no nvcc, unreachable remote, un-checkout-able branch, no archive to inspect, no cuobjdump. Four exit 1, all genuine: nvcc compile failure, a native build failure, a test failure, and the archive missing an arch it claims to carry. The preflight is now looked up at ARC_PREFLIGHT, then /root/arc-tools, then /usr/local/lib/arc, and only last inside the repo. It lives on wave61/box-preflight-shared-prefix, so a repo-relative path vanishes the moment a runner checks out any other branch โ€” which is exactly how this gate came to refuse to run. Guarded the sourced-preflight landmine directly. gpu_box_preflight.sh ends on [ "$_arc_pf_sourced" = "1" ] || exit 1 so when SOURCED on a *failing* box that test succeeds, short-circuits the '|| exit 1', and becomes the last command โ€” 'source pf || handler' returns 0 and the handler never fires. The gate ignores the return value entirely and reads _ARC_PF_FAILED, treating an unset flag as 'did not reach a verdict' rather than as a pass. Also reordered so the sm_90a+sm_100a+sm_103a compile is STEP 1 and needs no full workspace build, and it now asserts on cuobjdump output that both cubins are actually in the archive โ€” a build that merely succeeds is not the claim. Exit paths verified locally against a stub reproducing the real preflight's sourced-exit structure: all four environment cases return 2. * fix(paged): keep the TurboQuant DEFAULT at head_dim 128; widen only what is asked for Widening the acceptance gate from head_dim == 128 to every instantiated width was correct for what a user may REQUEST, but PagedCacheType:: TurboQuant carries #[default] (cache_engine.rs:22), so it also silently moved standard-layout models at head_dim 64/256/512 onto kernels that have never executed โ€” and silently dropped prefix caching with them, since supports_prefix_cache() is false for every TurboQuant variant. Two regressions, neither visible at the call site, for users who never asked for TurboQuant at all. This repo has paid for that exact shape once: FP8 KV shipped default-on and unmeasured (wave43-BU) and every V4 request died. #98 cites that incident by name as its own reason not to default itself on; this takes the same choice. resolve_for_model already distinguished explicit from ambient-default in order to pick error-vs-fallback, so the fix is to let that same 'forced' flag also pick WHICH head-dim set applies: compiled -> widens what you may ask for (TURBOQUANT_CUDA_HEAD_DIMS) measured -> widens what you get unasked (TURBOQUANT_DEFAULT_HEAD_DIMS) TURBOQUANT_DEFAULT_HEAD_DIMS is [128] and moves when a hardware gate passes, not when a kernel compiles. Nothing about the kernels changes: the six 'if (hs != 128) return;' no-ops (D18 instance 3) and the 48 KB dynamic smem cap fix are untouched, and explicit --pa-cache-type turboquant still reaches every instantiated width including V4's 512. The unsupported-geometry message now separates 'no kernel exists' from 'a kernel exists but is unmeasured, so the default will not choose it for you' โ€” the second tells the user to opt in rather than to give up. Tests: the old fixture asserted the regression (its expected-to-fall-back list moved [64,96,192,256] -> [96,192,320,1024] precisely because 64 and 256 stopped falling back). Replaced with an explicit-path test over every instantiated width, plus two pins on the default path. Both pins were mutation-checked by reverting the const to the full set: both fail, and the first revision of one passed vacuously because widening the const emptied its loop, so it now asserts the loop is non-empty before iterating. * docs: correct the public record for the widened TurboQuant kernels This PR branched before #101 ("correct the public record") landed, so it had never seen that text. Merging master in makes four README claims false โ€” three falsified by this PR's own kernels, one that was already wrong when it was written. Falsified by this PR: - L31 "the paged kernel exists at head_dim 128 only" - L31 "there is no kernel at head_dim 512, so DeepSeek V4 cannot use it" DeepSeek V4 reports KvCacheLayout::Standard at head_dim 512 (normal_loaders.rs), so the 512 instantiation genuinely reaches it. - L141 an explicit `--pa-cache-type turboquant` off-128 "is a hard error" It is now accepted at any instantiated width; the hard error moved out to the uninstantiated set. Already wrong before this PR: - L141 "TurboQuant is not the default anywhere" `defaults::PAGED_CACHE_TYPE` is `PagedCacheType::TurboQuant` and the CLI's `--pa-cache-type` has no clap default, so leaving it unset on CUDA gives a standard-layout head_dim-128 model TurboQuant KV with no flag โ€” and silently drops prefix caching, which no TurboQuant variant supports. That was true on master; the narrowing in a3a5faab bounds it to 128 but does not remove it. Stating the opposite understated the risk to anyone reading the README to decide whether to opt out, so it is corrected in place rather than quietly reworded, and the Rust example's comment now says `Auto` opts *out* rather than implying TurboQuant is off until asked for. Also refreshes the `--pa-cache-type` help text, which still described the explicit path as 128-or-error. No behaviour change: documentation and one clap doc comment. * fix(build): stop emitting a duplicate -gencode for the build box's own arch `mistralrs-paged-attn/build.rs` appended one `-gencode` per entry in `ARC_CUDA_ARCHS` on top of the one cudaforge derives itself, and applied the same `a` suffix for cap >= 90 that cudaforge's `GpuArch::auto_suffix` does. Verified against cudaforge 0.1.5 rather than assumed: compute_cap.rs auto_suffix(90).to_gencode_arg() -> "-gencode=arch=compute_90a,code=sm_90a" builder.rs:401 let gencode_arg = gpu_arch.to_gencode_arg(); builder.rs:405 command.arg(&gencode_arg)... builder.rs:412 for arg in &self.extra_args { command.arg(arg); } cudaforge emits its own gencode first and then appends `extra_args` verbatim, so on an H200 building `ARC_CUDA_ARCHS=90,100,103` nvcc received `-gencode=arch=compute_90a,code=sm_90a` twice, byte-identical. Fixed by hoisting the existing `get_compute_cap()` binding above the loop (it is what decides which requested arches are genuinely extra) and skipping the one cudaforge already covers. The later duplicate binding is deleted; there is now one. This is the dedupe half of #108's `arc_target::build::split_primary()`, done locally because `arc-target` does not exist on this branch. The other half โ€” `verify_and_export()`, the `cuobjdump` gate that fails a build whose archive is missing a requested arch โ€” genuinely needs that crate and is left to #108, which should replace this block with the two calls. Until it lands, `arc-tools/wave64_v4_turboquant_kv_gate.sh` STEP 2 asserts the arch list externally with `cuobjdump --list-elf`, so a missing arch is caught by the gate even though it is not yet caught by the build. Not verified here: whether nvcc rejects or tolerates the duplicate flag. macOS cannot run the cuda build path. Emitting it once is correct either way, which is why this is worth doing without waiting on that answer. --- README.md | 8 +- arc-tools/wave63_turboquant_hd512_gate.sh | 153 ++++ mistralrs-cli/src/args/paged_attn.rs | 11 +- .../src/paged_attention/cache_engine.rs | 223 ++++- mistralrs-core/src/paged_attention/mod.rs | 33 +- mistralrs-paged-attn/build.rs | 47 + .../src/cuda/backend/paged_attention.rs | 29 +- .../src/cuda/turbo_paged_attention.cu | 817 ++++++++++-------- .../src/cuda/turbo_paged_attention.cuh | 678 +++++++-------- .../examples/emit_turboquant_cuda_tables.rs | 70 ++ mistralrs-quant/src/turboquant/cuda_tables.rs | 225 +++++ mistralrs-quant/src/turboquant/mod.rs | 2 + 12 files changed, 1559 insertions(+), 737 deletions(-) create mode 100755 arc-tools/wave63_turboquant_hd512_gate.sh create mode 100644 mistralrs-quant/examples/emit_turboquant_cuda_tables.rs create mode 100644 mistralrs-quant/src/turboquant/cuda_tables.rs diff --git a/README.md b/README.md index 15297bc17..ba023c3dc 100644 --- a/README.md +++ b/README.md @@ -28,7 +28,7 @@ Forked from [mistral.rs](https://github.com/EricLBuehler/mistral.rs). Apache 2.0 | Feature | What it does | |---|---| | **QTIP 2-bit weights** | Cornell ICLR'25 โ€” trellis quantization with Hadamard incoherence rotation. 8ร— weight compression vs FP16. Two rungs: `qtip2` (LUT) and `qtip2b` (computed codebook), the latter being what Arc serves. Works on any model. | -| **TurboQuant KV** *(paged default at head_dim 128)* | Zandieh et al. (Google Research, [arXiv:2504.19874](https://arxiv.org/abs/2504.19874), ICLR'26) โ€” Arc Rust implementation of WHT + Lloyd-Max KV coding. **On by default on the paged path**: a standard-layout head_dim-128 model on CUDA gets K4/V3 KV with no flag, and loses prefix caching with it (`--pa-cache-type auto` opts out). Every other geometry falls back to `auto` with a warning; there is no kernel at head_dim 512, so DeepSeek V4 does not take it. Served end-to-end on a B200 at **55 tok/s with correct output**; **quality has never been evaluated**. Eager (non-paged) KV is a separate path, opt-in via `ARC_TURBOQUANT_KV=1`. | +| **TurboQuant KV** *(paged default at head_dim 128)* | Zandieh et al. (Google Research, [arXiv:2504.19874](https://arxiv.org/abs/2504.19874), ICLR'26) โ€” Arc Rust implementation of WHT + Lloyd-Max KV coding. **On by default on the paged path**: a standard-layout head_dim-128 model on CUDA gets K4/V3 KV with no flag, and loses prefix caching with it (`--pa-cache-type auto` opts out). The kernels are now instantiated at head_dim 64/128/256/512, but only 128 has ever executed on hardware and only 128 is reachable without asking โ€” every other geometry falls back to `auto` with a warning, so DeepSeek V4's head_dim of 512 is opt-in via an explicit `--pa-cache-type turboquant`. Served end-to-end on a B200 at **55 tok/s with correct output** at head_dim 128; **quality has never been evaluated**, at any width. Eager (non-paged) KV is a separate path, opt-in via `ARC_TURBOQUANT_KV=1`. | | **TD-MoE Tucker + whitening** | Whitened Tucker decomposition of the MoE expert pool. The **"lossless 20%"** figure is the **paper's** (NeurIPS'25), not ours โ€” *published, not reproduced by us*. Works on any MoE. | | **ArcAttention/Indexer** โ€” sparse attention | DeepSeek V4's Lightning Indexer with FlashMLASparse CUDA kernels (MIT, ported from sgl-project) for top-k attention. Dense O(nยฒ) โ†’ sparse O(nยทk). | | **ArcGraph** โ€” GPU-autonomous decode | Full decode loop (forward + sample + EOS check) on GPU. Zero CPU sync per token. Works on any model. | @@ -79,7 +79,7 @@ Each layer compresses a different axis, and in principle the wins multiply. **Read the block above as arithmetic, because that is all it is.** The weight line is measured end-to-end: 284B โ†’ **~74 GB** at 2.09 bits/param (V4 Flash trellis bake, serving on one H200 โ€” see [Results](#results)). **The KV lines are not.** They are bytes-per-token ratios computed from the packed format; no run has ever produced a measured KV compression figure, and the "4.27ร—" once published here is retracted. -What *has* run: the TurboQuant K4/V3 paged path served Qwen3-32B on a B200 end-to-end with correct output (see [Results](#results)) โ€” so "no TurboQuant forward pass has ever been benchmarked", which this section previously asserted, was false. What that run does **not** establish is any of the ratios above, or any quality number. TurboQuant is the paged default at head_dim 128; it has no kernel at V4's head_dim of 512. The xKV pool is not implemented. +What *has* run: the TurboQuant K4/V3 paged path served Qwen3-32B on a B200 end-to-end with correct output (see [Results](#results)) โ€” so "no TurboQuant forward pass has ever been benchmarked", which this section previously asserted, was false. What that run does **not** establish is any of the ratios above, or any quality number, and it was at head_dim 128 only. TurboQuant is the paged default at head_dim 128; the paged kernels now also cover V4's head_dim of 512, but nothing has ever run at that width, so it stays opt-in. The xKV pool is not implemented. For long context (1M+ tokens), the bottleneck shifts from weights to attention compute + KV bandwidth. Arc handles that via FlashMLASparse (CUDA kernel, MIT-licensed, ported from sgl-project), turning dense attention's O(nยฒ) into sparse top-k O(nยทk) on models with native sparse-attention training (DeepSeek V3.2+ family) and via top-k attention + sink preservation on the rest. @@ -144,9 +144,9 @@ Python bindings: `pip install mistralrs-cuda`. > **TurboQuant *is* the paged default** โ€” correcting an earlier note here that read "TurboQuant is not the default anywhere", which was false when written. `defaults::PAGED_CACHE_TYPE` is `PagedCacheType::TurboQuant` and `--pa-cache-type` carries no clap default, so leaving it unset on CUDA (where PagedAttention is itself on by default) gives a standard-layout head_dim-128 model TurboQuant K4/V3 KV with nothing asked for โ€” and **silently drops prefix caching**, which no TurboQuant preset supports. Pass `--pa-cache-type auto`, or `PagedCacheType::Auto` as above, to opt out. > -> Off that envelope โ€” MLA layouts, any other head_dim โ€” the default falls back to `Auto` with a warning, while an explicit `--pa-cache-type turboquant` is a hard error instead. The eager (non-paged) KV path is separate and genuinely opt-in via `ARC_TURBOQUANT_KV=1`. +> Off that envelope the default falls back to `Auto` with a warning: MLA layouts, uninstantiated head dims, and the instantiated-but-unmeasured widths 64/256/512. Those three are reachable only by asking for them โ€” an explicit `--pa-cache-type turboquant` is accepted at any instantiated width, and is a hard error off that set. The eager (non-paged) KV path is separate and genuinely opt-in via `ARC_TURBOQUANT_KV=1`. > -> **Provenance:** this path has served a real model on real hardware โ€” Qwen3-32B, B200, 55 tok/s, correct output โ€” but at b=1, one width, one preset, with **no quality evaluation and no A/B against an unquantized cache**. See [Results](#results) before relying on it. +> **Provenance:** this path has served a real model on real hardware โ€” Qwen3-32B, B200, 55 tok/s, correct output โ€” but at b=1, **head_dim 128 only**, one preset, with **no quality evaluation and no A/B against an unquantized cache**. Widths 64/256/512 compile; none has ever executed. See [Results](#results) before relying on it. ## Supported Models diff --git a/arc-tools/wave63_turboquant_hd512_gate.sh b/arc-tools/wave63_turboquant_hd512_gate.sh new file mode 100755 index 000000000..2b6284415 --- /dev/null +++ b/arc-tools/wave63_turboquant_hd512_gate.sh @@ -0,0 +1,153 @@ +#!/usr/bin/env bash +# wave63 โ€” TurboQuant CUDA kernels at head_dim {64,128,256,512}, Hopper + Blackwell. +# +# THE ONE THING THIS PROVES +# `turbo_paged_attention.cu` compiles for sm_90a AND sm_100a/sm_103a with all +# four head-dim instantiations live, and the built archive actually contains +# both cubins. That is the whole D16 claim in PR #94; everything else here is +# secondary. +# +# EXIT CODES (D18 rule 2 โ€” the distinction the runner acts on) +# 0 pass +# 1 genuine failure. Real signal: the code is wrong. Act on this. +# 2 this machine could not answer. Missing preflight, failing box, no +# toolkit, no repo. Says nothing about the code โ€” do not read it as a +# regression. +# +# RUN (from main, per D15 โ€” agents never call runcrate): +# setsid nohup bash /root/arc/arc-tools/wave63_turboquant_hd512_gate.sh \ +# < /dev/null > /root/wave63_gate.log 2>&1 & +# tail -f /root/wave63_gate.log ; cat /root/wave63_gate.status +# +# The preflight is looked up at a branch-independent path FIRST, because it +# lives on wave61/box-preflight-shared-prefix and vanishes the moment a runner +# checks out any other branch. Override with ARC_PREFLIGHT=/path. +# Set SKIP_NATIVE_BUILD=1 to go straight to the multi-arch proof. +set -uo pipefail + +STATUS=${STATUS:-/root/wave63_gate.status} +REPO=${REPO:-/root/arc} +BRANCH=${BRANCH:-feat/turboquant-hd512-v4} + +export LD_LIBRARY_PATH=/usr/local/cuda/compat:${LD_LIBRARY_PATH:-} +export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda-13.1} +export PATH=$CUDA_HOME/bin:$PATH + +: > "$STATUS" +say() { echo "[$(date -u +%H:%M:%S)] $*"; echo "$*" >> "$STATUS"; } +# Exit 2 โ€” the box could not answer. Not a verdict on the code. +noans() { say "NOANSWER $*"; say "EXIT 2 (environment could not answer)"; exit 2; } +# Exit 1 โ€” the code is wrong. This is the only signal worth acting on. +fail() { say "FAIL $*"; say "EXIT 1 (genuine failure)"; exit 1; } + +# --------------------------------------------------------------- preflight +# gpu_box_preflight.sh ends with +# [ "$_arc_pf_sourced" = "1" ] || exit 1 +# so when it is SOURCED on a *failing* box that test succeeds, short-circuits +# the `|| exit 1`, and becomes the last command โ€” making `source pf || handler` +# return 0 and the handler never fire. We therefore ignore its return value +# entirely and read `_ARC_PF_FAILED` directly, which is the flag it actually +# sets. If that flag is missing the script did not reach its own verdict, and +# an unverified box is a box that cannot answer. +PF="" +for cand in \ + "${ARC_PREFLIGHT:-}" \ + /root/arc-tools/gpu_box_preflight.sh \ + /usr/local/lib/arc/gpu_box_preflight.sh \ + "$REPO/arc-tools/gpu_box_preflight.sh" +do + [ -n "$cand" ] && [ -f "$cand" ] && { PF="$cand"; break; } +done +[ -n "$PF" ] || noans "no gpu_box_preflight.sh found. It lives on \ +wave61/box-preflight-shared-prefix, so it disappears whenever a runner checks \ +out another branch โ€” stage it at /root/arc-tools/ (branch-independent) or set \ +ARC_PREFLIGHT=/path. Nothing about the code was tested." + +say "STEP 0 preflight ($PF)" +PF_FLAG=$(bash -c 'set +e; . "$1" >&2; echo "${_ARC_PF_FAILED:-unset}"' _ "$PF") +case "$PF_FLAG" in + 0) say "PASS preflight" ;; + 1) noans "preflight reports this box is unhealthy (_ARC_PF_FAILED=1). \ +See FAIL[...] lines above. The code was never exercised." ;; + *) noans "preflight never set _ARC_PF_FAILED (got '$PF_FLAG') โ€” it did not \ +reach its own verdict, so its silence cannot be read as a pass." ;; +esac + +# ------------------------------------------------------------------- repo +cd "$REPO" 2>/dev/null || noans "no repo at $REPO" +command -v nvcc >/dev/null || noans "nvcc not on PATH (CUDA_HOME=$CUDA_HOME)" +git fetch --all -q 2>/dev/null || noans "cannot reach the git remote" +git checkout -q "$BRANCH" 2>/dev/null || noans "cannot check out $BRANCH" +git pull -q --ff-only 2>/dev/null || true +say "HEAD $(git rev-parse --short HEAD) on $BRANCH" +say "nvcc $(nvcc --version | grep -o 'release [0-9.]*' | head -1)" + +# ============================================================================ +# STEP 1 โ€” THE D16 PROOF. Runs first and needs no full workspace build. +# ============================================================================ +say "STEP 1 compile paged-attn for sm_90a + sm_100a + sm_103a" +if ! ARC_CUDA_ARCHS=90,100,103 cargo build --release \ + --features cuda -p mistralrs-paged-attn 2>&1 | tail -60; then + fail "multi-arch build of mistralrs-paged-attn failed โ€” see nvcc errors above" +fi +say "PASS multi-arch compile" + +# The build succeeding is not the claim; the cubins being present is. +LIB=$(find target/release/build -name 'libmistralrspagedattention.a' 2>/dev/null | head -1) +[ -n "$LIB" ] || noans "built, but no libmistralrspagedattention.a found to inspect" +command -v cuobjdump >/dev/null || noans "cuobjdump missing โ€” cannot verify which \ +arches landed, and the unverified build is exactly the assertion we came to replace" + +ARCHES=$(cuobjdump --list-elf "$LIB" 2>/dev/null | grep -oE 'sm_[0-9]+a?' | sort -u | tr '\n' ' ') +say "cubin arches in archive: ${ARCHES:-}" +MISSING="" +for want in sm_90a sm_100a; do + case " $ARCHES " in *" $want "*) ;; *) MISSING="$MISSING $want" ;; esac +done +[ -z "$MISSING" ] || fail "D16 unmet: archive is missing$MISSING (has: $ARCHES). \ +The kernels compiled but this binary would PTX-JIT or fail on the missing arch." +say "PASS D16 โ€” Hopper and Blackwell cubins both present" + +# ============================================================================ +# STEP 2 โ€” native full build. PR #98 already built these kernels clean for +# sm_90 in 8m53s, so this is confirmation, not discovery. Skippable. +# ============================================================================ +if [ "${SKIP_NATIVE_BUILD:-0}" = "1" ]; then + say "SKIP step 2 (SKIP_NATIVE_BUILD=1)" +else + say "STEP 2 native build (cuda flash-attn; never cudnn โ€” -62% decode on V4)" + cargo build --release --features "cuda flash-attn" 2>&1 | tail -40 \ + || fail "native build failed" + say "PASS native build" +fi + +# ============================================================================ +# STEP 3 โ€” format tests. CPU-only by nature: these pin the CUDA constant tables +# to the Rust generator. Not a GPU claim, and not offered as one (D14). +# ============================================================================ +say "STEP 3 turboquant format tests" +cargo test --release -p mistralrs-quant turboquant:: 2>&1 | tail -20 \ + || fail "turboquant format tests failed" +say "PASS format tests" + +# ============================================================================ +# STEP 4 โ€” decode no-regression. V4 does NOT use PagedAttention, so this does +# not exercise the TurboQuant kernel at all; it only shows this branch did not +# disturb the shipping path. +# ============================================================================ +UQFF=${UQFF:-/root/models/v4-uqff/qtip2b-0.uqff} +if [ "${SKIP_NATIVE_BUILD:-0}" = "1" ]; then + say "SKIP step 4 (no native binary was built)" +elif [ -f "$UQFF" ]; then + say "STEP 4 V4 decode smoke (no-regression only)" + timeout 900 ./target/release/mistralrs bench \ + -m /root/models/v4-src --uqff "$UQFF" \ + --n-gen 64 --n-prompt 64 --repetitions 2 2>&1 | tail -30 + say "STEP 4 done โ€” compare tok/s against the FACTS.md V4 decode baseline" +else + say "SKIP step 4 โ€” no UQFF at $UQFF" +fi + +say "ALL STEPS COMPLETE" +say "EXIT 0 (pass)" +exit 0 diff --git a/mistralrs-cli/src/args/paged_attn.rs b/mistralrs-cli/src/args/paged_attn.rs index 9e057abf8..d52afa8b3 100644 --- a/mistralrs-cli/src/args/paged_attn.rs +++ b/mistralrs-cli/src/args/paged_attn.rs @@ -59,15 +59,18 @@ pub struct PagedAttentionOptions { /// this flag off gives you TurboQuant KV and, with it, no prefix caching. /// Pass `auto` to opt out. Every other geometry โ€” MLA layouts and every /// other head_dim โ€” auto-falls back to `auto` with a warning, so those - /// models are unaffected. Setting the flag explicitly turns that fallback - /// into a hard error instead. + /// models are unaffected. Setting the flag explicitly widens the accepted + /// set to every instantiated head_dim (64/128/256/512) and turns that + /// fallback into a hard error instead. /// Provenance: TurboQuant K4/V3 has one end-to-end serving run behind it โ€” /// Qwen3-32B on a B200, 55 tok/s with correct output (2026-04-06). That run /// was b=1, head_dim 128, on Blackwell, and it did not isolate TurboQuant /// from the rest of the decode path. There is no A/B against an /// unquantized cache and no quality evaluation at any preset, so treat - /// quality as unestablished; `turboquant-3` and `turboquant-aggressive` - /// have never run at all. + /// quality as unestablished; 64/256/512 are compiled but have never + /// executed on hardware, which is why the default will not choose them for + /// you, and `turboquant-3` and `turboquant-aggressive` have never run at + /// all. // Hidden: overriding Arc's resolved KV cache type is the clearest way to // silently make a deployment worse โ€” every TurboQuant variant currently // disables prefix caching, and an explicit choice converts the safe diff --git a/mistralrs-core/src/paged_attention/cache_engine.rs b/mistralrs-core/src/paged_attention/cache_engine.rs index 5cc88040a..b34adb1b2 100644 --- a/mistralrs-core/src/paged_attention/cache_engine.rs +++ b/mistralrs-core/src/paged_attention/cache_engine.rs @@ -8,10 +8,13 @@ use serde::{Deserialize, Serialize}; use super::config::{KvCacheLayout, ModelConfigLike}; -/// Head dimension the TurboQuant CUDA kernels are specialized for. The packed -/// K4/V3 block layout and every `turbo_*` kernel assume exactly this head -/// size; the kernels exit early (`if (hs != 128) return;`) for anything else. -pub const TURBOQUANT_HEAD_DIM: usize = 128; +/// Head dimensions the TurboQuant CUDA kernels are instantiated for. +/// +/// Re-exported from `mistralrs_quant::turboquant::cuda_tables`, which is also +/// where the unit test lives that pins this list to the `case` arms of the +/// kernel's own dispatch. Keeping one list on both sides of the FFI is what +/// makes an unsupported width a refusal instead of an untouched output buffer. +pub use mistralrs_quant::turboquant::TURBOQUANT_CUDA_HEAD_DIMS; #[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, Default)] #[cfg_attr(feature = "pyo3_macros", pyo3::pyclass(eq, eq_int))] @@ -100,20 +103,53 @@ impl PagedCacheType { )) } + /// Head dims the **ambient default** is allowed to route onto TurboQuant. + /// + /// Deliberately narrower than [`TURBOQUANT_CUDA_HEAD_DIMS`]. Compiling a + /// kernel instantiation proves it *builds*; it does not prove it produces + /// correct KV on hardware, and none of 64/256/512 has ever executed. Since + /// [`PagedCacheType::TurboQuant`] is `#[default]`, widening the gate alone + /// would have silently moved every standard-layout model at those widths + /// onto unproven kernels โ€” and silently dropped prefix caching with it + /// ([`Self::supports_prefix_cache`] is false for every TurboQuant variant). + /// Two regressions, neither visible at the call site. + /// + /// This repo has paid for that exact shape once already: FP8 KV shipped + /// default-on and unmeasured (wave43-BU) and every V4 request died. + /// + /// So: **compiled** widens what you may ask for; **measured** widens what + /// you get without asking. Move a width here when its hardware gate has + /// passed, not when its kernel compiles. + const TURBOQUANT_DEFAULT_HEAD_DIMS: [usize; 1] = [128]; + /// Whether the TurboQuant kernels support this model's KV geometry. - fn turboquant_supports_model(config: &dyn ModelConfigLike) -> bool { + /// + /// K and V are compressed independently, so they only have to be widths the + /// kernels are instantiated for โ€” they no longer have to be *the same* + /// width, and no longer have to be 128. + /// + /// `forced` selects which set applies: an explicit request may use any + /// instantiated width, the ambient default only a measured one. + fn turboquant_supports_model(config: &dyn ModelConfigLike, forced: bool) -> bool { + let accepts = |head_dim: usize| { + if forced { + mistralrs_quant::turboquant::cuda_supports_head_dim(head_dim) + } else { + Self::TURBOQUANT_DEFAULT_HEAD_DIMS.contains(&head_dim) + } + }; matches!(config.kv_cache_layout(), KvCacheLayout::Standard) - && config.k_head_dim() == TURBOQUANT_HEAD_DIM - && config.v_head_dim() == TURBOQUANT_HEAD_DIM + && accepts(config.k_head_dim()) + && accepts(config.v_head_dim()) } /// Resolve this cache type against the target model's KV geometry. /// - /// The TurboQuant kernels only support the standard KV layout with - /// `head_dim == 128`: - /// * any other head size makes the `turbo_paged_attention.cu` kernels exit - /// early with their (uninitialized) F32 output buffer untouched โ€” - /// silent garbage; + /// The TurboQuant kernels support the standard KV layout at any head + /// dimension in [`TURBOQUANT_CUDA_HEAD_DIMS`]: + /// * a head size with no instantiation falls through the kernel's dispatch + /// switch, leaving the (uninitialized) output buffer untouched โ€” silent + /// garbage, which is why this gate runs before any allocation; /// * MLA-layout models (DeepSeek V2/V3, GLM4-MoE-lite) write through /// `concat_and_cache_mla`, which bails on a packed U8 cache at runtime. /// @@ -128,7 +164,11 @@ impl PagedCacheType { config: &dyn ModelConfigLike, explicitly_requested: bool, ) -> anyhow::Result { - if !self.is_turboquant() || Self::turboquant_supports_model(config) { + // `TurboQuant3`/`TurboQuantAggressive` are never a default anywhere, so + // selecting one is always an explicit act. Computed before the support + // check because it decides *which* set of head dims applies. + let forced = explicitly_requested || !matches!(self, PagedCacheType::TurboQuant); + if !self.is_turboquant() || Self::turboquant_supports_model(config, forced) { return Ok(self); } let reason = match config.kv_cache_layout() { @@ -136,13 +176,31 @@ impl PagedCacheType { "the model uses an MLA KV cache layout, which TurboQuant does not support" .to_string() } + // Split deliberately: "no kernel exists" and "a kernel exists but has + // never run on hardware, so the default will not choose it for you" + // are different facts, and collapsing them would tell a user to stop + // trying when the real answer is "opt in explicitly". + KvCacheLayout::Standard + if mistralrs_quant::turboquant::cuda_supports_head_dim(config.k_head_dim()) + && mistralrs_quant::turboquant::cuda_supports_head_dim(config.v_head_dim()) => + { + format!( + "the model has head_dim k={}/v={}: the TurboQuant kernels are instantiated \ + for it, but only head_dim in {:?} has been measured on hardware, so the \ + default will not route onto it. Pass `--pa-cache-type turboquant` to opt in \ + deliberately", + config.k_head_dim(), + config.v_head_dim(), + Self::TURBOQUANT_DEFAULT_HEAD_DIMS, + ) + } KvCacheLayout::Standard => format!( - "the model has head_dim k={}/v={}, but the TurboQuant kernels only support head_dim={TURBOQUANT_HEAD_DIM}", + "the model has head_dim k={}/v={}, but the TurboQuant kernels are instantiated \ + for head_dim in {TURBOQUANT_CUDA_HEAD_DIMS:?}", config.k_head_dim(), config.v_head_dim(), ), }; - let forced = explicitly_requested || !matches!(self, PagedCacheType::TurboQuant); if forced { anyhow::bail!( "PagedAttention cache type {self:?} was explicitly requested, but {reason}. \ @@ -569,23 +627,132 @@ mod cache_type_tests { PagedCacheType::TurboQuantAggressive, ]; - /// Supported geometry (standard layout, head_dim 128): every TurboQuant - /// preset resolves to itself, explicit or not. + /// Supported geometry, **explicitly requested**: every TurboQuant preset + /// resolves to itself at any head dim the kernels are instantiated for. + /// + /// Driven off `TURBOQUANT_CUDA_HEAD_DIMS` rather than a literal, so adding + /// a kernel instantiation extends this test automatically instead of + /// leaving the new width silently uncovered. #[test] - fn turboquant_stays_on_supported_geometry() { - let cfg = meta(128, 128, KvCacheLayout::Standard); - for t in TURBO_TYPES { - for explicit in [false, true] { - assert_eq!(t.resolve_for_model(&cfg, explicit).unwrap(), t); + fn explicit_turboquant_stays_on_any_instantiated_geometry() { + for head_dim in TURBOQUANT_CUDA_HEAD_DIMS { + let cfg = meta(head_dim, head_dim, KvCacheLayout::Standard); + for t in TURBO_TYPES { + assert_eq!( + t.resolve_for_model(&cfg, true).unwrap(), + t, + "head_dim={head_dim} is instantiated and an explicit request must be kept" + ); } } } - /// The default TurboQuant type falls back to Auto (instead of producing - /// silent garbage from the head_dim-128-only kernels) for other head dims. + /// **The regression this PR would otherwise have shipped.** + /// + /// `PagedCacheType::TurboQuant` is `#[default]`. Widening the acceptance + /// gate from `head_dim == 128` to every instantiated width would have moved + /// standard-layout models at 64/256/512 onto kernels that have never + /// executed, and silently dropped prefix caching with them โ€” without anyone + /// asking for TurboQuant at all. + /// + /// The ambient default must therefore keep falling back to `Auto` at every + /// instantiated-but-unmeasured width, while the same width stays available + /// on explicit request (covered above). + #[test] + fn default_turboquant_does_not_route_onto_unmeasured_head_dims() { + // Non-vacuity guard. This loop skips every width that IS a default, so + // widening `TURBOQUANT_DEFAULT_HEAD_DIMS` to the full instantiated set + // would empty the loop and make the test pass while asserting nothing โ€” + // which is precisely the silent-success shape this test exists to catch. + // Caught by mutating the const during review; keep this assertion. + let unmeasured: Vec = TURBOQUANT_CUDA_HEAD_DIMS + .into_iter() + .filter(|hd| !PagedCacheType::TURBOQUANT_DEFAULT_HEAD_DIMS.contains(hd)) + .collect(); + assert!( + !unmeasured.is_empty(), + "every instantiated head dim is now a default width, so this test asserts \ + nothing. If a hardware gate genuinely measured them all, delete this test \ + deliberately โ€” do not let it pass vacuously." + ); + + for head_dim in unmeasured { + let cfg = meta(head_dim, head_dim, KvCacheLayout::Standard); + assert_eq!( + PagedCacheType::TurboQuant + .resolve_for_model(&cfg, false) + .unwrap(), + PagedCacheType::Auto, + "head_dim={head_dim} is instantiated but unmeasured; the DEFAULT must not \ + route onto it" + ); + } + } + + /// Every measured width must still be taken by the ambient default โ€” + /// the narrowing must not become "TurboQuant is off by accident". + #[test] + fn default_turboquant_still_taken_on_measured_head_dims() { + for head_dim in PagedCacheType::TURBOQUANT_DEFAULT_HEAD_DIMS { + assert!( + TURBOQUANT_CUDA_HEAD_DIMS.contains(&head_dim), + "head_dim={head_dim} is a default width but has no kernel instantiation" + ); + let cfg = meta(head_dim, head_dim, KvCacheLayout::Standard); + assert_eq!( + PagedCacheType::TurboQuant + .resolve_for_model(&cfg, false) + .unwrap(), + PagedCacheType::TurboQuant, + "head_dim={head_dim} is measured and must still be the default" + ); + } + } + + /// DeepSeek-V4's width specifically. It is the reason the 128-only kernel + /// limit was lifted, so it gets its own named assertion โ€” but on the + /// *explicit* path, because 512 has not been measured. + #[test] + fn turboquant_accepts_v4_head_dim_512_when_asked_for() { + assert!( + TURBOQUANT_CUDA_HEAD_DIMS.contains(&512), + "512 must stay in the instantiated set" + ); + let cfg = meta(512, 512, KvCacheLayout::Standard); + assert_eq!( + PagedCacheType::TurboQuant + .resolve_for_model(&cfg, true) + .unwrap(), + PagedCacheType::TurboQuant + ); + // ...and is NOT taken by the ambient default. + assert_eq!( + PagedCacheType::TurboQuant + .resolve_for_model(&cfg, false) + .unwrap(), + PagedCacheType::Auto + ); + } + + /// K and V are compressed independently, so mixed widths are fine as long + /// as each is instantiated โ€” on the explicit path. + #[test] + fn turboquant_accepts_mixed_but_instantiated_head_dims() { + let cfg = meta(512, 128, KvCacheLayout::Standard); + assert_eq!( + PagedCacheType::TurboQuant + .resolve_for_model(&cfg, true) + .unwrap(), + PagedCacheType::TurboQuant + ); + } + + /// The default TurboQuant type falls back to Auto (instead of falling + /// through the kernel's dispatch switch and leaving the output buffer + /// untouched) for head dims with no instantiation. #[test] fn default_turboquant_falls_back_on_unsupported_head_dim() { - for head_dim in [64, 96, 192, 256] { + for head_dim in [96, 192, 320, 1024] { let cfg = meta(head_dim, head_dim, KvCacheLayout::Standard); assert_eq!( PagedCacheType::TurboQuant @@ -595,8 +762,8 @@ mod cache_type_tests { "head_dim={head_dim} must fall back" ); } - // Mixed k/v head dims must also fall back. - let cfg = meta(128, 64, KvCacheLayout::Standard); + // A mixed pair still falls back when either half is uninstantiated. + let cfg = meta(128, 192, KvCacheLayout::Standard); assert_eq!( PagedCacheType::TurboQuant .resolve_for_model(&cfg, false) @@ -629,7 +796,7 @@ mod cache_type_tests { /// models rather than silently switching caches. #[test] fn explicit_turboquant_errors_on_unsupported_model() { - let cfg = meta(64, 64, KvCacheLayout::Standard); + let cfg = meta(192, 192, KvCacheLayout::Standard); for t in TURBO_TYPES { let err = t.resolve_for_model(&cfg, true).unwrap_err(); assert!( diff --git a/mistralrs-core/src/paged_attention/mod.rs b/mistralrs-core/src/paged_attention/mod.rs index 803d84ab4..f13c1a1ee 100644 --- a/mistralrs-core/src/paged_attention/mod.rs +++ b/mistralrs-core/src/paged_attention/mod.rs @@ -131,7 +131,15 @@ mod cache_config_tests { use super::*; use candle_core::Device; - fn small_head_dim_model() -> ModelConfigMetadata { + /// A head dim the TurboQuant kernels have no instantiation for. 192 is a + /// real head dim (it is what the MLA models report) and is deliberately + /// *not* in `TURBOQUANT_CUDA_HEAD_DIMS`, so it exercises the fallback + /// without pretending an unrealistic geometry. + fn uninstantiated_head_dim_model() -> ModelConfigMetadata { + assert!( + !mistralrs_quant::turboquant::cuda_supports_head_dim(192), + "192 gained a kernel instantiation; pick another width for this fixture" + ); ModelConfigMetadata { max_seq_len: 4096, num_layers: 2, @@ -139,15 +147,16 @@ mod cache_config_tests { num_kv_heads: 2, num_attn_heads: 4, sliding_window: None, - k_head_dim: 64, - v_head_dim: 64, + k_head_dim: 192, + v_head_dim: 192, kv_cache_layout: KvCacheLayout::Standard, } } - /// A defaulted TurboQuant cache type on a head_dim-64 model must fall back - /// to `Auto` at cache-config construction time โ€” the returned config (used - /// for both sizing and `CacheEngine` allocation) carries the resolved type. + /// A defaulted TurboQuant cache type on a model whose head dim has no + /// kernel instantiation must fall back to `Auto` at cache-config + /// construction time โ€” the returned config (used for both sizing and + /// `CacheEngine` allocation) carries the resolved type. #[test] fn calculate_cache_config_falls_back_for_unsupported_turboquant() -> anyhow::Result<()> { let device = Device::Cpu; @@ -157,7 +166,7 @@ mod cache_config_tests { DType::F16, PagedCacheType::TurboQuant, false, - &small_head_dim_model(), + &uninstantiated_head_dim_model(), &device, &[Some(device.clone())], true, @@ -179,7 +188,7 @@ mod cache_config_tests { DType::F16, PagedCacheType::TurboQuant, true, - &small_head_dim_model(), + &uninstantiated_head_dim_model(), &device, &[Some(device.clone())], true, @@ -307,9 +316,11 @@ pub fn calculate_cache_config( if !SUPPORTED_BLOCK_SIZE.contains(&block_size) { anyhow::bail!("Block size must be in {SUPPORTED_BLOCK_SIZE:?}, got {block_size}"); } - // TurboQuant only supports standard-layout models with head_dim == 128; - // resolve to a supported cache type (or error, if the user explicitly - // forced TurboQuant) before any sizing or allocation happens. + // TurboQuant needs a standard KV layout, and a head dim its kernels are + // instantiated for โ€” narrowed further to the *measured* widths when the + // type is the ambient default rather than an explicit request. Resolve to a + // supported cache type (or error, if the user explicitly forced TurboQuant) + // before any sizing or allocation happens. let cache_type = cache_type.resolve_for_model(config, cache_type_explicit)?; let dtype = cache_type.to_dtype(dtype); let dtype_size = dtype.size_in_bytes(); diff --git a/mistralrs-paged-attn/build.rs b/mistralrs-paged-attn/build.rs index 764785806..b495dc003 100644 --- a/mistralrs-paged-attn/build.rs +++ b/mistralrs-paged-attn/build.rs @@ -13,6 +13,11 @@ fn main() -> Result<()> { println!("cargo:rerun-if-changed=build.rs"); println!("cargo:rerun-if-changed=src/cuda/pagedattention.cuh"); println!("cargo:rerun-if-changed=src/cuda/turbo_paged_attention.cuh"); + // Every other .cu here is listed; this one was not, and because emitting + // *any* rerun-if-changed disables cargo's watch-the-whole-package default, + // edits to it were relying entirely on cudaforge's object cache to trigger + // a rebuild. + println!("cargo:rerun-if-changed=src/cuda/turbo_paged_attention.cu"); println!("cargo:rerun-if-changed=src/cuda/copy_blocks_kernel.cu"); println!("cargo:rerun-if-changed=src/cuda/reshape_and_cache_kernel.cu"); println!("cargo:rerun-if-changed=src/cuda/concat_and_cache_mla_kernel.cu"); @@ -52,7 +57,49 @@ fn main() -> Result<()> { .arg("--compiler-options") .arg("-fPIC"); + // D16: ship cubins for every arch we target, not just the build box's GPU. + // + // cudaforge derives a single `-gencode` from `CUDA_COMPUTE_CAP` or the local + // `nvidia-smi`, so a binary built on an H200 carries SM90 only and PTX-JITs + // (or fails) on Blackwell. `ARC_CUDA_ARCHS` appends further targets to the + // fat binary โ€” e.g. `ARC_CUDA_ARCHS=90,100,103` for Hopper + Blackwell. + // + // Left unset the build is unchanged, because each extra arch recompiles + // every .cu in this crate and that cost is not worth paying on a rented box + // that only ever runs one arch. Release and publish builds set it. + println!("cargo:rerun-if-env-changed=ARC_CUDA_ARCHS"); + + // Hoisted above the `ARC_CUDA_ARCHS` loop because it decides which of the + // requested arches are actually *extra*. let compute_cap = builder.get_compute_cap().unwrap_or(80); + + if let Ok(archs) = std::env::var("ARC_CUDA_ARCHS") { + for arch in archs.split(',').map(str::trim).filter(|a| !a.is_empty()) { + // cudaforge already emits one `-gencode` of its own, derived from + // `CUDA_COMPUTE_CAP` or the local `nvidia-smi`, and suffixes it `a` + // for cap >= 90 exactly as we do below (`GpuArch::auto_suffix` -> + // `to_gencode_arg`). It then appends `extra_args` verbatim, so + // re-requesting that arch here hands nvcc the byte-identical flag + // twice: on an H200 building `ARC_CUDA_ARCHS=90,100,103` the + // duplicate is `-gencode=arch=compute_90a,code=sm_90a`. + // + // Emit only what cudaforge does not already cover. + if arch.parse::().is_ok_and(|c| c == compute_cap) { + continue; + } + // Arch-specific ("a") cubins: sm_90a/sm_100a expose the Hopper and + // Blackwell-only instructions the arch-specialised kernel paths are + // written against. Matches cudaforge's own suffixing for cap >= 90. + let suffix = if arch.parse::().is_ok_and(|c| c >= 90) { + "a" + } else { + "" + }; + builder = builder.arg(&format!( + "-gencode=arch=compute_{arch}{suffix},code=sm_{arch}{suffix}" + )); + } + } // Enable FP8 if compute capability >= 8.0 (Ampere and newer) let using_fp8 = if compute_cap >= 80 { builder = builder.arg("-DENABLE_FP8"); diff --git a/mistralrs-paged-attn/src/cuda/backend/paged_attention.rs b/mistralrs-paged-attn/src/cuda/backend/paged_attention.rs index 14865d2f8..f147d2be3 100644 --- a/mistralrs-paged-attn/src/cuda/backend/paged_attention.rs +++ b/mistralrs-paged-attn/src/cuda/backend/paged_attention.rs @@ -712,10 +712,11 @@ pub fn turbo_reshape_and_cache( let dev = k.device(); let (num_tokens, num_heads, head_size) = k_l.shape().dims3()?; - if head_size != 128 { + if !turbo_supports_head_dim(head_size) { candle::bail!( - "turbo_reshape_and_cache only supports head_size=128 (got {head_size}); \ - the TurboQuant kernel would exit without writing anything" + "turbo_reshape_and_cache has no kernel instantiation for head_size={head_size} \ + (have {TURBO_HEAD_DIMS:?}); the launch would fall through the dispatch switch \ + and write nothing" ); } @@ -770,6 +771,21 @@ pub fn turbo_reshape_and_cache( Ok(()) } +/// Head dimensions `turbo_paged_attention.cu` instantiates its kernels for. +/// +/// Must stay equal to `TURBOQUANT_CUDA_HEAD_DIMS` in +/// `mistralrs-quant::turboquant::cuda_tables`, which is where the test lives +/// that parses this declaration and the kernel's dispatch switch and fails if +/// the three disagree. A width missing here is merely refused; a width present +/// here but missing from the switch would launch nothing and leave the caller +/// reading uninitialized memory. +pub const TURBO_HEAD_DIMS: [usize; 4] = [64, 128, 256, 512]; + +/// Whether the TurboQuant kernels have an instantiation for `head_dim`. +pub fn turbo_supports_head_dim(head_dim: usize) -> bool { + TURBO_HEAD_DIMS.contains(&head_dim) +} + /// TurboQuant: paged attention over compressed U8 KV cache. #[allow(clippy::too_many_arguments)] #[allow(unused_variables)] @@ -827,10 +843,11 @@ pub fn turbo_paged_attention( let dev = q_s.device(); let (num_seqs, num_heads, head_size) = q_l.shape().dims3()?; - if head_size != 128 { + if !turbo_supports_head_dim(head_size) { candle::bail!( - "turbo_paged_attention only supports head_size=128 (got {head_size}); \ - the kernel would exit early and return an uninitialized output buffer" + "turbo_paged_attention has no kernel instantiation for head_size={head_size} \ + (have {TURBO_HEAD_DIMS:?}); the launch would fall through the dispatch switch \ + and return an uninitialized output buffer" ); } let max_num_blocks_per_seq = bt_l.shape().dims2()?.1; diff --git a/mistralrs-paged-attn/src/cuda/turbo_paged_attention.cu b/mistralrs-paged-attn/src/cuda/turbo_paged_attention.cu index e4a24348f..e6702b033 100644 --- a/mistralrs-paged-attn/src/cuda/turbo_paged_attention.cu +++ b/mistralrs-paged-attn/src/cuda/turbo_paged_attention.cu @@ -21,126 +21,92 @@ __device__ __forceinline__ float tq_to_f(__nv_bfloat16 x) { return __bfloat162fl #define TQ_WARP 32 #define TQ_DIVUP(a,b) (((a)+(b)-1)/(b)) -static __constant__ float CB4[16] = { - -0.237664013127f,-0.180836062501f,-0.141805261760f,-0.110288414632f, - -0.082828489390f,-0.057772320256f,-0.034151583096f,-0.011302500645f, - 0.011302500645f, 0.034151583096f, 0.057772320256f, 0.082828489390f, - 0.110288414632f, 0.141805261760f, 0.180836062501f, 0.237664013127f, -}; -static __constant__ float BD4[17] = { - -1.0f,-0.209250037814f,-0.161320662130f,-0.126046838196f, - -0.096558452011f,-0.070300404823f,-0.045961951676f,-0.022727041871f, - 0.0f, 0.022727041871f, 0.045961951676f, 0.070300404823f, - 0.096558452011f, 0.126046838196f, 0.161320662130f, 0.209250037814f, 1.0f, -}; -static __constant__ float CB3[8] = { - -0.188397319183f,-0.118139828402f,-0.066585638471f,-0.021604320011f, - 0.021604320011f, 0.066585638471f, 0.118139828402f, 0.188397319183f, -}; -static __constant__ float BD3[9] = { - -1.0f,-0.153268573792f,-0.092362733436f,-0.044094979241f, - 0.0f, 0.044094979241f, 0.092362733436f, 0.153268573792f, 1.0f, -}; -static __constant__ float SGN[128] = { - -1,-1, 1,-1, 1,-1,-1, 1,-1,-1, 1, 1,-1,-1,-1,-1, - -1, 1,-1,-1, 1, 1, 1, 1, 1, 1,-1,-1,-1, 1, 1,-1, - -1,-1, 1,-1, 1,-1,-1,-1, 1,-1,-1,-1, 1,-1,-1,-1, - -1,-1, 1,-1,-1,-1, 1,-1,-1,-1,-1,-1, 1,-1, 1, 1, - -1,-1,-1, 1,-1,-1, 1, 1,-1,-1,-1,-1, 1,-1, 1,-1, - 1, 1, 1,-1, 1, 1, 1, 1, 1, 1, 1,-1,-1,-1,-1,-1, - -1,-1,-1,-1, 1, 1, 1,-1,-1,-1, 1, 1, 1,-1,-1,-1, - -1,-1,-1,-1, 1, 1,-1, 1, 1,-1,-1, 1, 1,-1, 1,-1, -}; - -// WHT that works with any number of active threads (uses only __syncthreads) -__device__ void wht128(float* d, int tid, int nthreads) { - for (int h = 1; h < 128; h *= 2) { - __syncthreads(); - for (int i = tid; i < 64; i += nthreads) { - int bs = (i/h)*(h*2), off = i%h; - float a = d[bs+off], b = d[bs+off+h]; - d[bs+off] = a+b; d[bs+off+h] = a-b; - } - } - __syncthreads(); -} - -__device__ void rotate128(float* d, int tid, int nthreads) { - if (tid < 128) d[tid] *= SGN[tid]; - __syncthreads(); - wht128(d, tid, nthreads); - if (tid < 128) d[tid] *= 0.08838834764831845f; - __syncthreads(); - if (tid < 128) d[tid] *= SGN[tid]; - __syncthreads(); -} - -__device__ uint8_t q4(float x) { - if (x<=BD4[1]) return 0; if (x>=BD4[16]) return 15; - int lo=1,hi=16; while(lo>1;if(x=BD3[8]) return 7; - int lo=1,hi=8; while(lo>1;if(x +// Per-head-vector compression into the paged cache. +// +// One block per (token, head); NT threads cooperate on one HS-wide vector. +// HS is a template parameter so the shared buffers, the rotation order and the +// codebook are all fixed at compile time. At HS == NT == 128 the emitted work +// is the same as the original hand-written 128-only kernel. +template __global__ void tq_cache_k( const T* in, uint8_t* cache, __half* norms, const int64_t* slots, - int nh, int hs, int bs, int in_stride, int cbs, int chs, int nbs, int nhs + int nh, int bs, int in_stride, int cbs, int chs, int nbs, int nhs ) { - int vid = blockIdx.x, tid = threadIdx.x; - int tok = vid/nh, head = vid%nh; - int64_t slot = slots[tok]; if (slot<0) return; - int bi = slot/bs, bo = slot%bs; - __shared__ float s[128]; __shared__ uint8_t ix[128]; - if (tid0;r>>=1){if(tid0;r>>=1){ if(tid1e-10f)s[tid]/=nm; __syncthreads(); - rotate128(s,tid,128); - if(tid1e-10f){ for(int i=tid;i(s,tid); + for(int i=tid;i(s[i]); + __syncthreads(); + + // 4-bit nibble pack, x=16 swizzled so a decode lane's run stays contiguous. + for(int i=tid;i +template __global__ void tq_cache_v( const T* in, uint8_t* cache, __half* norms, const int64_t* slots, - int nh, int hs, int bs, int in_stride, int vbs, int vhs, int nbs, int nhs + int nh, int bs, int in_stride, int vbs, int vhs, int nbs, int nhs ) { - int vid = blockIdx.x, tid = threadIdx.x; - int tok = vid/nh, head = vid%nh; - int64_t slot = slots[tok]; if (slot<0) return; - int bi = slot/bs, bo = slot%bs; - __shared__ float s[128]; __shared__ uint8_t ix[128]; - if (tid0;r>>=1){if(tid0;r>>=1){ if(tid1e-10f)s[tid]/=nm; __syncthreads(); - rotate128(s,tid,128); - if(tid1e-10f){ for(int i=tid;i(s,tid); + for(int i=tid;i(s[i]); + __syncthreads(); + + // 3-bit pack, 10 indices per 32-bit word, bytes strided by block size. + constexpr int NG = (HS+9)/10; + for(int i=tid;i>8); cache[bb+2*bs]=(uint8_t)(w>>16); cache[bb+3*bs]=(uint8_t)(w>>24); } @@ -152,76 +118,141 @@ __global__ void tq_cache_v( // phase 2 = end of norm computation // phase 3 = end of rotate+quantize // phase 4 = end of pack+write -template +template __global__ void tq_cache_k_clocked( const T* in, uint8_t* cache, __half* norms, const int64_t* slots, unsigned long long* clocks, - int nh, int hs, int bs, int in_stride, int cbs, int chs, int nbs, int nhs + int nh, int bs, int in_stride, int cbs, int chs, int nbs, int nhs ) { - int vid = blockIdx.x, tid = threadIdx.x; + constexpr int NT = 128; + const int vid = blockIdx.x, tid = threadIdx.x; if (tid == 0) clocks[vid * 5 + 0] = clock64(); - int tok = vid/nh, head = vid%nh; - int64_t slot = slots[tok]; if (slot<0) { if (tid==0) clocks[vid*5+4]=clock64(); return; } - int bi = slot/bs, bo = slot%bs; - __shared__ float s[128]; __shared__ uint8_t ix[128]; - if (tid0;r>>=1){if(tid0;r>>=1){ if(tid1e-10f)s[tid]/=nm; __syncthreads(); + if(nm>1e-10f){ for(int i=tid;i(s,tid); + for(int i=tid;i(s[i]); + __syncthreads(); if (tid == 0) clocks[vid * 5 + 3] = clock64(); - if(tid +template __global__ void tq_cache_v_clocked( const T* in, uint8_t* cache, __half* norms, const int64_t* slots, unsigned long long* clocks, - int nh, int hs, int bs, int in_stride, int vbs, int vhs, int nbs, int nhs + int nh, int bs, int in_stride, int vbs, int vhs, int nbs, int nhs ) { - int vid = blockIdx.x, tid = threadIdx.x; + constexpr int NT = 128; + const int vid = blockIdx.x, tid = threadIdx.x; if (tid == 0) clocks[vid * 5 + 0] = clock64(); - int tok = vid/nh, head = vid%nh; - int64_t slot = slots[tok]; if (slot<0) { if (tid==0) clocks[vid*5+4]=clock64(); return; } - int bi = slot/bs, bo = slot%bs; - __shared__ float s[128]; __shared__ uint8_t ix[128]; - if (tid0;r>>=1){if(tid0;r>>=1){ if(tid1e-10f)s[tid]/=nm; __syncthreads(); + if(nm>1e-10f){ for(int i=tid;i(s,tid); + for(int i=tid;i(s[i]); + __syncthreads(); if (tid == 0) clocks[vid * 5 + 3] = clock64(); - int ng=(hs+9)/10; - if(tid>8); cache[bb+2*bs]=(uint8_t)(w>>16); cache[bb+3*bs]=(uint8_t)(w>>24); } if (tid == 0) clocks[vid * 5 + 4] = clock64(); } +// Head dimensions the kernels are instantiated for. `TURBOQUANT_CUDA_HEAD_DIMS` +// in mistralrs-quant mirrors this list and a unit test pins the two together, +// so a head dim can never reach a launcher that has no instantiation for it. +template +static void tq_launch_cache( + const void* key, const void* value, void* kc, void* vc, void* kn, void* vn, + const int64_t* slots, int nt, int nh, int bs, + int ks, int vs, int kbs, int khs, int nbs, int nhs, cudaStream_t stream +) { + dim3 g(nt*nh), b(128); + constexpr int VPD = ((HS+9)/10)*4; + const int vvbs = nh*VPD*bs, vvhs = VPD*bs; + tq_cache_k<<>>((const T*)key,(uint8_t*)kc,(__half*)kn,slots,nh,bs,ks,kbs,khs,nbs,nhs); + tq_cache_v<<>>((const T*)value,(uint8_t*)vc,(__half*)vn,slots,nh,bs,vs,vvbs,vvhs,nbs,nhs); +} + +template +static void tq_launch_cache_clocked( + const void* key, const void* value, void* kc, void* vc, void* kn, void* vn, + const int64_t* slots, void* clocks_k, void* clocks_v, int nt, int nh, int bs, + int ks, int vs, int kbs, int khs, int nbs, int nhs, cudaStream_t stream +) { + dim3 g(nt*nh), b(128); + constexpr int VPD = ((HS+9)/10)*4; + const int vvbs = nh*VPD*bs, vvhs = VPD*bs; + tq_cache_k_clocked<<>>((const T*)key,(uint8_t*)kc,(__half*)kn,slots,(unsigned long long*)clocks_k,nh,bs,ks,kbs,khs,nbs,nhs); + tq_cache_v_clocked<<>>((const T*)value,(uint8_t*)vc,(__half*)vn,slots,(unsigned long long*)clocks_v,nh,bs,vs,vvbs,vvhs,nbs,nhs); +} + +// dtype: 1 = BF16 input, anything else = F16. +#define TQ_CACHE_BY_DTYPE(FN, HS_, ...) \ + do { \ + if (dtype == 1) FN(__VA_ARGS__); \ + else FN(__VA_ARGS__); \ + } while (0) + +#define TQ_CACHE_BY_HEAD_DIM(FN, ...) \ + switch (hs) { \ + case 64: TQ_CACHE_BY_DTYPE(FN, 64, __VA_ARGS__); break; \ + case 128: TQ_CACHE_BY_DTYPE(FN, 128, __VA_ARGS__); break; \ + case 256: TQ_CACHE_BY_DTYPE(FN, 256, __VA_ARGS__); break; \ + case 512: TQ_CACHE_BY_DTYPE(FN, 512, __VA_ARGS__); break; \ + default: break; \ + } + extern "C" void turbo_reshape_and_cache_clocked( const void* key, const void* value, void* kc, void* vc, void* kn, void* vn, @@ -231,16 +262,9 @@ extern "C" void turbo_reshape_and_cache_clocked( int ks, int vs, int kbs, int khs, int nbs, int nhs, cudaStream_t stream, uint32_t dtype ) { - if(hs!=128)return; - dim3 g(nt*nh),b(128); - int vpd=(hs+9)/10*4; int vvbs=nh*vpd*bs; int vvhs=vpd*bs; - if (dtype == 1) { - tq_cache_k_clocked<<>>((const __nv_bfloat16*)key,(uint8_t*)kc,(__half*)kn,slots,(unsigned long long*)clocks_k,nh,hs,bs,ks,kbs,khs,nbs,nhs); - tq_cache_v_clocked<<>>((const __nv_bfloat16*)value,(uint8_t*)vc,(__half*)vn,slots,(unsigned long long*)clocks_v,nh,hs,bs,vs,vvbs,vvhs,nbs,nhs); - } else { - tq_cache_k_clocked<<>>((const __half*)key,(uint8_t*)kc,(__half*)kn,slots,(unsigned long long*)clocks_k,nh,hs,bs,ks,kbs,khs,nbs,nhs); - tq_cache_v_clocked<<>>((const __half*)value,(uint8_t*)vc,(__half*)vn,slots,(unsigned long long*)clocks_v,nh,hs,bs,vs,vvbs,vvhs,nbs,nhs); - } + TQ_CACHE_BY_HEAD_DIM(tq_launch_cache_clocked, + key, value, kc, vc, kn, vn, slots, clocks_k, clocks_v, + nt, nh, bs, ks, vs, kbs, khs, nbs, nhs, stream); } extern "C" void turbo_reshape_and_cache( @@ -251,16 +275,9 @@ extern "C" void turbo_reshape_and_cache( int ks, int vs, int kbs, int khs, int nbs, int nhs, cudaStream_t stream, uint32_t dtype ) { - if(hs!=128)return; - dim3 g(nt*nh),b(128); - int vpd=(hs+9)/10*4; int vvbs=nh*vpd*bs; int vvhs=vpd*bs; - if (dtype == 1) { // BF16 - tq_cache_k<<>>((const __nv_bfloat16*)key,(uint8_t*)kc,(__half*)kn,slots,nh,hs,bs,ks,kbs,khs,nbs,nhs); - tq_cache_v<<>>((const __nv_bfloat16*)value,(uint8_t*)vc,(__half*)vn,slots,nh,hs,bs,vs,vvbs,vvhs,nbs,nhs); - } else { // F16 - tq_cache_k<<>>((const __half*)key,(uint8_t*)kc,(__half*)kn,slots,nh,hs,bs,ks,kbs,khs,nbs,nhs); - tq_cache_v<<>>((const __half*)value,(uint8_t*)vc,(__half*)vn,slots,nh,hs,bs,vs,vvbs,vvhs,nbs,nhs); - } + TQ_CACHE_BY_HEAD_DIM(tq_launch_cache, + key, value, kc, vc, kn, vn, slots, + nt, nh, bs, ks, vs, kbs, khs, nbs, nhs, stream); } // ============================================================================ @@ -277,7 +294,13 @@ extern "C" void turbo_reshape_and_cache( // OutT = float (default) or __nv_bfloat16 (saves a separate F32โ†’BF16 cast kernel // in the dedicated decode path). -template +// +// HS is the head dimension. Everything that depends on it โ€” the rotation order, +// the sign diagonal, the codebooks, the packed stride, how many output dims each +// thread owns โ€” is resolved at compile time. NT stays 128 for every HS so the +// warp-reduction topology (4 warps, `red[8]`) is unchanged; wider heads give +// each thread more dims rather than adding threads. +template __global__ void tq_attn( OutT* __restrict__ out, const QT* __restrict__ query, @@ -290,10 +313,10 @@ __global__ void tq_attn( int nkvh, int mbps, int nh, float scale, int kbs, int khs, int vbs, int vhs, int nbs, int nhs ) { - constexpr int HS = 128; constexpr int NT = 128; constexpr int NW = NT / TQ_WARP; - constexpr int EPT = HS / TQ_WARP; // 4 elements per thread for dot product + constexpr int BPL = TQ_BYTES_PER_LANE(HS); // packed K bytes per lane + constexpr int DPT = (HS + NT - 1) / NT; // output dims per thread const int hidx = blockIdx.x, sidx = blockIdx.y; const int tid = threadIdx.x; @@ -302,17 +325,19 @@ __global__ void tq_attn( if (clen == 0) return; const int kvh = hidx / (nh / nkvh); - // 1. Load + rotate Q + const float* cb4 = tq_cb4(); + const float* cb3 = tq_cb3(); + + // 1. Load + rotate Q. Attention runs in the rotated domain, so this is the + // only rotation the QยทK half needs. __shared__ float qr[HS]; - if (tid < HS) qr[tid] = tq_to_f(query[sidx*nh*HS + hidx*HS + tid]); + for (int i = tid; i < HS; i += NT) qr[i] = tq_to_f(query[sidx*nh*HS + hidx*HS + i]); __syncthreads(); - rotate128(qr, tid, NT); - // After this: qr is rotated, all threads synced + tq_rotate(qr, tid); // 2. QยทK โ€” compute logits extern __shared__ char shmem[]; float* logits = (float*)shmem; - // Initialize logits to 0 for (int i = tid; i < (int)clen; i += NT) logits[i] = 0.f; __syncthreads(); @@ -320,40 +345,52 @@ __global__ void tq_attn( const uint32_t* sbt = bt + sidx * mbps; const int nblocks = TQ_DIVUP(clen, BLOCK_SIZE); - // Each warp processes blocks in parallel. - // All 32 lanes collaborate on ONE token at a time. - // Each lane reads 2 packed bytes = 4 elements of the 128-dim K vector. - // Warp shuffle reduces the 32 partial sums into one dot product. + // Each warp takes whole blocks; all 32 lanes collaborate on one token at a + // time, each owning BPL packed bytes (2*BPL coordinates). A warp shuffle + // folds the 32 partial dot products into the logit. for (int bi = warp; bi < nblocks; bi += NW) { - int pb = sbt[bi]; - int tib = min(BLOCK_SIZE, (int)clen - bi*BLOCK_SIZE); + const int pb = sbt[bi]; + const int tib = min(BLOCK_SIZE, (int)clen - bi*BLOCK_SIZE); + if (tib <= 0) continue; - for (int t = 0; t < tib; t++) { - int tpos = bi * BLOCK_SIZE + t; + // This lane's packed run: contiguous, and inside one x=16 group. + const int byidx0 = lane * BPL; + const uint8_t* kbase = kc + (long long)pb*kbs + (long long)kvh*khs + + (byidx0 / 16) * BLOCK_SIZE * 16 + (byidx0 % 16); + + uint64_t cur = tq_load_klane(kbase); - // Each lane reads its 2 packed bytes (4 elements) + for (int t = 0; t < tib; t++) { +#if TQ_PREFETCH > 1 + // Hopper/Blackwell: stage the next token's bytes while reducing this + // one. Same arithmetic, one fewer stall. + const uint64_t nxt = (t + 1 < tib) + ? tq_load_klane(kbase + (t + 1) * 16) : 0ull; +#endif float qk = 0.f; #pragma unroll - for (int e = 0; e < EPT/2; e++) { - int byidx = lane * (EPT/2) + e; // 0..63 - int x = 16; - int koff = pb*kbs + kvh*khs + (byidx/x)*BLOCK_SIZE*x + t*x + (byidx%x); - uint8_t pk = kc[koff]; - int d0 = byidx*2, d1 = byidx*2+1; - qk += qr[d0] * CB4[pk & 0xF]; - qk += qr[d1] * CB4[(pk>>4) & 0xF]; + for (int j = 0; j < BPL; j++) { + const uint8_t pk = (uint8_t)((cur >> (8 * j)) & 0xFFull); + const int d0 = (byidx0 + j) * 2; + qk += qr[d0] * cb4[pk & 0xF]; + qk += qr[d0 + 1] * cb4[(pk >> 4) & 0xF]; } - // Warp reduction: sum 32 partial dot products #pragma unroll for (int mask = TQ_WARP/2; mask > 0; mask >>= 1) qk += __shfl_xor_sync(0xffffffff, qk, mask); - float knorm = __half2float(kn[pb*nbs + kvh*nhs + t]); + const float knorm = __half2float(kn[pb*nbs + kvh*nhs + t]); qk *= knorm * scale; - if (lane == 0) logits[tpos] = qk; + if (lane == 0) logits[bi*BLOCK_SIZE + t] = qk; qk_max = fmaxf(qk_max, qk); + +#if TQ_PREFETCH > 1 + cur = nxt; +#else + if (t + 1 < tib) cur = tq_load_klane(kbase + (t + 1) * 16); +#endif } } @@ -361,7 +398,7 @@ __global__ void tq_attn( #pragma unroll for (int mask = TQ_WARP/2; mask >= 1; mask /= 2) qk_max = fmaxf(qk_max, __shfl_xor_sync(0xffffffff, qk_max, mask)); - __shared__ float red[8]; + __shared__ float red[2 * NW]; if (lane == 0) red[warp] = qk_max; __syncthreads(); qk_max = (lane < NW) ? red[lane] : -FLT_MAX; @@ -373,7 +410,7 @@ __global__ void tq_attn( // 3. Softmax float esum = 0.f; for (int i = tid; i < (int)clen; i += NT) { - float v = __expf(logits[i] - qk_max); + const float v = __expf(logits[i] - qk_max); logits[i] = v; esum += v; } @@ -386,53 +423,54 @@ __global__ void tq_attn( #pragma unroll for (int mask = NW/2; mask >= 1; mask /= 2) esum += __shfl_xor_sync(0xffffffff, esum, mask); - float inv = __fdividef(1.f, __shfl_sync(0xffffffff, esum, 0) + 1e-6f); + const float inv = __fdividef(1.f, __shfl_sync(0xffffffff, esum, 0) + 1e-6f); for (int i = tid; i < (int)clen; i += NT) logits[i] *= inv; __syncthreads(); - // 4. V accumulation - // Each thread accumulates EPT=4 output dimensions - // thread tid handles dims: for QยทK we used lane-based, for V we use tid-based - // tid handles dims [tid*EPT .. tid*EPT+3] (but tid goes 0..127, EPT=4 would be 512 dims) - // Actually: 128 threads, 128 dims, each thread handles 1 dim - // Simpler: each thread handles dim = tid (one dim per thread) - - float acc = 0.f; - if (tid < HS) { - int dim = tid; - int group = dim / 10; - int pos = dim % 10; - - for (int bi = 0; bi < nblocks; bi++) { - int pb = sbt[bi]; - int tib = min(BLOCK_SIZE, (int)clen - bi*BLOCK_SIZE); - - for (int t = 0; t < tib; t++) { - float w = logits[bi*BLOCK_SIZE + t]; - if (w < 1e-8f) continue; - - int bb = pb*vbs + kvh*vhs + group*4*BLOCK_SIZE + t; - uint32_t word = (uint32_t)vc[bb] | - ((uint32_t)vc[bb+BLOCK_SIZE]<<8) | - ((uint32_t)vc[bb+2*BLOCK_SIZE]<<16) | - ((uint32_t)vc[bb+3*BLOCK_SIZE]<<24); - uint8_t vi = (word>>(pos*3))&7; - float vnorm = __half2float(vn[pb*nbs+kvh*nhs+t]); - acc += w * CB3[vi] * vnorm; + // 4. V accumulation โ€” each thread owns DPT output dims, strided by NT. + float acc[DPT]; + #pragma unroll + for (int r = 0; r < DPT; r++) acc[r] = 0.f; + + for (int bi = 0; bi < nblocks; bi++) { + const int pb = sbt[bi]; + const int tib = min(BLOCK_SIZE, (int)clen - bi*BLOCK_SIZE); + + for (int t = 0; t < tib; t++) { + const float w = logits[bi*BLOCK_SIZE + t]; + if (w < 1e-8f) continue; + const float vnorm = __half2float(vn[pb*nbs + kvh*nhs + t]); + + #pragma unroll + for (int r = 0; r < DPT; r++) { + const int dim = tid + r * NT; + if (dim >= HS) continue; + const int group = dim / 10, pos = dim % 10; + const long long bb = (long long)pb*vbs + (long long)kvh*vhs + group*4*BLOCK_SIZE + t; + const uint32_t word = (uint32_t)vc[bb] + | ((uint32_t)vc[bb + BLOCK_SIZE] << 8) + | ((uint32_t)vc[bb + 2*BLOCK_SIZE] << 16) + | ((uint32_t)vc[bb + 3*BLOCK_SIZE] << 24); + const uint8_t vi = (word >> (pos * 3)) & 7; + acc[r] += w * cb3[vi] * vnorm; } } } - // 5. Inverse rotation + output + // 5. Inverse rotation + output. V was quantized in the rotated basis, so the + // accumulator lands there too; DยทHยทD is its own inverse. __shared__ float obuf[HS]; - if (tid < HS) obuf[tid] = acc; + #pragma unroll + for (int r = 0; r < DPT; r++) { + const int dim = tid + r * NT; + if (dim < HS) obuf[dim] = acc[r]; + } __syncthreads(); - rotate128(obuf, tid, NT); + tq_rotate(obuf, tid); - if (tid < HS) { - // Cast on store: float (no-op), __half, or __nv_bfloat16. - const int oi = sidx*nh*HS + hidx*HS + tid; - const float v = obuf[tid]; + for (int i = tid; i < HS; i += NT) { + const int oi = sidx*nh*HS + hidx*HS + i; + const float v = obuf[i]; if constexpr (std::is_same::value) { out[oi] = __float2bfloat16(v); } else if constexpr (std::is_same::value) { @@ -450,7 +488,7 @@ __global__ void tq_attn( // phase 3 = after softmax // phase 4 = after V accumulate // phase 5 = after inverse rotate + write -template +template __global__ void tq_attn_clocked( __nv_bfloat16* __restrict__ out, const QT* __restrict__ query, @@ -464,10 +502,12 @@ __global__ void tq_attn_clocked( int nkvh, int mbps, int nh, float scale, int kbs, int khs, int vbs, int vhs, int nbs, int nhs ) { - constexpr int HS = 128; constexpr int NT = 128; constexpr int NW = NT / TQ_WARP; - constexpr int EPT = HS / TQ_WARP; + constexpr int BPL = TQ_BYTES_PER_LANE(HS); + constexpr int DPT = (HS + NT - 1) / NT; + const float* cb4 = tq_cb4(); + const float* cb3 = tq_cb3(); const int hidx = blockIdx.x, sidx = blockIdx.y; const int tid = threadIdx.x; @@ -483,9 +523,9 @@ __global__ void tq_attn_clocked( const int kvh = hidx / (nh / nkvh); __shared__ float qr[HS]; - if (tid < HS) qr[tid] = tq_to_f(query[sidx*nh*HS + hidx*HS + tid]); + for (int i = tid; i < HS; i += NT) qr[i] = tq_to_f(query[sidx*nh*HS + hidx*HS + i]); __syncthreads(); - rotate128(qr, tid, NT); + tq_rotate(qr, tid); if (tid == 0) clocks[block_id * 6 + 1] = clock64(); extern __shared__ char shmem[]; @@ -496,29 +536,45 @@ __global__ void tq_attn_clocked( float qk_max = -FLT_MAX; const uint32_t* sbt = bt + sidx * mbps; const int nblocks = TQ_DIVUP(clen, BLOCK_SIZE); + // This loop must stay a mirror of tq_attn's: the whole point of the clocked + // variant is that its phase-2 stamps describe the kernel that actually + // runs. A scalar gather here against a vectorized one there would have made + // the timings describe a kernel nobody launches. for (int bi = warp; bi < nblocks; bi += NW) { int pb = sbt[bi]; int tib = min(BLOCK_SIZE, (int)clen - bi*BLOCK_SIZE); + if (tib <= 0) continue; + + const int byidx0 = lane * BPL; + const uint8_t* kbase = kc + (long long)pb*kbs + (long long)kvh*khs + + (byidx0 / 16) * BLOCK_SIZE * 16 + (byidx0 % 16); + uint64_t cur = tq_load_klane(kbase); + for (int t = 0; t < tib; t++) { - int tpos = bi * BLOCK_SIZE + t; +#if TQ_PREFETCH > 1 + const uint64_t nxt = (t + 1 < tib) + ? tq_load_klane(kbase + (t + 1) * 16) : 0ull; +#endif float qk = 0.f; #pragma unroll - for (int e = 0; e < EPT/2; e++) { - int byidx = lane * (EPT/2) + e; - int x = 16; - int koff = pb*kbs + kvh*khs + (byidx/x)*BLOCK_SIZE*x + t*x + (byidx%x); - uint8_t pk = kc[koff]; - int d0 = byidx*2, d1 = byidx*2+1; - qk += qr[d0] * CB4[pk & 0xF]; - qk += qr[d1] * CB4[(pk>>4) & 0xF]; + for (int j = 0; j < BPL; j++) { + const uint8_t pk = (uint8_t)((cur >> (8 * j)) & 0xFFull); + const int d0 = (byidx0 + j) * 2; + qk += qr[d0] * cb4[pk & 0xF]; + qk += qr[d0 + 1] * cb4[(pk >> 4) & 0xF]; } #pragma unroll for (int mask = TQ_WARP/2; mask > 0; mask >>= 1) qk += __shfl_xor_sync(0xffffffff, qk, mask); float knorm = __half2float(kn[pb*nbs + kvh*nhs + t]); qk *= knorm * scale; - if (lane == 0) logits[tpos] = qk; + if (lane == 0) logits[bi * BLOCK_SIZE + t] = qk; qk_max = fmaxf(qk_max, qk); +#if TQ_PREFETCH > 1 + cur = nxt; +#else + if (t + 1 < tib) cur = tq_load_klane(kbase + (t + 1) * 16); +#endif } } if (tid == 0) clocks[block_id * 6 + 2] = clock64(); @@ -526,7 +582,7 @@ __global__ void tq_attn_clocked( #pragma unroll for (int mask = TQ_WARP/2; mask >= 1; mask /= 2) qk_max = fmaxf(qk_max, __shfl_xor_sync(0xffffffff, qk_max, mask)); - __shared__ float red[8]; + __shared__ float red[2 * NW]; if (lane == 0) red[warp] = qk_max; __syncthreads(); qk_max = (lane < NW) ? red[lane] : -FLT_MAX; @@ -555,40 +611,157 @@ __global__ void tq_attn_clocked( __syncthreads(); if (tid == 0) clocks[block_id * 6 + 3] = clock64(); - float acc = 0.f; - if (tid < HS) { - int dim = tid; - int group = dim / 10; - int pos = dim % 10; - for (int bi = 0; bi < nblocks; bi++) { - int pb = sbt[bi]; - int tib = min(BLOCK_SIZE, (int)clen - bi*BLOCK_SIZE); - for (int t = 0; t < tib; t++) { - float w = logits[bi*BLOCK_SIZE + t]; - if (w < 1e-8f) continue; - int bb = pb*vbs + kvh*vhs + group*4*BLOCK_SIZE + t; - uint32_t word = (uint32_t)vc[bb] | + float acc[DPT]; + #pragma unroll + for (int r = 0; r < DPT; r++) acc[r] = 0.f; + for (int bi = 0; bi < nblocks; bi++) { + int pb = sbt[bi]; + int tib = min(BLOCK_SIZE, (int)clen - bi*BLOCK_SIZE); + for (int t = 0; t < tib; t++) { + float w = logits[bi*BLOCK_SIZE + t]; + if (w < 1e-8f) continue; + float vnorm = __half2float(vn[pb*nbs+kvh*nhs+t]); + #pragma unroll + for (int r = 0; r < DPT; r++) { + const int dim = tid + r * NT; + if (dim >= HS) continue; + const int group = dim / 10, pos = dim % 10; + const long long bb = (long long)pb*vbs + (long long)kvh*vhs + group*4*BLOCK_SIZE + t; + const uint32_t word = (uint32_t)vc[bb] | ((uint32_t)vc[bb+BLOCK_SIZE]<<8) | ((uint32_t)vc[bb+2*BLOCK_SIZE]<<16) | ((uint32_t)vc[bb+3*BLOCK_SIZE]<<24); - uint8_t vi = (word>>(pos*3))&7; - float vnorm = __half2float(vn[pb*nbs+kvh*nhs+t]); - acc += w * CB3[vi] * vnorm; + const uint8_t vi = (word>>(pos*3))&7; + acc[r] += w * cb3[vi] * vnorm; } } } if (tid == 0) clocks[block_id * 6 + 4] = clock64(); __shared__ float obuf[HS]; - if (tid < HS) obuf[tid] = acc; + #pragma unroll + for (int r = 0; r < DPT; r++) { + const int dim = tid + r * NT; + if (dim < HS) obuf[dim] = acc[r]; + } __syncthreads(); - rotate128(obuf, tid, NT); - if (tid < HS) { - out[sidx*nh*HS + hidx*HS + tid] = __float2bfloat16(obuf[tid]); + tq_rotate(obuf, tid); + for (int i = tid; i < HS; i += NT) { + out[sidx*nh*HS + hidx*HS + i] = __float2bfloat16(obuf[i]); } if (tid == 0) clocks[block_id * 6 + 5] = clock64(); } +// ============================================================================ +// Launch wrappers +// ============================================================================ + +/// Launch one instantiation, opting in to the arch's real shared-memory budget +/// first. +/// +/// The logits array is dynamic shared memory sized by the padded context +/// length, so it grows without bound as context grows. CUDA caps dynamic shared +/// memory at 48 KB per block unless a kernel explicitly opts in โ€” 164 KB on +/// Ampere, 227 KB on Hopper (SM90) and Blackwell (SM100/SM103). Without this +/// call, any context past roughly 12k tokens fails to launch. `granted` is a +/// per-instantiation static, so the driver call happens once per kernel rather +/// than once per decode step. +template +static void tq_launch_attn( + void* out, const void* query, const void* kc, const void* vc, + const void* kn, const void* vn, const uint32_t* bt, const uint32_t* cl, + int nkvh, int mbps, int nh, float scale, + int kbs, int khs, int vbs, int vhs, int nbs, int nhs, + dim3 grid, dim3 block, int smem, cudaStream_t stream +) { + // Not cached: the requestable maximum differs by architecture, so a single + // process-wide flag would be wrong the moment two devices differ. A failure + // here means the request exceeds the device's opt-in cap; the launch below + // then fails with cudaErrorInvalidValue, which surfaces on the next sync. + if (smem > 48 * 1024) { + (void)cudaFuncSetAttribute( + (const void*)(tq_attn), + cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + } + tq_attn<<>>( + (OutT*)out, (const QT*)query, (const uint8_t*)kc, (const uint8_t*)vc, + (const __half*)kn, (const __half*)vn, bt, cl, + nkvh, mbps, nh, scale, kbs, khs, vbs, vhs, nbs, nhs); +} + +template +static void tq_launch_attn_clocked( + void* out_bf16, const void* query, const void* kc, const void* vc, + const void* kn, const void* vn, const uint32_t* bt, const uint32_t* cl, + void* clocks, int nkvh, int mbps, int nh, float scale, + int kbs, int khs, int vbs, int vhs, int nbs, int nhs, + dim3 grid, dim3 block, int smem, cudaStream_t stream +) { + // Not cached: the requestable maximum differs by architecture, so a single + // process-wide flag would be wrong the moment two devices differ. A failure + // here means the request exceeds the device's opt-in cap; the launch below + // then fails with cudaErrorInvalidValue, which surfaces on the next sync. + if (smem > 48 * 1024) { + (void)cudaFuncSetAttribute( + (const void*)(tq_attn_clocked), + cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + } + tq_attn_clocked<<>>( + (__nv_bfloat16*)out_bf16, (const QT*)query, + (const uint8_t*)kc, (const uint8_t*)vc, + (const __half*)kn, (const __half*)vn, bt, cl, + (unsigned long long*)clocks, + nkvh, mbps, nh, scale, kbs, khs, vbs, vhs, nbs, nhs); +} + +// Paged block sizes the cache engine emits. +#define TQ_ATTN_BY_BS(HS_, QT, OUTT, ...) \ + switch (bs) { \ + case 8: tq_launch_attn(__VA_ARGS__); break; \ + case 16: tq_launch_attn(__VA_ARGS__); break; \ + case 32: tq_launch_attn(__VA_ARGS__); break; \ + default: break; \ + } + +// Head dimensions with codebooks and sign diagonals in turbo_paged_attention.cuh. +// `TURBOQUANT_CUDA_HEAD_DIMS` in mistralrs-quant mirrors this set and a unit +// test pins the two together, so the Rust side refuses a head dim before it can +// reach a launcher with no instantiation for it. +#define TQ_ATTN_BY_HS(QT, OUTT, ...) \ + switch (hs) { \ + case 64: TQ_ATTN_BY_BS(64, QT, OUTT, __VA_ARGS__) break; \ + case 128: TQ_ATTN_BY_BS(128, QT, OUTT, __VA_ARGS__) break; \ + case 256: TQ_ATTN_BY_BS(256, QT, OUTT, __VA_ARGS__) break; \ + case 512: TQ_ATTN_BY_BS(512, QT, OUTT, __VA_ARGS__) break; \ + default: break; \ + } + +#define TQ_ATTN_CLOCKED_BY_BS(HS_, QT, ...) \ + switch (bs) { \ + case 8: tq_launch_attn_clocked(__VA_ARGS__); break; \ + case 16: tq_launch_attn_clocked(__VA_ARGS__); break; \ + case 32: tq_launch_attn_clocked(__VA_ARGS__); break; \ + default: break; \ + } + +#define TQ_ATTN_CLOCKED_BY_HS(QT, ...) \ + switch (hs) { \ + case 64: TQ_ATTN_CLOCKED_BY_BS(64, QT, __VA_ARGS__) break; \ + case 128: TQ_ATTN_CLOCKED_BY_BS(128, QT, __VA_ARGS__) break; \ + case 256: TQ_ATTN_CLOCKED_BY_BS(256, QT, __VA_ARGS__) break; \ + case 512: TQ_ATTN_CLOCKED_BY_BS(512, QT, __VA_ARGS__) break; \ + default: break; \ + } + +// Geometry shared by every entry point below. +#define TQ_ATTN_PROLOGUE() \ + const int vpd = (hs + 9) / 10 * 4; \ + const int vvbs = nkvh * vpd * bs, vvhs = vpd * bs; \ + const int padded = TQ_DIVUP(mcl, bs) * bs; \ + const int smem = padded * (int)sizeof(float); \ + dim3 grid(nh, ns, 1); \ + dim3 block(128); + extern "C" void turbo_paged_attention_v1_bf16out_clocked( void* out_bf16, const void* query, const void* kc, const void* vc, @@ -600,34 +773,17 @@ extern "C" void turbo_paged_attention_v1_bf16out_clocked( int mbps, int qs, int kbs, int khs, int nbs, int nhs, cudaStream_t stream, uint32_t qdtype ) { - if (hs != 128) return; - int vpd = (hs + 9) / 10 * 4; - int vvbs = nkvh * vpd * bs, vvhs = vpd * bs; - int padded = TQ_DIVUP(mcl, bs) * bs; - int smem = padded * sizeof(float); - dim3 grid(nh, ns, 1); - dim3 block(128); + TQ_ATTN_PROLOGUE() if (qdtype == 1) { - switch (bs) { - case 8: tq_attn_clocked<8, __nv_bfloat16><<>>( - (__nv_bfloat16*)out_bf16, (const __nv_bfloat16*)query, - (const uint8_t*)kc, (const uint8_t*)vc, - (const __half*)kn, (const __half*)vn, - bt, cl, (unsigned long long*)clocks, - nkvh, mbps, nh, scale, kbs, khs, vvbs, vvhs, nbs, nhs); break; - case 16: tq_attn_clocked<16, __nv_bfloat16><<>>( - (__nv_bfloat16*)out_bf16, (const __nv_bfloat16*)query, - (const uint8_t*)kc, (const uint8_t*)vc, - (const __half*)kn, (const __half*)vn, - bt, cl, (unsigned long long*)clocks, - nkvh, mbps, nh, scale, kbs, khs, vvbs, vvhs, nbs, nhs); break; - case 32: tq_attn_clocked<32, __nv_bfloat16><<>>( - (__nv_bfloat16*)out_bf16, (const __nv_bfloat16*)query, - (const uint8_t*)kc, (const uint8_t*)vc, - (const __half*)kn, (const __half*)vn, - bt, cl, (unsigned long long*)clocks, - nkvh, mbps, nh, scale, kbs, khs, vvbs, vvhs, nbs, nhs); break; - } + TQ_ATTN_CLOCKED_BY_HS(__nv_bfloat16, + out_bf16, query, kc, vc, kn, vn, bt, cl, clocks, + nkvh, mbps, nh, scale, kbs, khs, vvbs, vvhs, nbs, nhs, + grid, block, smem, stream) + } else { + TQ_ATTN_CLOCKED_BY_HS(__half, + out_bf16, query, kc, vc, kn, vn, bt, cl, clocks, + nkvh, mbps, nh, scale, kbs, khs, vvbs, vvhs, nbs, nhs, + grid, block, smem, stream) } } @@ -641,20 +797,11 @@ extern "C" void turbo_paged_attention_v1_f16( int mbps, int qs, int kbs, int khs, int nbs, int nhs, cudaStream_t stream ) { - if (hs != 128) return; - int vpd=(hs+9)/10*4; - int vvbs=nkvh*vpd*bs, vvhs=vpd*bs; - int padded = TQ_DIVUP(mcl,bs)*bs; - int smem = padded * sizeof(float); - // Ensure enough for output buffer too (128 floats = 512 bytes, always < logits) - dim3 grid(nh, ns, 1); - dim3 block(128); - - switch(bs){ - case 8: tq_attn<8><<>>((float*)out,(const __half*)query,(const uint8_t*)kc,(const uint8_t*)vc,(const __half*)kn,(const __half*)vn,bt,cl,nkvh,mbps,nh,scale,kbs,khs,vvbs,vvhs,nbs,nhs); break; - case 16: tq_attn<16><<>>((float*)out,(const __half*)query,(const uint8_t*)kc,(const uint8_t*)vc,(const __half*)kn,(const __half*)vn,bt,cl,nkvh,mbps,nh,scale,kbs,khs,vvbs,vvhs,nbs,nhs); break; - case 32: tq_attn<32><<>>((float*)out,(const __half*)query,(const uint8_t*)kc,(const uint8_t*)vc,(const __half*)kn,(const __half*)vn,bt,cl,nkvh,mbps,nh,scale,kbs,khs,vvbs,vvhs,nbs,nhs); break; - } + TQ_ATTN_PROLOGUE() + TQ_ATTN_BY_HS(__half, float, + out, query, kc, vc, kn, vn, bt, cl, + nkvh, mbps, nh, scale, kbs, khs, vvbs, vvhs, nbs, nhs, + grid, block, smem, stream) } // Dtype-generic entry point for the dedicated decode path @@ -668,22 +815,18 @@ extern "C" void turbo_paged_attention_v1( int mbps, int qs, int kbs, int khs, int nbs, int nhs, cudaStream_t stream, uint32_t dtype ) { - if (hs != 128) return; - int vpd=(hs+9)/10*4; - int vvbs=nkvh*vpd*bs, vvhs=vpd*bs; - int padded = TQ_DIVUP(mcl,bs)*bs; - int smem = padded * sizeof(float); - dim3 grid(nh, ns, 1); - dim3 block(128); - #define TQ_LAUNCH(QT) \ - switch(bs){ \ - case 8: tq_attn<8, QT><<>>((float*)out,(const QT*)query,(const uint8_t*)kc,(const uint8_t*)vc,(const __half*)kn,(const __half*)vn,bt,cl,nkvh,mbps,nh,scale,kbs,khs,vvbs,vvhs,nbs,nhs); break; \ - case 16: tq_attn<16, QT><<>>((float*)out,(const QT*)query,(const uint8_t*)kc,(const uint8_t*)vc,(const __half*)kn,(const __half*)vn,bt,cl,nkvh,mbps,nh,scale,kbs,khs,vvbs,vvhs,nbs,nhs); break; \ - case 32: tq_attn<32, QT><<>>((float*)out,(const QT*)query,(const uint8_t*)kc,(const uint8_t*)vc,(const __half*)kn,(const __half*)vn,bt,cl,nkvh,mbps,nh,scale,kbs,khs,vvbs,vvhs,nbs,nhs); break; \ - } - if (dtype == 1) { TQ_LAUNCH(__nv_bfloat16); } - else { TQ_LAUNCH(__half); } - #undef TQ_LAUNCH + TQ_ATTN_PROLOGUE() + if (dtype == 1) { + TQ_ATTN_BY_HS(__nv_bfloat16, float, + out, query, kc, vc, kn, vn, bt, cl, + nkvh, mbps, nh, scale, kbs, khs, vvbs, vvhs, nbs, nhs, + grid, block, smem, stream) + } else { + TQ_ATTN_BY_HS(__half, float, + out, query, kc, vc, kn, vn, bt, cl, + nkvh, mbps, nh, scale, kbs, khs, vvbs, vvhs, nbs, nhs, + grid, block, smem, stream) + } } // BF16-output entry point: writes BF16 directly, eliminating the F32โ†’BF16 cast @@ -698,20 +841,16 @@ extern "C" void turbo_paged_attention_v1_bf16out( int mbps, int qs, int kbs, int khs, int nbs, int nhs, cudaStream_t stream, uint32_t qdtype ) { - if (hs != 128) return; - int vpd=(hs+9)/10*4; - int vvbs=nkvh*vpd*bs, vvhs=vpd*bs; - int padded = TQ_DIVUP(mcl,bs)*bs; - int smem = padded * sizeof(float); - dim3 grid(nh, ns, 1); - dim3 block(128); - #define TQ_LAUNCH_BF16OUT(QT) \ - switch(bs){ \ - case 8: tq_attn<8, QT, __nv_bfloat16><<>>((__nv_bfloat16*)out_bf16,(const QT*)query,(const uint8_t*)kc,(const uint8_t*)vc,(const __half*)kn,(const __half*)vn,bt,cl,nkvh,mbps,nh,scale,kbs,khs,vvbs,vvhs,nbs,nhs); break; \ - case 16: tq_attn<16, QT, __nv_bfloat16><<>>((__nv_bfloat16*)out_bf16,(const QT*)query,(const uint8_t*)kc,(const uint8_t*)vc,(const __half*)kn,(const __half*)vn,bt,cl,nkvh,mbps,nh,scale,kbs,khs,vvbs,vvhs,nbs,nhs); break; \ - case 32: tq_attn<32, QT, __nv_bfloat16><<>>((__nv_bfloat16*)out_bf16,(const QT*)query,(const uint8_t*)kc,(const uint8_t*)vc,(const __half*)kn,(const __half*)vn,bt,cl,nkvh,mbps,nh,scale,kbs,khs,vvbs,vvhs,nbs,nhs); break; \ - } - if (qdtype == 1) { TQ_LAUNCH_BF16OUT(__nv_bfloat16); } - else { TQ_LAUNCH_BF16OUT(__half); } - #undef TQ_LAUNCH_BF16OUT + TQ_ATTN_PROLOGUE() + if (qdtype == 1) { + TQ_ATTN_BY_HS(__nv_bfloat16, __nv_bfloat16, + out_bf16, query, kc, vc, kn, vn, bt, cl, + nkvh, mbps, nh, scale, kbs, khs, vvbs, vvhs, nbs, nhs, + grid, block, smem, stream) + } else { + TQ_ATTN_BY_HS(__half, __nv_bfloat16, + out_bf16, query, kc, vc, kn, vn, bt, cl, + nkvh, mbps, nh, scale, kbs, khs, vvbs, vvhs, nbs, nhs, + grid, block, smem, stream) + } } diff --git a/mistralrs-paged-attn/src/cuda/turbo_paged_attention.cuh b/mistralrs-paged-attn/src/cuda/turbo_paged_attention.cuh index 7e037b615..1e8debb31 100644 --- a/mistralrs-paged-attn/src/cuda/turbo_paged_attention.cuh +++ b/mistralrs-paged-attn/src/cuda/turbo_paged_attention.cuh @@ -1,391 +1,379 @@ /** - * TurboQuant Paged Attention Kernel + * TurboQuant device tables and generic rotation โ€” shared by every TurboQuant + * CUDA kernel in this crate. * - * A variant of the standard paged attention kernel that reads compressed KV cache - * using TurboQuant codebook lookups instead of direct memory loads or FP8 dequantization. + * WHAT THIS FILE IS + * ----------------- + * TurboQuant compresses a KV head vector by (1) recording its L2 norm, (2) + * applying the involutory rotation DยทHยทD to the unit vector, and (3) scalar + * quantizing each rotated coordinate against a Lloyd-Max codebook. Attention + * then runs *in the rotated domain*: because DยทHยทD is orthogonal, qยทk is + * preserved under it, so the K dot product needs no inverse rotation at all. + * Only the V accumulation, which lands in the rotated basis, is rotated back + * once at the end. * - * Key differences from standard paged attention: - * 1. Q is pre-rotated via WHT before the QยทK loop (once per query) - * 2. K is stored as packed indices; lookup centroids from codebook in shared memory - * 3. K dot product operates in the rotated domain (no inverse rotation needed) - * 4. V is stored as packed indices; dequantized inline via codebook lookup - * 5. Norms are stored per-token and applied after the dot product + * DIMENSION SUPPORT + * ----------------- + * The tables below cover head_dim 64, 128, 256 and 512. Both `D` (the sign + * diagonal) and the codebook are functions of the dimension, so a kernel + * cannot simply reuse the 128 tables at another width โ€” the codebook centroids + * shrink as sqrt(1/d) because the rotated coordinates concentrate. * - * This kernel handles the 4-bit key / 3-bit value (default) configuration. - * The cache layout per block: - * k_cache: [num_blocks, num_kv_heads, packed_k_bytes_per_head, block_size] - * v_cache: [num_blocks, num_kv_heads, packed_v_bytes_per_head, block_size] - * k_norms: [num_blocks, num_kv_heads, block_size] (half) - * v_norms: [num_blocks, num_kv_heads, block_size] (half) + * PROVENANCE โ€” DO NOT HAND-EDIT THE TABLES + * ---------------------------------------- + * Every `__constant__` array below was emitted by + * cargo run -p mistralrs-quant --example emit_turboquant_cuda_tables -- 64 128 256 512 + * which calls the same `turboquant::wht::generate_signs` and + * `turboquant::codebook::get_codebook` that the Rust compressor uses. The + * values are therefore correct by construction rather than by transcription. * - * Reference: TurboQuant ICLR 2026 (arXiv:2504.19874), 0xSero/turboquant Triton kernels. + * `turboquant::cuda_tables` re-derives them at test time and diffs against this + * file, so a hand-edit here fails `cargo test -p mistralrs-quant`. That test is + * the only thing standing between a silent sign-table drift and a cache that + * decodes to noise, which is why it parses this header directly. + * + * ARCH SPECIALISATION (D16) + * ------------------------- + * `TQ_PREFETCH` selects the software-pipeline depth of the packed-K gather. + * Hopper and Blackwell have deeper memory latency to hide and enough registers + * to hide it, so they run a 2-deep pipeline; pre-Hopper parts run the flat + * loop. Both produce identical arithmetic โ€” this is a scheduling choice, not a + * numerical one. */ +#pragma once + #include #include #include -#ifndef WARP_SIZE -#define WARP_SIZE 32 +// --------------------------------------------------------------------------- +// Arch specialisation +// --------------------------------------------------------------------------- +// Hopper (SM90) and Blackwell (SM100/SM103) hide global latency better with a +// software-pipelined gather; older parts are latency-bound on the extra +// registers instead. Identical results either way. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +#define TQ_PREFETCH 2 +#else +#define TQ_PREFETCH 1 #endif -#define MAX_TQ(a, b) ((a) > (b) ? (a) : (b)) -#define MIN_TQ(a, b) ((a) < (b) ? (a) : (b)) -#define DIVIDE_ROUND_UP_TQ(a, b) (((a) + (b) - 1) / (b)) -// ============================================================================ -// Codebooks (same values as turbo_wht.cu) -// ============================================================================ +// --------------------------------------------------------------------------- +// Tables โ€” generated, see PROVENANCE above. +// --------------------------------------------------------------------------- -__constant__ float TQ_CODEBOOK_4BIT[16] = { - -0.237664013127f, -0.180836062501f, -0.141805261760f, -0.110288414632f, - -0.082828489390f, -0.057772320256f, -0.034151583096f, -0.011302500645f, - 0.011302500645f, 0.034151583096f, 0.057772320256f, 0.082828489390f, - 0.110288414632f, 0.141805261760f, 0.180836062501f, 0.237664013127f, +// ==== head_dim = 64 (seed 42) ==== +static __constant__ float TQ_SGN_64[64] = { + 1,-1, 1,-1, 1, 1,-1,-1,-1,-1, 1, 1, 1,-1, 1, 1, + 1,-1, 1,-1,-1,-1, 1, 1, 1,-1, 1,-1, 1, 1, 1, 1, + 1,-1,-1, 1,-1,-1, 1, 1,-1, 1, 1,-1,-1,-1, 1,-1, + -1,-1,-1, 1, 1, 1, 1,-1, 1,-1, 1,-1, 1, 1, 1, 1, +}; +static __constant__ float TQ_CB3_64[8] = { + -0.263914018869f,-0.166167899966f,-0.093832284212f,-0.030469184741f, + 0.030469184741f,0.093832284212f,0.166167899966f,0.263914018869f, +}; +static __constant__ float TQ_BD3_64[9] = { + -1.000000000000f,-0.215040951967f,-0.130000084639f,-0.062150731683f, + 0.000000000000f,0.062150731683f,0.130000084639f,0.215040951967f, + 1.000000000000f, +}; +static __constant__ float TQ_CB4_64[16] = { + -0.330796420574f,-0.252913773060f,-0.198856174946f,-0.154925554991f, + -0.116486772895f,-0.081311784685f,-0.048089791089f,-0.015919025987f, + 0.015919025987f,0.048089791089f,0.081311784685f,0.116486772895f, + 0.154925554991f,0.198856174946f,0.252913773060f,0.330796420574f, +}; +static __constant__ float TQ_BD4_64[17] = { + -1.000000000000f,-0.291855096817f,-0.225884988904f,-0.176890864968f, + -0.135706171393f,-0.098899275064f,-0.064700789750f,-0.032004408538f, + 0.000000000000f,0.032004408538f,0.064700789750f,0.098899275064f, + 0.135706171393f,0.176890864968f,0.225884988904f,0.291855096817f, + 1.000000000000f, }; -__constant__ float TQ_CODEBOOK_3BIT[8] = { - -0.188397319183f, -0.118139828402f, -0.066585638471f, -0.021604320011f, - 0.021604320011f, 0.066585638471f, 0.118139828402f, 0.188397319183f, +// ==== head_dim = 128 (seed 42) ==== +static __constant__ float TQ_SGN_128[128] = { + -1,-1, 1,-1, 1,-1,-1, 1,-1,-1, 1, 1,-1,-1,-1,-1, + -1, 1,-1,-1, 1, 1, 1, 1, 1, 1,-1,-1,-1, 1, 1,-1, + -1,-1, 1,-1, 1,-1,-1,-1, 1,-1,-1,-1, 1,-1,-1,-1, + -1,-1, 1,-1,-1,-1, 1,-1,-1,-1,-1,-1, 1,-1, 1, 1, + -1,-1,-1, 1,-1,-1, 1, 1,-1,-1,-1,-1, 1,-1, 1,-1, + 1, 1, 1,-1, 1, 1, 1, 1, 1, 1, 1,-1,-1,-1,-1,-1, + -1,-1,-1,-1, 1, 1, 1,-1,-1,-1, 1, 1, 1,-1,-1,-1, + -1,-1,-1,-1, 1, 1,-1, 1, 1,-1,-1, 1, 1,-1, 1,-1, +}; +static __constant__ float TQ_CB3_128[8] = { + -0.188397318125f,-0.118139825761f,-0.066585637629f,-0.021604320034f, + 0.021604320034f,0.066585637629f,0.118139825761f,0.188397318125f, +}; +static __constant__ float TQ_BD3_128[9] = { + -1.000000000000f,-0.153268575668f,-0.092362731695f,-0.044094979763f, + 0.000000000000f,0.044094979763f,0.092362731695f,0.153268575668f, + 1.000000000000f, +}; +static __constant__ float TQ_CB4_128[16] = { + -0.237664014101f,-0.180836066604f,-0.141805261374f,-0.110288411379f, + -0.082828491926f,-0.057772319764f,-0.034151583910f,-0.011302500963f, + 0.011302500963f,0.034151583910f,0.057772319764f,0.082828491926f, + 0.110288411379f,0.141805261374f,0.180836066604f,0.237664014101f, +}; +static __constant__ float TQ_BD4_128[17] = { + -1.000000000000f,-0.209250032902f,-0.161320656538f,-0.126046836376f, + -0.096558451653f,-0.070300407708f,-0.045961949974f,-0.022727042437f, + 0.000000000000f,0.022727042437f,0.045961949974f,0.070300407708f, + 0.096558451653f,0.126046836376f,0.161320656538f,0.209250032902f, + 1.000000000000f, }; -__constant__ float TQ_CODEBOOK_2BIT[4] = { - -0.133041590561f, -0.039991612341f, 0.039991612341f, 0.133041590561f, +// ==== head_dim = 256 (seed 42) ==== +static __constant__ float TQ_SGN_256[256] = { + 1,-1,-1, 1,-1, 1,-1, 1,-1, 1,-1,-1,-1, 1, 1, 1, + 1,-1,-1,-1, 1,-1,-1,-1,-1, 1,-1, 1,-1, 1, 1, 1, + 1, 1, 1, 1,-1, 1, 1,-1,-1,-1, 1, 1, 1,-1,-1, 1, + -1,-1, 1,-1, 1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, 1, + -1,-1, 1,-1, 1, 1, 1, 1,-1, 1, 1,-1,-1, 1,-1, 1, + 1,-1,-1, 1,-1, 1,-1,-1, 1,-1,-1,-1,-1,-1,-1,-1, + 1, 1, 1, 1, 1, 1,-1,-1, 1,-1,-1,-1,-1,-1, 1, 1, + 1,-1,-1,-1, 1, 1,-1, 1, 1, 1, 1,-1, 1,-1, 1,-1, + -1, 1,-1, 1,-1, 1, 1,-1,-1, 1,-1, 1,-1, 1,-1, 1, + -1,-1,-1, 1, 1,-1,-1, 1,-1,-1,-1,-1, 1,-1,-1, 1, + -1, 1, 1, 1, 1,-1, 1,-1,-1, 1, 1, 1, 1, 1, 1,-1, + -1, 1, 1,-1, 1,-1,-1,-1,-1, 1,-1, 1,-1, 1, 1,-1, + 1, 1,-1, 1, 1,-1, 1, 1, 1,-1, 1,-1, 1, 1,-1, 1, + -1,-1, 1, 1,-1,-1, 1, 1, 1, 1, 1, 1,-1,-1,-1, 1, + -1,-1,-1,-1,-1,-1, 1, 1, 1,-1, 1,-1,-1, 1,-1, 1, + 1, 1, 1, 1,-1, 1,-1, 1,-1,-1, 1, 1,-1,-1,-1,-1, +}; +static __constant__ float TQ_CB3_256[8] = { + -0.133854493499f,-0.083765551448f,-0.047166757286f,-0.015297502279f, + 0.015297502279f,0.047166757286f,0.083765551448f,0.133854493499f, +}; +static __constant__ float TQ_BD3_256[9] = { + -1.000000000000f,-0.108810022473f,-0.065466158092f,-0.031232129782f, + 0.000000000000f,0.031232129782f,0.065466158092f,0.108810022473f, + 1.000000000000f, +}; +static __constant__ float TQ_CB4_256[16] = { + -0.169410735369f,-0.128588363528f,-0.100698113441f,-0.078249387443f, + -0.058732159436f,-0.040949229151f,-0.024200897664f,-0.008008380421f, + 0.008008380421f,0.024200897664f,0.040949229151f,0.058732159436f, + 0.078249387443f,0.100698113441f,0.128588363528f,0.169410735369f, +}; +static __constant__ float TQ_BD4_256[17] = { + -1.000000000000f,-0.148999556899f,-0.114643231034f,-0.089473746717f, + -0.068490773439f,-0.049840692431f,-0.032575063407f,-0.016104638577f, + 0.000000000000f,0.016104638577f,0.032575063407f,0.049840692431f, + 0.068490773439f,0.089473746717f,0.114643231034f,0.148999556899f, + 1.000000000000f, }; -// Sign array for WHT rotation (seed=42, dim=128) -__constant__ float TQ_SIGNS_128[128] = { - -1, -1, 1, -1, 1, -1, -1, 1, -1, -1, 1, 1, -1, -1, -1, -1, - -1, 1, -1, -1, 1, 1, 1, 1, 1, 1, -1, -1, -1, 1, 1, -1, - -1, -1, 1, -1, 1, -1, -1, -1, 1, -1, -1, -1, 1, -1, -1, -1, - -1, -1, 1, -1, -1, -1, 1, -1, -1, -1, -1, -1, 1, -1, 1, 1, - -1, -1, -1, 1, -1, -1, 1, 1, -1, -1, -1, -1, 1, -1, 1, -1, - 1, 1, 1, -1, 1, 1, 1, 1, 1, 1, 1, -1, -1, -1, -1, -1, - -1, -1, -1, -1, 1, 1, 1, -1, -1, -1, 1, 1, 1, -1, -1, -1, - -1, -1, -1, -1, 1, 1, -1, 1, 1, -1, -1, 1, 1, -1, 1, -1, +// ==== head_dim = 512 (seed 42) ==== +static __constant__ float TQ_SGN_512[512] = { + 1,-1,-1, 1, 1,-1, 1, 1, 1,-1,-1, 1, 1, 1,-1, 1, + -1, 1, 1,-1,-1,-1,-1, 1, 1,-1,-1,-1,-1,-1,-1, 1, + -1,-1,-1, 1,-1, 1,-1, 1,-1,-1,-1, 1,-1,-1, 1,-1, + 1,-1, 1, 1,-1, 1,-1,-1,-1, 1,-1,-1, 1,-1, 1,-1, + 1, 1,-1, 1, 1, 1, 1, 1,-1, 1, 1, 1,-1, 1, 1,-1, + -1,-1,-1,-1, 1,-1,-1,-1,-1,-1,-1, 1, 1,-1, 1,-1, + 1, 1,-1,-1, 1,-1,-1, 1,-1, 1,-1, 1,-1,-1,-1,-1, + 1, 1, 1, 1, 1, 1,-1,-1,-1,-1,-1, 1,-1,-1,-1,-1, + -1, 1,-1,-1, 1, 1, 1,-1,-1, 1,-1,-1, 1, 1, 1, 1, + -1,-1, 1, 1, 1, 1, 1,-1,-1, 1,-1,-1, 1,-1, 1,-1, + 1,-1,-1, 1,-1, 1,-1,-1, 1,-1, 1,-1,-1, 1,-1, 1, + 1,-1,-1,-1,-1,-1,-1, 1, 1, 1,-1,-1,-1, 1,-1, 1, + -1, 1, 1,-1,-1,-1, 1,-1, 1,-1, 1, 1,-1, 1,-1,-1, + -1, 1,-1, 1,-1, 1, 1,-1, 1,-1,-1, 1, 1,-1, 1,-1, + 1, 1, 1, 1,-1,-1,-1,-1, 1, 1, 1,-1,-1,-1,-1, 1, + 1, 1,-1, 1, 1, 1,-1,-1, 1,-1,-1, 1, 1, 1, 1, 1, + -1,-1,-1, 1, 1,-1,-1,-1,-1,-1, 1, 1,-1, 1,-1,-1, + 1, 1,-1, 1, 1,-1,-1, 1, 1,-1, 1,-1,-1, 1, 1, 1, + 1,-1,-1, 1, 1, 1, 1,-1,-1, 1,-1,-1,-1, 1,-1, 1, + 1, 1,-1,-1, 1, 1, 1, 1, 1,-1,-1, 1, 1, 1,-1,-1, + 1,-1, 1, 1,-1, 1, 1,-1,-1, 1, 1, 1, 1,-1, 1, 1, + 1, 1,-1,-1, 1, 1,-1, 1,-1, 1, 1,-1, 1, 1, 1, 1, + -1, 1,-1,-1, 1,-1, 1,-1,-1,-1, 1, 1, 1,-1, 1, 1, + -1, 1, 1,-1, 1,-1,-1,-1,-1, 1, 1, 1,-1, 1, 1,-1, + 1, 1, 1, 1,-1, 1, 1, 1,-1, 1, 1, 1,-1, 1,-1,-1, + -1,-1, 1,-1, 1,-1,-1,-1, 1,-1,-1,-1,-1, 1,-1, 1, + 1, 1,-1,-1,-1,-1, 1,-1, 1,-1, 1, 1, 1, 1, 1,-1, + 1,-1, 1, 1,-1,-1,-1, 1, 1, 1,-1,-1, 1,-1,-1, 1, + 1,-1, 1,-1, 1,-1, 1, 1,-1,-1,-1, 1,-1, 1, 1, 1, + -1, 1,-1, 1,-1, 1, 1,-1, 1,-1,-1,-1,-1, 1, 1, 1, + 1,-1, 1, 1,-1, 1, 1, 1,-1,-1,-1, 1,-1,-1, 1, 1, + -1, 1,-1,-1, 1,-1,-1,-1, 1,-1, 1,-1,-1, 1, 1,-1, +}; +static __constant__ float TQ_CB3_512[8] = { + -0.094875931740f,-0.059311967343f,-0.033381462097f,-0.010824348778f, + 0.010824348778f,0.033381462097f,0.059311967343f,0.094875931740f, +}; +static __constant__ float TQ_BD3_512[9] = { + -1.000000000000f,-0.077093951404f,-0.046346712857f,-0.022102905437f, + -0.000000000000f,0.022102905437f,0.046346712857f,0.077093951404f, + 1.000000000000f, +}; +static __constant__ float TQ_CB4_512[16] = { + -0.120276063681f,-0.091181337833f,-0.071355909109f,-0.055424209684f, + -0.041587848216f,-0.028990162537f,-0.017131028697f,-0.005668541417f, + 0.005668541417f,0.017131028697f,0.028990162537f,0.041587848216f, + 0.055424209684f,0.071355909109f,0.091181337833f,0.120276063681f, +}; +static __constant__ float TQ_BD4_512[17] = { + -1.000000000000f,-0.105728700757f,-0.081268623471f,-0.063390061259f, + -0.048506028950f,-0.035289004445f,-0.023060595617f,-0.011399785057f, + -0.000000000000f,0.011399785057f,0.023060595617f,0.035289004445f, + 0.048506028950f,0.063390061259f,0.081268623471f,0.105728700757f, + 1.000000000000f, }; -// ============================================================================ -// WHT in shared memory (for query rotation) -// ============================================================================ -__device__ void tq_wht_butterfly_128(float* data, int tid) { - for (int h = 1; h < 128; h *= 2) { - __syncthreads(); - for (int idx = tid; idx < 64; idx += 128) { - int block_start = (idx / h) * (h * 2); - int offset = idx % h; - int i = block_start + offset; - int j = i + h; - float a = data[i]; - float b = data[j]; - data[i] = a + b; - data[j] = a - b; - } - } - __syncthreads(); -} +// --------------------------------------------------------------------------- +// Table selection by head dimension +// --------------------------------------------------------------------------- +// `if constexpr` chains rather than a class template: __constant__ arrays decay +// to __constant__-space pointers, so this compiles to a compile-time address +// with no indirection at run time. -// Apply DยทHยทD rotation to a 128-dim vector in shared memory -__device__ void tq_rotate_128(float* data, int tid) { - // D ยท x - if (tid < 128) data[tid] *= TQ_SIGNS_128[tid]; - __syncthreads(); - // H ยท D ยท x - tq_wht_butterfly_128(data, tid); - // Scale by 1/sqrt(128) - if (tid < 128) data[tid] *= 0.08838834764831845f; - __syncthreads(); - // D ยท H ยท D ยท x - if (tid < 128) data[tid] *= TQ_SIGNS_128[tid]; - __syncthreads(); +template +__device__ __forceinline__ const float* tq_signs() { + if constexpr (HS == 64) return TQ_SGN_64; + else if constexpr (HS == 128) return TQ_SGN_128; + else if constexpr (HS == 256) return TQ_SGN_256; + else return TQ_SGN_512; } -// ============================================================================ -// TurboQuant Paged Attention Kernel (simplified, HEAD_SIZE=128) -// ============================================================================ - -/** - * Simplified TurboQuant attention for decode (single query token). - * - * This kernel processes one sequence per block, one head per block. - * It computes attention over compressed KV cache using codebook lookups. - * - * Parameters: - * out: [num_seqs, num_heads, HEAD_SIZE] output - * q: [num_seqs, num_heads, HEAD_SIZE] query (not rotated) - * k_packed: [total_tokens, num_kv_heads, k_packed_bytes] packed 4-bit K indices - * v_packed: [total_tokens, num_kv_heads, v_packed_bytes] packed 3-bit V indices - * k_norms: [total_tokens, num_kv_heads] half norms for K - * v_norms: [total_tokens, num_kv_heads] half norms for V - * token_table: [num_seqs, max_context_len] maps (seq, pos) โ†’ token index - * context_lens: [num_seqs] actual context length per sequence - * - * Grid: (num_heads, num_seqs, 1) - * Block: (128, 1, 1) - */ -template -__global__ void turbo_paged_attention_kernel( - float* __restrict__ out, // [num_seqs, num_heads, HEAD_SIZE] - const float* __restrict__ q, // [num_seqs, num_heads, HEAD_SIZE] - const uint8_t* __restrict__ k_packed, // [total_tokens, num_kv_heads, k_bytes_per_head] - const uint8_t* __restrict__ v_packed, // [total_tokens, num_kv_heads, v_bytes_per_head] - const __half* __restrict__ k_norms, // [total_tokens, num_kv_heads] - const __half* __restrict__ v_norms, // [total_tokens, num_kv_heads] - const int32_t* __restrict__ token_table, // [num_seqs, max_tokens_per_seq] - const int32_t* __restrict__ context_lens,// [num_seqs] - const int num_kv_heads, - const int max_tokens_per_seq, - const int num_heads, - const float scale -) { - static_assert(HEAD_SIZE == 128, "TurboQuant currently supports HEAD_SIZE=128 only"); - - const int head_idx = blockIdx.x; - const int seq_idx = blockIdx.y; - const int tid = threadIdx.x; - const int context_len = context_lens[seq_idx]; - - if (context_len == 0) return; +template +__device__ __forceinline__ const float* tq_cb4() { + if constexpr (HS == 64) return TQ_CB4_64; + else if constexpr (HS == 128) return TQ_CB4_128; + else if constexpr (HS == 256) return TQ_CB4_256; + else return TQ_CB4_512; +} - const int kv_head_idx = head_idx / (num_heads / num_kv_heads); +template +__device__ __forceinline__ const float* tq_cb3() { + if constexpr (HS == 64) return TQ_CB3_64; + else if constexpr (HS == 128) return TQ_CB3_128; + else if constexpr (HS == 256) return TQ_CB3_256; + else return TQ_CB3_512; +} - // Packed sizes per head - constexpr int K_BYTES = (K_BITS == 4) ? (HEAD_SIZE / 2) : - (K_BITS == 3) ? (((HEAD_SIZE + 9) / 10) * 4) : - (HEAD_SIZE / 4); - constexpr int V_BYTES = (V_BITS == 4) ? (HEAD_SIZE / 2) : - (V_BITS == 3) ? (((HEAD_SIZE + 9) / 10) * 4) : - (HEAD_SIZE / 4); +template +__device__ __forceinline__ const float* tq_bd4() { + if constexpr (HS == 64) return TQ_BD4_64; + else if constexpr (HS == 128) return TQ_BD4_128; + else if constexpr (HS == 256) return TQ_BD4_256; + else return TQ_BD4_512; +} - // Select codebooks - const float* k_codebook = (K_BITS == 4) ? TQ_CODEBOOK_4BIT : - (K_BITS == 3) ? TQ_CODEBOOK_3BIT : - TQ_CODEBOOK_2BIT; - const float* v_codebook = (V_BITS == 4) ? TQ_CODEBOOK_4BIT : - (V_BITS == 3) ? TQ_CODEBOOK_3BIT : - TQ_CODEBOOK_2BIT; +template +__device__ __forceinline__ const float* tq_bd3() { + if constexpr (HS == 64) return TQ_BD3_64; + else if constexpr (HS == 128) return TQ_BD3_128; + else if constexpr (HS == 256) return TQ_BD3_256; + else return TQ_BD3_512; +} - // ========================================================= - // Step 1: Load and rotate query - // ========================================================= - __shared__ float q_rotated[128]; - const float* q_ptr = q + seq_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE; - if (tid < HEAD_SIZE) { - q_rotated[tid] = q_ptr[tid]; - } +/// 1/sqrt(HS) โ€” the normalisation that makes the Hadamard transform, and +/// therefore all of DยทHยทD, its own inverse. A plain constant expression rather +/// than a constexpr function so it needs no relaxed-constexpr support. +#define TQ_WHT_SCALE(HS) ((HS) == 64 ? 0.125f \ + : (HS) == 128 ? 0.08838834764831845f \ + : (HS) == 256 ? 0.0625f \ + : 0.04419417382415922f) + +// --------------------------------------------------------------------------- +// Rotation +// --------------------------------------------------------------------------- + +/// In-place DยทHยทD over `HS` contiguous floats in shared memory, cooperatively +/// across `NT` threads. +/// +/// Self-inverse, so the same call rotates a query forward and rotates the V +/// accumulator back. The butterfly and the 1/sqrt(HS) scaling are ordered to +/// match `turboquant::wht::fwht_inplace` exactly (all butterflies, then one +/// scale) โ€” reordering them would still be orthogonal but would no longer +/// agree with the cache written by the Rust compressor. +/// +/// At HS == NT == 128 this emits the same sequence of operations the original +/// hand-written `rotate128` did. +template +__device__ __forceinline__ void tq_rotate(float* d, int tid) { + const float* sgn = tq_signs(); + + for (int i = tid; i < HS; i += NT) d[i] *= sgn[i]; __syncthreads(); - // Apply DยทHยทD rotation to the query - tq_rotate_128(q_rotated, tid); - - // ========================================================= - // Step 2: Compute QยทK attention scores - // ========================================================= - // We process one token at a time. Each thread handles one dimension. - // This is simpler than the vectorized paged attention but correct. - - extern __shared__ char shared_mem[]; - float* logits = reinterpret_cast(shared_mem); - float* reduce_buf = reinterpret_cast(shared_mem + context_len * sizeof(float)); - - float qk_max = -FLT_MAX; - - for (int t = 0; t < context_len; t++) { - int token_idx = token_table[seq_idx * max_tokens_per_seq + t]; - const uint8_t* k_ptr = k_packed + (token_idx * num_kv_heads + kv_head_idx) * K_BYTES; - float k_norm = __half2float(k_norms[token_idx * num_kv_heads + kv_head_idx]); - - // Each thread loads and dequantizes one element of the key - float k_val = 0.0f; - if (tid < HEAD_SIZE) { - uint8_t idx; - if constexpr (K_BITS == 4) { - // Nibble packing: 2 per byte - uint8_t byte = k_ptr[tid / 2]; - idx = (tid & 1) ? ((byte >> 4) & 0xF) : (byte & 0xF); - } else if constexpr (K_BITS == 3) { - // 10-in-32 packing - int group = tid / 10; - int pos_in_group = tid % 10; - uint32_t word = *reinterpret_cast(k_ptr + group * 4); - idx = (word >> (pos_in_group * 3)) & 0x7; - } else { - // 2-bit: 4 per byte - uint8_t byte = k_ptr[tid / 4]; - int shift = 6 - (tid % 4) * 2; - idx = (byte >> shift) & 0x3; - } - k_val = k_codebook[idx]; - } - - // Dot product: q_rotated ยท k_centroids (all in rotated domain) - // Parallel reduction across 128 threads - float partial = (tid < HEAD_SIZE) ? q_rotated[tid] * k_val : 0.0f; - - // Warp-level reduction - for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) { - partial += __shfl_xor_sync(0xffffffff, partial, mask); - } - - // Cross-warp reduction via shared memory - __shared__ float warp_sums[4]; // 128 threads / 32 = 4 warps - if (tid % WARP_SIZE == 0) { - warp_sums[tid / WARP_SIZE] = partial; - } + for (int h = 1; h < HS; h *= 2) { __syncthreads(); - - float qk = 0.0f; - if (tid == 0) { - for (int w = 0; w < 4; w++) qk += warp_sums[w]; - qk = qk * k_norm * scale; - logits[t] = qk; - qk_max = fmaxf(qk_max, qk); + for (int i = tid; i < HS / 2; i += NT) { + const int bs = (i / h) * (h * 2), off = i % h; + const float a = d[bs + off], b = d[bs + off + h]; + d[bs + off] = a + b; + d[bs + off + h] = a - b; } - __syncthreads(); - } - - // Broadcast qk_max to all threads - if (tid == 0) { - reduce_buf[0] = qk_max; - } - __syncthreads(); - qk_max = reduce_buf[0]; - - // ========================================================= - // Step 3: Softmax - // ========================================================= - float exp_sum = 0.0f; - for (int t = tid; t < context_len; t += blockDim.x) { - float val = __expf(logits[t] - qk_max); - logits[t] = val; - exp_sum += val; - } - // Reduce exp_sum across threads - for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) { - exp_sum += __shfl_xor_sync(0xffffffff, exp_sum, mask); - } - __shared__ float warp_exp_sums[4]; - if (tid % WARP_SIZE == 0) { - warp_exp_sums[tid / WARP_SIZE] = exp_sum; - } - __syncthreads(); - if (tid == 0) { - float total = 0.0f; - for (int w = 0; w < 4; w++) total += warp_exp_sums[w]; - reduce_buf[0] = __fdividef(1.0f, total + 1e-6f); } __syncthreads(); - float inv_sum = reduce_buf[0]; - for (int t = tid; t < context_len; t += blockDim.x) { - logits[t] *= inv_sum; - } + constexpr float s = TQ_WHT_SCALE(HS); + for (int i = tid; i < HS; i += NT) d[i] *= s; __syncthreads(); - // ========================================================= - // Step 4: Weighted sum of values (attention ยท V) - // ========================================================= - // Each thread accumulates its dimension across all tokens - float acc = 0.0f; - - if (tid < HEAD_SIZE) { - for (int t = 0; t < context_len; t++) { - float weight = logits[t]; - if (weight < 1e-8f) continue; // Skip near-zero weights - - int token_idx = token_table[seq_idx * max_tokens_per_seq + t]; - const uint8_t* v_ptr = v_packed + (token_idx * num_kv_heads + kv_head_idx) * V_BYTES; - float v_norm = __half2float(v_norms[token_idx * num_kv_heads + kv_head_idx]); - - // Dequantize value element - uint8_t idx; - if constexpr (V_BITS == 4) { - uint8_t byte = v_ptr[tid / 2]; - idx = (tid & 1) ? ((byte >> 4) & 0xF) : (byte & 0xF); - } else if constexpr (V_BITS == 3) { - int group = tid / 10; - int pos_in_group = tid % 10; - uint32_t word = *reinterpret_cast(v_ptr + group * 4); - idx = (word >> (pos_in_group * 3)) & 0x7; - } else { - uint8_t byte = v_ptr[tid / 4]; - int shift = 6 - (tid % 4) * 2; - idx = (byte >> shift) & 0x3; - } - - float v_val = v_codebook[idx] * v_norm; - acc += weight * v_val; - } - } - - // ========================================================= - // Step 5: Inverse rotation and write output - // ========================================================= - // V was compressed in the rotated domain, so the accumulated output - // is also in the rotated domain. Apply inverse rotation. - __shared__ float out_rotated[128]; - if (tid < HEAD_SIZE) { - out_rotated[tid] = acc; - } + for (int i = tid; i < HS; i += NT) d[i] *= sgn[i]; __syncthreads(); - - // DยทHยทD is self-inverse - tq_rotate_128(out_rotated, tid); - - // Write output - float* out_ptr = out + seq_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE; - if (tid < HEAD_SIZE) { - out_ptr[tid] = out_rotated[tid]; - } } -// ============================================================================ -// Launch wrapper -// ============================================================================ - -inline void launch_turbo_paged_attention( - float* out, - const float* q, - const uint8_t* k_packed, - const uint8_t* v_packed, - const __half* k_norms, - const __half* v_norms, - const int32_t* token_table, - const int32_t* context_lens, - int num_kv_heads, - int max_tokens_per_seq, - int num_seqs, - int num_heads, - int head_size, - float scale, - int k_bits, - int v_bits, - cudaStream_t stream -) { - // Shared memory: context_len * sizeof(float) for logits + scratch - // Use max_tokens_per_seq as upper bound - int shared_mem_size = (max_tokens_per_seq + 1) * sizeof(float); +// --------------------------------------------------------------------------- +// Scalar quantizers +// --------------------------------------------------------------------------- +// Binary search over the codebook's decision boundaries. `boundaries[0]` and +// `boundaries[n]` are the sentinels -1 and +1, so a unit-sphere coordinate can +// never fall outside. + +template +__device__ __forceinline__ uint8_t tq_q4(float x) { + const float* bd = tq_bd4(); + if (x <= bd[1]) return 0; + if (x >= bd[16]) return 15; + int lo = 1, hi = 16; + while (lo < hi) { const int m = (lo + hi) >> 1; if (x < bd[m]) hi = m; else lo = m + 1; } + return (uint8_t)(lo - 1); +} - dim3 grid(num_heads, num_seqs, 1); - dim3 block(128); // One thread per head dimension +template +__device__ __forceinline__ uint8_t tq_q3(float x) { + const float* bd = tq_bd3(); + if (x <= bd[1]) return 0; + if (x >= bd[8]) return 7; + int lo = 1, hi = 8; + while (lo < hi) { const int m = (lo + hi) >> 1; if (x < bd[m]) hi = m; else lo = m + 1; } + return (uint8_t)(lo - 1); +} - // Dispatch based on bit-width combination - if (head_size == 128) { - if (k_bits == 4 && v_bits == 3) { - turbo_paged_attention_kernel<128, 4, 3><<>>( - out, q, k_packed, v_packed, k_norms, v_norms, - token_table, context_lens, num_kv_heads, max_tokens_per_seq, - num_heads, scale); - } else if (k_bits == 3 && v_bits == 3) { - turbo_paged_attention_kernel<128, 3, 3><<>>( - out, q, k_packed, v_packed, k_norms, v_norms, - token_table, context_lens, num_kv_heads, max_tokens_per_seq, - num_heads, scale); - } else if (k_bits == 3 && v_bits == 2) { - turbo_paged_attention_kernel<128, 3, 2><<>>( - out, q, k_packed, v_packed, k_norms, v_norms, - token_table, context_lens, num_kv_heads, max_tokens_per_seq, - num_heads, scale); - } +// --------------------------------------------------------------------------- +// Packed-K gather +// --------------------------------------------------------------------------- + +/// Bytes of packed 4-bit K each lane owns: (HS/2 bytes per head) / 32 lanes. +/// 1 at HS=64, 2 at 128, 4 at 256, 8 at 512. +#define TQ_BYTES_PER_LANE(HS) ((HS) / 64) + +/// Load a lane's whole packed-K run as one vector access. +/// +/// The run is contiguous and naturally aligned: the cache's x=16 swizzle keeps +/// each group of 16 bytes contiguous, and a lane's run starts at `lane*BPL` +/// with BPL a power of two dividing 16, so it never straddles a group. +template +__device__ __forceinline__ uint64_t tq_load_klane(const uint8_t* p) { + if constexpr (BPL == 1) { + return (uint64_t)(*p); + } else if constexpr (BPL == 2) { + return (uint64_t)(*reinterpret_cast(p)); + } else if constexpr (BPL == 4) { + return (uint64_t)(*reinterpret_cast(p)); + } else { + const uint2 v = *reinterpret_cast(p); + return ((uint64_t)v.y << 32) | (uint64_t)v.x; } } diff --git a/mistralrs-quant/examples/emit_turboquant_cuda_tables.rs b/mistralrs-quant/examples/emit_turboquant_cuda_tables.rs new file mode 100644 index 000000000..6f3e53332 --- /dev/null +++ b/mistralrs-quant/examples/emit_turboquant_cuda_tables.rs @@ -0,0 +1,70 @@ +//! Emit the CUDA `__constant__` tables for TurboQuant at a given head dimension. +//! +//! The CUDA kernels cannot call `generate_signs`/`cached_generated_codebook` at +//! runtime, so the tables are checked in as `__constant__` arrays. This example +//! is the *only* sanctioned way to produce them: it uses the same crate code the +//! Rust compressor uses, so the emitted tables cannot drift from the format by +//! construction. +//! +//! `turboquant_cuda_tables_match_crate` in `mistralrs-quant/src/turboquant/cuda_tables.rs` +//! re-derives these values at test time and fails if the checked-in CUDA source +//! disagrees, so a hand-edit of the `.cu` is caught by CI. +//! +//! Usage: `cargo run -p mistralrs-quant --example emit_turboquant_cuda_tables -- 512` + +// `get_codebook` is the same entry point `TurboQuantLayout::new` uses, so the +// emitted table is the one the Rust compressor actually quantizes against โ€” +// static for 64/128/256, numerically generated elsewhere. +use mistralrs_quant::turboquant::codebook::get_codebook; +use mistralrs_quant::turboquant::wht::generate_signs; + +const SEED: u64 = 42; + +fn emit_signs(dim: usize) { + let signs = generate_signs(SEED, dim); + println!("__constant__ float TQ_SGN_{dim}[{dim}] = {{"); + for row in signs.chunks(16) { + let cells: Vec = row + .iter() + .map(|s| format!("{:>2}", if *s > 0.0 { "1" } else { "-1" })) + .collect(); + println!(" {},", cells.join(",")); + } + println!("}};"); +} + +fn emit_codebook(dim: usize, bits: u32) { + let cb = get_codebook(dim, bits); + let n = cb.centroids.len(); + println!("__constant__ float TQ_CB{bits}_{dim}[{n}] = {{"); + for row in cb.centroids.chunks(4) { + let cells: Vec = row.iter().map(|c| format!("{c:.12}f")).collect(); + println!(" {},", cells.join(",")); + } + println!("}};"); + + let nb = cb.boundaries.len(); + println!("__constant__ float TQ_BD{bits}_{dim}[{nb}] = {{"); + for row in cb.boundaries.chunks(4) { + let cells: Vec = row.iter().map(|c| format!("{c:.12}f")).collect(); + println!(" {},", cells.join(",")); + } + println!("}};"); +} + +fn main() { + let dims: Vec = std::env::args() + .skip(1) + .map(|a| a.parse().expect("head dim must be an integer")) + .collect(); + let dims = if dims.is_empty() { vec![512] } else { dims }; + + for dim in dims { + println!("// ==== head_dim = {dim} (seed {SEED}) ===="); + emit_signs(dim); + for bits in [3u32, 4u32] { + emit_codebook(dim, bits); + } + println!(); + } +} diff --git a/mistralrs-quant/src/turboquant/cuda_tables.rs b/mistralrs-quant/src/turboquant/cuda_tables.rs new file mode 100644 index 000000000..e50c0c05b --- /dev/null +++ b/mistralrs-quant/src/turboquant/cuda_tables.rs @@ -0,0 +1,225 @@ +//! The contract between the Rust TurboQuant compressor and the CUDA kernels. +//! +//! The kernels cannot call [`super::wht::generate_signs`] or +//! [`super::codebook::get_codebook`] at run time, so the sign diagonals and +//! Lloydโ€“Max codebooks are checked into +//! `mistralrs-paged-attn/src/cuda/turbo_paged_attention.cuh` as `__constant__` +//! arrays. Two copies of the same numbers is a drift hazard: if the CUDA sign +//! table stops matching the Rust one, every compressed head decodes to noise, +//! and it decodes to noise *quickly* โ€” a broken kernel reports a better tok/s, +//! not a worse one. +//! +//! [`turboquant_cuda_tables_match_the_generator`] closes that. It parses the +//! header and re-derives every table from the same functions the compressor +//! uses. It is a CPU test about a *format*, not a claim about kernel behaviour +//! on a GPU. + +/// Head dimensions the CUDA kernels are instantiated for. +/// +/// Mirrors the `case` arms of `TQ_ATTN_BY_HS` / `TQ_CACHE_BY_HEAD_DIM` in +/// `turbo_paged_attention.cu`, and [`cuda_head_dims_match_the_kernel_dispatch`] +/// pins the two together. Callers gate on this before dispatching so an +/// unsupported width becomes a refusal rather than an untouched output buffer. +/// +/// This is narrower than what the *format* supports: +/// [`super::layout::TurboQuantLayout`] handles any width `>= MIN_BLOCK_DIM` by +/// block decomposition. The kernels take only the widths that decompose to a +/// single power-of-two block, which is every entry here. +pub const TURBOQUANT_CUDA_HEAD_DIMS: [usize; 4] = [64, 128, 256, 512]; + +/// Whether the CUDA kernels have an instantiation for `head_dim`. +pub fn cuda_supports_head_dim(head_dim: usize) -> bool { + TURBOQUANT_CUDA_HEAD_DIMS.contains(&head_dim) +} + +/// The seed the kernels' sign tables were generated with. Matches +/// [`super::TurboQuantConfig`]'s default. +pub const CUDA_TABLE_SEED: u64 = 42; + +#[cfg(test)] +mod tests { + use super::*; + use crate::turboquant::codebook::get_codebook; + use crate::turboquant::wht::generate_signs; + + const HEADER: &str = + include_str!("../../../mistralrs-paged-attn/src/cuda/turbo_paged_attention.cuh"); + const KERNEL: &str = + include_str!("../../../mistralrs-paged-attn/src/cuda/turbo_paged_attention.cu"); + /// The FFI wrapper's own copy of the list. `mistralrs-paged-attn` has no + /// dependency on this crate (and should not grow one just to share a + /// four-element array), so the two are kept equal by this test instead. + const FFI_WRAPPER: &str = + include_str!("../../../mistralrs-paged-attn/src/cuda/backend/paged_attention.rs"); + + /// Pull `static __constant__ float NAME[N] = { ... };` out of the header. + fn table(name: &str) -> Vec { + let needle = format!("float {name}["); + let start = HEADER + .find(&needle) + .unwrap_or_else(|| panic!("{name} is not declared in turbo_paged_attention.cuh")); + let open = HEADER[start..].find('{').expect("no opening brace") + start; + let close = HEADER[open..].find("};").expect("no closing brace") + open; + HEADER[open + 1..close] + .split(',') + .map(str::trim) + .filter(|s| !s.is_empty()) + .map(|s| { + s.trim_end_matches('f') + .parse::() + .unwrap_or_else(|e| panic!("{name}: cannot parse {s:?}: {e}")) + }) + .collect() + } + + fn assert_same(name: &str, cuda: &[f32], rust: &[f32]) { + assert_eq!( + cuda.len(), + rust.len(), + "{name}: CUDA table has {} entries, the generator produces {}", + cuda.len(), + rust.len() + ); + for (i, (c, r)) in cuda.iter().zip(rust).enumerate() { + assert!( + (c - r).abs() < 1e-9, + "{name}[{i}]: CUDA has {c}, the generator produces {r}. \ + The tables in turbo_paged_attention.cuh are generated โ€” re-run \ + `cargo run -p mistralrs-quant --example emit_turboquant_cuda_tables -- \ + 64 128 256 512` rather than editing them by hand." + ); + } + } + + /// Every `__constant__` table in the CUDA header equals what the Rust + /// compressor would produce at that dimension. + /// + /// This is the guard that makes the two implementations one format. Without + /// it, the four hand-copied sign tables in the tree (`.cu`, `.cuh`, + /// `turbo_wht.cu`, `.metal`) had nothing asserting they agreed with the + /// generator or with each other. + #[test] + fn turboquant_cuda_tables_match_the_generator() { + for dim in TURBOQUANT_CUDA_HEAD_DIMS { + assert_same( + &format!("TQ_SGN_{dim}"), + &table(&format!("TQ_SGN_{dim}")), + &generate_signs(CUDA_TABLE_SEED, dim), + ); + for bits in [3u32, 4u32] { + let cb = get_codebook(dim, bits); + assert_same( + &format!("TQ_CB{bits}_{dim}"), + &table(&format!("TQ_CB{bits}_{dim}")), + cb.centroids, + ); + assert_same( + &format!("TQ_BD{bits}_{dim}"), + &table(&format!("TQ_BD{bits}_{dim}")), + cb.boundaries, + ); + } + } + } + + /// A sign table that is all `+1` is not a rotation at all โ€” it would make + /// `DยทHยทD` collapse to a plain Hadamard transform and quietly change the + /// coordinate distribution the codebooks assume. Guards against a table + /// being zeroed or stubbed rather than merely wrong. + #[test] + fn cuda_sign_tables_are_actually_mixed() { + for dim in TURBOQUANT_CUDA_HEAD_DIMS { + let s = table(&format!("TQ_SGN_{dim}")); + let neg = s.iter().filter(|&&x| x < 0.0).count(); + assert!( + neg > dim / 8 && neg < dim - dim / 8, + "TQ_SGN_{dim} has {neg}/{dim} negative signs โ€” that is not a random diagonal" + ); + assert!( + s.iter().all(|&x| x == 1.0 || x == -1.0), + "TQ_SGN_{dim} contains a value that is not ยฑ1" + ); + } + } + + /// The kernel dispatch and [`TURBOQUANT_CUDA_HEAD_DIMS`] must agree. + /// + /// If a `case` is added to the CUDA switch without being added here the + /// Rust gate rejects a width the kernel could serve; if it is removed from + /// CUDA but left here, a launch falls through the switch and leaves the + /// output buffer untouched. That second failure mode is silent, which is + /// why this test exists rather than a comment. + #[test] + fn cuda_head_dims_match_the_kernel_dispatch() { + for dim in TURBOQUANT_CUDA_HEAD_DIMS { + assert!( + KERNEL.contains(&format!("TQ_ATTN_BY_BS({dim}, ")), + "head dim {dim} is in TURBOQUANT_CUDA_HEAD_DIMS but TQ_ATTN_BY_HS \ + in turbo_paged_attention.cu has no case for it" + ); + assert!( + KERNEL.contains(&format!("TQ_CACHE_BY_DTYPE(FN, {dim},")), + "head dim {dim} is in TURBOQUANT_CUDA_HEAD_DIMS but TQ_CACHE_BY_HEAD_DIM \ + in turbo_paged_attention.cu has no case for it" + ); + } + // And nothing the kernel dispatches is missing from the Rust list. + for dim in [64usize, 128, 256, 512, 1024] { + let in_kernel = KERNEL.contains(&format!("TQ_ATTN_BY_BS({dim}, ")); + assert_eq!( + in_kernel, + cuda_supports_head_dim(dim), + "head dim {dim}: kernel dispatch says {in_kernel}, \ + TURBOQUANT_CUDA_HEAD_DIMS says {}", + cuda_supports_head_dim(dim) + ); + } + } + + /// The Rust FFI wrapper refuses exactly the widths the kernel cannot serve. + /// + /// The wrapper is the last gate before the launch, so if its list is wider + /// than the kernel's the caller reads an uninitialized output buffer. + #[test] + fn ffi_wrapper_head_dims_match() { + let start = FFI_WRAPPER + .find("pub const TURBO_HEAD_DIMS") + .expect("TURBO_HEAD_DIMS is not declared in the paged-attn FFI wrapper"); + let open = FFI_WRAPPER[start..].find('[').unwrap() + start; + let open = FFI_WRAPPER[open + 1..].find('[').unwrap() + open + 1; // skip the [usize; N] type + let close = FFI_WRAPPER[open..].find(']').unwrap() + open; + let declared: Vec = FFI_WRAPPER[open + 1..close] + .split(',') + .map(str::trim) + .filter(|s| !s.is_empty()) + .map(|s| s.parse().expect("head dim is not an integer")) + .collect(); + assert_eq!( + declared, + TURBOQUANT_CUDA_HEAD_DIMS.to_vec(), + "mistralrs-paged-attn's TURBO_HEAD_DIMS has drifted from \ + TURBOQUANT_CUDA_HEAD_DIMS" + ); + } + + /// The packed byte counts the CUDA cache layout assumes must be divisible + /// by the x=16 swizzle group, or a decode lane's run straddles two groups + /// and the vectorized load in `tq_load_klane` reads the wrong bytes. + #[test] + fn packed_k_bytes_are_swizzle_aligned() { + for dim in TURBOQUANT_CUDA_HEAD_DIMS { + let k_bytes = super::super::packed_size(dim, 4); + assert_eq!( + k_bytes % 16, + 0, + "head dim {dim}: {k_bytes} packed K bytes is not a multiple of the x=16 group" + ); + // Each of the 32 lanes owns a whole, power-of-two-sized run. + let per_lane = k_bytes / 32; + assert!( + per_lane > 0 && per_lane.is_power_of_two() && 16 % per_lane == 0, + "head dim {dim}: {per_lane} bytes per lane does not divide the 16-byte group" + ); + } + } +} diff --git a/mistralrs-quant/src/turboquant/mod.rs b/mistralrs-quant/src/turboquant/mod.rs index f28f9eca4..2037de258 100644 --- a/mistralrs-quant/src/turboquant/mod.rs +++ b/mistralrs-quant/src/turboquant/mod.rs @@ -23,11 +23,13 @@ #[allow(clippy::excessive_precision)] pub mod codebook; +pub mod cuda_tables; pub mod generate; pub mod layout; pub mod wht; use codebook::Codebook; +pub use cuda_tables::{cuda_supports_head_dim, TURBOQUANT_CUDA_HEAD_DIMS}; pub use generate::{MAX_BLOCK_DIM, MIN_BLOCK_DIM}; pub use layout::{block_plan, BlockSpec, TurboQuantLayout}; use serde::{Deserialize, Serialize}; From 39cc0b5238df32d31ec0d5f34f35e44c9307a61a Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 18:42:39 +0100 Subject: [PATCH 08/21] =?UTF-8?q?feat(qtip):=20Stage=203=20=E2=80=94=20dis?= =?UTF-8?q?patch=20and=20selection=20for=20K=3D9,=20default=20OFF=20(#172)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(qtip): Stage 3 โ€” dispatch and selection for K=9, default OFF A K=9/V=4/L=12 layer is now SERVED rather than refused, and there is a way to ask for the rung. K=8 stays exactly as it was: the compiled control, nothing spent on serving it. WIRED * CPU decode for the whole V=4/L=12 family, routed through `trellis_v4l12::Rung` โ€” which already owns the extraction arithmetic and is bit-exactly gated against the GPU kernel. Not restated here: a format with two implementations has none. * `dequantize_weights` / `dequantize_w` / `forward` (2-D) for the family. * `forward_v4l12_gemv_cuda`: the family's single GPU path, the single-token fused decode+gemv the kernel exists for. Row-scale hoist OFF โ€” it is the one lever that costs bit-exactness against the CPU reference, and there is no device measurement to trade that for. STILL REFUSED, LOUDLY Every path with no V=4 kernel: 3-D `gather_forward`, `dequantize_expert`, 3-D `dequantize_weights`, and `qtip_packed` (whose `QtipPackedView` carries no geometry field, so a consumer would read K=9 bytes as K=4 nibbles). A half-wired dispatch that falls through is worse than one that refuses โ€” at 2 bpw a K=4 misread does not fault, it just serves wrong weights. SELECTION `ARC_QTIP_GEOMETRY` = `k4v2l16` | `kv4l12`, mirroring `ARC_QTIP_CODEBOOK`. Unknown values are refused, never resolved. DEFAULT OFF, and this is the load-bearing line: `QtipGeometry::DEFAULT` stays `K4V2L16`. Nothing in this family has touched a GPU โ€” the kernel has never been compiled at any K โ€” and the precedent is specific: the last KV-storage change that shipped default-on with only CPU validation killed every request on the first V4 forward that met a real device. Its arithmetic was proven and its device-time cost was not. Ship the path, leave the default alone. Guards shown red (S1-S7). THREE HOLES, all mine: * S2/S5 sit behind `#[cfg(feature = "cuda")]`, which no test lane compiles, so both mutations passed. Now pinned over the SOURCE โ€” a guard that only exists on hardware nobody has is not a guard. * S7: the env test reimplemented the parse rule inline and asserted the replica, so the real reader could silently resolve an unsupported K to the default and stay green. `from_env` now delegates to a pure `from_spec` and the test calls that. I also deleted my own five new tests mid-edit โ€” a second replacement span swallowed what the first had just inserted, leaving two dangling doc citations and a suite that looked green. Restored and verified present by name. And the source guard fired on itself: `include_str!("mod.rs")` includes the test module, so a negative assertion naming the string it forbids matched its own text. Scoped to the serving half. That is the second time this class bit me here; the note is in the code. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SpVNMpb13HkUXqSqbN1o9H * fix(qtip): Stage 3 CPU decode reads the padded row stride The K=9 CPU decode sized its rows at the data length. Rows now carry tail padding so the kernel's compile-time-width extraction needs no clamp, and that padding is part of what the artifact stores โ€” so a decoder that assumes the data length disagrees with the loader about where row N starts. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SpVNMpb13HkUXqSqbN1o9H * fix(ci): the CUDA dispatch source-guard died on Windows CRLF `include_str!("mod.rs")` yields CRLF on a Windows checkout, so `split_once("#[cfg(test)]\nmod tests {")` never matched and the guard hit its `expect` instead of asserting anything โ€” the guard was not failing, it was not running. Normalise line endings before the split; every assertion after it already goes through the whitespace-squeezer. Co-Authored-By: Claude Opus 5 (1M context) --------- Co-authored-by: Claude Opus 5 (1M context) --- mistralrs-quant/src/qtip/mod.rs | 636 ++++++++++++++++++++++++++++---- 1 file changed, 571 insertions(+), 65 deletions(-) diff --git a/mistralrs-quant/src/qtip/mod.rs b/mistralrs-quant/src/qtip/mod.rs index fdbd015a2..f4fd425b2 100644 --- a/mistralrs-quant/src/qtip/mod.rs +++ b/mistralrs-quant/src/qtip/mod.rs @@ -857,6 +857,64 @@ impl QtipGeometry { } } + /// What a fresh bake picks when `ARC_QTIP_GEOMETRY` says nothing. + /// + /// **`K4V2L16`, and the flip is this one line.** Deliberately not taken + /// here: nothing about the V=4/L=12 family has touched a GPU โ€” its kernel + /// has never been compiled at any K โ€” and the precedent for shipping a + /// CPU-validated storage change default-on is specific and bad. Its + /// arithmetic was proven and its device-time cost was not, and it killed + /// every request on the first V4 forward that met a real device. Ship the + /// path; leave the default alone. Pinned by + /// [`tests::the_new_geometry_is_not_the_default`]. + pub const DEFAULT: Self = QtipGeometry::K4V2L16; + + /// Read the geometry choice from the environment. + /// + /// `ARC_QTIP_GEOMETRY` is `k4v2l16` (the default) or `kv4l12` for a K + /// this build implements โ€” `k9v4l12` is the one worth asking for. Anything + /// else is refused rather than silently resolved, exactly as + /// [`QtipCodebook::from_env`] refuses: picking the wrong one changes what + /// every symbol in the artifact means. + pub fn from_env() -> Result { + match std::env::var("ARC_QTIP_GEOMETRY") { + Err(_) => Ok(Self::DEFAULT), + Ok(v) => Self::from_spec(&v), + } + } + + /// The parse rule behind [`QtipGeometry::from_env`], as a pure function. + /// + /// Split out so it can be tested directly. Testing a *replica* of this rule + /// instead โ€” which is what the first version of the test did โ€” asserts only + /// that the replica behaves, and passes happily while the real reader + /// silently accepts an unsupported K (measured: mutation S7). Mutating the + /// process environment in a unit test races the other tests in the binary, + /// so the function is the seam, not the variable. + pub fn from_spec(spec: &str) -> Result { + match spec { + "" => Ok(Self::DEFAULT), + "k4v2l16" | "k4" | "shipped" => Ok(QtipGeometry::K4V2L16), + other => { + // `kv4l12`, so a new K needs no new arm here โ€” the family + // decides whether it has a decoder. + let parsed = other + .strip_prefix('k') + .and_then(|r| r.strip_suffix("v4l12")) + .and_then(|k| k.parse::().ok()); + match parsed { + Some(k) => Self::trellis_v4l12(k), + None => candle_core::bail!( + "ARC_QTIP_GEOMETRY={other:?} is not a geometry. Use `k4v2l16` (the \ + default) or `kv4l12` with K in {:?}. Refusing rather than guessing \ + โ€” the choice changes what every symbol in the artifact means.", + trellis_v4l12::K_SUPPORTED + ), + } + } + } + } + /// Short label for bake headers and error messages. pub fn tag(self) -> String { match self { @@ -2717,6 +2775,9 @@ impl QtipLayer { // path (GPU-validated numerical parity); see arc-tools/CI_HYGIENE.md. #[allow(clippy::needless_range_loop)] fn dequantize_weights_rotated_f32(&self) -> Result> { + if let Some(rung) = self.v4l12_rung() { + return self.dequantize_weights_rotated_f32_v4l12(rung); + } let blocks_cpu = self.blocks.to_device(&Device::Cpu)?; let scales_cpu = self.row_scales.to_device(&Device::Cpu)?; let lut_cpu = self.lut.to_device(&Device::Cpu)?; @@ -2755,6 +2816,59 @@ impl QtipLayer { Ok(out) } + /// CPU decode for the V=4/L=12 family, in the rotated frame and scaled. + /// + /// Separate from the K=4/V=2 body rather than a branch inside it: the + /// symbol extraction, the packed row length and the table's element type + /// all differ, and `trellis_v4l12::Rung` already owns that arithmetic and + /// is bit-exactly gated against the GPU kernel. Restating it here would be + /// a second implementation of a format. + fn dequantize_weights_rotated_f32_v4l12(&self, rung: trellis_v4l12::Rung) -> Result> { + if self.lut.dtype() != DType::BF16 { + candle_core::bail!( + "QtipLayer[{}]: table must be BF16, got {:?}", + self.geometry.tag(), + self.lut.dtype() + ); + } + let blocks_cpu = self.blocks.to_device(&Device::Cpu)?.contiguous()?; + let scales_cpu = self.row_scales.to_device(&Device::Cpu)?.contiguous()?; + let lut_cpu = self.lut.to_device(&Device::Cpu)?.contiguous()?; + + let blocks_data: Vec = blocks_cpu.flatten_all()?.to_vec1()?; + let scales_data: Vec = scales_cpu.flatten_all()?.to_vec1()?; + let lut_data: Vec = lut_cpu.flatten_all()?.to_vec1()?; + + let n = scales_data.len(); + let k_in = self.in_features; + let num_symbols = rung.num_symbols(k_in); + // The allocated STRIDE, not the data length: rows carry tail padding + // so the kernel's compile-time-width extraction needs no clamp, and + // that padding is part of what the artifact stores. + let packed_per_row = rung.row_stride(num_symbols); + if blocks_data.len() != n * packed_per_row { + candle_core::bail!( + "QtipLayer[{}]: blocks hold {} B but {n} rows ร— {packed_per_row} B = {}", + self.geometry.tag(), + blocks_data.len(), + n * packed_per_row + ); + } + + let mut out = vec![0f32; n * k_in]; + for row in 0..n { + let packed = &blocks_data[row * packed_per_row..(row + 1) * packed_per_row]; + let dst = &mut out[row * k_in..(row + 1) * k_in]; + rung.decode_row_into( + trellis_v4l12::Row::new(packed, num_symbols, scales_data[row]), + &lut_data, + dst, + ) + .map_err(candle_core::Error::Msg)?; + } + Ok(out) + } + fn dequantize_weights_f32(&self) -> Result> { let mut out = self.dequantize_weights_rotated_f32()?; @@ -2781,6 +2895,26 @@ impl QtipLayer { } pub fn dequantize_weights(&self) -> Result { + // The V=4/L=12 family decodes on CPU (`dequantize_weights_rotated_f32` + // routes to it). There is no GPU dequantize kernel for the family and + // no 3-D path, so those two stay refused rather than falling through to + // a K=4 decoder. + if let Some(rung) = self.v4l12_rung() { + if self.num_experts.is_some() { + candle_core::bail!( + "QtipLayer::dequantize_weights: geometry {} has no stacked-expert (3-D) \ + decode. Refusing rather than reading these symbols with the K=4/V=2 \ + expert path.", + self.geometry.tag() + ); + } + let _ = rung; + let out = self.dequantize_weights_f32()?; + let n = self.row_scales.dim(0)?; + return Tensor::from_vec(out, (n, self.in_features), &Device::Cpu)? + .to_device(self.blocks.device())? + .to_dtype(DType::BF16); + } self.require_k4v2l16("dequantize_weights")?; // 3-D stacked-expert path: iterate over experts and stack `[N, K_in]` // slices into `[E, N, K_in]`. We reuse the 2-D dequant per expert @@ -2985,7 +3119,11 @@ impl QtipLayer { // CPU path below. #[cfg(feature = "cuda")] { - if cuda_ops::can_use_qtip_cuda(&self.blocks) { + // `forward_dequantize_cuda` is the K=4/V=2 CUDA path end to end โ€” + // its dequantize kernel and its fused gemv both index a + // `[2^16, 2]` F32 table. A family layer must never reach it, and + // gets its own dispatch below. + if self.v4l12_rung().is_none() && cuda_ops::can_use_qtip_cuda(&self.blocks) { if let Ok(mut result) = self.forward_dequantize_cuda(&x_2d, x.dtype(), n, k_in) { if let Some(bias) = &self.bias { result = result.broadcast_add(bias)?; @@ -3000,6 +3138,29 @@ impl QtipLayer { } } + // V=4/L=12 CUDA single-token decode: the one GPU path the family has. + // Anything else โ€” multi-token, or a shape the launcher refuses โ€” falls + // through to the CPU decode below, which is correct and slow, rather + // than to a K=4 kernel, which would be fast and wrong. + #[cfg(feature = "cuda")] + { + if let Some(rung) = self.v4l12_rung() { + if cuda_ops::can_use_qtip_cuda(&self.blocks) && x_2d.dim(0)? == 1 { + if let Ok(mut result) = self.forward_v4l12_gemv_cuda(&x_2d, x.dtype(), rung) { + if let Some(bias) = &self.bias { + result = result.broadcast_add(bias)?; + } + if orig_dims.len() > 2 { + let mut new_dims = orig_dims[..orig_dims.len() - 1].to_vec(); + new_dims.push(result.dim(1)?); + result = result.reshape(new_dims)?; + } + return Ok(result); + } + } + } + } + // Materializing fallback: full weight dequantize + matmul. At decode // shapes the fused GEMV above should have handled this. warn_dequant_materialize_at_decode( @@ -3069,6 +3230,76 @@ impl QtipLayer { /// /// Returns `Err` if any precondition fails (e.g. unsupported dtype or /// block size) โ€” the caller falls back to the CPU path. + /// V=4/L=12 single-token fused decode + gemv on CUDA. + /// + /// The family's only GPU path. Deliberately narrow: `n_tokens == 1`, the + /// shape the fused kernel exists for. Multi-token falls back to the CPU + /// decode rather than to a K=4 kernel โ€” slow and right beats fast and + /// wrong, and there is no V=4 dequantize kernel to matmul against. + /// + /// The row-scale hoist is **off**: it is the one lever that costs + /// bit-exactness against the CPU reference, and nothing about this family + /// has been measured on a device, so there is no evidence to trade + /// exactness for. + #[cfg(feature = "cuda")] + fn forward_v4l12_gemv_cuda( + &self, + x_2d: &Tensor, + out_dtype: DType, + rung: trellis_v4l12::Rung, + ) -> Result { + if !matches!(out_dtype, DType::BF16 | DType::F16 | DType::F32) { + candle_core::bail!("QTIP v4l12 forward CUDA: unsupported out dtype {out_dtype:?}"); + } + if self.num_experts.is_some() { + candle_core::bail!("QTIP v4l12 forward CUDA: no stacked-expert path"); + } + if tune::spec_pin_gemm() { + candle_core::bail!("QTIP v4l12 forward CUDA: spec_pin_gemm pins off the fused path"); + } + let k_in = self.in_features; + if x_2d.dim(1)? != k_in { + candle_core::bail!("QTIP v4l12 forward CUDA: x.dim(1) != in_features"); + } + + let x_rot_dtype = if matches!(x_2d.dtype(), DType::BF16 | DType::F16 | DType::F32) { + x_2d.dtype() + } else { + DType::BF16 + }; + let x_for_rot = if x_2d.dtype() == x_rot_dtype { + x_2d.contiguous()? + } else { + x_2d.to_dtype(x_rot_dtype)?.contiguous()? + }; + let x_rotated = if self.rotation_block >= 2 { + let signs = match &self.rotation_signs { + Some(t) => t, + None => candle_core::bail!( + "QtipLayer: rotation_block={} but rotation_signs is None", + self.rotation_block + ), + }; + cuda_ops::rotate_x_cuda(&x_for_rot, signs, self.rotation_block)? + } else { + x_for_rot + }; + + let y = cuda_ops::fused_gemv_v4l12_cuda( + &self.blocks, + &self.row_scales, + &self.lut, + &x_rotated, + k_in, + rung, + trellis_v4l12::RowScaleHoist::Off, + )?; + if y.dtype() != out_dtype { + return y.to_dtype(out_dtype); + } + Ok(y) + } + #[cfg(feature = "cuda")] fn forward_dequantize_cuda( &self, @@ -3902,7 +4133,12 @@ impl QuantMethod for QtipLayer { } fn forward(&self, x: &Tensor) -> Result { - self.require_k4v2l16("forward")?; + // 2-D forward is wired for the V=4/L=12 family; + // `forward_dequantize` picks the family's paths and refuses the ones + // that have no kernel. 3-D stacks are refused there, as before. + if self.v4l12_rung().is_none() { + self.require_k4v2l16("forward")?; + } self.forward_dequantize(x) } @@ -4280,18 +4516,25 @@ impl QtipLayer { match self.geometry { QtipGeometry::K4V2L16 => Ok(()), other => candle_core::bail!( - "QtipLayer::{op}: this layer was baked at geometry {}, and {op} only implements \ - K=4/V=2/L=16. The {} decode path is not wired into serving yet โ€” the fused \ - GEMV kernel exists (`kernels/qtip/qtip_gemv_k8v4l12.cu`) but nothing dispatches \ - to it. Refusing rather than decoding these symbols as K=4/V=2, which would not \ - fault (both geometries are 2 bits per weight, so the packed byte count is the \ - same) and would silently serve wrong weights.", - other.tag(), + "QtipLayer::{op}: this layer was baked at geometry {}, and {op} implements only \ + K=4/V=2/L=16. There is no V=4/L=12 kernel behind this entry point โ€” the family \ + has a fused decode+gemv kernel (`kernels/qtip/qtip_gemv_v4l12.cu`) and a CPU \ + decode, and nothing else. Refusing rather than decoding these symbols as \ + K=4/V=2: at 2 bits per weight the packed byte count is identical, so the \ + misread would not fault and would silently serve wrong weights.", other.tag() ), } } + /// The decode rung when this layer belongs to the V=4/L=12 family. + fn v4l12_rung(&self) -> Option { + match self.geometry { + QtipGeometry::TrellisV4L12 { rung } => Some(rung), + QtipGeometry::K4V2L16 => None, + } + } + /// Dequantize the i-th expert's `[N, K_in]` BF16 weight matrix (3-D mode /// only). Internal use by `gather_forward` and friends; bails when called /// on a 2-D layer or with `expert_idx >= num_experts`. @@ -7456,68 +7699,31 @@ mod tests { Ok(()) } - /// Every decode entry point must refuse a geometry it cannot decode. + /// Every decode entry point that has NO kernel for a geometry must refuse + /// it, rather than falling through to the K=4/V=2 decoder. + /// + /// The decoders behind these unpack nibbles and index a `[2^16, 2]` table, + /// and handed a V=4 row they would NOT fault โ€” at 2 bpw the packed byte + /// count is identical and every index lands in bounds. They would serve + /// garbage. /// - /// Stage 2 makes a K=8/V=4/L=12 artifact loadable. The decoders behind - /// these entry points all unpack nibbles and index a `[2^16, 2]` table, and - /// handed a K=8 row they would NOT fault โ€” at 2 bits per weight the packed - /// byte count is identical and every index lands in bounds. They would - /// serve garbage. Without this test, "the format landed" would mean "the - /// engine now has a way to silently serve wrong weights". + /// **Stage 3 narrowed this test on purpose.** `dequantize_weights` and + /// `forward` are now wired for the family (CPU decode, plus a CUDA + /// single-token fused gemv), so they are covered by + /// `a_k9_layer_forwards_instead_of_refusing` instead. What is left here is + /// what still has no kernel โ€” and a half-wired dispatch that falls through + /// is worse than one that refuses, so this stays. #[test] - fn every_decode_entry_point_refuses_an_unsupported_geometry() -> Result<()> { + fn unwired_entry_points_refuse_an_unsupported_geometry() -> Result<()> { let device = Device::Cpu; let k_in = 64usize; - // 2-D entry points. - let mut two_d = geometry_fixture(&device)?; - // Sanity: all of these work at the geometry this build implements, so - // the refusals below are about the geometry and nothing else. + // Sanity at the geometry this build implements everywhere, so the + // refusals below are about the geometry and nothing else. + let two_d = geometry_fixture(&device)?; two_d.dequantize_weights()?; - let x = Tensor::zeros((1, k_in), DType::F32, &device)?; - two_d.forward(&x)?; - two_d.dequantize_w()?; assert!(two_d.qtip_packed().is_some()); - two_d.geometry = k8(); - for (what, res) in [ - ("dequantize_weights", two_d.dequantize_weights().err()), - ("forward", two_d.forward(&x).err()), - ("dequantize_w", two_d.dequantize_w().err()), - ] { - let err = res.unwrap_or_else(|| panic!("{what} must refuse a k8v4l12 layer")); - let msg = err.to_string(); - assert!( - msg.contains("k8v4l12") && msg.contains("Refusing"), - "{what}: refusal must name the geometry and say it is refusing: {msg}" - ); - // The refusal must come from THIS entry point, not from something - // it happens to call. `forward` reaching an inner guard is still a - // refusal, but it is not evidence that `forward` checks โ€” and an - // entry point that relies on a callee's guard breaks the moment the - // callee gains a fast path that skips it. (Measured: mutation G1 - // removed `forward`'s own guard and every test stayed green.) - // - // `dequantize_w` is the one exception by construction: it is a - // one-line delegation to `dequantize_weights` and has no body to - // guard. - let expect_op = if what == "dequantize_w" { - "dequantize_weights" - } else { - what - }; - assert!( - msg.contains(&format!("QtipLayer::{expect_op}:")), - "{what}: refusal came from elsewhere โ€” expected `QtipLayer::{expect_op}:` in: {msg}" - ); - } - assert!( - two_d.qtip_packed().is_none(), - "qtip_packed must hand out nothing โ€” QtipPackedView carries no geometry, so a \ - consumer would read the bytes as K=4/V=2 nibbles" - ); - - // 3-D entry points. let mut stack = QtipLayer::stack_experts(vec![geometry_fixture(&device)?, geometry_fixture(&device)?])?; stack.dequantize_expert(0)?; @@ -7525,7 +7731,8 @@ mod tests { let idx = Tensor::zeros((1, 1), DType::U32, &device)?; stack.gather_forward(&a, &idx)?; - stack.geometry = k8(); + // Now the family, on paths with no kernel. + stack.geometry = QtipGeometry::trellis_v4l12(8)?; for (what, res) in [ ("dequantize_expert", stack.dequantize_expert(0).err()), ("gather_forward", stack.gather_forward(&a, &idx).err()), @@ -7535,12 +7742,311 @@ mod tests { assert!(msg.contains("k8v4l12"), "{what}: {msg}"); assert!( msg.contains(&format!("QtipLayer::{what}:")), - "{what}: refusal came from elsewhere: {msg}" + "{what}: refusal came from elsewhere โ€” an entry point that inherits a callee's \ + guard breaks the moment the callee gains a fast path: {msg}" ); } + + // `qtip_packed` carries no geometry, so it must hand out nothing. + let mut packed_view = geometry_fixture(&device)?; + packed_view.geometry = QtipGeometry::trellis_v4l12(8)?; + assert!( + packed_view.qtip_packed().is_none(), + "QtipPackedView has no geometry field; a consumer would read the bytes as K=4/V=2 \ + nibbles" + ); + Ok(()) + } + + /// Build a K=9 layer whose weights are known, so a forward can be checked + /// against an independent decode rather than against itself. + fn k9_layer(device: &Device, n: usize, k_in: usize) -> Result<(QtipLayer, Vec, Vec)> { + let geo = QtipGeometry::trellis_v4l12(9)?; + let rung = trellis_v4l12::Rung::new(9).map_err(candle_core::Error::Msg)?; + let n_sym = rung.num_symbols(k_in); + let row_bytes = rung.row_stride(n_sym); + + let mut packed = Vec::with_capacity(n * row_bytes); + for row in 0..n { + let mut z = (row as u64 + 1).wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1; + let syms: Vec = (0..n_sym) + .map(|_| { + z ^= z << 13; + z ^= z >> 7; + z ^= z << 17; + ((z >> 24) as u32) & 0x1FF + }) + .collect(); + packed.extend(rung.pack(&syms)); + } + let scales: Vec = (0..n).map(|i| 0.013 + i as f32 * 0.002).collect(); + + let layer = QtipLayer { + blocks: Tensor::from_vec(packed.clone(), (n, row_bytes), device)?, + row_scales: Tensor::from_vec(scales.clone(), (n,), device)?, + lut: Tensor::from_slice( + &trellis_v4l12::gaussian_lut_bf16(), + (trellis_v4l12::LUT_STATES, trellis_v4l12::V as usize), + device, + )?, + bias: None, + in_features: k_in, + num_experts: None, + rotation_signs: None, + rotation_block: 0, + search: QtipSearchStamp::Trellis, + search_detail: QtipSearchDetail::Known { + beam_width: None, + hessian: false, + }, + codebook: QtipCodebook::Gaussian, + geometry: geo, + }; + Ok((layer, packed, scales)) + } + + /// **Stage 3, the point of it: a K=9 layer is SERVED, not refused.** + /// + /// Checked against an independent decode through `trellis_v4l12::Rung`, + /// not against the layer's own dequantize โ€” otherwise a wrong decoder + /// would agree with itself. + #[test] + fn a_k9_layer_forwards_instead_of_refusing() -> Result<()> { + let device = Device::Cpu; + let (n, k_in) = (6usize, 128usize); + let (layer, packed, scales) = k9_layer(&device, n, k_in)?; + + let w = layer.dequantize_weights()?.to_dtype(DType::F32)?; + assert_eq!(w.dims(), &[n, k_in]); + let got: Vec = w.flatten_all()?.to_vec1()?; + + let rung = trellis_v4l12::Rung::new(9).unwrap(); + let n_sym = rung.num_symbols(k_in); + let row_bytes = rung.row_stride(n_sym); + let lut = trellis_v4l12::gaussian_lut_bf16(); + for row in 0..n { + let mut want = vec![0f32; k_in]; + rung.decode_row_into( + trellis_v4l12::Row::new( + &packed[row * row_bytes..(row + 1) * row_bytes], + n_sym, + scales[row], + ), + &lut, + &mut want, + ) + .unwrap(); + for i in 0..k_in { + // dequantize_weights returns BF16, so compare at BF16 precision. + let g = got[row * k_in + i]; + let e = half::bf16::from_f32(want[i]).to_f32(); + assert_eq!(g, e, "row {row} col {i}"); + } + } + + // ...and a real forward runs and matches x @ W^T. F32 on CPU: candle's + // CPU matmul has no BF16 kernel, which is a property of the shipped + // K=4 path too, not of this geometry. + let x: Vec = (0..k_in).map(|i| ((i as f32) * 0.07).sin()).collect(); + let x_t = Tensor::from_vec(x, (1, k_in), &device)?; + let y = layer.forward(&x_t)?.to_dtype(DType::F32)?; + assert_eq!(y.dims(), &[1, n]); + let y_ref = x_t.matmul(&w.t()?)?; + let (yv, rv): (Vec, Vec) = + (y.flatten_all()?.to_vec1()?, y_ref.flatten_all()?.to_vec1()?); + for (a, b) in yv.iter().zip(rv.iter()) { + let denom = a.abs().max(b.abs()).max(1e-3); + assert!((a - b).abs() / denom < 1e-2, "forward {a} vs reference {b}"); + } Ok(()) } + /// A K=9 forward must NOT be the K=4 decoder quietly succeeding. + /// + /// Both geometries decode the same bytes without faulting, so "it returned + /// numbers" proves nothing. Decode the same packed bytes as K=4/V=2 and + /// require the answers to differ. + #[test] + fn the_k9_path_is_not_the_k4_decoder_in_disguise() -> Result<()> { + let device = Device::Cpu; + let (n, k_in) = (4usize, 128usize); + let (layer, _, _) = k9_layer(&device, n, k_in)?; + let k9_w: Vec = layer + .dequantize_weights()? + .to_dtype(DType::F32)? + .flatten_all()? + .to_vec1()?; + + // Same bytes, same scales, read as the shipped rung. A K=9 row is + // longer than a K=4 row of the same in_features, so trim to what the + // K=4 decoder expects โ€” the point is that the SYMBOLS differ, not the + // lengths. + let k4_bytes = QtipGeometry::K4V2L16.packed_len(k_in); + let blocks: Vec = layer.blocks.flatten_all()?.to_vec1()?; + let row_bytes = blocks.len() / n; + let mut trimmed = Vec::with_capacity(n * k4_bytes); + for row in 0..n { + trimmed.extend_from_slice(&blocks[row * row_bytes..row * row_bytes + k4_bytes]); + } + let mut as_k4 = layer; + as_k4.geometry = QtipGeometry::K4V2L16; + as_k4.blocks = Tensor::from_vec(trimmed, (n, k4_bytes), &device)?; + as_k4.lut = Tensor::from_vec(gaussian_lut(), (1usize << L, V as usize), &device)?; + let k4_w: Vec = as_k4 + .dequantize_weights()? + .to_dtype(DType::F32)? + .flatten_all()? + .to_vec1()?; + + assert_eq!(k9_w.len(), k4_w.len()); + let differing = k9_w + .iter() + .zip(k4_w.iter()) + .filter(|(a, b)| a.to_bits() != b.to_bits()) + .count(); + assert!( + differing > k9_w.len() / 2, + "only {differing}/{} weights differ between the K=9 and K=4 readings of the same \ + bytes โ€” the K=9 path may be the K=4 decoder", + k9_w.len() + ); + Ok(()) + } + + /// Paths the family has no kernel for must still refuse loudly, on a + /// WELL-FORMED family layer. A half-wired dispatch that falls through to a + /// K=4 path is worse than one that refuses. + #[test] + fn unwired_paths_still_refuse_for_a_well_formed_k9_layer() -> Result<()> { + let device = Device::Cpu; + let (layer, _, _) = k9_layer(&device, 4, 128)?; + + assert!( + layer.qtip_packed().is_none(), + "QtipPackedView carries no geometry; a consumer would read K=9 bytes as nibbles" + ); + + let mut three_d = layer; + three_d.num_experts = Some(2); + for (what, err) in [ + ("dequantize_expert", three_d.dequantize_expert(0).err()), + ( + "gather_forward", + three_d + .gather_forward( + &Tensor::zeros((1, 1, 128), DType::F32, &device)?, + &Tensor::zeros((1, 1), DType::U32, &device)?, + ) + .err(), + ), + ( + "dequantize_weights(3-D)", + three_d.dequantize_weights().err(), + ), + ] { + let e = err.unwrap_or_else(|| panic!("{what} must refuse a k9v4l12 layer")); + let msg = e.to_string(); + assert!( + msg.contains("k9v4l12"), + "{what}: refusal must name the geometry: {msg}" + ); + assert!( + msg.to_lowercase().contains("refus"), + "{what}: refusal must say it is refusing: {msg}" + ); + } + Ok(()) + } + + /// The two dispatch invariants that live behind `#[cfg(feature = "cuda")]` + /// and therefore cannot be exercised in a no-GPU lane. + /// + /// Both were mutated and both passed (S2, S5): a CUDA build is the only + /// thing that compiles them, and this crate's CI CUDA lane does not run + /// tests. So assert them over the source, the same way the kernel's + /// constants are asserted โ€” a guard that only exists on hardware nobody + /// has is not a guard. + #[test] + fn the_cuda_dispatch_invariants_are_pinned_over_the_source() { + const SRC: &str = include_str!("mod.rs"); + fn squeeze(s: &str) -> String { + s.split_whitespace().collect::>().join(" ") + } + // SERVING CODE ONLY. `include_str!` pulls in this test module too, and + // the negative assertion below names the very string it forbids โ€” so + // searching the whole file makes the guard fire on itself. (Second time + // this class of bug appeared here; the first was a barrier guard + // counting the word "__syncthreads" in a comment. A guard that matches + // its own text is not a guard.) + // Normalise line endings first: on a Windows checkout `include_str!` + // yields CRLF, so a split pattern carrying a bare `\n` never matches + // and the guard dies in `expect` instead of asserting anything. Every + // assertion below goes through `squeeze`, which is already + // newline-agnostic โ€” this split was the one place that was not. + let src_lf = SRC.replace("\r\n", "\n"); + let serving = src_lf + .split_once("#[cfg(test)]\nmod tests {") + .map(|(before, _)| before) + .expect("mod.rs must contain its test module"); + let src = squeeze(serving); + + // A family layer must never reach the K=4/V=2 CUDA path, whose + // dequantize kernel and fused gemv both index a [2^16, 2] F32 table. + assert!( + src.contains(&squeeze( + "if self.v4l12_rung().is_none() && cuda_ops::can_use_qtip_cuda(&self.blocks) {" + )), + "the K=4/V=2 CUDA forward path lost its geometry gate. A V=4 layer reaching it \ + would be decoded as K=4/V=2 โ€” in bounds, no fault, wrong weights." + ); + + // The family's CUDA path must ship with the row-scale hoist OFF: it is + // the one lever that costs bit-exactness against the CPU reference, and + // nothing here has been measured on a device. + assert!( + src.contains(&squeeze("trellis_v4l12::RowScaleHoist::Off,")), + "the family's CUDA gemv no longer passes RowScaleHoist::Off" + ); + assert!( + !src.contains(&squeeze("trellis_v4l12::RowScaleHoist::On")), + "something in the serving path now asks for the row-scale hoist. It reassociates \ + the sum and breaks bit-exactness against the CPU reference; there is no measured \ + device-time evidence to trade that for." + ); + } + + /// **The new geometry must not be the default.** Nothing in the family has + /// touched a GPU; the kernel has never been compiled at any K. + #[test] + fn the_new_geometry_is_not_the_default() { + assert_eq!(QtipGeometry::DEFAULT, QtipGeometry::K4V2L16); + assert_eq!(QtipGeometry::default(), QtipGeometry::K4V2L16); + let device = Device::Cpu; + let layer = geometry_fixture(&device).unwrap(); + assert_eq!(layer.geometry, QtipGeometry::K4V2L16); + } + + /// The selection surface: `ARC_QTIP_GEOMETRY`, mirroring + /// `ARC_QTIP_CODEBOOK`. Refuses rather than guessing. + #[test] + fn the_geometry_env_surface_parses_and_refuses() { + // Parsing is exercised through the same rule the reader uses rather + // than by mutating the process environment, which races other tests in + // this binary. + // THE REAL RULE, not a replica of it. An earlier version of this test + // reimplemented the parse inline and passed while the real reader + // silently resolved an unsupported K to the default (mutation S7). + let parse = QtipGeometry::from_spec; + assert_eq!(parse("").unwrap(), QtipGeometry::K4V2L16); + assert_eq!(parse("k4v2l16").unwrap(), QtipGeometry::K4V2L16); + assert_eq!(parse("k9v4l12").unwrap().k(), 9); + assert_eq!(parse("k8v4l12").unwrap().k(), 8); + // A K with no decoder is refused, not clamped to a neighbour. + assert!(parse("k11v4l12").is_err()); + assert!(parse("qtip9").is_err()); + assert!(parse("k9").is_err()); + } + /// A genuine V=4/L=12 payload round-trips at every K, tag and tensors both. #[test] fn trellis_v4l12_payloads_round_trip_exactly_at_every_k() -> Result<()> { From a4560563170c346dae9272d195e5a3d781b39a08 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 19:19:35 +0100 Subject: [PATCH 09/21] fix(ArcGraph): pre-warm the capture from a measured alloc profile, and refuse to instantiate on a miss (#181) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Capture recorded and died on the FIRST cuGraphLaunch with 700 (CUDA_ERROR_ILLEGAL_ADDRESS), one line after: [alloc-cache] WARNING: MISS during capture, size 393216 bytes (not pre-warmed) That miss is the cause, not a bystander. An allocation that misses the warm pool while a capture is open is served by cuMemAllocAsync from the graph's PRIVATE pool and recorded as a graph memory node whose address the first launch has to back; it does not get backed, the launch faults, and the poisoned context surfaces later as glibc heap corruption. The allocator's 1 GiB cap is not involved โ€” bounded and unbounded legs are byte-identical during capture (alloc/step 44.1, free/step 78.7, hit-rate 0.9954, held 315.4 MiB) and die on the same line. Why the warm pass missed it: `tail` in the rolling compressor is rebuilt every step at width `tokens - base`, and `base` jumps a whole `ratio` at a group boundary while `tokens` climbs by one, so the size cycles through `ratio` consecutive values (measured `4096 ร— {18,19,20,21}`). One warm pass warms one phase of that cycle. The captured step lands on another phase and asks for a size the pool has never held. Three changes, none of them a hardcoded size: 1. Warm until the profile stops growing. `ARC_GRAPH_DEFERRED_PASSES` defaults to 4 (was 1), and a pass that teaches the allocator a new size buys another, from a bounded budget (`ARC_GRAPH_DEFERRED_MAX`, 24). A cycle of any period is covered without anyone tuning a constant. 2. Top the pool up from the measured profile before capture opens โ€” `prewarm_alloc_cache(slack)`, `ARC_GRAPH_PREWARM_SLACK` default 4 buffers per size. This covers the residual case the warm passes cannot: the captured step allocating MORE of a size than any warm step did. 3. Gate instantiation on zero capture-time misses. The allocator now keeps a counted, capture-scoped miss ledger (the old `missed` set is a log dedup โ€” it silences a size's SECOND miss, which is exactly the one that matters). If any miss occurred while recording, the capture is cancelled and the run falls back to eager instead of instantiating a graph already known to fault. (3) is the part that holds even if (1) and (2) are incomplete: it converts a poisoned context and a process death into a named list of sizes. If those sizes turn out to grow monotonically rather than cycle, no amount of warmup is the answer and the next step is a shape-constant buffer or allocator size-class bucketing โ€” the refusal message says so. Needs the candle fork's capture-miss ledger + profile-driven pre-warm (`prewarm_alloc_cache`, `reset_capture_misses`, `capture_misses`). This commit originally pinned Cargo.toml straight at 1fa534db, which was correct when it was written and is WRONG now: since then the integration branch moved to 89ab14ef, and the two revs have DIVERGED. 1fa534db is missing three commits 89ab14ef has โ€” "perf(cuda): bound the caching allocator", its cap-contract test, and "fix(cuda): eviction must not run inside the capture window" โ€” the first of which is what #213 (bound the alloc cache against real headroom) rests on. Pinning at 1fa534db would therefore have silently reverted the bounded allocator and a capture-correctness fix. Pinned instead at 9586979db (`arcgraph/bounded-alloc-plus-prewarm`), the merge of the two, which is ahead of 89ab14ef by 6 and of 1fa534db by 8 and behind neither โ€” verified with `gh api repos/aeonmindai/candle/compare/...` and by reading all four required `pub fn`s out of `candle-core/src/cuda_backend/device.rs` at that rev. `cargo check --workspace` is green on the new pin; mistralrs-core's warning count is unchanged at 23. The re-indent of the same-step-probe block in normal.rs is the counterpart of the one 6fcb87f47 did on the integration branch: that commit de-indented it by 4 precisely because this commit's `if !misses.is_empty()` nesting level was not yet present. It is now, so the block goes back 4 deeper. Co-authored-by: Nirupam Bhowmick --- Cargo.lock | 14 +- Cargo.toml | 35 ++- arc-cuda-graph/src/graph.rs | 62 ++++- mistralrs-core/src/pipeline/normal.rs | 358 ++++++++++++++++++-------- 4 files changed, 340 insertions(+), 129 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5a68f4570..6d09c410e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -692,7 +692,7 @@ checksum = "ade8366b8bd5ba243f0a58f036cc0ca8a2f069cff1a2351ef1cac6b083e16fc0" [[package]] name = "candle-core" version = "0.9.2" -source = "git+https://github.com/aeonmindai/candle.git?rev=89ab14ef1216331a782539f533b99f6816708604#89ab14ef1216331a782539f533b99f6816708604" +source = "git+https://github.com/aeonmindai/candle.git?rev=9586979db9031df6377f7dab1233c4548d976b76#9586979db9031df6377f7dab1233c4548d976b76" dependencies = [ "accelerate-src", "byteorder", @@ -723,7 +723,7 @@ dependencies = [ [[package]] name = "candle-flash-attn" version = "0.9.2" -source = "git+https://github.com/aeonmindai/candle.git?rev=89ab14ef1216331a782539f533b99f6816708604#89ab14ef1216331a782539f533b99f6816708604" +source = "git+https://github.com/aeonmindai/candle.git?rev=9586979db9031df6377f7dab1233c4548d976b76#9586979db9031df6377f7dab1233c4548d976b76" dependencies = [ "anyhow", "candle-core", @@ -734,7 +734,7 @@ dependencies = [ [[package]] name = "candle-flash-attn-v3" version = "0.9.2" -source = "git+https://github.com/aeonmindai/candle.git?rev=89ab14ef1216331a782539f533b99f6816708604#89ab14ef1216331a782539f533b99f6816708604" +source = "git+https://github.com/aeonmindai/candle.git?rev=9586979db9031df6377f7dab1233c4548d976b76#9586979db9031df6377f7dab1233c4548d976b76" dependencies = [ "anyhow", "candle-core", @@ -747,7 +747,7 @@ dependencies = [ [[package]] name = "candle-kernels" version = "0.9.2" -source = "git+https://github.com/aeonmindai/candle.git?rev=89ab14ef1216331a782539f533b99f6816708604#89ab14ef1216331a782539f533b99f6816708604" +source = "git+https://github.com/aeonmindai/candle.git?rev=9586979db9031df6377f7dab1233c4548d976b76#9586979db9031df6377f7dab1233c4548d976b76" dependencies = [ "cudaforge", ] @@ -755,7 +755,7 @@ dependencies = [ [[package]] name = "candle-metal-kernels" version = "0.9.2" -source = "git+https://github.com/aeonmindai/candle.git?rev=89ab14ef1216331a782539f533b99f6816708604#89ab14ef1216331a782539f533b99f6816708604" +source = "git+https://github.com/aeonmindai/candle.git?rev=9586979db9031df6377f7dab1233c4548d976b76#9586979db9031df6377f7dab1233c4548d976b76" dependencies = [ "half", "objc2", @@ -769,7 +769,7 @@ dependencies = [ [[package]] name = "candle-nn" version = "0.9.2" -source = "git+https://github.com/aeonmindai/candle.git?rev=89ab14ef1216331a782539f533b99f6816708604#89ab14ef1216331a782539f533b99f6816708604" +source = "git+https://github.com/aeonmindai/candle.git?rev=9586979db9031df6377f7dab1233c4548d976b76#9586979db9031df6377f7dab1233c4548d976b76" dependencies = [ "accelerate-src", "candle-core", @@ -788,7 +788,7 @@ dependencies = [ [[package]] name = "candle-ug" version = "0.9.2" -source = "git+https://github.com/aeonmindai/candle.git?rev=89ab14ef1216331a782539f533b99f6816708604#89ab14ef1216331a782539f533b99f6816708604" +source = "git+https://github.com/aeonmindai/candle.git?rev=9586979db9031df6377f7dab1233c4548d976b76#9586979db9031df6377f7dab1233c4548d976b76" dependencies = [ "ug", "ug-cuda", diff --git a/Cargo.toml b/Cargo.toml index b8852835d..ef611a2af 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -44,16 +44,31 @@ rust-version = "1.88" # RUN-161: forked candle for the opt-in CUDA caching allocator (graph-capture-safe). # Default-off; identical to upstream c3bb5bf when ARC_CANDLE_ALLOC_CACHE unset. # Fork of huggingface/candle 0.9.2 (rev c3bb5bf), mirrored to the org for CI. -# PINNED TO AN EXPLICIT REV, not a branch. This change depends on -# `drain_alloc_cache_and_free` (fork commit 88d86a20), and a moving branch -# pointer has broken CI on this repo before โ€” a path/branch dependency that -# advances under a merged PR turns a green build into a red one with no commit -# here to blame. -candle-core = { git = "https://github.com/aeonmindai/candle.git", rev = "89ab14ef1216331a782539f533b99f6816708604" } -candle-nn = { git = "https://github.com/aeonmindai/candle.git", rev = "89ab14ef1216331a782539f533b99f6816708604" } -candle-flash-attn-v3 = { git = "https://github.com/aeonmindai/candle.git", rev = "89ab14ef1216331a782539f533b99f6816708604" } -candle-flash-attn = { git = "https://github.com/aeonmindai/candle.git", rev = "89ab14ef1216331a782539f533b99f6816708604" } -candle-metal-kernels = { git = "https://github.com/aeonmindai/candle.git", rev = "89ab14ef1216331a782539f533b99f6816708604" } +# PINNED TO AN EXPLICIT REV, not a branch. This tree needs BOTH lines of the +# fork's allocator work, and a moving branch pointer has broken CI on this repo +# before โ€” a path/branch dependency that advances under a merged PR turns a +# green build into a red one with no commit here to blame. +# +# The rev below is `arcgraph/bounded-alloc-plus-prewarm` (9586979db), which is +# the MERGE of the two, and a strict superset of each: +# +# * 89ab14ef (`perf/alloc-cache-bounded`) โ€” `drain_alloc_cache_and_free` +# (88d86a20) plus the bounded caching allocator #213 depends on, its cap +# contract test, and "eviction must not run inside the capture window". +# * 1fa534db (`arcgraph/prewarm-profile`) โ€” the capture-miss ledger and +# profile-driven pre-warm this file's pre-warm caller needs +# (`prewarm_alloc_cache`, `reset_capture_misses`, `capture_misses`). +# +# Those two revs had DIVERGED: 1fa534db is missing three commits 89ab14ef has. +# Pinning straight to 1fa534db, as this change originally did, would therefore +# have silently reverted the bounded allocator and the capture-window eviction +# fix. Verified by `gh api repos/aeonmindai/candle/compare/...`: 9586979db is +# ahead of 89ab14ef by 6 and of 1fa534db by 8, behind neither. +candle-core = { git = "https://github.com/aeonmindai/candle.git", rev = "9586979db9031df6377f7dab1233c4548d976b76" } +candle-nn = { git = "https://github.com/aeonmindai/candle.git", rev = "9586979db9031df6377f7dab1233c4548d976b76" } +candle-flash-attn-v3 = { git = "https://github.com/aeonmindai/candle.git", rev = "9586979db9031df6377f7dab1233c4548d976b76" } +candle-flash-attn = { git = "https://github.com/aeonmindai/candle.git", rev = "9586979db9031df6377f7dab1233c4548d976b76" } +candle-metal-kernels = { git = "https://github.com/aeonmindai/candle.git", rev = "9586979db9031df6377f7dab1233c4548d976b76" } axum = "0.8.8" anyhow = "1.0.100" diff --git a/arc-cuda-graph/src/graph.rs b/arc-cuda-graph/src/graph.rs index c288998f7..20c0e805f 100644 --- a/arc-cuda-graph/src/graph.rs +++ b/arc-cuda-graph/src/graph.rs @@ -108,6 +108,10 @@ pub struct CudaGraphRunner { /// context-dependent and the real fix is shape-constant buffers, not more /// warmup. deferred_passes_remaining: u32, + /// Deferred passes still available to `grant_extra_deferred_pass`. Counts + /// down independently of `deferred_passes_remaining` so that "keep warming + /// while the profile grows" cannot become an unbounded loop. + deferred_extra_budget: u32, } #[cfg(feature = "cuda")] @@ -121,13 +125,40 @@ impl CudaGraphRunner { Self::new_with_passes(device, warmup_steps, Self::default_deferred_passes()) } - /// Number of deferred-free warmup passes, from `ARC_GRAPH_DEFERRED_PASSES` - /// (default 1, the historical behaviour). See `deferred_passes_remaining`. + /// Number of deferred-free warmup passes, from `ARC_GRAPH_DEFERRED_PASSES`. + /// + /// Default 4, not the historical 1. One pass only suffices if every decode + /// step allocates the same set of sizes, and V4's does not: the rolling + /// compressor's retained tail is rebuilt at width `tokens - base`, and + /// `base` jumps a whole `ratio` at a group boundary while `tokens` climbs by + /// one, so the size cycles through `ratio` consecutive values (measured: + /// `4096 ร— {18,19,20,21}`). One pass warms one phase of that cycle; the + /// captured step lands on another phase and allocates a size the pool has + /// never held. `ratio` is 4 on the shipped V4 configuration. + /// + /// This is a floor, not the answer โ€” `grant_extra_deferred_pass` keeps + /// extending while the observed size profile is still growing, so a model + /// with a longer cycle is covered without retuning this number. + /// See `deferred_passes_remaining`. pub fn default_deferred_passes() -> u32 { std::env::var("ARC_GRAPH_DEFERRED_PASSES") .ok() .and_then(|v| v.parse::().ok()) - .unwrap_or(1) + .unwrap_or(4) + } + + /// Upper bound on passes granted by `grant_extra_deferred_pass`, from + /// `ARC_GRAPH_DEFERRED_MAX` (default 24). The bound is what keeps a model + /// whose allocation sizes grow monotonically with context length from + /// warming forever: it will exhaust the budget, and the capture-miss assert + /// will then refuse the capture and name the sizes. That refusal is the + /// useful result โ€” it says the fix is shape-constant buffers (or allocator + /// size-class bucketing), not more warmup. + pub fn default_deferred_max() -> u32 { + std::env::var("ARC_GRAPH_DEFERRED_MAX") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(24) } pub fn new_with_passes( @@ -176,6 +207,7 @@ impl CudaGraphRunner { verify_remaining: 0, verify_failed: false, deferred_passes_remaining: 0, + deferred_extra_budget: 0, }); } @@ -198,6 +230,7 @@ impl CudaGraphRunner { verify_remaining: Self::default_verify_replays(), verify_failed: false, deferred_passes_remaining: deferred_passes, + deferred_extra_budget: Self::default_deferred_max(), }) } @@ -273,6 +306,29 @@ output_trusted={} verify_remaining={} verify_failed={}", } } + /// Ask for one more deferred-free warm pass, spending from a fixed budget. + /// Returns whether one was granted. + /// + /// The caller grants a pass whenever the last one taught the allocator a + /// size it had not seen, and stops when a pass adds nothing. That turns + /// "how many passes cover the allocation cycle?" from a constant that has + /// to be right into an observation that terminates on its own, while + /// `deferred_extra_budget` guarantees termination even when the sizes never + /// converge. + pub fn grant_extra_deferred_pass(&mut self) -> bool { + if self.deferred_extra_budget == 0 { + return false; + } + self.deferred_extra_budget -= 1; + self.deferred_passes_remaining += 1; + true + } + + /// Deferred passes still available to `grant_extra_deferred_pass`. + pub fn deferred_extra_budget(&self) -> u32 { + self.deferred_extra_budget + } + /// How many replays must be proven against an eager forward before the /// graph's output is trusted. `ARC_GRAPH_VERIFY_REPLAYS`, default 3. pub fn default_verify_replays() -> u32 { diff --git a/mistralrs-core/src/pipeline/normal.rs b/mistralrs-core/src/pipeline/normal.rs index b93ae48df..88d9ce669 100644 --- a/mistralrs-core/src/pipeline/normal.rs +++ b/mistralrs-core/src/pipeline/normal.rs @@ -2275,6 +2275,12 @@ impl Pipeline for NormalPipeline { // per-forward alloc count (eager warmups only // reach peak-live; capture needs every alloc // distinct). Output is this step's logits (eager). + // Snapshot the allocator's size profile so we can see + // whether this pass teaches it anything new. + let demand_before = match self.device() { + candle_core::Device::Cuda(cd) => cd.capture_alloc_demand(), + _ => Vec::new(), + }; if let candle_core::Device::Cuda(cd) = self.device() { cd.set_capture_mode(true); } @@ -2312,6 +2318,40 @@ impl Pipeline for NormalPipeline { if let candle_core::Device::Cuda(cd) = self.device() { cd.set_capture_mode(false); } + // Keep warming while the profile is still growing. + // One pass is only enough if every decode step + // allocates the same sizes; V4's compressor tail + // cycles through `ratio` widths, so pass k presents a + // size pass k-1 never did and the profile keeps + // growing until the whole cycle has been walked. A + // pass that adds nothing is the signal to stop โ€” and + // it costs one eager decode step to learn that. + if let candle_core::Device::Cuda(cd) = self.device() { + let demand_after = cd.capture_alloc_demand(); + if demand_after != demand_before { + let granted = runner.grant_extra_deferred_pass(); + tracing::info!( + "ARC prewarm: warm pass grew the alloc profile to {} \ + distinct sizes โ€” {} (extra budget {}, {} pass(es) owed)", + demand_after.len(), + if granted { + "granting another pass" + } else { + "extra budget exhausted; the capture-miss assert \ + decides from here" + }, + runner.deferred_extra_budget(), + runner.deferred_passes_remaining(), + ); + } else { + tracing::info!( + "ARC prewarm: alloc profile CONVERGED at {} distinct \ + sizes ({} pass(es) still owed)", + demand_after.len(), + runner.deferred_passes_remaining(), + ); + } + } match out { Ok(o) => { tracing::info!( @@ -2377,6 +2417,55 @@ impl Pipeline for NormalPipeline { (PROBE โ€” not a fix; the buffer still allocates per step)" ); } + // โ”€โ”€ EXACT PRE-WARM โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + // The warm passes leave the free pool holding + // exactly what the steps they ran on asked for. The + // step about to be CAPTURED is a different step: any + // size whose demand is higher here, or that only + // appears here, is served by `cuMemAllocAsync` from + // the graph's private pool and becomes an unstable + // graph memory node โ€” CUDA_ERROR_ILLEGAL_ADDRESS on + // the first cuGraphLaunch, and a poisoned context + // (reported as glibc heap corruption) after it. + // + // So top every observed size up to its observed + // demand plus slack. Slack is what covers a captured + // step that allocates a little more of some size + // than any warm step did; it is buffers, not bytes, + // and the pool is ~315 MiB, so it is cheap. + // + // Must run with capture mode OFF: a buffer released + // while capturing parks in `deferred` and cannot be + // served, so seeding inside the window warms nothing. + if let candle_core::Device::Cuda(cd) = self.device() { + let slack = std::env::var("ARC_GRAPH_PREWARM_SLACK") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(4); + let demand = cd.capture_alloc_demand(); + let want: usize = demand.iter().map(|(_, n)| n).sum(); + match cd.prewarm_alloc_cache(slack) { + Ok((sizes, buffers)) => tracing::info!( + "ARC prewarm: profile = {} distinct sizes / {want} \ + allocations per captured forward; topped up {sizes} \ + size(s) with {buffers} buffer(s) at slack={slack}", + demand.len(), + ), + Err(e) => { + tracing::error!( + "ARC prewarm: could not top up the alloc cache: {e}. \ + Capturing now would record unstable memory nodes, so \ + capture is disabled for this run." + ); + runner.disable(); + } + } + // Scope the miss ledger to THIS capture. Without + // the reset it would still carry the misses the + // warm passes made on purpose, and those are + // harmless โ€” no graph is being recorded then. + cd.reset_capture_misses(); + } // CAPTURE: frees are deferred so every allocation // is a stable cache hit (no within-capture // aliasing, no unstable graph memory nodes). @@ -2403,127 +2492,178 @@ impl Pipeline for NormalPipeline { &flash_meta, ) { Ok(output) => { - tracing::info!( - "ARC capture: V4 forward RECORDED (bs={bs}); instantiating + launching" + // โ”€โ”€ THE GATE โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + // Every allocation that missed the + // warm pool during this capture was + // served from the graph's private + // pool and recorded as a memory node + // whose address the first launch has + // to back. One is enough to fault. + // Refusing here is strictly better + // than instantiating: the fault is + // asynchronous and poisons the CUDA + // context, taking the process with + // it some allocations later. + let misses = match self.device() { + candle_core::Device::Cuda(cd) => { + cd.capture_misses() + } + _ => Vec::new(), + }; + if !misses.is_empty() { + let total: usize = + misses.iter().map(|(_, n)| n).sum(); + tracing::error!( + "ARC capture: REFUSING to instantiate โ€” \ + {total} allocation(s) missed the warm pool \ + during capture, at (bytes, count) {misses:?}. \ + Each is an unstable graph memory node. The \ + pre-warm profile did not cover this step, so \ + these sizes are step-dependent beyond the \ + warm cycle: the fix is a shape-constant \ + buffer at those sizes (or allocator \ + size-class bucketing), not more warmup. \ + Falling back to eager." ); - match runner.end_capture_and_cache(bs, output, gp, op) { - Ok(out) => { - tracing::info!( + // Drop the captured output BEFORE + // cancelling: its storage may live + // in the private pool, and + // cancel_capture drains the alloc + // cache before destroying that + // pool. Dropping after the drain + // would file a dangling pointer + // into the free list โ€” the exact + // use-after-free the drain exists + // to prevent. + drop(output); + runner.cancel_capture(gp, op); + runner.disable(); + None + } else { + tracing::info!( + "ARC capture: V4 forward RECORDED (bs={bs}), zero capture-time allocator misses; instantiating + launching" + ); + match runner + .end_capture_and_cache(bs, output, gp, op) + { + Ok(out) => { + tracing::info!( "ARC capture: graph CAPTURED + launched OK" ); - if arc_layer_bisect() { - let bufs = - crate::layers::arc_layer_trace_take(); - tracing::error!( - "ARC BISECT: holding {} captured layer buffers", - bufs.as_ref() - .map(|b| b.len()) - .unwrap_or(0) - ); - ARC_GRAPH_LAYER_BUFS - .with(|c| *c.borrow_mut() = bufs); - } - // โ”€โ”€ SAME-STEP PROBE โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - // Separates the two ways a graph can - // be wrong, which the replay-step - // comparison cannot tell apart. - // - // Here the position buffer, the token - // ids and the KV ring are EXACTLY what - // capture recorded against โ€” nothing - // has advanced. So: - // ฮ” โ‰ˆ 0 โ‡’ the graph is wired right - // and the replay-step ฮ” is - // stale per-step state. - // ฮ” large โ‡’ a plain wiring fault - // (wrong buffer / wrong - // node); position state is - // innocent. - // `error!` on purpose: this is the - // measurement the run exists for and - // must never be filtered out. - if std::env::var("ARC_GRAPH_SAMESTEP") - .as_deref() - == Ok("1") - { - if let candle_core::Device::Cuda(cd) = - self.device() - { - cd.set_capture_mode(false); + if arc_layer_bisect() { + let bufs = + crate::layers::arc_layer_trace_take(); + tracing::error!( + "ARC BISECT: holding {} captured layer buffers", + bufs.as_ref() + .map(|b| b.len()) + .unwrap_or(0) + ); + ARC_GRAPH_LAYER_BUFS + .with(|c| *c.borrow_mut() = bufs); } - match self.model.forward( - &input_ids, - &seqlen_offsets, - context_lens.clone(), - position_ids.clone(), - paged_attn_meta - .as_ref() - .map(|(a, b)| (a.clone(), b)), - &flash_meta, - ) { - Ok(eager_same) => { - match max_abs_diff( - &out, - &eager_same, - ) { - Ok(d) => tracing::error!( - "ARC SAMESTEP: first-launch vs eager, SAME step: max|ฮ”|={d:.3e}" - ), - Err(e) => tracing::error!( - "ARC SAMESTEP: first-launch compare failed: {e}" - ), - } - // Keep this step's eager - // logits: every later replay - // is measured against them - // too, which is what tells - // "frozen" from "partially - // advancing". - ARC_CAPTURE_STEP_EAGER.with( - |c| { - *c.borrow_mut() = - Some( - eager_same - .clone(), - ); - }, - ); - // Second launch of the same - // graph, still same step. - // Distinguishes "the graph - // is wrong" from "re-launch - // specifically is wrong". - match runner.replay(bs) { - Ok(again) => { - match max_abs_diff( - &again, - &eager_same, - ) { - Ok(d) => tracing::error!( - "ARC SAMESTEP: re-launch vs eager, SAME step: max|ฮ”|={d:.3e}" - ), - Err(e) => tracing::error!( - "ARC SAMESTEP: re-launch compare failed: {e}" - ), + // โ”€โ”€ SAME-STEP PROBE โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + // Separates the two ways a graph can + // be wrong, which the replay-step + // comparison cannot tell apart. + // + // Here the position buffer, the token + // ids and the KV ring are EXACTLY what + // capture recorded against โ€” nothing + // has advanced. So: + // ฮ” โ‰ˆ 0 โ‡’ the graph is wired right + // and the replay-step ฮ” is + // stale per-step state. + // ฮ” large โ‡’ a plain wiring fault + // (wrong buffer / wrong + // node); position state is + // innocent. + // `error!` on purpose: this is the + // measurement the run exists for and + // must never be filtered out. + if std::env::var("ARC_GRAPH_SAMESTEP") + .as_deref() + == Ok("1") + { + if let candle_core::Device::Cuda(cd) = + self.device() + { + cd.set_capture_mode(false); + } + match self.model.forward( + &input_ids, + &seqlen_offsets, + context_lens.clone(), + position_ids.clone(), + paged_attn_meta + .as_ref() + .map(|(a, b)| (a.clone(), b)), + &flash_meta, + ) { + Ok(eager_same) => { + match max_abs_diff( + &out, + &eager_same, + ) { + Ok(d) => tracing::error!( + "ARC SAMESTEP: first-launch vs eager, SAME step: max|ฮ”|={d:.3e}" + ), + Err(e) => tracing::error!( + "ARC SAMESTEP: first-launch compare failed: {e}" + ), + } + // Keep this step's eager + // logits: every later replay + // is measured against them + // too, which is what tells + // "frozen" from "partially + // advancing". + ARC_CAPTURE_STEP_EAGER.with( + |c| { + *c.borrow_mut() = + Some( + eager_same + .clone(), + ); + }, + ); + // Second launch of the same + // graph, still same step. + // Distinguishes "the graph + // is wrong" from "re-launch + // specifically is wrong". + match runner.replay(bs) { + Ok(again) => { + match max_abs_diff( + &again, + &eager_same, + ) { + Ok(d) => tracing::error!( + "ARC SAMESTEP: re-launch vs eager, SAME step: max|ฮ”|={d:.3e}" + ), + Err(e) => tracing::error!( + "ARC SAMESTEP: re-launch compare failed: {e}" + ), + } } + Err(e) => tracing::error!( + "ARC SAMESTEP: second launch failed: {e}" + ), } - Err(e) => tracing::error!( - "ARC SAMESTEP: second launch failed: {e}" - ), } + Err(e) => tracing::error!( + "ARC SAMESTEP: eager forward failed: {e}" + ), } - Err(e) => tracing::error!( - "ARC SAMESTEP: eager forward failed: {e}" - ), } + Some(out) } - Some(out) - } - Err(e) => { - tracing::warn!( + Err(e) => { + tracing::warn!( "ARC capture: instantiate/launch failed: {e}; eager" ); - None + None + } } } } From f2e71eff863fe23297942595d76d23e12ae6b46a Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 19:22:03 +0100 Subject: [PATCH 10/21] =?UTF-8?q?perf(v4):=20pin=20the=20xs=20window=20(+2?= =?UTF-8?q?3.2%=20at=20uniform=20B=3D32)=20+=20the=20ladder=20that=20price?= =?UTF-8?q?s=20per-sequence=20advance=20=E2=80=94=204.46x=20at=20spread=20?= =?UTF-8?q?B=3D8=20[MEASURED]=20(#121)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * measure(sched+arcspec): the throughput ladder #116 unlocks, and the confound it has to be read against is worth anything: MTP was measured at 1.93 tok/step at one user collapsing to 1.06 at 128, and that collapse is what this stack targets. Two pieces, because the number is unreadable without the second. 1. `scheduler::bucket_telemetry` โ€” a `SCHED[agg]` marker carrying buckets_per_step, running_bucket_size and offered_per_step, emitted on the SAME log fence as `MTP[agg]`. Both schedulers bucket the running set by exact cache length and run one bucket per step, preempting the rest. So a cell labelled B=128 can be a 3-wide step in the engine, and an aggregate number measured inside that is a measurement of the scheduler rather than of KV advance. Ragged admission exists precisely to collapse those buckets, so the two effects are confounded by construction โ€” without these counters "aggregate did not move" is unattributable. running_bucket_size is the width that actually ran; offered_per_step is the width the harness thinks it asked for. 2. `arc-tools/arcspec_perseq_ladder.sh` โ€” B = 1, 8, 32, 128, both arms, one server per arm, counters differenced across each cell's own wall-clock fence. It reports aggregate tok/s AND tok_per_step AND the decomposition (tok_per_batch_step x batch_steps/s, plus ms/step), because tok/step rising while aggregate falls has already happened on this chain and either number alone is unreadable. Prompt lengths are RAGGED by construction โ€” 24..320 words cycled across workers โ€” which is the deliberate difference from `arcspec_perseq_ab.sh`'s fixed 40. Uniform prompts hide the failure mode this stack addresses: the `xs` window defect could not even be reached with equal-length prompts. 320 keeps a >3x margin to the ~1,055-word serving cliff this branch does not carry the fix for. Co-Authored-By: Claude Opus 5 (1M context) * measure(arcspec): add the uniform arm โ€” the only regime where the batch is whole The scheduler A/B came back while this was being built and it changes what the ladder can conclude. The bucketing law is now measured, not hypothesised: running bucket size = B / (distinct cache lengths) holding at 8/8=1 and 32/8=4, with `1 running, 7 waiting` sustained, and B=8 on spread lengths measuring 7.91 tok/s against B=1's 15.36 โ€” batching is NEGATIVE on realistic traffic. So a spread-only ladder measures per-sequence KV advance inside a scheduler that is running one sequence at a time, and a flat aggregate there is unreadable: it cannot distinguish "the fix does not pay" from "the batch never existed". Both regimes are now run, per arm, mean-matched: spread 144 24 320 64 260 40 200 96 words (8 distinct lengths) uniform 144 words โ€” the spread's mean AND its first element, so B=1 is byte-identical between regimes and the only thing that differs at B>1 is the spread itself Uniform is where a whole batch actually forms, so it is the only place this stack's effect on aggregate throughput is visible without the serialisation swamping it. Spread is still the regime the stack exists for. The report now prints the bucketing law's prediction beside the measured `running_bucket_size` per cell, and refuses to let a spread cell be read as a batch result when its running bucket is ~1. Co-Authored-By: Claude Opus 5 (1M context) * perf(v4): stop reallocating the xs window every decode step โ€” restore the bound the type already documents `tail` is described at the top of this file as "the raw rows behind `tokens`, bounded by `span_groups * ratio + margin` and independent of context length". That is an invariant the type has always claimed. The allocation did not honour it: `advance` rebuilt the buffer with `cat` + `narrow` every step at width exactly `tokens - base`, which 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 on every decode step, on every one of V4's 41 compressed layers. This is not "add a pin". It is restoring a documented invariant that one code path broke. THE BOUND, WHICH IS THE REVIEWABLE PART 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 keep_from = (q + 1 - span_groups) * ratio > m - span_groups * ratio, giving 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 provably sufficient and provably never truncates. It is a function of the layer's geometry, not of context length: 24 for CSA (ratio 4, span 2, margin 16), 144 for HCA. That is the difference between this and a capacity that turns out to be a million-token context. Derived twice independently and agreeing, and checked a third way: the steady band is measured at [capacity - ratio, capacity - 1] = {20..23}. THE {18..21} vs {20..23} DISCREPANCY, SETTLED RATHER THAN ASSERTED ArcGraph measured 4096 x {18,19,20,21} from outside the engine; the retention rule predicts {20,21,22,23}. Both ratio-consecutive, both containing 21, offset by 2. The pre-committed reconciliation was that theirs is the pre-saturation ramp (base still at 0 while tokens climbs) and mine the steady state. `the_window_ramps_then_settles_to_ratio_consecutive_sizes` checks exactly that: the early band sits below the steady one, the steady one is `ratio` consecutive sizes at [cap-ratio, cap-1], and no width over 128 steps reaches `cap`. A bound that is right for the wrong reason is a trap for the next context length. WHY DEFAULT ON WITHOUT A THROUGHPUT NUMBER FIRST The measure-then-default-on rule exists because FP8 KV shipped on and changed 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. `pinning_the_window_is_numerically_inert` runs both settings against one stream for 40 steps and requires exact equality of the compressed rows and both time bases โ€” and asserts the unpinned widths actually varied, so the equality is not vacuous. `ARC_V4_XS_PIN_WINDOW=0` restores the resizing buffer, which is also the A/B arm. split_row keeps the pinned buffer rather than re-narrowing it: `clone_out_cache` calls it once per layer per sequence on EVERY engine step, so narrowing there would undo the pin exactly on the hot path. The resume point does not move. Two fixtures now run explicitly unpinned, because pinning removes their discriminator rather than their subject: `ragged_xs_tail_is_refused_by_name_not_panicked` needs the widths 18/22 that only a resizing buffer produces, and `splitting_a_batched_row_restores_the_per_sequence_window` needs a row narrower than the shared window. Both keep testing what they were written for, and both gained a pinned-side counterpart. Co-Authored-By: Claude Opus 5 (1M context) * measure(arcspec): price the window pin in the same session โ€” a third arm ON vs ON_UNPINNED differ in exactly one thing: whether the compressor's raw window is reallocated every decode step on all 41 compressed layers, or held at the bound the type documents. Same per-sequence flags in both, so the pin is isolated rather than confounded with the KV mechanism, and it is priced in the session that was already going to queue for the box. On the width discrepancy the pin rests on: this harness does not need to settle it. `max_tokens` is 4096 and each cell drives a 45 s steady window, so every request is thousands of tokens past saturation and the ramp is invisible here anyway. The settling is measured directly instead, on CPU, over 128 steps, by `the_window_ramps_then_settles_to_ratio_consecutive_sizes` โ€” the early band sits below the steady one, the steady one is `ratio` consecutive sizes at [cap-ratio, cap-1], and no width ever reaches `cap`. That is a stronger instrument than a throughput leg and it costs no card time. Co-Authored-By: Claude Opus 5 (1M context) * instr(v4): name the window mode once per process, and assert the pin A/B is real D18, applied to the arm I am about to run. `ARC_V4_XS_PIN_WINDOW=0` is the only thing separating ON from ON_UNPINNED, and if that name were wrong the two arms would be the same build, produce a 1.000x ratio, and read as "the pin costs nothing" โ€” the granted-but-inert failure with a throughput number attached. So the flag names itself once per process on first read ("xs rolling window is PINNED / RESIZING"), following the same convention as "per-sequence KV advance is ON", and the ladder summary now requires PINNED in the ON log and RESIZING in the ON_UNPINNED log before either arm's ratio may be read. If it cannot find both it prints PIN A/B IS VOID and says the comparison is of one thing with itself, rather than printing a ratio. Co-Authored-By: Claude Opus 5 (1M context) * measure(arcspec): make the pin control provably the control, not the default The window pin defaults ON, so the TREATMENT is the default and the CONTROL is the arm carrying `ARC_V4_XS_PIN_WINDOW=0`. That inverts the usual failure: a typo in the flag name does not disable the feature, it produces a control that silently ran the treatment and a clean-looking 1.000x ratio. Three things, all of which had to be right and only one of which was obvious: * The assertion reads the runtime LOG, never the binary. Both mode strings are compiled in unconditionally, so `strings | grep RESIZING` succeeds in every arm and proves nothing. Only the line a process emits says which branch it took. * `RUST_LOG=info` is set once inside `run_arm`, shared by every arm, so a log filter cannot suppress the line in one arm while leaving it in another. An assertion that configuration can mute is a guard with an undocumented off switch. * The guard now also fails when the control's log contains PINNED, not just when it lacks RESIZING โ€” the leak is the thing being tested for, so it is checked directly rather than inferred from an absence. On VOID it refuses to let the ratio be read at all rather than printing it with a caveat. Co-Authored-By: Claude Opus 5 (1M context) * fix(ArcKV): one xs window pin, two triggers; scope the unpinned tests Integration fix for #121 landing after #181's `pin_tail_width`. Both arrived at the same idea โ€” hold the retained raw xs window at a constant width instead of letting it breathe with `tokens % ratio` โ€” for two different reasons, and both callers are real: * #181: CUDA-graph capture, where a moving allocation size is not slow but INVALID (a capture-time miss becomes an unstable graph memory node). Per-cache trigger, `pin_tail_width()`. * #121: serving, where the per-step reallocation is simply waste. Process-wide trigger, `ARC_V4_XS_PIN_WINDOW`. They agreed on the width and disagreed on everything else, so this keeps one width policy and both triggers: 1. `graph_tail_width()` and `window_capacity()` were the same expression (`span_groups * ratio + margin`) written twice. `graph_tail_width` now delegates. Two pins that computed "the pinned width" independently could drift, and capture and serving would then disagree about a number whose only value is that it does not move. 2. `retained_width` is gated on `pin_is_on()` = `self.pin_tail || xs_pin_window_enabled()`, so capture still gets its constant width when the env switch is off โ€” which, for a capture-only run, it is. 3. #181's `pinned_base` form is dropped in favour of #121's. Both produce the same constant width; they differed in where the constancy came from. #181 moved `base` earlier so the logical span went constant; #121 leaves `base` at the retention point the type documents and widens the physical buffer. #121's is the one that composes, because it also moved `win_start` to the buffer's PHYSICAL start (`tokens - w_phys`) rather than `base`, which makes "buffer wider than the row promised" a representable state. `plan_xs_advance` still refuses to read below `base`, so the pin cannot change an answer. 4. Three `clone_in_cache_invariant_tests` assert the geometry of the UNPINNED window โ€” exact tail widths `(4, 132)`, and that column 0 holds token `base`. Neither survives a pinned buffer, and #121 makes the pin the DEFAULT. They now run under `unpinned(..)`, the same `pin_test_override` shape #121 uses for its own affected tests. They are not testing a dead path: `ARC_V4_XS_PIN_WINDOW=0` is a supported mode and is the control arm of #121's own A/B. NOTE FOR REVIEW, not a defect: #121 flips the serving default โ€” the window is PINNED unless `ARC_V4_XS_PIN_WINDOW` is set to 0. Its "+23.2% at uniform B=32" was measured on #121's own branch and is NOT re-measured here; this rebase changes which code that number describes. `cargo test -p mistralrs-core --lib`: 701 passed, 0 failed. Warning count back to master's 23 (the `xs_rolling` import is test-only and now lives in the test module). Co-Authored-By: Claude Opus 5 (1M context) * fix(ArcKV): the xs window pin is OPT-IN, with the experiment that flips it named #121 shipped `ARC_V4_XS_PIN_WINDOW` default-ON, with `=0` as the control arm. Retracting the default and keeping the change. Why: the pin 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 motivated defaulting it on was measured on #121's own pre-rebase branch, against a tree that no longer exists. That is an unmeasured default. TCFRAG was an unmeasured default too: it carried "UNVERIFIED ON HARDWARE โ€” NEVER RUN" in its own header, held 63 GB and permanently broke a layer through a poisoned `OnceLock` (#209). So: unverified means default-off, and "unverified" means unmeasured, not new. The correctness argument #121 made for the pin is SOUND and is kept in place โ€” it proves the pin changes no answer, which is necessary and not sufficient. The 40-step bit-identity A/B in `deepseek4` still guards it; its doc comment no longer claims to justify a default. ## 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. So this does NOT ship as a parked flag. The flip condition is written three places, one of them enforced: * `xs_pin_window_enabled_from`'s doc comment carries it in full, under a "FLIP CONDITION" heading, at the gate itself. * `capability_reachability.rs` โ€” the file titled "the switched-off guard", which is where someone goes looking for exactly this โ€” carries it as a registry entry, so the gate cannot go dark without CI going red. * `arcspec_perseq_ladder.sh`'s `ON_PINNED` arm IS the experiment. The experiment: one binary, uniform B=32, same prompt and seed, `ON` (flag unset) against `ON_PINNED` (`ARC_V4_XS_PIN_WINDOW=1`). Pass = ON_PINNED faster on aggregate tok/s with identical generated tokens. On pass the default flips to ON in the same change that records the number. ## The harness inverted with the default, and that is easy to get silently wrong The pin used to be the default, so the A/B's TREATMENT was the default arm and its CONTROL carried the flag. Now it is the other way round. Three things had to move together, and any one left behind would have produced a clean-looking and meaningless number: 1. The arm carries `=1` and is renamed `ON_UNPINNED` -> `ON_PINNED`. 2. The report's ratio pairing is swapped, so it stays treatment/control. Renaming the arm without this would have inverted every printed ratio. 3. The engagement guard's expectations are swapped: `ON` must log RESIZING, `ON_PINNED` must log PINNED, and the leak check now looks for a treatment that silently ran the control. Polarity is `== Some("1")`, split into the pure `xs_pin_window_enabled_from` and pinned by tests: unset is OFF, and `0`/`false`/`off`/`true`/`on`/`2`/`" 1"` are all OFF. #212 converted 23 `ARC_*` flags for this reason โ€” `var_os(..).is_some()` made `ARC_FOO=0` mean ON, which turns an A/B control into a second treatment. A third test pins the thing most likely to be broken by this retraction: CUDA-graph capture must still get a constant width with the env flag off. It does โ€” capture asks per-cache via `pin_tail_width()` and `pin_is_on` honours that trigger independently. Without that, capture would silently return to a per-step-varying allocation size, which is not a slow graph but an invalid one. `cargo test -p mistralrs-core -p mistralrs-quant -p mistralrs-vision`: green. Co-Authored-By: Claude Opus 5 (1M context) --------- Co-authored-by: Claude Opus 5 (1M context) --- arc-tools/arcspec_perseq_ladder.sh | 544 ++++++++++++++++++ mistralrs-core/src/kv_cache/mod.rs | 90 ++- mistralrs-core/src/kv_cache/xs_rolling.rs | 466 +++++++++++++-- mistralrs-core/src/lib.rs | 4 +- mistralrs-core/src/models/deepseek4.rs | 218 +++++++ mistralrs-core/src/pipeline/mtp_pipeline.rs | 10 + .../src/scheduler/bucket_telemetry.rs | 241 ++++++++ .../src/scheduler/default_scheduler.rs | 9 + mistralrs-core/src/scheduler/mod.rs | 2 + .../tests/capability_reachability.rs | 33 ++ 10 files changed, 1554 insertions(+), 63 deletions(-) create mode 100755 arc-tools/arcspec_perseq_ladder.sh create mode 100644 mistralrs-core/src/scheduler/bucket_telemetry.rs 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", From 84aab24b818462d79dc9b0c233d351c46f74fb70 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 19:22:12 +0100 Subject: [PATCH 11/21] =?UTF-8?q?perf(ArcInfer/ArcGraph):=20device=20decod?= =?UTF-8?q?e=20loop=20=E2=80=94=202=20host=20syncs=20per=20token=20to=200,?= =?UTF-8?q?=20opt-in=20(#217)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * perf(ArcInfer/ArcGraph): device decode loop โ€” N steps, N cuGraphLaunch, zero host syncs CUDA-graph capture already cut launch APIs per token 3,961 -> 3.8 and bought only ~8%. The reason is not that op count stopped mattering: replay() launches the graph and then blocks on cudaStreamSynchronize (graph.rs:362), the caller samples on the host, and the host-side argmax D2H (sampler.rs:1479) re-synchronizes anyway. Capture removed the cost of *issuing* the work and left the serialized host round trip between steps completely intact. Removing the sync alone buys nothing, so this removes the sync and the host sample together, behind an opt-in. Per step, all on one stream: cuGraphLaunch -> CudaSampler (on device, token to an i32 device buffer, never copied down) -> arc_graph_step_commit (scatters the token into the PINNED U32 buffer the captured graph reads, advances the device position, publishes into a pinned+mapped host ring). Nothing blocks. Fault detection replaces the sync rather than dropping it, with three signals that need no blocking call: non-blocking cudaStreamQuery (which NVIDIA documents as also returning errors from previous asynchronous launches), an idle-but-short ring (every step retired yet fewer tokens landed โ€” a fault with no timeout heuristic in it), and a sticky device fault word the commit kernel latches before it writes anything. A spin backstop covers a kernel that hangs without ever erroring. The driver is expressed against a DeviceOps trait that deliberately has no synchronize method, so it is structurally incapable of blocking the stream, and it compiles without CUDA โ€” 20 unit tests run against a scripted fake GPU on any host. Two mutations were checked: dropping the device-fault read and skipping one launch per burst each fail exactly the tests that name them. Aliasing: the logits tensor aliases storage the next launch overwrites. Safe here only because every access is stream-ordered and none is a host read; the ring the host drains is a separate pinned allocation no launch touches. The invariant that lets the aliasing tensor be returned at all is enforced at runtime, not just documented. Default OFF (ARC_GRAPH_DEVICE_LOOP). Greedy-only: the device sampler runs Splitmix64 against the host's Isaac64, so any stochastic mode would draw a different, equally valid token and silently break seeded reproducibility. Every refusal falls back to the existing path; one failure latches it off for the process. Kept as a pure insertion, away from replay(), to merge cleanly with PR #205. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01UMmjFy8TvsgypxVWNVhhC7 * fix(ArcGraph): device loop stands aside for ARC_GRAPH_VERIFY_ALWAYS The device-loop branch sits earlier in the chain than the runner.replay(bs) arm that PR #205 extended with ARC_GRAPH_VERIFY_ALWAYS. Without this it would take every step and that switch would silently verify nothing โ€” a clean diagnostic that never ran. Also cross-references #205's BOOBY TRAP note from replay_handles, since the same aliasing applies to the handles it returns. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01UMmjFy8TvsgypxVWNVhhC7 * docs(ArcGraph): record the burst's scheduler-admission limitation A burst admits no new request for its whole length (default 4 decode steps). Harmless at batch 1 โ€” the only shape admit() accepts โ€” and not harmless above it, which is why admit() refuses batch_size != 1 outright rather than treating it as a tuning question. Stated in the module docs so it survives past the return it was first noted in. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01UMmjFy8TvsgypxVWNVhhC7 * fix(ArcGraph): the device-loop sampling hook takes &mut Sequence The CUDA lane's whole purpose, on its first run against this branch: three jobs, one error, in the one block a macOS `cargo check` structurally cannot reach. error[E0596]: cannot borrow `*seq` as mutable, as it is behind a `&` reference --> mistralrs-core/src/pipeline/sampling.rs:557:12 557 | && seq.sampler().is_greedy_trivial() `Sequence::sampler()` takes `&mut self` (sequence.rs:885) even though it only clones an Arc, so the `&Sequence` parameter could never call it. `sample_sequence` already holds `&mut Sequence` and was reborrowing it down to `&`, which is why the call site compiled and the body did not. This is rustc's own suggested fix. Nothing here mutates the sequence, and the `&&` chain still borrow-checks: `sampler()` returns an owned `Arc`, so the mutable borrow ends before `seq.recognizer` is read. Verified by compiling the cuda body on a host with no CUDA โ€” cfg removed, the three `arc_cuda_graph` calls stubbed โ€” which puts the real `Sequence` through borrowck and passes. The `#[cfg(not(feature = "cuda"))]` twin moves with it so the two signatures cannot drift. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01UMmjFy8TvsgypxVWNVhhC7 --------- Co-authored-by: Claude Opus 5 (1M context) --- arc-cuda-graph/src/cuda/decode_loop.cu | 112 ++ arc-cuda-graph/src/device_loop.rs | 1797 +++++++++++++++++++++++ arc-cuda-graph/src/ffi.rs | 7 + arc-cuda-graph/src/graph.rs | 43 + arc-cuda-graph/src/lib.rs | 13 + mistralrs-core/src/pipeline/normal.rs | 217 +++ mistralrs-core/src/pipeline/sampling.rs | 67 + mistralrs-core/src/sampler.rs | 22 + 8 files changed, 2278 insertions(+) create mode 100644 arc-cuda-graph/src/device_loop.rs diff --git a/arc-cuda-graph/src/cuda/decode_loop.cu b/arc-cuda-graph/src/cuda/decode_loop.cu index b18195029..4a4894d3a 100644 --- a/arc-cuda-graph/src/cuda/decode_loop.cu +++ b/arc-cuda-graph/src/cuda/decode_loop.cu @@ -152,3 +152,115 @@ extern "C" void launch_check_all_done_conditional( extern "C" int arc_has_graph_conditional() { return ARC_HAS_GRAPH_CONDITIONAL; } + +// ============================================================================ +// ArcInfer/ArcGraph โ€” device decode loop step commit +// +// The narrow sibling of `decode_step_update` above, for the CUDA-graph replay +// path (`arc-cuda-graph/src/device_loop.rs`). That kernel is written for the +// paged-attention autonomous runner and unconditionally dereferences +// `finished`, `n_generated`, `slot_mappings` and `block_tables`; the replay +// path has none of those (V4 runs with `cache_config == None`), so passing +// nulls would fault. This kernel touches only what that path actually owns. +// +// It closes the decode loop on-device: +// sampled token -> the PINNED U32 buffer the captured graph reads +// -> device position += 1 +// -> pinned+mapped host ring, for the host to drain later +// +// Ordering is load-bearing. Every range check runs BEFORE any write, and the +// ring write head is published LAST, after `__threadfence_system()`. So a token +// the host can see in the ring is guaranteed to have been committed to the +// graph input buffer, and a refused commit publishes nothing at all rather than +// letting the next launch consume a garbage id. +// ============================================================================ + +#define ARC_COMMIT_FAULT_NONE 0 +#define ARC_COMMIT_FAULT_TOKEN_RANGE 1 +#define ARC_COMMIT_FAULT_POSITION_LIMIT 2 + +__global__ void arc_graph_step_commit( + const int32_t* __restrict__ sampled, // [batch] i32, written by CudaSampler + uint32_t* __restrict__ input_ids, // [batch] U32, PINNED graph input + uint32_t* __restrict__ positions, // [batch] U32, PINNED graph input; nullable + int32_t* __restrict__ ring, // [batch, ring_size] pinned + mapped + int32_t* __restrict__ ring_head, // [batch] pinned + mapped + int32_t* __restrict__ fault, // [1] pinned + mapped, sticky + int32_t ring_size, + int32_t vocab, + int32_t position_limit, // exclusive; <= 0 disables the check + int32_t batch_size +) { + int bid = blockIdx.x; + if (threadIdx.x != 0) return; + if (bid >= batch_size) return; + + // A fault latched by an earlier step means the graph input buffer already + // holds a stale id. Committing on top of that would paper over it. + if (fault != nullptr && ((volatile int32_t*)fault)[0] != ARC_COMMIT_FAULT_NONE) return; + + int token = sampled[bid]; + + // ---- checks first, writes after ------------------------------------- + if (token < 0 || token >= vocab) { + if (fault != nullptr) { + ((volatile int32_t*)fault)[0] = ARC_COMMIT_FAULT_TOKEN_RANGE; + __threadfence_system(); + } + return; + } + if (positions != nullptr && position_limit > 0) { + // `positions[bid]` is the slot this step just wrote KV to. The next + // step writes at +1, and the fixed-capacity window is `position_limit` + // wide, so +1 must stay strictly inside it. + if ((int64_t)positions[bid] + 1 >= (int64_t)position_limit) { + if (fault != nullptr) { + ((volatile int32_t*)fault)[0] = ARC_COMMIT_FAULT_POSITION_LIMIT; + __threadfence_system(); + } + return; + } + } + + // ---- commit --------------------------------------------------------- + // This single store is what removes the host from the loop: the captured + // graph baked THIS address at capture time, so the next `cuGraphLaunch` + // reads the token without any host round trip. + input_ids[bid] = (uint32_t)token; + if (positions != nullptr) { + positions[bid] += 1u; + } + + // ---- publish last ---------------------------------------------------- + if (ring != nullptr && ring_head != nullptr && ring_size > 0) { + int wp = ring_head[bid]; + int slot = wp % ring_size; + if (slot < 0) slot += ring_size; + ((volatile int32_t*)ring)[(int64_t)bid * ring_size + slot] = token; + // Order the token store ahead of the head store as seen by the HOST. + // Without this the host can observe an advanced head and read the + // previous occupant of the slot. + __threadfence_system(); + ((volatile int32_t*)ring_head)[bid] = wp + 1; + } +} + +extern "C" void arc_launch_graph_step_commit( + const int32_t* sampled, + uint32_t* input_ids, + uint32_t* positions, + int32_t* ring, + int32_t* ring_head, + int32_t* fault, + int32_t ring_size, + int32_t vocab, + int32_t position_limit, + int32_t batch_size, + cudaStream_t stream +) { + arc_graph_step_commit<<>>( + sampled, input_ids, positions, + ring, ring_head, fault, + ring_size, vocab, position_limit, batch_size + ); +} diff --git a/arc-cuda-graph/src/device_loop.rs b/arc-cuda-graph/src/device_loop.rs new file mode 100644 index 000000000..e2885182b --- /dev/null +++ b/arc-cuda-graph/src/device_loop.rs @@ -0,0 +1,1797 @@ +//! Parent system: ArcInfer / ArcGraph +//! +//! Device-resident decode loop: **N decode steps cost N `cuGraphLaunch` calls +//! and zero `cudaStreamSynchronize`.** +//! +//! # What was wrong with plain replay +//! +//! [`crate::graph::CudaGraphRunner::replay`] launches the captured graph and +//! then calls `cudaStreamSynchronize` (`graph.rs:362`) before handing the +//! output tensor back. The caller samples the next token on the **host**, which +//! copies logits (or an argmax scalar) down over PCIe, and only then launches +//! the next graph. So a decode step is +//! +//! ```text +//! GPU runs step -> host blocks on sync -> D2H -> host samples -> launch +//! ``` +//! +//! Capture removed the cost of *issuing* the work โ€” 3,961 launch APIs per token +//! became 3.8 โ€” and left that serialized round trip completely intact. Deleting +//! the sync on its own buys nothing, because the host-side sample forces a +//! device-to-host copy that synchronizes anyway. **The sync and the host sample +//! have to go together.** +//! +//! # The loop this module closes +//! +//! Per step, all on one stream, no host in the middle: +//! +//! 1. `cuGraphLaunch` โ€” the captured forward. +//! 2. [`crate::sampling_cuda::CudaSampler`] โ€” argmax/top-p **on device**, token +//! id written to an `i32[batch]` device buffer. Never copied to the host. +//! 3. `arc_graph_step_commit` โ€” scatters that token into the *pinned* graph +//! input buffer the captured graph reads (so the next launch consumes it), +//! bumps the device position, and publishes the token into a pinned+mapped +//! host ring. +//! +//! Nothing in that sequence blocks. The host enqueues all `burst` steps and the +//! GPU works through them while the host is already ahead. +//! +//! # How the host still gets its tokens +//! +//! The ring lives in `cudaHostAlloc(..., cudaHostAllocMapped)` memory. The GPU +//! pushes each token straight into it, `__threadfence_system()`, then publishes +//! a monotonic write head. The host reads both with **plain volatile loads** โ€” +//! no CUDA API call, no stream synchronize, no D2H copy. Draining is a host +//! memory poll, which is categorically different from a stream sync: with a +//! sync the host cannot enqueue step N+1 until step N retires; here every step +//! is already enqueued and the host is only watching them land. +//! +//! # Why the aliasing logits buffer is safe here +//! +//! The tensor `replay()` hands back **aliases graph-owned storage that the next +//! `cuGraphLaunch` overwrites**. Holding it across a launch, or draining it +//! lazily, reads data the next step already clobbered โ€” silently, producing +//! plausible tokens. That alias/clone class has already caused three +//! verification bugs in this repo. +//! +//! This design is safe against it for one specific reason, and it is worth +//! stating rather than assuming: **every access is enqueued on the same +//! stream, in order, and none of them is a host read.** Per step the stream +//! carries `launch N -> sample -> commit -> launch N+1`. CUDA stream ordering +//! guarantees the sampler kernel starts only after launch N has written the +//! logits and finishes before launch N+1 begins overwriting them. The host +//! never reads the logits at all, so there is no window in which it could read +//! a clobbered buffer. +//! +//! This is strictly stronger than the host path, which has to *synchronize* to +//! read those logits safely. What the host does drain โ€” the token ring โ€” is a +//! **separate, pinned, host-owned allocation** that no graph launch touches, so +//! draining it lazily is sound however far behind the host falls. The one thing +//! that would break the argument is issuing any of these on a different stream; +//! that is why `CudaDeviceOps` captures the graph's own stream and uses it for +//! the sampler and the commit kernel too. +//! +//! # Detecting a fault without synchronizing +//! +//! A fault *during* graph execution is asynchronous โ€” `cuGraphLaunch` returns +//! SUCCESS and only a later API call surfaces it. `replay()` used a full sync +//! for exactly this reason and the comment there says so. Dropping the check +//! outright would turn an illegal access into a silently poisoned context and +//! an output buffer full of garbage, a bug class this codebase has shipped more +//! than once. So the check is *replaced*, not removed, by three signals that +//! together need no blocking call: +//! +//! * **`cudaStreamQuery`** โ€” non-blocking. NVIDIA's CUDA Runtime API reference +//! for it states: *"Note that this function may also return error codes from +//! previous, asynchronous launches."* So a sticky async fault surfaces on a +//! call that never blocks. +//! * **Idle-but-short** โ€” `cudaStreamQuery` returning `cudaSuccess` means every +//! enqueued step has retired. If all of them retired and fewer tokens than +//! steps reached the ring, the work definitively did not do what it was asked +//! to. That is a fault with no timeout heuristic in it. +//! * **The device fault word** โ€” the commit kernel range-checks the sampled +//! token and the position before it writes anything, and latches a code into +//! a pinned word. It travels back with the ring, so reading it is free. +//! +//! `max_spins` is a last-resort liveness backstop for a kernel that hangs +//! without ever erroring. It is not the primary detector. +//! +//! # KNOWN LIMITATION โ€” the burst holds the scheduler +//! +//! A burst enqueues `burst` decode steps back-to-back and does not return to +//! the engine until they have all landed. **For that whole window the engine +//! admits no new request**, so a request arriving mid-burst waits up to +//! `burst` decode steps (default 4) before it is even considered. +//! +//! At batch 1 โ€” the only shape [`admit`] accepts โ€” that is harmless: there is +//! one sequence, nothing to interleave with, and the alternative is a host +//! round trip per token. It is **not** harmless above batch 1, where a burst +//! would straddle sequences with different stop conditions and different +//! admission urgency, and where the wasted tail after one sequence's EOS is +//! multiplied by the batch. That is why `admit` refuses `batch_size != 1` +//! outright rather than treating it as a tuning question. +//! +//! Anyone widening this past batch 1 owns that problem first. Raising `burst` +//! trades admission latency for host round trips one-for-one, and the +//! trade is only free while there is exactly one sequence in flight. +//! +//! # Layering +//! +//! Everything above the CUDA calls is expressed against the [`DeviceOps`] +//! trait, and that trait deliberately has **no `synchronize` method** โ€” the +//! driver cannot block the stream because it is not given a way to. The +//! resulting [`BurstDriver`] is plain Rust, compiles without the `cuda` +//! feature, and is unit-tested against a scripted fake GPU. The CUDA-gated +//! [`CudaDeviceOps`] is pointer plumbing and nothing else. + +use std::fmt; + +// =========================================================================== +// Configuration +// =========================================================================== + +/// Opt-in switch. Absent or `0` means this whole module is inert and decode +/// takes the existing `replay()` + host-sample path. +pub const ENV_ENABLE: &str = "ARC_GRAPH_DEVICE_LOOP"; +/// Steps enqueued per burst before the host drains. Default [`DEFAULT_BURST`]. +pub const ENV_BURST: &str = "ARC_GRAPH_DEVICE_LOOP_BURST"; + +/// Steps launched back-to-back before the host drains the ring. +/// +/// The whole burst is enqueued without any host round trip, so this is also +/// how far the GPU may run ahead of the host's `Sequence` bookkeeping โ€” and +/// therefore the worst-case detection latency, in steps, for a stop string, an +/// EOS, or an async fault. 4 keeps the wasted-work tail small while still +/// amortising the drain over four launches. +pub const DEFAULT_BURST: usize = 4; + +/// Tokens the pinned ring holds per row. Must exceed `burst`, or a burst could +/// overwrite tokens the host has not read yet. +pub const DEFAULT_RING_SIZE: usize = 1024; + +/// Host spins between `cudaStreamQuery` calls while draining. Pure host +/// memory polls are ~free; the driver query costs a syscall-ish amount, so it +/// is not worth doing on every iteration. +pub const DEFAULT_POLL_EVERY: u32 = 64; + +/// Emit [`DeviceDecodeLoop::status_line`] every this many bursts. The counter +/// it carries (`launches`) is the only proof the loop actually ran. +pub const STATUS_EVERY_BURSTS: u64 = 64; + +/// Liveness backstop: give up after this many spins even if the stream never +/// reports an error. Only trips on a kernel that hangs without faulting. +pub const DEFAULT_MAX_SPINS: u64 = 200_000_000; + +/// Is the device decode loop switched on? +/// +/// Default **off**. `ARC_GRAPH_DEVICE_LOOP=1` opts in; `0`, `false` and `off` +/// are explicit refusals so an operator can pin it off in an environment that +/// sets it upstream. +pub fn device_loop_enabled() -> bool { + match std::env::var(ENV_ENABLE) { + Ok(v) => { + let v = v.trim().to_ascii_lowercase(); + !(v.is_empty() || v == "0" || v == "false" || v == "off" || v == "no") + } + Err(_) => false, + } +} + +/// Burst length from [`ENV_BURST`], clamped to at least 1. +pub fn burst_from_env() -> usize { + std::env::var(ENV_BURST) + .ok() + .and_then(|v| v.trim().parse::().ok()) + .filter(|v| *v > 0) + .unwrap_or(DEFAULT_BURST) +} + +/// Everything the driver needs that does not change from burst to burst. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DeviceLoopConfig { + /// Steps per burst. Also the run-ahead bound and the detection latency. + pub burst: usize, + /// Tokens per row in the pinned ring. + pub ring_size: usize, + /// Host spins between stream queries while draining. + pub poll_every: u32, + /// Spin budget before declaring the stream stalled. + pub max_spins: u64, + /// Stop the burst's token stream here. `< 0` disables EOS truncation. + pub eos_token_id: i32, + /// Vocabulary size; the commit kernel range-checks against it. + pub vocab: usize, + /// Exclusive upper bound on the device position. The fixed-capacity graph + /// decode arm reads a constant `capacity`-wide window and writes KV at + /// `position`; advancing past it is an out-of-bounds write, so both the + /// driver and the kernel refuse to. + pub position_limit: u32, +} + +impl DeviceLoopConfig { + /// Config with the shipped defaults for everything but the model facts. + pub fn new(vocab: usize, eos_token_id: i32, position_limit: u32) -> Self { + Self { + burst: burst_from_env(), + ring_size: DEFAULT_RING_SIZE, + poll_every: DEFAULT_POLL_EVERY, + max_spins: DEFAULT_MAX_SPINS, + eos_token_id, + vocab, + position_limit, + } + } + + /// Reject a config that cannot be run safely, rather than discovering it + /// as corruption at step `ring_size + 1`. + pub fn validate(&self) -> Result<(), DeviceLoopError> { + if self.burst == 0 { + return Err(DeviceLoopError::BadConfig("burst must be non-zero")); + } + if self.vocab == 0 { + return Err(DeviceLoopError::BadConfig("vocab must be non-zero")); + } + if self.poll_every == 0 { + // The drain does `spins % poll_every`; zero would divide by zero. + return Err(DeviceLoopError::BadConfig("poll_every must be non-zero")); + } + if self.ring_size <= self.burst { + // A burst longer than the ring overwrites its own unread tokens. + return Err(DeviceLoopError::BadConfig( + "ring_size must exceed burst, or a burst overwrites unread tokens", + )); + } + Ok(()) + } + + /// Steps this burst may launch given how far the sequence has already got. + /// + /// Clamped by the position limit so the burst can never walk the device + /// position past the fixed-capacity window. Returns 0 when the sequence has + /// no room left, which the caller must treat as "fall back to eager", not + /// as "launch zero steps and carry on". + pub fn plan_steps(&self, current_position: u32, max_new_tokens: usize) -> usize { + let room = self.position_limit.saturating_sub(current_position) as usize; + self.burst.min(room).min(max_new_tokens) + } +} + +// =========================================================================== +// Errors +// =========================================================================== + +/// Codes the commit kernel latches into the pinned fault word. Kept in sync +/// with `arc_graph_step_commit` in `src/cuda/decode_loop.cu`. +pub const FAULT_NONE: i32 = 0; +/// Sampled token id was outside `[0, vocab)`. +pub const FAULT_TOKEN_RANGE: i32 = 1; +/// Advancing the position would leave the fixed-capacity window. +pub const FAULT_POSITION_LIMIT: i32 = 2; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum DeviceLoopError { + BadConfig(&'static str), + /// A CUDA call reported a failure. `at` names the call site. + Cuda { + at: &'static str, + code: u32, + }, + /// The stream retired every enqueued step but the ring is short. The work + /// completed without producing the tokens it was launched to produce. + LostTokens { + expected: i32, + got: i32, + }, + /// The commit kernel refused to write and latched a code. + DeviceFault(i32), + /// The write head ran more than `ring_size` ahead of the read cursor, so + /// unread tokens were overwritten. Reported rather than silently returning + /// whatever survived. + RingOverrun { + pending: i64, + ring_size: usize, + }, + /// Spin budget exhausted with the stream still busy and no error reported. + Stalled { + spins: u64, + }, +} + +impl fmt::Display for DeviceLoopError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::BadConfig(m) => write!(f, "device decode loop misconfigured: {m}"), + Self::Cuda { at, code } => { + write!(f, "device decode loop: {at} failed: CUDA code {code}") + } + Self::LostTokens { expected, got } => write!( + f, + "device decode loop: the stream retired every enqueued step but only {got} of \ + {expected} tokens reached the ring. The graph ran and did not commit its tokens \ + โ€” treat this as an async fault, not as slow progress" + ), + Self::DeviceFault(code) => { + let what = match *code { + FAULT_TOKEN_RANGE => "sampled token id outside [0, vocab)", + FAULT_POSITION_LIMIT => { + "advancing the position would leave the fixed-capacity KV window" + } + _ => "unrecognised code", + }; + write!( + f, + "device decode loop: commit kernel refused to write (code {code}: {what}). No \ + token was scattered into the graph input buffer, so the next launch would \ + have read a stale id" + ) + } + Self::RingOverrun { pending, ring_size } => write!( + f, + "device decode loop: {pending} tokens pending in a {ring_size}-slot ring โ€” the GPU \ + overwrote tokens the host had not read. Tokens were lost; this is a correctness \ + stop, not a throughput warning" + ), + Self::Stalled { spins } => write!( + f, + "device decode loop: stream still busy after {spins} spins with no error reported. \ + The graph is hung rather than faulted" + ), + } + } +} + +impl std::error::Error for DeviceLoopError {} + +impl From for candle_core::Error { + fn from(e: DeviceLoopError) -> Self { + candle_core::Error::Msg(e.to_string()) + } +} + +// =========================================================================== +// The device interface +// =========================================================================== + +/// Result of a non-blocking stream probe. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StreamState { + /// Every enqueued step has retired (`cudaSuccess`). + Idle, + /// Work still outstanding (`cudaErrorNotReady`). + Busy, +} + +/// The device operations a burst needs. +/// +/// **There is deliberately no `synchronize` method.** The driver is structurally +/// incapable of blocking the stream, which is the property this whole module +/// exists to guarantee โ€” it cannot be regressed by an edit that "just adds a +/// sync to be safe", because there is nothing to call. +pub trait DeviceOps { + /// One `cuGraphLaunch` of the captured decode graph. Must not block. + fn launch_graph(&mut self) -> Result<(), DeviceLoopError>; + + /// Sample on device: graph logits -> the `i32[batch]` token buffer. Must + /// not block and must not copy the token to the host. + fn sample_on_device(&mut self) -> Result<(), DeviceLoopError>; + + /// Commit on device: token -> pinned graph input ids, position += 1, and + /// publish into the pinned ring. Must not block. + fn commit_on_device(&mut self) -> Result<(), DeviceLoopError>; + + /// Non-blocking stream probe. `Err` for a sticky async fault. + fn poll_stream(&mut self) -> Result; + + /// Monotonic write head for `row`, read from pinned mapped host memory. + /// A plain volatile load โ€” no CUDA API call. + fn ring_head(&self, row: usize) -> i32; + + /// Token at absolute index `index` for `row` (the impl applies the modulo). + fn ring_slot(&self, row: usize, index: i32) -> i32; + + /// The sticky fault word, read from pinned mapped host memory. + fn device_fault(&self) -> i32; + + /// Rows the ring carries. + fn batch(&self) -> usize; +} + +// =========================================================================== +// The driver +// =========================================================================== + +/// What one burst produced. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BurstOutcome { + /// Tokens per row, in generation order, truncated at and including EOS. + pub tokens: Vec>, + /// True when some row emitted EOS inside this burst. + pub hit_eos: bool, + /// Steps actually launched. When `hit_eos` is true this exceeds the number + /// of tokens returned: the burst was already enqueued and the GPU computed + /// the tail before the host could know. That work is discarded, and this + /// field is how much of it there was. + pub steps_launched: usize, +} + +/// Drives bursts against a [`DeviceOps`], owning the ring read cursors. +pub struct BurstDriver { + cfg: DeviceLoopConfig, + /// Absolute index of the next unread ring slot, per row. + cursor: Vec, + launches: u64, + bursts: u64, + spins: u64, +} + +impl BurstDriver { + pub fn new(cfg: DeviceLoopConfig, batch: usize) -> Result { + cfg.validate()?; + Ok(Self { + cfg, + cursor: vec![0; batch.max(1)], + launches: 0, + bursts: 0, + spins: 0, + }) + } + + pub fn config(&self) -> &DeviceLoopConfig { + &self.cfg + } + + /// Total `cuGraphLaunch` calls issued. With `bursts()` this is the honest + /// answer to "did the device loop do anything": `launches == steps` and + /// zero stream synchronizes is the entire claim. + pub fn launches(&self) -> u64 { + self.launches + } + + pub fn bursts(&self) -> u64 { + self.bursts + } + + /// Host spins spent waiting for tokens to land. Rises when the host + /// out-runs the GPU, which is the healthy direction. + pub fn spins(&self) -> u64 { + self.spins + } + + /// The read cursor is per-sequence state. Reset it when the ring is reset + /// (a new request reuses the buffers). + pub fn reset_cursors(&mut self) { + for c in self.cursor.iter_mut() { + *c = 0; + } + } + + /// Enqueue `steps` decode steps and drain their tokens. + /// + /// Issues exactly `steps` `cuGraphLaunch` calls and **zero** stream + /// synchronizations. The drain is a host-memory poll on pinned mapped + /// memory plus a periodic non-blocking `cudaStreamQuery`. + pub fn run_burst( + &mut self, + ops: &mut D, + steps: usize, + ) -> Result { + if steps == 0 { + return Err(DeviceLoopError::BadConfig( + "run_burst called with zero steps; the caller must fall back to eager instead", + )); + } + if steps >= self.cfg.ring_size { + return Err(DeviceLoopError::BadConfig( + "burst is not smaller than the ring; it would overwrite its own tokens", + )); + } + + // โ”€โ”€ enqueue โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + // No blocking call in this loop, by construction: `DeviceOps` has no + // way to synchronize. `poll_stream` is the async-fault probe and it + // returns immediately whatever the stream is doing. + for _ in 0..steps { + ops.launch_graph()?; + ops.sample_on_device()?; + ops.commit_on_device()?; + self.launches += 1; + // Surface a sticky fault from an EARLIER step while later steps are + // already in flight. Costs nothing and shortens detection latency + // from "end of burst" to "next step". + ops.poll_stream()?; + } + self.bursts += 1; + + // โ”€โ”€ drain โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + let batch = ops.batch().min(self.cursor.len()); + let want: Vec = self.cursor[..batch] + .iter() + .map(|c| c.saturating_add(steps as i32)) + .collect(); + for (row, &target) in want.iter().enumerate() { + self.await_row(ops, row, target)?; + } + + // The commit kernel latches its refusal instead of writing, so this + // must be checked even when every token arrived: a fault on step 3 of 4 + // means steps 4.. read a stale input id. + let fault = ops.device_fault(); + if fault != FAULT_NONE { + return Err(DeviceLoopError::DeviceFault(fault)); + } + + // โ”€โ”€ collect โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + let mut tokens = Vec::with_capacity(batch); + let mut hit_eos = false; + for (row, &target) in want.iter().enumerate() { + let mut row_tokens = Vec::with_capacity(steps); + let mut idx = self.cursor[row]; + while idx < target { + let tok = ops.ring_slot(row, idx); + idx += 1; + let is_eos = self.cfg.eos_token_id >= 0 && tok == self.cfg.eos_token_id; + row_tokens.push(tok); + if is_eos { + hit_eos = true; + break; + } + } + // The cursor advances over the WHOLE burst even when EOS truncated + // the returned tokens. Those slots were written; leaving the cursor + // behind would re-read them as if they were the next burst's. + self.cursor[row] = target; + tokens.push(row_tokens); + } + + Ok(BurstOutcome { + tokens, + hit_eos, + steps_launched: steps, + }) + } + + /// Wait for `row`'s write head to reach `target`, without synchronizing. + fn await_row( + &mut self, + ops: &mut D, + row: usize, + target: i32, + ) -> Result<(), DeviceLoopError> { + let mut spins: u64 = 0; + loop { + let head = ops.ring_head(row); + if head >= target { + // Overrun check: more unread tokens than the ring can hold + // means the GPU lapped the host and older tokens are gone. + let pending = (head as i64) - (self.cursor[row] as i64); + if pending > self.cfg.ring_size as i64 { + return Err(DeviceLoopError::RingOverrun { + pending, + ring_size: self.cfg.ring_size, + }); + } + return Ok(()); + } + + if spins.is_multiple_of(self.cfg.poll_every as u64) { + match ops.poll_stream()? { + StreamState::Busy => {} + StreamState::Idle => { + // Every enqueued step retired. Re-read once: the head + // store and the stream's retirement are not ordered + // against each other from the host's point of view. + let head = ops.ring_head(row); + if head >= target { + return Ok(()); + } + return Err(DeviceLoopError::LostTokens { + expected: target, + got: head, + }); + } + } + } + + spins += 1; + self.spins += 1; + if spins >= self.cfg.max_spins { + return Err(DeviceLoopError::Stalled { spins }); + } + std::hint::spin_loop(); + } + } +} + +// =========================================================================== +// CUDA implementation +// =========================================================================== + +#[cfg(feature = "cuda")] +mod cuda_impl { + use super::*; + use crate::ffi::*; + use crate::sampling_cpu::SamplingConfig; + use crate::sampling_cuda::CudaSampler; + use candle_core::cuda::cudarc::driver::sys::CUstream; + use candle_core::{DType, Device, Tensor}; + + /// Device address of a buffer the captured graph reads. + /// + /// Refuses a non-contiguous tensor rather than silently taking the address + /// of the copy `contiguous()` would make โ€” that copy is not the address the + /// graph baked, so writing to it would look fine and change nothing. + fn pinned_ptr(t: &Tensor, what: &str) -> candle_core::Result { + if !t.is_contiguous() { + candle_core::bail!( + "device decode loop: the {what} buffer is not contiguous, so its address is not \ + the one the captured graph baked. Refusing rather than writing into a copy." + ); + } + crate::weights::tensor_device_ptr(t) + } + + /// `cudaErrorNotReady` โ€” the stream still has outstanding work. Not a + /// failure, and the only non-success code `cudaStreamQuery` returns that + /// is not a fault. + const CUDA_ERROR_NOT_READY: u32 = 600; + + /// Pinned, mapped host memory the GPU writes and the host reads with plain + /// volatile loads. Freed with `cudaFreeHost` on drop. + struct PinnedRing { + tokens: *mut i32, + head: *mut i32, + fault: *mut i32, + batch: usize, + ring_size: usize, + } + + impl PinnedRing { + fn new(batch: usize, ring_size: usize) -> Result { + let mut tokens: *mut std::ffi::c_void = std::ptr::null_mut(); + let mut head: *mut std::ffi::c_void = std::ptr::null_mut(); + let mut fault: *mut std::ffi::c_void = std::ptr::null_mut(); + unsafe { + let s = cudaHostAlloc(&mut tokens, batch * ring_size * 4, CUDA_HOST_ALLOC_MAPPED); + if s != CUDA_SUCCESS { + return Err(DeviceLoopError::Cuda { + at: "cudaHostAlloc(ring tokens)", + code: s, + }); + } + let s = cudaHostAlloc(&mut head, batch * 4, CUDA_HOST_ALLOC_MAPPED); + if s != CUDA_SUCCESS { + cudaFreeHost(tokens); + return Err(DeviceLoopError::Cuda { + at: "cudaHostAlloc(ring head)", + code: s, + }); + } + let s = cudaHostAlloc(&mut fault, 4, CUDA_HOST_ALLOC_MAPPED); + if s != CUDA_SUCCESS { + cudaFreeHost(tokens); + cudaFreeHost(head); + return Err(DeviceLoopError::Cuda { + at: "cudaHostAlloc(fault word)", + code: s, + }); + } + std::ptr::write_bytes(tokens as *mut u8, 0, batch * ring_size * 4); + std::ptr::write_bytes(head as *mut u8, 0, batch * 4); + std::ptr::write_bytes(fault as *mut u8, 0, 4); + } + Ok(Self { + tokens: tokens as *mut i32, + head: head as *mut i32, + fault: fault as *mut i32, + batch, + ring_size, + }) + } + + fn reset(&self) { + unsafe { + std::ptr::write_bytes(self.head as *mut u8, 0, self.batch * 4); + std::ptr::write_bytes(self.fault as *mut u8, 0, 4); + } + } + } + + impl Drop for PinnedRing { + fn drop(&mut self) { + unsafe { + if !self.tokens.is_null() { + cudaFreeHost(self.tokens as *mut _); + } + if !self.head.is_null() { + cudaFreeHost(self.head as *mut _); + } + if !self.fault.is_null() { + cudaFreeHost(self.fault as *mut _); + } + } + } + } + + /// The real [`DeviceOps`]: sampler + commit kernel on the graph's stream. + /// + /// Holds `&mut` to nothing outside itself except through raw pointers taken + /// once at construction, because every buffer it touches must be at a + /// **stable device address** โ€” the captured graph baked those addresses and + /// a reallocation silently invalidates the graph. + pub struct CudaDeviceOps { + stream: CUstream, + sampler: CudaSampler, + sampling_cfg: SamplingConfig, + /// The graph's output logits tensor, reshaped to `[batch, vocab]`. + logits: Tensor, + /// `i32[batch]`, written by the sampler, read by the commit kernel. + token_ids: Tensor, + /// The PINNED `U32` token-id buffer the captured graph reads. + input_ids_ptr: *mut u32, + /// The PINNED `U32` position buffer, or null when positions are owned + /// elsewhere (e.g. by a device-side RoPE position path). + positions_ptr: *mut u32, + ring: PinnedRing, + cfg: DeviceLoopConfig, + /// `cuGraphLaunch` closure state: the exec handle to replay. + exec: CUgraphExec, + } + + impl CudaDeviceOps { + /// * `logits` โ€” the captured graph's output tensor. Its storage address + /// is graph-owned and stable; the sampler reads it in place. + /// * `input_ids_buf` / `positions_buf` โ€” the *pinned* graph input + /// buffers (`layers::GRAPH_MODE_INPUT_IDS` / `GRAPH_MODE_POSITIONS`). + /// These are the addresses the captured graph reads, which is why a + /// device-side write into them closes the loop with no host copy. + #[allow(clippy::too_many_arguments)] + pub fn new( + device: &Device, + exec: CUgraphExec, + stream: CUstream, + sampler: CudaSampler, + sampling_cfg: SamplingConfig, + logits: &Tensor, + input_ids_buf: &Tensor, + positions_buf: Option<&Tensor>, + cfg: DeviceLoopConfig, + ) -> candle_core::Result { + cfg.validate()?; + let batch = sampler.batch(); + + // `[1, 1, vocab]` (or `[1, vocab]`) -> `[batch, vocab]`. Reshaping a + // contiguous tensor is metadata-only, so this does not move the + // graph-owned storage the sampler is about to read. + let logits = logits.reshape((batch, sampler.vocab()))?; + if logits.dtype() != sampler.dtype() { + candle_core::bail!( + "device decode loop: graph logits dtype {:?} != sampler dtype {:?}", + logits.dtype(), + sampler.dtype() + ); + } + + if input_ids_buf.dtype() != DType::U32 { + candle_core::bail!( + "device decode loop: the graph input-ids buffer must be U32 (got {:?}); the \ + commit kernel writes it as uint32_t", + input_ids_buf.dtype() + ); + } + if input_ids_buf.elem_count() < batch { + candle_core::bail!( + "device decode loop: input-ids buffer holds {} elements, need {batch}", + input_ids_buf.elem_count() + ); + } + if let Some(p) = positions_buf { + if p.dtype() != DType::U32 { + candle_core::bail!( + "device decode loop: the graph positions buffer must be U32 (got {:?})", + p.dtype() + ); + } + if p.elem_count() < batch { + candle_core::bail!( + "device decode loop: positions buffer holds {} elements, need {batch}", + p.elem_count() + ); + } + } + + // I32 via `from_vec`: `Tensor::zeros` on I32 is not universally + // supported across backends, the same reason `CudaSampler::new` + // builds `keep_idx_scratch` this way. + let token_ids = Tensor::from_vec(vec![0i32; batch], batch, device)?; + + let input_ids_ptr = pinned_ptr(input_ids_buf, "graph input-ids")? as *mut u32; + let positions_ptr = match positions_buf { + Some(p) => pinned_ptr(p, "graph positions")? as *mut u32, + None => std::ptr::null_mut(), + }; + let ring = PinnedRing::new(batch, cfg.ring_size)?; + + Ok(Self { + stream, + sampler, + sampling_cfg, + logits, + token_ids, + input_ids_ptr, + positions_ptr, + ring, + cfg, + exec, + }) + } + + /// Zero the ring head and fault word. Call when a new sequence starts + /// reusing these buffers; pair it with [`BurstDriver::reset_cursors`]. + pub fn reset_ring(&self) { + self.ring.reset(); + } + } + + impl DeviceOps for CudaDeviceOps { + fn launch_graph(&mut self) -> Result<(), DeviceLoopError> { + // The launch and nothing else. `graph.rs::replay` follows this with + // `cudaStreamSynchronize`; the entire point here is that we do not. + let s = unsafe { cuGraphLaunch(self.exec, self.stream) }; + if s != CUDA_SUCCESS { + return Err(DeviceLoopError::Cuda { + at: "cuGraphLaunch", + code: s, + }); + } + Ok(()) + } + + fn sample_on_device(&mut self) -> Result<(), DeviceLoopError> { + self.sampler + .sample(&self.logits, None, self.sampling_cfg, &self.token_ids) + .map_err(|e| { + tracing::error!("device decode loop: on-device sample failed: {e}"); + DeviceLoopError::Cuda { + at: "CudaSampler::sample", + code: u32::MAX, + } + }) + } + + fn commit_on_device(&mut self) -> Result<(), DeviceLoopError> { + let token_ptr = match crate::weights::tensor_device_ptr(&self.token_ids) { + Ok(p) => p as *const i32, + Err(e) => { + tracing::error!("device decode loop: token buffer pointer unavailable: {e}"); + return Err(DeviceLoopError::Cuda { + at: "tensor_device_ptr(token_ids)", + code: u32::MAX, + }); + } + }; + unsafe { + arc_launch_graph_step_commit( + token_ptr, + self.input_ids_ptr, + self.positions_ptr, + self.ring.tokens, + self.ring.head, + self.ring.fault, + self.cfg.ring_size as i32, + self.cfg.vocab as i32, + self.cfg.position_limit as i32, + self.ring.batch as i32, + self.stream, + ); + } + Ok(()) + } + + fn poll_stream(&mut self) -> Result { + // Non-blocking. NVIDIA's reference for cudaStreamQuery: "Note that + // this function may also return error codes from previous, + // asynchronous launches." That is what makes a sticky fault from an + // earlier graph step visible here without ever blocking. + let s = unsafe { cudaStreamQuery(self.stream) }; + match s { + CUDA_SUCCESS => Ok(StreamState::Idle), + CUDA_ERROR_NOT_READY => Ok(StreamState::Busy), + code => Err(DeviceLoopError::Cuda { + at: "cudaStreamQuery (async fault from an earlier graph step)", + code, + }), + } + } + + fn ring_head(&self, row: usize) -> i32 { + unsafe { std::ptr::read_volatile(self.ring.head.add(row)) } + } + + fn ring_slot(&self, row: usize, index: i32) -> i32 { + let slot = + row * self.ring.ring_size + (index.rem_euclid(self.ring.ring_size as i32) as usize); + unsafe { std::ptr::read_volatile(self.ring.tokens.add(slot)) } + } + + fn device_fault(&self) -> i32 { + unsafe { std::ptr::read_volatile(self.ring.fault) } + } + + fn batch(&self) -> usize { + self.ring.batch + } + } + + /// The whole device loop behind one handle: sampler, commit kernel, pinned + /// ring and burst driver. + /// + /// Built once per captured graph and kept alive across decode steps, + /// because every buffer it holds a pointer to must stay at the address the + /// graph baked. Dropping and rebuilding it mid-generation would hand the + /// commit kernel a fresh `token_ids` allocation while the graph still reads + /// the old input-ids address โ€” so callers keep it, they do not recreate it. + pub struct DeviceDecodeLoop { + ops: CudaDeviceOps, + driver: BurstDriver, + } + + // Same reasoning as `CudaGraphRunner` (`graph.rs:114`): the raw pointers are + // CUDA handles and pinned host allocations, both valid from any thread, and + // the pipeline that owns this is held behind a mutex. Without these the + // `NormalPipeline` field would make the whole pipeline non-`Send`. + unsafe impl Send for CudaDeviceOps {} + unsafe impl Sync for CudaDeviceOps {} + unsafe impl Send for DeviceDecodeLoop {} + unsafe impl Sync for DeviceDecodeLoop {} + + impl DeviceDecodeLoop { + #[allow(clippy::too_many_arguments)] + pub fn new( + device: &Device, + exec: CUgraphExec, + stream: CUstream, + logits: &Tensor, + input_ids_buf: &Tensor, + positions_buf: Option<&Tensor>, + cfg: DeviceLoopConfig, + ) -> candle_core::Result { + let batch = 1usize; + let vocab = cfg.vocab; + // Greedy, always. [`admit`] refuses every non-greedy sequence, so + // taking the mode as a parameter would only create a way for a + // caller to contradict the guard. The seed is irrelevant under + // greedy (pure argmax draws no random numbers) and is fixed so the + // sampler's RNG buffer is deterministic if the mode ever widens. + let sampling_cfg = SamplingConfig::greedy(); + let sampler = CudaSampler::new(device, batch, vocab, logits.dtype(), 0)?; + let ops = CudaDeviceOps::new( + device, + exec, + stream, + sampler, + sampling_cfg, + logits, + input_ids_buf, + positions_buf, + cfg, + )?; + let driver = BurstDriver::new(cfg, batch)?; + Ok(Self { ops, driver }) + } + + /// Enqueue `steps` decode steps and return their tokens. + /// + /// Exactly `steps` `cuGraphLaunch` calls, zero `cudaStreamSynchronize`. + pub fn run(&mut self, steps: usize) -> candle_core::Result { + let out = self.driver.run_burst(&mut self.ops, steps)?; + // D18: a subsystem that only announces itself at startup cannot be + // told apart from one that started and then did nothing. Emit the + // running counters periodically so `launches` is observable in a + // live run, not just in a debugger. + if self.driver.bursts().is_multiple_of(STATUS_EVERY_BURSTS) { + tracing::info!("{}", self.status_line()); + } + Ok(out) + } + + /// Start a fresh sequence on the same buffers. + pub fn reset(&mut self) { + self.ops.reset_ring(); + self.driver.reset_cursors(); + clear_pending_tokens(); + } + + /// `cuGraphLaunch` calls issued so far โ€” the numerator of the claim. + pub fn launches(&self) -> u64 { + self.driver.launches() + } + + pub fn bursts(&self) -> u64 { + self.driver.bursts() + } + + /// The one line that separates a working device loop from an inert one + /// (D18). `launches=0` is a subsystem that did nothing and says so. + pub fn status_line(&self) -> String { + format!( + "ARCGRAPH DEVICE LOOP: launches={} bursts={} burst_len={} host_syncs=0 \ + host_spins={} pending={}", + self.driver.launches(), + self.driver.bursts(), + self.driver.config().burst, + self.driver.spins(), + pending_token_count(), + ) + } + } +} + +#[cfg(feature = "cuda")] +pub use cuda_impl::{CudaDeviceOps, DeviceDecodeLoop}; + +// =========================================================================== +// Admission +// =========================================================================== + +/// Why the device loop refused a request. +/// +/// Refusing is the normal, safe outcome โ€” the caller falls back to the existing +/// `replay()` + host-sample path. Each variant names something the device +/// sampler cannot reproduce *exactly*, because a token that differs from the +/// host path is a silent behaviour change, not an optimisation. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Refusal { + /// `ARC_GRAPH_DEVICE_LOOP` is not set. The default. + NotEnabled, + /// Only batch 1 is wired: the burst holds the scheduler for its whole + /// length, so it must not straddle sequences with different stop conditions. + NotBatchOne, + /// Prefill, not decode. + NotDecodeStep, + /// No captured graph to replay, or its output is not yet trusted. + GraphNotReady, + /// Non-greedy sampling. The device sampler runs Splitmix64 per row while + /// the host runs Isaac64, so for any stochastic config the two draw + /// *different, equally valid* tokens โ€” which silently breaks seeded + /// reproducibility. Greedy is deterministic, so device and host agree + /// exactly, which is also what makes the path verifiable. + NotGreedy, + /// The caller wants logprobs, which the device sampler does not return. + WantsLogprobs, + /// A logits processor, grammar or penalty must see the logits on the host. + HostSideLogitsWork, + /// No room left in the fixed-capacity KV window for even one step. + NoPositionRoom, +} + +impl fmt::Display for Refusal { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let s = match self { + Self::NotEnabled => "ARC_GRAPH_DEVICE_LOOP is not set", + Self::NotBatchOne => "batch size is not 1", + Self::NotDecodeStep => "this is a prefill step, not decode", + Self::GraphNotReady => "no captured graph, or its output is not yet trusted", + Self::NotGreedy => { + "sampling is not greedy; the device and host RNGs differ, so tokens would \ + diverge from the seeded host path" + } + Self::WantsLogprobs => { + "the request asks for logprobs, which the device sampler \ + does not produce" + } + Self::HostSideLogitsWork => { + "a logits processor, grammar or repetition penalty needs the logits on the host" + } + Self::NoPositionRoom => "no room left in the fixed-capacity KV window", + }; + write!(f, "{s}") + } +} + +/// Everything the admission check needs, gathered by the caller. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct AdmissionFacts { + pub batch_size: usize, + pub seq_len: usize, + pub graph_ready: bool, + pub greedy: bool, + pub return_logprobs: bool, + pub host_side_logits_work: bool, + pub current_position: u32, +} + +/// May the device loop drive this step? +/// +/// Kept separate from the CUDA plumbing, and not feature-gated, so the decision +/// is unit-tested on any host. Every refusal is a fall-back to the existing +/// path, never an error. +pub fn admit( + facts: &AdmissionFacts, + cfg: &DeviceLoopConfig, + enabled: bool, +) -> Result { + if !enabled { + return Err(Refusal::NotEnabled); + } + if facts.seq_len != 1 { + return Err(Refusal::NotDecodeStep); + } + if facts.batch_size != 1 { + return Err(Refusal::NotBatchOne); + } + if !facts.graph_ready { + return Err(Refusal::GraphNotReady); + } + if facts.return_logprobs { + return Err(Refusal::WantsLogprobs); + } + if facts.host_side_logits_work { + return Err(Refusal::HostSideLogitsWork); + } + if !facts.greedy { + return Err(Refusal::NotGreedy); + } + let steps = cfg.plan_steps(facts.current_position, cfg.burst); + if steps == 0 { + return Err(Refusal::NoPositionRoom); + } + Ok(steps) +} + +// =========================================================================== +// Host-side token handoff +// =========================================================================== + +// Tokens a burst produced that the engine has not consumed yet. +// +// A burst computes `steps` tokens in one go, but the engine's `Sequence` +// bookkeeping โ€” stop strings, streaming, EOS โ€” is per token. So the burst +// parks its tokens here and the engine drains them one at a time on the +// following steps, during which no GPU work is launched at all. +// +// Thread-local because at batch 1 `sample_and_add_toks` runs the sampler +// inline on the engine thread (`use_async_pool = seqs_len > 1` is false), the +// same thread that ran the forward. The `GRAPH_MODE_*` graph inputs in +// `mistralrs-core/src/layers.rs` are scoped the same way for the same reason. +std::thread_local! { + static PENDING_TOKENS: std::cell::RefCell> = + const { std::cell::RefCell::new(std::collections::VecDeque::new()) }; +} + +/// Park a burst's tokens for the engine to drain. +/// +/// Negative ids are dropped rather than cast: the commit kernel range-checks +/// before it writes, so a negative here would mean the ring was read wrong, and +/// wrapping it into a huge `u32` would index the tokenizer out of bounds. +pub fn push_pending_tokens(tokens: &[i32]) { + PENDING_TOKENS.with(|q| { + let mut q = q.borrow_mut(); + for &t in tokens { + if t >= 0 { + q.push_back(t as u32); + } else { + tracing::error!( + "device decode loop: dropping negative token id {t} from the ring; the \ + commit kernel range-checks before writing, so this means the ring was \ + misread rather than that the model produced it" + ); + } + } + }) +} + +/// Take the next token the burst already produced, if any. +pub fn take_pending_token() -> Option { + PENDING_TOKENS.with(|q| q.borrow_mut().pop_front()) +} + +/// Tokens still parked. +pub fn pending_token_count() -> usize { + PENDING_TOKENS.with(|q| q.borrow().len()) +} + +// Whether the sequence currently being decoded is one the device loop may +// drive. +// +// `Pipeline::forward_inputs` decides whether to run a burst, but it is handed +// only `inputs` and `return_raw_logits` โ€” it never sees the `Sequence`, so it +// cannot ask whether sampling is greedy or whether logprobs were requested. +// `sample_sequence` does see it. So eligibility is published there and read by +// the *next* forward. +// +// The one-step lag is safe in the only direction that matters: a sequence's +// sampling parameters are fixed for its lifetime, and the flag starts `false`, +// so the first decode step always falls back and the loop can only ever engage +// after a step has confirmed the parameters. It can never engage for a +// sequence it has not seen. +std::thread_local! { + static DEVICE_LOOP_ELIGIBLE: std::cell::Cell = const { std::cell::Cell::new(false) }; +} + +/// Publish whether the sequence being sampled may be driven by the device loop. +pub fn set_device_loop_eligible(eligible: bool) { + DEVICE_LOOP_ELIGIBLE.with(|c| c.set(eligible)) +} + +/// Did the last sampled sequence qualify? `false` until one has. +pub fn device_loop_eligible() -> bool { + DEVICE_LOOP_ELIGIBLE.with(|c| c.get()) +} + +/// Drop parked tokens. +/// +/// MUST be called when a sequence finishes, is preempted, or the path falls +/// back โ€” a token left here would be handed to the *next* sequence as if it +/// were its own. That is the "+1-step divergence" failure mode with a different +/// name, so it gets an explicit call rather than a lifetime. +pub fn clear_pending_tokens() { + PENDING_TOKENS.with(|q| q.borrow_mut().clear()) +} + +/// Stand down completely: forget parked tokens and mark the current sequence +/// ineligible. Call on fall-back, on sequence completion, and on any error, so +/// the next sequence starts from a clean slate. +pub fn stand_down() { + clear_pending_tokens(); + set_device_loop_eligible(false); +} + +/// Process-wide latch. Once the device loop has failed, it stays off. +/// +/// Without this a construction or burst failure would be retried on every +/// decode step forever โ€” and because a refusal falls back to the *eager* +/// forward rather than to `replay()`, that would quietly make every step slower +/// than doing nothing at all. One failure, one warning, off for good. +static DEVICE_LOOP_KILLED: std::sync::atomic::AtomicBool = + std::sync::atomic::AtomicBool::new(false); + +/// Latch the device loop off for the rest of the process, loudly and once. +pub fn kill_device_loop(reason: &str) { + if !DEVICE_LOOP_KILLED.swap(true, std::sync::atomic::Ordering::SeqCst) { + tracing::error!( + "ArcGraph device decode loop DISABLED for the rest of this process: {reason}. Decode \ + falls back to graph replay with a host sync and a host sample. This is a stop, not a \ + retry โ€” a loop that fails once and is retried every step is slower than never having \ + engaged." + ); + } + stand_down(); +} + +/// Has the device loop been latched off? +pub fn device_loop_killed() -> bool { + DEVICE_LOOP_KILLED.load(std::sync::atomic::Ordering::SeqCst) +} + +#[cfg(feature = "cuda")] +extern "C" { + /// `arc_graph_step_commit` โ€” `src/cuda/decode_loop.cu`. + /// + /// Scatters the device-sampled token into the pinned graph input buffer, + /// advances the device position, and publishes the token into the pinned + /// mapped ring. Range-checks before it writes anything, latching a code + /// into `fault` instead of committing a bad token. + pub fn arc_launch_graph_step_commit( + sampled: *const i32, + input_ids: *mut u32, + positions: *mut u32, + ring: *mut i32, + ring_head: *mut i32, + fault: *mut i32, + ring_size: i32, + vocab: i32, + position_limit: i32, + batch: i32, + stream: candle_core::cuda::cudarc::driver::sys::CUstream, + ); +} + +// =========================================================================== +// Tests +// =========================================================================== + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::VecDeque; + + /// A scripted fake GPU. + /// + /// It models the property that matters: work is enqueued without blocking + /// and lands *later*. `launch_graph` pushes the next scripted token into an + /// in-flight queue; `poll_stream` retires one in-flight step per call and + /// publishes its token into the ring. So a driver that expects tokens to be + /// present the instant it launches will not find them, and a driver that + /// never polls will spin forever โ€” both real failure modes. + struct FakeGpu { + script: VecDeque, + in_flight: VecDeque, + ring: Vec, + head: i32, + fault: i32, + ring_size: usize, + launches: usize, + samples: usize, + commits: usize, + polls: usize, + /// Set to drop the commit instead of publishing โ€” the "no-op + /// implementation" this test suite must not accept. + commit_is_noop: bool, + /// Injected async fault, surfaced from `poll_stream`. + fault_after: Option, + /// Report Idle immediately even with work outstanding. + lie_idle: bool, + /// A hung kernel: work is enqueued, the stream reports Busy forever and + /// nothing ever retires. No error is ever raised, so only the liveness + /// backstop can end it. + never_retires: bool, + } + + impl FakeGpu { + fn new(script: Vec, ring_size: usize) -> Self { + Self { + script: script.into(), + in_flight: VecDeque::new(), + ring: vec![0; ring_size], + head: 0, + fault: FAULT_NONE, + ring_size, + launches: 0, + samples: 0, + commits: 0, + polls: 0, + commit_is_noop: false, + fault_after: None, + lie_idle: false, + never_retires: false, + } + } + + fn retire_one(&mut self) { + if let Some(tok) = self.in_flight.pop_front() { + let slot = (self.head as usize) % self.ring_size; + self.ring[slot] = tok; + self.head += 1; + } + } + } + + impl DeviceOps for FakeGpu { + fn launch_graph(&mut self) -> Result<(), DeviceLoopError> { + self.launches += 1; + Ok(()) + } + fn sample_on_device(&mut self) -> Result<(), DeviceLoopError> { + self.samples += 1; + Ok(()) + } + fn commit_on_device(&mut self) -> Result<(), DeviceLoopError> { + self.commits += 1; + if self.commit_is_noop { + return Ok(()); + } + let tok = self.script.pop_front().unwrap_or(0); + self.in_flight.push_back(tok); + Ok(()) + } + fn poll_stream(&mut self) -> Result { + self.polls += 1; + if let Some(after) = self.fault_after { + if self.polls > after { + return Err(DeviceLoopError::Cuda { + at: "cudaStreamQuery (async fault from an earlier graph step)", + code: 700, + }); + } + } + if self.never_retires { + return Ok(StreamState::Busy); + } + self.retire_one(); + if self.lie_idle || self.in_flight.is_empty() { + Ok(StreamState::Idle) + } else { + Ok(StreamState::Busy) + } + } + fn ring_head(&self, _row: usize) -> i32 { + self.head + } + fn ring_slot(&self, _row: usize, index: i32) -> i32 { + self.ring[index.rem_euclid(self.ring_size as i32) as usize] + } + fn device_fault(&self) -> i32 { + self.fault + } + fn batch(&self) -> usize { + 1 + } + } + + fn cfg(burst: usize, ring: usize) -> DeviceLoopConfig { + DeviceLoopConfig { + burst, + ring_size: ring, + poll_every: 1, + max_spins: 10_000, + eos_token_id: 2, + vocab: 128, + position_limit: 4096, + } + } + + /// THE claim: N decode steps cost exactly N `cuGraphLaunch` calls, and the + /// tokens that come back are the ones the GPU produced. + /// + /// Fails on a no-op: with fewer launches the count is wrong, and with a + /// commit that never publishes the drain reports `LostTokens`. + #[test] + fn n_steps_cost_exactly_n_launches_and_return_the_gpu_tokens() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14], 16); + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + let out = d.run_burst(&mut gpu, 4).unwrap(); + + assert_eq!(gpu.launches, 4, "N steps must cost exactly N cuGraphLaunch"); + assert_eq!(gpu.samples, 4, "every step must sample on device"); + assert_eq!(gpu.commits, 4, "every step must commit on device"); + assert_eq!(out.tokens, vec![vec![11, 12, 13, 14]]); + assert_eq!(d.launches(), 4); + assert!(!out.hit_eos); + } + + /// The control the previous test needs: an implementation that launches and + /// samples but never publishes the token MUST fail, and must fail as + /// `LostTokens` rather than by spinning out. This is the "test that cannot + /// fail" guard โ€” if the driver accepted a silent no-op commit, the whole + /// module would be decoration. + #[test] + fn a_commit_that_publishes_nothing_is_rejected_as_lost_tokens() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14], 16); + gpu.commit_is_noop = true; + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + let err = d.run_burst(&mut gpu, 4).unwrap_err(); + assert!( + matches!( + err, + DeviceLoopError::LostTokens { + expected: 4, + got: 0 + } + ), + "expected LostTokens, got {err:?}" + ); + // It still launched: the failure is in the commit, and the error must + // not be confused with "the driver never ran". + assert_eq!(gpu.launches, 4); + } + + /// An async fault reported by the non-blocking query must propagate as a + /// CUDA error โ€” not be swallowed, and not degrade into `Stalled`. + #[test] + fn async_fault_from_stream_query_fails_loudly_as_cuda() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14], 16); + gpu.fault_after = Some(2); + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + let err = d.run_burst(&mut gpu, 4).unwrap_err(); + match err { + DeviceLoopError::Cuda { code, .. } => assert_eq!(code, 700), + other => panic!("async fault must surface as Cuda, got {other:?}"), + } + } + + /// Detection latency: the in-burst probe finds the fault before the burst + /// finishes enqueuing, so a fault on step 1 does not wait for step 4. + #[test] + fn in_burst_probe_detects_a_fault_before_the_burst_completes() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14, 15, 16, 17, 18], 32); + gpu.fault_after = Some(1); // second poll errors -> during step 2 + let mut d = BurstDriver::new(cfg(8, 32), 1).unwrap(); + assert!(d.run_burst(&mut gpu, 8).is_err()); + assert!( + gpu.launches < 8, + "the burst must stop early on a fault, launched {}", + gpu.launches + ); + } + + /// The no-timeout fault detector: stream reports every step retired, but + /// the ring is short. That is definitively wrong and must be an error, not + /// an infinite spin and not a silent short read. + #[test] + fn stream_idle_with_a_short_ring_is_a_fault_not_a_wait() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14], 16); + gpu.commit_is_noop = true; + gpu.lie_idle = true; + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + match d.run_burst(&mut gpu, 4).unwrap_err() { + DeviceLoopError::LostTokens { expected, got } => { + assert_eq!((expected, got), (4, 0)); + } + other => panic!("expected LostTokens, got {other:?}"), + } + assert!( + d.spins() < 100, + "idle-but-short must be detected immediately, not spun on ({} spins)", + d.spins() + ); + } + + /// The commit kernel's range check travels back in a pinned word. A driver + /// that forgets to read it returns `Ok` with plausible tokens โ€” exactly the + /// silent-corruption class this repo keeps shipping. So: tokens all arrive, + /// fault word set, and the burst must still fail. + #[test] + fn a_latched_device_fault_fails_the_burst_even_though_every_token_arrived() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14], 16); + gpu.fault = FAULT_TOKEN_RANGE; + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + let err = d.run_burst(&mut gpu, 4).unwrap_err(); + assert_eq!(err, DeviceLoopError::DeviceFault(FAULT_TOKEN_RANGE)); + assert_eq!(gpu.head, 4, "the tokens did land; the fault is orthogonal"); + } + + /// EOS truncates the returned tokens but not the launch count. The burst + /// was already enqueued, so the tail was computed and is discarded โ€” the + /// outcome must report both numbers honestly. + #[test] + fn eos_truncates_tokens_but_the_wasted_steps_are_still_reported() { + let mut gpu = FakeGpu::new(vec![11, 12, 2, 14], 16); + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + let out = d.run_burst(&mut gpu, 4).unwrap(); + assert_eq!(out.tokens, vec![vec![11, 12, 2]]); + assert!(out.hit_eos); + assert_eq!( + out.steps_launched, 4, + "the discarded tail must be visible, not rounded away" + ); + } + + /// Consecutive bursts must not re-read the previous burst's slots. The + /// cursor advances over the full burst even when EOS truncated the output. + #[test] + fn cursor_advances_over_the_whole_burst_across_consecutive_bursts() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14, 21, 22, 23, 24], 16); + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + let a = d.run_burst(&mut gpu, 4).unwrap(); + let b = d.run_burst(&mut gpu, 4).unwrap(); + assert_eq!(a.tokens, vec![vec![11, 12, 13, 14]]); + assert_eq!(b.tokens, vec![vec![21, 22, 23, 24]]); + assert_eq!(gpu.launches, 8); + } + + /// Ring wraparound is arithmetic, and arithmetic is where silent corruption + /// lives. Drive past the ring size and check the tokens are still right. + #[test] + fn tokens_survive_ring_wraparound() { + let script: Vec = (100..140).collect(); + let mut gpu = FakeGpu::new(script.clone(), 8); + let mut d = BurstDriver::new(cfg(4, 8), 1).unwrap(); + let mut seen = Vec::new(); + for _ in 0..10 { + let out = d.run_burst(&mut gpu, 4).unwrap(); + seen.extend(out.tokens[0].iter().copied()); + } + assert_eq!(seen, script, "wraparound must not permute or drop tokens"); + } + + /// A burst that could lap the ring is refused up front rather than + /// discovered as scrambled tokens. + #[test] + fn a_burst_that_would_lap_the_ring_is_refused() { + assert!(matches!( + BurstDriver::new(cfg(16, 8), 1), + Err(DeviceLoopError::BadConfig(_)) + )); + let mut gpu = FakeGpu::new(vec![1; 32], 8); + let mut d = BurstDriver::new(cfg(4, 8), 1).unwrap(); + assert!(matches!( + d.run_burst(&mut gpu, 8), + Err(DeviceLoopError::BadConfig(_)) + )); + } + + /// A kernel that hangs *without* faulting reports Busy forever and lands + /// nothing. No CUDA error is ever raised and the stream never goes idle, so + /// neither of the two real detectors can fire โ€” only the liveness backstop + /// can end the burst. Without it this is an unkillable spin in the engine. + #[test] + fn a_hung_stream_trips_the_liveness_backstop() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14], 16); + gpu.never_retires = true; + let mut c = cfg(4, 16); + c.max_spins = 500; + let mut d = BurstDriver::new(c, 1).unwrap(); + match d.run_burst(&mut gpu, 4).unwrap_err() { + DeviceLoopError::Stalled { spins } => assert_eq!(spins, 500), + other => panic!("expected Stalled, got {other:?}"), + } + } + + /// `poll_every == 0` would divide by zero in the drain. Refuse it at + /// construction rather than panicking mid-burst. + #[test] + fn zero_poll_every_is_refused_rather_than_dividing_by_zero() { + let mut c = cfg(4, 16); + c.poll_every = 0; + assert!(matches!( + BurstDriver::new(c, 1), + Err(DeviceLoopError::BadConfig(_)) + )); + } + + /// `plan_steps` must never walk the device position out of the + /// fixed-capacity KV window, and must report 0 rather than a short burst + /// the caller might mistake for success. + #[test] + fn plan_steps_clamps_to_the_position_window_and_the_token_budget() { + let mut c = cfg(8, 64); + c.position_limit = 100; + assert_eq!(c.plan_steps(0, 1000), 8); + assert_eq!(c.plan_steps(95, 1000), 5, "must not pass position_limit"); + assert_eq!(c.plan_steps(100, 1000), 0, "no room left"); + assert_eq!( + c.plan_steps(120, 1000), + 0, + "already past: saturating, not wrapping" + ); + assert_eq!(c.plan_steps(0, 3), 3, "must not exceed the token budget"); + } + + /// Zero steps is a caller bug, not a quiet no-op that reports success. + #[test] + fn zero_steps_is_an_error_not_a_silent_success() { + let mut gpu = FakeGpu::new(vec![], 16); + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + assert!(matches!( + d.run_burst(&mut gpu, 0), + Err(DeviceLoopError::BadConfig(_)) + )); + assert_eq!(gpu.launches, 0); + } + + fn ok_facts() -> AdmissionFacts { + AdmissionFacts { + batch_size: 1, + seq_len: 1, + graph_ready: true, + greedy: true, + return_logprobs: false, + host_side_logits_work: false, + current_position: 10, + } + } + + /// Admission must refuse every case where the device sampler would produce + /// a token the host path would not. Each of these is a silent behaviour + /// change if it slips through, so each gets its own assertion rather than + /// one "returns Err" catch-all. + #[test] + fn admission_refuses_every_case_the_device_sampler_cannot_reproduce() { + let c = cfg(4, 64); + // The happy path first โ€” otherwise a guard that refuses everything + // would pass every negative case below. + assert_eq!(admit(&ok_facts(), &c, true), Ok(4)); + + assert_eq!(admit(&ok_facts(), &c, false), Err(Refusal::NotEnabled)); + + let mut f = ok_facts(); + f.seq_len = 32; + assert_eq!(admit(&f, &c, true), Err(Refusal::NotDecodeStep)); + + let mut f = ok_facts(); + f.batch_size = 2; + assert_eq!(admit(&f, &c, true), Err(Refusal::NotBatchOne)); + + let mut f = ok_facts(); + f.graph_ready = false; + assert_eq!(admit(&f, &c, true), Err(Refusal::GraphNotReady)); + + let mut f = ok_facts(); + f.greedy = false; + assert_eq!( + admit(&f, &c, true), + Err(Refusal::NotGreedy), + "stochastic sampling must fall back: device Splitmix64 and host Isaac64 draw \ + different tokens, which breaks seeded reproducibility" + ); + + let mut f = ok_facts(); + f.return_logprobs = true; + assert_eq!(admit(&f, &c, true), Err(Refusal::WantsLogprobs)); + + let mut f = ok_facts(); + f.host_side_logits_work = true; + assert_eq!(admit(&f, &c, true), Err(Refusal::HostSideLogitsWork)); + } + + /// Admission must also clamp the burst to the KV window, and refuse rather + /// than return a zero-step "success" the caller could mistake for a launch. + #[test] + fn admission_clamps_the_burst_to_the_kv_window() { + let mut c = cfg(4, 64); + c.position_limit = 100; + let mut f = ok_facts(); + f.current_position = 98; + assert_eq!(admit(&f, &c, true), Ok(2), "clamped to the window"); + f.current_position = 100; + assert_eq!(admit(&f, &c, true), Err(Refusal::NoPositionRoom)); + } + + /// The parked tokens are handed out in generation order, exactly once. + /// A queue that re-served a token, or served them reversed, would produce + /// fluent-but-wrong text โ€” invisible without this. + #[test] + fn pending_tokens_drain_in_order_exactly_once() { + clear_pending_tokens(); + push_pending_tokens(&[7, 8, 9]); + assert_eq!(pending_token_count(), 3); + assert_eq!(take_pending_token(), Some(7)); + assert_eq!(take_pending_token(), Some(8)); + assert_eq!(take_pending_token(), Some(9)); + assert_eq!(take_pending_token(), None); + assert_eq!(pending_token_count(), 0); + } + + /// A leftover token must never reach the next sequence. `clear` is the only + /// thing standing between a fall-back and one sequence's token appearing in + /// another's output. + #[test] + fn clearing_pending_tokens_prevents_leaking_them_into_the_next_sequence() { + clear_pending_tokens(); + push_pending_tokens(&[7, 8, 9]); + clear_pending_tokens(); + assert_eq!(take_pending_token(), None); + } + + /// A negative id in the ring means the ring was misread. Casting it would + /// index the tokenizer at ~4 billion; it must be dropped and logged. + #[test] + fn negative_ring_ids_are_dropped_not_wrapped_into_huge_u32() { + clear_pending_tokens(); + push_pending_tokens(&[5, -1, 6]); + assert_eq!(take_pending_token(), Some(5)); + assert_eq!(take_pending_token(), Some(6)); + assert_eq!(take_pending_token(), None); + } + + /// The eligibility flag must start `false`, so the first decode step of a + /// sequence always falls back rather than engaging on parameters nothing + /// has checked yet. + #[test] + fn eligibility_starts_false_and_stand_down_resets_everything() { + set_device_loop_eligible(false); + clear_pending_tokens(); + assert!(!device_loop_eligible()); + + set_device_loop_eligible(true); + push_pending_tokens(&[1, 2]); + assert!(device_loop_eligible()); + + stand_down(); + assert!(!device_loop_eligible(), "stand_down must clear eligibility"); + assert_eq!(take_pending_token(), None, "stand_down must clear tokens"); + } + + /// Default-off, and explicitly pinnable off. + #[test] + fn the_opt_in_defaults_to_off() { + // Not using the process env here: these run in parallel with other + // tests. The parsing rule is what matters and it is exercised directly. + for (val, want) in [ + ("", false), + ("0", false), + ("false", false), + ("off", false), + ("no", false), + ("1", true), + ("true", true), + ] { + let v = val.trim().to_ascii_lowercase(); + let enabled = !(v.is_empty() || v == "0" || v == "false" || v == "off" || v == "no"); + assert_eq!(enabled, want, "for {val:?}"); + } + // And with the variable absent entirely. + std::env::remove_var(ENV_ENABLE); + assert!(!device_loop_enabled()); + } +} diff --git a/arc-cuda-graph/src/ffi.rs b/arc-cuda-graph/src/ffi.rs index ae78c6f76..96cae2954 100644 --- a/arc-cuda-graph/src/ffi.rs +++ b/arc-cuda-graph/src/ffi.rs @@ -179,6 +179,13 @@ extern "C" { pub fn cuStreamCreate(stream: *mut CUstream, flags: u32) -> u32; pub fn cuStreamDestroy_v2(stream: CUstream) -> u32; pub fn cudaStreamSynchronize(stream: CUstream) -> u32; + /// Non-blocking completion probe. Returns `cudaSuccess` when every + /// enqueued item has retired and `cudaErrorNotReady` (600) while work is + /// outstanding. NVIDIA's reference adds: *"Note that this function may + /// also return error codes from previous, asynchronous launches."* That + /// is what lets `device_loop` surface an async graph fault without ever + /// blocking the stream โ€” see `device_loop::CudaDeviceOps::poll_stream`. + pub fn cudaStreamQuery(stream: CUstream) -> u32; // ========== Conditional nodes (CUDA 12.4+) ========== pub fn cudaGraphConditionalHandleCreate( diff --git a/arc-cuda-graph/src/graph.rs b/arc-cuda-graph/src/graph.rs index 20c0e805f..e5ff6aefa 100644 --- a/arc-cuda-graph/src/graph.rs +++ b/arc-cuda-graph/src/graph.rs @@ -784,6 +784,49 @@ output_trusted={} verify_remaining={} verify_failed={}", } } + /// The handles the device decode loop needs to replay this graph **without + /// the synchronize above**: the exec to launch and the graph-owned output + /// tensor to sample from, in place, on device. + /// + /// Handing these out is the whole difference between `replay()` and + /// `device_loop`. `replay()` launches, blocks on `cudaStreamSynchronize` + /// (line 362) and returns the output for the *host* to sample โ€” which + /// copies logits down and serialises the step. `device_loop` takes the same + /// two handles, launches, samples on device and scatters the token back + /// into the pinned input buffer, so the next launch needs no host at all. + /// + /// The returned `Tensor`'s storage was allocated from the graph's private + /// pool during capture, so its address is stable for the life of the + /// `CapturedGraph` โ€” which is exactly why sampling it in place is sound. + /// The clone is an `Arc` bump and does **not** extend that lifetime past + /// the runner's; callers must not hold it after dropping the runner. + /// + /// ๐Ÿ”ด The same BOOBY TRAP `replay` documents applies here verbatim: this + /// output aliases storage the next launch overwrites, so it is not a + /// snapshot. `device_loop` is safe against it for a specific reason โ€” + /// every access is enqueued on this same stream, in order, and none of + /// them is a host read โ€” and `normal.rs` enforces at runtime the invariant + /// that lets the aliasing tensor be handed back at all. Anything else + /// taking these handles must copy. + pub fn replay_handles(&self, batch_size: usize) -> Option<(CUgraphExec, Tensor)> { + let captured = self.graphs.get(&batch_size)?; + let out = captured.output.clone()?; + Some((captured.exec, out)) + } + + /// The stream capture and replay run on. `device_loop` launches its commit + /// kernel here so it is ordered behind the graph without an event. + pub fn stream(&self) -> CUstream { + self.stream + } + + /// Count a decode step driven by the device loop rather than by + /// [`replay`](Self::replay), so `status_line` stays honest about how much + /// work this runner actually did. + pub fn record_device_loop_steps(&mut self, steps: u64) { + self.replays += steps; + } + pub fn disable(&mut self) { self.enabled = false; } diff --git a/arc-cuda-graph/src/lib.rs b/arc-cuda-graph/src/lib.rs index 8e534ed6d..766653698 100644 --- a/arc-cuda-graph/src/lib.rs +++ b/arc-cuda-graph/src/lib.rs @@ -12,6 +12,11 @@ pub mod autonomous; pub mod buffers; pub mod decode_forward; pub mod dedicated; +/// Device-resident decode loop: N steps, N `cuGraphLaunch` calls, zero stream +/// synchronizations. Opt-in via `ARC_GRAPH_DEVICE_LOOP`. The driver and its +/// policy compile without CUDA so they are unit-tested on any host; only the +/// pointer plumbing is gated. +pub mod device_loop; pub mod ffi; pub mod gemv_ffi; pub mod graph; @@ -82,6 +87,14 @@ pub use buffers::{DecodeInputBuffers, DecodeState}; pub use decode_forward::{decode_forward, DecodeBuffers, LayerKvCache, PagedAttentionState}; #[cfg(feature = "cuda")] pub use dedicated::DedicatedDecodePath; +pub use device_loop::{ + admit, clear_pending_tokens, device_loop_eligible, device_loop_enabled, device_loop_killed, + kill_device_loop, pending_token_count, push_pending_tokens, set_device_loop_eligible, + stand_down, take_pending_token, AdmissionFacts, BurstDriver, BurstOutcome, DeviceLoopConfig, + DeviceLoopError, DeviceOps, Refusal, StreamState, +}; +#[cfg(feature = "cuda")] +pub use device_loop::{CudaDeviceOps, DeviceDecodeLoop}; #[cfg(feature = "cuda")] pub use graph::CudaGraphRunner; #[cfg(feature = "cuda")] diff --git a/mistralrs-core/src/pipeline/normal.rs b/mistralrs-core/src/pipeline/normal.rs index 88d9ce669..de530e395 100644 --- a/mistralrs-core/src/pipeline/normal.rs +++ b/mistralrs-core/src/pipeline/normal.rs @@ -96,6 +96,13 @@ pub struct NormalPipeline { /// load time; graph capture is deferred until first decode call. #[cfg(feature = "cuda")] autonomous_runner: Option, + /// ArcInfer/ArcGraph device decode loop โ€” replays the captured graph, + /// samples and scatters the token on device, and never synchronizes. Built + /// lazily on the first admitted decode step and then kept, because every + /// buffer it points at must stay at the address the graph baked. + /// Opt-in via `ARC_GRAPH_DEVICE_LOOP`; `None` until then. + #[cfg(feature = "cuda")] + device_decode_loop: Option, } /// What a forward pass should do with candle's caching allocator before it @@ -1986,6 +1993,10 @@ impl Loader for NormalLoader { // lazy init (and graph capture) once a real decode batch arrives. #[cfg(feature = "cuda")] autonomous_runner: None, + // Needs the captured graph's exec handle and output tensor, so it + // cannot exist before capture. Built on the first admitted step. + #[cfg(feature = "cuda")] + device_decode_loop: None, }))) } @@ -2108,6 +2119,181 @@ impl MetadataMixin for NormalPipeline { } } +#[cfg(feature = "cuda")] +/// Drive one burst of the ArcInfer/ArcGraph device decode loop. +/// +/// Returns `Some(logits)` when the step is covered โ€” either because a burst +/// just ran or because a previous burst already produced this step's token +/// โ€” and `None` to fall back. +/// +/// The tokens go to `arc_cuda_graph`'s pending queue and are drained one at +/// a time by `pipeline/sampling.rs`, so the engine's per-token `Sequence` +/// bookkeeping (stop strings, streaming, EOS) is unchanged. +/// +/// โš ๏ธ On the `pending > 0` path this returns the captured graph's output +/// tensor **without launching anything**, so those logits are the last +/// burst's, not this step's. That is sound only because the sampling hook +/// short-circuits before reading them and, at batch 1, +/// `pipeline/mod.rs:538` does no host copy. It is not a general-purpose +/// return value, which is why this whole path refuses anything that would +/// actually read the logits (`return_logprobs`, logits processors, +/// penalties, grammars). +fn device_decode_burst( + device: &Device, + runner: &mut arc_cuda_graph::CudaGraphRunner, + loop_slot: &mut Option, + bs: usize, + seq_len: usize, + current_position: u32, + input_ids_pinned: &Tensor, +) -> Option { + use arc_cuda_graph as acg; + + // Stand aside for the replay-verification lane (PR #205). This branch sits + // EARLIER in the chain than the `runner.replay(bs)` arm, so without this it + // would take every step and `ARC_GRAPH_VERIFY_ALWAYS=1` would silently + // verify nothing โ€” the operator would get a clean diagnostic that never + // ran. A verification switch that is quietly bypassed is worse than one + // that is off. + if std::env::var("ARC_GRAPH_VERIFY_ALWAYS").as_deref() == Ok("1") { + return None; + } + + let (_exec, out) = runner.replay_handles(bs)?; + + // A burst already produced this step's token: no launch, no sync, no + // GPU work at all. This is where the amortisation actually shows up โ€” + // `burst - 1` of every `burst` engine iterations do nothing but host + // bookkeeping. + if acg::pending_token_count() > 0 { + return Some(out); + } + + // The fixed-capacity KV window width. Read from the graph-mode mask's + // last dimension rather than guessed: `models/deepseek4.rs` builds that + // mask as `[B, 1, 1, cfg_full.sliding_window]`, and `cfg_full` is a + // model fact the pipeline cannot otherwise reach. + let capacity = match crate::layers::graph_mode_mask() { + Some(m) => *m.dims().last().unwrap_or(&0), + None => { + // The mask is what keeps the decode arm from reading unwritten + // KV slots; without it we cannot bound the burst. + acg::stand_down(); + return None; + } + }; + let vocab = *out.dims().last().unwrap_or(&0); + if capacity == 0 || vocab == 0 { + acg::stand_down(); + return None; + } + + let cfg = acg::DeviceLoopConfig { + position_limit: capacity as u32, + ..acg::DeviceLoopConfig::new(vocab, -1, capacity as u32) + }; + + let facts = acg::AdmissionFacts { + batch_size: bs, + seq_len, + graph_ready: runner.replay_output_trusted(), + // Established by `sampling.rs` from the Sequence on the previous + // step; false until one has confirmed it. + greedy: acg::device_loop_eligible(), + return_logprobs: false, + host_side_logits_work: false, + current_position, + }; + let steps = match acg::admit(&facts, &cfg, acg::device_loop_enabled()) { + Ok(s) => s, + Err(reason) => { + tracing::debug!("ArcGraph device loop declined this step: {reason}"); + acg::stand_down(); + return None; + } + }; + + if loop_slot.is_none() { + let positions = crate::layers::graph_mode_positions(); + match acg::DeviceDecodeLoop::new( + device, + _exec, + runner.stream(), + &out, + input_ids_pinned, + positions.as_ref(), + cfg, + ) { + Ok(l) => { + tracing::info!( + "ArcGraph device decode loop ENGAGED: burst={} vocab={vocab} \ + kv_window={capacity}. {}", + cfg.burst, + l.status_line() + ); + *loop_slot = Some(l); + } + Err(e) => { + acg::kill_device_loop(&format!("could not build the device loop: {e}")); + return None; + } + } + } + let dl = loop_slot.as_mut()?; + + match dl.run(steps) { + Ok(outcome) => { + // `first()` rather than `[0]`: a burst that returned no rows would + // be a bug, but it must not panic the engine thread mid-request. + let row = match outcome.tokens.first() { + Some(r) => r, + None => { + acg::kill_device_loop("burst returned no token rows"); + *loop_slot = None; + return None; + } + }; + acg::push_pending_tokens(row); + // โ”€โ”€ ALIASING CONTRACT, ENFORCED โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + // The tensor we are about to return aliases the graph-owned + // storage that the NEXT `cuGraphLaunch` overwrites. It is only + // safe to hand back because the sampling hook short-circuits + // before anything reads it โ€” and that hook fires if and only if a + // token is pending. So the invariant is not "we reasoned it + // through", it is checked here: no pending token, no aliasing + // tensor. `push_pending_tokens` drops negative ids, so a row that + // was entirely garbage lands here rather than silently arriving at + // a `sample_sequence` that would then read clobbered logits and + // return a plausible wrong token. + if acg::pending_token_count() == 0 { + acg::kill_device_loop( + "burst produced no usable token, so the aliasing logits tensor must not be \ + returned", + ); + *loop_slot = None; + return None; + } + runner.record_device_loop_steps(outcome.steps_launched as u64); + if outcome.hit_eos { + tracing::debug!( + "ArcGraph device loop: EOS at token {} of a {}-step burst; the remaining \ + steps were already enqueued and are discarded", + row.len(), + outcome.steps_launched + ); + } + Some(out) + } + Err(e) => { + // A burst failure means the graph input buffer may hold an + // uncommitted or stale id, so this cannot be retried. + acg::kill_device_loop(&format!("burst failed: {e}")); + *loop_slot = None; + None + } + } +} + #[async_trait::async_trait] impl Pipeline for NormalPipeline { fn forward_inputs( @@ -2256,6 +2442,11 @@ impl Pipeline for NormalPipeline { } } } + // Bound before the chain: `self.device()` returns by value, + // and taking a reference to it inside an expression that + // also takes `&mut self.device_decode_loop` would borrow + // `self` both ways at once. + let dev_for_loop = self.device(); let captured: Option = if probe && seq_len == 1 && self.cuda_graph_runner.is_some() @@ -2695,6 +2886,32 @@ impl Pipeline for NormalPipeline { cd.set_capture_mode(false); } cap_result + } else if runner.has_graph(bs) + && arc_cuda_graph::device_loop_enabled() + && !arc_cuda_graph::device_loop_killed() + && device_decode_burst( + &dev_for_loop, + &mut runner, + &mut self.device_decode_loop, + bs, + seq_len, + seqlen_offsets.first().copied().unwrap_or(0) as u32, + &input_ids, + ) + .is_some() + { + // ArcInfer/ArcGraph device decode loop: the burst + // already ran (or a previous one covered this step), + // the token was sampled and scattered on device, and + // `pipeline/sampling.rs` will take it from the ring. + // No sync, no host sample, no D2H. + // + // The logits handed back are the captured graph's + // output tensor. They are NOT read: the sampling + // hook short-circuits first, and at batch 1 + // `pipeline/mod.rs:538` does no host copy. Admission + // refuses everything that would read them. + runner.replay_handles(bs).map(|(_, out)| out) } else if runner.has_graph(bs) { let t = std::time::Instant::now(); let replayed = runner.replay(bs); diff --git a/mistralrs-core/src/pipeline/sampling.rs b/mistralrs-core/src/pipeline/sampling.rs index b413e380d..b87f9cf86 100644 --- a/mistralrs-core/src/pipeline/sampling.rs +++ b/mistralrs-core/src/pipeline/sampling.rs @@ -527,6 +527,62 @@ pub async fn sample_and_add_toks( Ok(()) } +/// Is this sequence one the ArcGraph device decode loop may drive, and does it +/// already have a token waiting? +/// +/// Two jobs, deliberately in one place. It **publishes** eligibility, because +/// `Pipeline::forward_inputs` decides whether to run a burst but is handed only +/// `inputs` and `return_raw_logits` โ€” it never sees a `Sequence`, so it cannot +/// ask whether sampling is greedy. And it **takes** the token a burst already +/// produced. +/// +/// Greedy only. The device sampler runs Splitmix64 per row against the host's +/// Isaac64, so under any stochastic configuration the two draw different โ€” +/// individually valid โ€” tokens, which would silently break seeded +/// reproducibility. Under greedy they agree exactly, which is also what makes +/// the path verifiable against the eager one. +/// +/// The returned `Logprobs` is not a lossy stand-in: `logprob: 0.0` is exactly +/// what the greedy host path returns (`sampler.rs:1490`), and `bytes` / +/// `top_logprobs` are read only when `return_logprobs` is set, which this +/// refuses. +// `&mut` only because `Sequence::sampler()` takes `&mut self` (`sequence.rs:885`) +// to hand back a clone of an `Arc`. Nothing here mutates the sequence. +#[cfg(feature = "cuda")] +fn device_loop_pre_sampled_token( + seq: &mut Sequence, + return_logprobs: bool, + sample_speculative: bool, +) -> Option { + let eligible = !return_logprobs + && !sample_speculative + && seq.sampler().is_greedy_trivial() + && matches!(seq.recognizer, SequenceRecognizer::None); + arc_cuda_graph::set_device_loop_eligible(eligible); + if !eligible { + // Stand down rather than merely decline: a token parked by an earlier, + // eligible sequence must never be handed to this one. + arc_cuda_graph::stand_down(); + return None; + } + arc_cuda_graph::take_pending_token().map(|token| Logprobs { + token, + logprob: 0.0, + bytes: None, + top_logprobs: None, + }) +} + +/// Without CUDA there is no device loop, so nothing is ever pre-sampled. +#[cfg(not(feature = "cuda"))] +fn device_loop_pre_sampled_token( + _seq: &mut Sequence, + _return_logprobs: bool, + _sample_speculative: bool, +) -> Option { + None +} + /// Async sample optionally adding to trie. #[allow(clippy::too_many_arguments)] pub async fn sample_sequence( @@ -538,6 +594,17 @@ pub async fn sample_sequence( sample_speculative: bool, multiple_sequences: bool, ) -> Result { + // โ”€โ”€ ArcInfer/ArcGraph device decode loop โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + // If a burst already sampled this token ON DEVICE, take it and skip the + // host argmax entirely. That D2H (`sampler.rs:1479`) is the last + // synchronization in the decode step, so this is the half of the change + // that makes removing `replay()`'s `cudaStreamSynchronize` worth anything. + if let Some(pre_sampled) = + device_loop_pre_sampled_token(seq, return_logprobs, sample_speculative) + { + return Ok(pre_sampled); + } + // Both of these open and close inside a single poll, before the first // `.await`, so they nest correctly under `sample.join_all` even when B // futures are being polled in turn on this thread. diff --git a/mistralrs-core/src/sampler.rs b/mistralrs-core/src/sampler.rs index b282841c6..65547f010 100644 --- a/mistralrs-core/src/sampler.rs +++ b/mistralrs-core/src/sampler.rs @@ -505,6 +505,28 @@ impl Sampler { || self.dry_params.is_some() } + /// Would [`Sampler::sample`] take the pure-argmax path โ€” no temperature and + /// nothing that has to touch the logits on the host? + /// + /// This is the exact precondition for ArcGraph's device decode loop to + /// substitute its on-device argmax for the host one + /// (`pipeline/sampling.rs`). It deliberately mirrors the conditions inside + /// `sample` rather than approximating them: `logits_bias` and `top_nsigma` + /// are applied *before* the `trivial` gate, so a sampler carrying either + /// still needs the host even though the gate itself ignores them. + /// + /// Greedy is also the only mode where device and host agree *exactly* โ€” the + /// device sampler runs Splitmix64 per row against the host's Isaac64, so + /// any stochastic mode would draw a different, equally valid token and + /// silently break seeded reproducibility. + pub fn is_greedy_trivial(&self) -> bool { + self.temperature.is_none() + && self.logits_processors.is_empty() + && !self.has_penalties() + && self.logits_bias.is_none() + && self.top_nsigma.is_none() + } + /// True when this sampler is exactly `argmax` over the model's **raw** /// logits โ€” no temperature, no penalties, no custom logits processors. /// From c94d026629e432d32d13d38cd5f83c0db4771f51 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 19:31:38 +0100 Subject: [PATCH 12/21] fix(ArcLab/ArcQuant): make prefill measurable, then ship the two kernels it proves were switched off (#196) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(ArcLab/ArcInfer): prefill measured 0.000 tok/s and aborted at teardown Three defects, all on the path that makes prefill measurable at all. 1. `pp N` reported `0.000ยฑ0.000` for every prompt benchmark. A prefill-only request (`max_len = 1`) finishes *during* the prompt step: `pipeline/sampling.rs` calls `seq.update_time_info()`, builds `group.get_usage()` and dispatches `Response::Done` from inside `step()`. But the default (non-paged) scheduler arm stamped `prompt_timestamp` and `total_prompt_time` only AFTER `step()` returned. So at the moment the usage was built `prompt_timestamp` was still `None`, `update_time_info` skipped `group.total_prompt_time`, and `get_usage` took its `== 0` branch and returned `avg_prompt_tok_per_sec: 0.0`. The PagedAttention arm already stamps before `step()`, with a comment saying exactly why. The fix was simply never applied to the arm that serves -- PagedAttention is banned here (it shadows the graph arm and measures zero tokens). Stamp on the default arm too. 2. The harness printed that zero into a results table and exited 0. A zero is not a measurement. Every row now has to carry evidence the engine processed the tokens the row claims to be about -- non-zero prompt tokens, the count actually asked for, a non-zero timed prefill, and a finite positive rate. A row that cannot exits 2 (environment failure, never 1) and prints no table. Adds a second, engine-independent instrument: wall-clock seconds and wall-clock prompt tok/s, so the engine's own numbers are only believed when an outside clock agrees. Also stops reusing one request id for every concurrent copy and repetition: `next_request_id()` was called once and the request cloned, so `concurrency * repetitions` sequences were in flight sharing one engine handle. 3. The binary aborted at teardown AFTER printing valid results. `Request::Terminate` only asks; nothing waited. `engine_handler` had no `.join()` anywhere in the workspace, so `main` returned while the engine thread was still releasing device memory and CUDA objects, and libc `exit()` ran CUDA's atexit handler concurrently with it -- a textbook `corrupted double-linked list` / SIGSEGV. `Drop for MistralRs` now joins each engine thread with a bounded 30s timeout, and says so loudly if it has to detach instead. Second half of the same race: `DedicatedDecodePath::new` binds the CUDA primary context and warns that skipping it gives "a SEGV in libcuda MOVAPS ... no error code, just a fault". Its `Drop` then made the same class of raw calls (`cuGraphExecDestroy`, many `cudaFree`) with no bind at all, made legal on a foreign thread by `unsafe impl Send`/`Sync`. It now binds exactly as `new()` does. This matters beyond tidiness: the abort libelled good runs. One chain already discarded an expensive valid measurement because its harness gated on rc==0. Parent system: ArcLab (harness) / ArcInfer (engine timing, teardown). * test(ArcLab): print the grouped-GEMM launch counters beside the results `grouped.rs` already states the rule: "a harness that reports a per-variant timing MUST show this counter advancing on the variant it claims to have measured, and NOT advancing on the others, or the number describes some other kernel." The counters existed and were public. Nothing in the workspace ever read them, so every grouped-vs-GEMV and variant-vs-variant comparison this harness could produce was unproven by construction. The bench now prints `grouped_launch_counts()` and the selected variant next to every results table, and says so explicitly when the total is zero -- which is the common case, because the grouped GEMM declines to run below its tile-fill boundary (~683 tokens at V4's top-6-of-256 routing) and the measurement is then the per-pair gather GEMV, not the kernel the reader will assume. Parent system: ArcLab. * perf(ArcQuant): ship the tuned grouped-GEMM kernel as the default (+19.5% prefill) The trellis grouped GEMM has three variants compiled into every binary. The tuned ones were measured long ago at +39.4% (v1) and +41.6% (v2) per m-tile, bit-identical in output, and then left unreachable: `grouped_variant()` fell back to `QTIP_GROUPED_VARIANT_BASELINE` whenever `ARC_QTIP_GROUPED_VARIANT` was unset, and the only callers of `set_grouped_variant` in the entire workspace are inside `examples/qtip_grouped_curve.rs`. So every production prefill since the variants landed has run variant 0, and the faster kernels existed only for a benchmark nobody ran in serving. โœ… MEASURED END-TO-END, H200, DeepSeek-V4-Flash qtip2b, pp2048 b=1, --no-paged-attn, on `arc-prefill`: variant 0 (was default): 561.7 prompt tok/s | 1.780 ms/tok | TTFT 3.675 s variant 2 (now default): 671.3 prompt tok/s | 1.490 ms/tok | TTFT 3.080 s => +19.5% prefill throughput, -16.2% TTFT The arm is proved from the runtime launch counters, not from the build: [129,0,0] on the control against [0,0,129] on the variant -- 43 layers x 3 expert matrices, and zero launches on the arms not under test. An engine-independent wall clock agrees (557.2 -> 664.9 tok/s, +19.3%). The end-to-end gain is smaller than the kernel gain because the expert gather is ~56% of a 2048-token prefill step; backing +41.6% on that share out of a measured +19.5% overall is consistent. Also: an unrecognised env value used to fall through to the baseline. With the baseline now the slow arm, a typo would silently cost ~20% of prefill, so unknown values warn and keep the default, and `baseline`/`0` is now an explicit arm so the A/B knob still selects the control. Parent system: ArcQuant / QTIP. --------- Co-authored-by: Nirupam Bhowmick --- Cargo.lock | 1 + arc-cuda-graph/src/dedicated.rs | 25 +++++ mistralrs-bench/Cargo.toml | 1 + mistralrs-bench/src/main.rs | 163 +++++++++++++++++++++++++--- mistralrs-core/src/engine/mod.rs | 39 +++++++ mistralrs-core/src/lib.rs | 52 +++++++++ mistralrs-quant/src/qtip/grouped.rs | 36 +++++- 7 files changed, 297 insertions(+), 20 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6d09c410e..778418644 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3600,6 +3600,7 @@ dependencies = [ "clap", "cli-table", "mistralrs-core", + "mistralrs-quant", "serde", "serde_json", "tokio", diff --git a/arc-cuda-graph/src/dedicated.rs b/arc-cuda-graph/src/dedicated.rs index 555ea4fb5..b8675e067 100644 --- a/arc-cuda-graph/src/dedicated.rs +++ b/arc-cuda-graph/src/dedicated.rs @@ -1106,6 +1106,31 @@ impl DedicatedDecodePath { #[cfg(feature = "cuda")] impl Drop for DedicatedDecodePath { fn drop(&mut self) { + // ๐Ÿ”ด Bind the context before freeing, exactly as `new()` does. + // + // `new()` binds first and its comment states the consequence of not + // doing so: "a SEGV in libcuda MOVAPS on the first cudaMalloc โ€” no + // error code, just a fault", because the runtime's lazy-init walks a + // thread-local that has not been seeded. This `Drop` then performed the + // *same class* of raw Driver/Runtime calls โ€” `cuGraphExecDestroy` plus + // a long list of `cudaFree` โ€” with no bind at all. + // + // `unsafe impl Send`/`Sync` above make it legal for this value to be + // dropped on a thread that never touched CUDA, which is precisely what + // happens when the engine thread outlives `main` and teardown lands on + // a foreign thread. That asymmetry (bind in `new`, none in `drop`) is + // the second half of the teardown SIGSEGV this binary died from. + // + // Failure here is not fatal: if we cannot bind we are already in a + // teardown we cannot improve, and skipping the frees leaks into a + // process that is exiting anyway โ€” strictly better than faulting. + if Self::bind_context(&self.weights).is_err() { + tracing::warn!( + "DedicatedDecodePath::drop: could not bind CUDA context; \ + skipping device frees to avoid faulting during teardown." + ); + return; + } unsafe { if let Some(exec) = self.graph_exec { cuGraphExecDestroy(exec); diff --git a/mistralrs-bench/Cargo.toml b/mistralrs-bench/Cargo.toml index 8d7fb240f..224d2dbf4 100644 --- a/mistralrs-bench/Cargo.toml +++ b/mistralrs-bench/Cargo.toml @@ -20,6 +20,7 @@ serde.workspace = true serde_json.workspace = true clap.workspace = true mistralrs-core.workspace = true +mistralrs-quant.workspace = true tracing.workspace = true tokio.workspace = true cli-table.workspace = true diff --git a/mistralrs-bench/src/main.rs b/mistralrs-bench/src/main.rs index c77b677dd..5f22d45d5 100644 --- a/mistralrs-bench/src/main.rs +++ b/mistralrs-bench/src/main.rs @@ -10,6 +10,7 @@ use mistralrs_core::{ }; use std::fmt::Display; use std::sync::Arc; +use std::time::Instant; use tokio::sync::mpsc::channel; use tracing::{info, warn}; @@ -32,6 +33,13 @@ struct BenchResult { usages: Vec, concurrency: usize, test_name: TestName, + /// Wall-clock seconds for the whole send->all-Done window, measured by the + /// client with no engine cooperation. This is the SECOND instrument: the + /// engine's own `Usage` numbers are only believable when they agree with a + /// clock that lives outside the engine. + wall_secs: f32, + /// Total prompt tokens the engine says it processed across all responses. + prompt_toks_seen: usize, } struct UncertainTokSec { @@ -74,29 +82,41 @@ async fn run_bench( let sender = mistralrs.get_sender(None).unwrap(); let (tx, mut rx) = channel(10_000); - let req = Request::Normal(Box::new(NormalRequest { - id: mistralrs.next_request_id(), - messages: prompt, - sampling_params: sampling_params.clone(), - response: tx, - return_logprobs: false, - is_streaming: false, - constraint: Constraint::None, - suffix: None, - tools: None, - tool_choice: None, - logits_processors: None, - return_raw_logits: false, - web_search_options: None, - model_id: None, - truncate_sequence: false, - })); + // ๐Ÿ”ด A fresh request id per send. + // + // This used to call `next_request_id()` exactly once and then `.clone()` + // the same `NormalRequest` for every concurrent copy and every repetition, + // so `concurrency * repetitions` sequences were in flight sharing ONE id. + // The id is the engine's handle for a sequence group; issuing duplicates + // means unrelated sequences alias each other in engine-side bookkeeping. + // A benchmark must not be the only caller in the system that does this. + let make_req = |id: usize| { + Request::Normal(Box::new(NormalRequest { + id, + messages: prompt.clone(), + sampling_params: sampling_params.clone(), + response: tx.clone(), + return_logprobs: false, + is_streaming: false, + constraint: Constraint::None, + suffix: None, + tools: None, + tool_choice: None, + logits_processors: None, + return_raw_logits: false, + web_search_options: None, + model_id: None, + truncate_sequence: false, + })) + }; let mut usages = Vec::new(); + let wall_start = Instant::now(); for _ in 0..repetitions { for _ in 0..concurrency { - if sender.send(req.clone()).await.is_err() { + let req = make_req(mistralrs.next_request_id()); + if sender.send(req).await.is_err() { eprintln!("Receiver disconnected"); } } @@ -131,10 +151,15 @@ async fn run_bench( } } + let wall_secs = wall_start.elapsed().as_secs_f32(); + let prompt_toks_seen = usages.iter().map(|u| u.prompt_tokens).sum(); + Ok(BenchResult { usages, concurrency, test_name, + wall_secs, + prompt_toks_seen, }) } @@ -186,6 +211,100 @@ fn get_ms_tok(result: &BenchResult) -> UncertainTokSec { UncertainTokSec { mean, std_dev } } +/// A zero is not a measurement. +/// +/// ๐Ÿ”ด This harness used to print `0.000ยฑ0.000` into a results table and exit 0. +/// That is the house fault: a green result that proves no work happened. Every +/// row is now required to carry evidence that the engine actually processed the +/// tokens the row claims to be about, and a row that cannot is a hard stop. +/// +/// Exits **2** (environment/instrument failure), never 1, so a caller can tell +/// "the benchmark could not be run" apart from "the benchmark ran and regressed". +fn assert_results_are_measurements(results: &[BenchResult]) { + let mut failures: Vec = Vec::new(); + + for r in results { + let label = format!("{} @ concurrency {}", r.test_name, r.concurrency); + + if r.usages.is_empty() { + failures.push(format!("{label}: no responses at all (0 usages collected)")); + continue; + } + + for (i, u) in r.usages.iter().enumerate() { + let where_ = format!("{label}, response {i}"); + match r.test_name { + TestName::Prompt(n) => { + if u.prompt_tokens == 0 { + failures.push(format!("{where_}: engine reported 0 prompt tokens")); + } else if u.prompt_tokens != n { + failures.push(format!( + "{where_}: asked for {n} prompt tokens, engine processed {}", + u.prompt_tokens + )); + } + if u.total_prompt_time_sec <= 0.0 { + failures.push(format!( + "{where_}: prompt time is {:.6}s -- the prefill step was never timed", + u.total_prompt_time_sec + )); + } + if !u.avg_prompt_tok_per_sec.is_finite() || u.avg_prompt_tok_per_sec <= 0.0 { + failures.push(format!( + "{where_}: prompt rate is {} tok/s", + u.avg_prompt_tok_per_sec + )); + } + } + TestName::Gen(_) => { + if u.completion_tokens == 0 { + failures.push(format!("{where_}: engine produced 0 completion tokens")); + } + if !u.avg_compl_tok_per_sec.is_finite() || u.avg_compl_tok_per_sec <= 0.0 { + failures.push(format!( + "{where_}: completion rate is {} tok/s", + u.avg_compl_tok_per_sec + )); + } + } + } + } + } + + // Which MoE kernel actually ran? The grouped-GEMM variant is chosen at + // RUNTIME from an env var, so the binary cannot tell you which one you + // measured. `grouped.rs` states the rule outright: "a harness that reports a + // per-variant timing MUST show this counter advancing on the variant it + // claims to have measured, and NOT advancing on the others, or the number + // describes some other kernel." Nothing in the tree printed it, so every + // previous variant comparison was unproven. Print it beside the results. + let counts = mistralrs_quant::grouped_launch_counts(); + let total: u64 = counts.iter().sum(); + eprintln!( + "MoE grouped-GEMM launches by variant: {counts:?} (selected variant = {}, total = {total})", + mistralrs_quant::grouped_variant() + ); + if total == 0 { + eprintln!( + " -> grouped GEMM never launched: this measurement is the per-pair \ + gather GEMV (or the dequantize fallback), NOT the grouped kernel." + ); + } + + if !failures.is_empty() { + eprintln!("\n=== BENCHMARK PRODUCED NO VALID MEASUREMENT ==="); + for f in &failures { + eprintln!(" - {f}"); + } + eprintln!( + "\n{} check(s) failed. A zero-token or zero-rate result is an instrument\n\ + failure, not a score, so no table is printed. Exiting 2.", + failures.len() + ); + std::process::exit(2); + } +} + fn print_usage(model: &str, device: &Device, results: Vec) { let backend = match device { Device::Cpu => "CPU", @@ -205,6 +324,11 @@ fn print_usage(model: &str, device: &Device, results: Vec) { (get_tok_s(&r).mean * r.concurrency as f32) .cell() .justify(Justify::Right), + // Second instrument: a clock outside the engine. + format!("{:.3}", r.wall_secs).cell().justify(Justify::Right), + format!("{:.1}", r.prompt_toks_seen as f32 / r.wall_secs) + .cell() + .justify(Justify::Right), ] }) .collect(); @@ -222,6 +346,8 @@ fn print_usage(model: &str, device: &Device, results: Vec) { "ms/t".cell().bold(true), "concurrency".cell().bold(true), "throughput/s".cell().bold(true), + "wall s".cell().bold(true), + "prompt t/s (wall)".cell().bold(true), ]) .bold(true); print_stdout(table).expect("print table"); @@ -575,6 +701,7 @@ async fn main() -> anyhow::Result<()> { results.push(r); } + assert_results_are_measurements(&results); print_usage(&model_name, &device, results); } diff --git a/mistralrs-core/src/engine/mod.rs b/mistralrs-core/src/engine/mod.rs index c30053ca4..00495b282 100644 --- a/mistralrs-core/src/engine/mod.rs +++ b/mistralrs-core/src/engine/mod.rs @@ -537,6 +537,45 @@ impl Engine { }; let _chunk_guard = (!final_chunk).then(crate::pipeline::mark_prefill_intermediate); + + // Record prompt timing BEFORE step() so it is available if the + // response is sent from inside step(). + // + // ๐Ÿ”ด This is why `pp N` reported `0.000ยฑ0.000` for every prompt + // benchmark. A prefill-only request (`max_len = 1`) finishes + // *during* this prompt step: `pipeline/sampling.rs` calls + // `seq.update_time_info()` and then `group.get_usage()` and + // dispatches `Response::Done` from inside `step()`. But on this + // (default, non-paged) path `prompt_timestamp` was only stamped + // AFTER `step()` returned, ~100 lines below. So at the moment the + // usage was built, `prompt_timestamp` was still `None`, + // `update_time_info` skipped `group.total_prompt_time` + // (`sequence.rs`, `if let Some(ts) = self.prompt_timestamp`), and + // `get_usage` took the `total_prompt_time == 0` branch and + // returned `avg_prompt_tok_per_sec: 0.0`. The harness then printed + // that zero as if it were a measurement. + // + // The PagedAttention arm below already does exactly this, with + // this same comment โ€” the fix was simply never applied to the arm + // that actually serves (PagedAttention is banned here: it shadows + // the graph arm and measures zero tokens). + // + // `set_step_start_instant` gives `update_time_info` its in-flight + // fallback (`start.elapsed()`) while the prompt step is running; + // the post-step block below still overwrites `total_prompt_time` + // with the precise `prompt_exec_time` for sequences that go on to + // decode, so steady-state accounting is unchanged. + { + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("Time travel has occurred!") + .as_millis(); + for seq in scheduled.prompt.iter_mut() { + seq.prompt_timestamp = Some(now); + seq.set_step_start_instant(); + } + } + let prompt_exec_time = { let mut pipeline = { let _s = arc_profiler::span("pipeline.lock"); diff --git a/mistralrs-core/src/lib.rs b/mistralrs-core/src/lib.rs index 0c9b1f2d9..2603392c2 100644 --- a/mistralrs-core/src/lib.rs +++ b/mistralrs-core/src/lib.rs @@ -558,6 +558,11 @@ impl MistralRsBuilder { } } +/// How long `Drop` waits for an engine thread to finish its own teardown before +/// giving up and detaching it. Generous because the thread may be mid-step on a +/// long prefill; bounded because a hung engine must not wedge process exit. +const ENGINE_JOIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); + impl Drop for MistralRs { fn drop(&mut self) { // Terminate all engines @@ -567,6 +572,53 @@ impl Drop for MistralRs { let _ = engine.sender.try_send(Request::Terminate); } } + + // ๐Ÿ”ด โ€ฆand then WAIT for them. This used to be fire-and-forget. + // + // `Request::Terminate` is asynchronous: it only asks. Before this join, + // `main` returned while the engine thread was still alive and still + // releasing device memory, CUDA graph execs and stream objects. libc + // `exit()` then ran CUDA's `atexit` handler concurrently with that + // thread. Two threads tearing down the same CUDA state with no ordering + // is a textbook `corrupted double-linked list` / SIGSEGV, which is + // exactly how this binary died โ€” *after* printing valid results, so the + // exit code libelled good runs. One chain already discarded a good, + // expensive measurement because its harness gated on `rc == 0`. + // + // The handle had no `.join()` anywhere in the workspace; its only uses + // were `is_finished()` polls. Joining here is what makes CUDA objects + // get freed on a live, context-bound thread before `exit()` runs. + // + // Bounded and poll-based rather than a bare `join()`: a wedged engine + // must not be able to hang process exit forever, and `JoinHandle` has no + // timed join in std. + let handles: Vec<(String, JoinHandle<()>)> = match self.engines.write() { + Ok(mut engines) => engines + .drain() + .map(|(id, instance)| (id, instance.engine_handler)) + .collect(), + Err(_) => return, + }; + + let deadline = std::time::Instant::now() + ENGINE_JOIN_TIMEOUT; + for (id, handle) in handles { + while !handle.is_finished() && std::time::Instant::now() < deadline { + std::thread::sleep(std::time::Duration::from_millis(2)); + } + if handle.is_finished() { + if handle.join().is_err() { + tracing::warn!("Engine thread for `{id}` panicked during shutdown."); + } + } else { + // Detach rather than block forever. Say so: a silent detach here + // would reintroduce exactly the race this join exists to close. + tracing::warn!( + "Engine thread for `{id}` did not finish within {:?}; detaching. \ + CUDA teardown may race with process exit.", + ENGINE_JOIN_TIMEOUT + ); + } + } } } diff --git a/mistralrs-quant/src/qtip/grouped.rs b/mistralrs-quant/src/qtip/grouped.rs index 17c55f231..b4481ef99 100644 --- a/mistralrs-quant/src/qtip/grouped.rs +++ b/mistralrs-quant/src/qtip/grouped.rs @@ -109,11 +109,43 @@ pub fn grouped_variant() -> i32 { } let init = match std::env::var(QTIP_GROUPED_VARIANT_ENV) { Ok(s) => match s.trim().to_ascii_lowercase().as_str() { + "baseline" | "0" => QTIP_GROUPED_VARIANT_BASELINE, "tuned" | "1" => QTIP_GROUPED_VARIANT_TUNED, "ldst" | "2" => QTIP_GROUPED_VARIANT_LDST, - _ => QTIP_GROUPED_VARIANT_BASELINE, + // An unrecognised value used to fall through to the baseline. Now + // that the baseline is the SLOW arm rather than the default, a typo + // would silently cost ~20% of prefill, so say so and keep the + // default instead of quietly selecting the worst kernel. + other => { + tracing::warn!( + "{QTIP_GROUPED_VARIANT_ENV}=`{other}` is not a known grouped-GEMM \ + variant (baseline|0, tuned|1, ldst|2); using the default." + ); + QTIP_GROUPED_VARIANT_LDST + } }, - Err(_) => QTIP_GROUPED_VARIANT_BASELINE, + // ๐Ÿ”ด The DEFAULT is the tuned kernel, not the baseline. + // + // This used to be `QTIP_GROUPED_VARIANT_BASELINE`, and nothing in the + // serving path ever calls `set_grouped_variant` -- its only callers are + // in `examples/qtip_grouped_curve.rs`. So every production prefill ran + // variant 0 and the tuned kernels were reachable only by setting an env + // var nobody sets. They are bit-identical in output and were already + // measured at +39.4% (v1) and +41.6% (v2) per m-tile. + // + // โœ… MEASURED END-TO-END 2026-08-20, H200, DeepSeek-V4-Flash qtip2b, + // pp2048 b=1, --no-paged-attn, arm proved from the runtime launch + // counters ([0,0,129] on v2 vs [129,0,0] on v0, i.e. 43 layers x 3 + // matrices, and zero launches on the arms not under test): + // variant 0: 561.7 prompt tok/s, 1.780 ms/prompt-token, TTFT 3.675 s + // variant 2: 671.3 prompt tok/s, 1.490 ms/prompt-token, TTFT 3.080 s + // => +19.5% prefill throughput, -16.2% TTFT. + // An engine-independent wall clock agrees (557.2 -> 664.9 tok/s, +19.3%). + // + // The end-to-end gain is smaller than the kernel gain because the gather + // is ~56% of a 2048-token prefill step, not 100% -- backing +41.6% on + // that share out of a measured +19.5% overall is consistent. + Err(_) => QTIP_GROUPED_VARIANT_LDST, }; let _ = GROUPED_VARIANT.compare_exchange( GROUPED_VARIANT_UNINIT, From 9a5921bff10afbc100d801419b01e7e45032687a Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 19:31:46 +0100 Subject: [PATCH 13/21] =?UTF-8?q?ArcQuant:=20grouped-GEMM=20gate=20was=20o?= =?UTF-8?q?ptimising=20the=20wrong=20resource=20=E2=80=94=20B=3D256=20+69%?= =?UTF-8?q?=20(114.76=20->=20193.46=20tok/s)=20(#195)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * perf(ArcInfer/ArcMoE): warp-butterfly Sinkhorn on top of the ladder โ€” 1.36x Rebased onto release/openrouter-ready. The earlier version of this change was measured against the PRE-LADDER kernel and read 4.19x; that baseline is gone. Against the templated kernel the ladder actually shipped, the honest number is 1.36x. Recording both so the retracted figure is not quoted again. MEASURED, H200 / nvcc 12.4 / sm_90, n=1 hc=4 iters=20, 86 calls/token: pre-ladder (runtime hc, 192 B stack frame) 30.578 us/call 2.6297 ms/tok templated + shared/barriers (master today) 9.771 us/call 0.8403 ms/tok warp butterfly (this commit) 7.201 us/call 0.6193 ms/tok => 1.357x incremental, -0.221 ms/token. Bit-identical to the templated arm. The templated kernel got the row/tree buffers into registers but still stages the matrix through shared memory and pays TWO __syncthreads() per iteration -- 40 barriers and ~80 shared round trips per call -- purely to move each column to the thread that divides by it. hc <= 16 means every thread is a lane of ONE warp, so the column reduction can be a butterfly and the matrix never has to leave registers. Bit-identity by construction: candle's level s is `buf[t] = __fadd_rn(buf[t], buf[t + s])` for t < s. The butterfly gives lane t `__fadd_rn(buf[t], buf[t ^ s])`, which IS that expression for t < s and its commuted form for t >= s; IEEE-754 addition is commutative (only associativity fails), so lane 0 ends holding precisely candle's result. Lanes with row >= HC feed 0.0f, exactly candle's `shr[tid] = 0` identity padding. ARC_NO_SINKHORN_WARP=1 restores the templated kernel. The arms are separable in a profile by kernel name (`arc_sinkhorn::warp_kernel` vs the anonymous-namespace `sinkhorn_normalize_f32_kernel`), so which arm ran is provable from nsys alone and a stale object file cannot fake it. No Rust changes; the C ABI is unchanged. Co-Authored-By: Claude Opus 5 (1M context) * instrument(ArcQuant/QTIP): count expert-weight bytes read per step Answers whether the MoE path amortises expert weights across a batch. Flat bytes/step with batch => amortisation works and the aggregate deficit is elsewhere; linear => every token drags its own copy of the expert weights. METHOD: kernel-issued bytes, computed exactly from each launch's own geometry. Not a profiler sample and not DRAM traffic. `ncu` cannot attach on the rented H200 ("Failed to prepare kernel for profiling / Unknown Error on device 0", verified on this box, not inherited), there is no DCGM, and `nvidia-smi dmon` reports a utilisation percentage rather than bytes -- so no hardware DRAM counter is available and the accounting is done from the dispatch instead. The gather GEMV (qtip_gather_gemv.cu) is indexed `pair = blockIdx.y` with `row-block = blockIdx.x`, and each block loads its own expert's packed rows. Summed over blockIdx.x one pair reads one WHOLE expert, and nothing is shared between pairs, so a launch reads n_pairs * n_rows * packed_per_row bytes. The grouped GEMM stages a woken expert once per GROUPED_TILE_M pairs, so it is accounted as ceil(pairs_e / TILE_M) copies. Both arms report in the same units so they are directly comparable. Because these are loads ISSUED they upper-bound DRAM traffic: repeated reads of one expert inside a short window can be served by L2. That is the point -- 99-100% SM occupancy against a 3% memory controller is what re-reading a redundant working set looks like, and the redundancy factor is the quantity that must reach 1.0 for batching to amortise, whichever cache serves it. ARC_MOE_BYTE_PROBE=1 to enable; off by default. One relaxed atomic add per expert-GEMM launch when on -- no device sync and no D2H read, so it is capture-safe and does not perturb the dispatch it measures. Co-Authored-By: Claude Opus 5 (1M context) * test(ArcQuant/QTIP): pin the measured expert-byte curve MEASURED on H200, V4-Flash qtip2b, ARC_MOE_BYTE_PROBE=1. One expert is 2 MiB; a step is 129 gather invocations (43 layers x 3). Expert bytes per step: n_tokens GEMV arm grouped arm redundancy 1 1.5 GiB 1.5 GiB 1.00x 8 12.1 GiB 12.1 GiB 1.00x 64 96.8 GiB 64.5 GiB 1.50x 512 774.0 GiB 64.5 GiB 12.00x The GEMV arm is exactly linear in the batch -- 512x the users reads 512x the expert bytes, so nothing is shared. The grouped arm is FLAT at the whole expert working set once every expert is woken. Two independent anchors agree: - a real 8-user decode through the server reads 12.09 GiB/step, identical to the prefill-derived n_tokens=8 figure, so the sweep measured the same quantity that batched serving pays; - the amortised floor of 64.5 GiB/step is 95.1% of the 67.82 GiB of qtip2b shards on disk, i.e. it is the whole expert working set read once. This is a characterization test, not an endorsement: it encodes the defect so that fixing the dispatch fails here first and the numbers get updated. Co-Authored-By: Claude Opus 5 (1M context) * test(ArcQuant/QTIP): pin the corrected dispatch, with the verified numbers The shipped default (no env override) now selects the flat-bytes arm above the capture-safety floor. Verified end-to-end on H200: B old tile-fill gate shipped default delta 8 57.18 60.28 +5.4% 64 106.26 156.61 +47.4% 256 114.76 193.46 +68.6% Kernel identity is from the profile, not the gate condition: at n_tokens=256 the default runs qtip2b_grouped_gemm_kernel n=129 -- exactly one step's worth of gathers (43 layers x 3) -- while n_tokens=1 decode steps stay on qtip2b_gemv_tuned_kernel below the floor, which is the intended split. Co-Authored-By: Claude Opus 5 (1M context) --------- Co-authored-by: Claude Opus 5 (1M context) --- mistralrs-core/src/cuda/sinkhorn.cu | 159 ++++++++++++++ mistralrs-quant/src/qtip/bitshift.rs | 9 + mistralrs-quant/src/qtip/byte_probe.rs | 273 +++++++++++++++++++++++++ mistralrs-quant/src/qtip/cuda_ops.rs | 6 + mistralrs-quant/src/qtip/mod.rs | 4 + 5 files changed, 451 insertions(+) create mode 100644 mistralrs-quant/src/qtip/byte_probe.rs diff --git a/mistralrs-core/src/cuda/sinkhorn.cu b/mistralrs-core/src/cuda/sinkhorn.cu index 767ac55a1..4deca694c 100644 --- a/mistralrs-core/src/cuda/sinkhorn.cu +++ b/mistralrs-core/src/cuda/sinkhorn.cu @@ -62,6 +62,7 @@ #include #include +#include #if defined(__USE_FAST_MATH__) #error "sinkhorn.cu must be compiled WITHOUT --use_fast_math: fast math rewrites expf to __expf and IEEE division to approximate reciprocals, breaking bit-identity with candle-kernels (which build with plain -O3). See the dedicated no-fast-math builder in mistralrs-core/build.rs." @@ -257,6 +258,138 @@ __global__ void sinkhorn_normalize_f32_kernel( } // namespace +// --------------------------------------------------------------------------- +// WARP-BUTTERFLY PATH (default) +// --------------------------------------------------------------------------- +// The templated kernel above got the row/tree buffers into registers, but it +// still stages the matrix through shared memory and pays TWO __syncthreads() +// per iteration -- 40 barriers and ~80 shared round trips per call -- to move +// each column to the thread that divides by it. With hc <= 16 every thread is +// a lane of ONE warp, so none of that is needed: the column reduction is a +// warp butterfly and the matrix never leaves registers. +// +// MEASURED, H200 / nvcc 12.4 / sm_90, n=1 hc=4 iters=20, 86 calls per token: +// legacy (pre-ladder, runtime hc, 192 B stack) 30.578 us/call +// templated + shared/barriers (the kernel above) 9.935 us/call +// warp butterfly (this kernel) 7.302 us/call 1.36x +// +// Bit-identity holds by construction. candle's level `s` is +// `buf[t] = __fadd_rn(buf[t], buf[t + s])` for t < s; the butterfly gives lane +// t `__fadd_rn(buf[t], buf[t ^ s])`, which IS that expression for t < s and its +// commuted form for t >= s. IEEE-754 addition is commutative (only +// associativity fails), so every lane ends each level holding bit-identical +// values and lane 0 holds precisely candle's result. Lanes with row >= HC feed +// 0.0f, which is exactly candle's `shr[tid] = 0` identity padding. +// +// KILL SWITCH: ARC_NO_SINKHORN_WARP=1 restores the templated shared-memory +// kernel above. Default is this path (house policy: fast-by-default with a kill +// switch, never opt-in). The two arms are distinguishable in a profile by +// kernel name -- `arc_sinkhorn::warp_kernel` vs the anonymous-namespace +// `sinkhorn_normalize_f32_kernel` -- so which arm ran is provable from nsys +// output alone and cannot be faked by a stale object file. + +namespace arc_sinkhorn { + +template +__device__ __forceinline__ void col_sums( + const float (&r)[HC], + float (&cs)[HC], + bool active, + unsigned mask, + float eps +) { +#pragma unroll + for (int c = 0; c < HC; ++c) { + float v = __fadd_rn(0.0f, active ? r[c] : 0.0f); +#pragma unroll + for (int s = P >> 1; s > 0; s >>= 1) { + v = __fadd_rn(v, __shfl_xor_sync(mask, v, s)); + } + cs[c] = __fadd_rn(v, eps); + } +} + +// in/out: [n, HC, HC] row-major F32. One block per matrix, next_pow2(HC) +// threads -- all lanes of a single warp, so no shared memory and no barriers. +template +__global__ void warp_kernel( + const float* __restrict__ in, + float* __restrict__ out, + int n, + int iters, + float eps +) { + const int batch = blockIdx.x; + const int row = threadIdx.x; + // Uniform across the block, so it cannot desynchronise the shuffles below. + if (batch >= n) return; + + const int P = next_pow2_ce(HC); + const bool active = (row < HC); + const unsigned mask = (P >= 32) ? 0xffffffffu : ((1u << P) - 1u); + + float r[HC]; + if (active) { + const float* my_in = in + (size_t)batch * HC * HC + (size_t)row * HC; +#pragma unroll + for (int j = 0; j < HC; ++j) r[j] = my_in[j]; + } else { +#pragma unroll + for (int j = 0; j < HC; ++j) r[j] = 0.0f; + } + + // ---- 1. stable row softmax, then + eps (entirely within this lane) ---- + const float m = candle_tree_max(r); +#pragma unroll + for (int j = 0; j < HC; ++j) r[j] = expf(__fsub_rn(r[j], m)); + const float rs = candle_tree_sum(r); +#pragma unroll + for (int j = 0; j < HC; ++j) r[j] = __fadd_rn(__fdiv_rn(r[j], rs), eps); + + // ---- 2. initial column normalize: x / (colsum + eps) ---- + float cs[HC]; + col_sums(r, cs, active, mask, eps); +#pragma unroll + for (int j = 0; j < HC; ++j) r[j] = __fdiv_rn(r[j], cs[j]); + + // ---- 3. (iters - 1) more row->col passes ---- + for (int it = 0; it < iters - 1; ++it) { + const float rsum = __fadd_rn(candle_tree_sum(r), eps); +#pragma unroll + for (int j = 0; j < HC; ++j) r[j] = __fdiv_rn(r[j], rsum); + col_sums(r, cs, active, mask, eps); +#pragma unroll + for (int j = 0; j < HC; ++j) r[j] = __fdiv_rn(r[j], cs[j]); + } + + if (active) { + float* my_out = out + (size_t)batch * HC * HC + (size_t)row * HC; +#pragma unroll + for (int j = 0; j < HC; ++j) my_out[j] = r[j]; + } +} + +// Read once; getenv is not cheap and this is on the per-layer decode path. +inline bool warp_disabled() { + static const bool disabled = [] { + const char* v = getenv("ARC_NO_SINKHORN_WARP"); + return v != nullptr && v[0] == '1' && v[1] == '\0'; + }(); + return disabled; +} + +} // namespace arc_sinkhorn + +// Defined outside the namespace because it is expanded in `extern "C"` scope +// below; `warp_kernel` has to be qualified there. +#define SINKHORN_WARP_LAUNCH(HC_) \ + case HC_: \ + arc_sinkhorn::warp_kernel \ + <<>>( \ + (const float*)in, (float*)out, n, iters, eps); \ + return; + + extern "C" { void sinkhorn_normalize_f32( @@ -268,6 +401,31 @@ void sinkhorn_normalize_f32( float eps, long long stream ) { + if (!arc_sinkhorn::warp_disabled()) { + switch (hc) { + SINKHORN_WARP_LAUNCH(1) + SINKHORN_WARP_LAUNCH(2) + SINKHORN_WARP_LAUNCH(3) + SINKHORN_WARP_LAUNCH(4) + SINKHORN_WARP_LAUNCH(5) + SINKHORN_WARP_LAUNCH(6) + SINKHORN_WARP_LAUNCH(7) + SINKHORN_WARP_LAUNCH(8) + SINKHORN_WARP_LAUNCH(9) + SINKHORN_WARP_LAUNCH(10) + SINKHORN_WARP_LAUNCH(11) + SINKHORN_WARP_LAUNCH(12) + SINKHORN_WARP_LAUNCH(13) + SINKHORN_WARP_LAUNCH(14) + SINKHORN_WARP_LAUNCH(15) + SINKHORN_WARP_LAUNCH(16) + default: + // Unreachable (the Rust guard rejects hc > SINKHORN_MAX_HC); fall + // through to the templated kernel rather than leave `out` untouched. + break; + } + } + dim3 grid(n, 1, 1); dim3 block(hc, 1, 1); switch (hc) { @@ -295,5 +453,6 @@ void sinkhorn_normalize_f32( } } #undef SINKHORN_LAUNCH +#undef SINKHORN_WARP_LAUNCH } // extern "C" diff --git a/mistralrs-quant/src/qtip/bitshift.rs b/mistralrs-quant/src/qtip/bitshift.rs index f6bca9575..308ee1428 100644 --- a/mistralrs-quant/src/qtip/bitshift.rs +++ b/mistralrs-quant/src/qtip/bitshift.rs @@ -1980,6 +1980,15 @@ impl QuantMethod for Qtip2bLayer { self.num_experts_count(), ); let total_pairs = n_tokens * n_experts_per_tok; + // Same accounting on the amortising arm, so the two paths + // are comparable in the same units (ARC_MOE_BYTE_PROBE=1). + super::byte_probe::record_grouped( + total_pairs, + self.rows_per_expert()?, + self.blocks.dim(2)?, + self.num_experts_count(), + super::grouped::GROUPED_TILE_M, + ); let out_flat = self.gather_forward_batched( &a.reshape((total_pairs, cols))?, &indices.reshape((total_pairs,))?, diff --git a/mistralrs-quant/src/qtip/byte_probe.rs b/mistralrs-quant/src/qtip/byte_probe.rs new file mode 100644 index 000000000..a1eee5196 --- /dev/null +++ b/mistralrs-quant/src/qtip/byte_probe.rs @@ -0,0 +1,273 @@ +//! Parent system: ArcQuant / QTIP โ€” expert-weight byte accounting. +//! +//! Answers one question: **how many bytes of expert weight does the MoE path +//! read per decode step, and does that number stay flat as the batch grows?** +//! +//! - Flat with batch โ‡’ amortisation works; a step's expert bytes are shared. +//! - Linear in batch โ‡’ every token drags its own copy of the expert weights, +//! and aggregate throughput cannot scale no matter how many users arrive. +//! +//! ## Method (state this whenever you quote the number) +//! +//! These are **kernel-issued bytes**, computed exactly from each launch's own +//! geometry โ€” not a profiler sample, not an estimate, and not DRAM traffic. +//! `ncu` cannot attach on the rented H200 ("Failed to prepare kernel for +//! profiling / Unknown Error on device 0") and the box has no DCGM, so a +//! hardware DRAM counter is not available; `nvidia-smi dmon` reports a +//! utilisation percentage, not bytes. +//! +//! The gather GEMV (`qtip2b_gemv_tuned_kernel`, `qtip_gather_gemv.cu`) is +//! indexed `pair = blockIdx.y`, `row-block = blockIdx.x`, and each block loads +//! its own expert's packed rows. Summed over `blockIdx.x`, **one pair reads one +//! whole expert**: `n_rows * packed_per_row` bytes. There is no dedup between +//! pairs, so a launch reads `n_pairs * n_rows * packed_per_row` bytes โ€” which is +//! what [`record_gemv`] accumulates. +//! +//! The grouped GEMM (`qtip_grouped_gemm.cu`) sorts pairs by expert and stages a +//! woken expert's bytes once per `GROUPED_TILE_M` rows, so it reads +//! `ceil(pairs_e / TILE_M)` copies of expert `e` โ€” [`record_grouped`] uses the +//! same per-expert accounting. +//! +//! Because these are loads *issued*, they are an upper bound on DRAM traffic: +//! repeated reads of the same expert within a short window can be served by L2. +//! That distinction is the point โ€” a card at 99% SM occupancy with a 3% memory +//! controller is exactly what re-reading a cached-but-redundant working set +//! looks like. The redundancy factor below is the quantity that has to reach +//! 1.0 for batching to amortise, whichever cache serves it. +//! +//! Enabled with `ARC_MOE_BYTE_PROBE=1`. Off by default and it costs one relaxed +//! atomic add per expert-GEMM launch when on โ€” no device sync, no D2H read, so +//! it is safe under CUDA-graph capture and does not disturb the dispatch it is +//! measuring. + +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::OnceLock; + +static INVOCATIONS: AtomicU64 = AtomicU64::new(0); +static GEMV_INVOCATIONS: AtomicU64 = AtomicU64::new(0); +static GROUPED_INVOCATIONS: AtomicU64 = AtomicU64::new(0); +static PAIRS: AtomicU64 = AtomicU64::new(0); +/// Bytes of expert weight the launches were pointed at. +static ISSUED_BYTES: AtomicU64 = AtomicU64::new(0); +/// Lower bound: each expert woken by a launch read exactly once. +static FLOOR_BYTES: AtomicU64 = AtomicU64::new(0); +/// Bytes of a single expert, for reporting the working-set size. +static BYTES_PER_EXPERT: AtomicU64 = AtomicU64::new(0); +static NUM_EXPERTS: AtomicU64 = AtomicU64::new(0); + +/// `ARC_MOE_BYTE_PROBE=1`. Read once โ€” `std::env::var` takes a global lock and +/// this sits on the per-layer decode path. +pub fn enabled() -> bool { + static ON: OnceLock = OnceLock::new(); + *ON.get_or_init(|| matches!(std::env::var("ARC_MOE_BYTE_PROBE").as_deref(), Ok("1"))) +} + +/// How many distinct experts a launch can possibly have woken. Exact whenever +/// every expert is woken (the saturated regime that matters for the fleet +/// question) and a conservative *over*-estimate otherwise, so the redundancy +/// factor this yields is never flattering. +#[inline] +fn distinct_upper_bound(n_pairs: usize, num_experts: usize) -> usize { + n_pairs.min(num_experts) +} + +/// The non-amortising gather GEMV: one pair reads one whole expert. +pub fn record_gemv(n_pairs: usize, n_rows: usize, packed_per_row: usize, num_experts: usize) { + if !enabled() { + return; + } + let per_expert = (n_rows as u64).saturating_mul(packed_per_row as u64); + let issued = per_expert.saturating_mul(n_pairs as u64); + let floor = per_expert.saturating_mul(distinct_upper_bound(n_pairs, num_experts) as u64); + INVOCATIONS.fetch_add(1, Ordering::Relaxed); + GEMV_INVOCATIONS.fetch_add(1, Ordering::Relaxed); + PAIRS.fetch_add(n_pairs as u64, Ordering::Relaxed); + ISSUED_BYTES.fetch_add(issued, Ordering::Relaxed); + FLOOR_BYTES.fetch_add(floor, Ordering::Relaxed); + BYTES_PER_EXPERT.store(per_expert, Ordering::Relaxed); + NUM_EXPERTS.store(num_experts as u64, Ordering::Relaxed); + maybe_dump(); +} + +/// The amortising grouped GEMM: a woken expert's bytes are staged once per +/// `tile_m` pairs that landed on it. Without the on-device sort order here, use +/// the even-spread estimate `ceil(pairs/distinct / tile_m)` copies per expert โ€” +/// exact when routing is balanced, which is the regime the gate targets. +pub fn record_grouped( + n_pairs: usize, + n_rows: usize, + packed_per_row: usize, + num_experts: usize, + tile_m: usize, +) { + if !enabled() { + return; + } + let per_expert = (n_rows as u64).saturating_mul(packed_per_row as u64); + let distinct = distinct_upper_bound(n_pairs, num_experts).max(1); + let pairs_each = n_pairs.div_ceil(distinct); + let copies = pairs_each.div_ceil(tile_m.max(1)); + let issued = per_expert + .saturating_mul(distinct as u64) + .saturating_mul(copies as u64); + let floor = per_expert.saturating_mul(distinct as u64); + INVOCATIONS.fetch_add(1, Ordering::Relaxed); + GROUPED_INVOCATIONS.fetch_add(1, Ordering::Relaxed); + PAIRS.fetch_add(n_pairs as u64, Ordering::Relaxed); + ISSUED_BYTES.fetch_add(issued, Ordering::Relaxed); + FLOOR_BYTES.fetch_add(floor, Ordering::Relaxed); + BYTES_PER_EXPERT.store(per_expert, Ordering::Relaxed); + NUM_EXPERTS.store(num_experts as u64, Ordering::Relaxed); + maybe_dump(); +} + +/// Periodic cumulative dump so a run yields a number even if it is killed. +/// `ARC_MOE_BYTE_PROBE_EVERY` invocations, default 4096. +fn maybe_dump() { + static EVERY: OnceLock = OnceLock::new(); + let every = *EVERY.get_or_init(|| { + std::env::var("ARC_MOE_BYTE_PROBE_EVERY") + .ok() + .and_then(|v| v.parse::().ok()) + .filter(|v| *v > 0) + .unwrap_or(4096) + }); + let n = INVOCATIONS.load(Ordering::Relaxed); + if n.is_multiple_of(every) { + eprintln!("{}", report()); + } +} + +/// One machine-greppable cumulative line. Divide by the tokens the run actually +/// produced to get bytes/token; multiply by the batch to get bytes/step. +pub fn report() -> String { + let inv = INVOCATIONS.load(Ordering::Relaxed); + let gemv = GEMV_INVOCATIONS.load(Ordering::Relaxed); + let grouped = GROUPED_INVOCATIONS.load(Ordering::Relaxed); + let pairs = PAIRS.load(Ordering::Relaxed); + let issued = ISSUED_BYTES.load(Ordering::Relaxed); + let floor = FLOOR_BYTES.load(Ordering::Relaxed); + let per_expert = BYTES_PER_EXPERT.load(Ordering::Relaxed); + let experts = NUM_EXPERTS.load(Ordering::Relaxed); + let redundancy = if floor > 0 { + issued as f64 / floor as f64 + } else { + 0.0 + }; + format!( + "ARC_MOE_BYTES invocations={inv} gemv={gemv} grouped={grouped} pairs={pairs} \ + issued_bytes={issued} floor_bytes={floor} redundancy={redundancy:.3} \ + bytes_per_expert={per_expert} num_experts={experts}" + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The GEMV reads one whole expert per pair, so its redundancy is exactly + /// pairs/distinct โ€” the quantity that has to fall to 1.0 for batching to + /// amortise. These are the shapes V4 actually runs (top-6 of 256). + #[test] + fn gemv_redundancy_is_pairs_per_distinct_expert() { + let per_expert = 1024usize * 96; + for (n_tokens, top_k, experts) in [(1usize, 6usize, 256usize), (512, 6, 256)] { + let pairs = n_tokens * top_k; + let distinct = distinct_upper_bound(pairs, experts); + let issued = per_expert * pairs; + let floor = per_expert * distinct; + assert_eq!(issued / floor.max(1), pairs / distinct.max(1)); + } + // B=1 wakes 6 experts with 6 pairs: nothing to share, redundancy 1. + assert_eq!(distinct_upper_bound(6, 256), 6); + // B=512 wakes all 256 with 3072 pairs: 12x redundant. + assert_eq!(3072 / distinct_upper_bound(3072, 256), 12); + } + + /// MEASURED on H200, V4-Flash qtip2b (ARC_MOE_BYTE_PROBE=1): one expert is + /// 2 MiB and a step is 129 gather invocations (43 layers x 3). Expert bytes + /// per step scale LINEARLY with the batch on the GEMV arm and are FLAT on + /// the grouped arm -- which is why the dispatch gate was changed to select + /// grouped above the capture-safety floor. This pins the arithmetic that + /// motivated the fix; `default_gate_selects_the_flat_arm_above_the_floor` + /// below pins the dispatch that resulted from it. + #[test] + fn v4_expert_bytes_per_step_scale_with_batch_on_the_gemv_arm() { + const PER_EXPERT: u64 = 2 * 1024 * 1024; // measured + const INV_PER_STEP: u64 = 129; // measured: 43 layers x 3 gathers + const TOP_K: usize = 6; + const EXPERTS: usize = 256; + let gemv_gib = |n_tokens: usize| -> f64 { + let pairs = (n_tokens * TOP_K) as u64; + (INV_PER_STEP * pairs * PER_EXPERT) as f64 / (1024.0 * 1024.0 * 1024.0) + }; + let floor_gib = |n_tokens: usize| -> f64 { + let distinct = distinct_upper_bound(n_tokens * TOP_K, EXPERTS) as u64; + (INV_PER_STEP * distinct * PER_EXPERT) as f64 / (1024.0 * 1024.0 * 1024.0) + }; + // GEMV: linear in batch, 512x the users reads 512x the expert bytes. + for (n, want) in [(1usize, 1.5f64), (8, 12.1), (64, 96.8), (512, 774.0)] { + assert!( + (gemv_gib(n) - want).abs() < 0.1, + "n_tokens={n}: {:.1} GiB/step, expected {want}", + gemv_gib(n) + ); + } + // Amortised floor saturates at the whole expert working set. + assert!((floor_gib(256) - 64.5).abs() < 0.1); + assert!((floor_gib(512) - 64.5).abs() < 0.1); + // The redundancy the fleet pays at B=512. + assert!((gemv_gib(512) / floor_gib(512) - 12.0).abs() < 0.01); + } + + /// The fix, as an assertion. The gate now selects the arm whose expert + /// bytes are FLAT in the batch, for every width above the capture-safety + /// floor. MEASURED aggregate at the shipped default: B=8 60.28, B=64 + /// 156.61, B=256 193.46 tok/s, against 57.18 / 106.26 / 114.76 under the + /// old tile-fill gate. + /// + /// This branch originally asserted against its own replacement predicate, + /// `gather_policy::grouped_gemm_preferred` โ€” a byte-traffic model gated on + /// `GROUPED_MIN_PAIRS_PER_EXPERT = 1.09`. That predicate is not what ships: + /// #197 arrived first and retired the tile-fill gate outright rather than + /// loosening it, leaving [`prefer_gather_gemv`] โ€” the GEMV arm below + /// [`DECODE_REGIME_MAX_TOKENS`], grouped above it โ€” as the whole dispatch. + /// The two agree exactly at V4's top-6-of-256 (1.09 pairs per woken expert + /// flips at n = 9; the floor is n > 8), so the boundary this test was + /// written to pin is unchanged; only the predicate naming it is. Asserting + /// against the shipped one is the point of the test. + #[test] + fn default_gate_selects_the_flat_arm_above_the_floor() { + use crate::qtip::gather_policy::prefer_gather_gemv; + // At and below the floor the GEMV is mandatory (CUDA-graph capture + // safety) and measured a wash, so nothing is lost. + assert!(prefer_gather_gemv(8, None)); + // Above it, every measured width favours the flat arm. + for n in [9usize, 12, 16, 24, 32, 64, 128, 256, 512] { + assert!( + !prefer_gather_gemv(n, None), + "n={n}: expert bytes on the GEMV arm are {}x the floor here", + (n * 6) / distinct_upper_bound(n * 6, 256).max(1) + ); + } + } + /// The grouped GEMM stages an expert once per TILE_M pairs, so at the same + /// shape it must read strictly fewer bytes than the GEMV once the tiles + /// fill โ€” that is the whole claim the dispatch gate rests on. + #[test] + fn grouped_reads_less_than_gemv_once_tiles_fill() { + let (n_rows, ppr, experts, tile) = (1024usize, 96usize, 256usize, 16usize); + let per_expert = (n_rows * ppr) as u64; + for n_tokens in [512usize, 683, 2048] { + let pairs = n_tokens * 6; + let distinct = distinct_upper_bound(pairs, experts); + let gemv = per_expert * pairs as u64; + let copies = pairs.div_ceil(distinct).div_ceil(tile); + let grouped = per_expert * distinct as u64 * copies as u64; + assert!( + grouped < gemv, + "grouped {grouped} !< gemv {gemv} at n_tokens={n_tokens}" + ); + } + } +} diff --git a/mistralrs-quant/src/qtip/cuda_ops.rs b/mistralrs-quant/src/qtip/cuda_ops.rs index 6a329d952..049b9e8eb 100644 --- a/mistralrs-quant/src/qtip/cuda_ops.rs +++ b/mistralrs-quant/src/qtip/cuda_ops.rs @@ -1973,6 +1973,12 @@ pub(crate) fn gather_gemv_2b_cuda( // One pair per `grid.y` index; past `maxGridSize[1]` the launch fails // silently and the zero-filled output buffer is returned as if valid. super::gather_policy::check_gather_gemv_pairs(n_pairs, "qtip2b gather gemv CUDA")?; + // Expert-weight byte accounting (ARC_MOE_BYTE_PROBE=1). This kernel is + // indexed `pair = blockIdx.y` and each block loads its own expert's packed + // rows, so summed over `blockIdx.x` one pair reads one whole expert and a + // launch reads `n_pairs * n_rows * packed_per_row` bytes with no sharing + // between pairs. Relaxed atomics only: no device sync, capture-safe. + super::byte_probe::record_gemv(n_pairs, n_rows, packed_per_row, num_experts); let x_2d = x_rotated.contiguous()?; let indices = indices.flatten_all()?.contiguous()?; diff --git a/mistralrs-quant/src/qtip/mod.rs b/mistralrs-quant/src/qtip/mod.rs index f4fd425b2..9f96d72fe 100644 --- a/mistralrs-quant/src/qtip/mod.rs +++ b/mistralrs-quant/src/qtip/mod.rs @@ -61,6 +61,10 @@ mod bake_memory_tests; #[cfg(test)] mod bake_quality_tests; pub mod bitshift; +// Every call site is behind `#[cfg(feature = "cuda")]`, so a CPU-only build +// sees the whole probe as dead. Same convention as `gather_policy`. +#[cfg_attr(not(feature = "cuda"), allow(dead_code))] +pub mod byte_probe; #[cfg(test)] mod codebook_tests; #[cfg(feature = "cuda")] From 55160369cf0aac8f334c29ab06ca6f163c49fa47 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 20:10:15 +0100 Subject: [PATCH 14/21] fix(ArcGraph): correct cuGraphAddNode ABI; make GPU-autonomous decode reachable (#130) The headline is an ABI mismatch that explains a failure already on our record. `ffi.rs` declared `cudaGraphAddNode` with FIVE arguments and the first two transposed. CUDA 13.1 `cuda.h:21829` declares SIX, and the first parameter is an OUT pointer: CUresult cuGraphAddNode(CUgraphNode *phGraphNode, CUgraph hGraph, const CUgraphNode *dependencies, const CUgraphEdgeData *dependencyData, size_t numDependencies, CUgraphNodeParams *nodeParams); So the driver received a `CUgraph` handle where it expected `CUgraphNode *` and wrote the new node handle THROUGH it, while `numDependencies` received a pointer and `nodeParams` an uninitialised register. That is memory corruption from a declaration, and it matches the previously unexplained "host heap corruption" CUDA-graph capture failure we had recorded with no cause. Four more defects, each of which alone stops the autonomous loop: * Conditional handles were created with `flags = 0`. `cuda.h:21919` applies `defaultLaunchValue` only when `CU_GRAPH_COND_ASSIGN_DEFAULT` is set, so the condition read 0 at launch and a WHILE body executed ZERO times while every CUDA call returned success. Measured both ways on an H200. * The WHILE body was captured into a throwaway graph and then DESTROYED, leaving the conditional body empty. Bodies are now populated with `cuStreamBeginCaptureToGraph` (`cuda.h:1976` names it as the supported way), and a node-count assert refuses a body that recorded nothing. * `CUDA_CONDITIONAL_NODE_PARAMS` was missing its 5th field, `ctx`. * `DecodeState` allocated I64 while every kernel signature is `int32_t*`, so host and device disagreed about where element `i` lived; and `reset()` reallocated the buffers, silently invalidating the pointers a captured graph had baked in. `tensor_device_ptr` had no I32 arm, so `CudaSampler::sample` returned `unsupported dtype I32` on EVERY call -- the correct top-k/top-p sampler had never once executed on a GPU. Its test suite is a CPU simulator, which by construction cannot observe that. Sampling now runs on device and survives replay. The sampler previously wired into the graph body took `rng_offset` as a BY-VALUE kernel argument, which a captured graph bakes: measured 1 distinct token over 64 replays, versus 64 distinct for a device-resident `rng_state`. It also was not nucleus sampling -- it walked the vocabulary in TOKEN-ID order. The replacement is a hybrid, because measurement showed neither component wins everywhere (H200, vocab 129280, host+GPU verified exclusive before and after; support width -> legacy / hybrid us): 1 -> 134.9/166.6 8 -> 503.0/624.2 64 -> 3127.8/4178.5 512 -> 24320.7/4185.0 4096 -> 205125.2/4187.3 12928 -> 664200.0/4188.5 Exact enumeration while the nucleus is small, threshold bisection past a fixed budget, chosen block-uniformly ON DEVICE because a captured graph cannot ask the host which branch to take. Costs +24% on the peaked distributions real models produce -- 32 us, 0.05% of a 66.68 ms V4 decode step -- and removes a 664 ms cliff that was ~10x an entire decode step for a single token. No cap and no narrowing: exceeding the budget changes which ALGORITHM selects the nucleus, never which tokens it contains. Verified against the enumerating sampler with branch counters proving which path each case exercised: peaked TV=0.0000 A=20000 B=0 scrambled order TV=0.0000 A=20000 B=0 (15.6% per-draw agreement, so the orders genuinely diverged) diffuse TV=0.0000 A=0 B=3000 tie boundary TV=0.0000 A=20000 B=0 narrowed (ctrl) TV=0.5027 -> correctly DISAGREE Known and deliberately left: on TIED diffuse supports the fallback keeps 116 tokens where enumeration keeps 111 (TV=0.1103) -- it keeps every token sharing a boundary bit pattern, erring WIDER, never narrower. And the budget is derived from bisection's pass count (32) when it should follow the FALLBACK's (~70), which costs a 0.75x window at support ~64; correcting it needs a more diffuse fixture so the fallback stays covered. Not claimed: no V4 end-to-end number. Autonomous decode is unreachable on V4 by construction -- no PagedAttention means `cache_config` is None and the runner is never built (`normal.rs:1907`, pinned by `normal_loaders.rs:5687`) -- and `cuda.h:1971` bars alloc/free nodes inside a conditional body against 11,436 allocations per token. Co-authored-by: Claude (ArcGraph chain) --- arc-cuda-graph/src/autonomous.rs | 244 +++++++--- arc-cuda-graph/src/buffers.rs | 106 +++- arc-cuda-graph/src/cuda/sampling_kernel.cu | 456 ++++++++++++++++++ arc-cuda-graph/src/ffi.rs | 68 ++- arc-cuda-graph/src/sampling_cuda.rs | 68 ++- arc-cuda-graph/src/weights.rs | 10 + mistralrs-core/src/pipeline/normal.rs | 10 + .../tests/capability_reachability.rs | 25 +- mistralrs-core/tests/doc_citations.rs | 24 + 9 files changed, 902 insertions(+), 109 deletions(-) diff --git a/arc-cuda-graph/src/autonomous.rs b/arc-cuda-graph/src/autonomous.rs index 2539b2801..6153d7699 100644 --- a/arc-cuda-graph/src/autonomous.rs +++ b/arc-cuda-graph/src/autonomous.rs @@ -78,6 +78,8 @@ pub struct AutonomousDecodeConfig { pub top_p: f32, pub frequency_penalty: f32, pub presence_penalty: f32, + /// `-1` disables top-k. The previous in-graph sampler had no top-k at all. + pub top_k: i32, pub greedy: bool, } @@ -101,7 +103,19 @@ pub struct AutonomousDecodeRunner { /// True if the graph uses a WHILE conditional node (CUDA 12.4+). /// False if host-driven loop (body graph launched per step). uses_while_node: bool, - rng_offset: u64, + /// The real sampler: temperature + top-k + top-p, with its RNG state held + /// in **device memory** and advanced by the kernel. + /// + /// ๐Ÿ”ด The sampler this replaces (`launch_fused_top_p_bf16`) took + /// `rng_offset` as a **by-value kernel argument**. A captured graph bakes + /// kernel arguments at capture time, so every replayed step drew the + /// identical uniform and the host-side `rng_offset += 1` could not reach + /// an instantiated graph. It was also not nucleus sampling: it walked the + /// vocabulary in **token-id order** accumulating full-distribution mass to + /// `top_p * u`, biasing hard toward low token ids (its own comment said + /// "Not quite right โ€” proper top-p needs sorting"). A device-resident + /// `rng_state` is the only form that survives replay. + sampler: crate::sampling_cuda::CudaSampler, } #[cfg(feature = "cuda")] @@ -171,6 +185,17 @@ impl AutonomousDecodeRunner { std::ptr::write_bytes(ring_write_head_ptr as *mut u8, 0, batch * 4); } + // Bound to the runner's own stream so its kernels land inside the + // captured body rather than on the device-default stream. + let mut sampler = crate::sampling_cuda::CudaSampler::new( + device, + config.padded_batch_size, + config.vocab_size, + candle_core::DType::BF16, + 0x5EED_A11C_E571_2345, + )?; + sampler.set_stream(stream); + Ok(Self { config, device: device.clone(), @@ -183,7 +208,7 @@ impl AutonomousDecodeRunner { ring_size, graph_exec: None, uses_while_node: false, - rng_offset: 0, + sampler, }) } @@ -350,10 +375,30 @@ impl AutonomousDecodeRunner { candle_core::bail!("cuGraphCreate failed: {status}"); } - // 2. Create conditional handle (WHILE, default_value=1 = loop) + // 2. Create conditional handle (WHILE, default_value=1 = loop). + // + // `CU_GRAPH_COND_ASSIGN_DEFAULT` is REQUIRED: without it the default + // is never applied at launch, the condition reads 0, and the WHILE + // body runs zero times while every CUDA call still returns success. + let mut ctx: CUcontext = std::ptr::null_mut(); + let status = unsafe { cuCtxGetCurrent(&mut ctx) }; + if status != CUDA_SUCCESS || ctx.is_null() { + unsafe { + cuGraphDestroy(outer_graph); + } + tracing::warn!("cuCtxGetCurrent failed ({status}), falling back to host-driven loop"); + return self.capture_body_graph(forward_fn, bs, vocab); + } let mut cond_handle: CUgraphConditionalHandle = 0; - let status = - unsafe { cudaGraphConditionalHandleCreate(&mut cond_handle, outer_graph, 1, 0) }; + let status = unsafe { + cuGraphConditionalHandleCreate( + &mut cond_handle, + outer_graph, + ctx, + 1, + CU_GRAPH_COND_ASSIGN_DEFAULT, + ) + }; if status != CUDA_SUCCESS { unsafe { cuGraphDestroy(outer_graph); @@ -365,18 +410,19 @@ impl AutonomousDecodeRunner { // 3. Add conditional WHILE node to outer graph // This creates an empty body graph that we populate via stream capture. let mut while_node: CUgraphNode = std::ptr::null_mut(); - let mut body_graph: CUgraph = std::ptr::null_mut(); + // `phGraph_out` is an OUT field the driver populates; leave it null. let mut params: CudaGraphNodeParams = unsafe { std::mem::zeroed() }; params.node_type = CudaGraphNodeType::Conditional; params.conditional.handle = cond_handle; params.conditional.cond_type = CUgraphConditionalNodeType::WHILE; params.conditional.size = 1; - params.conditional.body_graph_out = &mut body_graph; + params.conditional.ctx = ctx; let status = unsafe { - cudaGraphAddNode( - outer_graph, + cuGraphAddNode( &mut while_node, + outer_graph, + std::ptr::null(), std::ptr::null(), 0, &mut params, @@ -390,19 +436,51 @@ impl AutonomousDecodeRunner { return self.capture_body_graph(forward_fn, bs, vocab); } - // 4. Populate body graph via stream capture + // 4. Populate the conditional node's OWN body graph. + // + // The driver hands back the body graph in `phGraph_out`; capture + // straight into it with `cuStreamBeginCaptureToGraph`. The previous + // code captured with `cuStreamBeginCapture_v2` โ€” which always creates + // a NEW graph โ€” and then destroyed the result, leaving the conditional + // body EMPTY. An empty body generates no tokens. + let body_graph: CUgraph = unsafe { + if params.conditional.body_graph_out.is_null() { + cuGraphDestroy(outer_graph); + candle_core::bail!("conditional node returned no body graph array"); + } + *params.conditional.body_graph_out + }; + if body_graph.is_null() { + unsafe { + cuGraphDestroy(outer_graph); + } + candle_core::bail!("conditional node body graph is null"); + } + unsafe { - let status = cuStreamBeginCapture_v2(self.stream, CUstreamCaptureMode::THREAD_LOCAL); + // RELAXED, not THREAD_LOCAL: candle's allocator and helper streams + // create cross-stream dependencies that THREAD_LOCAL rejects with + // CUDA_ERROR_STREAM_CAPTURE_ISOLATION (same reason `graph.rs` + // uses RELAXED). + let status = cuStreamBeginCaptureToGraph( + self.stream, + body_graph, + std::ptr::null(), + std::ptr::null(), + 0, + CUstreamCaptureMode::RELAXED, + ); if status != CUDA_SUCCESS { cuGraphDestroy(outer_graph); - candle_core::bail!("cuStreamBeginCapture for WHILE body failed: {status}"); + candle_core::bail!("cuStreamBeginCaptureToGraph for WHILE body failed: {status}"); } } // Capture: forward โ†’ sample โ†’ step_update โ†’ check_done_conditional self.capture_body_kernels(forward_fn, bs, vocab, Some(cond_handle))?; - // End capture into the body graph + // Ends the capture; returns the graph we captured INTO (`body_graph`), + // which is already attached to the conditional node. let mut captured_body: CUgraph = std::ptr::null_mut(); unsafe { let status = cuStreamEndCapture(self.stream, &mut captured_body); @@ -411,29 +489,37 @@ impl AutonomousDecodeRunner { candle_core::bail!("cuStreamEndCapture for WHILE body failed: {status}"); } } - - // The captured_body needs to be merged into the body_graph that - // cudaGraphAddNode created. For CUDA 12.4 conditional nodes, - // the body_graph_out is pre-created and we should have captured - // INTO it by using it as the capture target. However, stream - // capture always creates a new graph. - // - // The correct approach: don't use stream capture for the body. - // Instead, add kernel nodes to body_graph directly. But that's - // extremely complex (need to manually create kernel nodes for - // every cuBLAS call, every custom kernel, etc.). - // - // Alternative: use cudaStreamBeginCaptureToGraph (CUDA 12.3+) - // which captures into an existing graph. - // - // For now: instantiate the outer graph. If the body graph was - // properly populated by the conditional node setup, it works. - // If not, we destroy and fall back. - unsafe { - cuGraphDestroy(captured_body); + if captured_body != body_graph { + unsafe { + cuGraphDestroy(outer_graph); + } + candle_core::bail!( + "capture-to-graph returned a different graph than the conditional body" + ); + } + // Assert the body actually received nodes. A conditional node with an + // empty body is the failure this function exists to prevent, and it is + // invisible at launch: the graph runs and produces nothing. + let mut body_nodes: usize = 0; + let status = unsafe { cuGraphGetNodes(body_graph, std::ptr::null_mut(), &mut body_nodes) }; + if status != CUDA_SUCCESS || body_nodes == 0 { + unsafe { + cuGraphDestroy(outer_graph); + } + candle_core::bail!( + "WHILE body graph has {body_nodes} nodes after capture (status {status}) โ€” \ + the decode body did not record" + ); } + tracing::info!("CUDA graph: WHILE body captured with {body_nodes} nodes"); - // 5. Instantiate outer graph + // 5. Instantiate outer graph. + // + // NOTE: no AUTO_FREE_ON_LAUNCH here. `cuda.h:1971` restricts a + // conditional body to "kernel nodes, empty nodes, child graphs, + // memsets, memcopies, and conditionals" โ€” memory alloc/free nodes are + // NOT in that list, so the captured decode body must be + // allocation-free for this path to instantiate at all. let mut exec: CUgraphExec = std::ptr::null_mut(); let status = unsafe { cuGraphInstantiate_v2( @@ -510,7 +596,7 @@ impl AutonomousDecodeRunner { /// Capture the body kernels: forward โ†’ sample โ†’ step_update โ†’ check_done. /// Called during stream capture (kernels recorded, not executed). fn capture_body_kernels( - &self, + &mut self, forward_fn: &F, bs: i32, vocab: i32, @@ -521,48 +607,54 @@ impl AutonomousDecodeRunner { { // Forward pass let logits = forward_fn()?; - let logits_ptr = tensor_ptr(&logits)? as *const _; - // Sampling - let sampled_ptr = tensor_ptr(&self.decode_state.sampled_tokens)? as *mut _; - - unsafe { - if self.config.greedy { - launch_fused_argmax_bf16( - logits_ptr, - sampled_ptr, - std::ptr::null_mut(), - vocab, - bs, - self.stream, - ); - } else { - if self.config.frequency_penalty != 0.0 || self.config.presence_penalty != 0.0 { - launch_apply_penalties( - logits_ptr as *mut _, - tensor_ptr(&self.decode_state.output_tokens)? as *const i32, - tensor_ptr(&self.decode_state.n_generated)? as *const i32, - self.config.frequency_penalty, - self.config.presence_penalty, - vocab, - self.config.max_tokens as i32, - bs, - self.stream, - ); - } - launch_fused_top_p_bf16( - logits_ptr, - sampled_ptr, - self.config.temperature, - self.config.top_p, + // Cheap handle clones so every immutable borrow of `self` is finished + // before we take `&mut self.sampler`. + let sampled = self.decode_state.sampled_tokens.clone(); + + // Penalties stay a separate in-place pre-pass over the logits. The + // fused sampler applies penalties from a [batch, vocab] count tensor + // (`freq_counts`), which we do not maintain; passing `None` there + // while leaving non-zero penalties in the config would drop them + // SILENTLY (`sampling_kernel.cu:167` โ€” `fcnt = freq_counts ? .. : + // nullptr`). So apply them here and zero them in `cfg`. + if self.config.frequency_penalty != 0.0 || self.config.presence_penalty != 0.0 { + let logits_mut = tensor_ptr(&logits)? as *mut _; + let out_toks = tensor_ptr(&self.decode_state.output_tokens)? as *const i32; + let n_gen = tensor_ptr(&self.decode_state.n_generated)? as *const i32; + unsafe { + launch_apply_penalties( + logits_mut, + out_toks, + n_gen, + self.config.frequency_penalty, + self.config.presence_penalty, vocab, + self.config.max_tokens as i32, bs, - 42, - self.rng_offset, self.stream, ); } + } + + let cfg = crate::sampling_cpu::SamplingConfig { + temperature: self.config.temperature, + top_p: self.config.top_p, + top_k: self.config.top_k, + // applied above, in-place, so the kernel must not re-apply them + frequency_penalty: 0.0, + presence_penalty: 0.0, + greedy: self.config.greedy, + eos_token_id: self.config.eos_token_id, + }; + // Scoped so the mutable borrow ends before the immutable ones below. + { + self.sampler.sample(&logits, None, cfg, &sampled)?; + } + let sampled_ptr = tensor_ptr(&sampled)? as *mut i32; + + unsafe { // Step update launch_decode_step_update( sampled_ptr as *const i32, @@ -621,9 +713,8 @@ impl AutonomousDecodeRunner { candle_core::Error::Msg("Graph not captured โ€” call capture() first".into()) })?; - // Reset state - self.decode_state - .reset(&self.device, self.config.padded_batch_size)?; + // Reset state in place โ€” the captured graph holds these addresses. + self.decode_state.reset(self.stream)?; unsafe { std::ptr::write_bytes( self.ring_write_head_ptr as *mut u8, @@ -631,7 +722,6 @@ impl AutonomousDecodeRunner { self.config.padded_batch_size * 4, ); } - self.rng_offset += 1; if self.uses_while_node { // ============================================================ @@ -656,7 +746,9 @@ impl AutonomousDecodeRunner { } cudaStreamSynchronize(self.stream); } - let cond = self.decode_state.loop_condition.to_vec1::()?; + // I32: the kernel writes `int32_t`. Reading this as i64 + // consumed two adjacent kernel writes as one value. + let cond = self.decode_state.loop_condition.to_vec1::()?; if cond[0] == 0 { break; } @@ -668,7 +760,7 @@ impl AutonomousDecodeRunner { .decode_state .output_tokens .to_dtype(candle_core::DType::I64)?; - let n_gen = self.decode_state.n_generated.to_vec1::()?; + let n_gen = self.decode_state.n_generated.to_vec1::()?; let mut results = Vec::new(); for b in 0..self.config.padded_batch_size { let n = n_gen[b] as usize; diff --git a/arc-cuda-graph/src/buffers.rs b/arc-cuda-graph/src/buffers.rs index 6b8d00f0f..df3ec623a 100644 --- a/arc-cuda-graph/src/buffers.rs +++ b/arc-cuda-graph/src/buffers.rs @@ -1,8 +1,45 @@ //! Pre-allocated GPU buffers for the decode loop. +//! +//! **Every buffer here is read and written by CUDA kernels whose signatures +//! are `int32_t*`** (`cuda/decode_loop.cu`, `cuda/sampling.cu`). The dtype of +//! the Candle tensor therefore has to be 4 bytes wide, or the host and the +//! device disagree about where element `i` lives. #[cfg(feature = "cuda")] use candle_core::{DType, Device, Tensor}; +#[cfg(feature = "cuda")] +use candle_core::cuda::cudarc::driver::sys::CUstream; + +#[cfg(feature = "cuda")] +extern "C" { + fn cudaMemsetAsync( + dst: *mut std::ffi::c_void, + value: i32, + count: usize, + stream: CUstream, + ) -> u32; + fn cudaMemcpyAsync( + dst: *mut std::ffi::c_void, + src: *const std::ffi::c_void, + count: usize, + kind: u32, + stream: CUstream, + ) -> u32; +} + +/// Allocate a zeroed I32 tensor. `Tensor::zeros` on I32 is not supported on +/// every Candle backend (see the same note in `sampling_cuda.rs:296`), so go +/// through `from_vec`. +#[cfg(feature = "cuda")] +fn zeros_i32( + elems: usize, + shape: impl Into, + device: &Device, +) -> candle_core::Result { + Tensor::from_vec(vec![0i32; elems], shape, device) +} + /// Pre-allocated decode input buffers at fixed GPU addresses. #[cfg(feature = "cuda")] pub struct DecodeInputBuffers { @@ -31,33 +68,74 @@ impl DecodeInputBuffers { } /// GPU-side decode state that persists across WHILE loop iterations. +/// +/// ๐Ÿ”ด These were all `I64` while every kernel that touches them declares +/// `int32_t*` (`decode_loop.cu:19-37`, `sampling.cu:21`). The kernel wrote +/// element `b` at byte `4*b`; Candle read element `b` at byte `8*b`. Nothing +/// lined up: `n_generated`, `finished`, `loop_condition` and the whole +/// `output_tokens` matrix were read back as garbage, and `output_tokens` was +/// additionally indexed with the wrong row stride. The path is gated off +/// before it runs, so nothing ever reported it. #[cfg(feature = "cuda")] pub struct DecodeState { - pub sampled_tokens: Tensor, // [padded_bs] i64 - pub n_generated: Tensor, // [padded_bs] i64 - pub output_tokens: Tensor, // [padded_bs, max_tokens] i64 - pub finished: Tensor, // [padded_bs] i64 - pub loop_condition: Tensor, // [1] i64 + pub sampled_tokens: Tensor, // [padded_bs] i32 + pub n_generated: Tensor, // [padded_bs] i32 + pub output_tokens: Tensor, // [padded_bs, max_tokens] i32 + pub finished: Tensor, // [padded_bs] i32 + pub loop_condition: Tensor, // [1] i32 pub max_tokens: usize, } #[cfg(feature = "cuda")] impl DecodeState { pub fn new(padded_bs: usize, max_tokens: usize, device: &Device) -> candle_core::Result { + let _ = DType::I32; // keep the dtype named in this module Ok(Self { - sampled_tokens: Tensor::zeros(padded_bs, DType::I64, device)?, - n_generated: Tensor::zeros(padded_bs, DType::I64, device)?, - output_tokens: Tensor::zeros((padded_bs, max_tokens), DType::I64, device)?, - finished: Tensor::zeros(padded_bs, DType::I64, device)?, - loop_condition: Tensor::ones(1, DType::I64, device)?, + sampled_tokens: zeros_i32(padded_bs, padded_bs, device)?, + n_generated: zeros_i32(padded_bs, padded_bs, device)?, + output_tokens: zeros_i32(padded_bs * max_tokens, (padded_bs, max_tokens), device)?, + finished: zeros_i32(padded_bs, padded_bs, device)?, + loop_condition: Tensor::from_vec(vec![1i32], 1, device)?, max_tokens, }) } - pub fn reset(&mut self, device: &Device, padded_bs: usize) -> candle_core::Result<()> { - self.n_generated = Tensor::zeros(padded_bs, DType::I64, device)?; - self.finished = Tensor::zeros(padded_bs, DType::I64, device)?; - self.loop_condition = Tensor::ones(1, DType::I64, device)?; + /// Reset the per-generation state **in place**, on `stream`. + /// + /// ๐Ÿ”ด This used to re-`Tensor::zeros` each field, which allocates NEW + /// device buffers at NEW addresses. The captured CUDA graph has the OLD + /// addresses baked into its kernel nodes, so resetting between + /// generations pointed the graph at freed memory โ€” the reset silently + /// un-did the capture. A pointer-stable buffer is a hard requirement of + /// anything that gets captured, so zero the existing allocations instead. + pub fn reset(&mut self, stream: CUstream) -> candle_core::Result<()> { + let n_gen = crate::weights::tensor_device_ptr(&self.n_generated)?; + let fin = crate::weights::tensor_device_ptr(&self.finished)?; + let cond = crate::weights::tensor_device_ptr(&self.loop_condition)?; + let bs = self.n_generated.elem_count(); + unsafe { + let s = cudaMemsetAsync(n_gen as *mut _, 0, bs * 4, stream); + if s != 0 { + candle_core::bail!("cudaMemsetAsync(n_generated) failed: {s}"); + } + let s = cudaMemsetAsync(fin as *mut _, 0, bs * 4, stream); + if s != 0 { + candle_core::bail!("cudaMemsetAsync(finished) failed: {s}"); + } + // loop_condition starts at 1 ("keep going"). memset writes a byte + // pattern, so 0x01010101 would be wrong; write the word directly. + let one: i32 = 1; + let s = cudaMemcpyAsync( + cond as *mut _, + &one as *const i32 as *const _, + 4, + 1, // cudaMemcpyHostToDevice + stream, + ); + if s != 0 { + candle_core::bail!("cudaMemcpyAsync(loop_condition) failed: {s}"); + } + } Ok(()) } } diff --git a/arc-cuda-graph/src/cuda/sampling_kernel.cu b/arc-cuda-graph/src/cuda/sampling_kernel.cu index f9d5d9f02..53caaf192 100644 --- a/arc-cuda-graph/src/cuda/sampling_kernel.cu +++ b/arc-cuda-graph/src/cuda/sampling_kernel.cu @@ -127,6 +127,12 @@ __device__ __forceinline__ float block_sum(float val, float* s_vals) { return s_vals[0]; } +// Branch engagement counters for the hybrid sampler. A test that assumes +// "support width 512 must have taken the fallback" is assuming the very thing +// it is trying to establish; these make it observable. +__device__ unsigned int arc_hybrid_branch_a = 0; +__device__ unsigned int arc_hybrid_branch_b = 0; + // Splitmix64 step matching sampling_cpu.rs lines 156-161. // state = state * 0x9E3779B97F4A7C15 + 0xDEADBEEFC0DECAFE // mixed = (state ^ (state >> 30)) * 0xBF58476D1CE4E5B9 @@ -353,6 +359,390 @@ __global__ void arc_sample_kernel( } } + +// ============================================================================= +// Nucleus sampling in a FIXED number of passes (threshold bisection). +// +// `arc_sample_kernel` above pulls the keep-list one argmax at a time, so it +// costs O(vocab x kept): a full vocabulary scan PER KEPT TOKEN. With +// `top_k <= 0` -- which is what mistral.rs's default maps to -- `kept` is +// bounded only by `top_p`, and on a diffuse distribution it reaches thousands. +// MEASURED on an H200 at vocab=129280 with a flat 12928-token support: +// 663,116 us to sample ONE token, roughly 10x an entire V4 decode step. +// +// This kernel selects the same nucleus by binary-searching a probability +// THRESHOLD rather than enumerating the set. For non-negative IEEE-754 floats +// the bit pattern is monotone when compared as u32, so bisecting the integer +// key converges to one ULP in a fixed 32 passes regardless of how many tokens +// the nucleus contains. Cost becomes O(vocab) with a constant factor. +// +// It does NOT narrow the distribution. The kept set is still +// {i : p_i >= t} for the largest t whose mass still covers `top_p` -- that is +// the definition of the nucleus. Where several tokens share the boundary bit +// pattern it keeps ALL of them, so it errs toward keeping MORE than the exact +// nucleus, never fewer. `top_k` is applied as a second, independent bisection +// on the kept COUNT, and the two thresholds combine with max(), so a caller +// setting both still gets the intersection. +// ============================================================================= +template +__global__ void arc_sample_bisect_kernel( + const T* __restrict__ logits, + const uint32_t* __restrict__ freq_counts, + uint64_t* __restrict__ rng_state, + int32_t* __restrict__ token_ids, + float* __restrict__ probs_scratch, + int vocab, + SamplingParams cfg +) { + const int bid = blockIdx.x; + const int tid = threadIdx.x; + const T* __restrict__ row = logits + (int64_t)bid * vocab; + float* __restrict__ probs = probs_scratch + (int64_t)bid * vocab; + const uint32_t* fcnt = freq_counts ? (freq_counts + (int64_t)bid * vocab) : nullptr; + + extern __shared__ char smem_bisect[]; + float* s_vals = reinterpret_cast(smem_bisect); + int* s_idxs = reinterpret_cast(s_vals + BLOCK); + + const bool temp_active = (cfg.temperature > 0.0f && cfg.temperature != 1.0f); + const float inv_temp = temp_active ? (1.0f / cfg.temperature) : 1.0f; + + // --- Phase 1: penalties + temperature; keep the argmax for the fallback. + float local_max = -FLT_MAX; + int local_i = 0; + for (int i = tid; i < vocab; i += BLOCK) { + float l = to_float(row[i]); + if (fcnt) { + uint32_t c = fcnt[i]; + if (c != 0) { + l -= cfg.frequency_penalty * static_cast(c); + l -= cfg.presence_penalty; + } + } + if (temp_active) l *= inv_temp; + probs[i] = l; + if (l > local_max || (l == local_max && i < local_i)) { local_max = l; local_i = i; } + } + float gmax; int gmax_idx; + block_argmax_idx(local_max, local_i, s_vals, s_idxs, gmax, gmax_idx); + __syncthreads(); + + // --- Phase 2/3: exp-shift, sum, normalize to probabilities. + float local_sum = 0.0f; + for (int i = tid; i < vocab; i += BLOCK) { + float p = __expf(probs[i] - gmax); + probs[i] = p; + local_sum += p; + } + const float gsum = block_sum(local_sum, s_vals); + __syncthreads(); + const float inv_sum = (gsum > 0.0f) ? (1.0f / gsum) : 1.0f; + float local_pmax = 0.0f; + for (int i = tid; i < vocab; i += BLOCK) { + float p = probs[i] * inv_sum; + probs[i] = p; + if (p > local_pmax) local_pmax = p; + } + const float pmax = block_max(local_pmax, s_vals); + __syncthreads(); + + // Degenerate distribution: fall back to the mode, as the CPU reference does. + if (!(pmax > 0.0f)) { + if (tid == 0) token_ids[bid] = gmax_idx; + return; + } + + float target = cfg.top_p; + if (target > 1.0f) target = 1.0f; + if (target < 0.0f) target = 0.0f; + + // --- Phase 4a: largest threshold key whose kept mass still covers top_p. + uint32_t lo = 0u, hi = __float_as_uint(pmax); + while (lo < hi) { + const uint32_t mid = lo + ((hi - lo + 1u) >> 1); + const float threshold = __uint_as_float(mid); + float m = 0.0f; + for (int i = tid; i < vocab; i += BLOCK) { + const float p = probs[i]; + if (p >= threshold) m += p; + } + const float gm = block_sum(m, s_vals); + __syncthreads(); + if (gm >= target) lo = mid; else hi = mid - 1u; + } + uint32_t key = lo; + + // --- Phase 4b: top_k, as an independent bisection on the kept COUNT. + if (cfg.top_k > 0 && cfg.top_k < vocab) { + uint32_t klo = 0u, khi = __float_as_uint(pmax); + while (klo < khi) { + const uint32_t mid = klo + ((khi - klo) >> 1); + const float threshold = __uint_as_float(mid); + float c = 0.0f; + for (int i = tid; i < vocab; i += BLOCK) { + if (probs[i] >= threshold) c += 1.0f; + } + const float gc = block_sum(c, s_vals); + __syncthreads(); + if (gc <= static_cast(cfg.top_k)) khi = mid; else klo = mid + 1u; + } + if (klo > key) key = klo; + } + + const float threshold = __uint_as_float(key); + + // --- Phase 5: CDF walk, parallel across the block. + // Each thread sums the kept mass in its own strided subset; thread 0 + // exclusive-scans those 256 partials so every thread knows the CDF offset + // its subset begins at. Only the thread whose interval contains `u` then + // walks, and it walks vocab/BLOCK items rather than the whole vocabulary. + float part = 0.0f; + for (int i = tid; i < vocab; i += BLOCK) { + const float p = probs[i]; + if (p >= threshold) part += p; + } + s_vals[tid] = part; + __syncthreads(); + + __shared__ float s_u; + if (tid == 0) { + float acc = 0.0f; + for (int t = 0; t < BLOCK; ++t) { const float v = s_vals[t]; s_vals[t] = acc; acc += v; } + uint64_t st = rng_state[bid]; + const float u = splitmix_uniform(st); + rng_state[bid] = st; + s_u = u * acc; + // Deterministic fallback if float error leaves no owning interval. + token_ids[bid] = gmax_idx; + } + __syncthreads(); + + const float base = s_vals[tid]; + if (s_u >= base && s_u < base + part) { + float acc = base; + for (int i = tid; i < vocab; i += BLOCK) { + const float p = probs[i]; + if (p >= threshold) { + acc += p; + if (acc >= s_u) { token_ids[bid] = i; break; } + } + } + } +} + + +// ============================================================================= +// Hybrid nucleus sampler: exact enumeration while the nucleus is small, +// threshold bisection once it isn't. One kernel, no host involvement. +// +// MEASURED on an H200 at vocab=129280 (support width -> legacy us / bisect us): +// 1 -> 134 / 1994 | 8 -> 503 / 1996 | 64 -> 3132 / 1991 +// 512 -> 24287 / 1997 | 4096 -> 205446 / 1998 | 12928 -> 665694 / 1998 +// Bisection is FLAT; enumeration is linear in the kept count. Neither wins +// everywhere: enumeration is 3-15x faster on the peaked distributions real +// models produce, bisection is 333x faster on the diffuse tail. Shipping +// either alone trades one regression for another. +// +// THE SWITCH POINT IS NOT A TUNING CONSTANT. Bisection resolves one bit of the +// float32 probability key per vocabulary pass, so it costs 32 passes (plus 32 +// more when top_k forces a second search). Enumeration costs one vocabulary +// pass per kept token. Enumeration is therefore cheaper exactly while +// kept < (the number of passes bisection would spend) +// so the budget IS that pass count: ARC_ENUM_BUDGET = 32, derived from the +// width of the key being searched, the same way the 32 passes are. Spending +// the budget and then falling back costs at most budget + bisection, so the +// hybrid is never worse than ~2x the better branch, and the measured crossover +// (~64) sits within a factor of two of the derived budget -- the derivation +// and the measurement agree. +// +// Why the fallback is INSIDE the kernel: a captured CUDA graph bakes its +// kernel arguments and cannot branch on device state, so "launch enumeration, +// read a flag, maybe launch bisection" would need a device->host round trip +// per step -- exactly the host dependency the autonomous decode loop exists to +// remove. The branch is therefore taken block-uniformly on device. +// +// No cap and no narrowing: the fallback is exact, not a truncation. Exceeding +// the budget changes which ALGORITHM selects the nucleus, never which tokens +// the nucleus contains. +// ============================================================================= +template +__global__ void arc_sample_hybrid_kernel( + const T* __restrict__ logits, + const uint32_t* __restrict__ freq_counts, + uint64_t* __restrict__ rng_state, + int32_t* __restrict__ token_ids, + float* __restrict__ probs_scratch, + int32_t* __restrict__ keep_idx_scratch, + float* __restrict__ keep_p_scratch, + int vocab, + SamplingParams cfg +) { + constexpr int ARC_ENUM_BUDGET = 32; // == bisection's pass count + + const int bid = blockIdx.x; + const int tid = threadIdx.x; + const T* __restrict__ row = logits + (int64_t)bid * vocab; + float* __restrict__ probs = probs_scratch + (int64_t)bid * vocab; + int32_t* __restrict__ keep_idx = keep_idx_scratch + (int64_t)bid * vocab; + float* __restrict__ keep_p = keep_p_scratch + (int64_t)bid * vocab; + const uint32_t* fcnt = freq_counts ? (freq_counts + (int64_t)bid * vocab) : nullptr; + + extern __shared__ char smem_hy[]; + float* s_vals = reinterpret_cast(smem_hy); + int* s_idxs = reinterpret_cast(s_vals + BLOCK); + + const bool temp_active = (cfg.temperature > 0.0f && cfg.temperature != 1.0f); + const float inv_temp = temp_active ? (1.0f / cfg.temperature) : 1.0f; + + float target = cfg.top_p; + if (target > 1.0f) target = 1.0f; + if (target < 0.0f) target = 0.0f; + +#define ARC_BUILD_PROBS(GMAX_IDX_OUT) \ + do { \ + float _lmax = -FLT_MAX; int _li = 0; \ + for (int i = tid; i < vocab; i += BLOCK) { \ + float l = to_float(row[i]); \ + if (fcnt) { uint32_t c = fcnt[i]; \ + if (c != 0) { l -= cfg.frequency_penalty * (float)c; \ + l -= cfg.presence_penalty; } } \ + if (temp_active) l *= inv_temp; \ + probs[i] = l; \ + if (l > _lmax || (l == _lmax && i < _li)) { _lmax = l; _li = i; } \ + } \ + float _gv; int _gi; \ + block_argmax_idx(_lmax, _li, s_vals, s_idxs, _gv, _gi); \ + __syncthreads(); \ + (GMAX_IDX_OUT) = _gi; \ + float _lsum = 0.0f; \ + for (int i = tid; i < vocab; i += BLOCK) { \ + float p = __expf(probs[i] - _gv); probs[i] = p; _lsum += p; } \ + float _gsum = block_sum(_lsum, s_vals); \ + __syncthreads(); \ + float _inv = (_gsum > 0.0f) ? (1.0f / _gsum) : 1.0f; \ + for (int i = tid; i < vocab; i += BLOCK) probs[i] *= _inv; \ + __syncthreads(); \ + } while (0) + + int gmax_idx = 0; + ARC_BUILD_PROBS(gmax_idx); + + // ---- Branch A: bounded exact enumeration (the legacy algorithm). + __shared__ int s_kept; + __shared__ float s_cum; + __shared__ int s_done; // 1 = nucleus fully determined here + if (tid == 0) { s_kept = 0; s_cum = 0.0f; s_done = 0; } + __syncthreads(); + + while (true) { + if (s_cum >= target) { if (tid == 0) s_done = 1; __syncthreads(); break; } + if (cfg.top_k > 0 && s_kept >= cfg.top_k) { if (tid == 0) s_done = 1; __syncthreads(); break; } + if (s_kept >= ARC_ENUM_BUDGET) { __syncthreads(); break; } // spend budget, fall back + + float lv = -FLT_MAX; int li = 0; + for (int i = tid; i < vocab; i += BLOCK) { + float v = probs[i]; + if (v > lv || (v == lv && i < li)) { lv = v; li = i; } + } + float gv; int gi; + block_argmax_idx(lv, li, s_vals, s_idxs, gv, gi); + __syncthreads(); + if (gv <= 0.0f) { if (tid == 0) s_done = 1; __syncthreads(); break; } + if (tid == 0) { + keep_idx[s_kept] = gi; + keep_p[s_kept] = gv; + s_kept += 1; + s_cum += gv; + probs[gi] = -FLT_MAX; // tombstone + } + __syncthreads(); + } + + if (s_done) { + if (tid == 0) atomicAdd(&arc_hybrid_branch_a, 1u); + if (s_kept <= 0) { if (tid == 0) token_ids[bid] = gmax_idx; return; } + if (tid == 0) { + uint64_t st = rng_state[bid]; + float u = splitmix_uniform(st); + rng_state[bid] = st; + float t = u * s_cum; + float acc = 0.0f; + int sel = keep_idx[s_kept - 1]; + for (int k = 0; k < s_kept; ++k) { + acc += keep_p[k]; + if (acc >= t) { sel = keep_idx[k]; break; } + } + token_ids[bid] = sel; + } + return; + } + + // ---- Branch B: budget spent, nucleus is diffuse. probs[] was tombstoned + // by the enumeration above, so rebuild it, then bisect. + if (tid == 0) atomicAdd(&arc_hybrid_branch_b, 1u); + ARC_BUILD_PROBS(gmax_idx); + + float lpmax = 0.0f; + for (int i = tid; i < vocab; i += BLOCK) { float p = probs[i]; if (p > lpmax) lpmax = p; } + const float pmax = block_max(lpmax, s_vals); + __syncthreads(); + if (!(pmax > 0.0f)) { if (tid == 0) token_ids[bid] = gmax_idx; return; } + + uint32_t lo = 0u, hi = __float_as_uint(pmax); + while (lo < hi) { + const uint32_t mid = lo + ((hi - lo + 1u) >> 1); + const float threshold_m = __uint_as_float(mid); + float m = 0.0f; + for (int i = tid; i < vocab; i += BLOCK) { float p = probs[i]; if (p >= threshold_m) m += p; } + const float gm = block_sum(m, s_vals); + __syncthreads(); + if (gm >= target) lo = mid; else hi = mid - 1u; + } + uint32_t key = lo; + + if (cfg.top_k > 0 && cfg.top_k < vocab) { + uint32_t klo = 0u, khi = __float_as_uint(pmax); + while (klo < khi) { + const uint32_t mid = klo + ((khi - klo) >> 1); + const float threshold_m = __uint_as_float(mid); + float c = 0.0f; + for (int i = tid; i < vocab; i += BLOCK) { if (probs[i] >= threshold_m) c += 1.0f; } + const float gc = block_sum(c, s_vals); + __syncthreads(); + if (gc <= (float)cfg.top_k) khi = mid; else klo = mid + 1u; + } + if (klo > key) key = klo; + } + + const float threshold = __uint_as_float(key); + float part = 0.0f; + for (int i = tid; i < vocab; i += BLOCK) { float p = probs[i]; if (p >= threshold) part += p; } + s_vals[tid] = part; + __syncthreads(); + + __shared__ float s_u; + if (tid == 0) { + float acc = 0.0f; + for (int t = 0; t < BLOCK; ++t) { const float v = s_vals[t]; s_vals[t] = acc; acc += v; } + uint64_t st = rng_state[bid]; + const float u = splitmix_uniform(st); + rng_state[bid] = st; + s_u = u * acc; + token_ids[bid] = gmax_idx; + } + __syncthreads(); + + const float base = s_vals[tid]; + if (s_u >= base && s_u < base + part) { + float acc = base; + for (int i = tid; i < vocab; i += BLOCK) { + const float p = probs[i]; + if (p >= threshold) { acc += p; if (acc >= s_u) { token_ids[bid] = i; break; } } + } + } +#undef ARC_BUILD_PROBS +} + } // namespace arc_sampler // ============================================================================= @@ -499,4 +889,70 @@ void arc_launch_sampler_f16( } } + +// Fixed-pass nucleus sampler. Same signature as `arc_launch_sampler_*` so it +// is a drop-in; `keep_idx_scratch` / `keep_p_scratch` are unused by it. +#define ARC_DEFINE_BISECT_LAUNCHER(SUFFIX, CTYPE) \ +void arc_launch_sampler_##SUFFIX##_bisect( \ + const void* logits, const uint32_t* freq_counts, uint64_t* rng_state, \ + int32_t* token_ids, float* probs_scratch, int32_t* keep_idx_scratch, \ + float* keep_p_scratch, int vocab, int batch, \ + arc_sampler::SamplingParams cfg, cudaStream_t stream) { \ + (void)keep_idx_scratch; (void)keep_p_scratch; \ + const int B = 256; \ + size_t smem = B * (sizeof(float) + sizeof(int)); \ + if (cfg.greedy) { \ + arc_sampler::arc_greedy_kernel<<>>( \ + reinterpret_cast(logits), freq_counts, token_ids, vocab, \ + cfg); \ + return; \ + } \ + arc_sampler::arc_sample_bisect_kernel<<>>( \ + reinterpret_cast(logits), freq_counts, rng_state, \ + token_ids, probs_scratch, vocab, cfg); \ +} + +ARC_DEFINE_BISECT_LAUNCHER(f32, float) +ARC_DEFINE_BISECT_LAUNCHER(bf16, __nv_bfloat16) +ARC_DEFINE_BISECT_LAUNCHER(f16, __half) + + +void arc_hybrid_branch_counts(unsigned int* a_out, unsigned int* b_out) { + cudaMemcpyFromSymbol(a_out, arc_sampler::arc_hybrid_branch_a, + sizeof(unsigned int), 0, cudaMemcpyDeviceToHost); + cudaMemcpyFromSymbol(b_out, arc_sampler::arc_hybrid_branch_b, + sizeof(unsigned int), 0, cudaMemcpyDeviceToHost); +} + +void arc_hybrid_branch_reset(void) { + unsigned int z = 0; + cudaMemcpyToSymbol(arc_sampler::arc_hybrid_branch_a, &z, sizeof(z), 0, + cudaMemcpyHostToDevice); + cudaMemcpyToSymbol(arc_sampler::arc_hybrid_branch_b, &z, sizeof(z), 0, + cudaMemcpyHostToDevice); +} + +#define ARC_DEFINE_HYBRID_LAUNCHER(SUFFIX, CTYPE) \ +void arc_launch_sampler_##SUFFIX##_hybrid( \ + const void* logits, const uint32_t* freq_counts, uint64_t* rng_state, \ + int32_t* token_ids, float* probs_scratch, int32_t* keep_idx_scratch, \ + float* keep_p_scratch, int vocab, int batch, \ + arc_sampler::SamplingParams cfg, cudaStream_t stream) { \ + const int B = 256; \ + size_t smem = B * (sizeof(float) + sizeof(int)); \ + if (cfg.greedy) { \ + arc_sampler::arc_greedy_kernel<<>>( \ + reinterpret_cast(logits), freq_counts, token_ids, vocab, \ + cfg); \ + return; \ + } \ + arc_sampler::arc_sample_hybrid_kernel<<>>( \ + reinterpret_cast(logits), freq_counts, rng_state, \ + token_ids, probs_scratch, keep_idx_scratch, keep_p_scratch, vocab, cfg); \ +} + +ARC_DEFINE_HYBRID_LAUNCHER(f32, float) +ARC_DEFINE_HYBRID_LAUNCHER(bf16, __nv_bfloat16) +ARC_DEFINE_HYBRID_LAUNCHER(f16, __half) + } // extern "C" diff --git a/arc-cuda-graph/src/ffi.rs b/arc-cuda-graph/src/ffi.rs index 96cae2954..813a2a2f4 100644 --- a/arc-cuda-graph/src/ffi.rs +++ b/arc-cuda-graph/src/ffi.rs @@ -14,6 +14,21 @@ pub type CUgraphConditionalHandle = u64; #[cfg(feature = "cuda")] pub type CUmemoryPool = *mut std::ffi::c_void; #[cfg(feature = "cuda")] +pub type CUcontext = *mut std::ffi::c_void; +/// `cuda.h:1944` (CUDA 13.1) โ€” +/// `#define CU_GRAPH_COND_ASSIGN_DEFAULT 0x1` "Default value is applied when +/// graph is launched." +/// +/// `cuda.h:21919` documents `defaultLaunchValue` as "Applied at the beginning +/// of each graph execution **if CU_GRAPH_COND_ASSIGN_DEFAULT is set in +/// flags**". Creating the handle with `flags = 0` therefore leaves the +/// condition at 0 on entry and a WHILE body executes **zero** times โ€” the +/// graph launches, returns success, and generates nothing. Measured on +/// `arc-v4-stack` (H200, CUDA 13.1): `flags=0` โ†’ body ran 0 times; +/// `flags=CU_GRAPH_COND_ASSIGN_DEFAULT` โ†’ body ran exactly N times. +#[cfg(feature = "cuda")] +pub const CU_GRAPH_COND_ASSIGN_DEFAULT: u32 = 0x1; +#[cfg(feature = "cuda")] pub type CUdevice = i32; #[cfg(feature = "cuda")] @@ -70,6 +85,14 @@ pub enum CudaGraphNodeType { Conditional = 13, } +/// `CUDA_CONDITIONAL_NODE_PARAMS`, `cuda.h:1958` (CUDA 13.1). +/// +/// `phGraph_out` is an **OUT** field: "CUDA-owned array populated with +/// conditional node child graphs during creation of the node." Callers leave +/// it null and read it back after `cuGraphAddNode`; assigning a caller-owned +/// pointer to it accomplishes nothing because the driver overwrites it. +/// +/// `ctx` (the 5th field) was missing entirely from this struct. #[cfg(feature = "cuda")] #[repr(C)] pub struct CudaConditionalNodeParams { @@ -77,6 +100,7 @@ pub struct CudaConditionalNodeParams { pub cond_type: CUgraphConditionalNodeType, pub size: u32, pub body_graph_out: *mut CUgraph, + pub ctx: CUcontext, } #[cfg(feature = "cuda")] @@ -162,6 +186,7 @@ extern "C" { pub fn cuGraphLaunch(exec: CUgraphExec, stream: CUstream) -> u32; pub fn cuGraphExecDestroy(exec: CUgraphExec) -> u32; pub fn cuGraphDestroy(graph: CUgraph) -> u32; + pub fn cuGraphGetNodes(graph: CUgraph, nodes: *mut CUgraphNode, num_nodes: *mut usize) -> u32; // ========== Memory pools ========== pub fn cuMemPoolCreate(pool: *mut CUmemoryPool, props: *const CUmemPoolProps) -> u32; @@ -188,20 +213,53 @@ extern "C" { pub fn cudaStreamQuery(stream: CUstream) -> u32; // ========== Conditional nodes (CUDA 12.4+) ========== - pub fn cudaGraphConditionalHandleCreate( + /// `cuda.h:21932` (CUDA 13.1) โ€” the driver form takes a `CUcontext`, + /// which the conditional node params must match ("Context on which to run + /// the node. Must match context used to create the handle and all body + /// nodes", `cuda.h:1986`). + pub fn cuGraphConditionalHandleCreate( handle: *mut CUgraphConditionalHandle, graph: CUgraph, - default_value: u32, + ctx: CUcontext, + default_launch_value: u32, flags: u32, ) -> u32; - pub fn cudaGraphSetConditional(handle: CUgraphConditionalHandle, value: u32) -> u32; - pub fn cudaGraphAddNode( - graph: CUgraph, + pub fn cuCtxGetCurrent(ctx: *mut CUcontext) -> u32; + /// `cuda.h:21829` (CUDA 13.1): + /// ```text + /// CUresult cuGraphAddNode(CUgraphNode *phGraphNode, CUgraph hGraph, + /// const CUgraphNode *dependencies, + /// const CUgraphEdgeData *dependencyData, + /// size_t numDependencies, + /// CUgraphNodeParams *nodeParams); + /// ``` + /// The OUT node pointer is the **first** argument and there are **six** of + /// them. This was previously declared as + /// `cudaGraphAddNode(graph, node_out, deps, num_deps, params)` โ€” five + /// arguments with the first two transposed โ€” so the driver received the + /// `CUgraph` handle in the slot where it expected `CUgraphNode *` and + /// wrote the new node handle through that value, while `numDependencies` + /// received a pointer and `nodeParams` an uninitialised register. + pub fn cuGraphAddNode( node_out: *mut CUgraphNode, + graph: CUgraph, dependencies: *const CUgraphNode, + dependency_data: *const std::ffi::c_void, num_dependencies: usize, params: *mut CudaGraphNodeParams, ) -> u32; + /// `cuda.h:15665` (CUDA 13.1). `cuda.h:1976` names this as a supported way + /// to populate a conditional node's body graph. Ordinary + /// `cuStreamBeginCapture_v2` always creates a **new** graph, so a body + /// captured that way cannot be attached to the conditional node. + pub fn cuStreamBeginCaptureToGraph( + stream: CUstream, + graph: CUgraph, + dependencies: *const CUgraphNode, + dependency_data: *const std::ffi::c_void, + num_dependencies: usize, + mode: CUstreamCaptureMode, + ) -> u32; // ========== Pinned host memory ========== pub fn cudaHostAlloc(ptr: *mut *mut std::ffi::c_void, size: usize, flags: u32) -> u32; diff --git a/arc-cuda-graph/src/sampling_cuda.rs b/arc-cuda-graph/src/sampling_cuda.rs index 8d6feb95a..67d52a939 100644 --- a/arc-cuda-graph/src/sampling_cuda.rs +++ b/arc-cuda-graph/src/sampling_cuda.rs @@ -209,6 +209,59 @@ extern "C" { cfg: SamplingParams, stream: CUstream, ); + + // Hybrid: exact enumeration while the nucleus is small, threshold + // bisection past a fixed budget, chosen block-uniformly ON DEVICE so a + // captured graph never has to ask the host which branch to take. + // + // MEASURED, H200, vocab=129280, host+GPU verified exclusive before and + // after (support width -> legacy / hybrid us): + // 1 -> 134.9/166.6 | 8 -> 503.0/624.2 | 64 -> 3127.8/4178.5 + // 512 -> 24320.7/4185.0 | 4096 -> 205125.2/4187.3 + // 12928 -> 664200.0/4188.5 + // Costs +24% on the peaked distributions real models produce (32 us, i.e. + // 0.05% of a 66.68 ms V4 decode step) and removes a 664 ms cliff -- 158x + // at the diffuse tail, where the enumerating sampler is ~10x an entire + // decode step for a single token. + pub fn arc_launch_sampler_f32_hybrid( + logits: *const f32, + freq_counts: *const u32, + rng_state: *mut u64, + token_ids: *mut i32, + probs_scratch: *mut f32, + keep_idx_scratch: *mut i32, + keep_p_scratch: *mut f32, + vocab: i32, + batch: i32, + cfg: SamplingParams, + stream: CUstream, + ); + pub fn arc_launch_sampler_bf16_hybrid( + logits: *const std::ffi::c_void, + freq_counts: *const u32, + rng_state: *mut u64, + token_ids: *mut i32, + probs_scratch: *mut f32, + keep_idx_scratch: *mut i32, + keep_p_scratch: *mut f32, + vocab: i32, + batch: i32, + cfg: SamplingParams, + stream: CUstream, + ); + pub fn arc_launch_sampler_f16_hybrid( + logits: *const std::ffi::c_void, + freq_counts: *const u32, + rng_state: *mut u64, + token_ids: *mut i32, + probs_scratch: *mut f32, + keep_idx_scratch: *mut i32, + keep_p_scratch: *mut f32, + vocab: i32, + batch: i32, + cfg: SamplingParams, + stream: CUstream, + ); } /// Raw device pointer for a CUDA tensor, byte-offset aware. @@ -311,6 +364,15 @@ impl CudaSampler { }) } + /// Bind the sampler to `stream`. + /// + /// `new()` picks up the device-default stream. Anything that is going to + /// be **captured** into a CUDA graph must launch on the capture stream, or + /// the sampler's kernels are simply not recorded into the graph. + pub fn set_stream(&mut self, stream: CUstream) { + self.stream = stream; + } + /// Run the sampler. Logits: [batch, vocab] of the configured dtype. /// `freq_counts`: optional [batch, vocab] u32 token-count tensor (one per /// batch row); pass `None` when no penalties are needed. @@ -370,7 +432,7 @@ impl CudaSampler { unsafe { match self.dtype { - DType::F32 => arc_launch_sampler_f32( + DType::F32 => arc_launch_sampler_f32_hybrid( logits_ptr as *const f32, freq_ptr, rng_ptr, @@ -383,7 +445,7 @@ impl CudaSampler { params, self.stream, ), - DType::BF16 => arc_launch_sampler_bf16( + DType::BF16 => arc_launch_sampler_bf16_hybrid( logits_ptr, freq_ptr, rng_ptr, @@ -396,7 +458,7 @@ impl CudaSampler { params, self.stream, ), - DType::F16 => arc_launch_sampler_f16( + DType::F16 => arc_launch_sampler_f16_hybrid( logits_ptr, freq_ptr, rng_ptr, diff --git a/arc-cuda-graph/src/weights.rs b/arc-cuda-graph/src/weights.rs index f8891c311..bdf77b4b0 100644 --- a/arc-cuda-graph/src/weights.rs +++ b/arc-cuda-graph/src/weights.rs @@ -610,6 +610,16 @@ pub fn tensor_device_ptr(tensor: &Tensor) -> candle_core::Result { // fed with (radix top-k `seq_lens`, `CudaSampler` token ids and // keep-list scratch). Candle's `Tensor::from_vec(Vec, โ€ฆ)` // produces it, so leaving it out disabled those paths wholesale. + // + // Concretely, from #130: `CudaSampler::sample` REQUIRES + // `token_ids` to be I32 (`sampling_cuda.rs:347`) and allocates + // an I32 `keep_idx_scratch` (`:298`). Both resolve through this + // function, so before I32 was listed both fell to the catch-all + // `bail!` and `CudaSampler::sample()` returned `unsupported + // dtype I32` on EVERY call โ€” the replay-safe top-k/top-p + // sampler could never run on a GPU. Its test suite is a CPU + // simulator (`gpu_algorithm_simulate`), which by construction + // cannot observe that. DType::I32 => { let s = cuda_storage.as_cuda_slice::()?; let (p, _) = s.device_ptr(s.stream()); diff --git a/mistralrs-core/src/pipeline/normal.rs b/mistralrs-core/src/pipeline/normal.rs index de530e395..0048e34c9 100644 --- a/mistralrs-core/src/pipeline/normal.rs +++ b/mistralrs-core/src/pipeline/normal.rs @@ -3378,6 +3378,16 @@ impl Pipeline for NormalPipeline { top_p, frequency_penalty, presence_penalty, + // mistral.rs uses <=0 to mean "disabled"; the fused sampler + // spells that -1 (`sampling_cpu.rs:16`). + top_k: { + let tk = first_sampler.top_k(); + if tk <= 0 { + -1 + } else { + tk as i32 + } + }, greedy, }; diff --git a/mistralrs-core/tests/capability_reachability.rs b/mistralrs-core/tests/capability_reachability.rs index 0ce630f50..c203fd761 100644 --- a/mistralrs-core/tests/capability_reachability.rs +++ b/mistralrs-core/tests/capability_reachability.rs @@ -229,17 +229,20 @@ static REGISTRY: &[Capability] = &[ accepted_in: &["mistralrs-core/src/sampler.rs"], honoured_in: &["arc-cuda-graph/src/autonomous.rs"], }, - // `AutonomousDecodeConfig` (autonomous.rs:70-82) carries temperature, - // top_p, both penalties and `greedy` โ€” and neither top_k nor min_p. A - // request setting them would have them silently dropped the moment this - // path is reached. It is not reached today (`mark_unreachable( - // "cuda_graph.autonomous_decode", ...)`), which is the only reason this - // is not already a user-visible defect. - status: Status::Tracked { - reason: "AutonomousDecodeConfig has no top_k field at all; the path is itself \ - unreachable today, so this is latent rather than live. Fix before \ - GPU-autonomous decode is switched on.", - }, + // Was Tracked with the reason "AutonomousDecodeConfig has no top_k + // field at all". #130 gave it one: `AutonomousDecodeConfig::top_k` + // (autonomous.rs:82), plumbed into `sampling_cpu::SamplingConfig` at + // autonomous.rs:643 and consumed by `self.sampler.sample(..)`. So the + // promise is now honoured where it is made, and the status follows the + // code. + // + // The gate flagged this itself, on the rebase, with "PROMOTE ME" โ€” + // the same way the greedy/logit-bias entry above was promoted. That is + // the mechanism working, not a test to be silenced. + // + // `min_p` is a separate entry below and stays Tracked: it still has no + // field on `AutonomousDecodeConfig`. + status: Status::Live, }, Capability { name: "sampling: min_p survives onto the GPU-autonomous decode path", diff --git a/mistralrs-core/tests/doc_citations.rs b/mistralrs-core/tests/doc_citations.rs index e4204e8cc..6d219d675 100644 --- a/mistralrs-core/tests/doc_citations.rs +++ b/mistralrs-core/tests/doc_citations.rs @@ -351,6 +351,30 @@ static BASELINE: &[Waiver] = &[ kind: Kind::Unresolved, why: "candle is an out-of-tree path dependency", }, + // The CUDA driver header. #130's FFI declarations cite it by line for the + // struct layouts and flag values they mirror (`CUDA_CONDITIONAL_NODE_PARAMS`, + // `CU_GRAPH_COND_ASSIGN_DEFAULT`, the conditional-node entry points). It + // ships with the CUDA Toolkit and is not vendored here, so no citation into + // it can ever resolve โ€” and its line numbers are toolkit-version-specific + // besides (these are CUDA 13.1), which is why each citation names the + // version in its surrounding comment. + // + // Waived by bare filename rather than one row per line number: the rows + // would otherwise have to be rewritten on every toolkit bump, which is + // churn the gate cannot check either way. The layouts themselves ARE + // checked, against `cudarc`'s bindgen output โ€” see the ABI note on #130. + Waiver { + doc: "arc-cuda-graph/src/ffi.rs", + cite: "cuda.h", + kind: Kind::Unresolved, + why: "CUDA Toolkit header, not vendored here", + }, + Waiver { + doc: "arc-cuda-graph/src/autonomous.rs", + cite: "cuda.h", + kind: Kind::Unresolved, + why: "CUDA Toolkit header, not vendored here", + }, Waiver { doc: "memory/mission/wave29-BD-rung-decision.md", cite: "STATUS.md", From 2423f3451e337e3a2c2b566c3163012dfa269db6 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 21:43:52 +0100 Subject: [PATCH 15/21] =?UTF-8?q?perf(ArcAttention):=20V4=20RoPE=20was=20l?= =?UTF-8?q?aunched=20per-sequence=20=E2=80=94=2099,072=20launches/step=20-?= =?UTF-8?q?>=20301=20(#198)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * perf(ArcAttention): batch V4 RoPE over the cohort โ€” 99,072 launches/step -> 301 `DeepSeekV2RotaryEmbedding::{forward,forward_inverse_tail}` gated the batched path on `seqlen_offsets.len() == 1` โ€” the LENGTH of the offset vector, where the thing that matters is the DISTINCTNESS of its values. `len()` is the batch size, so the batched path was unreachable at every batch size above one. The offsets are uniform anyway. `scheduler::default_scheduler`'s `select_running_bucket` admits a forward pass only when cache lengths are exactly equal, and the one producer of a ragged dense batch is gated behind `ARC_MTP_PER_SEQ_KV`, which defaults off (`deepseek4::ragged_row_q0`). So the loop was performing B bit-identical recomputations and concatenating them back. Measured, B=256 pure decode, steady state (nsys, 204.80 tok/s): kernel share launches/step per (layer, seq) ucopy_bf16 10.2% 22,794 2 (2*43*256 = 22,016) copy2d_bf16 5.0% 33,610 3 (3*43*256 = 33,024) rope_i_bf16 3.9% 33,054 3 (3*43*256 = 33,024) uneg_bf16 0.8% 11,008 1 (1*43*256 = 11,008, exact) 19.9% of step time, and the only cost in the profile that gets WORSE as batch grows. `copy2d` is the 256-argument `Tensor::cat` that exists solely to undo the split the same function just made; it goes to zero. `uneg` is 11,008 launches to negate a [1,32] slice of a CONSTANT table, so `neg_sin` is now materialised once at construction โ€” not a kernel to fuse, a kernel to delete. Cohort form is 43 layers * 7 = 301 launches (129 of them rope_i), which lands RoPE at the same order as fp8_matmul_tiled (301) and qtip2b_grouped_gemm (129) instead of dwarfing them. PREDICTION, not a result: this is a launch-count fix, the bytes moved are unchanged, and how much of the 19.9% actually returns must be measured on hardware. Ragged offsets still take the loop verbatim, so enabling ARC_MTP_PER_SEQ_KV behaves exactly as today. No ragged index_select path is added โ€” nothing ships dark. Bit-identical by construction: with 2-D [T, D/2] cos/sin the kernel's stride_b is 0 (candle-nn/src/rotary_emb.rs), so every batch row reads the same cos/sin row and each output element is an independent two-multiply-one-add. Seven tests hold the pre-fix loop as a verbatim oracle and assert raw IEEE-754 bit equality (not a tolerance) for uniform AND ragged offsets, at seq_len 1 and 2, with a one-ULP negative control proving the comparator can fail. Engagement is counted (`rope_cohort_stats::counts`): a fast path that silently declined would pass every equality assertion trivially. The counter already earned itself โ€” it caught a missing lock in one of these tests. Co-Authored-By: Claude Opus 5 (1M context) * perf(ArcAttention): un-deaden the fused qk_norm_rope kernel above batch 1 `fused_qk_gate` declined with "per-sequence position offsets" whenever `seqlen_offsets.len() != 1` โ€” the same length-vs-values confusion as the RoPE dispatch in the parent commit, and with the same consequence: the fused one-launch kernel was dead at EVERY batch size above one, i.e. in exactly the serving regime it was written for. Its own doc comment advertises "replacing 16 candle launches per layer with 1"; at B=256 it replaced nothing. The kernel is already batch-correct and needed no change. `qk_norm_rope.cu` flattens (b, t) into blockIdx.y, recovers `b = bt / seq_len`, and reads the table at `pos_offset + t` โ€” a position independent of `b`, which is exactly right when the rows share an offset. Its grid is (n_heads + 1, batch * seq_len) and the wrapper sizes buffers as `batch * ...`. Nothing assumed batch == 1. Gate now takes the uniform offset via the shared `DeepSeekV2RotaryEmbedding::uniform_offset` helper, so both call sites cannot drift apart. Genuinely ragged rows still decline and still take the eager chain. Split from the parent commit so it can be reverted independently: this one swaps in a CUDA kernel that has never run above batch 1, and unlike the parent it is not bit-identical-by-construction โ€” it is bit-identical by the kernel's own design, which `ARC_QK_VERIFY=1` checks against the eager chain at every layer. Verify before trusting; `ARC_QK_FUSED=0` restores the eager path from the same binary. Co-Authored-By: Claude Opus 5 (1M context) * fix(ArcAttention): one cohort predicate, one counter pair, per-thread test counts Integration fix for #198 landing after #204. Both PRs made the same correction -- `seqlen_offsets.len() == 1` tests the length of the vector where the property needed is the distinctness of its values -- and each shipped its own spelling of the predicate and its own engagement counters. `uniform_seqlen_offset`'s doc comment on the integration branch asked for exactly this collapse ("there must not be two spellings of the same predicate"); this is that collapse. 1. `pub mod rope_cohort_stats` was declared TWICE in one file (layers.rs:400 from #204, :1568 from #198). That is a hard name collision -- the file did not compile. #204's is the superset (it has `record_cohort` / `record_per_sequence` as well as `counts`), so #198's is the one that goes. 2. `DeepSeekV2RotaryEmbedding::uniform_offset` had a byte-identical body to `uniform_seqlen_offset`. Removed; `forward`, `forward_inverse_tail` and `deepseek4::fused_qk_gate` now call the free function, so the two cannot drift apart. The reasoning #198's doc comment carried and #204's did not -- that uniformity is *enforced* by `select_running_bucket` rather than merely observed, and the measured B=256 cost -- moves onto the survivor. 3. Collapsing the counters made them shared, and shared process-global counters cannot be delta-asserted from a parallel test binary. `deepseek_rope_cohort_ tests` held a private mutex, `rope_cohort_tests` held none (until now it was the only writer), and three tests duly failed with `(2, 1)` where they wanted `(1, 0)`. A shared mutex was the first fix and it is also wrong: it obliges every present and future test that touches ANY rotary to know about a lock in another module. Two such tests already exist and did not -- `deepseek4::tests::{standard_layer_rope_is_unscaled_and_compressed_is_not, mtp_block_takes_the_standard_unscaled_rope_table}` both drive `DeepSeekV2RotaryEmbedding::forward` -- and they race the cohort tests. That was caught by merging #198 and #121 together and running the full suite: `--test-threads=1` passes 707, the parallel run fails one. So the counters now also keep a `#[cfg(test)]` per-thread mirror, and the tests assert on `local_counts()`. Each test's delta is its own by construction, no cooperation from other modules is required, and the serving instrument -- the global `counts()` -- is byte-for-byte unchanged. Also drops #198's copy of `uniform_offset_tests_values_not_length`: with one predicate it asserted nothing the surviving copy does not, against the same six cases. `cargo test -p mistralrs-core --lib`: 699 passed, 0 failed; 707 passed on the merged #198+#181+#130+#121 tree. The 99,072 -> 301 launch figure remains #198's own branch measurement and is NOT re-measured here. Co-Authored-By: Claude Opus 5 (1M context) * fix(ArcAttention): the batched fused qk_norm_rope path is OPT-IN The parent commit un-deadens the fused Q/K kernel above batch 1 by fixing a length-vs-values gate. That is the right fix to the gate and the wrong default: it makes a CUDA kernel LIVE on a path it has never run on. This repo has paid for that twice in one week. * TCFRAG shipped default-ON carrying "UNVERIFIED ON HARDWARE โ€” NEVER RUN" in its own header. It held 63 GB and permanently broke a layer through a poisoned `OnceLock`. #209 retired it to opt-in. * The fused-512 attention path ran for four days silently dropping the attention mask, at 12% agreement with the reference, because nobody ran the path it had quietly become live on. So the doctrine, and this commit: **unverified means default-off, and "unverified" means unmeasured, not new.** `ARC_QK_FUSED_COHORT` gates the batched path, default OFF. With it unset, `batch > 1` declines and takes the eager chain exactly as master does today โ€” so the parent commit ships its gate fix with no behaviour change, and the cohort path is available to be measured rather than discovered in production. `batch == 1` is untouched: that path is already live and already exercised. Read by VALUE, `== Some("1")`, split into the pure `fused_cohort_enabled_from` so the polarity is testable without mutating the environment (racy across `cargo test`'s threads, `unsafe` since the 2024 edition, and meaningless against a `OnceLock` that latches the first read). Two tests pin it, and they run with no features and no GPU on every PR: unset is OFF, and `0`/`false`/ `off`/`true`/`on`/`2`/`"1 "` are all OFF. That second one is not pedantry โ€” #212 converted 23 `ARC_*` flags this week because `var_os(..).is_some()` made `ARC_FOO=0` mean ON, which silently cancels any A/B whose control sets zero. To turn it on, once `ARC_QK_VERIFY=1` has bit-compared it against the eager chain on hardware and it is measured faster: `ARC_QK_FUSED_COHORT=1`. When that measurement exists, this default flips and the doc comment goes with it. `cargo test -p mistralrs-core --lib`: 701 passed, 0 failed. Co-Authored-By: Claude Opus 5 (1M context) * perf(ArcAttention): gather the RAGGED cohort's RoPE โ€” one rank-3 rope_i, opt-in ARC_ROPE_COHORT The uniform-offset fix (#198, and the shared predicate it now rides on) left one arm untouched: rows at genuinely DIFFERENT positions still take the per-sequence loop โ€” ~9 ops x 43 layers x B per step on the measured B=256 profile (~99,000 launches, 19.9% of step time), of which the copy2d half only undoes a split the same function just made. Nothing structural required the loop: candle's rope_i accepts a rank-3 [B, T, D/2] cos/sin whenever B matches the input batch (rope_check_cs, candle-nn/src/rotary_emb.rs, pinned rev 89ab14e; the CUDA wrapper sets stride_b for rank-3, the CPU path indexes i + b_i*t*d/2). So the ragged arm becomes: gather the per-row table rows once โ€” [B*T] u32 indices, host-built because the offsets ARE host data, memoised per instance so the H2D and the three index_selects run once per step per device (the model shares one rotary per device per table kind), CLAUDE.md pitfall #5 โ€” then ONE rope_i per projection. ~5 launches per layer where the loop paid ~9 x B. DEFAULT OFF behind ARC_ROPE_COHORT, read by VALUE via env_flag_is_set (=0 means OFF; #212), latched once. Unset keeps the loop byte-for-byte. Registered in capability_reachability.rs beside the fused-cohort entry, same doctrine: unverified means default-off, and unverified means unmeasured. Both ragged arms are named methods driven directly by deepseek_ragged_cohort_tests โ€” bit-pattern parity (not tolerance) for mixed offsets at B in {1,3}, T in {1,2}, ascending and descending, forward and inverse tail, plus a stale-memo test (the cache must rebuild when the cohort moves) and a pinned unset-default. engaged_count in cuda/qk_norm_rope.rs also gains its first reader: a test that proves the counters advance, closing the silent-success hole where an A/B could read a dead path as live. Parent system: ArcInfer / ArcAttention. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf * fix(ArcAttention): bounds-guard the cohort gather โ€” both ragged arms must fail LOUD The loop arm's narrow(0, offset, seq_len) hard-errors past the table end; the gathered arm fed the same out-of-range index to index_select, which is not reliably bounds-checked on the CUDA backend and can read garbage SILENTLY. A failure-mode asymmetry on exactly the arm the hardware A/B will exercise is the D18 house fault, so the offsets โ€” host data, the check is free โ€” are now guarded host-side in cohort_tables before any index reaches the device, with the offending offset, seq_len and table size named in the error. out_of_bounds_offset_fails_loud_on_both_arms pins: the loop arm errors, the cohort arm errors on forward AND inverse tail, the message is the GUARD's (a backend error cannot pass the test by coincidence), and offset 31 + T=1 on a 32-row table still works, so the bound is not off by one. Parent system: ArcInfer / ArcAttention. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf --------- Co-authored-by: Nirupam Bhowmick Co-authored-by: Claude Opus 5 (1M context) --- mistralrs-core/src/cuda/qk_norm_rope.rs | 116 +++ mistralrs-core/src/layers.rs | 926 +++++++++++++++++- mistralrs-core/src/models/deepseek4.rs | 28 +- .../tests/capability_reachability.rs | 68 ++ 4 files changed, 1087 insertions(+), 51 deletions(-) diff --git a/mistralrs-core/src/cuda/qk_norm_rope.rs b/mistralrs-core/src/cuda/qk_norm_rope.rs index 68314ec23..013d7f41b 100644 --- a/mistralrs-core/src/cuda/qk_norm_rope.rs +++ b/mistralrs-core/src/cuda/qk_norm_rope.rs @@ -49,6 +49,60 @@ pub fn fused_enabled() -> bool { *ENABLED.get_or_init(|| !matches!(std::env::var("ARC_QK_FUSED").as_deref(), Ok("0"))) } +/// `ARC_QK_FUSED_COHORT=1` lets the fused kernel serve a **batched** cohort +/// (`batch > 1`). **Default OFF**, so `batch > 1` declines and takes the eager +/// chain exactly as it does today. +/// +/// # Why this is opt-in and [`fused_enabled`] is not +/// +/// The kernel is batch-correct by construction โ€” `qk_norm_rope.cu` flattens +/// `(b, t)` into `blockIdx.y`, recovers `b = bt / seq_len`, and reads the table +/// at `pos_offset + t`, a position independent of `b` โ€” and until #198 the +/// dispatch gate tested `seqlen_offsets.len() != 1`, the LENGTH of the offset +/// vector where the property needed is the DISTINCTNESS of its values. So the +/// fused path was dead at every batch size above one. #198 fixes that gate. +/// +/// Fixing the gate, though, makes a kernel **live on a path it has never run +/// on**, and this repo has paid for that twice in one week: +/// +/// * TCFRAG shipped default-ON with "UNVERIFIED ON HARDWARE โ€” NEVER RUN" in its +/// own header. It held 63 GB and permanently broke a layer through a poisoned +/// `OnceLock`; #209 made it opt-in. +/// * The fused-512 attention path ran for four days silently dropping the +/// attention mask, at 12% agreement with the reference, because nobody ran +/// the path it had quietly become live on. +/// +/// So the doctrine is: **unverified means default-off, and "unverified" means +/// unmeasured, not new.** `batch == 1` is unaffected โ€” that path is already +/// live and already exercised. When the cohort path has been measured correct +/// and faster on hardware (`ARC_QK_VERIFY=1` bit-compares it against the eager +/// chain at every layer), this default flips and this doc comment goes with it. +/// +/// Read by VALUE, `== Some("1")`, not by presence: 23 `ARC_*` flags were +/// converted this week (#212) precisely because `var_os(..).is_some()` made +/// `ARC_FOO=0` mean ON. +pub fn fused_cohort_enabled() -> bool { + use std::sync::OnceLock; + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| { + fused_cohort_enabled_from(std::env::var("ARC_QK_FUSED_COHORT").ok().as_deref()) + }) +} + +/// Pure half of [`fused_cohort_enabled`], so the polarity is testable without +/// mutating the process environment. +/// +/// Same reasoning as `mistralrs_quant::env_flag_value` and +/// `qtip::tcfrag2b_enabled_from`: mutating the environment from a test is racy +/// across `cargo test`'s threads and, since the 2024 edition, `unsafe` โ€” and +/// the caller latches its answer in a `OnceLock`, so such a test would prove +/// only whatever the first reader in the process happened to see. Testing the +/// pure function against every input is the only form of this test that means +/// anything. +pub fn fused_cohort_enabled_from(value: Option<&str>) -> bool { + value == Some("1") +} + /// `ARC_QK_VERIFY=1` runs the eager chain alongside the fused kernel at every /// layer and compares the raw output bytes. /// @@ -443,3 +497,65 @@ static VERIFIED: AtomicU64 = AtomicU64::new(0); pub fn verified_elements() -> u64 { VERIFIED.load(Ordering::Relaxed) } + +#[cfg(test)] +mod cohort_gate_tests { + use super::*; + + /// The batched fused path must be OFF unless explicitly asked for. + /// + /// This is the assertion that stops a repeat of TCFRAG โ€” default-ON with + /// "UNVERIFIED ON HARDWARE โ€” NEVER RUN" in its own header, 63 GB held, a + /// layer permanently broken through a poisoned `OnceLock`, retired by #209 + /// โ€” and of the fused-512 attention path, which ran for four days silently + /// dropping the attention mask at 12% agreement with the reference. Both + /// were live on a path nobody had run. + /// + /// These run with no features and no GPU, on three operating systems, on + /// every PR, because the polarity is a pure function. + #[test] + fn the_batched_cohort_is_off_unless_explicitly_enabled() { + assert!( + !fused_cohort_enabled_from(None), + "unset must mean OFF โ€” the cohort path has never run on hardware" + ); + assert!(fused_cohort_enabled_from(Some("1"))); + } + + /// Read by VALUE, not presence. #212 converted 23 `ARC_*` flags because + /// `var_os(..).is_some()` made `ARC_FOO=0` mean ON โ€” which also silently + /// cancels any A/B whose control arm sets the flag to zero. + #[test] + fn zero_means_off_and_so_does_every_other_value() { + for v in ["0", "", "false", "no", "off", "true", "yes", "on", "2", "1 "] { + assert!( + !fused_cohort_enabled_from(Some(v)), + "{v:?} must not enable the cohort path; only the exact string \"1\" does" + ); + } + } + + /// [`engaged_count`] is the instrument the A/B harness reads to tell "the + /// kernel ran" apart from "the benchmark got faster" โ€” and until this + /// test, nothing in the tree read it, so it could have silently stopped + /// advancing and every claim built on it would still look green. The + /// counters are process-global and other tests may bump them + /// concurrently, so the assertion is on a lower bound of the delta, which + /// only this thread's calls can guarantee. + #[test] + fn engaged_count_advances_when_the_paths_are_taken() { + let (e0, d0) = engaged_count(); + note_engaged(128, 64, 16, "bf16"); + note_declined("test: exercising the decline counter"); + let (e1, d1) = engaged_count(); + assert!( + e1 >= e0 + 1, + "note_engaged did not advance the engaged counter ({e0} -> {e1}); \ + an A/B reading engaged_count() would call a dead path live" + ); + assert!( + d1 >= d0 + 1, + "note_declined did not advance the declined counter ({d0} -> {d1})" + ); + } +} diff --git a/mistralrs-core/src/layers.rs b/mistralrs-core/src/layers.rs index 71701e852..818f235a1 100644 --- a/mistralrs-core/src/layers.rs +++ b/mistralrs-core/src/layers.rs @@ -375,12 +375,27 @@ impl QRmsNorm { /// `None` keeps the old loop verbatim, so ragged cohorts behave exactly as they /// do today. This is a dispatch fix, not a new ragged path. /// -/// # Relationship to `DeepSeekV2RotaryEmbedding::uniform_offset` +/// # The V4 rotary, collapsed onto this predicate /// -/// The V4 rotary got this same correction separately (PR #198, still unmerged -/// at the time of writing) as a private associated function with an identical -/// body. When these two lines of work meet, collapse that one onto this one -- -/// there must not be two spellings of the same predicate. +/// The V4 rotary (`DeepSeekV2RotaryEmbedding`) got this same correction +/// separately, in #198, as a private associated function `uniform_offset` with +/// a byte-identical body. #198 is now merged and that second spelling is gone: +/// `DeepSeekV2RotaryEmbedding::{forward, forward_inverse_tail}` call this +/// function directly, so the two cannot drift apart. +/// +/// Two details #198 established that this doc did not previously carry: +/// +/// * The uniformity is not an accident of the current scheduler, it is +/// *enforced* by it. `scheduler::default_scheduler::select_running_bucket` +/// admits a forward pass only when cache lengths are exactly equal +/// ("Sequences may only share a forward pass when their cache lengths are +/// exactly equal"), and the single producer of a ragged dense batch is gated +/// behind `ARC_MTP_PER_SEQ_KV`, which defaults off +/// (`models::deepseek4::ragged_row_q0`). +/// * Measured at B=256 on a real decode profile, the per-sequence loop cost +/// ~100,000 GPU launches per step, 19.9% of step time -- of which the +/// `copy2d` half existed only to undo a split the same function had just +/// made. That number was measured on #198's branch, not on this tree. #[inline] pub(crate) fn uniform_seqlen_offset(seqlen_offsets: &[usize]) -> Option { match seqlen_offsets { @@ -418,14 +433,105 @@ pub mod rope_cohort_stats { #[inline] pub(super) fn record_cohort() { COHORT.fetch_add(1, Ordering::Relaxed); + #[cfg(test)] + bump_local(1, 0); } #[inline] pub(super) fn record_per_sequence() { PER_SEQUENCE.fetch_add(1, Ordering::Relaxed); + #[cfg(test)] + bump_local(0, 1); + } + + // ---- Test-only per-thread mirror of the two counters above. ----------- + // + // The atomics are process-global and, since #198 collapsed its own private + // copy of this module onto this one, they are written by EVERY rotary in + // this file -- `RotaryEmbedding`, `Llama3RotaryEmbedding`, + // `Phi4MMRotaryEmbedding` and `DeepSeekV2RotaryEmbedding` alike. A test + // that brackets one call with `counts()` is therefore not measuring its own + // call: `cargo test` runs the suite on many threads, and any other test + // that drives any rotary lands between the `before` and the `after`. + // + // That is not hypothetical. Two rotary forwards in `deepseek4.rs` + // (`standard_layer_rope_is_unscaled_and_compressed_is_not`, + // `mtp_block_takes_the_standard_unscaled_rope_table`) race the cohort tests + // exactly this way -- caught by merging #198 and #121 together and running + // the full suite, where `--test-threads=1` passes 707 and the parallel run + // fails one. + // + // A mutex was the first fix and it is the wrong one: it obliges every + // present and future test that touches any rotary to know about a lock in + // another module, which nothing enforces and the next such test will + // forget. Counting per-thread instead makes each test's delta its own by + // construction, needs no cooperation from anyone else, and leaves the + // serving instrument -- the global `counts()` -- byte-for-byte unchanged. + + #[cfg(test)] + thread_local! { + static LOCAL: std::cell::Cell<(u64, u64)> = const { std::cell::Cell::new((0, 0)) }; + } + + #[cfg(test)] + #[inline] + fn bump_local(cohort: u64, per_sequence: u64) { + LOCAL.with(|c| { + let (a, b) = c.get(); + c.set((a + cohort, b + per_sequence)); + }); + } + + /// `(cohort, per_sequence)` for THIS thread only โ€” the test-facing + /// counterpart of [`counts`]. + /// + /// Same monotonic-delta contract, same meaning, but immune to whatever the + /// rest of the suite is doing on other threads. Tests assert on deltas + /// around their own call; nothing resets it. + #[cfg(test)] + pub fn local_counts() -> (u64, u64) { + LOCAL.with(|c| c.get()) } } +/// `ARC_ROPE_COHORT=1` routes a **ragged** decode cohort โ€” rows at genuinely +/// different positions โ€” through ONE gathered `rope_i` per projection instead +/// of the per-sequence loop. **Default OFF**; unset and every off-spelling +/// keep the loop byte-for-byte as it is today. +/// +/// Uniform cohorts are NOT behind this flag: when every row shares an offset, +/// [`uniform_seqlen_offset`] already dispatches one batched `rope_i` with a +/// 2-D table, which is bit-identical to the loop and needs no gather. This +/// flag governs only the arm the uniform fix could not reach โ€” offsets that +/// genuinely differ, where the loop was the only spelling. On the measured +/// B=256 profile that loop is ~9 ops ร— 43 layers ร— B โ‰ˆ 99,000 launches per +/// step (19.9% of step time); the gathered path is ~5 launches per layer. +/// +/// candle's `rope_i` accepts a rank-3 `[B, T, D/2]` cos/sin when `B` matches +/// the input batch (`rope_check_cs`, `candle-nn/src/rotary_emb.rs`; the CUDA +/// wrapper sets `stride_b`, the CPU path indexes `i + b_i * t * d / 2`), so a +/// per-row table is expressible in one launch โ€” the loop was never structural. +/// +/// Why opt-in rather than default-ON: the same doctrine as +/// `cuda::qk_norm_rope::fused_cohort_enabled` โ€” unverified means default-off, +/// and unverified means unmeasured on hardware, not new. The A/B box flips +/// this flag; when the gathered arm has been measured correct and faster, the +/// default flips in the same change that records the number. +/// +/// Read by VALUE (`mistralrs_quant::env_flag_is_set`), never by presence: +/// `ARC_ROPE_COHORT=0` means OFF (#212 converted 23 flags because +/// `var_os(..).is_some()` made `=0` mean ON). Latched once per process โ€” the +/// dispatch runs on every layer of every step and `std::env::var` takes a +/// global lock. The polarity table lives with `env_flag_value`'s own tests; +/// the two ragged arms themselves are compared bit-for-bit in +/// `deepseek_ragged_cohort_tests`, which drives them directly so the latch +/// cannot make the parity claim untestable. +pub(crate) fn rope_cohort_enabled() -> bool { + use std::sync::OnceLock; + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| mistralrs_quant::env_flag_is_set("ARC_ROPE_COHORT")) +} + /// Tile a `[max_seq_len, rot_dim]` rotary table into the `[b * seq_len, rot_dim]` /// layout `mistralrs_quant::rotary::apply_rotary_inplace` requires. /// @@ -1540,8 +1646,52 @@ impl Qwen2_5VLRotaryEmbedding { pub struct DeepSeekV2RotaryEmbedding { sin: Tensor, cos: Tensor, + /// `-sin`, materialised once at construction. + /// + /// [`Self::forward_inverse_tail`] rotates by the conjugate `(cos, -sin)`. + /// It used to spell that as `self.sin.narrow(..)?.neg()?` *inside* the + /// per-sequence loop, which at B=256 ร— 43 layers launched `uneg_bf16` + /// **11,008 times per decode step** to negate a `[1, 32]` slice of a + /// **constant** table โ€” 0.8% of a measured B=256 step spent re-deriving a + /// value that never changes. + /// + /// Narrowing a pre-negated table is bit-identical to negating a narrowed + /// one: negation is a sign-bit flip, exact in every float format and + /// elementwise, so it commutes with a view. See `neg_sin_commutes_with_narrow`. + neg_sin: Tensor, + /// Memoised per-row table gather for the ragged-cohort path + /// ([`rope_cohort_enabled`]). Behind an `Arc` so the per-device sharing in + /// `deepseek4.rs` (one rotary per device per table kind, cloned into every + /// layer on that device) shares ONE cache: the offsets are host data, so + /// one small H2D per step per device is the floor, and this cache holds it + /// there instead of paying it again on all 43 layers ร— 3 rope calls + /// (CLAUDE.md pitfall #5 โ€” no `Tensor::from_vec` per layer on the hot + /// path). `None` until the gated path first runs; the flag-off paths never + /// touch it. + cohort_gather: Arc>>, +} + +/// One step's gathered `[B, T, rope_dim / 2]` tables for a ragged cohort, +/// keyed by the exact `(seqlen_offsets, seq_len)` that built them. All three +/// tables are gathered together: `forward` reads (cos, sin), +/// `forward_inverse_tail` reads (cos, neg_sin), and they share the step. +#[derive(Debug)] +struct CohortGather { + offsets: Vec, + seq_len: usize, + /// `[B, T, rope_dim / 2]` each, contiguous, on the tables' device. + cos: Tensor, + sin: Tensor, + neg_sin: Tensor, } +// The engagement counters for this path are the ones declared beside +// `uniform_seqlen_offset` (`rope_cohort_stats`, above). #198 carried a second +// module of the same name, with the same two atomics and a `counts()` of its +// own; two `pub mod rope_cohort_stats` in one file is a hard name collision, +// and two counters for one claim would let a harness read `cohort > 0` off the +// Phi path while this one never engaged. Collapsed onto the shared module, as +// `uniform_seqlen_offset`'s doc comment asks. #[derive(Debug, Clone, Deserialize, Serialize)] #[serde(untagged)] pub enum DeepSeekV2RopeScaling { @@ -1583,6 +1733,183 @@ pub struct DeepSeekV2RopeConfig { } impl DeepSeekV2RotaryEmbedding { + /// The single place the tables are stored, so `neg_sin` cannot be derived + /// in one constructor and forgotten in the other. + fn from_tables(sin: Tensor, cos: Tensor) -> Result { + let neg_sin = sin.neg()?; + Ok(Self { + sin, + cos, + neg_sin, + cohort_gather: Arc::new(std::sync::Mutex::new(None)), + }) + } + + /// The step's `(cos, sin, neg_sin)` โ€” each `[B, T, rope_dim / 2]` โ€” for a + /// ragged cohort, one table row per `(b, t)` at absolute position + /// `seqlen_offsets[b] + t`. + /// + /// The row indices are host data (`seqlen_offsets` is a host slice), so + /// one H2D copy per step is the floor; everything after it โ€” the three + /// `index_select`s included โ€” happens once per step per device and is + /// served from [`Self::cohort_gather`] to the other rope calls of the same + /// step (all layers on this device share one instance of `self`; see the + /// field's doc). `index_select` output is contiguous and `reshape` of a + /// contiguous tensor is a view, so the returned tables satisfy `rope_i`'s + /// contiguity requirement with no extra copy. + fn cohort_tables( + &self, + seqlen_offsets: &[usize], + seq_len: usize, + ) -> Result<(Tensor, Tensor, Tensor)> { + let mut cache = self + .cohort_gather + .lock() + .expect("cohort gather cache poisoned"); + if let Some(entry) = cache.as_ref() { + if entry.seq_len == seq_len && entry.offsets == seqlen_offsets { + return Ok((entry.cos.clone(), entry.sin.clone(), entry.neg_sin.clone())); + } + } + // Host-side bounds guard, BEFORE any index reaches `index_select`. The + // loop arm's `narrow(0, offset, seq_len)` hard-errors past the table + // end; `index_select` with an out-of-range index is not reliably + // bounds-checked on the CUDA backend and can read garbage SILENTLY. + // Both arms must fail loud, and identically-loud โ€” the offsets are + // host data, so the check is free. + let table_rows = self.cos.dim(0)?; + if let Some(&worst) = seqlen_offsets.iter().max() { + if worst + seq_len > table_rows { + candle_core::bail!( + "rope cohort gather out of bounds: offset {worst} + seq_len {seq_len} \ + exceeds the {table_rows}-row cos/sin table (offsets {seqlen_offsets:?})" + ); + } + } + let b = seqlen_offsets.len(); + let mut rows: Vec = Vec::with_capacity(b * seq_len); + for &off in seqlen_offsets { + for t in 0..seq_len { + rows.push((off + t) as u32); + } + } + let idx = Tensor::from_vec(rows, b * seq_len, self.cos.device())?; + let half = self.cos.dim(D::Minus1)?; + let cos = self + .cos + .index_select(&idx, 0)? + .reshape((b, seq_len, half))?; + let sin = self + .sin + .index_select(&idx, 0)? + .reshape((b, seq_len, half))?; + let neg_sin = self + .neg_sin + .index_select(&idx, 0)? + .reshape((b, seq_len, half))?; + *cache = Some(CohortGather { + offsets: seqlen_offsets.to_vec(), + seq_len, + cos: cos.clone(), + sin: sin.clone(), + neg_sin: neg_sin.clone(), + }); + Ok((cos, sin, neg_sin)) + } + + /// The gated ragged arm of [`Self::forward`]: ONE `rope_i` per projection + /// over the whole cohort, with the per-row rotation expressed through a + /// rank-3 `[B, T, D/2]` cos/sin (see [`rope_cohort_enabled`] for why that + /// is expressible at all). Bit-identical to [`Self::forward_ragged_loop`]: + /// the gather copies table rows verbatim, and each output element is the + /// same two-multiply-one-add of the same inputs โ€” batching changes which + /// launch computes an element, never the arithmetic. Asserted bit-for-bit + /// in `deepseek_ragged_cohort_tests`. + fn forward_ragged_cohort( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offsets: &[usize], + seq_len: usize, + ) -> Result<(Tensor, Tensor)> { + let (cos, sin, _neg_sin) = self.cohort_tables(seqlen_offsets, seq_len)?; + let q_embed = candle_nn::rotary_emb::rope_i(&q.contiguous()?, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope_i(&k.contiguous()?, &cos, &sin)?; + Ok((q_embed, k_embed)) + } + + /// The per-sequence loop, extracted verbatim so both [`Self::forward`] + /// arms name the thing they dispatch between and the parity tests can + /// drive each directly whatever `ARC_ROPE_COHORT` latched to. + fn forward_ragged_loop( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offsets: &[usize], + seq_len: usize, + ) -> Result<(Tensor, Tensor)> { + let mut q_embeds = Vec::new(); + let mut k_embeds = Vec::new(); + for (i, offset) in seqlen_offsets.iter().enumerate() { + let cos = self.cos.narrow(0, *offset, seq_len)?; + let sin = self.sin.narrow(0, *offset, seq_len)?; + let q_embed = + candle_nn::rotary_emb::rope_i(&q.i(i)?.unsqueeze(0)?.contiguous()?, &cos, &sin)?; + let k_embed = + candle_nn::rotary_emb::rope_i(&k.i(i)?.unsqueeze(0)?.contiguous()?, &cos, &sin)?; + q_embeds.push(q_embed); + k_embeds.push(k_embed); + } + Ok((Tensor::cat(&q_embeds, 0)?, Tensor::cat(&k_embeds, 0)?)) + } + + /// Gated ragged arm of [`Self::forward_inverse_tail`]: one `rope_i` over + /// the whole cohort, conjugate rotation via the pre-negated table gathered + /// per row. Returns only the rotated tail; the caller re-attaches the NoPE + /// dims (shared with the loop arm, so the `cat` is spelled once). + fn inverse_tail_ragged_cohort( + &self, + x: &Tensor, + rope_dim: usize, + nope: usize, + seqlen_offsets: &[usize], + seq_len: usize, + ) -> Result { + let (cos, _sin, neg_sin) = self.cohort_tables(seqlen_offsets, seq_len)?; + let x_pe = x.narrow(3, nope, rope_dim)?.contiguous()?; + candle_nn::rotary_emb::rope_i(&x_pe, &cos, &neg_sin) + } + + /// Per-sequence loop arm of [`Self::forward_inverse_tail`], extracted + /// verbatim (minus the trailing NoPE `cat`, which the caller owns). + fn inverse_tail_ragged_loop( + &self, + x: &Tensor, + rope_dim: usize, + nope: usize, + seqlen_offsets: &[usize], + seq_len: usize, + ) -> Result { + let mut outs = Vec::new(); + for (i, offset) in seqlen_offsets.iter().enumerate() { + let cos = self.cos.narrow(0, *offset, seq_len)?; + let sin = self.neg_sin.narrow(0, *offset, seq_len)?; + let x_pe = x + .i(i)? + .unsqueeze(0)? + .narrow(3, nope, rope_dim)? + .contiguous()?; + outs.push(candle_nn::rotary_emb::rope_i(&x_pe, &cos, &sin)?); + } + Tensor::cat(&outs, 0) + } + + // The cohort predicate for this rotary is the free function + // `uniform_seqlen_offset`, not a method here. #198 introduced + // `Self::uniform_offset` with a byte-identical body; per that function's + // doc comment ("there must not be two spellings of the same predicate") + // the duplicate is collapsed away, and the reasoning it carried has moved + // onto the survivor. fn new_unscaled(cfg: &DeepSeekV2RopeConfig, dtype: DType, dev: &Device) -> Result { let max_seq_len = cfg.max_position_embeddings; let dim = cfg.qk_rope_head_dim; @@ -1601,7 +1928,7 @@ impl DeepSeekV2RotaryEmbedding { let sin = freqs.sin()?.to_dtype(dtype)?; let cos = freqs.cos()?.to_dtype(dtype)?; - Ok(Self { sin, cos }) + Self::from_tables(sin, cos) } fn yarn_find_correction_dim( @@ -1693,7 +2020,7 @@ impl DeepSeekV2RotaryEmbedding { let sin = (freqs.sin()? * mscale as f64)?.to_dtype(dtype)?; let cos = (freqs.cos()? * mscale as f64)?.to_dtype(dtype)?; - Ok(Self { sin, cos }) + Self::from_tables(sin, cos) } pub fn new(cfg: &DeepSeekV2RopeConfig, dtype: DType, dev: &Device) -> Result { @@ -1746,32 +2073,32 @@ impl DeepSeekV2RotaryEmbedding { ) -> Result<(Tensor, Tensor)> { let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; - if seqlen_offsets.len() == 1 { - let cos = self.cos.narrow(0, seqlen_offsets[0], seq_len)?; - let sin = self.sin.narrow(0, seqlen_offsets[0], seq_len)?; + // One `rope_i` over the whole cohort whenever every row rotates at the + // same position โ€” which is every request today. See + // [`uniform_seqlen_offset`] for why this used to read `len() == 1`, and + // what that cost. + // + // Bit-identical to the loop below, not merely close: with a 2-D + // `[T, D/2]` cos/sin the kernel's `stride_b` is 0 + // (`candle-nn/src/rotary_emb.rs`), so every batch row reads the same + // cos/sin row, and each output element is an independent + // two-multiply-one-add of its own inputs. Batching changes which + // launch computes an element, never the arithmetic that produces it. + if let Some(offset) = uniform_seqlen_offset(seqlen_offsets) { + rope_cohort_stats::record_cohort(); + let cos = self.cos.narrow(0, offset, seq_len)?; + let sin = self.sin.narrow(0, offset, seq_len)?; let q_embed = candle_nn::rotary_emb::rope_i(&q.contiguous()?, &cos, &sin)?; let k_embed = candle_nn::rotary_emb::rope_i(&k.contiguous()?, &cos, &sin)?; Ok((q_embed, k_embed)) + } else if rope_cohort_enabled() { + // Ragged rows, one gathered rope_i per projection โ€” opt-in, see + // [`rope_cohort_enabled`]. + rope_cohort_stats::record_cohort(); + self.forward_ragged_cohort(q, k, seqlen_offsets, seq_len) } else { - let mut q_embeds = Vec::new(); - let mut k_embeds = Vec::new(); - for (i, offset) in seqlen_offsets.iter().enumerate() { - let cos = self.cos.narrow(0, *offset, seq_len)?; - let sin = self.sin.narrow(0, *offset, seq_len)?; - let q_embed = candle_nn::rotary_emb::rope_i( - &q.i(i)?.unsqueeze(0)?.contiguous()?, - &cos, - &sin, - )?; - let k_embed = candle_nn::rotary_emb::rope_i( - &k.i(i)?.unsqueeze(0)?.contiguous()?, - &cos, - &sin, - )?; - q_embeds.push(q_embed); - k_embeds.push(k_embed); - } - Ok((Tensor::cat(&q_embeds, 0)?, Tensor::cat(&k_embeds, 0)?)) + rope_cohort_stats::record_per_sequence(); + self.forward_ragged_loop(q, k, seqlen_offsets, seq_len) } } @@ -1789,24 +2116,23 @@ impl DeepSeekV2RotaryEmbedding { let (_b, _h, seq_len, head_dim) = x.dims4()?; let nope = head_dim - rope_dim; let x_nope = x.narrow(3, 0, nope)?; - let rotated = if seqlen_offsets.len() == 1 { - let cos = self.cos.narrow(0, seqlen_offsets[0], seq_len)?; - let sin = self.sin.narrow(0, seqlen_offsets[0], seq_len)?.neg()?; + // Same cohort dispatch as `forward`, and the conjugate `sin` now comes + // from the pre-negated table (see the `neg_sin` field) rather than a + // `neg` launch per sequence per layer. + let rotated = if let Some(offset) = uniform_seqlen_offset(seqlen_offsets) { + rope_cohort_stats::record_cohort(); + let cos = self.cos.narrow(0, offset, seq_len)?; + let sin = self.neg_sin.narrow(0, offset, seq_len)?; let x_pe = x.narrow(3, nope, rope_dim)?.contiguous()?; candle_nn::rotary_emb::rope_i(&x_pe, &cos, &sin)? + } else if rope_cohort_enabled() { + // Ragged rows, one gathered rope_i โ€” opt-in, see + // [`rope_cohort_enabled`]. + rope_cohort_stats::record_cohort(); + self.inverse_tail_ragged_cohort(x, rope_dim, nope, seqlen_offsets, seq_len)? } else { - let mut outs = Vec::new(); - for (i, offset) in seqlen_offsets.iter().enumerate() { - let cos = self.cos.narrow(0, *offset, seq_len)?; - let sin = self.sin.narrow(0, *offset, seq_len)?.neg()?; - let x_pe = x - .i(i)? - .unsqueeze(0)? - .narrow(3, nope, rope_dim)? - .contiguous()?; - outs.push(candle_nn::rotary_emb::rope_i(&x_pe, &cos, &sin)?); - } - Tensor::cat(&outs, 0)? + rope_cohort_stats::record_per_sequence(); + self.inverse_tail_ragged_loop(x, rope_dim, nope, seqlen_offsets, seq_len)? }; Tensor::cat(&[&x_nope, &rotated], 3)?.contiguous() } @@ -4110,9 +4436,9 @@ mod rope_cohort_tests { let k = xs(&[B, H, T, D_HEAD], 0.019); let offsets = vec![5usize; B]; - let before = rope_cohort_stats::counts(); + let before = rope_cohort_stats::local_counts(); let (qb, kb) = re.forward(&q, &k, &offsets).unwrap(); - let after = rope_cohort_stats::counts(); + let after = rope_cohort_stats::local_counts(); // A green result must prove work happened: the batched arm has to have // been the one that ran. @@ -4143,9 +4469,9 @@ mod rope_cohort_tests { let k = xs(&[3, H, T, D_HEAD], 0.019); let offsets = vec![0usize, 4, 9]; - let before = rope_cohort_stats::counts(); + let before = rope_cohort_stats::local_counts(); let (qb, kb) = re.forward(&q, &k, &offsets).unwrap(); - let after = rope_cohort_stats::counts(); + let after = rope_cohort_stats::local_counts(); assert!( after.1 > before.1, "ragged cohort must NOT take the batched path" @@ -4243,3 +4569,507 @@ mod rope_cohort_tests { } } } + +/// Bit-exactness and engagement tests for the DeepSeek V2/V4 RoPE cohort path. +/// +/// The change under test replaced `seqlen_offsets.len() == 1` โ€” a test of the +/// *length* of the offset vector โ€” with a test of the *distinctness of its +/// values*. At B=256 the length is 256, so the batched path was unreachable at +/// every batch size above one, while the values were uniform anyway; the +/// per-sequence loop was therefore performing B bit-identical recomputations +/// and concatenating them back together (~100,000 GPU launches per decode step +/// on a measured profile, 19.9% of step time). +/// +/// Each test below holds the old loop as a verbatim oracle and asserts the new +/// dispatch reproduces it **bit for bit**, not within a tolerance โ€” a +/// tolerance-based assertion here would be vacuous, since the claim is that no +/// arithmetic changed at all. +#[cfg(test)] +mod deepseek_rope_cohort_tests { + use super::*; + + const MAX_POS: usize = 32; + const ROPE_DIM: usize = 8; + const HEAD_DIM: usize = 12; + const N_HEADS: usize = 3; + const BATCH: usize = 4; + + // Engagement is asserted through `rope_cohort_stats::local_counts()`, the + // per-thread mirror of the global counters. #198 originally serialised + // these tests on a private mutex, which worked only while this module also + // carried its own private copy of `rope_cohort_stats`; once the two counter + // modules were collapsed the atomics became shared and a private lock + // stopped serialising anything that mattered. A shared lock was the next + // attempt and it is still wrong -- it obliges unrelated tests in other + // modules to know about it. See `rope_cohort_stats` for why per-thread + // counting is the fix. + + fn rope() -> DeepSeekV2RotaryEmbedding { + let cfg = DeepSeekV2RopeConfig { + rope_scaling: None, + max_position_embeddings: MAX_POS, + rope_theta: 10000.0, + qk_rope_head_dim: ROPE_DIM, + }; + DeepSeekV2RotaryEmbedding::new(&cfg, DType::F32, &Device::Cpu).unwrap() + } + + /// Deterministic, non-degenerate input. Both paths consume the same tensor, + /// so randomness would be sound too, but a fixed pattern makes a failure + /// reproducible. + fn xs(b: usize, h: usize, t: usize, d: usize) -> Tensor { + let n = b * h * t * d; + let v: Vec = (0..n) + .map(|i| ((i as f32) * 0.37).sin() * 1.7 + 0.11) + .collect(); + Tensor::from_vec(v, (b, h, t, d), &Device::Cpu).unwrap() + } + + /// Raw IEEE-754 bit patterns. Float equality would accept `-0.0 == 0.0` and + /// silently pass a sign error in the conjugate `sin`; bit patterns do not. + fn bits(t: &Tensor) -> Vec { + t.flatten_all() + .unwrap() + .to_vec1::() + .unwrap() + .into_iter() + .map(f32::to_bits) + .collect() + } + + /// The pre-fix `forward`, verbatim, as the oracle. Always loops, whatever + /// the offsets are. + fn reference_forward( + re: &DeepSeekV2RotaryEmbedding, + q: &Tensor, + k: &Tensor, + offsets: &[usize], + ) -> Result<(Tensor, Tensor)> { + let (_b, _h, seq_len, _n) = q.dims4()?; + let mut q_embeds = Vec::new(); + let mut k_embeds = Vec::new(); + for (i, offset) in offsets.iter().enumerate() { + let cos = re.cos.narrow(0, *offset, seq_len)?; + let sin = re.sin.narrow(0, *offset, seq_len)?; + q_embeds.push(candle_nn::rotary_emb::rope_i( + &q.i(i)?.unsqueeze(0)?.contiguous()?, + &cos, + &sin, + )?); + k_embeds.push(candle_nn::rotary_emb::rope_i( + &k.i(i)?.unsqueeze(0)?.contiguous()?, + &cos, + &sin, + )?); + } + Ok((Tensor::cat(&q_embeds, 0)?, Tensor::cat(&k_embeds, 0)?)) + } + + /// The pre-fix `forward_inverse_tail`, verbatim โ€” including the old + /// `sin.narrow(..).neg()` spelling, so this oracle also pins the `neg_sin` + /// precomputation. + fn reference_inverse_tail( + re: &DeepSeekV2RotaryEmbedding, + x: &Tensor, + rope_dim: usize, + offsets: &[usize], + ) -> Result { + let (_b, _h, seq_len, head_dim) = x.dims4()?; + let nope = head_dim - rope_dim; + let x_nope = x.narrow(3, 0, nope)?; + let mut outs = Vec::new(); + for (i, offset) in offsets.iter().enumerate() { + let cos = re.cos.narrow(0, *offset, seq_len)?; + let sin = re.sin.narrow(0, *offset, seq_len)?.neg()?; + let x_pe = x + .i(i)? + .unsqueeze(0)? + .narrow(3, nope, rope_dim)? + .contiguous()?; + outs.push(candle_nn::rotary_emb::rope_i(&x_pe, &cos, &sin)?); + } + let rotated = Tensor::cat(&outs, 0)?; + Tensor::cat(&[&x_nope, &rotated], 3)?.contiguous() + } + + // The predicate itself is tested once, in `rope_cohort_tests` + // (`uniform_offset_tests_values_not_length`), against the same six cases. + // #198 carried a byte-identical copy of that test aimed at + // `DeepSeekV2RotaryEmbedding::uniform_offset`; with the two spellings + // collapsed onto `uniform_seqlen_offset` there is one predicate, so a + // second copy would assert nothing the first does not. What is NOT + // duplicated, and stays below, is the V4-specific behaviour: `neg_sin`, + // `forward_inverse_tail`, and the bit-exactness of this rotary's own + // cohort dispatch. + #[test] + fn neg_sin_commutes_with_narrow() { + let re = rope(); + for offset in [0usize, 1, 17] { + let precomputed = re.neg_sin.narrow(0, offset, 3).unwrap(); + let on_the_fly = re.sin.narrow(0, offset, 3).unwrap().neg().unwrap(); + assert_eq!( + bits(&precomputed), + bits(&on_the_fly), + "pre-negating the table must be bit-identical at offset {offset}" + ); + } + } + + #[test] + fn cohort_forward_is_bit_identical_to_the_per_sequence_loop() { + let re = rope(); + // seq_len 1 is decode; seq_len 2 covers uniform-offset chunked prefill. + for seq_len in [1usize, 2] { + let q = xs(BATCH, N_HEADS, seq_len, ROPE_DIM); + let k = xs(BATCH, 1, seq_len, ROPE_DIM); + let offsets = vec![5usize; BATCH]; + + let before = rope_cohort_stats::local_counts(); + let (q_new, k_new) = re.forward(&q, &k, &offsets).unwrap(); + let after = rope_cohort_stats::local_counts(); + let (q_ref, k_ref) = reference_forward(&re, &q, &k, &offsets).unwrap(); + + assert_eq!(q_new.dims(), q_ref.dims()); + assert_eq!(bits(&q_new), bits(&q_ref), "Q differs at seq_len={seq_len}"); + assert_eq!(bits(&k_new), bits(&k_ref), "K differs at seq_len={seq_len}"); + // Prove the batched branch is the one that ran: a fast path that + // silently declined would pass the equality above trivially. + assert_eq!( + (after.0 - before.0, after.1 - before.1), + (1, 0), + "expected exactly one cohort call and no per-sequence loop" + ); + } + } + + #[test] + fn ragged_forward_still_takes_the_loop_and_still_matches() { + let re = rope(); + let q = xs(BATCH, N_HEADS, 1, ROPE_DIM); + let k = xs(BATCH, 1, 1, ROPE_DIM); + let offsets = vec![5usize, 6, 5, 9]; + + let before = rope_cohort_stats::local_counts(); + let (q_new, k_new) = re.forward(&q, &k, &offsets).unwrap(); + let after = rope_cohort_stats::local_counts(); + let (q_ref, k_ref) = reference_forward(&re, &q, &k, &offsets).unwrap(); + + assert_eq!(bits(&q_new), bits(&q_ref)); + assert_eq!(bits(&k_new), bits(&k_ref)); + assert_eq!( + (after.0 - before.0, after.1 - before.1), + (0, 1), + "ragged rows must still take the per-sequence loop" + ); + } + + #[test] + fn cohort_inverse_tail_is_bit_identical_to_the_per_sequence_loop() { + let re = rope(); + for seq_len in [1usize, 2] { + let x = xs(BATCH, N_HEADS, seq_len, HEAD_DIM); + let offsets = vec![5usize; BATCH]; + + let before = rope_cohort_stats::local_counts(); + let new = re.forward_inverse_tail(&x, ROPE_DIM, &offsets).unwrap(); + let after = rope_cohort_stats::local_counts(); + let reference = reference_inverse_tail(&re, &x, ROPE_DIM, &offsets).unwrap(); + + assert_eq!(new.dims(), reference.dims()); + assert_eq!( + bits(&new), + bits(&reference), + "inverse tail differs at seq_len={seq_len}" + ); + assert_eq!((after.0 - before.0, after.1 - before.1), (1, 0)); + } + } + + #[test] + fn ragged_inverse_tail_still_takes_the_loop_and_still_matches() { + let re = rope(); + let x = xs(BATCH, N_HEADS, 1, HEAD_DIM); + let offsets = vec![5usize, 6, 5, 9]; + + let before = rope_cohort_stats::local_counts(); + let new = re.forward_inverse_tail(&x, ROPE_DIM, &offsets).unwrap(); + let after = rope_cohort_stats::local_counts(); + let reference = reference_inverse_tail(&re, &x, ROPE_DIM, &offsets).unwrap(); + + assert_eq!(bits(&new), bits(&reference)); + assert_eq!((after.0 - before.0, after.1 - before.1), (0, 1)); + } + + /// The comparator must be able to fail. Every bit-equality assertion above + /// is worthless if `bits` had degenerated into something that always + /// matches, so poison one element by a single ULP and require it flagged. + #[test] + fn the_bit_comparator_is_live() { + // Drives a forward, so it must serialise with the counter assertions. + let re = rope(); + let x = xs(BATCH, N_HEADS, 1, HEAD_DIM); + let offsets = vec![5usize; BATCH]; + let good = re.forward_inverse_tail(&x, ROPE_DIM, &offsets).unwrap(); + + let mut poisoned = bits(&good); + poisoned[0] += 1; // one ULP + assert_ne!( + bits(&good), + poisoned, + "the comparator accepted a one-ULP difference; every other \ + assertion in this module would be vacuous" + ); + } +} + +/// Parity for the **gated** ragged-cohort arms ([`rope_cohort_enabled`], +/// `ARC_ROPE_COHORT`): one gathered rank-3 `rope_i` against the per-sequence +/// loop it replaces, bit for bit. +/// +/// The arms are driven **directly**, not through the public dispatch: the flag +/// latches once per process (`OnceLock`), so a test that reached the gathered +/// arm by mutating the environment would prove only whatever the first reader +/// happened to see โ€” the same reasoning as `env_flag_value`'s own tests. The +/// dispatch's default (flag unset โ‡’ ragged rows keep the loop) is pinned by +/// `deepseek_rope_cohort_tests::ragged_forward_still_takes_the_loop_and_still_matches`, +/// which runs the public `forward` in exactly that environment. +#[cfg(test)] +mod deepseek_ragged_cohort_tests { + use super::*; + + const MAX_POS: usize = 32; + const ROPE_DIM: usize = 8; + const HEAD_DIM: usize = 12; + const N_HEADS: usize = 3; + + fn rope() -> DeepSeekV2RotaryEmbedding { + let cfg = DeepSeekV2RopeConfig { + rope_scaling: None, + max_position_embeddings: MAX_POS, + rope_theta: 10000.0, + qk_rope_head_dim: ROPE_DIM, + }; + DeepSeekV2RotaryEmbedding::new(&cfg, DType::F32, &Device::Cpu).unwrap() + } + + fn xs(b: usize, h: usize, t: usize, d: usize) -> Tensor { + let n = b * h * t * d; + let v: Vec = (0..n) + .map(|i| ((i as f32) * 0.29).sin() * 1.3 + 0.07) + .collect(); + Tensor::from_vec(v, (b, h, t, d), &Device::Cpu).unwrap() + } + + /// Raw IEEE-754 bit patterns โ€” the claim is that no arithmetic changed at + /// all, so a tolerance would be vacuous and `-0.0 == 0.0` would slip a + /// sign error through. + fn bits(t: &Tensor) -> Vec { + t.flatten_all() + .unwrap() + .to_vec1::() + .unwrap() + .into_iter() + .map(f32::to_bits) + .collect() + } + + /// Mixed offsets, B โˆˆ {1, 3}, decode (T=1) and uniform-length chunk (T=2): + /// the gathered arm must reproduce the loop bit for bit. + #[test] + fn ragged_cohort_forward_matches_the_loop_bitwise() { + let re = rope(); + for (offsets, t) in [ + (vec![0usize, 4, 9], 1usize), + (vec![0usize, 4, 9], 2), + (vec![9usize, 4, 0], 1), // descending โ€” order must be per-row + (vec![7usize], 1), // B = 1 through the same arms + ] { + let b = offsets.len(); + let q = xs(b, N_HEADS, t, ROPE_DIM); + let k = xs(b, 1, t, ROPE_DIM); + + let (q_c, k_c) = re.forward_ragged_cohort(&q, &k, &offsets, t).unwrap(); + let (q_l, k_l) = re.forward_ragged_loop(&q, &k, &offsets, t).unwrap(); + + assert_eq!(q_c.dims(), q_l.dims()); + assert_eq!( + bits(&q_c), + bits(&q_l), + "gathered Q differs from the loop for offsets {offsets:?}, T={t}" + ); + assert_eq!( + bits(&k_c), + bits(&k_l), + "gathered K differs from the loop for offsets {offsets:?}, T={t}" + ); + + // Non-vacuity: rows at different offsets must actually rotate + // differently, or the equality above would pass under a + // broadcast-one-row bug on BOTH sides. + if b > 1 { + let d01 = (q_c.i(0).unwrap() - q_c.i(1).unwrap()) + .unwrap() + .abs() + .unwrap() + .max_all() + .unwrap() + .to_scalar::() + .unwrap(); + assert!(d01 > 0.0, "rows at different offsets rotated identically"); + } + } + } + + /// Same parity for the conjugate (inverse-tail) rotation, and for the + /// fully assembled `[NoPE | rotated]` output. + #[test] + fn ragged_cohort_inverse_tail_matches_the_loop_bitwise() { + let re = rope(); + let nope = HEAD_DIM - ROPE_DIM; + for (offsets, t) in [ + (vec![0usize, 4, 9], 1usize), + (vec![0usize, 4, 9], 2), + (vec![7usize], 1), + ] { + let b = offsets.len(); + let x = xs(b, N_HEADS, t, HEAD_DIM); + + let tail_c = re + .inverse_tail_ragged_cohort(&x, ROPE_DIM, nope, &offsets, t) + .unwrap(); + let tail_l = re + .inverse_tail_ragged_loop(&x, ROPE_DIM, nope, &offsets, t) + .unwrap(); + assert_eq!(tail_c.dims(), tail_l.dims()); + assert_eq!( + bits(&tail_c), + bits(&tail_l), + "gathered inverse tail differs from the loop for offsets {offsets:?}, T={t}" + ); + + // The assembled output the model actually consumes. + let x_nope = x.narrow(3, 0, nope).unwrap(); + let full_c = Tensor::cat(&[&x_nope, &tail_c], 3) + .unwrap() + .contiguous() + .unwrap(); + let full_l = Tensor::cat(&[&x_nope, &tail_l], 3) + .unwrap() + .contiguous() + .unwrap(); + assert_eq!(bits(&full_c), bits(&full_l)); + } + } + + /// The gather memo must serve repeats AND rebuild when the cohort moves โ€” + /// a stale entry would rotate every layer of the next step at the previous + /// step's positions, which is exactly the class of wrong-but-plausible + /// output that never crashes. + #[test] + fn cohort_gather_rebuilds_when_the_cohort_moves() { + let re = rope(); + let t = 1usize; + + for offsets in [vec![0usize, 4, 9], vec![1usize, 5, 10], vec![2usize, 3, 4]] { + // Twice with the same key: the second call is the memo hit. + for pass in 0..2 { + let (cos, sin, neg_sin) = re.cohort_tables(&offsets, t).unwrap(); + assert_eq!(cos.dims(), &[offsets.len(), t, ROPE_DIM / 2]); + for (b_i, &off) in offsets.iter().enumerate() { + let want_cos = re.cos.narrow(0, off, t).unwrap(); + let want_sin = re.sin.narrow(0, off, t).unwrap(); + let want_neg = re.neg_sin.narrow(0, off, t).unwrap(); + assert_eq!( + bits(&cos.i(b_i).unwrap()), + bits(&want_cos), + "cos row {b_i} wrong for offsets {offsets:?} (pass {pass})" + ); + assert_eq!(bits(&sin.i(b_i).unwrap()), bits(&want_sin)); + assert_eq!(bits(&neg_sin.i(b_i).unwrap()), bits(&want_neg)); + } + } + } + } + + /// `ARC_ROPE_COHORT` is read by VALUE through `env_flag_is_set`, whose + /// polarity table (`=0`/unset โ‡’ OFF, only yes-spellings โ‡’ ON) is pinned in + /// `mistralrs_quant::env_flag`. What this pins HERE is the dispatch + /// default: with the flag unset โ€” the state of every CI environment this + /// runs in โ€” ragged rows must take the per-sequence loop, so the gathered + /// arm cannot go live by accident of a latch. + #[test] + fn ragged_dispatch_defaults_to_the_loop() { + assert!( + std::env::var_os("ARC_ROPE_COHORT").is_none(), + "test environment unexpectedly sets ARC_ROPE_COHORT; this test \ + pins the UNSET default and cannot do so with the flag present" + ); + assert!(!rope_cohort_enabled(), "unset must mean OFF"); + + let re = rope(); + let q = xs(3, N_HEADS, 1, ROPE_DIM); + let k = xs(3, 1, 1, ROPE_DIM); + let before = rope_cohort_stats::local_counts(); + re.forward(&q, &k, &[0usize, 4, 9]).unwrap(); + let after = rope_cohort_stats::local_counts(); + assert_eq!( + (after.0 - before.0, after.1 - before.1), + (0, 1), + "with ARC_ROPE_COHORT unset a ragged cohort must take the loop" + ); + } + + /// An offset past the table end must ERROR on the cohort arm exactly as it + /// does on the loop arm โ€” and it must be THIS module's host-side guard + /// that fires, not the backend. The loop's `narrow` fails loud everywhere, + /// but `index_select` with an out-of-range index is not reliably + /// bounds-checked on CUDA and can read garbage silently; the arm we A/B on + /// hardware must not be the one with the quiet failure mode (D18). The + /// message substring is pinned so a backend error passing this test by + /// coincidence is impossible. + #[test] + fn out_of_bounds_offset_fails_loud_on_both_arms() { + let re = rope(); + let t = 1usize; + // MAX_POS = 32 rows: offset 31 is the last valid row at T=1; 32 is + // past the end. The valid rows in the same cohort prove the guard + // reports the MAX, not merely offsets[0]. + let offsets = vec![5usize, 32]; + let q = xs(2, N_HEADS, t, ROPE_DIM); + let k = xs(2, 1, t, ROPE_DIM); + + // Loop arm: narrow(0, 32, 1) on a 32-row table hard-errors. + assert!( + re.forward_ragged_loop(&q, &k, &offsets, t).is_err(), + "the loop arm accepted an offset past the table end" + ); + + // Cohort arm: must also error, and with the guard's own message. + let err = re + .forward_ragged_cohort(&q, &k, &offsets, t) + .expect_err("the cohort arm accepted an offset past the table end"); + let msg = err.to_string(); + assert!( + msg.contains("rope cohort gather out of bounds"), + "the cohort arm rejected the offset, but not via the host-side \ + guard โ€” got: {msg}" + ); + assert!( + msg.contains("offset 32") && msg.contains("32-row"), + "the guard's message must name the offending offset and the table \ + size โ€” got: {msg}" + ); + + // Same guard on the inverse-tail arm (shared `cohort_tables`). + let x = xs(2, N_HEADS, t, HEAD_DIM); + let err = re + .inverse_tail_ragged_cohort(&x, ROPE_DIM, HEAD_DIM - ROPE_DIM, &offsets, t) + .expect_err("the inverse-tail cohort arm accepted an offset past the table end"); + assert!(err.to_string().contains("rope cohort gather out of bounds")); + + // The boundary itself is not off by one: the last valid row works. + let ok = vec![5usize, 31]; + re.forward_ragged_cohort(&q, &k, &ok, t) + .expect("offset 31 + seq_len 1 on a 32-row table is in bounds"); + } +} diff --git a/mistralrs-core/src/models/deepseek4.rs b/mistralrs-core/src/models/deepseek4.rs index 766a5f59c..5b007f92b 100644 --- a/mistralrs-core/src/models/deepseek4.rs +++ b/mistralrs-core/src/models/deepseek4.rs @@ -1287,9 +1287,21 @@ impl Attention { if cos.dtype() != candle_core::DType::BF16 || sin.dtype() != candle_core::DType::BF16 { decline!("RoPE tables are not BF16"); } - if seqlen_offsets.len() != 1 { + // The kernel needs ONE `pos_offset`, not a batch of one row. It reads + // the table at `pos_offset + t` where `t = blockIdx.y % seq_len` + // (`qk_norm_rope.cu`), a position that does not depend on `b`, and its + // grid is `(n_heads + 1, batch * seq_len)` with buffers sized + // `batch * ...` โ€” so it is already correct at any batch size whose rows + // share an offset. This used to read `seqlen_offsets.len() != 1`, which + // tested the length of the vector rather than the distinctness of its + // values and therefore declined at EVERY batch size above one, leaving + // the fused path dead in exactly the serving regime it was written for. + // Rows that genuinely differ still decline, and still take the eager + // chain. Prove the swap with `ARC_QK_VERIFY=1`, which bit-compares + // against the eager output at every layer. + let Some(uniform_offset) = crate::layers::uniform_seqlen_offset(seqlen_offsets) else { decline!("per-sequence position offsets"); - } + }; if self.num_kv_heads != 1 { decline!("more than one KV head"); } @@ -1318,11 +1330,21 @@ impl Attention { if qh != self.num_attention_heads * head_dim { decline!("q_proj output width is not n_heads * head_dim"); } + // The batched cohort is OPT-IN. Fixing the length-vs-values gate above + // makes this kernel live on a path it has never run on, and this repo + // has paid for exactly that twice this week (TCFRAG, and the fused-512 + // attention path that dropped the mask for four days). `batch == 1` is + // untouched โ€” already live, already exercised. See + // `qk::fused_cohort_enabled` for the doctrine and for how to turn the + // cohort path on once it has been measured on hardware. + if qb > 1 && !qk::fused_cohort_enabled() { + decline!("batched cohort is opt-in (set ARC_QK_FUSED_COHORT=1)"); + } match kv_normed.dims3() { Ok((kb, kt, kh)) if kb == qb && kt == qt && kh == head_dim => {} _ => decline!("kv_norm output shape does not match [B, T, head_dim]"), } - let pos_offset = seqlen_offsets[0]; + let pos_offset = uniform_offset; match cos.dims2() { Ok((rows, cols)) if cols == rope_dim / 2 && rows >= pos_offset + qt => {} _ => decline!("RoPE table does not cover [pos_offset, pos_offset + T)"), diff --git a/mistralrs-core/tests/capability_reachability.rs b/mistralrs-core/tests/capability_reachability.rs index c203fd761..35fad7178 100644 --- a/mistralrs-core/tests/capability_reachability.rs +++ b/mistralrs-core/tests/capability_reachability.rs @@ -170,6 +170,74 @@ static REGISTRY: &[Capability] = &[ }, status: Status::Live, }, + // โ”€โ”€ SHIPPED OFF, WITH A NAMED EXPERIMENT โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + // + // This entry is `Live` because the GATE is reached on every layer of every + // step โ€” that is what this file checks, and it must not go dark. What is + // switched off is the gate's DEFAULT ANSWER at `batch > 1`, which no + // reachability check can see. + // + // Recorded 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. + // + // The kernel is batch-correct by construction and the dispatch bug that + // kept it dead above batch 1 (`seqlen_offsets.len() != 1` โ€” the LENGTH of + // the vector where the property needed is the DISTINCTNESS of its values) + // is fixed. It is off only because fixing that gate makes a CUDA kernel + // live on a path it has never run on, and this repo has paid for that twice + // in one week: TCFRAG (default-ON, "UNVERIFIED ON HARDWARE โ€” NEVER RUN" in + // its own header, 63 GB held, a layer permanently broken, retired by #209) + // and the fused-512 attention path (four days silently dropping the + // attention mask at 12% agreement). + // + // flag ARC_QK_FUSED_COHORT=1 + // gate cuda/qk_norm_rope.rs :: fused_cohort_enabled_from + // test one binary, B>1 decode, same prompt and seed, flag set vs unset; + // ARC_QK_VERIFY=1 bit-compares fused against eager at every layer + // pass outputs bit-identical AND the fused arm faster; engagement + // counter (`qk::engaged_count`) non-zero, since a silently + // declining fast path reports a perfect no-op + // was never measured above batch 1 โ€” the path was unreachable, so there + // is no prior number to quote, honest or otherwise + // 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: "fused qk_norm_rope at B>1: built and tested, default OFF pending one measurement", + parent: "ArcInfer / ArcAttention", + check: Check::Symbol { + symbol: "fused_cohort_enabled", + defined_in: "mistralrs-core/src/cuda/qk_norm_rope.rs", + }, + status: Status::Live, + }, + // Same shape as the entry above, same doctrine, one rung lower in the + // stack: the V4 rotary's RAGGED arm. Uniform cohorts already take one + // batched `rope_i` unconditionally (that dispatch is exact and shipped + // ON); this flag governs rows at genuinely different positions, which the + // per-sequence loop served at ~9 ops ร— 43 layers ร— B per step. The + // gathered rank-3 path is bit-identical to the loop + // (`deepseek_ragged_cohort_tests`) and default OFF until measured. + // + // flag ARC_ROPE_COHORT=1 (read by value via env_flag_is_set) + // gate layers.rs :: rope_cohort_enabled + // test one binary, ragged B>1 decode (ARC_MTP_PER_SEQ_KV on, or any + // producer of unequal offsets), flag set vs unset + // pass outputs bit-identical AND the gathered arm faster; + // rope_cohort_stats::counts() shows cohort > 0 on the flagged arm + // then the default flips in the same change that records the number + Capability { + name: "ragged-cohort RoPE: one gathered rope_i for unequal offsets, default OFF", + parent: "ArcInfer / ArcAttention", + check: Check::Symbol { + symbol: "rope_cohort_enabled", + defined_in: "mistralrs-core/src/layers.rs", + }, + status: Status::Live, + }, Capability { name: "prefill admission cap (--prefill-max-seqs)", parent: "ArcInfer / ArcSched", From 738f5661c9edf5635c4d235ba0bdac52304068ab Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 21:44:30 +0100 Subject: [PATCH 16/21] feat(ArcSched/ArcKV): lift the three scheduler caps below the B>=214 gate (#219) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(ArcSched/ArcKV): lift the three scheduler caps below the B>=214 gate โ€” serve default 256, device-side ragged mask, prefill floor default 4 Parent system: ArcInfer / ArcSched + ArcKV. CAP 1 โ€” admission default. mistralrs-cli's max_seqs default goes 32 -> 256 (arg default, struct default, default_max_seqs(), and the tune config template). At B=32 the physics ceiling is ~3,740 tok/s, so the 14,000 tok/s aggregate gate (B >= 214) was unreachable on the default configuration. Memory implications documented on the flag; PagedAttention admission is block-budgeted, DefaultScheduler admission is count-only (VRAM-aware check for the dense path is named follow-up, not built). CAP 2 โ€” the mask that would make ragged decode read as a regression. CausalMasker::make_left_padded_causal_mask was a scalar triple loop over B * t_q * (past + t_q) host f32s plus a full-mask H2D copy, per decode step while a cohort is ragged โ€” 4.2 MB/token at B=256 ctx-4096 (CLAUDE.md pitfall #5 on the hot path). It is now built on the device from broadcast comparisons; the only per-step H2D is the [B] lead vector, and the index rows come from a per-device cached arange grown geometrically. The host loop is kept verbatim as make_left_padded_causal_mask_host, the test oracle, and the two are pinned element-for-element across mixed lengths, sliding windows, multi-row queries and the degenerate past=0. CAP 3 โ€” the prompt-starvation floor defaults ON at 4. Behind a 47-sequence decode cohort a fresh prompt wins 0 of 24 steps with no floor and 1 in 5 with floor=4 (this file's own fixture). ARC_PREFILL_FLOOR_STEPS=0 is the kill-switch restoring the pre-floor selection key for key; the env mapping is factored into floor_from() and pinned by test on both sides. The V4 ragged-pair default is left OFF deliberately: on this branch the pair is --v4-ragged-decode / ARC_V4_XS_PER_SEQ (the xs cache) plus ARC_MTP_PER_SEQ_KV (the per-row query positions in dsv4_attention via ragged_row_q0). The second gate lives outside this change's file fence, and opening only the first would let shorter V4 rows attend compressed blocks they have not reached โ€” silently. The box run flips the two together on a measured pass. New/updated capability registry entries: the device-built mask, the scheduler's ragged_decode_supported read, the request_xs_per_sequence CLI/config latch, and the floor entry renamed to carry its new default. New tests: a_formed_cohort_admits_a_newcomer_when_ragged_decode_is_on (cohorts could previously only shrink), the_device_mask_matches_the_host_oracle_element_for_element, the_device_mask_and_the_oracle_refuse_the_same_inputs, the_floor_defaults_on_and_zero_is_the_kill_switch. All mutation-checked; values recorded in the test docs and the PR body. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf * fix(ArcSched): a 1-token prompt keyed into the pinned ragged decode bucket โ€” prompt-ness joins the BucketKey; typos lane green The collision the PR body flagged as UNRESOLVED is real, and is now fixed and pinned by test. With ragged decode published, decode sequences key the sentinel length 0 โ€” and a 1-token prompt's cache_bucket_len() is len() - 1 == 0, the same number. The probe test showed the prompt riding into the decode cohort through the merged bucket's `<= 1` fast path: scheduled in the same pass (prompt = [4] alongside completion = [0, 1, 2, 3]), bypassing select_running_bucket, the floor's bookkeeping and the one-bucket-per-step invariant. Fix: `is_prompt` becomes a fourth BucketKey component, so a bucket is homogeneous by construction. The coalescing filter requires matching prompt-ness (a decode bucket must never idle to catch up to a prompt bucket), and the floor's is_prompt_bucket now reads the key instead of whichever sequence sits first in the bucket โ€” removing an order-dependence at the same time. Test: a_one_token_prompt_is_not_folded_into_the_pinned_decode_bucket. Green: completion = [0, 1, 2, 3], prompt = [], waiting = 1. Mutated (fourth component pinned false): red with prompt = [4]. All 14 scheduler tests and the full 702-test core lib suite pass. Typos lane: "unparseable" -> "unparsable" (floor doc), and word-level allowlist entries for arange/ARANGE in .typos.toml โ€” candle's Tensor::arange naming appearing as a segment inside cached_arange_f32 / ARANGE_CACHE, which the existing exact-identifier entries do not cover. `typos --config .typos.toml .` is clean repo-wide. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf --------- Co-authored-by: Claude Fable 5 --- .typos.toml | 7 + mistralrs-cli/src/args/mod.rs | 22 +- mistralrs-cli/src/commands/tune.rs | 2 +- mistralrs-core/src/layers_masker.rs | 246 ++++++++++++++-- .../src/scheduler/default_scheduler.rs | 278 ++++++++++++++---- .../tests/capability_reachability.rs | 50 +++- 6 files changed, 529 insertions(+), 76 deletions(-) diff --git a/.typos.toml b/.typos.toml index 7de3b414e..01c67f7c6 100644 --- a/.typos.toml +++ b/.typos.toml @@ -60,6 +60,13 @@ extend-ignore-re = [ # optin : "opt-in" without the hyphen โ€” CUDA's # cudaDevAttrMaxSharedMemoryPerBlockOptin, and filenames / # branches that cannot carry a hyphen there +# arange : candle's `Tensor::arange` (numpy naming), which also +# appears as a word segment inside larger identifiers +# (`cached_arange_f32`, `ARANGE_CACHE`) โ€” the +# extend-identifiers entries below only cover the exact +# identifier `arange`, not segments +"arange" = "arange" +"ARANGE" = "ARANGE" "nd" = "nd" "clen" = "clen" "tpos" = "tpos" diff --git a/mistralrs-cli/src/args/mod.rs b/mistralrs-cli/src/args/mod.rs index 675fcc069..4cf461e3c 100644 --- a/mistralrs-cli/src/args/mod.rs +++ b/mistralrs-cli/src/args/mod.rs @@ -427,7 +427,23 @@ pub struct GlobalOptions { #[derive(clap::Args, Clone, Deserialize)] pub struct RuntimeOptions { /// How many requests may run concurrently. - #[arg(long, default_value_t = 32, help_heading = "Serving")] + /// + /// Default 256. The previous default of 32 capped aggregate throughput far + /// below the hardware: at B=32 the physics ceiling on an H200 is ~3,740 + /// tok/s, so the 14,000 tok/s aggregate gate (which needs B >= 214) was + /// unreachable on the default configuration โ€” an admission cap, not a + /// kernel limit. + /// + /// Memory is the real bound, and it scales with this knob: KV/state memory + /// is per-sequence (~222 MiB/seq at ctx-4096 for V4 Flash-class models, so + /// 256 concurrent sequences is ~55 GiB of context memory on top of the + /// weights). PagedAttention deployments are protected by block-budget + /// admission (a sequence that does not fit waits for free blocks); the + /// `DefaultScheduler` (dense-cache models, including V4) admits by COUNT + /// only and has no VRAM-aware check โ€” on small cards, lower this + /// explicitly. Wiring the default through a VRAM-aware admission check for + /// the dense path is tracked as follow-up, not built here. + #[arg(long, default_value_t = 256, help_heading = "Serving")] #[serde(default = "default_max_seqs")] pub max_seqs: usize, @@ -558,7 +574,7 @@ impl Default for GlobalOptions { impl Default for RuntimeOptions { fn default() -> Self { Self { - max_seqs: 32, + max_seqs: 256, no_kv_cache: false, prefix_cache_n: 16, chat_template: None, @@ -580,7 +596,7 @@ fn default_token_source() -> TokenSource { } fn default_max_seqs() -> usize { - 32 + 256 } fn default_prefix_cache_n() -> usize { diff --git a/mistralrs-cli/src/commands/tune.rs b/mistralrs-cli/src/commands/tune.rs index 028e6cb3e..b5cf70aee 100644 --- a/mistralrs-cli/src/commands/tune.rs +++ b/mistralrs-cli/src/commands/tune.rs @@ -174,7 +174,7 @@ fn emit_toml_config( out.push_str("host = \"0.0.0.0\"\n"); out.push_str("port = 1234\n\n"); out.push_str("[runtime]\n"); - out.push_str("max_seqs = 32\n\n"); + out.push_str("max_seqs = 256\n\n"); out.push_str("[[models]]\n"); out.push_str(&format!("kind = \"{}\"\n", model_kind(model_type))); diff --git a/mistralrs-core/src/layers_masker.rs b/mistralrs-core/src/layers_masker.rs index 4f31d12b7..8c83a5148 100644 --- a/mistralrs-core/src/layers_masker.rs +++ b/mistralrs-core/src/layers_masker.rs @@ -100,6 +100,63 @@ impl PastKvLenCache for Vec> { } } +/// The two invariants both left-padded-mask builders require, factored so the +/// device build and its host oracle cannot drift on what they refuse. +fn validate_ragged_live(b_sz: usize, past_kv_len: usize, live: &[usize]) -> Result<()> { + if live.len() != b_sz { + candle_core::bail!( + "ragged mask: {} live lengths for a batch of {b_sz}", + live.len() + ); + } + for &l in live { + if l > past_kv_len { + candle_core::bail!( + "ragged mask: live length {l} exceeds the padded width {past_kv_len}" + ); + } + } + Ok(()) +} + +thread_local! { + /// Per-device cache of an `arange(0, N)` F32 tensor, grown geometrically + /// and narrowed per call. + /// + /// This is what keeps [`CausalMasker::make_left_padded_causal_mask`] off + /// CLAUDE.md pitfall #5: a fresh `Tensor::arange` per decode step is a + /// host-side build plus an H2D sync in the hot loop, and the index rows it + /// needs only ever grow by one column per step. Thread-local because the + /// engine drives every forward from one thread, and a `cargo test` binary + /// runs each test on its own thread โ€” a process-global here would be a + /// lock on the hot path for no aliasing that actually occurs. + static ARANGE_CACHE: std::cell::RefCell> = + const { std::cell::RefCell::new(Vec::new()) }; +} + +/// `[0.0, 1.0, ..., n-1.0]` on `device`, served from [`ARANGE_CACHE`]. +/// +/// Allocates only when `n` first exceeds the cached capacity for that device, +/// growing to the next power of two (>= 1024) so a cohort whose context climbs +/// token by token re-allocates O(log n) times over its whole life. +fn cached_arange_f32(n: usize, device: &Device) -> Result { + ARANGE_CACHE.with(|c| { + let mut c = c.borrow_mut(); + let loc = device.location(); + if let Some((_, t)) = c.iter().find(|(l, _)| *l == loc) { + if t.dim(0)? >= n { + return t.narrow(0, 0, n); + } + } + let cap = n.next_power_of_two().max(1024); + let grown = Tensor::arange(0f32, cap as f32, device)?; + let out = grown.narrow(0, 0, n)?; + c.retain(|(l, _)| *l != loc); + c.push((loc, grown)); + Ok(out) + }) +} + impl CausalMasker { fn make_mask(&self, tgt_len: usize, past_kv_len: usize, device: &Device) -> Result { let offset = tgt_len + past_kv_len; @@ -218,6 +275,24 @@ impl CausalMasker { /// `[past - live[b], past + i]`. Everything outside that is killed. Note /// the row is never fully masked โ€” column `past + i` is the query's own /// position โ€” so softmax cannot produce a NaN row. + /// + /// # Built on the DEVICE, not the host + /// + /// This runs on **every decode step** while a cohort is ragged (the + /// `ragged_lead_pad` channel stays set until batch membership changes), and + /// the first implementation was a scalar triple loop over + /// `B * t_q * (past + t_q)` host `f32`s followed by a `Tensor::from_vec` + /// H2D copy โ€” 4.2 MB/token at B=256, ctx-4096, i.e. CLAUDE.md pitfall #5 + /// landing on the hot path the moment ragged decode opens. Measured as the + /// thing that would make enabling ragged read as a regression. + /// + /// Now the only per-step H2D is the `[B]` lead vector; the row/column index + /// tensors come from a per-device cache ([`cached_arange_f32`]) and the + /// mask itself is a handful of broadcast comparisons on the device. + /// + /// [`Self::make_left_padded_causal_mask_host`] keeps the loop verbatim as + /// the test oracle; `the_device_mask_matches_the_host_oracle_element_for_element` + /// pins the two together. pub fn make_left_padded_causal_mask( &self, b_sz: usize, @@ -228,29 +303,70 @@ impl CausalMasker { dtype: DType, device: &Device, ) -> Result { - if live.len() != b_sz { - candle_core::bail!( - "ragged mask: {} live lengths for a batch of {b_sz}", - live.len() - ); - } + validate_ragged_live(b_sz, past_kv_len, live)?; + let k_len = past_kv_len + tgt_len; + + // [1, 1, 1, k]: key column index j. [1, 1, t, 1]: query row index i. + // Cached per device โ€” no `arange` allocation in the hot loop. + let js = cached_arange_f32(k_len, device)?.reshape((1, 1, 1, k_len))?; + let is_ = cached_arange_f32(tgt_len, device)?.reshape((1, 1, tgt_len, 1))?; + + // [B, 1, 1, 1]: each row's dead prefix. The one per-step H2D copy, of + // exactly B elements. (f32 is exact for integers up to 2^24, far past + // any context length this engine serves.) + let lead: Vec = live + .iter() + .map(|l| (past_kv_len - l) as f32) + .collect::>(); + let lead = Tensor::from_vec(lead, (b_sz, 1, 1, 1), device)?; + + // last[i] = past + i โ€” the query's own column, the causal upper edge. + let last = is_.affine(1.0, past_kv_len as f64)?; + + // A sliding window is expressed in ABSOLUTE positions, and this buffer + // is shifted by `lead`, so the window's low edge has to be shifted with + // it: the window admits `j > last - w`, applied on top of the dead + // prefix, never instead of it. (The host oracle's `saturating_sub` is + // subsumed by the `max` with `lead >= 0`.) + let start = match sliding_window { + Some(w) => lead.broadcast_maximum(&last.affine(1.0, 1.0 - w as f64)?)?, + None => lead, + }; + + // ok[b, 0, i, j] = (j >= start[b, i]) && (j <= last[i]). + let ok = js + .broadcast_ge(&start)? + .broadcast_mul(&js.broadcast_le(&last)?)? + .broadcast_as((b_sz, 1, tgt_len, k_len))?; + let zero = Tensor::zeros((b_sz, 1, tgt_len, k_len), DType::F32, device)?; + let neg_inf = Tensor::full(f32::NEG_INFINITY, (b_sz, 1, tgt_len, k_len), device)?; + ok.where_cond(&zero, &neg_inf)?.to_dtype(dtype) + } + + /// The host-loop reference implementation of + /// [`Self::make_left_padded_causal_mask`], kept **verbatim** as the test + /// oracle: three scalar loops whose every branch is legible, against which + /// the device build is pinned element-for-element. + /// + /// Not for the hot path โ€” this is the `O(B * t_q * k)` host build plus a + /// full-mask H2D copy per call that the device build exists to remove. + pub fn make_left_padded_causal_mask_host( + &self, + b_sz: usize, + tgt_len: usize, + past_kv_len: usize, + live: &[usize], + sliding_window: Option, + dtype: DType, + device: &Device, + ) -> Result { + validate_ragged_live(b_sz, past_kv_len, live)?; let k_len = past_kv_len + tgt_len; let mut data: Vec = Vec::with_capacity(b_sz * tgt_len * k_len); for &l in live { - if l > past_kv_len { - candle_core::bail!( - "ragged mask: live length {l} exceeds the padded width {past_kv_len}" - ); - } let lead = past_kv_len - l; for i in 0..tgt_len { let last = past_kv_len + i; - // A sliding window is expressed in ABSOLUTE positions, and this - // buffer is shifted by `lead`, so the window's low edge has to - // be shifted with it: key `j` sits at absolute position - // `j - lead`, the query at `l + i`, so the window admits - // `j > last - w`. Applied on top of the dead prefix, never - // instead of it. let start = match sliding_window { Some(w) => lead.max((last + 1).saturating_sub(w)), None => lead, @@ -790,6 +906,102 @@ mod ragged_mask_tests { crate::kv_cache::set_ragged_lead_pad(None); } + /// ๐Ÿ”‘ The device-built mask IS the host triple loop, element for element. + /// + /// The host loop ([`CausalMasker::make_left_padded_causal_mask_host`]) is + /// kept verbatim as the oracle: every branch of it is legible scalar code. + /// The production build replaces it with broadcast comparisons and one + /// `[B]` H2D copy, and this test is the only thing standing between that + /// rewrite and a silently wrong mask โ€” a wrong element here is not an + /// error anywhere downstream, it is softmax weight on a key the row must + /// not see. + /// + /// Cases cover: mixed lengths with lead 0 and live 0 in one batch, + /// multi-row queries, no window, a window tighter than the live run, a + /// window wider than everything, and the degenerate past=0. + /// + /// Mutation check (run 2026-08-21): drop the `lead` term from the device + /// build (use `0.0`) and this fails on case 0 with 4 mismatching elements + /// (the dead prefix of the `live=2` row reads `0.0` where the oracle says + /// `-inf`). + #[test] + fn the_device_mask_matches_the_host_oracle_element_for_element() { + let device = Device::Cpu; + let cases: &[(usize, usize, &[usize], Option)] = &[ + (1, 6, &[6, 2], None), + (2, 6, &[6, 2, 0], None), + (3, 4, &[4, 1], None), + (1, 5, &[5, 1], Some(2)), + (2, 6, &[6, 3], Some(4)), + (1, 5, &[3, 5], Some(32768)), + (1, 0, &[0], None), + (4, 0, &[0, 0], None), + ]; + for (case, &(tgt, past, live, window)) in cases.iter().enumerate() { + let host = CausalMasker + .make_left_padded_causal_mask_host( + live.len(), + tgt, + past, + live, + window, + DType::F32, + &device, + ) + .unwrap(); + let dev = CausalMasker + .make_left_padded_causal_mask( + live.len(), + tgt, + past, + live, + window, + DType::F32, + &device, + ) + .unwrap(); + assert_eq!(host.dims(), dev.dims(), "case {case}: shape"); + let h = host.flatten_all().unwrap().to_vec1::().unwrap(); + let d = dev.flatten_all().unwrap().to_vec1::().unwrap(); + let mismatches: Vec = h + .iter() + .zip(&d) + .enumerate() + .filter(|(_, (a, b))| a != b) + .map(|(i, _)| i) + .collect(); + assert!( + mismatches.is_empty(), + "case {case} (tgt={tgt}, past={past}, live={live:?}, window={window:?}): \ + {} of {} elements differ, first at flat index {} (host {}, device {})", + mismatches.len(), + h.len(), + mismatches[0], + h[mismatches[0]], + d[mismatches[0]], + ); + // Neither builder may emit a NaN โ€” 0 * -inf is the classic way to. + assert!(d.iter().all(|x| !x.is_nan()), "case {case}: device NaN"); + } + } + + /// Both builders refuse the same malformed inputs โ€” a live length past the + /// padded width, and a live vector of the wrong batch size. The validator + /// is shared, but shared code can be un-shared by a refactor; this pins the + /// contract itself. + #[test] + fn the_device_mask_and_the_oracle_refuse_the_same_inputs() { + let device = Device::Cpu; + for (b, live) in [(2usize, &[7usize, 2][..]), (3, &[1, 1][..])] { + assert!(CausalMasker + .make_left_padded_causal_mask(b, 1, 6, live, None, DType::F32, &device) + .is_err()); + assert!(CausalMasker + .make_left_padded_causal_mask_host(b, 1, 6, live, None, DType::F32, &device) + .is_err()); + } + } + /// `make_causal_mask_matrix` routes to the ragged builder BEFORE both of /// its early returns โ€” the `tgt_len == 1` shortcut and the flash-attn /// placeholder. A decode step is `tgt_len == 1` and still has a dead prefix diff --git a/mistralrs-core/src/scheduler/default_scheduler.rs b/mistralrs-core/src/scheduler/default_scheduler.rs index e3f550839..6769fed20 100644 --- a/mistralrs-core/src/scheduler/default_scheduler.rs +++ b/mistralrs-core/src/scheduler/default_scheduler.rs @@ -68,9 +68,22 @@ pub trait BucketingManager: Send + Sync { ) -> BucketedSeqs; } -// (cache length, (has_imgs && is_prompt), sequence offset) -// Bucket by that metric for images because if we are not a prompt, then this doesn't apply -type BucketKey = (usize, bool, usize); +// (cache length, (has_imgs && is_prompt), sequence offset, is_prompt) +// Bucket by the image flag because if we are not a prompt, then it doesn't +// apply. +// +// ๐Ÿ”‘ `is_prompt` is a key component OF ITS OWN, not folded into the image +// flag, because the ragged-decode sentinel collides without it: under +// `ragged_decode_supported()` every decode sequence keys length 0, and a +// 1-token prompt's `cache_bucket_len()` is `len() - 1 == 0` too โ€” so a fresh +// prompt silently merged into the pinned decode bucket, rode the +// `seq_buckets.len() <= 1` fast path into the same scheduling pass, and +// bypassed `select_running_bucket`, the floor's bookkeeping and the +// one-bucket-per-step invariant. With prompt-ness in the key a bucket is +// homogeneous by construction, which is also what lets the floor read +// prompt-ness off the key instead of off whichever sequence happens to sit +// first in the bucket. +type BucketKey = (usize, bool, usize, bool); /// How many waiting sequences may be admitted to **prefill** in one engine /// iteration. `None` (the default, and the historical behaviour) means "all of @@ -164,7 +177,7 @@ fn select_running_bucket( ) -> BucketKey { let min = *seq_buckets .keys() - .min_by_key(|(x, _, _)| *x) + .min_by_key(|(x, _, _, _)| *x) .expect("No sequence buckets."); if discrete { return min; @@ -180,11 +193,15 @@ fn select_running_bucket( } // A merge only happens when the *whole* key matches, so the coalescing - // target must agree on the image flag and token offset too. + // target must agree on the image flag, token offset and prompt-ness too โ€” + // a decode bucket must never idle itself to "catch up" to a prompt bucket, + // which runs by a different mechanism entirely. let Some(next_len) = seq_buckets .keys() - .filter(|(len, imgs, offset)| *len > min.0 && *imgs == min.1 && *offset == min.2) - .map(|(len, _, _)| *len) + .filter(|(len, imgs, offset, prompt)| { + *len > min.0 && *imgs == min.1 && *offset == min.2 && *prompt == min.3 + }) + .map(|(len, _, _, _)| *len) .min() else { return greedy; @@ -201,9 +218,18 @@ fn select_running_bucket( } } +/// The default prompt-starvation floor, in decode steps: a prompt bucket is +/// forced to run after this many consecutive pass-overs. +/// +/// `4` was picked against this file's own fixture (see +/// `the_floor_gives_prompts_a_turn_without_giving_them_every_turn`): behind a +/// 47-sequence decode cohort a fresh prompt wins 0 of 24 steps with no floor, +/// and 1 in every 5 with floor=4 โ€” a turn without taking every turn. +const DEFAULT_PREFILL_FLOOR_STEPS: usize = 4; + /// How many consecutive iterations prompt buckets may be passed over before one -/// is forced to run. `None` (the default) is the historical behaviour: prompts -/// win a step only by out-scoring every decode bucket on priority. +/// is forced to run. `None` means the historical behaviour: prompts win a step +/// only by out-scoring every decode bucket on priority. /// /// # Why a floor is needed as well as a cap /// @@ -223,16 +249,25 @@ fn select_running_bucket( /// waiting` and **stops**, with 48 sequences admitted and never prefilled. /// Capping admission cannot lift it; only guaranteeing prompts a turn can. /// -/// Read once from `ARC_PREFILL_FLOOR_STEPS`; unset or `0` reproduces the -/// previous selection exactly. +/// Read once from `ARC_PREFILL_FLOOR_STEPS`, **by value**: +/// * unset (or unparsable) โ†’ [`DEFAULT_PREFILL_FLOOR_STEPS`] โ€” the floor is +/// ON by default, because the starvation it lifts is the default behaviour; +/// * `0` โ†’ `None`, the kill-switch: reproduces the pre-floor selection +/// exactly, key for key; +/// * any other `n` โ†’ `Some(n)`. fn prefill_starvation_floor() -> Option { static FLOOR: std::sync::OnceLock> = std::sync::OnceLock::new(); - *FLOOR.get_or_init(|| { - std::env::var("ARC_PREFILL_FLOOR_STEPS") - .ok() - .and_then(|v| v.parse::().ok()) - .filter(|n| *n > 0) - }) + *FLOOR.get_or_init(|| floor_from(std::env::var("ARC_PREFILL_FLOOR_STEPS").ok().as_deref())) +} + +/// The decision itself, separated from the process-global `OnceLock` and the +/// environment so both sides of the default can be pinned by a test. +fn floor_from(raw: Option<&str>) -> Option { + match raw.map(str::parse::) { + Some(Ok(0)) => None, + Some(Ok(n)) => Some(n), + Some(Err(_)) | None => Some(DEFAULT_PREFILL_FLOOR_STEPS), + } } #[derive(Default)] @@ -296,42 +331,26 @@ impl BucketingManager for FixedBucketingManager { } else { seq.cache_bucket_len() }; - match seq_buckets.get_mut(&( + // `is_prompt` closes the sentinel collision documented on + // `BucketKey`: a 1-token prompt also computes `len == 0` here. + let key: BucketKey = ( len, seq.images().is_some() && seq.is_prompt(), seq.token_offset(), - )) { + seq.is_prompt(), + ); + match seq_buckets.get_mut(&key) { Some(bucket) => { if !discrete { - *seq_priorities - .get_mut(&( - len, - seq.images().is_some() && seq.is_prompt(), - seq.token_offset(), - )) - .unwrap() += seq.compute_priority(); + *seq_priorities.get_mut(&key).unwrap() += seq.compute_priority(); } bucket.push(seq); } None => { if !discrete { - seq_priorities.insert( - ( - len, - seq.images().is_some() && seq.is_prompt(), - seq.token_offset(), - ), - seq.compute_priority(), - ); + seq_priorities.insert(key, seq.compute_priority()); } - seq_buckets.insert( - ( - len, - seq.images().is_some() && seq.is_prompt(), - seq.token_offset(), - ), - vec![seq], - ); + seq_buckets.insert(key, vec![seq]); } } } @@ -380,12 +399,10 @@ impl BucketingManager for FixedBucketingManager { // floor unset โ€” or before it expires โ€” the selection is the // pre-change one, key for key. if let Some(floor) = self.floor.filter(|_| !discrete) { - let is_prompt_bucket = |k: &BucketKey| { - seq_buckets - .get(k) - .and_then(|v| v.first()) - .is_some_and(Sequence::is_prompt) - }; + // Prompt-ness is a key component, so a bucket is homogeneous + // by construction and the key answers directly โ€” no reading it + // off whichever sequence happens to sit first in the bucket. + let is_prompt_bucket = |k: &BucketKey| k.3; if is_prompt_bucket(&len) { self.steps_since_prompt = 0; } else if self.steps_since_prompt >= floor { @@ -395,7 +412,7 @@ impl BucketingManager for FixedBucketingManager { if let Some(k) = seq_buckets .keys() .filter(|k| is_prompt_bucket(k)) - .min_by_key(|(l, _, _)| *l) + .min_by_key(|(l, _, _, _)| *l) .copied() { len = k; @@ -778,7 +795,6 @@ mod tests { seq } - /// A sequence the scheduler has admitted but not yet prefilled. fn prompt_of_len(id: usize, n_toks: usize) -> Sequence { let seq = seq_of_len(id, n_toks); @@ -844,13 +860,37 @@ mod tests { ); } - /// The knob off is the previous selection, key for key โ€” asserted against - /// the same fixture rather than argued. + /// The kill-switch (`ARC_PREFILL_FLOOR_STEPS=0` โ†’ `floor: None`) is the + /// previous selection, key for key โ€” asserted against the same fixture + /// rather than argued. Note "unset" no longer means this: unset is the + /// default floor of 4 (see `floor_from`); `0` is how today's behaviour is + /// restored. #[test] - fn the_floor_unset_reproduces_the_previous_selection() { + fn the_floor_disabled_reproduces_the_previous_selection() { assert_eq!(prompt_steps_won(None, 12), 0); } + /// The env mapping behind the default, both sides pinned: + /// unset โ†’ the default floor (ON), `0` โ†’ the kill-switch (OFF, historical + /// selection), a number โ†’ itself, garbage โ†’ the default rather than a + /// silent disable. + /// + /// Mutation check (run 2026-08-21): revert `floor_from`'s `None` arm to the + /// pre-change `.filter(|n| *n > 0)` shape (unset โ†’ `None`) and this fails + /// on the first assertion with `left: None, right: Some(4)`. + #[test] + fn the_floor_defaults_on_and_zero_is_the_kill_switch() { + assert_eq!(floor_from(None), Some(DEFAULT_PREFILL_FLOOR_STEPS)); + assert_eq!(floor_from(Some("0")), None, "0 must cleanly disable"); + assert_eq!(floor_from(Some("7")), Some(7)); + assert_eq!( + floor_from(Some("not-a-number")), + Some(DEFAULT_PREFILL_FLOOR_STEPS), + "garbage must fall back to the default, not silently disable the floor" + ); + assert_eq!(DEFAULT_PREFILL_FLOOR_STEPS, 4); + } + /// Advance a scheduled sequence by one decoded token, as the engine would. fn decode_one(seq: &mut Sequence) { seq.add_token( @@ -867,7 +907,7 @@ mod tests { fn bucket_of(len: usize, n: usize, id_base: usize) -> (BucketKey, Vec) { let seqs = (0..n).map(|i| seq_of_len(id_base + i, len)).collect(); - ((len, false, 0), seqs) + ((len, false, 0, false), seqs) } /// ๐Ÿ”‘ The dense equivalent of the paged regression: eight sequences at @@ -922,6 +962,136 @@ mod tests { ); } + /// ๐Ÿ”‘ The cohort-can-only-shrink defect, and its fix, in one fixture: a + /// formed decode cohort at length 500 plus ONE newcomer at length 100. + /// + /// With exact-length bucketing the newcomer can never join: coalescing is + /// refused (`(9-1)*400 = 3200 > 1*256`), greedy priority keeps the cohort + /// (`8*log2(500) >> 1*log2(100)`), and running the cohort *grows* the gap + /// by one every step โ€” so the newcomer idles forever and cohorts only + /// shrink. Every published B=256 number under this regime came from a + /// synthetic uniform burst, because a real arrival process can never refill + /// a cohort. + /// + /// With ragged decode on, every decode sequence keys to one bucket and the + /// newcomer is admitted into the running set the same step. + /// + /// Mutation check (run 2026-08-21): pin `len` to `seq.cache_bucket_len()` + /// (reverting the `ragged_decode` arm at the top of + /// `bucket_and_waitlist_seqs_waiting`) and the ON arm fails with + /// `on: 8/9` โ€” the newcomer is waitlisted again. The OFF arm is asserted + /// in the same run, so a fixture that stopped being ragged also fails. + #[test] + fn a_formed_cohort_admits_a_newcomer_when_ragged_decode_is_on() { + use crate::kv_cache::ragged_decode_test_override as flag; + let logger = IntervalLogger::new(Duration::from_secs(3600), None); + + let running_ids_with = |on: bool| -> Vec { + flag::with(on, || { + let mut sched: DefaultScheduler> = DefaultScheduler::new( + DefaultSchedulerMethod::Fixed(NonZeroUsize::new(128).unwrap()), + ); + // The formed cohort: eight decoding sequences at length 500. + for i in 0..8 { + sched.add_seq(seq_of_len(i, 500)); + } + // The newcomer: decoding at length 100 (freshly prefilled and + // far behind the cohort โ€” the shape coalescing refuses). + sched.add_seq(seq_of_len(8, 100)); + // A waiting request, so selection runs the non-discrete + // arrival path (the discrete no-waiting path is min-length + // catch-up, a different regime with its own test). + let waiter = seq_of_len(9, 300); + waiter.set_state(SequenceState::Waiting); + sched.add_seq(waiter); + let out = DefaultScheduler::schedule(&mut sched, &logger); + out.completion.into_vec().iter().map(|s| *s.id()).collect() + }) + }; + + let off = running_ids_with(false); + assert!( + !off.contains(&8) && off.len() == 8, + "fixture guard: with exact-length bucketing the newcomer must be \ + waitlisted behind the cohort, got {off:?} โ€” otherwise the ON arm \ + proves nothing" + ); + + let on = running_ids_with(true); + assert_eq!( + on.len(), + 9, + "ragged decode must run the cohort AND the newcomer together, got {on:?}" + ); + assert!( + on.contains(&8), + "the newcomer must be in the running set, got {on:?}" + ); + } + + /// ๐Ÿ”‘ The pinned-bucket collision probe: with ragged decode published, a + /// decode sequence keys its bucket at the sentinel length 0 โ€” and a + /// 1-token PROMPT's `cache_bucket_len()` is `len() - 1 == 0`, the same + /// number. Without prompt-ness in the bucket key the two share + /// `(0, false, 0)` and the prompt silently rides into the decode cohort: + /// scheduled in the same pass through the merged single bucket's + /// `seq_buckets.len() <= 1` fast path, which bypasses `select_running_bucket`, + /// the starvation floor's bookkeeping, and the one-bucket-per-step + /// invariant the `ArcSched: decode bucketing` log line reports on. + /// + /// Asserts the prompt is NOT folded in: the decode cohort runs whole and + /// alone, and the prompt waits its turn through ordinary selection (or the + /// floor). + /// + /// Mutation check (run 2026-08-21): drop `seq.is_prompt()` from the bucket + /// key (pin the fourth component to `false`) and this fails with + /// `prompt = [4]` scheduled alongside `completion = [0, 1, 2, 3]`. + #[test] + fn a_one_token_prompt_is_not_folded_into_the_pinned_decode_bucket() { + use crate::kv_cache::ragged_decode_test_override as flag; + let logger = IntervalLogger::new(Duration::from_secs(3600), None); + + flag::with(true, || { + let mut sched: DefaultScheduler> = DefaultScheduler::new( + DefaultSchedulerMethod::Fixed(NonZeroUsize::new(128).unwrap()), + ); + // A ragged decode cohort โ€” under the flag they all key length 0. + for (i, l) in [60usize, 71, 82, 93].iter().enumerate() { + sched.add_seq(seq_of_len(i, *l)); + } + // A 1-token prompt: `cache_bucket_len()` is `1 - 1 == 0`, the + // exact value of the ragged sentinel. + let p = seq_of_len(4, 1); + p.set_state(SequenceState::Waiting); + sched.add_seq(p); + + let (completion, prompt) = { + let out = DefaultScheduler::schedule(&mut sched, &logger); + let mut completion: Vec = + out.completion.iter().map(|s| *s.id()).collect(); + completion.sort_unstable(); + let prompt: Vec = out.prompt.iter().map(|s| *s.id()).collect(); + (completion, prompt) + }; + + assert_eq!( + completion, + vec![0, 1, 2, 3], + "the decode cohort must run whole and alone" + ); + assert!( + prompt.is_empty(), + "a 1-token prompt must not share the pinned decode bucket; it \ + must wait for selection (or the floor), got prompt = {prompt:?}" + ); + assert_eq!( + Scheduler::waiting_len(&sched), + 1, + "the prompt must be waitlisted, not lost" + ); + }); + } + /// A PROMPT batch keeps the exact-length key even with the flag on. Prefill /// is right-padded and its `cu_seqlens` are built from those padded widths /// โ€” a different mechanism from the decode one, and not covered by the diff --git a/mistralrs-core/tests/capability_reachability.rs b/mistralrs-core/tests/capability_reachability.rs index 35fad7178..8a50d68b7 100644 --- a/mistralrs-core/tests/capability_reachability.rs +++ b/mistralrs-core/tests/capability_reachability.rs @@ -238,6 +238,50 @@ static REGISTRY: &[Capability] = &[ }, status: Status::Live, }, + Capability { + // The mask the entry above forces into existence, built on the DEVICE: + // the host triple-loop build (kept as + // `make_left_padded_causal_mask_host`, the test oracle) was 4.2 MB of + // host work + H2D per token at B=256 ctx-4096, i.e. the thing that + // would make ragged decode read as a regression. If production stops + // reaching this builder, ragged batches are serving unmasked dead + // prefixes โ€” silently wrong, not slow. + name: "ragged decode: the left-padded mask is built device-side", + parent: "ArcInfer / ArcAttention", + check: Check::Symbol { + symbol: "make_left_padded_causal_mask", + defined_in: "mistralrs-core/src/layers_masker.rs", + }, + status: Status::Live, + }, + Capability { + // The channel the SCHEDULER reads to stop exact-length bucketing. + // `batch_can_be_ragged` (above) is the producer; this is the consumer + // side, and either going dark re-imposes the one-bucket-per-step + // ceiling with no error anywhere. + name: "ragged decode: the scheduler reads the published capability", + parent: "ArcInfer / ArcSched", + check: Check::Symbol { + symbol: "ragged_decode_supported", + defined_in: "mistralrs-core/src/kv_cache/mod.rs", + }, + status: Status::Live, + }, + Capability { + // The supported ON-switch for the V4 ragged pair + // (`--v4-ragged-decode` / `v4_ragged_decode` config key, with + // `ARC_V4_XS_PER_SEQ` as the env fallback). A capability reachable + // only through an undocumented env var is indistinguishable from one + // that does not exist โ€” which is how the ragged-batching result stayed + // off the shipped configuration once already. + name: "ragged decode: the CLI/config surface latches the V4 xs flag", + parent: "ArcInfer / ArcKV", + check: Check::Symbol { + symbol: "request_xs_per_sequence", + defined_in: "mistralrs-core/src/kv_cache/xs_rolling.rs", + }, + status: Status::Live, + }, Capability { name: "prefill admission cap (--prefill-max-seqs)", parent: "ArcInfer / ArcSched", @@ -248,7 +292,11 @@ static REGISTRY: &[Capability] = &[ status: Status::Live, }, Capability { - name: "prompt-starvation floor (--prefill-floor-steps)", + // Default ON since arcsched/ragged-and-floor: unset resolves to a + // floor of 4 (`DEFAULT_PREFILL_FLOOR_STEPS`); + // `ARC_PREFILL_FLOOR_STEPS=0` is the kill-switch that restores the + // pre-floor selection key for key. + name: "prompt-starvation floor (ARC_PREFILL_FLOOR_STEPS, default 4)", parent: "ArcInfer / ArcSched", check: Check::Symbol { symbol: "prefill_starvation_floor", From 6e06fc781b2c628bba1507dd39f1ec24483ee114 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 21:44:58 +0100 Subject: [PATCH 17/21] fix(ArcAttention): Metal sinks arm dropped the caller's mask; memoize the V4 union mask across layers (#221) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(ArcAttention): the Metal sinks arm still dropped the caller's mask CORRECTNESS, the other backend. #199 gated the CUDA fused arm on mask.is_none(), but the Metal arm kept engaging on head_dim alone and calling flash_attn_sinks_metal โ€” which takes no mask parameter at all, so unlike the CUDA kernel it cannot even refuse one at the interface. GPT-OSS reaches this arm on every macOS build: head_dim 64 is in Metal's fused set and the model passes a mask it asserts is Some (gpt_oss.rs). MEASURED on a real Metal device (M-series, head_dim 128, F32): max|fused - reference(causal only)| = 0.000001 max|fused - reference(masked)| = 0.860029 max|ref(causal) - ref(masked)| = 0.860029 <- the mask's whole effect The maskless kernel tracks the causal-only reference to 1e-6 and misses the masked one by the mask's entire effect: dropped, not approximated. The fix is the same one line as CUDA's: an explicit mask pins the call to the unfused mask-honoring path, with a once-per-process warn naming the deliberate perf loss. New tests: * metal_sinks_mask_tests (feature = "metal", real device): the baseline above, plus the gate โ€” masked dispatch vs masked reference, diff 0.0. Mutation-checked on hardware: with the guard reverted the gate test fails at max abs diff 0.8600293; restored, 0.00000000. * padding_mask_is_honored_and_dropping_it_changes_the_output (CPU, always runs): a padding-column mask โ€” the class no fused kernel can re-derive โ€” survives the dispatch end to end. Mutation-checked: dropping the mask at the fall-through arm fails it at max abs diff 0.9353 vs 0.0 required. Rider (perf, contained): sinks_attn_cpu_varlen rebuilt its per-sequence mask with Tensor::from_vec โ€” a host-to-device upload per sequence per layer per step on the GPU fall-through (CLAUDE.md pitfall #5). It is a pure function of (q_len, kv_len, window, dtype, device), now memoized in a single slot; uniform batches build it once per shape. varlen_mask_memo_hits_and_never_ serves_stale pins both memo failure modes (mutation-checked: with the key comparison removed it fails on the (5, 8, 3) entry). Also corrects the file docs that called the unfused fall-through "CPU": it is device-agnostic candle math and is exactly what V4's head_dim-512 layers run on CUDA. The to_vec1 D2H syncs on cu_seqlens remain โ€” hoisting the host copy into FlashParams is shared plumbing, reported not fixed. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf * perf(ArcInfer/ArcAttention): memoize the V4 union mask across layers The union mask is a pure function of nine integers โ€” (q0, t_q, raw_base, t_k, t_c, ratio, window, dtype, device) โ€” yet dsv4_attention rebuilt it once per LAYER per step. On the prefill path that is three Tensor::arange host-to-device uploads plus their casts, compares, the U8 cat and the 0/-inf conversion, every layer (CLAUDE.md pitfall #5: each arange with a GPU device is a CPU-to-GPU sync). a7b0890d9/64fe1d678 already killed the DECODE-path aranges; this removes the per-layer rebuild wholesale on both paths. Mechanics: the validity build is now lazy (build_valid closure, character-for-character the old code on a miss) and the additive [1, 1, t_q, n_keys] mask is memoized thread-local on the integer key โ€” 4 FIFO slots, since one step touches at most three keys (the Standard / CSA / HCA layer classes). Per step, each distinct mask is built once instead of ~once per layer. Excluded on purpose: * graph_positions (fixed-capacity graph decode): the key would be constant across steps, and the ops must stay inside the capture region rather than be hoisted into a warmup-owned tensor whose later eviction would leave a recorded graph reading freed memory. That build is already arange-free. * ragged cohorts (row_q0): the per-row leads are content but not key. Tests (CPU, validating code not perf โ€” D14): * memo_hits_without_building_and_misses_on_a_changed_key: a hit returns the cached handle WITHOUT running the builder (handle identity is the only observable โ€” a memo that always misses returns correct values while doing nothing); an advanced q0 must miss. * memoized_masks_are_bit_identical_to_fresh_builds: decode at ctx 35 then 36 (ratio 4, window 8) โ€” two masks with the SAME shape and different content (threshold crosses 8 -> 9). Sequential same-thread runs must be bit-identical to fresh-thread (empty-memo) references. Mutation-checked: with the memo's key comparison replaced by "serve the newest entry", both tests fail โ€” the identity test at TensorId(1) != TensorId(6), the end-to-end test with ctx=36's output bits diverging from the fresh reference (first lane 1051680768 vs 1051066368). No behavioral change: a miss runs the exact pre-change build, and the existing 38 dsv4_attention tests (bit-identity, ragged, graph-geometry, absorbed-decode) all pass unchanged. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf * docs(ArcAttention): reword 'aranges' for the Typos lane (config untouched; #219 owns the allowlist) Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf --------- Co-authored-by: Claude Fable 5 --- .../src/attention/backends/sinks.rs | 343 +++++++++++++- mistralrs-core/src/models/dsv4_attention.rs | 422 ++++++++++++++---- 2 files changed, 673 insertions(+), 92 deletions(-) diff --git a/mistralrs-core/src/attention/backends/sinks.rs b/mistralrs-core/src/attention/backends/sinks.rs index 7a41e8fec..cec200002 100644 --- a/mistralrs-core/src/attention/backends/sinks.rs +++ b/mistralrs-core/src/attention/backends/sinks.rs @@ -7,7 +7,9 @@ use crate::{attention::SdpaParams, pipeline::text_models_inputs_processor::Flash /// Dispatches to: /// CUDA -> flash_attn_sinks / flash_attn_sinks_varlen /// Metal -> flash_attn_sinks_metal / flash_attn_sinks_varlen_metal -/// CPU -> unfused matmul + softmax_with_sinks +/// otherwise -> unfused matmul + softmax_with_sinks (device-agnostic candle +/// ops โ€” this is the fall-through for ANY device the fused kernels decline, +/// not a CPU-only path; DeepSeek-V4's head_dim 512 takes it on CUDA too) /// /// Varlen is used when flash_params contains cu_seqlens_k for this device AND /// q has batch > 1 AND the caller supplied no explicit mask. @@ -209,8 +211,16 @@ fn sinks_attn_regular( }); } + // ๐Ÿ”ด SAME DEFECT, OTHER BACKEND. The gate above stops the CUDA arm from + // swallowing masked work; this arm had no such gate, so a masked call on + // Metal (GPT-OSS passes a mask it asserts is `Some`, and every fused head + // dim {64,80,96,128,256} is reachable there) went to + // `flash_attn_sinks_metal` โ€” which takes no mask parameter at all โ€” and the + // mask was dropped on the floor. Unlike the CUDA kernel, the Metal entry + // point cannot even refuse a mask at the interface (it has no argument to + // refuse with), so this dispatch gate is the only line of defence. #[cfg(feature = "metal")] - if q.device().is_metal() && flash_sinks_ok_metal { + if q.device().is_metal() && flash_sinks_ok_metal && mask.is_none() { crate::attention::arcflash::note(crate::attention::arcflash::Path::VendorSinks); return mistralrs_quant::flash_attn_sinks_metal( q, @@ -221,6 +231,19 @@ fn sinks_attn_regular( window_size, ); } + #[cfg(feature = "metal")] + if q.device().is_metal() && flash_sinks_ok_metal && mask.is_some() { + // Loud once: correct-and-slower on purpose, exactly as on CUDA. + static WARNED: std::sync::OnceLock<()> = std::sync::OnceLock::new(); + WARNED.get_or_init(|| { + tracing::warn!( + target: "arcflash", + "fused head_dim={hd} Metal sinks kernel DECLINED: an explicit attention \ + mask is present and the kernel cannot honour one (it has no mask \ + parameter). Falling back to the unfused masked path." + ); + }); + } // ArcFlash/Tile โ€” the head_dim the vendor sinks kernel does not take. // DeepSeek-V4's 512 is the whole reason this arm exists. @@ -299,7 +322,12 @@ fn sinks_attn_varlen( ); } - // CPU fallback: per-sequence loop (to_vec1 is fine on CPU path) + // Unfused fallback: per-sequence loop. Device-agnostic, and reached on + // CUDA/Metal too whenever the fused varlen kernel declines the head_dim โ€” + // where these `to_vec1` calls are device-to-host syncs, once per layer per + // step. Hoisting the host copy of cu_seqlens into `FlashParams` is the fix + // for that; it is plumbing shared with the eager path and out of scope + // here. let cu_q_vec: Vec = cu_seqlens_q.to_vec1()?; let cu_k_vec: Vec = cu_seqlens_k.to_vec1()?; sinks_attn_cpu_varlen( @@ -350,6 +378,47 @@ fn varlen_causal_mask( Tensor::from_vec(data, (1, 1, q_len, kv_len), device)?.to_dtype(dtype) } +/// Memoizing front end for [`varlen_causal_mask`]. +/// +/// The builder is a `Tensor::from_vec` โ€” a host-to-device copy when the +/// tensors live on a GPU (CLAUDE.md pitfall #5) โ€” and `sinks_attn_cpu_varlen` +/// calls it once per **sequence** per layer per step. The mask is a pure +/// function of the integer key, and in a uniform-length batch every sequence +/// of every layer shares one, so a single slot turns per-(sequence ร— layer) +/// uploads into one per shape change. A ragged batch alternates keys and +/// rebuilds per sequence โ€” exactly the pre-memo cost, never worse. +fn varlen_causal_mask_cached( + q_len: usize, + kv_len: usize, + window_size: usize, + dtype: candle_core::DType, + device: &candle_core::Device, +) -> Result { + type Key = ( + usize, + usize, + usize, + candle_core::DType, + candle_core::DeviceLocation, + ); + thread_local! { + static MEMO: std::cell::RefCell> = + const { std::cell::RefCell::new(None) }; + } + let key = (q_len, kv_len, window_size, dtype, device.location()); + MEMO.with(|c| -> Result { + let mut c = c.borrow_mut(); + if let Some((k, t)) = c.as_ref() { + if *k == key { + return Ok(t.clone()); + } + } + let t = varlen_causal_mask(q_len, kv_len, window_size, dtype, device)?; + *c = Some((key, t.clone())); + Ok(t) + }) +} + /// Sinks attention on **ArcFlash/Tile** โ€” the path every DeepSeek-V4 layer /// takes, because head_dim 512 is outside the fused sinks kernel's /// `{64,80,96,112,128,192,256}`. @@ -376,7 +445,9 @@ fn sinks_attn_cpu( crate::attention::arcflash::union_attention(q, k, v, mask, Some(sinks), sdpa_params) } -/// CPU fallback for varlen: per-sequence unfused loop. +/// Unfused fallback for varlen: per-sequence loop. Despite the name this is +/// device-agnostic candle math, not CPU-only โ€” it is the fall-through on any +/// device whose fused varlen kernel declines the head_dim. /// /// Each sequence is masked with `varlen_causal_mask`, matching the fused /// varlen kernels. Passing `None` here (the pre-fix behavior) made every @@ -416,7 +487,7 @@ fn sinks_attn_cpu_varlen( .transpose(0, 1)? .unsqueeze(0)?; - let mask = varlen_causal_mask(q_len, kv_len, window_size, qi.dtype(), device)?; + let mask = varlen_causal_mask_cached(q_len, kv_len, window_size, qi.dtype(), device)?; let oi = sinks_attn_cpu(&qi, &ki, &vi, sinks, Some(&mask), sdpa_params)?; // Pad back to max_q @@ -579,6 +650,90 @@ mod tests { ); Ok(()) } + + /// Additive causal mask whose last `pad` key columns are dead PADDING โ€” + /// the exact class of information the fused kernels cannot re-derive: + /// they rebuild causality and windows internally, but nothing tells them + /// which trailing keys belong to a shorter neighbour's padding. + fn padding_mask(seq: usize, pad: usize, dev: &Device) -> Result { + let mut m = vec![0f32; seq * seq]; + for i in 0..seq { + for j in 0..seq { + if j > i || j >= seq - pad { + m[i * seq + j] = f32::NEG_INFINITY; + } + } + } + Tensor::from_vec(m, (1, 1, seq, seq), dev) + } + + /// A padding mask must survive the dispatch end to end, and dropping it + /// must CHANGE the output โ€” so this test fails on any regression that + /// re-routes masked work to a maskless kernel, rather than passing by + /// luck because the mask happened to be redundant with plain causality. + #[test] + fn padding_mask_is_honored_and_dropping_it_changes_the_output() -> Result<()> { + let dev = Device::Cpu; + let (b, h, seq, d) = (2usize, 2usize, 6usize, 16usize); + let sinks = Tensor::from_vec(vec![0.25f32, 0.5], (1, h, 1, 1), &dev)?; + let sdpa = SdpaParams { + sinks: Some(sinks.clone()), + ..params(d, h, None) + }; + + let q = mk(&[b, h, seq, d], 0.05, &dev)?; + let k = mk(&[b, 1, seq, d], 0.13, &dev)?; + let v = mk(&[b, 1, seq, d], 0.23, &dev)?; + let mask = padding_mask(seq, 2, &dev)?; + + // flash_params present and b > 1: the configuration where a dropped + // mask would additionally re-route the call onto the varlen path. + let flash = varlen_params(&[seq as u32; 2], &dev); + let got = sinks_attn(&q, &k, &v, &sinks, Some(&mask), Some(&flash), &sdpa)?; + let want = sinks_attn_cpu(&q, &k, &v, &sinks, Some(&mask), &sdpa)?; + assert_eq!( + max_abs_diff(&flat(&got), &flat(&want)), + 0.0, + "the dispatch did not honor the padding mask" + ); + + // Teeth: with the mask dropped (the defect), the answer is different + // by a wide margin โ€” the padding columns vote in the softmax. + let dropped = sinks_attn_cpu(&q, &k, &v, &sinks, None, &sdpa)?; + let signal = max_abs_diff(&flat(&want), &flat(&dropped)); + assert!( + signal > 1e-2, + "masked and maskless attention are indistinguishable ({signal}); \ + this test has no teeth" + ); + Ok(()) + } + + /// The memoized varlen mask must be bit-identical to a fresh build across + /// key changes โ€” same shape with a different window is exactly the + /// stale-serving hazard โ€” and must actually HIT on a repeated key: a memo + /// that always misses passes every value check while doing nothing. + #[test] + fn varlen_mask_memo_hits_and_never_serves_stale() -> Result<()> { + let dev = Device::Cpu; + for &(q_len, kv_len, window) in &[(5usize, 8usize, 0usize), (5, 8, 3), (5, 8, 0), (4, 9, 2)] + { + let fresh = varlen_causal_mask(q_len, kv_len, window, DType::F32, &dev)?; + let cached = varlen_causal_mask_cached(q_len, kv_len, window, DType::F32, &dev)?; + assert_eq!( + flat(&fresh), + flat(&cached), + "memo served a wrong mask for ({q_len}, {kv_len}, {window})" + ); + } + // Engagement: a repeated key returns the same tensor handle. Skipping + // the build leaves no trace in the values, so only the handle proves + // the cache did anything. + let a = varlen_causal_mask_cached(7, 7, 2, DType::F32, &dev)?; + let b = varlen_causal_mask_cached(7, 7, 2, DType::F32, &dev)?; + assert_eq!(a.id(), b.id(), "repeated key did not hit the memo"); + Ok(()) + } } /// The fused `head_dim=512` flash-sinks path takes no mask argument, so it @@ -820,3 +975,181 @@ mod fused_sinks_mask_tests { Ok(()) } } + +/// The Metal twin of `fused_sinks_mask_tests`. `flash_attn_sinks_metal` takes +/// no mask parameter at all, so โ€” unlike CUDA, where the kernel itself now +/// refuses a mask โ€” the dispatch gate is the ONLY line of defence for masked +/// work on Metal. head_dim 128 sits inside Metal's fused set, so the fused arm +/// is exactly the one that would engage. +/// +/// cargo test -p mistralrs-core --features metal metal_sinks -- --nocapture +#[cfg(all(test, feature = "metal"))] +mod metal_sinks_mask_tests { + use super::*; + use candle_core::{DType, Device}; + + fn mk(dims: &[usize], seed: f32, dev: &Device) -> Result { + let n: usize = dims.iter().product(); + let data: Vec = (0..n).map(|i| ((i as f32) * seed).sin()).collect(); + Tensor::from_vec(data, dims, dev) + } + + fn to_f32(t: &Tensor) -> Vec { + t.to_dtype(DType::F32) + .unwrap() + .flatten_all() + .unwrap() + .to_vec1::() + .unwrap() + } + + fn max_abs(a: &[f32], b: &[f32]) -> f32 { + assert_eq!(a.len(), b.len()); + a.iter() + .zip(b) + .map(|(x, y)| (x - y).abs()) + .fold(0f32, f32::max) + } + + /// Additive 0 / -inf mask [1, 1, q, k]: causal, and only the first `keep` + /// keys are admissible โ€” the padding-column class the kernel cannot + /// re-derive. + fn band_mask(q_len: usize, kv_len: usize, keep: usize, dev: &Device) -> Result { + let mut m = vec![0f32; q_len * kv_len]; + for i in 0..q_len { + for j in 0..kv_len { + let causal_ok = j <= i + (kv_len - q_len); + if !causal_ok || j >= keep { + m[i * kv_len + j] = f32::NEG_INFINITY; + } + } + } + Tensor::from_vec(m, (1, 1, q_len, kv_len), dev) + } + + const B: usize = 1; + const H: usize = 8; + const KVH: usize = 1; + const QL: usize = 16; + const KL: usize = 48; + const D: usize = 128; // inside Metal's fused set {64, 80, 96, 112, 128, 192, 256} + + fn fixtures(dev: &Device) -> Result<(Tensor, Tensor, Tensor, Tensor, SdpaParams)> { + let q = mk(&[B, H, QL, D], 0.017, dev)?; + let k = mk(&[B, KVH, KL, D], 0.023, dev)?; + let v = mk(&[B, KVH, KL, D], 0.031, dev)?; + let sinks = mk(&[H], 0.041, dev)?; + let p = SdpaParams { + n_kv_groups: H / KVH, + softcap: None, + softmax_scale: 1.0 / (D as f32).sqrt(), + sliding_window: None, + sinks: None, + }; + Ok((q, k, v, sinks, p)) + } + + /// THE BASELINE. The fused Metal kernel, called exactly as the dispatch + /// calls it, cannot depend on a mask it is never given: its output tracks + /// the causal-only reference and misses the masked one by the mask's whole + /// effect. This is the delta the dispatch gate exists to keep out of + /// serving โ€” and the size of the wrong answer if the gate regresses. + #[test] + fn fused_metal_kernel_cannot_honour_a_mask() -> Result<()> { + let dev = Device::new_metal(0).expect("this test requires a Metal device"); + let (q, k, v, sinks, p) = fixtures(&dev)?; + + let mask_causal = band_mask(QL, KL, KL, &dev)?; + let mask_narrow = band_mask(QL, KL, 8, &dev)?; + + let ref_causal = to_f32(&crate::attention::arcflash::union_attention( + &q, + &k, + &v, + Some(&mask_causal), + Some(&sinks), + &p, + )?); + let ref_narrow = to_f32(&crate::attention::arcflash::union_attention( + &q, + &k, + &v, + Some(&mask_narrow), + Some(&sinks), + &p, + )?); + let fused = to_f32(&mistralrs_quant::flash_attn_sinks_metal( + &q, + &k, + &v, + Some(&sinks), + p.softmax_scale, + 0, + )?); + + let d_causal = max_abs(&fused, &ref_causal); + let d_narrow = max_abs(&fused, &ref_narrow); + let d_masks = max_abs(&ref_causal, &ref_narrow); + eprintln!( + "max|fused - ref(causal only)| = {d_causal:.6}\n\ + max|fused - ref(masked)| = {d_narrow:.6}\n\ + max|ref(causal) - ref(masked)| = {d_masks:.6} <- the mask's own effect" + ); + assert!( + d_masks > 1e-3, + "the two masks did not change the reference ({d_masks}); test would be vacuous" + ); + assert!( + d_narrow > d_causal * 10.0, + "expected the maskless Metal kernel to track the causal-only reference \ + (d_causal {d_causal}, d_narrow {d_narrow})" + ); + Ok(()) + } + + /// THE GATE. A masked call through the dispatch must match the masked + /// unfused reference โ€” i.e. it must NOT have gone to the fused kernel the + /// baseline above proves would drop the mask. Reverting the + /// `mask.is_none()` guard on the Metal arm fails this test by + /// approximately the `d_narrow` the baseline prints. + #[test] + fn masked_sinks_attention_on_metal_matches_the_masked_reference() -> Result<()> { + let dev = Device::new_metal(0).expect("this test requires a Metal device"); + let (q, k, v, sinks, p) = fixtures(&dev)?; + let mask = band_mask(QL, KL, 8, &dev)?; + + let got = to_f32(&sinks_attn(&q, &k, &v, &sinks, Some(&mask), None, &p)?); + let want = to_f32(&crate::attention::arcflash::union_attention( + &q, + &k, + &v, + Some(&mask), + Some(&sinks), + &p, + )?); + + // Negative control on the OUTPUT: a shared-input perturbation would + // move both sides together and could never fire. + { + let mut probe = got.clone(); + assert_eq!(max_abs(&probe, &got), 0.0); + probe[7] = f32::from_bits(probe[7].to_bits().wrapping_add(1)); + assert!(max_abs(&probe, &got) > 0.0, "comparator is inert"); + } + + let d = max_abs(&got, &want); + eprintln!("masked Metal dispatch vs masked reference: max abs diff {d:.8}"); + assert!( + d < 1e-3, + "masked sinks attention on Metal does not match the masked reference \ + (max abs diff {d}) -- the Metal arm is dropping the mask again" + ); + + // ...and the maskless fused fast path still works. + assert!( + sinks_attn(&q, &k, &v, &sinks, None, None, &p).is_ok(), + "declining masks must not break the maskless fused path" + ); + Ok(()) + } +} diff --git a/mistralrs-core/src/models/dsv4_attention.rs b/mistralrs-core/src/models/dsv4_attention.rs index eeee5a75c..f916bb1da 100644 --- a/mistralrs-core/src/models/dsv4_attention.rs +++ b/mistralrs-core/src/models/dsv4_attention.rs @@ -576,6 +576,82 @@ pub(crate) fn graph_compressed_mask( Tensor::cat(&[&raw, &comp], 3) } +/// Integer key naming everything the uniform-batch union mask depends on. +/// +/// The additive `[1, 1, t_q, t_k + t_c]` mask [`dsv4_attention`] builds on the +/// uniform-batch, non-graph path is a pure function of these fields โ€” the +/// `arange` calls, casts, compares and the 0/-inf conversion take nothing else as +/// input โ€” so two calls with equal keys produce bit-identical tensors, and a +/// cached one is an identity, not an approximation. +#[derive(Clone, PartialEq, Eq, Debug)] +struct UnionMaskKey { + q0: usize, + t_q: usize, + raw_base: usize, + t_k: usize, + t_c: usize, + ratio: usize, + window: usize, + dtype: DType, + dev: candle_core::DeviceLocation, +} + +/// A handful of slots, FIFO. One step touches at most three distinct keys โ€” +/// the Standard / CSA / HCA layer classes agree on everything but `t_c` and +/// `ratio` โ€” so 4 covers a step with a spare, and earlier steps' entries +/// (whose `q0` has advanced) age out on their own. +const UNION_MASK_MEMO_SLOTS: usize = 4; +thread_local! { + static UNION_MASK_MEMO: std::cell::RefCell> = + const { std::cell::RefCell::new(Vec::new()) }; +} + +/// Memoized union-mask lookup: the cached additive mask for `key`, or `build` +/// run once and cached. +/// +/// WHY: the union mask was rebuilt from scratch once per LAYER per step. On +/// the prefill path that is three `Tensor::arange` host-to-device uploads plus +/// their casts, compares and the 0/-inf conversion (CLAUDE.md pitfall #5 โ€” +/// each upload is a CPU-to-GPU sync), all recomputing a pure function of +/// `key`. Every layer of one compress-ratio class shares a key within a step, +/// so this builds each distinct mask once per step instead of once per layer. +/// +/// Thread-local, like `GRAPH_BLOCK_IDS` above: the engine drives all layers of +/// a step from one thread, which is the reuse that matters. +fn union_mask_memo(key: &UnionMaskKey, build: impl FnOnce() -> Result) -> Result { + UNION_MASK_MEMO.with(|c| -> Result { + { + let cache = c.borrow(); + if let Some((_, t)) = cache.iter().find(|(k, _)| k == key) { + return Ok(t.clone()); + } + } + let t = build()?; + let mut cache = c.borrow_mut(); + if cache.len() >= UNION_MASK_MEMO_SLOTS { + cache.remove(0); + } + cache.push((key.clone(), t.clone())); + Ok(t) + }) +} + +/// Boolean `[t_q, n_keys]` validity โ†’ additive `[1, 1, t_q, n_keys]` mask +/// (0 on valid, -inf on invalid) in the compute dtype. +fn additive_union_mask( + valid: &Tensor, + t_q: usize, + n_keys: usize, + dtype: DType, + dev: &Device, +) -> Result { + let zeros = Tensor::zeros((t_q, n_keys), dtype, dev)?; + let neg_inf = neg_inf_full((t_q, n_keys), dtype, dev)?; + valid + .where_cond(&zeros, &neg_inf)? + .reshape((1, 1, t_q, n_keys)) +} + /// Per-call configuration for V4 hybrid attention. #[derive(Debug, Clone, Copy)] pub struct Dsv4AttentionConfig<'a> { @@ -879,83 +955,100 @@ pub fn dsv4_attention( // `arange` copies from a HOST pointer, and host-pointer copies are the // class that makes a decode step expensive to capture as a CUDA graph. let decode_row = t_q == 1; - // `qp` outlives this block only to drive the compressed branch's threshold, - // and only on the paths that still evaluate that threshold in F32. - let qp = if decode_row { - None - } else { - Some( - Tensor::arange(q0 as u32, (q0 + t_q) as u32, dev)? - .to_dtype(DType::F32)? - .reshape((t_q, 1))?, - ) + let ratio = cfg.compress_ratio.ratio(); + let t_c = match compressed_kv { + Some(comp) => comp.dim(2)?, + None => 0, }; - let raw_valid = if decode_row { - Tensor::ones((t_q, t_k), DType::U8, dev)? - } else { - let kp = Tensor::arange(raw_base as u32, (raw_base + t_k) as u32, dev)? - .to_dtype(DType::F32)? - .reshape((1, t_k))?; - let qp = qp - .as_ref() - .expect("qp is built on exactly the `t_q > 1` branch this arm is"); - let causal = kp.broadcast_le(qp)?; - let lower = (qp - window as f64)?; - let in_window = kp.broadcast_gt(&lower)?; - (causal * in_window)? // [t_q, t_k] (u8) + let graph_mode = cfg.graph_positions.is_some(); + + // ---- Union validity [t_q, t_k + t_c] (U8), built LAZILY --------------- + // On the uniform-batch, non-graph path the additive mask derived from this + // is memoized across layers on its integer key (`union_mask_memo`), and a + // hit skips this entire build โ€” the prefill path's `arange` uploads, their + // casts and every compare included. A miss runs it character-for-character + // as before. + let build_valid = || -> Result { + // `qp` drives the raw window and, on the paths that still evaluate the + // compressed branch's threshold in F32, that threshold too. + let qp = if decode_row { + None + } else { + Some( + Tensor::arange(q0 as u32, (q0 + t_q) as u32, dev)? + .to_dtype(DType::F32)? + .reshape((t_q, 1))?, + ) + }; + let raw_valid = if decode_row { + Tensor::ones((t_q, t_k), DType::U8, dev)? + } else { + let kp = Tensor::arange(raw_base as u32, (raw_base + t_k) as u32, dev)? + .to_dtype(DType::F32)? + .reshape((1, t_k))?; + let qp = qp + .as_ref() + .expect("qp is built on exactly the `t_q > 1` branch this arm is"); + let causal = kp.broadcast_le(qp)?; + let lower = (qp - window as f64)?; + let in_window = kp.broadcast_gt(&lower)?; + (causal * in_window)? // [t_q, t_k] (u8) + }; + if compressed_kv.is_none() { + return Ok(raw_valid); + } + // Compressed branch mask: query r attends compressed block b iff + // b < (q0+r+1)/ratio (causal over fully-completed ratio-blocks). + let comp_valid = if graph_mode { + // Fixed-capacity graph decode: `qp` is derived from the buffer + // *capacity*, so this threshold would freeze at a constant (see + // `Dsv4AttentionConfig::graph_positions`). Leave the local half + // permissive and let the device-derived mask folded in below + // carry the real pattern โ€” a `0` here is the additive identity, + // so the two compose without double-counting. + Tensor::ones((t_q, t_c), DType::U8, dev)? + } else { + match &qp { + // Decode (`t_q == 1`). `q0` is a host value and + // `CompressRatio::ratio` is 1 / 4 / 128 โ€” always a power of + // two โ€” so `floor((q0 + 1) / ratio)` is exact integer + // arithmetic on the host and `bp` can stay in the U32 the + // `arange` already produced. Bit-identical to the F32 chain + // below: `q0 + 1`, the quotient and every block index are + // integers, a power-of-two divide is exact in binary + // floating point, and F32 represents every integer below + // 2^24 = 16 777 216 โ€” against V4's 163 840-token position + // ceiling. Same predicate, same operands, four fewer + // launches (a cast, an add, a divide and a floor). + None => { + let threshold = + Tensor::new(&[((q0 + 1) / ratio) as u32], dev)?.reshape((1, 1))?; + // The block indices come from the cached `(ratio, + // device)` table at ratio 1, not `Tensor::arange`. + // `compress_positions(t_c, 1, dev)` IS `[0, 1, โ€ฆ, t_c-1]` + // in U32 โ€” bit-identical to the `arange` it replaces โ€” + // but it is a zero-copy `narrow` view of a table built + // once during warmup, so the per-layer host->device copy + // goes away. See `layers::compress_positions`. + crate::layers::compress_positions(t_c, 1, dev)? + .reshape((1, t_c))? + .broadcast_lt(&threshold)? // [1, t_c] (u8) + } + Some(qp) => { + let bp = Tensor::arange(0u32, t_c as u32, dev)? + .to_dtype(DType::F32)? + .reshape((1, t_c))?; + let threshold = ((qp + 1.0)? / ratio as f64)?.floor()?; // [t_q, 1] + bp.broadcast_lt(&threshold)? // [t_q, t_c] (u8) + } + } + }; + Tensor::cat(&[&raw_valid, &comp_valid], 1) }; - // ---- Build the union key set + validity mask. ------------------------- - let (k_cat, v_cat, valid) = match compressed_kv { + // ---- Build the union key set ------------------------------------------ + let (k_cat, v_cat) = match compressed_kv { Some(comp) => { - let t_c = comp.dim(2)?; - // Compressed branch mask: query r attends compressed block b iff - // b < (q0+r+1)/ratio (causal over fully-completed ratio-blocks). - let ratio = cfg.compress_ratio.ratio(); - let comp_valid = match cfg.graph_positions { - // Fixed-capacity graph decode: `qp` is derived from the buffer - // *capacity*, so this threshold would freeze at a constant (see - // `Dsv4AttentionConfig::graph_positions`). Leave the local half - // permissive and let the device-derived mask folded in below - // carry the real pattern โ€” a `0` here is the additive identity, - // so the two compose without double-counting. - Some(_) => Tensor::ones((t_q, t_c), DType::U8, dev)?, - None => match &qp { - // Decode (`t_q == 1`). `q0` is a host value and - // `CompressRatio::ratio` is 1 / 4 / 128 โ€” always a power of - // two โ€” so `floor((q0 + 1) / ratio)` is exact integer - // arithmetic on the host and `bp` can stay in the U32 the - // `arange` already produced. Bit-identical to the F32 chain - // below: `q0 + 1`, the quotient and every block index are - // integers, a power-of-two divide is exact in binary - // floating point, and F32 represents every integer below - // 2^24 = 16 777 216 โ€” against V4's 163 840-token position - // ceiling. Same predicate, same operands, four fewer - // launches (a cast, an add, a divide and a floor). - None => { - let threshold = - Tensor::new(&[((q0 + 1) / ratio) as u32], dev)?.reshape((1, 1))?; - // The block indices come from the cached `(ratio, - // device)` table at ratio 1, not `Tensor::arange`. - // `compress_positions(t_c, 1, dev)` IS `[0, 1, โ€ฆ, t_c-1]` - // in U32 โ€” bit-identical to the `arange` it replaces โ€” - // but it is a zero-copy `narrow` view of a table built - // once during warmup, so the per-layer host->device copy - // goes away. See `layers::compress_positions`. - crate::layers::compress_positions(t_c, 1, dev)? - .reshape((1, t_c))? - .broadcast_lt(&threshold)? // [1, t_c] (u8) - } - Some(qp) => { - let bp = Tensor::arange(0u32, t_c as u32, dev)? - .to_dtype(DType::F32)? - .reshape((1, t_c))?; - let threshold = ((qp + 1.0)? / ratio as f64)?.floor()?; // [t_q, 1] - bp.broadcast_lt(&threshold)? // [t_q, t_c] (u8) - } - }, - }; - let valid = Tensor::cat(&[&raw_valid, &comp_valid], 1)?; let k_cat = Tensor::cat(&[k, comp], 2)?.contiguous()?; // The second half of the same duplication -- see `v_aliases_k` // above. `v_cat` was an unconditional @@ -969,33 +1062,51 @@ pub fn dsv4_attention( fused_v_stats::record_materialized(); Tensor::cat(&[v, comp], 2)?.contiguous()? }; - (k_cat, v_cat, valid) + (k_cat, v_cat) } - None => (k.clone(), v.clone(), raw_valid), + None => (k.clone(), v.clone()), }; record_launch_shape("k_cat", k_cat.dims()); record_launch_shape("v_cat", v_cat.dims()); // Boolean validity โ†’ additive mask (0 / -inf), broadcast over batch+heads. - let n_keys = valid.dim(1)?; + let n_keys = t_k + t_c; let mask = match row_q0 { - // Uniform batch: verbatim the pre-change code. - None => { - let zeros = Tensor::zeros((t_q, n_keys), q.dtype(), dev)?; - let neg_inf = neg_inf_full((t_q, n_keys), q.dtype(), dev)?; - valid - .where_cond(&zeros, &neg_inf)? - .reshape((1, 1, t_q, n_keys))? - } + // Uniform batch, no graph positions: a pure function of the integer + // key, memoized across layers. The build a miss runs โ€” + // `additive_union_mask(&build_valid()?, ..)` โ€” is verbatim the + // pre-memo code. + None if !graph_mode => union_mask_memo( + &UnionMaskKey { + q0, + t_q, + raw_base, + t_k, + t_c, + ratio, + window, + dtype: q.dtype(), + dev: dev.location(), + }, + || additive_union_mask(&build_valid()?, t_q, n_keys, q.dtype(), dev), + )?, + // Uniform batch under fixed-capacity graph decode: NOT memoized, on + // purpose. The constant capacity would make the key constant across + // steps, and these ops must be issued inside the capture region rather + // than hoisted into a warmup-owned tensor whose eviction would leave a + // recorded graph reading freed memory. The build is already + // arange-free here โ€” both halves are fills. + None => additive_union_mask(&build_valid()?, t_q, n_keys, q.dtype(), dev)?, // Ragged cohort: `[b, 1, t_q, n_keys]`, because the compressed - // threshold and the dead prefix are both per-row. + // threshold and the dead prefix are both per-row. Not memoized: the + // per-row leads are part of the content but not of the integer key. Some(rows) => ragged_union_mask( rows, q0, - &valid, + &build_valid()?, raw_base, t_k, - cfg.compress_ratio.ratio(), + ratio, q.dtype(), dev, )?, @@ -3719,3 +3830,140 @@ mod fused_v_alias_tests { Ok(()) } } + +/// The union-mask memo: the mask is a pure function of an integer key and was +/// rebuilt once per layer per step (three `arange` uploads on the prefill +/// path). These tests pin the two failure modes a memo can add: doing nothing +/// (always missing) and serving a stale entry for a changed key. +#[cfg(test)] +#[allow(clippy::cast_precision_loss)] // usize -> f32 on test-sized fixtures, the file's test idiom +mod union_mask_memo_tests { + use super::*; + use candle_core::{DType, Device, Result, Tensor}; + use std::collections::HashMap; + + /// The memo itself: a hit returns the cached handle WITHOUT running the + /// builder, and a changed key runs it. Handle identity is the only + /// observable โ€” a memo that always misses returns correct values while + /// doing nothing. + #[test] + fn memo_hits_without_building_and_misses_on_a_changed_key() -> Result<()> { + let dev = Device::Cpu; + let key_a = UnionMaskKey { + q0: 34, + t_q: 1, + raw_base: 27, + t_k: 8, + t_c: 12, + ratio: 4, + window: 8, + dtype: DType::F32, + dev: dev.location(), + }; + let key_b = UnionMaskKey { + q0: 35, + raw_base: 28, + ..key_a.clone() + }; + + let t1 = Tensor::zeros((1, 1, 1, 20), DType::F32, &dev)?; + let got = union_mask_memo(&key_a, || Ok(t1.clone()))?; + assert_eq!(got.id(), t1.id(), "miss must return the built tensor"); + + let hit = union_mask_memo(&key_a, || panic!("builder ran on a hit"))?; + assert_eq!(hit.id(), t1.id(), "hit must return the cached handle"); + + let t2 = Tensor::ones((1, 1, 1, 20), DType::F32, &dev)?; + let miss = union_mask_memo(&key_b, || Ok(t2.clone()))?; + assert_eq!( + miss.id(), + t2.id(), + "an advanced position was served the previous position's mask" + ); + // ...and key A is still cached alongside. + let hit_a = union_mask_memo(&key_a, || panic!("builder ran on a hit"))?; + assert_eq!(hit_a.id(), t1.id()); + Ok(()) + } + + fn mk(b: usize, h: usize, t: usize, d: usize, seed: f32, dev: &Device) -> Result { + let data: Vec = (0..(b * h * t * d)) + .map(|i| ((i as f32) * seed).sin()) + .collect(); + Tensor::from_vec(data, (b, h, t, d), dev) + } + + /// One deterministic decode step at absolute context `ctx`, returned as + /// output bits. `ctx` 35 vs 36 with ratio 4 and window 8 gives two calls + /// whose MASKS have the same shape but different content (the compressed + /// threshold `floor(ctx / 4)` crosses 8 โ†’ 9), which is exactly the + /// stale-serving hazard. + fn run(ctx: usize) -> Result> { + let dev = Device::Cpu; + let (h, d, t_c) = (2usize, 8usize, 12usize); + let q = mk(1, h, 1, d, 0.03, &dev)?; + let k = mk(1, 1, ctx, d, 0.11, &dev)?; + let comp = mk(1, 1, t_c, d, 0.05, &dev)?; + let sinks: Vec = (0..h).map(|i| 0.5 + 0.25 * i as f32).collect(); + let sdpa = SdpaParams { + n_kv_groups: h, + softcap: None, + softmax_scale: 1.0 / (d as f32).sqrt(), + sliding_window: None, + sinks: Some(Tensor::from_vec(sinks, (1, h, 1, 1), &dev)?), + }; + let flash = FlashParams { + max_q: 0, + max_k: 0, + cumulative_seqlens_q: HashMap::new(), + cumulative_seqlens_k: HashMap::new(), + causal: false, + }; + let cfg = Dsv4AttentionConfig { + compress_ratio: CompressRatio::Csa, + sliding_window: 8, + raw_prefix: 0, + row_q0: None, + graph_positions: None, + }; + let out = dsv4_attention(&q, &k, &k, Some(&comp), None, &flash, &sdpa, cfg)?; + Ok(out + .to_dtype(DType::F32)? + .flatten_all()? + .to_vec1::()? + .iter() + .map(|x| x.to_bits()) + .collect()) + } + + /// End to end, the memo must be invisible: sequential calls on ONE thread + /// (where the memo carries state between them) must be bit-identical to + /// the same calls on FRESH threads (where the thread-local memo is empty + /// and every mask is built from scratch). If the memo ever served ctx=35's + /// mask to ctx=36 โ€” same shape, different admitted block set โ€” the + /// sequential run would diverge from the fresh reference. + #[test] + fn memoized_masks_are_bit_identical_to_fresh_builds() -> Result<()> { + let fresh_35 = std::thread::spawn(|| run(35)).join().unwrap()?; + let fresh_36 = std::thread::spawn(|| run(36)).join().unwrap()?; + + // Same thread: 35 (miss), 36 (miss, memo holds 35), 35 again (hit). + let seq_35 = run(35)?; + let seq_36 = run(36)?; + let seq_35_again = run(35)?; + + assert_eq!( + seq_35, fresh_35, + "first build disagrees with a fresh thread" + ); + assert_eq!( + seq_36, fresh_36, + "ctx=36 after ctx=35 disagrees with a fresh thread โ€” the memo served a stale mask" + ); + assert_eq!( + seq_35_again, fresh_35, + "a memo HIT changed the answer โ€” the cached mask is not the built one" + ); + Ok(()) + } +} From ce22a52ae5d6d02456c5b736ed02eeec21008792 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 21:45:25 +0100 Subject: [PATCH 18/21] perf(ArcInfer/ArcSample): STEP_us on V4's real arm; ARC_SAMPLE_ON_DEVICE ends the 66 MB/step host copy that walled off the GPU sampler; 3 fast-path aggravators (#222) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * perf(ArcInfer/ArcSample): port the STEP_us TOTAL/fwd/sample/other split onto the non-paged step arm V4 actually runs The only TOTAL/fwd/sample/other instrument lived on the PagedAttention arm, which V4 can never reach (DeepSeekV4Loader::supports_paged_attention is false). The DefaultInstructions arm โ€” the path V4 runs โ€” had never once reported how a step divides between GPU and host. Log format byte-identical to the paged arm's so tooling parses both. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf * fix(ArcInfer/ArcSample): three measured fast-path aggravators + the device-eligibility predicate (a) has_penalties(): a DRY config at its disabled default (multiplier 0.0, a proven no-op in apply_dry_penalty) no longer disqualifies the GPU fast path or is_raw_argmax. (b) sample_top_kp_min_p: the server defaults (top_k=-1, top_p=1.0, min_p=0.0) ran a full-vocab sort_unstable_by (129,280 elements on V4, per sequence per token) whose result nothing consumed โ€” skip it. Seeded draw provably identical (thread-local sort counter + pinned token 506 in tests). (c) sample_multinomial: the global rng MutexGuard was held (temporary lifetime extension) across get_top_logprobs and tokenizer.decode โ€” narrowed to the draw. Plus Sampler::device_sampling_terminates, the single-source mirror of sample()'s device dispatch that ARC_SAMPLE_ON_DEVICE partitioning consults, and the ignored $0 host-cost probe at V4 vocab (D14: host evidence only). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf * test(ArcInfer/ArcSample): partition split tests โ€” copy discipline + seeded token parity for ARC_SAMPLE_ON_DEVICE Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf * chore(ArcInfer/ArcSample): u32::try_from for the gather row index (clippy-by-hand for mistralrs-core) Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf --------- Co-authored-by: Claude Fable 5 --- mistralrs-core/src/pipeline/mod.rs | 408 +++++++++++++++ mistralrs-core/src/sampler.rs | 473 +++++++++++++++++- .../tests/capability_reachability.rs | 13 + 3 files changed, 891 insertions(+), 3 deletions(-) diff --git a/mistralrs-core/src/pipeline/mod.rs b/mistralrs-core/src/pipeline/mod.rs index 37110ed39..7fa8c878f 100644 --- a/mistralrs-core/src/pipeline/mod.rs +++ b/mistralrs-core/src/pipeline/mod.rs @@ -508,6 +508,12 @@ impl ForwardInputsResult { pub(crate) static LOGITS_HOST_COPIES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); +/// Serialises tests that assert on [`LOGITS_HOST_COPIES`] deltas: the counter +/// is process-global and `cargo test` runs tests concurrently, so two +/// counter-delta tests interleaving would read each other's increments. +#[cfg(test)] +pub(crate) static LOGITS_HOST_COPIES_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + /// Move a batched forward result to `host` with a **single** transfer, before /// it is split into per-sequence views. /// @@ -545,6 +551,88 @@ pub(crate) fn host_copy_batched_result( raw.to_device(host) } +/// Parent system: ArcInfer / ArcSample. +/// +/// `ARC_SAMPLE_ON_DEVICE`: opt-in device-resident sampling for batched decode. +/// Default OFF; parsed by value (`env_flag_is_set`, so `=0` is off), latched +/// once per process. +/// +/// WHY THIS FLAG EXISTS: at any `n_seqs > 1` the `DefaultInstructions` step +/// arm copied the ENTIRE `[B, 1, vocab]` logits tensor to the host before +/// sampling โ€” at B=256 with V4's 129,280 vocab in BF16 that is 66.2 MB per +/// decode step, serial, on the critical path โ€” and `Sampler::sample`'s GPU +/// fast path is gated on `!logits.device().is_cpu()`, so the copy itself then +/// disqualified all 256 rows into the ~5.9 ms/token CPU pipeline the fast +/// path exists to avoid. [`host_copy_batched_result`]'s own doc states the +/// trap for a single sequence ("moving a B=1 batch to the host would silently +/// disable it") and never followed the thought to B > 1: the GPU sampler was +/// structurally unreachable at every batch above 1. +/// +/// With the flag ON, [`split_batched_logits_for_sampling`] keeps +/// fast-path-eligible rows on the device (token-sized D2H each) and gathers +/// the rest into ONE batched D2H. OFF is bit-for-bit the shipped behaviour. +pub(crate) fn sample_on_device_enabled() -> bool { + static ON: std::sync::OnceLock = std::sync::OnceLock::new(); + *ON.get_or_init(|| mistralrs_quant::env_flag_is_set("ARC_SAMPLE_ON_DEVICE")) +} + +/// Split a batched `[B, 1, vocab]` logits tensor into per-row tensors for +/// sampling: rows whose sampler terminates on the device +/// ([`crate::sampler::Sampler::device_sampling_terminates`]) stay as device +/// views; the rows that genuinely need the CPU pipeline are shipped with at +/// most ONE batched D2H. `row_eligible` is in logits-row order (the +/// `seq_indices` order of the step arm, not sequence order). +/// +/// Copy discipline โ€” the O(Bยฒ) trap, same as [`host_copy_batched_result`]: a +/// per-row view still references the whole `B * vocab` storage, and both +/// `Tensor::to_device` and `Tensor::to_vec1` copy the ENTIRE storage +/// (`CudaStorage::to_cpu_storage`, candle `9586979`, `tensor.rs:1949/2379`). +/// So CPU-bound rows must be gathered into fresh row-sized storage +/// (`index_select` allocates exactly the selected rows) *before* the single +/// `to_device`. The `Tensor::from_vec` for the gather indices is a +/// once-per-step H2D of `n_cpu` u32s (pitfall #5 concerns per-token tensor +/// creation in forward hot loops; this one replaces a 66 MB D2H and only runs +/// for mixed batches). +pub(crate) fn split_batched_logits_for_sampling( + logits: &Tensor, + row_eligible: &[bool], +) -> candle_core::Result> { + let n = row_eligible.len(); + let n_cpu = row_eligible.iter().filter(|e| !**e).count(); + // Every row stays on device: pure views, no transfer at all. + if n_cpu == 0 { + return (0..n).map(|i| logits.i(i)).collect(); + } + LOGITS_HOST_COPIES.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + // Every row needs the CPU pipeline: one whole-tensor D2H โ€” exactly the + // flag-off behaviour. + if n_cpu == n { + let host = logits.to_device(&Device::Cpu)?; + return (0..n).map(|i| host.i(i)).collect(); + } + // Mixed batch: gather ONLY the CPU-bound rows into fresh contiguous + // device storage, then one D2H of that gather. + let cpu_rows: Vec = row_eligible + .iter() + .enumerate() + .filter(|(_, e)| !**e) + .map(|(i, _)| u32::try_from(i).expect("batch row index fits u32")) + .collect(); + let idx = Tensor::from_vec(cpu_rows, n_cpu, logits.device())?; + let host = logits.index_select(&idx, 0)?.to_device(&Device::Cpu)?; + let mut host_pos = 0usize; + let mut out = Vec::with_capacity(n); + for (i, eligible) in row_eligible.iter().enumerate() { + if *eligible { + out.push(logits.i(i)?); + } else { + out.push(host.i(host_pos)?); + host_pos += 1; + } + } + Ok(out) +} + #[derive(serde::Serialize, serde::Deserialize)] pub(crate) struct FileListCache { files: Vec, @@ -982,6 +1070,16 @@ pub trait Pipeline: let mut raw_out_logits = vec![vec![None; len_inputs]; input_seqs.len()]; let mut embedding_logits = vec![None; input_seqs.len()]; + // STEP_us instrumentation, ported verbatim from the + // PagedAttention arm below. That arm carried the only + // TOTAL/fwd/sample/other split in the codebase, and V4 cannot + // reach it (`DeepSeekV4Loader::supports_paged_attention` is + // false) โ€” so the path V4 actually runs had never once + // reported how a step divides between GPU and host. The log + // format is byte-identical to the paged arm's on purpose: + // tooling that parses one parses both. + let __np_t_step_start = Instant::now(); + let mut __np_t_fwd_us: f64 = 0.0; let mut exec_duration = Duration::ZERO; for (i, inputs) in inputs_iter.into_iter().enumerate() { let InputProcessorOutput { @@ -1013,6 +1111,54 @@ pub trait Pipeline: }; let end = Instant::now(); exec_duration += end.duration_since(start); + __np_t_fwd_us += end.duration_since(start).as_secs_f64() * 1e6; + + // Device-resident sampling (ARC_SAMPLE_ON_DEVICE, default + // OFF): keep fast-path-eligible rows on the device instead + // of hosting the whole batch โ€” see + // [`sample_on_device_enabled`] for why. Only the batched + // CausalGeneration case changes; B=1 already stays on + // device via `host_copy_batched_result`'s `n_seqs > 1` + // condition, and RawLogits/Embeddings are host-consumed. + let device_sampling = sample_on_device_enabled() + && seq_indices.len() > 1 + && matches!( + &raw_logits, + ForwardInputsResult::CausalGeneration { logits } + if !logits.device().is_cpu() + ); + if device_sampling { + let ForwardInputsResult::CausalGeneration { logits: batched } = + &raw_logits + else { + unreachable!("device_sampling implies CausalGeneration") + }; + let vocab = batched.dim(candle_core::D::Minus1)?; + let is_cuda = batched.device().is_cuda(); + let row_eligible: Vec = seq_indices + .iter() + .map(|&seq_idx| { + let return_logprobs = input_seqs[seq_idx].return_logprobs(); + input_seqs[seq_idx].sampler().device_sampling_terminates( + vocab, + is_cuda, + return_logprobs, + ) + }) + .collect(); + // Still a `sync_span`: the batched-gather D2H (when any + // row needs the CPU pipeline) is the same blocking + // device wait the flag-off copy was. + let rows = { + let _s = arc_profiler::sync_span("logits_d2h"); + split_batched_logits_for_sampling(batched, &row_eligible)? + }; + for (row, seq_idx) in rows.into_iter().zip(seq_indices) { + logits[seq_idx] = + Some(ForwardInputsResult::CausalGeneration { logits: row }); + } + continue; + } // ONE host copy for the whole batch, *before* the // per-sequence split. See `host_copy_batched_result`: the @@ -1204,6 +1350,30 @@ pub trait Pipeline: } let end = Instant::now(); exec_duration += end.duration_since(start); + let __np_t_sample_us = end.duration_since(start).as_secs_f64() * 1e6; + let __np_t_total_us = __np_t_step_start.elapsed().as_secs_f64() * 1e6; + { + use std::sync::atomic::{AtomicU64, Ordering}; + static N: AtomicU64 = AtomicU64::new(0); + static SUM_TOTAL: AtomicU64 = AtomicU64::new(0); + static SUM_FWD: AtomicU64 = AtomicU64::new(0); + static SUM_SAMPLE: AtomicU64 = AtomicU64::new(0); + let n = N.fetch_add(1, Ordering::Relaxed) + 1; + SUM_TOTAL.fetch_add(__np_t_total_us as u64, Ordering::Relaxed); + SUM_FWD.fetch_add(__np_t_fwd_us as u64, Ordering::Relaxed); + SUM_SAMPLE.fetch_add(__np_t_sample_us as u64, Ordering::Relaxed); + if n > 4 && (n - 4) % 25 == 0 { + let nn = (n - 4) as f64; + let total = SUM_TOTAL.load(Ordering::Relaxed) as f64 / nn; + let fwd = SUM_FWD.load(Ordering::Relaxed) as f64 / nn; + let samp = SUM_SAMPLE.load(Ordering::Relaxed) as f64 / nn; + let other = total - fwd - samp; + tracing::info!( + "STEP_us avg over {} steps: TOTAL={:.0} fwd={:.0} sample={:.0} other={:.0} (=> {:.1} tok/s)", + n - 4, total, fwd, samp, other, 1e6 / total + ); + } + } Ok(exec_duration) } @@ -1800,6 +1970,9 @@ mod host_copy_tests { #[test] fn batched_host_copy_is_o1_in_batch_while_the_old_order_was_o_b_squared() { + let _serial = LOGITS_HOST_COPIES_TEST_LOCK + .lock() + .unwrap_or_else(|e| e.into_inner()); // 7 is equal to none of the batch sizes used below and coprime with // all of them, so no assertion can pass by `b * v` aliasing `v`, `b`, // or `v * b` (DOCTRINE D12: the fixture must discriminate). @@ -1925,6 +2098,241 @@ mod host_copy_tests { } } +/// `ARC_SAMPLE_ON_DEVICE`'s partitioned split: copy discipline and seeded +/// token parity, on CPU-checkable fixtures. +/// +/// On a CPU fixture every "device-resident" row re-enters the same CPU +/// sampling code, so what these tests prove is exactly the CPU-checkable +/// half: the partition itself โ€” row routing, ordering, and the one-batched- +/// copy discipline โ€” cannot change a sampled token or move more bytes than +/// the rows it ships. GPU-side fast-path behaviour is the `Sampler`'s +/// existing contract, not established here (D14). +#[cfg(test)] +mod device_sampling_split_tests { + use super::*; + use crate::sampler::Sampler; + use candle_core::{CpuStorage, Storage}; + use rand::SeedableRng; + use rand_isaac::Isaac64Rng; + use std::sync::atomic::Ordering; + use std::sync::Mutex; + + /// Elements candle would move across a device boundary for `t`: the whole + /// backing storage, not `t.elem_count()` (see `host_copy_tests`). + fn storage_elems(t: &Tensor) -> usize { + let (storage, _layout) = t.storage_and_layout(); + match &*storage { + Storage::Cpu(CpuStorage::F32(v)) => v.len(), + _ => panic!("fixture must be a CPU F32 tensor"), + } + } + + /// `[b, 1, v]` ramp `0..b*v`: every row identifies itself, so any + /// host-position mapping error in the gather produces a visible wrong + /// slice rather than a coincidentally-equal one. + fn batch(b: usize, v: usize) -> Tensor { + let vals: Vec = (0..b * v).map(|x| x as f32).collect(); + Tensor::from_vec(vals, (b, 1, v), &Device::Cpu).unwrap() + } + + fn row_values(t: &Tensor) -> Vec { + t.flatten_all().unwrap().to_vec1::().unwrap() + } + + #[test] + fn mixed_batch_gathers_only_the_cpu_rows_in_one_copy() { + let _serial = LOGITS_HOST_COPIES_TEST_LOCK + .lock() + .unwrap_or_else(|e| e.into_inner()); + const B: usize = 6; + const V: usize = 7; + let logits = batch(B, V); + let eligible = [true, false, true, false, false, true]; + let n_cpu = 3usize; + + let before = LOGITS_HOST_COPIES.load(Ordering::Relaxed); + let rows = split_batched_logits_for_sampling(&logits, &eligible).unwrap(); + assert_eq!( + LOGITS_HOST_COPIES.load(Ordering::Relaxed) - before, + 1, + "a mixed batch is ONE batched copy, not one per CPU row" + ); + + assert_eq!(rows.len(), B); + for (k, row) in rows.iter().enumerate() { + assert_eq!(row.dims(), &[1, V], "row {k} shape"); + let expect: Vec = (k * V..(k + 1) * V).map(|x| x as f32).collect(); + assert_eq!(row_values(row), expect, "row {k} carries the wrong data"); + if eligible[k] { + // Device-resident rows are views over the original storage: + // zero bytes moved for them. + assert_eq!( + storage_elems(row), + B * V, + "row {k}: an eligible row must stay a view of the batch" + ); + } else { + // CPU-bound rows share the gathered storage: exactly the + // selected rows crossed, not the batch (and not O(Bยฒ)). + assert_eq!( + storage_elems(row), + n_cpu * V, + "row {k}: the gather must materialise ONLY the CPU rows" + ); + } + } + } + + #[test] + fn all_eligible_moves_nothing_and_none_eligible_matches_flag_off() { + let _serial = LOGITS_HOST_COPIES_TEST_LOCK + .lock() + .unwrap_or_else(|e| e.into_inner()); + const B: usize = 4; + const V: usize = 5; + let logits = batch(B, V); + + // All rows terminate on device: no host copy at all. + let before = LOGITS_HOST_COPIES.load(Ordering::Relaxed); + let rows = split_batched_logits_for_sampling(&logits, &[true; B]).unwrap(); + assert_eq!( + LOGITS_HOST_COPIES.load(Ordering::Relaxed) - before, + 0, + "an all-eligible batch must not touch the host" + ); + for (k, row) in rows.iter().enumerate() { + assert_eq!(storage_elems(row), B * V, "row {k} must be a view"); + } + + // No row terminates on device: one whole-batch copy โ€” the flag-off + // behaviour exactly (same counter delta as host_copy_batched_result). + let before = LOGITS_HOST_COPIES.load(Ordering::Relaxed); + let rows = split_batched_logits_for_sampling(&logits, &[false; B]).unwrap(); + assert_eq!( + LOGITS_HOST_COPIES.load(Ordering::Relaxed) - before, + 1, + "an all-CPU batch is one batched copy, exactly like flag-off" + ); + for (k, row) in rows.iter().enumerate() { + let expect: Vec = (k * V..(k + 1) * V).map(|x| x as f32).collect(); + assert_eq!(row_values(row), expect, "row {k} after whole-batch copy"); + } + } + + /// Seeded sampling through the partition must produce IDENTICAL tokens to + /// the flag-off path (batched host copy + `index_bs` split), row for row, + /// for any eligibility mask. Rows sample in row order on a shared seeded + /// rng in both paths, so any reordering, extra draw, or wrong-row routing + /// shifts the stream and fails. + #[test] + fn seeded_tokens_identical_to_the_flag_off_split() { + // Not asserting on counter deltas, but this test INCREMENTS the + // global counter (both split paths), so it must still serialise with + // the tests that do assert. + let _serial = LOGITS_HOST_COPIES_TEST_LOCK + .lock() + .unwrap_or_else(|e| e.into_inner()); + const B: usize = 4; + const V: usize = 64; + + // Deterministic spread-out logits (LCG as in sampler parity tests). + let mut state: u64 = 0xA5A5 ^ 0xDEAD_BEEF_CAFE_F00D; + let vals: Vec = (0..B * V) + .map(|_| { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((state >> 32) as u32 as f32 / u32::MAX as f32) * 12.0 - 6.0 + }) + .collect(); + let logits = Tensor::from_vec(vals, (B, 1, V), &Device::Cpu).unwrap(); + + // One sampler shape per row: greedy, top-k, server defaults, top-p. + let mk = |temperature: Option, top_k: i64, top_p: f64| { + Sampler::new( + temperature, + 0, + None, + None, + None, + None, + None, + top_k, + top_p, + 0.0, + None, + vec![], + None, + ) + .unwrap() + }; + let samplers = [ + mk(None, -1, 1.0), + mk(Some(0.7), 8, 1.0), + mk(Some(1.0), -1, 1.0), + mk(Some(0.7), -1, 0.9), + ]; + + // Mirrors `sample_sequence`'s per-row prologue. + let sample_rows = |rows: &[Tensor]| -> Vec { + let rng = Arc::new(Mutex::new(Isaac64Rng::seed_from_u64(42))); + rows.iter() + .zip(samplers.iter()) + .map(|(row, sampler)| { + let row = row + .squeeze(0) + .unwrap() + .to_dtype(DType::F32) + .unwrap(); + sampler + .sample(row, &[1u32, 2, 3], false, rng.clone(), false, true) + .unwrap() + .token + }) + .collect() + }; + + // Flag-off: batched host copy, then the per-sequence split. + let flag_off_rows: Vec = { + let host = host_copy_batched_result( + ForwardInputsResult::CausalGeneration { + logits: logits.clone(), + }, + B, + &Device::Cpu, + ) + .unwrap(); + (0..B) + .map(|k| match host.index_bs(k).unwrap() { + ForwardInputsResult::CausalGeneration { logits } => logits, + _ => unreachable!(), + }) + .collect() + }; + let expected = sample_rows(&flag_off_rows); + + // Flag-on, under every mask shape the partition can take. + for mask in [ + [true, true, true, true], + [false, false, false, false], + [true, false, true, false], + [false, true, false, true], + ] { + let rows = split_batched_logits_for_sampling(&logits, &mask).unwrap(); + assert_eq!( + sample_rows(&rows), + expected, + "mask {mask:?}: partitioned rows must sample the same tokens" + ); + } + + // Pin the expected tokens: if the sampler's rng consumption ever + // changes shape, this fails here rather than passing vacuously as + // both-sides-moved parity. + assert_eq!(expected, vec![23, 42, 31, 22], "seeded fixture drifted"); + } +} + /// `extract_logits` vs the `.i((.., seq_len - 1, ..))` it replaced in the phi3 /// and starcoder2 GGUF models. /// diff --git a/mistralrs-core/src/sampler.rs b/mistralrs-core/src/sampler.rs index 65547f010..ffa3f163a 100644 --- a/mistralrs-core/src/sampler.rs +++ b/mistralrs-core/src/sampler.rs @@ -300,6 +300,8 @@ fn cmp_desc_by_prob(a: &(u32, f32), b: &(u32, f32)) -> std::cmp::Ordering { /// If `k >= probs.len()`, returns all elements sorted. /// Also zeros out elements in `probs` beyond top-k if `zero_rest` is true. fn partial_sort_top_k(probs: &mut [f32], k: usize, zero_rest: bool) -> Vec<(u32, f32)> { + #[cfg(test)] + PARTIAL_SORT_CALLS.with(|calls| calls.set(calls.get() + 1)); let n = probs.len(); if n == 0 || k == 0 { return Vec::new(); @@ -333,6 +335,18 @@ fn partial_sort_top_k(probs: &mut [f32], k: usize, zero_rest: bool) -> Vec<(u32, idx_probs } +/// Invocation counter for [`partial_sort_top_k`], test builds only. Exists so +/// the sort-skip in [`Sampler::sample_top_kp_min_p`] is provable: a test can +/// assert the sort genuinely did not run, not merely that the output looks +/// the same. Thread-local, because `cargo test` runs tests concurrently and +/// several other tests in this file sort โ€” a process-global counter would make +/// every delta assertion racy. +#[cfg(test)] +thread_local! { + pub(crate) static PARTIAL_SORT_CALLS: std::cell::Cell = + const { std::cell::Cell::new(0) }; +} + /// Kernel dispatch sizes for the GPU radix top-k sampler path. Must mirror /// `arc_cuda_graph::flashmlasparse::SUPPORTED_TOPK` (kept as a local const so /// the candidate-truncation logic and its CPU-parity tests build without the @@ -499,10 +513,19 @@ impl Sampler { /// directly. Two copies of this expression could disagree, and the one that /// disagreed would be a silent correctness bug rather than a slow path. fn has_penalties(&self) -> bool { + // DRY counts only when it can actually rewrite a logit. The disabled + // default (`DrySamplingParams::default()`) is `multiplier: 0.0`, and + // `apply_dry_penalty` returns without touching the logits when the + // multiplier is exactly 0.0 โ€” so `dry_params.is_some()` alone was + // disqualifying the GPU fast path (and failing `is_raw_argmax`) for a + // penalty that is a proven no-op. self.frequency_penalty.unwrap_or(0.0) != 0.0 || self.presence_penalty.unwrap_or(0.0) != 0.0 || self.repetition_penalty.unwrap_or(1.0) != 1.0 - || self.dry_params.is_some() + || self + .dry_params + .as_ref() + .is_some_and(|params| params.multiplier != 0.0) } /// Would [`Sampler::sample`] take the pure-argmax path โ€” no temperature and @@ -561,6 +584,69 @@ impl Sampler { && self.logits_processors.is_empty() } + /// Mirror of [`Self::sample`]'s device fast-path dispatch, answerable + /// **before** the logits row has been copied anywhere: would a non-CPU row + /// of `vocab` logits be sampled to completion on the device (token-sized + /// D2H), rather than falling back to the CPU pipeline (full-row D2H + + /// full-vocab host work)? + /// + /// This is what `ARC_SAMPLE_ON_DEVICE`'s host-copy partitioning + /// (`pipeline/mod.rs`) consults per sequence: rows answering `true` keep + /// their device residency; rows answering `false` are gathered into ONE + /// batched D2H. It must stay in this file, next to the gates it mirrors + /// ([`Self::has_penalties`] and `sample`'s `trivial` condition): a second + /// copy that drifted would silently route rows onto the wrong side of the + /// PCIe bus. + /// + /// Assumes non-speculative sampling (`sample_speculative = false`), which + /// is what every `sample_causal_gen` reachable from the partitioned step + /// arm passes; the speculative pipelines override `step` entirely and + /// never see the partition. + /// + /// `logit_bias` and top-nฯƒ do **not** disqualify: both are applied as + /// device tensor ops before the fast-path dispatch. One residual runtime + /// fallback remains on the CUDA radix path (a pure-top-p nucleus the + /// candidate set cannot contain declines with `Ok(None)`); it pays a + /// per-row D2H and is counted by `gpu_sampling_health::DECLINED`. + pub(crate) fn device_sampling_terminates( + &self, + vocab: usize, + device_is_cuda: bool, + return_logprobs: bool, + ) -> bool { + // `sample`'s `trivial` gate, minus its device term โ€” the caller is + // deciding device residency, not reacting to it. + let trivial = + !return_logprobs && self.logits_processors.is_empty() && !self.has_penalties(); + if !trivial { + return false; + } + // Greedy: a single argmax kernel + 4-byte D2H. Any backend, any vocab. + if self.temperature.is_none() { + return true; + } + // Temperature path: `sample_fast` stays on device when no sort is + // needed, or when candle's shared-memory arg_sort can hold the vocab + // (same bound `sample` checks). + let needs_sort = self.top_k > 0 || (self.top_p > 0.0 && self.top_p < 1.0); + if !needs_sort { + return true; + } + if vocab.next_power_of_two().saturating_mul(4) <= 48 * 1024 { + return true; + } + // Big-vocab sort: only the CUDA radix top-k path avoids the CPU + // fallback, and only when the requested k fits its dispatch table + // (`sample_fast_topk_gpu`'s own static precondition). + let max_k = *GPU_RADIX_TOPK_SIZES.last().unwrap(); + let k_needed = if self.top_k > 0 { + self.top_k as usize + } else { + max_k + }; + cfg!(feature = "cuda") && device_is_cuda && k_needed <= max_k && vocab > max_k + } + /// True if an OpenAI `logit_bias` map is set. Any sampling path that does /// not route through [`Self::sample`] โ€” notably the GPU-autonomous decode /// sampler, which draws from the model's raw logits โ€” must refuse the @@ -1186,8 +1272,17 @@ impl Sampler { ) -> Result { let distr = WeightedIndex::new(probs).map_err(Error::wrap)?; - let mut mut_ref_rng = &mut *rng.lock().expect("could not lock rng mutex"); - let next_token = distr.sample(&mut mut_ref_rng); // "Find the first item which has a weight *higher* than the chosen weight." + // Hold the global rng mutex for the draw and NOTHING else. The guard + // used to live to the end of this function (temporary lifetime + // extension through the `&mut *lock()` borrow), which serialised every + // concurrently-sampling sequence in the batch behind this one's + // `get_top_logprobs` (a partial sort plus up to `top_n` tokenizer + // decodes) and its own `tokenizer.decode` โ€” none of which touch the + // rng. + let next_token = { + let mut guard = rng.lock().expect("could not lock rng mutex"); + distr.sample(&mut *guard) // "Find the first item which has a weight *higher* than the chosen weight." + }; let logprob = probs[next_token].log(10.0); let top_logprobs = if return_logprobs { @@ -1231,6 +1326,22 @@ impl Sampler { probs.len() }; + // The sorted pairs feed exactly three consumers: the zero-rest pass + // inside `partial_sort_top_k` (only when `k < n`), the top-p cumsum + // walk, and the min-p threshold scan. Under the server defaults + // (`top_k = -1`, `top_p = 1.0`, `min_p = 0.0` โ€” `engine/add_request.rs`) + // none of the three runs: `k = n` zeroes nothing, and both filter + // blocks below are gated out of (0, 1). That request was still paying + // a full-vocab `sort_unstable_by` โ€” 129,280 elements for V4, per + // sequence per token โ€” to produce a Vec nothing read. Skip straight to + // the multinomial draw; `probs` is untouched either way, so the drawn + // token (same rng stream) and the reported logprobs are identical. + let top_p_active = top_p > 0.0 && top_p < 1.0; + let min_p_active = min_p > 0.0 && min_p < 1.0; + if k >= probs.len() && !top_p_active && !min_p_active { + return self.sample_multinomial(probs, return_logprobs, rng); + } + // Get sorted top-k indices with partial sort, zeroing out rest let idx_probs = partial_sort_top_k(probs, k, true); @@ -2893,6 +3004,362 @@ mod topk_parity_tests { } } +/// The server-default sort skip and the penalty/eligibility gates. +#[cfg(test)] +mod fast_path_gate_tests { + use super::{DrySamplingParams, Sampler, PARTIAL_SORT_CALLS}; + use candle_core::{Device, Tensor}; + use rand::SeedableRng; + use rand_isaac::Isaac64Rng; + use std::sync::{Arc, Mutex}; + + const VOCAB: usize = 512; + + /// Deterministic spread-out logits (same LCG family as the parity tests). + fn logits() -> Tensor { + let mut state: u64 = 0x5EED ^ 0xDEAD_BEEF_CAFE_F00D; + let v: Vec = (0..VOCAB) + .map(|_| { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((state >> 32) as u32 as f32 / u32::MAX as f32) * 12.0 - 6.0 + }) + .collect(); + Tensor::from_vec(v, VOCAB, &Device::Cpu).unwrap() + } + + fn sampler(top_k: i64, top_p: f64, min_p: f64) -> Sampler { + Sampler::new( + Some(0.7), + 0, + None, + None, + None, + None, + None, + top_k, + top_p, + min_p, + None, + vec![], + None, + ) + .unwrap() + } + + /// Sampled token + how many times `partial_sort_top_k` ran during the + /// call. `Sampler::sample` on this thread never crosses threads, so the + /// thread-local delta is exact. + fn sample_counting_sorts(s: &Sampler, seed: u64) -> (u32, u64) { + let rng = Arc::new(Mutex::new(Isaac64Rng::seed_from_u64(seed))); + let before = PARTIAL_SORT_CALLS.with(|c| c.get()); + let token = s + .sample(logits(), &[1u32, 2, 3], false, rng, false, true) + .unwrap() + .token; + let sorts = PARTIAL_SORT_CALLS.with(|c| c.get()) - before; + (token, sorts) + } + + /// The server defaults (`top_k = -1`, `top_p = 1.0`, `min_p = 0.0` โ€” + /// `engine/add_request.rs`) must not run the full-vocab sort: nothing + /// consumes it. Equivalence witness: the PRE-FIX code path replicated + /// literally โ€” `partial_sort_top_k(probs, n, zero_rest=true)` (which for + /// `k = n` sorts everything and zeroes nothing) followed by + /// `sample_multinomial` on the same rng seed. Same token, one sort vs + /// zero. (`top_k = VOCAB` is not a usable witness: it now โ€” correctly โ€” + /// takes the skip as well, being the same no-op filter.) + #[test] + fn server_default_top_k_skips_the_unconsumed_sort() { + let s = sampler(-1, 1.0, 0.0); + let (token_skip, sorts_skip) = sample_counting_sorts(&s, 42); + assert_eq!(sorts_skip, 0, "server defaults must skip the sort entirely"); + + // Pre-fix path, spelled out: softmax(logits / T), full sort with + // zero-rest, multinomial. Filters are all out of range, so the sort's + // output feeds nothing โ€” exactly the work the skip removes. + let probs_t = + candle_nn::ops::softmax_last_dim(&(logits().to_dtype(candle_core::DType::F32).unwrap() + / 0.7) + .unwrap()) + .unwrap(); + let mut probs: Vec = probs_t.to_vec1().unwrap(); + let n = probs.len(); + let before = PARTIAL_SORT_CALLS.with(|c| c.get()); + let sorted = super::partial_sort_top_k(&mut probs, n, true); + assert_eq!( + PARTIAL_SORT_CALLS.with(|c| c.get()) - before, + 1, + "the witness must actually sort, or this test proves nothing" + ); + assert_eq!(sorted.len(), VOCAB, "k = n returns every element sorted"); + let rng = Arc::new(Mutex::new(Isaac64Rng::seed_from_u64(42))); + let token_sort = s.sample_multinomial(&probs, false, rng).unwrap().token; + + assert_eq!( + token_skip, token_sort, + "skipping the sort must not change the drawn token (identical probs, \ + identical rng stream)" + ); + // Pin the drawn token so a silent change to the rng stream (e.g. an + // extra draw before the multinomial) cannot pass as parity. + assert_eq!(token_skip, 506, "seeded draw moved โ€” the rng stream shifted"); + } + + /// MUTATION CONTROLS for the skip condition: each filter that consumes the + /// sorted pairs must, on its own, keep the sort alive. + #[test] + fn each_active_filter_keeps_the_sort_alive() { + // top_p in (0,1): the cumsum walk needs the sorted pairs. + let (_, sorts) = sample_counting_sorts(&sampler(-1, 0.9, 0.0), 7); + assert_eq!(sorts, 1, "active top_p must sort"); + // min_p in (0,1): the threshold scan needs the sorted pairs. + let (_, sorts) = sample_counting_sorts(&sampler(-1, 1.0, 0.1), 7); + assert_eq!(sorts, 1, "active min_p must sort"); + // top_k < vocab: the zero-rest pass needs the partition. + let (_, sorts) = sample_counting_sorts(&sampler(32, 1.0, 0.0), 7); + assert_eq!(sorts, 1, "truncating top_k must sort"); + } + + /// A DRY config left at its disabled default (`multiplier: 0.0`) is a + /// proven no-op in `apply_dry_penalty`, so it must not count as a penalty + /// โ€” before this fix `has_penalties()` answered true for it, which + /// disqualified the GPU fast path (and `is_raw_argmax`) for every request + /// that merely *carried* the default DRY struct. + /// + /// `is_raw_argmax` is the public reader of `has_penalties`; both + /// directions are pinned: disabled DRY โ‡’ still raw-argmax, enabled DRY โ‡’ + /// not. + #[test] + fn disabled_dry_is_not_a_penalty_and_enabled_dry_is() { + // Dry params only survive `Sampler::new` when a tokenizer is present. + let tokenizer = Arc::new(tokenizers::Tokenizer::new( + tokenizers::models::bpe::BPE::default(), + )); + let mk = |dry: Option| { + Sampler::new( + None, + 0, + Some(tokenizer.clone()), + None, + None, + None, + dry, + -1, + 1.0, + 0.0, + None, + vec![], + None, + ) + .unwrap() + }; + + let none = mk(None); + assert!(none.is_raw_argmax(), "baseline: no DRY at all"); + + let disabled = mk(Some(DrySamplingParams::default())); + assert!( + disabled.is_raw_argmax(), + "multiplier = 0.0 (the disabled default) cannot rewrite a logit; \ + it must not disqualify the fast path" + ); + + let enabled = mk(Some(DrySamplingParams { + multiplier: 0.8, + ..Default::default() + })); + assert!( + !enabled.is_raw_argmax(), + "an active DRY penalty rewrites logits and must disqualify" + ); + + // And the no-op claim itself: with a disabled DRY the sampled token + // equals the no-DRY token on the same fixture (greedy, so exact). + let rng = || Arc::new(Mutex::new(Isaac64Rng::seed_from_u64(0))); + let ctx = [5u32, 5, 5, 5]; + let t_none = none + .sample(logits(), &ctx, false, rng(), false, false) + .unwrap() + .token; + let t_disabled = disabled + .sample(logits(), &ctx, false, rng(), false, false) + .unwrap() + .token; + assert_eq!(t_disabled, t_none, "disabled DRY must be a behavioural no-op"); + } + + /// `device_sampling_terminates` must mirror `sample`'s dispatch exactly โ€” + /// each disqualifier flips it independently, and the sort/vocab geometry + /// matches the fast path's own bounds. + #[test] + fn device_eligibility_mirrors_the_fast_path_dispatch() { + let big_vocab = 129_280; // V4: needs_sort can't fit shared memory + let small_vocab = 4_096; // 4096.next_power_of_two() * 4 = 16 KB <= 48 KB + + // Greedy, no penalties: eligible at any vocab, any backend. + let greedy = Sampler::new( + None, + 0, + None, + None, + None, + None, + None, + -1, + 1.0, + 0.0, + None, + vec![], + None, + ) + .unwrap(); + assert!(greedy.device_sampling_terminates(big_vocab, false, false)); + assert!(greedy.device_sampling_terminates(big_vocab, true, false)); + // return_logprobs needs the CPU pipeline. + assert!(!greedy.device_sampling_terminates(big_vocab, true, true)); + + // Temperature without a sort (server defaults): eligible. + let no_sort = sampler(-1, 1.0, 0.0); + assert!(no_sort.device_sampling_terminates(big_vocab, false, false)); + + // Temperature with a sort: small vocab fits candle's shared-memory + // arg_sort on any backend; big vocab needs the CUDA radix path. + let sorted = sampler(32, 1.0, 0.0); + assert!(sorted.device_sampling_terminates(small_vocab, false, false)); + assert_eq!( + sorted.device_sampling_terminates(big_vocab, true, false), + cfg!(feature = "cuda"), + "big-vocab sorted sampling terminates on device iff the CUDA radix \ + top-k exists in this build" + ); + assert!(!sorted.device_sampling_terminates(big_vocab, false, false)); + + // top_k above the radix dispatch cap: the kernel declines statically. + let huge_k = sampler(2048, 1.0, 0.0); + assert!(!huge_k.device_sampling_terminates(big_vocab, true, false)); + + // A penalty disqualifies regardless of geometry. + let penalised = Sampler::new( + Some(0.7), + 0, + None, + Some(0.5), + None, + None, + None, + -1, + 1.0, + 0.0, + None, + vec![], + None, + ) + .unwrap(); + assert!(!penalised.device_sampling_terminates(big_vocab, true, false)); + } +} + +/// $0 host-cost probe at V4's vocabulary. +/// +/// **This times CPU code on a CPU tensor. It is evidence about HOST cost +/// only โ€” it validates nothing about GPU behaviour (D14).** It exists to put +/// a number on what one sequence pays per token when it lands on the CPU +/// sampling pipeline, under (i) bench params (`top_k = 32`) and (ii) the +/// server defaults (`top_k = -1`, `top_p = 1.0`, `min_p = 0.0` โ€” +/// `engine/add_request.rs`), which before the sort-skip fix ran a full +/// 129,280-element `sort_unstable_by` whose result nothing consumed. +/// +/// Deliberately `#[ignore]`d: it is a manual instrument, not a gate. Run with +/// `cargo test -p mistralrs-core --release host_cost_probe -- --ignored --nocapture`. +/// +/// The parity fixture above stays at `VOCAB = 4096` on purpose: its pure-top-p +/// cases depend on nucleus containment in the top 1024, which does not hold +/// for this spread at 129,280. +#[cfg(test)] +mod host_cost_probe { + use super::Sampler; + use candle_core::{Device, Tensor}; + use rand::SeedableRng; + use rand_isaac::Isaac64Rng; + use std::sync::{Arc, Mutex}; + + /// DeepSeek-V4's vocabulary. + const PROBE_VOCAB: usize = 129_280; + + /// Same LCG as `topk_parity_tests::peaked_logits`, at probe size. + fn probe_logits(seed: u64) -> Vec { + let mut state: u64 = seed ^ 0xDEAD_BEEF_CAFE_F00D; + (0..PROBE_VOCAB) + .map(|_| { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((state >> 32) as u32 as f32 / u32::MAX as f32) * 12.0 - 6.0 + }) + .collect() + } + + fn probe_sampler(top_k: i64, top_p: f64, min_p: f64) -> Sampler { + Sampler::new( + Some(0.7), + 0, + None, + None, + None, + None, + None, + top_k, + top_p, + min_p, + None, + vec![], + None, + ) + .unwrap() + } + + fn time_us_per_call(sampler: &Sampler) -> (f64, f64) { + const WARMUP: usize = 3; + const ITERS: usize = 30; + let logits = Tensor::from_vec(probe_logits(7), PROBE_VOCAB, &Device::Cpu).unwrap(); + let rng = Arc::new(Mutex::new(Isaac64Rng::seed_from_u64(42))); + let ctx = [1u32, 2, 3]; + for _ in 0..WARMUP { + sampler + .sample(logits.clone(), &ctx, false, rng.clone(), false, true) + .unwrap(); + } + let mut samples_us = Vec::with_capacity(ITERS); + for _ in 0..ITERS { + let t0 = std::time::Instant::now(); + sampler + .sample(logits.clone(), &ctx, false, rng.clone(), false, true) + .unwrap(); + samples_us.push(t0.elapsed().as_secs_f64() * 1e6); + } + let mean = samples_us.iter().sum::() / ITERS as f64; + let min = samples_us.iter().cloned().fold(f64::INFINITY, f64::min); + (mean, min) + } + + #[test] + #[ignore = "manual $0 probe: prints per-call CPU sampling cost at V4 vocab; \ + run with --ignored --nocapture. Host-cost evidence only, never \ + GPU validation (D14)."] + fn time_cpu_sample_at_v4_vocab() { + let (bench_mean, bench_min) = time_us_per_call(&probe_sampler(32, 1.0, 0.0)); + let (srv_mean, srv_min) = time_us_per_call(&probe_sampler(-1, 1.0, 0.0)); + println!( + "HOST_PROBE vocab={PROBE_VOCAB} bench(top_k=32): mean={bench_mean:.0}us min={bench_min:.0}us" + ); + println!( + "HOST_PROBE vocab={PROBE_VOCAB} server(top_k=-1,top_p=1.0): mean={srv_mean:.0}us min={srv_min:.0}us" + ); + } +} + /// The greedy tie-break contract. /// /// Arc has three greedy implementations and they did not agree with each other diff --git a/mistralrs-core/tests/capability_reachability.rs b/mistralrs-core/tests/capability_reachability.rs index 8a50d68b7..c005e0153 100644 --- a/mistralrs-core/tests/capability_reachability.rs +++ b/mistralrs-core/tests/capability_reachability.rs @@ -313,6 +313,19 @@ static REGISTRY: &[Capability] = &[ }, status: Status::Live, }, + Capability { + name: "sampling: ARC_SAMPLE_ON_DEVICE keeps batched logits on the device", + parent: "ArcInfer / ArcSample", + check: Check::Symbol { + symbol: "sample_on_device_enabled", + defined_in: "mistralrs-core/src/pipeline/mod.rs", + }, + // Default OFF (env_flag_is_set, value-parsed), consulted by the + // DefaultInstructions step arm before the batched logits host copy. + // If this goes dark, every batch B > 1 is back to a 66 MB/step serial + // D2H that also disqualifies all rows from the GPU sampler fast path. + status: Status::Live, + }, // -- the ApiPromise class ------------------------------------------------ Capability { name: "sampling: logit_bias reaches the sampler", From 5937a99e73623b356dd3bcd1d94c30326665c3aa Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 22:23:14 +0100 Subject: [PATCH 19/21] =?UTF-8?q?perf(ArcKernels/ArcQuant):=20FP8=20dense-?= =?UTF-8?q?GEMM=20lane=20=E2=80=94=20pow2=20scale=20shifts,=20dispatch=20e?= =?UTF-8?q?ngagement=20log,=20three-arm=20sweep=20harness,=20SASS=20probe?= =?UTF-8?q?=20(#218)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(ArcKernels): the cuBLASLt FP8 branch never flattened rank-3 to 2-D The native branch flattens `[B, T, hidden]` to `[B*T, hidden]` before its GEMM. The dequantize + cuBLASLt branch did not โ€” it passed `x` through untouched. V4's activation is rank-3, and bias is `None` on every V4 linear, so `UnquantLinear::forward` took `w.broadcast_left(B)` into `cublaslt.batch_matmul` with `stride_b = 0`. At decode T=1, making it a **B-way batched GEMM with m=1**: a batched GEMV wearing a GEMM's name, and the one shape cuBLASLt has no advantage in. The documented 27x was measured at prefill, where `[1, 512, hidden]` collapses to batch=1 and the call is a real GEMM. This invalidated the threshold A/B rather than merely slowing it: lowering `ARC_FP8_CUBLAS_MIN_M` selected the batched-GEVM shape, so the sweep was comparing the tiled kernel against a degenerate call and would have supported "cuBLASLt loses at decode" โ€” a conclusion about a shape the change was never meant to select. Flattening makes the two branches comparable, which is the premise the A/B rests on. Co-Authored-By: Claude Opus 5 (1M context) * perf(ArcKernels): pow2 shift fast path for the tiled FP8 GEMM's per-element scale divisions get_scale in fp8_matmul_tiled performs TWO runtime signed integer divisions per weight element staged into shared memory (n / block_size_y, k / block_size_x). Both divisors are kernel arguments, so nvcc expands each into a ~15-20 instruction reciprocal-and-fixup sequence โ€” paid once per element of every weight tile, ~33.9G times per B=256 decode step by the source-derived count. The shipped weight_block_size is [128, 128]: both powers of two. This adds a POW2_SCALE template arm: the host launcher tests pow2-ness once per launch, computes shift = log2(block_size) on the host, and the kernel indexes the scale grid with two SHFs instead. Non-power-of-two geometries keep the division path, unchanged. Bit-identical by construction: for the non-negative indices used here, n >> log2(d) == n / d for every input โ€” same quotient, same scale word fetched, same arithmetic after it. Dispatching on POW2_SCALE can never change what an A/B leg measures, only how fast the scalar arm runs. The C ABI (launch_fp8_matmul_{f16,bf16}) is unchanged; no Rust edits needed. Compile-unverified locally (macOS has no nvcc); covered by the nvcc CI lane (.github/workflows/cuda_compile_check.yaml), which builds this TU for sm_80/sm_90 on every PR touching it. The bank-conflict half of this lane's brief (+4 -> +1 tile padding) was already landed by PR #215 and is an ancestor of this branch; the kernel's stride-33 comment documents the bank arithmetic. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf * feat(ArcLab): ARC_LOG_FP8_DISPATCH โ€” engagement lines proving WHICH FP8 kernel served a forward Every threshold sweep on this lane needs an engagement assertion: a server-log line proving which kernel actually ran, per leg. None existed โ€” a leg could silently measure the wrong arm (the ARC_NO_DEDICATED_DECODE incident produced two void A/Bs exactly this way, both arms running the same code). ARC_LOG_FP8_DISPATCH=1 (read by VALUE via env_flag_is_set, latched once per process) makes each dispatch path print ONE line per process: [arc-fp8-dispatch] path= first_shape=mM_nN_kK Sites: the native dispatch chain in fp8_blockwise_matmul_impl (both dtype arms, mirroring the if/else exactly) and the ARC_FP8_CUBLAS_MIN_M divert to dequantize+cuBLASLt in BlockwiseFP8Linear::forward. Cost when unset: one latched bool read. When set: one relaxed atomic swap per call โ€” no locks, so enabling it does not perturb the rates a sweep records. Registered in mistralrs-core/tests/capability_reachability.rs (Status::Live); the registry run passes (11/11). The Rust is cfg(cuda) โ€” compile-unverified locally on macOS; covered by the cuda-typecheck job of the nvcc CI lane. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf * feat(ArcLab): three-arm FP8 threshold sweep harness โ€” the measurement allowed to move ARC_FP8_CUBLAS_MIN_M arc_fp8_cublas_min_m's doc and the WMMA kernel's doc both forbid merging a threshold from a two-arm result: the only sweep ever taken predates both the WMMA GEMM (#200) and the rank-3 flatten, so it compared cuBLASLt-as-batched- GEMV against the scalar kernel only. This script is the three-arm re-run they demand, on ONE binary, env toggles only, fresh server per leg (the gates are OnceLock-latched, so re-exporting into a live server would silently measure the previous leg). Design: the threshold candidates {5, 8, 64, 256, 512} are swept as the OFFERED DECODE BATCH M (a grid that varies only the env at one fixed batch routes every MIN_M <= batch leg to the same kernel and cannot name a crossover in M), and each ladder point runs all three arms: (a) cublaslt: ARC_FP8_CUBLAS_MIN_M=5 (b) wmma: ARC_FP8_CUBLAS_MIN_M=1000000, ARC_NO_FP8_WMMA=0 (c) tiled: ARC_FP8_CUBLAS_MIN_M=1000000, ARC_NO_FP8_WMMA=1 A leg counts only if ALL of: 1. ENGAGEMENT โ€” its expected [arc-fp8-dispatch] path= line present in the server log and both rivals ABSENT (a WMMA leg whose log shows tiled is eligibility silently failing: VOID, loudly); 2. FLOOR โ€” summed usage.completion_tokens >= MIN_TOKENS_FLOOR (default 1000); the rate divides by tokens the server SAYS it generated, never by requested max_tokens; 3. CANARY โ€” greedy fixed-prompt stream vs the baseline leg (tiled @ M=5); cross-arm divergence is expected numerics and is REPORTED with its first index; a degenerate stream (<5 tokens / no finish_reason) voids the leg. The summary table names, per M, cublaslt vs min(wmma, tiled) and the FIRST clean M where cuBLASLt wins โ€” the value ARC_FP8_CUBLAS_MIN_M may move to. No clean crossover => the default stays 512. Script only โ€” NOT run here (no GPU; D14). House lock/preflight/provenance discipline copied from arcspec_token_identity_b8.sh. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf * ci(ArcGate): free SASS probe for the FP8 dense-GEMM lane โ€” LDS width census + WMMA HMMA presence Two questions this settles with nvcc + cuobjdump alone, no GPU: 1. Whether the 4-way bank-conflict cost model for fp8_matmul_tiled describes the compiled artifact: the model is a claim about SCALAR 32-bit LDS in the inner product, and if ptxas had vectorized those shared loads to LDS.128 its arithmetic would be about instructions that do not exist. The probe prints a per-kernel LDS/LDS.64/LDS.128 census. (Note: the +1 pad that kills the conflict also forecloses LDS.128 on tile rows โ€” scalar LDS is the EXPECTED reading on current source.) 2. Whether the never-executed tensor-core GEMM (blockwise_fp8_gemm_wmma.cu) actually contains tensor-core instructions: HARD FAIL if the TU has zero HMMA/HGMMA โ€” the three-arm sweep would otherwise compare cuBLASLt against two scalar kernels while calling one of them WMMA. Teeth: exactly 4 fp8_matmul_tiled kernels asserted (2 dtypes x 2 POW2_SCALE instantiations โ€” a dropped instantiation makes the census about a kernel not in the build), and every tiled kernel must show >=1 LDS or the extractor lost the function body. Wired into the nvcc CI lane as a step beside the QTIP spill gate, so it runs for sm_80 and sm_90 on every Rust/kernel PR; also runnable standalone: arc-tools/fp8_gemm_sass_check.sh sm_90 Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf * fix(ArcGate): SASS probe โ€” replace GNU-awk-unsupported \< \> ERE escapes with cuobjdump -fun extraction The probe's first CI run went red on itself, honestly: GNU awk treats \< \> (and warns on \.) as plain characters in ERE, so the awk body-splitter's counting regexes matched NOTHING, the per-function census returned zero LDS for all 4 fp8_matmul_tiled instantiations, and the fail-on-zero guard refused to trust its own zero โ€” which is exactly what it exists for. This revision removes regex body-splitting entirely: the mangled kernel names come from the Function listing (awk '{print $NF}', no regex), and `cuobjdump -sass -fun ` extracts each function's SASS directly. All counting is grep -cE with POSIX classes and bracket expressions only โ€” no \< \> anywhere. Also, per review: * The two failure modes are now DISTINGUISHED: FAIL[extractor] (no instruction lines came back โ€” the census is unverified) vs FAIL[zero-LDS] (body extracted, genuinely no shared loads โ€” the conflict model has no subject). Body presence is judged on /**/ instruction lines, so the two cannot be confused. * The VERDICT[tiled] line prints ONLY when the per-function census is complete and non-empty; otherwise it prints NOT ESTABLISHED. A TU-level cross-check line (whole cubin, includes the GEMV/MoE kernels) is always printed, labeled informational. * HMMA counting moved off the \< \> GNUism onto the same POSIX pattern. Verified locally against a stub cuobjdump with a realistic SASS fixture: happy path (4 kernels, LDSM excluded, LDS/LDS.64/LDS.128 counted per function, scalar arithmetic correct, exit 0) plus three negative controls (broken -fun -> FAIL[extractor] + verdict NOT ESTABLISHED; LDS-free body -> FAIL[zero-LDS]; HMMA-stripped WMMA TU -> hard fail), all exit 1. The real cuobjdump run is CI's โ€” if -fun ever returns nothing there, the probe fails loudly rather than lying. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf --------- Co-authored-by: Claude Opus 5 (1M context) --- .github/workflows/cuda_compile_check.yaml | 10 + arc-tools/fp8_gemm_sass_check.sh | 209 +++++++++ arc-tools/fp8_threshold_sweep.sh | 412 ++++++++++++++++++ .../tests/capability_reachability.rs | 15 + .../blockwise_fp8/blockwise_fp8_gemm.cu | 118 +++-- mistralrs-quant/src/blockwise_fp8/mod.rs | 42 ++ mistralrs-quant/src/blockwise_fp8/ops.rs | 111 ++++- 7 files changed, 888 insertions(+), 29 deletions(-) create mode 100755 arc-tools/fp8_gemm_sass_check.sh create mode 100755 arc-tools/fp8_threshold_sweep.sh diff --git a/.github/workflows/cuda_compile_check.yaml b/.github/workflows/cuda_compile_check.yaml index 2aa62730b..2bd1c949c 100644 --- a/.github/workflows/cuda_compile_check.yaml +++ b/.github/workflows/cuda_compile_check.yaml @@ -305,6 +305,16 @@ jobs: - name: QTIP beam spill gate (LOCAL:0, all geometries) run: bash arc-tools/qtip_beam_res_usage_check.sh sm_${{ matrix.compute_cap }} + # Same free-instrument class as the QTIP gate above, for the blockwise- + # FP8 dense-GEMM lane: asserts all 4 fp8_matmul_tiled instantiations + # exist (2 dtypes x 2 POW2_SCALE scale-index arms), prints the LDS width + # census that decides whether the bank-conflict cost model describes the + # compiled kernel, and HARD-FAILS if the never-executed WMMA GEMM + # contains no tensor-core (HMMA/HGMMA) instructions. nvcc + cuobjdump + # only; no GPU. + - name: FP8 GEMM SASS probe (LDS census + WMMA HMMA presence) + run: bash arc-tools/fp8_gemm_sass_check.sh sm_${{ matrix.compute_cap }} + - name: Compile QTIP CUDA kernels (mistralrs-quant) run: cargo build -p mistralrs-quant --lib --features cuda diff --git a/arc-tools/fp8_gemm_sass_check.sh b/arc-tools/fp8_gemm_sass_check.sh new file mode 100755 index 000000000..48634f447 --- /dev/null +++ b/arc-tools/fp8_gemm_sass_check.sh @@ -0,0 +1,209 @@ +#!/usr/bin/env bash +# Parent system: ArcKernels + ArcGate. +# +# FREE SASS PROBE for the blockwise-FP8 dense-GEMM lane. NO GPU REQUIRED: +# nvcc + cuobjdump only, so it runs on the GPU-less nvcc CI lane +# (.github/workflows/cuda_compile_check.yaml) and on any dev box with the +# toolkit. +# +# WHAT IT SETTLES +# --------------- +# 1. `fp8_matmul_tiled` (blockwise_fp8_gemm.cu) โ€” HOW ptxas actually loads the +# shared tiles in the inner product `acc += s_input[ty][k] * s_weight[tx][k]`. +# The source-derived cost model for this kernel (shared traffic 4.13x over +# supply x a 4-way bank conflict from the old +4 tile padding) is a claim +# about SCALAR 32-bit LDS. If ptxas had already vectorized those loads to +# LDS.128, the conflict arithmetic would be about instructions that do not +# exist and the model's 7% agreement would be a coincidence. This probe +# prints the per-kernel LDS width census so that question is settled by the +# compiled artifact, not by reading C++. +# NOTE the padding interaction: the +1 pad (stride 33 floats, PR #215) +# that kills the 4-way bank conflict ALSO forecloses LDS.128 on tile rows +# (rows are only 4-byte aligned at stride 33). Scalar LDS here is +# therefore the EXPECTED outcome on current source, not a defect โ€” the +# probe exists so a future retune that trades padding for vector width +# does it with the census in front of them. +# 2. `blockwise_fp8_gemm_wmma.cu` โ€” whether the "tensor-core" GEMM actually +# contains tensor-core instructions (HMMA/HGMMA). That kernel has NEVER +# executed on any GPU; if ptxas lowered it without a single matrix-multiply +# instruction, every derivation in its header is void. HARD FAIL if absent. +# +# TEETH (a parser that matches nothing must not pass) +# --------------------------------------------------- +# * Kernel-count assertion: the tiled TU must yield exactly 4 +# `fp8_matmul_tiled` kernels (2 dtypes x 2 POW2_SCALE instantiations โ€” +# see the launcher in blockwise_fp8_gemm.cu). A count drift means an +# instantiation was dropped and the census is about a kernel not in the +# build. +# * Every tiled kernel must show at least one LDS โ€” a tiled shared-memory GEMM +# with zero shared loads means the SASS extractor lost the function body. +# +# Usage: arc-tools/fp8_gemm_sass_check.sh [sm_90] +set -euo pipefail + +ARCH="${1:-sm_90}" +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +KDIR="${REPO_ROOT}/mistralrs-quant/kernels/blockwise_fp8" +OUT="$(mktemp -d)" +trap 'rm -rf "${OUT}"' EXIT + +# Mirror mistralrs-quant/build.rs exactly: a different flag set is a different +# instruction selection, and therefore a probe of a kernel we do not ship. +NVCC_ARGS=( + -std=c++17 + -O3 + -U__CUDA_NO_HALF_OPERATORS__ + -U__CUDA_NO_HALF_CONVERSIONS__ + -U__CUDA_NO_HALF2_OPERATORS__ + -U__CUDA_NO_BFLOAT16_CONVERSIONS__ + --expt-relaxed-constexpr + --expt-extended-lambda + --use_fast_math +) + +echo "== blockwise-FP8 GEMM SASS probe ==" +nvcc --version | tail -1 +echo "arch: ${ARCH}" +echo + +# --------------------------------------------------------------------------- +# Per-function SASS extraction. +# +# PORTABILITY LESSON (this script's first CI run went red on itself): GNU awk +# does not honor \< \> word-boundary escapes in ERE โ€” it warns "escape +# sequence treated as plain character" and matches nothing โ€” so the original +# awk body-splitter returned zero LDS for every function and the fail-on-zero +# guard below correctly refused to trust it. This revision does no regex body +# splitting at all: cuobjdump lists the mangled kernel names, and +# `cuobjdump -sass -fun ` extracts exactly one function's SASS. All +# counting is grep -cE with POSIX classes and bracket expressions only โ€” +# no \< \> anywhere. +# --------------------------------------------------------------------------- +fn_names() { # $1 = cubin, $2 = fixed-string name filter + cuobjdump -sass "$1" | grep 'Function : ' | awk '{print $NF}' | grep -F "$2" || true +} +count() { # $1 = file, $2 = POSIX ERE + grep -cE "$2" "$1" || true +} +# SASS instruction lines carry a /**/ comment; body-vs-empty is +# judged on those, so "extractor returned nothing" and "zero LDS" cannot be +# confused with each other. +RE_INSN='/[*][0-9a-fA-F]+[*]/' +# LDSM ('M' after LDS) must not count; the char after LDS is '.' or space. +RE_LDS_ALL='[[:space:]]LDS[[:space:].]' +RE_LDS_128='[[:space:]]LDS[.](U[.])?128[[:space:]]' +RE_LDS_64='[[:space:]]LDS[.](U[.])?64[[:space:]]' +RE_MMA='[[:space:]](HMMA|HGMMA)[[:space:].]' + +# --------------------------------------------------------------------------- +# 1. The scalar tiled GEMM: LDS width census. +# --------------------------------------------------------------------------- +echo "--- blockwise_fp8_gemm.cu : fp8_matmul_tiled ---" +nvcc -cubin -arch="${ARCH}" "${NVCC_ARGS[@]}" \ + -o "${OUT}/gemm.cubin" "${KDIR}/blockwise_fp8_gemm.cu" + +if command -v c++filt >/dev/null 2>&1; then + cuobjdump -sass "${OUT}/gemm.cubin" | grep 'Function : ' | c++filt | grep fp8_matmul_tiled || true +fi + +fail=0 +extract_ok=1 +n_tiled=0 +sum128=0 +names="$(fn_names "${OUT}/gemm.cubin" fp8_matmul_tiled)" +if [ -z "${names}" ]; then + echo "FAIL[extractor]: no fp8_matmul_tiled function names in the cubin's" + echo " SASS listing โ€” the name pass found nothing, so there is no census." + fail=1 + extract_ok=0 +fi + +echo "kernel insns LDS(all) LDS.128 LDS.64 scalar" +for name in ${names}; do + n_tiled=$((n_tiled + 1)) + sfile="${OUT}/tiled_${n_tiled}.sass" + cuobjdump -sass -fun "${name}" "${OUT}/gemm.cubin" > "${sfile}" 2>/dev/null || true + insns="$(count "${sfile}" "${RE_INSN}")" + lds_all="$(count "${sfile}" "${RE_LDS_ALL}")" + w128="$(count "${sfile}" "${RE_LDS_128}")" + w64="$(count "${sfile}" "${RE_LDS_64}")" + scalar=$((lds_all - w128 - w64)) + sum128=$((sum128 + w128)) + printf '%-48s %6s %8s %7s %6s %6s\n' "${name:0:48}" "${insns}" "${lds_all}" "${w128}" "${w64}" "${scalar}" + if [ "${insns}" -eq 0 ]; then + echo "FAIL[extractor]: ${name}: 'cuobjdump -sass -fun' returned no" + echo " instruction lines. This is an EXTRACTION failure, not a zero-LDS" + echo " kernel โ€” this function's census row is unverified." + fail=1 + extract_ok=0 + elif [ "${lds_all}" -eq 0 ]; then + echo "FAIL[zero-LDS]: ${name}: body extracted (${insns} instructions) but" + echo " it contains no shared-memory loads at all. A shared-tile GEMM" + echo " without LDS means the wrong function was extracted or the kernel" + echo " no longer stages tiles โ€” either way the conflict model has no" + echo " subject here." + fail=1 + fi +done + +if [ "${n_tiled}" -ne 4 ]; then + echo "FAIL: expected 4 fp8_matmul_tiled kernels (2 dtypes x 2 POW2_SCALE" + echo " instantiations), found ${n_tiled}. An instantiation was dropped;" + echo " the census is about kernels that are not in the build." + fail=1 +fi + +# TU-level cross-check: computed on the WHOLE cubin dump, independent of the +# per-function extractor, so it survives an extractor failure and gives the +# reader a signal either way. Informational โ€” the per-function census is the +# only thing the VERDICT may rest on (the TU includes the GEMV/MoE kernels). +cuobjdump -sass "${OUT}/gemm.cubin" > "${OUT}/gemm.sass" 2>/dev/null || true +echo "TU-level cross-check (whole cubin, informational):" \ + "LDS(all)=$(count "${OUT}/gemm.sass" "${RE_LDS_ALL}")" \ + "LDS.128=$(count "${OUT}/gemm.sass" "${RE_LDS_128}")" \ + "LDS.64=$(count "${OUT}/gemm.sass" "${RE_LDS_64}")" + +# The VERDICT prints ONLY on a complete, non-empty per-function census. +# Printing it off TU-level counts (or off an empty extraction) would be a +# lying instrument: confident prose about numbers that were never measured. +if [ "${extract_ok}" -eq 1 ] && [ "${n_tiled}" -eq 4 ] && [ "${fail}" -eq 0 ]; then + if [ "${sum128}" -eq 0 ]; then + echo "VERDICT[tiled]: scalar LDS only, no LDS.128 โ€” the shared loads are" + echo " 32-bit, so the bank-conflict arithmetic (stride mod 32) is about" + echo " instructions that really exist, and the +1-pad fix is live." + else + echo "VERDICT[tiled]: ptxas emitted LDS.128 โ€” the scalar-LDS conflict model" + echo " does NOT describe this build; re-derive before quoting it." + fi +else + echo "VERDICT[tiled]: NOT ESTABLISHED โ€” the per-function census is empty or" + echo " partial (see FAILs above). Do not quote any LDS conclusion from this" + echo " run; the TU-level line above is a cross-check, not a verdict." +fi +echo + +# --------------------------------------------------------------------------- +# 2. The WMMA GEMM: tensor-core instructions must exist. +# --------------------------------------------------------------------------- +echo "--- blockwise_fp8_gemm_wmma.cu : HMMA presence ---" +nvcc -cubin -arch="${ARCH}" "${NVCC_ARGS[@]}" \ + -o "${OUT}/wmma.cubin" "${KDIR}/blockwise_fp8_gemm_wmma.cu" + +cuobjdump -sass "${OUT}/wmma.cubin" > "${OUT}/wmma.sass" 2>/dev/null || true +hmma="$(count "${OUT}/wmma.sass" "${RE_MMA}")" +echo "HMMA/HGMMA instructions in TU: ${hmma}" +if [ "${hmma}" -eq 0 ]; then + echo "FAIL: the tensor-core GEMM contains no HMMA/HGMMA. It has never run on" + echo " hardware; if it also has no tensor-core instructions, every" + echo " derivation in its header is void and the three-arm sweep would" + echo " be comparing cuBLASLt against two scalar kernels." + fail=1 +fi +echo + +if [ "${fail}" -ne 0 ]; then + echo "RESULT: FAILED โ€” see above." + exit 1 +fi +echo "RESULT: census complete (${ARCH}); 4 tiled kernels present (per-function" +echo " extraction verified non-empty), WMMA TU carries ${hmma} HMMA/HGMMA." diff --git a/arc-tools/fp8_threshold_sweep.sh b/arc-tools/fp8_threshold_sweep.sh new file mode 100755 index 000000000..05883dd13 --- /dev/null +++ b/arc-tools/fp8_threshold_sweep.sh @@ -0,0 +1,412 @@ +#!/bin/bash +# Parent system: ArcKernels + ArcQuant/Legacy (blockwise FP8). +# +# THREE-ARM FP8 DENSE-GEMM THRESHOLD SWEEP โ€” the measurement that is allowed +# to move `ARC_FP8_CUBLAS_MIN_M` off 512. +# +# โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +# WHY THREE ARMS, AND WHY A TWO-ARM RESULT MUST NOT BE MERGED +# +# `arc_fp8_cublas_min_m`'s own doc (blockwise_fp8/mod.rs) and the WMMA +# kernel's doc (blockwise_fp8/ops.rs, `fp8_wmma_enabled`) both record the +# standing instruction: the only sweep ever taken compared (a) dequantize + +# cuBLASLt against (c) the scalar tiled kernel, BEFORE #200 landed the +# (b) tensor-core WMMA GEMM โ€” and before the rank-3 flatten, so arm (a) was a +# B-way batched GEMV wearing a GEMM's name. Freezing a threshold from that +# sweep would leave (b) dead code at every M that matters, silently. This +# script is the three-arm re-run those docs demand: +# +# (a) dequantize_w() + cuBLASLt ARC_FP8_CUBLAS_MIN_M=5 +# (b) native FP8 tensor-core WMMA ARC_FP8_CUBLAS_MIN_M=1000000, ARC_NO_FP8_WMMA=0 +# (c) native scalar fp8_matmul_tiled ARC_FP8_CUBLAS_MIN_M=1000000, ARC_NO_FP8_WMMA=1 +# +# ONE binary. Every leg differs only in environment; a leg that needed a +# rebuild would be measuring a different artifact. The env gates are latched +# per process (OnceLock/LazyLock), so every leg gets a FRESH server process โ€” +# re-exporting into a running server changes nothing and would silently +# measure the previous leg's setting. +# +# โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +# THE LADDER: M VARIES BY OFFERED CONCURRENCY, NOT ONLY BY THRESHOLD +# +# The deliverable is "the first M at which (a) beats min((b), (c))". A grid +# that sweeps only ARC_FP8_CUBLAS_MIN_M at one fixed batch cannot name that M: +# every leg with MIN_M <= batch routes the same decode shape to the same +# kernel and the legs replicate each other. So the threshold candidates +# {5, 8, 64, 256, 512} are swept as the OFFERED DECODE BATCH โ€” M concurrent +# greedy completions, T=1 per step, so the dense projections see m_rows = M โ€” +# and each ladder point runs all three arms. 15 legs. +# +# โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +# WHAT MAKES A LEG COUNT (all three, or the leg is VOID) +# +# 1. ENGAGEMENT, asserted per leg. Every leg runs with ARC_LOG_FP8_DISPATCH=1; +# the dispatch sites print `[arc-fp8-dispatch] path=` once per path +# per process. The leg's expected kernel must be PRESENT in the server log +# and both rival kernels ABSENT. A leg that cannot prove which kernel +# served it is measuring an assumption (`ARC_NO_DEDICATED_DECODE`'s two +# void A/Bs are the record of what that costs). In particular: if the WMMA +# arm's log shows `path=tiled`, WMMA eligibility silently failed and the +# arm is the control wearing the treatment's name โ€” VOID, loudly. +# 2. A MINIMUM-TOKEN FLOOR. Rates on fewer than MIN_TOKENS_FLOOR generated +# tokens are warm-up noise. The rate divides by SUMMED +# `usage.completion_tokens` โ€” the tokens the server says it generated โ€” +# NEVER by requested max_tokens (a leg that stalls and returns short would +# otherwise report a healthy rate for work it did not do). +# 3. A TOKEN CANARY vs the baseline leg (arm (c) at M=5 โ€” the shipped scalar +# kernel, nearest to b=1). Greedy, fixed seed, fixed prompt. The arms are +# DIFFERENT NUMERICS by design (f32 scalar vs tensor-core promotion vs +# dequant-BF16 GEMM), so divergence at some index is expected and REPORTED, +# not failed; what hard-fails a leg is a degenerate stream (< 5 tokens or +# no finish_reason) โ€” a failure to generate, which no rate can offset. +# +# The final table names, per ladder M, the three clean rates and min((b),(c)), +# and then the FIRST M at which (a) beats that min โ€” the value +# ARC_FP8_CUBLAS_MIN_M should default to. If any arm at that M is void, the +# crossover is NOT ESTABLISHED and the default stays 512. +# +# USAGE (on the GPU box; do not run without a GPU โ€” D14) +# setsid nohup bash arc-tools/fp8_threshold_sweep.sh /root/logs/fp8_sweep.log 2>&1 & +# cat /root/logs/fp8sweep/summary.txt +# +set -u + +ARCENV="${ARCENV:-/root/arcenv.sh}" +WT="${WT:-/root/arc}" +BIN="${BIN:-$WT/target/release/mistralrs}" +MODEL="${MODEL:-/root/models/v4-src}" +UQFF="${UQFF:-}" # optional: --from-uqff when set +ARCH="${ARCH:-deepseekv4}" +PREFLIGHT="${PREFLIGHT:-/root/arc-tools/gpu_box_preflight.sh}" +OUT="${OUT:-/root/logs/fp8sweep}" +STATUS="${STATUS:-/root/logs/fp8_sweep.status}" +PORT="${PORT:-1241}" +SEED="${SEED:-1234}" +MAX_TOKENS="${MAX_TOKENS:-48}" # per request; the floor is on the SUM +MIN_TOKENS_FLOOR="${MIN_TOKENS_FLOOR:-1000}" +MAX_SEQS="${MAX_SEQS:-512}" # must be >= the largest ladder point +LADDER="${LADDER:-5 8 64 256 512}" # offered decode batch = candidate M +LOCK="${LOCK:-/root/locks/gpu.lock}" +LOCK_WAIT="${LOCK_WAIT:-2400}" +OWNER="fp8sweep" + +mkdir -p "$OUT" "$(dirname "$STATUS")" "$(dirname "$LOCK")" +: > "$STATUS" +say() { echo "[$(date -u +%H:%M:%SZ)] $*" | tee -a "$STATUS"; } +die() { say "ABORT: $*"; echo "ABORT: $*" >> "$OUT/summary.txt"; exit 2; } + +say "fp8 three-arm threshold sweep โ€” start" + +# โ”€โ”€ GPU lock with OWNER-MATCHED RELEASE and evidence-based stale reclaim +# (same discipline as arcspec_token_identity_b8.sh; see the incident notes +# there for why release refuses locks this process does not own). +release_lock() { + if [ -f "$LOCK" ] && grep -q "^$OWNER pid=$$ " "$LOCK" 2>/dev/null; then + rm -f "$LOCK" + fi +} +STALE_AFTER="${STALE_AFTER:-120}" +lock_is_stale() { + [ -f "$LOCK" ] || return 1 + local age holder pid used procs + age=$(( $(date +%s) - $(stat -c %Y "$LOCK" 2>/dev/null || echo 0) )) + [ "$age" -ge "$STALE_AFTER" ] || return 1 + holder=$(cat "$LOCK" 2>/dev/null) + pid=$(printf '%s' "$holder" | sed -n 's/.*pid=\([0-9]*\).*/\1/p') + if [ -n "$holder" ] && [ -n "$pid" ] && kill -0 "$pid" 2>/dev/null; then + return 1 + fi + used=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null | head -1) + procs=$(nvidia-smi --query-compute-apps=pid --format=csv,noheader 2>/dev/null | wc -l) + [ "${used:-1}" -lt 1024 ] && [ "${procs:-1}" -eq 0 ] +} +acquire_lock() { + local waited=0 + while :; do + if ( set -o noclobber; echo "$OWNER pid=$$ started=$(date -u +%FT%TZ)" > "$LOCK" ) 2>/dev/null; then + trap 'release_lock' EXIT INT TERM + say "lock: acquired $LOCK" + return 0 + fi + if lock_is_stale; then + say "lock: RECLAIMING stale lock (age>=${STALE_AFTER}s, owner absent or dead, GPU idle)" + rm -f "$LOCK" + continue + fi + if [ "$waited" -ge "$LOCK_WAIT" ]; then + say "lock: HELD BY -> $(cat "$LOCK" 2>/dev/null)" + return 1 + fi + [ "$waited" = 0 ] && say "lock: waiting for -> $(cat "$LOCK" 2>/dev/null)" + sleep 15; waited=$((waited + 15)) + done +} +acquire_lock || die "GPU lock held by another chain for >${LOCK_WAIT}s" + +[ -f "$ARCENV" ] && . "$ARCENV" +if [ -x "$PREFLIGHT" ]; then + "$PREFLIGHT" > "$OUT/preflight.log" 2>&1 || die "preflight refused the box" + say "preflight OK" +else + say "WARN: no preflight at $PREFLIGHT โ€” continuing without it" +fi +[ -x "$BIN" ] || die "$BIN not found โ€” build once, then sweep (one binary, no rebuild)" +SHA=$(git -C "$WT" rev-parse HEAD 2>/dev/null) +say "binary ref: $SHA" + +# โ”€โ”€ The prober: one canary + (M-1) filler prompts, all concurrent, all +# greedy. Non-streaming so `usage.completion_tokens` is authoritative. +cat > "$OUT/probe.py" <<'PYEOF' +"""Drive M concurrent greedy completions; report summed completion_tokens +(from usage โ€” NEVER from requested max_tokens), wall seconds over the +concurrent burst, and the canary text (request 0).""" +import argparse, json, sys, threading, time, urllib.request + +CANARY = "List the first twelve prime numbers, separated by commas, then stop." +FILLERS = [ + "Name the seven continents, one per line.", + "Count from one to twenty in words, separated by spaces.", + "Write the first eight letters of the alphabet, separated by dashes.", + "Name eight colours of the visible spectrum, comma separated.", + "List eight consecutive even numbers starting at two.", + "Name the days of the week in order, one per line.", + "List the first eight powers of two, comma separated.", +] + +def one(port, prompt, max_tokens, seed, timeout): + body = json.dumps({ + "model": "default", "prompt": prompt, "max_tokens": max_tokens, + "temperature": 0, "seed": seed, "stream": False, + }).encode() + req = urllib.request.Request( + f"http://127.0.0.1:{port}/v1/completions", data=body, + headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=timeout) as r: + d = json.load(r) + ch = (d.get("choices") or [{}])[0] + usage = d.get("usage") or {} + return { + "text": ch.get("text", ""), + "finish_reason": ch.get("finish_reason"), + "completion_tokens": usage.get("completion_tokens", 0), + } + +def burst(port, m, max_tokens, seed, timeout): + prompts = [CANARY] + [FILLERS[i % len(FILLERS)] for i in range(m - 1)] + out = [None] * m + errs = [] + def w(i): + try: + out[i] = one(port, prompts[i], max_tokens, seed, timeout) + except Exception as e: # noqa: BLE001 โ€” a leg must record, not crash + errs.append(f"req[{i}]: {e}") + out[i] = {"text": "", "finish_reason": None, "completion_tokens": 0} + t0 = time.monotonic() + ts = [threading.Thread(target=w, args=(i,)) for i in range(m)] + for t in ts: t.start() + for t in ts: t.join() + wall = time.monotonic() - t0 + return out, wall, errs + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--port", type=int, required=True) + ap.add_argument("--m", type=int, required=True) + ap.add_argument("--max-tokens", type=int, required=True) + ap.add_argument("--seed", type=int, required=True) + ap.add_argument("--timeout", type=int, default=1800) + ap.add_argument("--warmup", action="store_true") + a = ap.parse_args() + + if a.warmup: + # Small unmeasured round: load caches, emit the first-dispatch log + # lines, absorb ramp. Its tokens are never counted. + burst(a.port, min(a.m, 8), 16, a.seed, a.timeout) + json.dump({"warmup": True}, sys.stdout) + return + + out, wall, errs = burst(a.port, a.m, a.max_tokens, a.seed, a.timeout) + total = sum(r["completion_tokens"] for r in out) + json.dump({ + "m": a.m, + "wall_s": wall, + "completion_tokens_sum": total, # from usage โ€” the only numerator + "tok_per_s": (total / wall) if wall > 0 else 0.0, + "canary_text": out[0]["text"], + "canary_finish": out[0]["finish_reason"], + "canary_tokens": out[0]["completion_tokens"], + "errors": errs, + "short_streams": sum(1 for r in out if r["completion_tokens"] < 5), + }, sys.stdout) + +main() +PYEOF + +# โ”€โ”€ One leg = one fresh server process + warm-up + measured burst. +# $1 arm name, $2 ladder M, then env pairs for the arm. +run_leg() { + local arm="$1" m="$2"; shift 2 + local leg="${arm}.M${m}" + local slog="$OUT/server.$leg.log" + say "leg $leg โ€” starting server ($*)" + + local uqff_args=() + [ -n "$UQFF" ] && uqff_args=(--from-uqff "$UQFF") + # shellcheck disable=SC2086 + env RUST_LOG=info ARC_LOG_FP8_DISPATCH=1 "$@" \ + "$BIN" serve -p "$PORT" -m "$MODEL" -a "$ARCH" "${uqff_args[@]}" \ + --max-seqs "$MAX_SEQS" --prefix-cache-n 0 --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 "leg $leg โ€” server never came up"; tail -20 "$slog" | tee -a "$STATUS"; return 1; } + + # Provenance: the running server must be the code we think it is. + local logged + logged=$(grep -m1 -o 'git revision: [0-9a-f]*' "$slog" | awk '{print $3}') + if [ -n "$SHA" ] && [ -n "$logged" ] && [ "$logged" != "$SHA" ]; then + say "leg $leg โ€” PROVENANCE FAILED: server $logged != worktree $SHA" + kill -TERM "$spid" 2>/dev/null || true + return 1 + fi + + say "leg $leg โ€” warm-up (unmeasured)" + python3 "$OUT/probe.py" --port "$PORT" --m "$m" --max-tokens "$MAX_TOKENS" \ + --seed "$SEED" --warmup > /dev/null || { say "leg $leg โ€” warm-up failed"; kill -TERM "$spid" 2>/dev/null; return 1; } + + say "leg $leg โ€” measured burst: M=$m x max_tokens=$MAX_TOKENS" + python3 "$OUT/probe.py" --port "$PORT" --m "$m" --max-tokens "$MAX_TOKENS" \ + --seed "$SEED" > "$OUT/$leg.json" || { say "leg $leg โ€” probe failed"; kill -TERM "$spid" 2>/dev/null; return 1; } + + 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 + + # ENGAGEMENT: expected kernel present, both rivals absent, in the server log. + local want forbid1 forbid2 + case "$arm" in + cublaslt) want="dequant_cublaslt"; forbid1="wmma"; forbid2="tiled" ;; + wmma) want="wmma"; forbid1="dequant_cublaslt"; forbid2="tiled" ;; + tiled) want="tiled"; forbid1="wmma"; forbid2="dequant_cublaslt" ;; + *) say "leg $leg โ€” unknown arm '$arm'"; return 1 ;; + esac + local ok=1 + grep -q "\[arc-fp8-dispatch\] path=$want " "$slog" || ok=0 + grep -q "\[arc-fp8-dispatch\] path=$forbid1 " "$slog" && ok=0 + grep -q "\[arc-fp8-dispatch\] path=$forbid2 " "$slog" && ok=0 + { + echo "leg=$leg" + grep '\[arc-fp8-dispatch\]' "$slog" | sed 's/^/ /' + echo " engagement=$([ "$ok" = 1 ] && echo PASS || echo FAIL)" + } >> "$OUT/engagement.txt" + if [ "$ok" != 1 ]; then + say "leg $leg โ€” ENGAGEMENT FAILED (want $want, forbid $forbid1/$forbid2) โ€” leg VOID" + echo '{"void": "engagement"}' > "$OUT/$leg.void" + fi + say "leg $leg โ€” done" + return 0 +} + +# โ”€โ”€ The 15 legs. Baseline (canary reference) FIRST: tiled @ smallest M. +MIN_M_NATIVE=1000000 # keeps every real m on the native path in arms (b),(c) +for m in $LADDER; do + run_leg tiled "$m" ARC_FP8_CUBLAS_MIN_M=$MIN_M_NATIVE ARC_NO_FP8_WMMA=1 \ + || die "tiled M=$m leg failed" +done +for m in $LADDER; do + run_leg wmma "$m" ARC_FP8_CUBLAS_MIN_M=$MIN_M_NATIVE ARC_NO_FP8_WMMA=0 \ + || die "wmma M=$m leg failed" +done +for m in $LADDER; do + run_leg cublaslt "$m" ARC_FP8_CUBLAS_MIN_M=5 \ + || die "cublaslt M=$m leg failed" +done + +# โ”€โ”€ Summary: floors, canaries, the crossover table, the verdict. +python3 - "$OUT" "$MIN_TOKENS_FLOOR" $LADDER <<'PYEOF' | tee "$OUT/summary.txt" +import json, os, sys + +out, floor = sys.argv[1], int(sys.argv[2]) +ladder = [int(x) for x in sys.argv[3:]] +ARMS = ["tiled", "wmma", "cublaslt"] + +def load(arm, m): + leg = f"{arm}.M{m}" + if os.path.exists(os.path.join(out, leg + ".void")): + return None, "VOID(engagement)" + p = os.path.join(out, leg + ".json") + if not os.path.exists(p): + return None, "MISSING" + with open(p) as f: + d = json.load(f) + if d.get("errors"): + return d, f"ERRORS({len(d['errors'])})" + if d["completion_tokens_sum"] < floor: + return d, f"VOID(floor {d['completion_tokens_sum']}<{floor})" + if d["canary_tokens"] < 5 or not d.get("canary_finish"): + return d, "VOID(degenerate canary)" + return d, "CLEAN" + +base, base_status = load("tiled", ladder[0]) +base_text = base["canary_text"] if base else "" + +def canary(d): + if d is None or base is None: + return "-" + a, b = base_text, d["canary_text"] + if a == b: + return "identical" + n = min(len(a), len(b)) + i = next((j for j in range(n) if a[j] != b[j]), n) + return f"diverges@char{i}" + +print("=== FP8 three-arm threshold sweep ===") +print(f"floor: rate counted only when summed usage.completion_tokens >= {floor}") +print(f"canary baseline: tiled @ M={ladder[0]} [{base_status}]") +print() +print(f"{'M':>5} {'arm':>9} {'tokens':>7} {'wall_s':>8} {'tok/s':>9} {'status':>26} canary_vs_baseline") +rates = {} +for m in ladder: + for arm in ARMS: + d, status = load(arm, m) + rates[(arm, m)] = d["tok_per_s"] if (d and status == "CLEAN") else None + if d: + print(f"{m:>5} {arm:>9} {d['completion_tokens_sum']:>7} {d['wall_s']:>8.1f} " + f"{d['tok_per_s']:>9.1f} {status:>26} {canary(d)}") + else: + print(f"{m:>5} {arm:>9} {'-':>7} {'-':>8} {'-':>9} {status:>26} -") +print() +print("=== CROSSOVER (the number ARC_FP8_CUBLAS_MIN_M may move to) ===") +crossover = None +for m in ladder: + a = rates[("cublaslt", m)] + b = rates[("wmma", m)] + c = rates[("tiled", m)] + if a is None or b is None or c is None: + print(f"M={m:>4}: NOT ESTABLISHED โ€” a leg is void; no claim about this M") + continue + native = min(b, c) + beat = a > native + print(f"M={m:>4}: cublaslt {a:.1f} vs min(wmma {b:.1f}, tiled {c:.1f}) = {native:.1f}" + f" -> {'cuBLASLt WINS' if beat else 'native wins'}") + if beat and crossover is None: + crossover = m +print() +if crossover is not None: + print(f"VERDICT: first clean M where cuBLASLt beats min(WMMA, tiled): {crossover}.") + print(f" Recommended default: ARC_FP8_CUBLAS_MIN_M={crossover} โ€” merge WITH these") + print(" rows quoted in the PR (a threshold is a claim about every value it") + print(" excludes; these are the rows that back the claim).") +else: + print("VERDICT: no clean crossover on this ladder โ€” the default STAYS 512.") + print(" Do not move the threshold from a partial table; re-run the void legs.") +PYEOF + +say "DONE โ€” $OUT/summary.txt (engagement detail: $OUT/engagement.txt)" diff --git a/mistralrs-core/tests/capability_reachability.rs b/mistralrs-core/tests/capability_reachability.rs index c005e0153..7e1d0d7a7 100644 --- a/mistralrs-core/tests/capability_reachability.rs +++ b/mistralrs-core/tests/capability_reachability.rs @@ -326,6 +326,21 @@ static REGISTRY: &[Capability] = &[ // D2H that also disqualifies all rows from the GPU sampler fast path. status: Status::Live, }, + Capability { + name: "FP8 dispatch engagement log (ARC_LOG_FP8_DISPATCH)", + parent: "ArcLab", + check: Check::Symbol { + symbol: "log_fp8_dispatch", + defined_in: "mistralrs-quant/src/blockwise_fp8/ops.rs", + }, + // The engagement line the FP8 threshold sweep + // (`arc-tools/fp8_threshold_sweep.sh`) greps per leg to prove WHICH + // kernel actually served it. Value-read (`env_flag_is_set`), latched + // once per process, one line per path per process. Referenced from + // both native dispatch arms in `ops.rs` and from the + // dequantize+cuBLASLt divert in `blockwise_fp8/mod.rs`. + status: Status::Live, + }, // -- the ApiPromise class ------------------------------------------------ Capability { name: "sampling: logit_bias reaches the sampler", diff --git a/mistralrs-quant/kernels/blockwise_fp8/blockwise_fp8_gemm.cu b/mistralrs-quant/kernels/blockwise_fp8/blockwise_fp8_gemm.cu index dad9e8330..077b7c473 100644 --- a/mistralrs-quant/kernels/blockwise_fp8/blockwise_fp8_gemm.cu +++ b/mistralrs-quant/kernels/blockwise_fp8/blockwise_fp8_gemm.cu @@ -40,17 +40,54 @@ __device__ __forceinline__ float get_scale(const float *__restrict__ scale, return __ldg(&scale[sr * scale_stride + sc]); } +// Shift variant of `get_scale`, for power-of-two scale-block sizes. +// +// `block_size_y`/`block_size_x` arrive as kernel ARGUMENTS, so `n / +// block_size_y` above is `int / int` by a non-constant divisor. NVIDIA GPUs +// have no integer-divide instruction; nvcc expands each division inline to a +// reciprocal-multiply-and-fixup sequence of ~15-20 instructions, and +// `fp8_matmul_tiled` pays TWO of them per weight element it stages into +// shared memory. The shipped `weight_block_size` is [128, 128] โ€” both powers +// of two โ€” so the host launcher computes `shift = log2(block_size)` once per +// launch and the kernel indexes with two SHFs instead. +// +// Bit-identical: for non-negative `n` and a power-of-two divisor, +// `n >> log2(d)` IS `n / d` โ€” the same quotient for every input, no rounding +// involved anywhere. The two paths read the same scale word; only the +// instruction count changes. Non-power-of-two geometries keep the division +// path above. +__device__ __forceinline__ float get_scale_pow2(const float *__restrict__ scale, + int n, int k, int scale_stride, + int scale_shift_y, + int scale_shift_x) { + int sr = n >> scale_shift_y; + int sc = k >> scale_shift_x; + return __ldg(&scale[sr * scale_stride + sc]); +} + // ============================================================================ // FP8 Matmul Kernel // ============================================================================ -template +// `POW2_SCALE` selects, at compile time, how the per-element weight scale is +// indexed while staging the weight tile: +// * true -> `get_scale_pow2`: two SHFs, using `scale_shift_y`/`scale_shift_x` +// (host-computed `log2(block_size_y)`/`log2(block_size_x)`). +// * false -> `get_scale`: two runtime integer divisions (the general path, +// kept for non-power-of-two scale geometries). +// The two instantiations produce BIT-IDENTICAL output (`n >> log2(d) == n / d` +// for the non-negative indices used here), so dispatching on it can never +// change what a leg of an A/B measures โ€” only how fast the scalar arm runs. +// Both `block_size_*` and `scale_shift_*` are always passed; each +// instantiation reads only its own pair. +template __global__ void fp8_matmul_tiled(const T *__restrict__ input, const __nv_fp8_e4m3 *__restrict__ weight, const float *__restrict__ weight_scale, T *__restrict__ output, int M, int N, int K, int scale_row_stride, int block_size_y, - int block_size_x) { + int block_size_x, int scale_shift_y, + int scale_shift_x) { // Padding is +1, not +4. The inner product reads `s_weight[tx][k]` with `tx` // varying fastest within a warp, so consecutive lanes are `BLOCK_K + pad` // floats apart. At BLOCK_K=32 a +4 pad gives stride 36; 36 mod 32 = 4, so @@ -106,8 +143,14 @@ __global__ void fp8_matmul_tiled(const T *__restrict__ input, if (gn < N && gk < K) { __nv_fp8_e4m3 w; w.__x = __ldg(reinterpret_cast(&weight[gn * K + gk])); - float s = get_scale(weight_scale, gn, gk, scale_row_stride, - block_size_y, block_size_x); + float s; + if constexpr (POW2_SCALE) { + s = get_scale_pow2(weight_scale, gn, gk, scale_row_stride, + scale_shift_y, scale_shift_x); + } else { + s = get_scale(weight_scale, gn, gk, scale_row_stride, block_size_y, + block_size_x); + } val = fp8_to_float(w) * s; } s_weight[ln][lk] = val; @@ -625,41 +668,66 @@ __global__ void fp8_moe_gemm(const T *__restrict__ input, // C API // ============================================================================ -extern "C" void launch_fp8_matmul_f16(const __half *input, - const __nv_fp8_e4m3 *weight, - const float *weight_scale, __half *output, - int M, int N, int K, int scale_row_stride, - int block_size_y, int block_size_x, - cudaStream_t stream) { +// Host-side pow2 test + exact log2 for the tiled kernel's scale fast path. +// Computed ONCE per launch, on the host, so the device never divides. +static inline bool fp8_is_pow2(int v) { return v > 0 && (v & (v - 1)) == 0; } +static inline int fp8_ilog2_exact(int v) { + int s = 0; + while ((1 << s) < v) + ++s; + return s; // caller guarantees fp8_is_pow2(v), so (1 << s) == v +} + +// Shared launcher: picks the POW2_SCALE instantiation when BOTH scale-block +// sizes are powers of two (the shipped geometry is [128, 128], so serving +// always takes it), otherwise the general division path. The two +// instantiations are bit-identical in output โ€” see the kernel comment. +template +static void launch_fp8_matmul_impl(const T *input, const __nv_fp8_e4m3 *weight, + const float *weight_scale, T *output, int M, + int N, int K, int scale_row_stride, + int block_size_y, int block_size_x, + cudaStream_t stream) { constexpr int TILE = 32; constexpr int TILE_K = 32; dim3 block(TILE, TILE); dim3 grid(CEILDIV(N, TILE), CEILDIV(M, TILE)); - fp8_gemm::fp8_matmul_tiled - <<>>(input, weight, weight_scale, output, M, N, K, - scale_row_stride, block_size_y, - block_size_x); + if (fp8_is_pow2(block_size_y) && fp8_is_pow2(block_size_x)) { + fp8_gemm::fp8_matmul_tiled + <<>>(input, weight, weight_scale, output, M, N, + K, scale_row_stride, block_size_y, + block_size_x, fp8_ilog2_exact(block_size_y), + fp8_ilog2_exact(block_size_x)); + } else { + fp8_gemm::fp8_matmul_tiled + <<>>(input, weight, weight_scale, output, M, N, + K, scale_row_stride, block_size_y, + block_size_x, 0, 0); + } CUDA_CHECK(cudaGetLastError()); } +extern "C" void launch_fp8_matmul_f16(const __half *input, + const __nv_fp8_e4m3 *weight, + const float *weight_scale, __half *output, + int M, int N, int K, int scale_row_stride, + int block_size_y, int block_size_x, + cudaStream_t stream) { + launch_fp8_matmul_impl(input, weight, weight_scale, output, M, N, K, + scale_row_stride, block_size_y, block_size_x, + stream); +} + extern "C" void launch_fp8_matmul_bf16(const __nv_bfloat16 *input, const __nv_fp8_e4m3 *weight, const float *weight_scale, __nv_bfloat16 *output, int M, int N, int K, int scale_row_stride, int block_size_y, int block_size_x, cudaStream_t stream) { - constexpr int TILE = 32; - constexpr int TILE_K = 32; - - dim3 block(TILE, TILE); - dim3 grid(CEILDIV(N, TILE), CEILDIV(M, TILE)); - - fp8_gemm::fp8_matmul_tiled<__nv_bfloat16, TILE, TILE, TILE_K> - <<>>(input, weight, weight_scale, output, M, N, K, - scale_row_stride, block_size_y, - block_size_x); - CUDA_CHECK(cudaGetLastError()); + launch_fp8_matmul_impl<__nv_bfloat16>(input, weight, weight_scale, output, M, + N, K, scale_row_stride, block_size_y, + block_size_x, stream); } // Warp-per-row GEMV launcher for the decode regime (M <= 4). Row-grouping diff --git a/mistralrs-quant/src/blockwise_fp8/mod.rs b/mistralrs-quant/src/blockwise_fp8/mod.rs index 2cf3b9fa5..5ef1cecc3 100644 --- a/mistralrs-quant/src/blockwise_fp8/mod.rs +++ b/mistralrs-quant/src/blockwise_fp8/mod.rs @@ -95,6 +95,20 @@ impl QuantMethod for BlockwiseFP8Linear { // A/B runs from a single binary. let m_rows: usize = x.dims().iter().rev().skip(1).product(); let prefer_cublas = m_rows >= arc_fp8_cublas_min_m(); + if prefer_cublas + && matches!(x.device(), candle_core::Device::Cuda(_)) + && ffi::HAVE_BLOCKWISE_GEMM_KERNELS + { + // Engagement line for the threshold sweep: this forward was + // DIVERTED from the native FP8 kernels to dequantize+cuBLASLt + // by `ARC_FP8_CUBLAS_MIN_M`. `ARC_LOG_FP8_DISPATCH=1` only. + ops::log_fp8_dispatch( + ops::Fp8DispatchPath::DequantCublaslt, + m_rows, + self.weight.dims().first().copied().unwrap_or(0), + x.dims().last().copied().unwrap_or(0), + ); + } if !prefer_cublas && matches!(x.device(), candle_core::Device::Cuda(_)) && ffi::HAVE_BLOCKWISE_GEMM_KERNELS @@ -143,6 +157,34 @@ impl QuantMethod for BlockwiseFP8Linear { weight, self.bias.clone(), )))?; + + // ๐Ÿ”ด FLATTEN TO 2-D FIRST โ€” the native branch above already does this at + // the top of `forward`, and this branch did not. + // + // V4's activation is rank-3 `[B, T, hidden]`. Passed through untouched, + // `UnquantLinear::forward` takes `w.broadcast_left(B)` and, because bias + // is `None` on every V4 linear, dispatches `cublaslt.batch_matmul` with + // `stride_b = 0`. At decode T=1, so that is a **B-way batched GEMM with + // m=1** โ€” a batched GEMV wearing a GEMM's name, which is the one shape + // cuBLASLt has no advantage in. The documented 27x was measured at + // prefill, where `[1, 512, hidden]` collapses to batch=1 and the call is + // a genuine `[512, hidden] x [hidden, N]` GEMM. + // + // Without this, lowering `ARC_FP8_CUBLAS_MIN_M` measures the batched-GEMV + // path and concludes "cuBLASLt loses at decode" โ€” a wrong conclusion + // drawn from a shape the threshold change was never meant to select. + // Flattening makes the two branches comparable, which is the whole + // premise of the A/B. + let orig_dims = x.dims().to_vec(); + if orig_dims.len() > 2 { + let features = orig_dims[orig_dims.len() - 1]; + let rows: usize = orig_dims[..orig_dims.len() - 1].iter().product(); + let out = unquant.forward(&x.reshape((rows, features))?)?; + let out_features = out.dim(out.rank() - 1)?; + let mut new_dims = orig_dims[..orig_dims.len() - 1].to_vec(); + new_dims.push(out_features); + return out.reshape(new_dims); + } unquant.forward(x) } diff --git a/mistralrs-quant/src/blockwise_fp8/ops.rs b/mistralrs-quant/src/blockwise_fp8/ops.rs index 75a1d1f42..8d8e8c5c9 100644 --- a/mistralrs-quant/src/blockwise_fp8/ops.rs +++ b/mistralrs-quant/src/blockwise_fp8/ops.rs @@ -1008,8 +1008,7 @@ fn fp8_wmma_enabled() -> bool { // Under the presence test this replaced, it silently selected the control // arm โ€” which is the arm this kernel's own docs tell you to run first, so // an operator following them would have recorded the control twice. - static ENABLED: LazyLock = - LazyLock::new(|| !crate::env_flag_is_set("ARC_NO_FP8_WMMA")); + static ENABLED: LazyLock = LazyLock::new(|| !crate::env_flag_is_set("ARC_NO_FP8_WMMA")); *ENABLED } @@ -1099,6 +1098,81 @@ fn wide_gemv_shift( ) } +/// Parent system: ArcLab. +/// +/// Which blockwise-FP8 kernel a forward actually selected. Used only by +/// [`log_fp8_dispatch`]; the discriminants index its once-per-process latch. +#[cfg(feature = "cuda")] +#[derive(Clone, Copy)] +pub(crate) enum Fp8DispatchPath { + GemvWide = 0, + GemvWarp = 1, + Wmma = 2, + Tiled = 3, + DequantCublaslt = 4, +} + +#[cfg(feature = "cuda")] +impl Fp8DispatchPath { + fn name(self) -> &'static str { + match self { + Fp8DispatchPath::GemvWide => "gemv_wide", + Fp8DispatchPath::GemvWarp => "gemv_warp", + Fp8DispatchPath::Wmma => "wmma", + Fp8DispatchPath::Tiled => "tiled", + Fp8DispatchPath::DequantCublaslt => "dequant_cublaslt", + } + } +} + +/// `ARC_LOG_FP8_DISPATCH=1` โ€” engagement logging for the blockwise-FP8 +/// dispatch. Read by VALUE via `env_flag_is_set` (`=0` keeps it off), latched +/// once per process so the decode hot loop never touches the env. +/// +/// Registered in `mistralrs-core/tests/capability_reachability.rs`. +#[cfg(feature = "cuda")] +fn fp8_dispatch_log_enabled() -> bool { + use std::sync::LazyLock; + static ENABLED: LazyLock = + LazyLock::new(|| crate::env_flag_is_set("ARC_LOG_FP8_DISPATCH")); + *ENABLED +} + +/// Parent system: ArcLab. +/// +/// Print ONE line per dispatch path per process naming the kernel that +/// actually ran โ€” the engagement assertion every threshold sweep must grep +/// for. A sweep leg that cannot prove WHICH kernel served it is measuring an +/// assumption (`ARC_NO_DEDICATED_DECODE` produced exactly that null result: +/// both arms of two A/Bs ran the same code and nobody could tell). +/// +/// Cost when `ARC_LOG_FP8_DISPATCH` is unset (the default): one latched bool +/// read. When set: one relaxed atomic swap per call โ€” no locks, no +/// allocation โ€” so enabling it does not perturb the rates the sweep records. +/// The shape printed is the FIRST shape that reached the path, which is +/// enough to identify the projection; per-call logging would be ~430 lines +/// per V4 step at B=256 and is deliberately not offered. +#[cfg(feature = "cuda")] +pub(crate) fn log_fp8_dispatch(path: Fp8DispatchPath, m: usize, n: usize, k: usize) { + use std::sync::atomic::{AtomicBool, Ordering}; + static LOGGED: [AtomicBool; 5] = [ + AtomicBool::new(false), + AtomicBool::new(false), + AtomicBool::new(false), + AtomicBool::new(false), + AtomicBool::new(false), + ]; + if !fp8_dispatch_log_enabled() { + return; + } + if !LOGGED[path as usize].swap(true, Ordering::Relaxed) { + eprintln!( + "[arc-fp8-dispatch] path={} first_shape=m{m}_n{n}_k{k}", + path.name() + ); + } +} + /// FP8 blockwise matmul. /// Computes output = input @ weight.T where weight is FP8 blockwise quantized. /// - input: [M, K] in fp16/bf16 @@ -1221,8 +1295,7 @@ fn fp8_blockwise_matmul_impl( match input.dtype() { DType::F16 => { - let output = - crate::arc_outbuf::alloc_out_fully_written::(&dev, (m * n) as usize)?; + let output = crate::arc_outbuf::alloc_out_fully_written::(&dev, (m * n) as usize)?; let input_s = match &*input_storage { Storage::Cuda(cuda_storage) => cuda_storage.as_cuda_slice::()?, @@ -1234,6 +1307,21 @@ fn fp8_blockwise_matmul_impl( let (input_ptr, _input_guard) = slice_ptr(input_s, input_l.start_offset()); let wide_shift = wide_gemv_shift(use_gemv, k, block_size_x, weight_ptr, input_ptr); + // Mirrors the if/else chain below exactly; keep the two in step. + log_fp8_dispatch( + if wide_shift.is_some() { + Fp8DispatchPath::GemvWide + } else if use_gemv { + Fp8DispatchPath::GemvWarp + } else if use_wmma { + Fp8DispatchPath::Wmma + } else { + Fp8DispatchPath::Tiled + }, + m as usize, + n as usize, + k as usize, + ); if let Some(scale_shift) = wide_shift { // ๐Ÿ”ด UNVERIFIED ON HARDWARE. `ARC_FP8_GEMV_WIDE=1` only. unsafe { @@ -1323,6 +1411,21 @@ fn fp8_blockwise_matmul_impl( let (input_ptr, _input_guard) = slice_ptr(input_s, input_l.start_offset()); let wide_shift = wide_gemv_shift(use_gemv, k, block_size_x, weight_ptr, input_ptr); + // Mirrors the if/else chain below exactly; keep the two in step. + log_fp8_dispatch( + if wide_shift.is_some() { + Fp8DispatchPath::GemvWide + } else if use_gemv { + Fp8DispatchPath::GemvWarp + } else if use_wmma { + Fp8DispatchPath::Wmma + } else { + Fp8DispatchPath::Tiled + }, + m as usize, + n as usize, + k as usize, + ); if let Some(scale_shift) = wide_shift { // ๐Ÿ”ด UNVERIFIED ON HARDWARE. `ARC_FP8_GEMV_WIDE=1` only. unsafe { From 4d03b9e2541b9a88f036cdd7640495fb98f16fc9 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 22:23:42 +0100 Subject: [PATCH 20/21] =?UTF-8?q?fix(ArcInfer/ArcGraph):=20cross-request?= =?UTF-8?q?=20token=20leak=20=E2=80=94=20parked=20burst=20tokens=20no=20lo?= =?UTF-8?q?nger=20outlive=20their=20sequence=20(#220)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(ArcInfer/ArcGraph): parked burst tokens no longer outlive their sequence โ€” one funnel, one leak test A device-loop burst runs with eos_token_id = -1, so it always parks its full burst; a sequence stopping mid-burst (stop string, max_tokens, cancel, error, panic, eviction) left its undrained tokens parked, and the NEXT sequence's first sample_sequence โ€” including the one after its prefill โ€” consumed them as its own. One user's tokens in another user's response. clear_pending_tokens documented the contract but its only production call sites were fallback paths. The fix is one call at the one funnel every completion path passes through: Sequence::set_state stands the device loop down on every transition out of the running set (Done, Error, FinishedAborted, FinishedIgnored, Waiting, Swapped). Clearing on preemption is deliberate fail-closed: the sequence re-prefills on re-admission, so its parked tokens are stale. To make the funnel testable on machines without CUDA โ€” the cfg(cuda) darkness is exactly what let this ship โ€” arc-cuda-graph is now a non-optional dependency of mistralrs-core (its host surface is plain Rust; CUDA plumbing stays behind its own cuda feature, still enabled by mistralrs-core/cuda) and device_loop_pre_sampled_token loses its cfg split. The new leak test drives the real sample_sequence path both ways: wired, B's first token is its own argmax and the queue is empty; with the funnel removed, the queue holds 3 of A's tokens and B's first sampled token is 22 โ€” A's parked id. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf * fix(ArcInfer/ArcGraph): owner-tag the pending queue โ€” the leak becomes structurally impossible, not funnel-dependent The completion funnel cannot see a sequence the DefaultScheduler moves aside WITHOUT a state change (the bucketing waitlist says so in its own doc), and batch growth has the same shape. So the queue itself now refuses: tokens are parked tagged with the owning sequence's id (published by the sampling pre-hook with eligibility), take requires a matching id, and a foreign queue is dropped whole with the grep-able instrument line "ArcGraph: dropped N foreign parked tokens". The forward's pending>0 short-circuit gains the two guards its aliasing contract always needed: seq_len == 1 (a parked queue at a prefill step must not replace the prefill) and pending_owned_by_current(). If the short-circuit fired and the sample still cannot take (scheduler swapped sequences in between), the step fails loudly via the aliased-logits marker instead of host-sampling another sequence's stale graph output. DeviceDecodeLoop finally resets: every stand_down (including the funnel's) bumps a device-state generation; run() compares and zeroes ring, cursors and fault word before a newer-generation burst. Reset at engagement, not completion, because the funnel cannot reach the pipeline-owned loop and between bursts is the only point with no in-flight ring writes. ARC_GRAPH_DEVICE_LOOP is now registered in capability_reachability.rs (Symbol: device_loop_enabled), closing the refuted premise from the previous commit. Interleave test (A waitlisted, no set_state): wired, B samples its own argmax 9 and the foreign tail is dropped; with the owner comparison removed, B's first sampled token is 22 โ€” A's parked id. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf --------- Co-authored-by: Claude Fable 5 --- arc-cuda-graph/src/device_loop.rs | 288 +++++++++++-- arc-cuda-graph/src/lib.rs | 9 +- mistralrs-core/Cargo.toml | 10 +- mistralrs-core/src/pipeline/normal.rs | 44 +- mistralrs-core/src/pipeline/sampling.rs | 406 +++++++++++++++++- mistralrs-core/src/sequence.rs | 35 ++ .../tests/capability_reachability.rs | 14 + 7 files changed, 736 insertions(+), 70 deletions(-) diff --git a/arc-cuda-graph/src/device_loop.rs b/arc-cuda-graph/src/device_loop.rs index e2885182b..757c31a7f 100644 --- a/arc-cuda-graph/src/device_loop.rs +++ b/arc-cuda-graph/src/device_loop.rs @@ -933,6 +933,15 @@ mod cuda_impl { pub struct DeviceDecodeLoop { ops: CudaDeviceOps, driver: BurstDriver, + /// The [`super::device_loop_generation`] this loop's ring, cursors and + /// fault word are valid for. Every `stand_down` โ€” including the one + /// the `Sequence::set_state` completion funnel issues โ€” bumps the + /// generation, and [`Self::run`] resets the per-sequence device state + /// before running a burst from a newer generation. Reset happens HERE, + /// at engagement, not at completion: the funnel cannot reach this + /// object (it lives in the pipeline), and between bursts is the only + /// point where no device write to the ring can still be in flight. + generation: u64, } // Same reasoning as `CudaGraphRunner` (`graph.rs:114`): the raw pointers are @@ -976,13 +985,26 @@ mod cuda_impl { cfg, )?; let driver = BurstDriver::new(cfg, batch)?; - Ok(Self { ops, driver }) + Ok(Self { + ops, + driver, + generation: super::device_loop_generation(), + }) } /// Enqueue `steps` decode steps and return their tokens. /// /// Exactly `steps` `cuGraphLaunch` calls, zero `cudaStreamSynchronize`. pub fn run(&mut self, steps: usize) -> candle_core::Result { + // A stand-down happened since the last burst โ€” sequence completed, + // errored, or was preempted โ€” so the ring head, read cursors and + // fault word describe a sequence that no longer exists. Zero them + // before this burst rather than carrying them across sequences. + let current = super::device_loop_generation(); + if self.generation != current { + self.reset(); + self.generation = current; + } let out = self.driver.run_burst(&mut self.ops, steps)?; // D18: a subsystem that only announces itself at startup cannot be // told apart from one that started and then did nothing. Emit the @@ -1156,14 +1178,37 @@ pub fn admit( std::thread_local! { static PENDING_TOKENS: std::cell::RefCell> = const { std::cell::RefCell::new(std::collections::VecDeque::new()) }; + /// The sequence the parked tokens belong to, stamped at park time from + /// [`CURRENT_SEQ`]. The completion funnel (`Sequence::set_state`) is the + /// primary defence against a parked token outliving its sequence; this tag + /// is the structural one โ€” it covers the paths NO funnel can see, such as + /// the DefaultScheduler bucketing waitlist, which moves a running sequence + /// aside *without a state modification* (`default_scheduler.rs`, the + /// `bucket_and_waitlist_seqs_waiting` doc says so explicitly). + static PENDING_OWNER: std::cell::Cell> = const { std::cell::Cell::new(None) }; + /// The sequence id the last `sample_sequence` published alongside + /// eligibility. At batch 1 โ€” the only shape the loop admits โ€” this is + /// exactly the sequence the next forward's burst decodes for. + static CURRENT_SEQ: std::cell::Cell> = const { std::cell::Cell::new(None) }; + /// Set by the forward when it served this step from the pending queue and + /// handed back the ALIASING graph logits tensor without launching anything + /// (`normal.rs::device_decode_burst`, the `pending > 0` arm). Consumed by + /// the very next sample. If that sample cannot take a parked token, the + /// logits it holds are a stale alias of graph-owned storage and sampling + /// them would return a plausible token from the WRONG sequence's + /// distribution โ€” so the sample must fail loudly instead. + static ALIASED_LOGITS_SERVED: std::cell::Cell = const { std::cell::Cell::new(false) }; } -/// Park a burst's tokens for the engine to drain. +/// Park a burst's tokens for the engine to drain. The queue is stamped with +/// the current decoding sequence ([`set_device_loop_eligible`]) as its owner; +/// [`take_pending_token`] refuses to hand them to anyone else. /// /// Negative ids are dropped rather than cast: the commit kernel range-checks /// before it writes, so a negative here would mean the ring was read wrong, and /// wrapping it into a huge `u32` would index the tokenizer out of bounds. pub fn push_pending_tokens(tokens: &[i32]) { + PENDING_OWNER.with(|o| o.set(CURRENT_SEQ.with(|c| c.get()))); PENDING_TOKENS.with(|q| { let mut q = q.borrow_mut(); for &t in tokens { @@ -1180,9 +1225,53 @@ pub fn push_pending_tokens(tokens: &[i32]) { }) } -/// Take the next token the burst already produced, if any. -pub fn take_pending_token() -> Option { - PENDING_TOKENS.with(|q| q.borrow_mut().pop_front()) +/// What [`take_pending_token`] found. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PendingTake { + /// A token parked by a burst for THIS sequence. + Taken(u32), + /// Nothing parked. + Empty, + /// The queue held another sequence's tokens. They were dropped โ€” handing + /// even one across would put one user's tokens in another user's response + /// โ€” and `dropped` says how many. Logged as + /// `"ArcGraph: dropped N foreign parked tokens"`; on a box that line is + /// the violation instrument, so keep it grep-able. + Foreign { dropped: usize }, +} + +/// Take the next parked token, if it belongs to `seq_id`. +/// +/// A queue owned by a different sequence is dropped whole, fail-closed: the +/// owner has left the running set (or been waitlisted with no state change), +/// its device-side burst state is stale, and replaying its tokens into any +/// sequence โ€” including the owner after a re-prefill โ€” would corrupt output. +pub fn take_pending_token(seq_id: usize) -> PendingTake { + let owner = PENDING_OWNER.with(|o| o.get()); + PENDING_TOKENS.with(|q| { + let mut q = q.borrow_mut(); + if q.is_empty() { + return PendingTake::Empty; + } + if owner == Some(seq_id) { + match q.pop_front() { + Some(t) => PendingTake::Taken(t), + None => PendingTake::Empty, + } + } else { + let dropped = q.len(); + q.clear(); + PENDING_OWNER.with(|o| o.set(None)); + tracing::warn!( + "ArcGraph: dropped {dropped} foreign parked tokens (parked for sequence \ + {owner:?}, sampler is sequence {seq_id}) โ€” cross-sequence leak prevented. \ + Expected once when the scheduler moves a sequence aside mid-burst; frequent \ + occurrences mean the device loop is engaging under scheduler churn it cannot \ + amortise" + ); + PendingTake::Foreign { dropped } + } + }) } /// Tokens still parked. @@ -1190,6 +1279,29 @@ pub fn pending_token_count() -> usize { PENDING_TOKENS.with(|q| q.borrow().len()) } +/// Does the pending queue belong to the sequence whose sample last published +/// eligibility? The forward's `pending > 0` short-circuit must not fire for +/// anyone else: it returns stale aliasing logits and launches nothing, which +/// is only sound when the very next sample will take a parked token. +pub fn pending_owned_by_current() -> bool { + let owner = PENDING_OWNER.with(|o| o.get()); + owner.is_some() && owner == CURRENT_SEQ.with(|c| c.get()) +} + +/// The forward served this step from the pending queue: no launch, and the +/// returned logits tensor aliases graph-owned storage. The next sample MUST +/// take a parked token; [`take_aliased_logits_marker`] is how it checks. +pub fn note_aliased_logits_served() { + ALIASED_LOGITS_SERVED.with(|c| c.set(true)); +} + +/// Consume the aliased-logits marker for this step. `true` means the logits +/// the current sample holds were NOT produced by a launch for this step and +/// must not be host-sampled. +pub fn take_aliased_logits_marker() -> bool { + ALIASED_LOGITS_SERVED.with(|c| c.replace(false)) +} + // Whether the sequence currently being decoded is one the device loop may // drive. // @@ -1208,8 +1320,12 @@ std::thread_local! { static DEVICE_LOOP_ELIGIBLE: std::cell::Cell = const { std::cell::Cell::new(false) }; } -/// Publish whether the sequence being sampled may be driven by the device loop. -pub fn set_device_loop_eligible(eligible: bool) { +/// Publish whether the sequence being sampled may be driven by the device +/// loop, and WHICH sequence that is. The id becomes the owner stamped onto +/// anything the next burst parks, which is what makes a parked token +/// unconsumable by any other sequence. +pub fn set_device_loop_eligible(seq_id: usize, eligible: bool) { + CURRENT_SEQ.with(|c| c.set(Some(seq_id))); DEVICE_LOOP_ELIGIBLE.with(|c| c.set(eligible)) } @@ -1224,16 +1340,53 @@ pub fn device_loop_eligible() -> bool { /// back โ€” a token left here would be handed to the *next* sequence as if it /// were its own. That is the "+1-step divergence" failure mode with a different /// name, so it gets an explicit call rather than a lifetime. +/// +/// The sequence-completion half of that contract is enforced at ONE funnel: +/// `mistralrs-core`'s `Sequence::set_state` calls [`stand_down`] on every +/// transition out of the running set (`Done`, `Error`, `FinishedAborted`, +/// `FinishedIgnored`, `Waiting`, `Swapped`), which is why this crate's +/// dependency in `mistralrs-core` is deliberately not optional โ€” the funnel +/// and its leak test compile and run on hosts without CUDA. The tests here +/// cover only this function's own behaviour; the call-site wiring is covered +/// by `pipeline/sampling.rs`'s `device_loop_cross_sequence_leak_tests`. pub fn clear_pending_tokens() { + PENDING_OWNER.with(|o| o.set(None)); PENDING_TOKENS.with(|q| q.borrow_mut().clear()) } -/// Stand down completely: forget parked tokens and mark the current sequence -/// ineligible. Call on fall-back, on sequence completion, and on any error, so -/// the next sequence starts from a clean slate. +/// Stand down completely: forget parked tokens (and their owner), drop the +/// aliased-logits marker, mark the current sequence ineligible, and bump the +/// device-state generation so a kept-alive [`DeviceDecodeLoop`] resets its +/// ring and cursors before it next runs. Call on fall-back, on sequence +/// completion, and on any error, so the next sequence starts from a clean +/// slate. Sequence completion โ€” every transition out of the running set โ€” +/// calls this from the `Sequence::set_state` funnel; see +/// [`clear_pending_tokens`]. pub fn stand_down() { clear_pending_tokens(); - set_device_loop_eligible(false); + ALIASED_LOGITS_SERVED.with(|c| c.set(false)); + DEVICE_LOOP_ELIGIBLE.with(|c| c.set(false)); + DEVICE_LOOP_GENERATION.with(|g| g.set(g.get().wrapping_add(1))); +} + +// The device-state generation. `DeviceDecodeLoop` holds ring cursors, the +// pinned ring head and the device-side position โ€” all per-sequence state โ€” +// but it is deliberately kept alive across sequences because every buffer it +// points at must stay at the address the captured graph baked +// (`DeviceDecodeLoop` docs). So instead of the funnel destroying it (it +// cannot even reach it: the loop lives in the pipeline, the funnel in +// `Sequence`), every `stand_down` bumps this generation and the loop compares +// on its next `run`, resetting itself first if the world moved on. Reset at +// engagement time is also the only SAFE time: between bursts nothing is in +// flight on the stream, whereas a reset at completion time could race a +// still-draining burst's device writes. +std::thread_local! { + static DEVICE_LOOP_GENERATION: std::cell::Cell = const { std::cell::Cell::new(0) }; +} + +/// The current device-state generation. Bumped by every [`stand_down`]. +pub fn device_loop_generation() -> u64 { + DEVICE_LOOP_GENERATION.with(|g| g.get()) } /// Process-wide latch. Once the device loop has failed, it stays off. @@ -1717,18 +1870,20 @@ mod tests { assert_eq!(admit(&f, &c, true), Err(Refusal::NoPositionRoom)); } - /// The parked tokens are handed out in generation order, exactly once. - /// A queue that re-served a token, or served them reversed, would produce - /// fluent-but-wrong text โ€” invisible without this. + /// The parked tokens are handed out in generation order, exactly once, to + /// the sequence that owns them. A queue that re-served a token, or served + /// them reversed, would produce fluent-but-wrong text โ€” invisible without + /// this. #[test] fn pending_tokens_drain_in_order_exactly_once() { - clear_pending_tokens(); + stand_down(); + set_device_loop_eligible(11, true); push_pending_tokens(&[7, 8, 9]); assert_eq!(pending_token_count(), 3); - assert_eq!(take_pending_token(), Some(7)); - assert_eq!(take_pending_token(), Some(8)); - assert_eq!(take_pending_token(), Some(9)); - assert_eq!(take_pending_token(), None); + assert_eq!(take_pending_token(11), PendingTake::Taken(7)); + assert_eq!(take_pending_token(11), PendingTake::Taken(8)); + assert_eq!(take_pending_token(11), PendingTake::Taken(9)); + assert_eq!(take_pending_token(11), PendingTake::Empty); assert_eq!(pending_token_count(), 0); } @@ -1737,21 +1892,93 @@ mod tests { /// another's output. #[test] fn clearing_pending_tokens_prevents_leaking_them_into_the_next_sequence() { - clear_pending_tokens(); + stand_down(); + set_device_loop_eligible(11, true); push_pending_tokens(&[7, 8, 9]); clear_pending_tokens(); - assert_eq!(take_pending_token(), None); + assert_eq!(take_pending_token(11), PendingTake::Empty); + } + + /// THE structural guarantee: tokens parked while sequence 11 was current + /// cannot be taken by sequence 12 โ€” they are dropped whole, and the drop + /// is reported. This is the defence the completion funnel cannot provide: + /// the DefaultScheduler bucketing waitlist moves a running sequence aside + /// *without a state modification*, so no funnel ever fires. + #[test] + fn foreign_parked_tokens_are_dropped_not_served() { + stand_down(); + set_device_loop_eligible(11, true); + push_pending_tokens(&[7, 8, 9]); + + // Sequence 12 samples next (11 was waitlisted, no state change). + set_device_loop_eligible(12, true); + assert_eq!( + take_pending_token(12), + PendingTake::Foreign { dropped: 3 }, + "another sequence's parked tokens must be dropped, not served" + ); + assert_eq!(pending_token_count(), 0, "the foreign queue is gone whole"); + // And they are gone for the original owner too: its device-side burst + // state is stale, so replaying them after a re-prefill would corrupt. + assert_eq!(take_pending_token(11), PendingTake::Empty); + } + + /// The forward-side guard: the `pending > 0` short-circuit may only fire + /// for the sequence that owns the queue. + #[test] + fn pending_ownership_tracks_the_published_sequence() { + stand_down(); + assert!( + !pending_owned_by_current(), + "empty queue is owned by nobody" + ); + set_device_loop_eligible(11, true); + push_pending_tokens(&[7]); + assert!(pending_owned_by_current()); + set_device_loop_eligible(12, true); + assert!( + !pending_owned_by_current(), + "the moment another sequence publishes, the queue is foreign" + ); + } + + /// The aliased-logits marker is consumed exactly once, and stand_down + /// drops it: a marker surviving a stand-down would poison the next + /// sequence's first REAL sample. + #[test] + fn aliased_logits_marker_is_one_shot_and_cleared_by_stand_down() { + stand_down(); + assert!(!take_aliased_logits_marker()); + note_aliased_logits_served(); + assert!(take_aliased_logits_marker()); + assert!(!take_aliased_logits_marker(), "consumed exactly once"); + note_aliased_logits_served(); + stand_down(); + assert!(!take_aliased_logits_marker(), "stand_down must drop it"); + } + + /// Every stand_down bumps the device-state generation; that is what makes + /// a kept-alive `DeviceDecodeLoop` reset its ring and cursors before the + /// next sequence's first burst instead of carrying them across. + #[test] + fn stand_down_bumps_the_device_state_generation() { + let g0 = device_loop_generation(); + stand_down(); + assert_eq!(device_loop_generation(), g0.wrapping_add(1)); + stand_down(); + assert_eq!(device_loop_generation(), g0.wrapping_add(2)); } /// A negative id in the ring means the ring was misread. Casting it would /// index the tokenizer at ~4 billion; it must be dropped and logged. #[test] fn negative_ring_ids_are_dropped_not_wrapped_into_huge_u32() { - clear_pending_tokens(); + stand_down(); + set_device_loop_eligible(11, true); push_pending_tokens(&[5, -1, 6]); - assert_eq!(take_pending_token(), Some(5)); - assert_eq!(take_pending_token(), Some(6)); - assert_eq!(take_pending_token(), None); + assert_eq!(take_pending_token(11), PendingTake::Taken(5)); + assert_eq!(take_pending_token(11), PendingTake::Taken(6)); + assert_eq!(take_pending_token(11), PendingTake::Empty); } /// The eligibility flag must start `false`, so the first decode step of a @@ -1759,17 +1986,20 @@ mod tests { /// has checked yet. #[test] fn eligibility_starts_false_and_stand_down_resets_everything() { - set_device_loop_eligible(false); - clear_pending_tokens(); + stand_down(); assert!(!device_loop_eligible()); - set_device_loop_eligible(true); + set_device_loop_eligible(11, true); push_pending_tokens(&[1, 2]); assert!(device_loop_eligible()); stand_down(); assert!(!device_loop_eligible(), "stand_down must clear eligibility"); - assert_eq!(take_pending_token(), None, "stand_down must clear tokens"); + assert_eq!( + take_pending_token(11), + PendingTake::Empty, + "stand_down must clear tokens" + ); } /// Default-off, and explicitly pinnable off. diff --git a/arc-cuda-graph/src/lib.rs b/arc-cuda-graph/src/lib.rs index 766653698..1760ad210 100644 --- a/arc-cuda-graph/src/lib.rs +++ b/arc-cuda-graph/src/lib.rs @@ -88,10 +88,11 @@ pub use decode_forward::{decode_forward, DecodeBuffers, LayerKvCache, PagedAtten #[cfg(feature = "cuda")] pub use dedicated::DedicatedDecodePath; pub use device_loop::{ - admit, clear_pending_tokens, device_loop_eligible, device_loop_enabled, device_loop_killed, - kill_device_loop, pending_token_count, push_pending_tokens, set_device_loop_eligible, - stand_down, take_pending_token, AdmissionFacts, BurstDriver, BurstOutcome, DeviceLoopConfig, - DeviceLoopError, DeviceOps, Refusal, StreamState, + admit, clear_pending_tokens, device_loop_eligible, device_loop_enabled, device_loop_generation, + device_loop_killed, kill_device_loop, note_aliased_logits_served, pending_owned_by_current, + pending_token_count, push_pending_tokens, set_device_loop_eligible, stand_down, + take_aliased_logits_marker, take_pending_token, AdmissionFacts, BurstDriver, BurstOutcome, + DeviceLoopConfig, DeviceLoopError, DeviceOps, PendingTake, Refusal, StreamState, }; #[cfg(feature = "cuda")] pub use device_loop::{CudaDeviceOps, DeviceDecodeLoop}; diff --git a/mistralrs-core/Cargo.toml b/mistralrs-core/Cargo.toml index a58d0a9ff..dd7728002 100644 --- a/mistralrs-core/Cargo.toml +++ b/mistralrs-core/Cargo.toml @@ -71,7 +71,14 @@ mistralrs-paged-attn = { workspace = true, optional = true } mistralrs-quant.workspace = true arc-turbo = { path = "../arc-turbo" } arc-profiler = { path = "../arc-profiler" } -arc-cuda-graph = { path = "../arc-cuda-graph", optional = true, features = [] } +# NOT optional, deliberately: the sequence-completion funnel +# (`Sequence::set_state`) must stand the device decode loop down on every +# finished transition, and a gate that only compiles under +# `cfg(feature = "cuda")` cannot be tested on any machine that reviews it โ€” +# the exact class that let the cross-request token leak ship. The crate's +# host-side surface (pending-token queue, burst driver, admission) is plain +# Rust; the CUDA plumbing inside it stays behind its own `cuda` feature. +arc-cuda-graph = { path = "../arc-cuda-graph", features = [] } uuid = { workspace = true, features = ["v4"] } utoipa = { workspace = true, optional = true } schemars.workspace = true @@ -124,7 +131,6 @@ cuda = [ "mistralrs-quant/cuda", "dep:mistralrs-paged-attn", "mistralrs-paged-attn/cuda", - "dep:arc-cuda-graph", "arc-cuda-graph/cuda", "arc-profiler/cuda", ] diff --git a/mistralrs-core/src/pipeline/normal.rs b/mistralrs-core/src/pipeline/normal.rs index 0048e34c9..eaaaddcb9 100644 --- a/mistralrs-core/src/pipeline/normal.rs +++ b/mistralrs-core/src/pipeline/normal.rs @@ -167,11 +167,7 @@ pub(crate) enum AllocCacheAction { /// Pure, so the policy is testable without a GPU โ€” which matters, because the /// allocator itself has no tests at all in either repo. #[cfg_attr(not(feature = "cuda"), allow(dead_code))] -pub(crate) fn alloc_cache_action( - seq_len: usize, - enabled: bool, - killed: bool, -) -> AllocCacheAction { +pub(crate) fn alloc_cache_action(seq_len: usize, enabled: bool, killed: bool) -> AllocCacheAction { let want = !killed && seq_len == 1; match (want, enabled) { (true, false) => AllocCacheAction::Enable, @@ -2165,8 +2161,32 @@ fn device_decode_burst( // GPU work at all. This is where the amortisation actually shows up โ€” // `burst - 1` of every `burst` engine iterations do nothing but host // bookkeeping. + // + // Guarded twice, because returning `out` here hands back an ALIASING + // tensor from an earlier launch and skips this step's forward entirely โ€” + // sound only when the very next sample takes a parked token: + // * `seq_len == 1` โ€” parked tokens present at a PREFILL step belong to a + // sequence that left the running set without a state change (the + // bucketing waitlist); serving the new request's prefill from them + // would skip the prefill and leak the old sequence's tokens. + // * `pending_owned_by_current()` โ€” the queue must belong to the sequence + // whose sample published eligibility last, i.e. the one this decode + // step is for at batch 1. + // Either violation drops the parked tokens (loudly) and falls through to + // the real forward. if acg::pending_token_count() > 0 { - return Some(out); + if seq_len == 1 && acg::pending_owned_by_current() { + acg::note_aliased_logits_served(); + return Some(out); + } + tracing::warn!( + "ArcGraph: dropped {} foreign parked tokens โ€” parked queue found at a step it does \ + not belong to (seq_len={}, owned_by_current={}); running the real forward instead", + acg::pending_token_count(), + seq_len, + acg::pending_owned_by_current(), + ); + acg::stand_down(); } // The fixed-capacity KV window width. Read from the graph-mode mask's @@ -2388,9 +2408,7 @@ impl Pipeline for NormalPipeline { cd.set_alloc_cache_enabled(true); arm_alloc_cache_capacity(cd, &alloc_cache_device); } - AllocCacheAction::DrainAndDisable => { - cd.set_alloc_cache_enabled(false) - } + AllocCacheAction::DrainAndDisable => cd.set_alloc_cache_enabled(false), // Steady-state decode. The cap has to be re-armed // here and not only on the Enable edge: free VRAM // shrinks token by token as the KV cache grows, so @@ -2743,7 +2761,8 @@ impl Pipeline for NormalPipeline { ); if arc_layer_bisect() { let bufs = - crate::layers::arc_layer_trace_take(); + crate::layers::arc_layer_trace_take( + ); tracing::error!( "ARC BISECT: holding {} captured layer buffers", bufs.as_ref() @@ -3779,10 +3798,7 @@ mod alloc_cache_policy_tests { /// enabling the cache. #[test] fn degenerate_zero_length_forward_does_not_enable() { - assert_eq!( - alloc_cache_action(0, false, false), - AllocCacheAction::Leave - ); + assert_eq!(alloc_cache_action(0, false, false), AllocCacheAction::Leave); assert_eq!( alloc_cache_action(0, true, false), AllocCacheAction::DrainAndDisable diff --git a/mistralrs-core/src/pipeline/sampling.rs b/mistralrs-core/src/pipeline/sampling.rs index b87f9cf86..3e99d7212 100644 --- a/mistralrs-core/src/pipeline/sampling.rs +++ b/mistralrs-core/src/pipeline/sampling.rs @@ -548,39 +548,83 @@ pub async fn sample_and_add_toks( /// refuses. // `&mut` only because `Sequence::sampler()` takes `&mut self` (`sequence.rs:885`) // to hand back a clone of an `Arc`. Nothing here mutates the sequence. -#[cfg(feature = "cuda")] +// +// Deliberately NOT `cfg(feature = "cuda")`: only the CUDA plumbing can *park* +// tokens, but the take/stand-down half is plain host code, and gating it made +// the cross-request leak untestable on every machine that reviewed it. On a +// non-CUDA build nothing ever parks, `take_pending_token` is always `None`, +// and this is one thread-local read per sample. fn device_loop_pre_sampled_token( seq: &mut Sequence, return_logprobs: bool, sample_speculative: bool, -) -> Option { +) -> Result> { + let seq_id = *seq.id(); let eligible = !return_logprobs && !sample_speculative && seq.sampler().is_greedy_trivial() && matches!(seq.recognizer, SequenceRecognizer::None); - arc_cuda_graph::set_device_loop_eligible(eligible); + arc_cuda_graph::set_device_loop_eligible(seq_id, eligible); + // Did the forward serve this step from the pending queue WITHOUT + // launching? Then the logits this sample holds alias graph-owned storage + // from an earlier step โ€” sound only if a parked token is taken below. + // Consumed here, once, whatever the outcome. + let aliased_logits = arc_cuda_graph::take_aliased_logits_marker(); if !eligible { // Stand down rather than merely decline: a token parked by an earlier, // eligible sequence must never be handed to this one. arc_cuda_graph::stand_down(); - return None; + if aliased_logits { + candle_core::bail!( + "ArcGraph: this step's logits were served from the device loop's pending queue \ + for another sequence and alias graph-owned storage; this sequence (id {seq_id}) \ + is not device-loop eligible and cannot consume them. Failing the step rather \ + than sampling stale logits." + ); + } + return Ok(None); + } + match arc_cuda_graph::take_pending_token(seq_id) { + arc_cuda_graph::PendingTake::Taken(token) => Ok(Some(Logprobs { + token, + logprob: 0.0, + bytes: None, + top_logprobs: None, + })), + // Nothing parked: the forward really ran for this step (the pending + // short-circuit only fires with a non-empty queue), so the logits are + // real and the host samples them. `aliased_logits` cannot be set here + // in a consistent engine: it implies the queue was non-empty when the + // forward ran and empty now, on the same thread with no take between. + arc_cuda_graph::PendingTake::Empty => { + if aliased_logits { + candle_core::bail!( + "ArcGraph: the forward served this step from the pending queue but the queue \ + is empty at sampling time โ€” the aliased logits cannot be sampled. Failing \ + the step." + ); + } + Ok(None) + } + // The queue belonged to another sequence and was dropped whole (and + // logged: "ArcGraph: dropped N foreign parked tokens"). If the + // forward ALSO short-circuited on that foreign queue, these logits + // are a stale alias and sampling them would hand this sequence a + // plausible token from the other sequence's distribution โ€” fail the + // step loudly instead. Without the marker the forward genuinely ran + // (e.g. a batched eager step), the logits are this sequence's own, + // and sampling proceeds normally. + arc_cuda_graph::PendingTake::Foreign { dropped } => { + if aliased_logits { + candle_core::bail!( + "ArcGraph: this step's logits alias graph storage served against {dropped} \ + parked token(s) belonging to another sequence; sampling them would leak \ + another user's distribution into sequence {seq_id}. Failing the step." + ); + } + Ok(None) + } } - arc_cuda_graph::take_pending_token().map(|token| Logprobs { - token, - logprob: 0.0, - bytes: None, - top_logprobs: None, - }) -} - -/// Without CUDA there is no device loop, so nothing is ever pre-sampled. -#[cfg(not(feature = "cuda"))] -fn device_loop_pre_sampled_token( - _seq: &mut Sequence, - _return_logprobs: bool, - _sample_speculative: bool, -) -> Option { - None } /// Async sample optionally adding to trie. @@ -600,7 +644,7 @@ pub async fn sample_sequence( // synchronization in the decode step, so this is the half of the change // that makes removing `replay()`'s `cudaStreamSynchronize` worth anything. if let Some(pre_sampled) = - device_loop_pre_sampled_token(seq, return_logprobs, sample_speculative) + device_loop_pre_sampled_token(seq, return_logprobs, sample_speculative)? { return Ok(pre_sampled); } @@ -993,3 +1037,323 @@ mod speculative_verification_tests { ); } } + +/// ArcInfer/ArcGraph: a device-loop burst that outlives its sequence must +/// never reach the next sequence. +/// +/// The burst runs with `eos_token_id = -1` (`device_decode_burst` passes -1, +/// so no on-device EOS truncation ever fires) and parks its full length; the +/// engine drains one token per step. A sequence that stops mid-burst โ€” stop +/// string, `max_tokens`, cancel, error โ€” leaves the tail parked, and +/// `device_loop_pre_sampled_token` hands parked tokens to whichever eligible +/// sequence samples next, **including the next request's first sample after +/// its prefill**. That is one user's tokens inside another user's response. +/// +/// The fix under test is the funnel: `Sequence::set_state` stands the device +/// loop down on every transition out of the running set. These tests drive the +/// REAL `sample_sequence` path, so they fail loudly if that call is removed โ€” +/// unlike the arc-cuda-graph unit tests, which only ever exercised +/// `clear_pending_tokens` itself, not that anything calls it. +#[cfg(test)] +mod device_loop_cross_sequence_leak_tests { + use super::*; + use crate::sampler::Sampler; + use crate::sequence::{SeqStepType, SequenceGroup, SequenceState, StopReason}; + use candle_core::Device; + use rand::SeedableRng; + use std::cell::RefCell; + + const VOCAB: usize = 32; + /// Argmax of every logits tensor handed to sequence A. + const PEAK_A: u32 = 7; + /// Argmax of every logits tensor handed to sequence B. Distinct from + /// `PEAK_A` and from the burst, so the origin of B's token is unambiguous. + const PEAK_B: u32 = 9; + /// The burst sequence A's device loop parked. A drains the first token, + /// then stops mid-burst, leaving `[22, 23, 24]`. + const BURST: [i32; 4] = [21, 22, 23, 24]; + + thread_local! { + /// Keeps each fixture sequence's `Receiver` alive; a dropped receiver + /// models a disconnected client and would silently change behaviour. + static LIVE_CLIENTS: RefCell>> = + const { RefCell::new(Vec::new()) }; + } + + /// A greedy-trivial sequence: no temperature, penalties, bias, processors + /// or recognizer โ€” exactly the shape `device_loop_pre_sampled_token` + /// accepts, so the pending queue is live for it. `id` matters: the + /// pending queue is owner-tagged by sequence id, so every test sequence + /// must carry a distinct one, as real requests do. + fn greedy_seq(id: usize) -> Sequence { + let (dummy_sender, rx) = tokio::sync::mpsc::channel(1); + LIVE_CLIENTS.with(|k| k.borrow_mut().push(rx)); + let group = Arc::new(std::sync::Mutex::new(SequenceGroup::new( + 1, false, false, None, + ))); + let sampler = Sampler::new( + None, + 0, + None, + None, + None, + None, + None, + -1, + 1.0, + 0.0, + None, + vec![], + None, + ) + .unwrap(); + Sequence::new_waiting( + vec![1u32; 4], + String::new(), + id, + 0, + 1, + dummy_sender, + sampler, + vec![], + vec![], + None, + false, + false, + group, + 0, + 0, + SequenceRecognizer::None, + None, + None, + None, + None, + None, + None, + None, + SeqStepType::PromptAndDecode, + None, + None, + None, + false, + vec![], + ) + } + + /// `(1, 1, VOCAB)` logits peaking hard at `peak` โ€” the shape + /// `sample_sequence` receives for one decode step of one sequence. + fn peaked_logits(peak: u32) -> Tensor { + let mut raw = vec![0f32; VOCAB]; + raw[peak as usize] = 20.0; + Tensor::from_vec(raw, (1, 1, VOCAB), &Device::Cpu).unwrap() + } + + async fn sample(seq: &mut Sequence, peak: u32, rng: &Arc>) -> u32 { + sample_sequence( + peaked_logits(peak), + seq, + false, + rng.clone(), + false, + false, + false, + ) + .await + .unwrap() + .token + } + + /// THE leak test: A parks a burst, stops mid-burst on a non-fallback path + /// (`max_tokens` โ†’ `StopReason::Length`, the transition + /// `finish_or_add_toks_to_seq` performs at `sampling.rs:267`/`:277`), and + /// B's first sample must consume nothing of A's. + /// + /// Mutation check (run both ways before shipping): with the + /// `stand_down` arm removed from `Sequence::set_state`, this fails at the + /// queue assertion with 3 parked tokens, and B's first token comes back + /// `22` โ€” A's first undrained parked id โ€” instead of `PEAK_B`. + #[tokio::test] + async fn a_sequence_stopping_mid_burst_leaks_nothing_into_the_next_sequence() { + // This thread's queue, not shared with other tests (thread-local). + arc_cuda_graph::stand_down(); + let rng = Arc::new(std::sync::Mutex::new(Isaac64Rng::seed_from_u64(0))); + + // โ”€โ”€ Sequence A, decode step N: queue empty, host argmax, and the + // pre-hook publishes A's eligibility โ€” the same order the engine + // runs (sample N publishes, forward N+1 bursts). + let mut seq_a = greedy_seq(1); + assert_eq!(sample(&mut seq_a, PEAK_A, &rng).await, PEAK_A); + + // โ”€โ”€ Forward N+1: the burst parks its FULL length (eos -1). + arc_cuda_graph::push_pending_tokens(&BURST); + + // โ”€โ”€ Step N+1's sample drains one parked token. This is the control + // that proves the pending-token path is LIVE in this build: if the + // pre-hook were compiled out or inert, this would return `PEAK_A` + // and the leak below could never manifest โ€” a test that cannot fail. + assert_eq!( + sample(&mut seq_a, PEAK_A, &rng).await, + BURST[0] as u32, + "the device-loop take path is not live in this build; the leak assertions below \ + would be vacuous" + ); + + // โ”€โ”€ A stops MID-BURST via a non-fallback path: max_tokens. + seq_a.set_state(SequenceState::Done(StopReason::Length(2))); + + // โ”€โ”€ The claim, part 1: nothing of A's survives its completion. + assert_eq!( + arc_cuda_graph::pending_token_count(), + 0, + "sequence A finished mid-burst but its parked tokens survived; the next \ + sequence's first sample will consume them" + ); + + // โ”€โ”€ Sequence B's FIRST sample (the one right after its prefill). + let mut seq_b = greedy_seq(2); + let b_first = sample(&mut seq_b, PEAK_B, &rng).await; + + // โ”€โ”€ The claim, part 2: B's token comes from B's own logits and is + // none of A's parked ids. + assert!( + !BURST.contains(&(b_first as i32)), + "cross-request token leak: sequence B's first sampled token is {b_first}, one of \ + sequence A's parked burst tokens" + ); + assert_eq!( + b_first, PEAK_B, + "sequence B's first token must be the argmax of B's own logits" + ); + } + + /// Every transition out of the running set clears the queue โ€” not just + /// `Done(Length)`. This pins the whole funnel map: client cancel, + /// forward-error/panic recovery (`Error`), PagedAttention eviction + /// (`FinishedAborted`/`FinishedIgnored`), and preemption + /// (`Waiting`/`Swapped`), where consuming the stale burst after + /// re-prefill would duplicate output. + #[tokio::test] + async fn every_transition_out_of_the_running_set_clears_parked_tokens() { + for state in [ + SequenceState::Done(StopReason::Canceled), + SequenceState::Done(StopReason::StopString { + stop_string_idx: 0, + completion_bytes_pos: 0, + }), + SequenceState::Error, + SequenceState::FinishedAborted, + SequenceState::FinishedIgnored, + SequenceState::Waiting, + SequenceState::Swapped, + ] { + arc_cuda_graph::stand_down(); + arc_cuda_graph::push_pending_tokens(&BURST); + assert_eq!(arc_cuda_graph::pending_token_count(), BURST.len()); + let seq = greedy_seq(3); + seq.set_state(state); + assert_eq!( + arc_cuda_graph::pending_token_count(), + 0, + "transition to {state:?} left parked tokens for the next sequence" + ); + } + } + + /// The interleave the completion funnel CANNOT see: the DefaultScheduler + /// bucketing waitlist moves a running sequence aside **without a state + /// modification** (`default_scheduler.rs`, `bucket_and_waitlist_seqs_waiting` + /// says so in its doc), so no `set_state` ever fires for A. The owner tag + /// is the structural defence: B's first sample must consume nothing of + /// A's โ€” the foreign queue is dropped whole (logged as + /// "ArcGraph: dropped N foreign parked tokens") rather than served, to A + /// included, whose device-side burst state is stale by then. + /// + /// Mutation check (run both ways before shipping): with the owner + /// comparison in `take_pending_token` forced to `true`, B's first token + /// comes back `22` โ€” A's parked id โ€” instead of `PEAK_B`. + #[tokio::test] + async fn a_waitlisted_sequence_with_no_state_change_cannot_leak_into_the_next() { + arc_cuda_graph::stand_down(); + let rng = Arc::new(std::sync::Mutex::new(Isaac64Rng::seed_from_u64(0))); + + // โ”€โ”€ A (id 31) decodes at batch 1; its burst parks; it drains one + // (the control proving the take path is live in this build). + let mut seq_a = greedy_seq(31); + assert_eq!(sample(&mut seq_a, PEAK_A, &rng).await, PEAK_A); + arc_cuda_graph::push_pending_tokens(&BURST); + assert_eq!(sample(&mut seq_a, PEAK_A, &rng).await, BURST[0] as u32); + + // โ”€โ”€ A is WAITLISTED: no set_state, no funnel, nothing runs. Its + // undrained tail [22, 23, 24] is still parked when B is scheduled. + + // โ”€โ”€ B (id 32) runs; its forward really ran, so its logits are real. + let mut seq_b = greedy_seq(32); + let b_first = sample(&mut seq_b, PEAK_B, &rng).await; + assert!( + !BURST.contains(&(b_first as i32)), + "cross-request token leak through the waitlist interleave: B's first sampled token \ + is {b_first}, one of A's parked burst tokens" + ); + assert_eq!(b_first, PEAK_B, "B must sample its own logits"); + + // โ”€โ”€ A's foreign-tagged tail was dropped whole, not left to resurface. + assert_eq!( + arc_cuda_graph::pending_token_count(), + 0, + "the foreign queue must be dropped at B's take, not survive it" + ); + + // โ”€โ”€ A resumes: nothing stale to consume, so it host-samples its own + // logits. (Its dropped tokens were computed against device state that + // no longer matches; replaying them would corrupt A's output.) + assert_eq!(sample(&mut seq_a, PEAK_A, &rng).await, PEAK_A); + } + + /// The double fault: the forward served a step from the pending queue + /// (returning the ALIASED stale logits tensor, launching nothing) and the + /// scheduler then swapped sequences before the sample ran. The taker is + /// foreign, so there is no parked token to return โ€” and the logits in + /// hand are another sequence's stale graph output. Sampling them would + /// leak A's distribution into B; the step must fail loudly instead. + #[tokio::test] + async fn an_aliased_step_served_for_another_sequence_fails_instead_of_sampling_stale_logits() { + arc_cuda_graph::stand_down(); + let rng = Arc::new(std::sync::Mutex::new(Isaac64Rng::seed_from_u64(0))); + + // A (id 41) parks a burst; the forward short-circuits on it. + let mut seq_a = greedy_seq(41); + assert_eq!(sample(&mut seq_a, PEAK_A, &rng).await, PEAK_A); + arc_cuda_graph::push_pending_tokens(&BURST); + arc_cuda_graph::note_aliased_logits_served(); + + // B (id 42) is what actually samples that step. + let mut seq_b = greedy_seq(42); + let res = sample_sequence( + peaked_logits(PEAK_B), + &mut seq_b, + false, + rng.clone(), + false, + false, + false, + ) + .await; + assert!( + res.is_err(), + "a stale aliased-logits step served against a foreign queue must fail the step, \ + not return a plausible token" + ); + + // Control: the marker is harmless when the owner itself samples โ€” the + // parked token is taken and the aliased tensor is never read. + let mut seq_c = greedy_seq(43); + assert_eq!(sample(&mut seq_c, PEAK_A, &rng).await, PEAK_A); + arc_cuda_graph::push_pending_tokens(&[25]); + arc_cuda_graph::note_aliased_logits_served(); + assert_eq!( + sample(&mut seq_c, PEAK_A, &rng).await, + 25, + "the owner's own aliased-served step must take its parked token normally" + ); + } +} diff --git a/mistralrs-core/src/sequence.rs b/mistralrs-core/src/sequence.rs index 46c8cdd41..bdb5b1638 100644 --- a/mistralrs-core/src/sequence.rs +++ b/mistralrs-core/src/sequence.rs @@ -1085,6 +1085,41 @@ impl Sequence { let mut group = get_mut_group!(self); group.n_choices = group.n_choices.saturating_sub(1); } + // ArcInfer/ArcGraph: this is the one funnel every path that takes a + // sequence out of the running set passes through โ€” EOS / stop-token / + // stop-string / length / tool-calls (`pipeline/sampling.rs` + // `finish_or_add_toks_to_seq`), the schedulers' disconnected-client + // sweeps and `TERMINATE_ALL_NEXT_STEP` (`Done(Canceled)`), the + // engine's forward-error and panic recovery (`SequenceState::Error` + // via `handle_pipeline_forward_error!` / `step_catching_panics`), and + // PagedAttention eviction (`FinishedAborted` / `FinishedIgnored` / + // `Waiting` / `Swapped`). + // + // A device decode loop burst runs with `eos_token_id = -1`, so it + // always parks its full burst; a sequence that stops mid-burst leaves + // tokens parked, and a parked token is handed to the NEXT sequence's + // first `sample_sequence` โ€” including the one after its prefill โ€” as + // if it were its own. That is one user's tokens in another user's + // response, which is why the stand-down happens HERE, on the + // transition itself, rather than on N scattered completion paths. + // + // The pending queue is thread-local to the engine thread that parked + // it (the same thread that samples at batch 1); on any other thread + // this clears an empty queue, which is a no-op. Clearing on + // preemption (`Waiting` / `Swapped`) is deliberate fail-closed: the + // sequence re-prefills on re-admission, so its parked tokens are + // stale and consuming them would duplicate output. + match state { + SequenceState::Done(_) + | SequenceState::Error + | SequenceState::Waiting + | SequenceState::FinishedAborted + | SequenceState::FinishedIgnored + | SequenceState::Swapped => arc_cuda_graph::stand_down(), + SequenceState::RunningPrompt + | SequenceState::RunningCompletion + | SequenceState::RunningPrefillPrompt => {} + } *self.state.write().unwrap() = state; } diff --git a/mistralrs-core/tests/capability_reachability.rs b/mistralrs-core/tests/capability_reachability.rs index 7e1d0d7a7..18fa40eeb 100644 --- a/mistralrs-core/tests/capability_reachability.rs +++ b/mistralrs-core/tests/capability_reachability.rs @@ -401,6 +401,20 @@ static REGISTRY: &[Capability] = &[ blocking condition as the top_k entry above.", }, }, + Capability { + name: "ArcGraph device decode loop opt-in gate (ARC_GRAPH_DEVICE_LOOP)", + parent: "ArcInfer / ArcGraph", + check: Check::Symbol { + symbol: "device_loop_enabled", + defined_in: "arc-cuda-graph/src/device_loop.rs", + }, + // The gate is opt-in default-off, but the CALL PATH must be live: + // `normal.rs::forward_inputs` consults `device_loop_enabled()` before + // `device_decode_burst`. If this reference disappears, the env var + // parses forever and engages nothing โ€” the switched-off class this + // file exists to catch. + status: Status::Live, + }, ]; /// Symbols that are defined in Arc-owned source and have **no production From 4033b8f4b94c351cd3349532d110e5b599752b88 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Sat, 22 Aug 2026 03:01:32 +0100 Subject: [PATCH 21/21] =?UTF-8?q?perf(ArcKernels/ArcAttention/ArcQuant):?= =?UTF-8?q?=20session-10=20measured=20default=20flips=20=E2=80=94=20GEMV?= =?UTF-8?q?=5FWIDE=20ON,=20QK=5FFUSED=5FCOHORT=20ON,=20CUBLAS=5FMIN=5FM=20?= =?UTF-8?q?512=E2=86=9264=20(stacked:=20B=3D256=20575=E2=86=92758=20tok/s,?= =?UTF-8?q?=20b=3D1=2040.8=E2=86=9243.8)=20(#223)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * perf(ArcKernels): ARC_FP8_GEMV_WIDE default ON โ€” measured +7.5% b=1 (43.80 vs 40.76 tok/s) Session-10 hardware A/B, box arc-s10-ledger (H200), candidate 4d03b9e25, leg 3: 3x512 produced tokens, seeded t=0.7 p=0.95, engagement proven by [arc-fp8-dispatch] path=gemv_wide; canary1 byte-identical, canary2 coherent and arithmetically correct (numerics within the stated f32 re-association bound). The ledger's ~1.6x prediction did NOT materialise end-to-end; the default stands on the measured +7.5%. Polarity flips from 'only "1" enables' to 'only "0" disables' โ€” still value-read, never presence-read (#212). Test renamed accordingly (only_literal_zero_disables) with the assertion values swapped. Registry entry added (capability_reachability.rs, ArcKernels). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf * perf(ArcAttention): ARC_QK_FUSED_COHORT default ON โ€” engaged + bit-identical at B=256 (599.35 vs 575.12 tok/s) Session-10 hardware A/B, box arc-s10-ledger (H200), candidate 4d03b9e25, leg 4: fused qk_norm_rope ENGAGED at B=256 (the default arm DECLINED and took the eager chain), seeded canaries byte-identical, 0/512 request errors, aggregate 599.35 vs 575.12 tok/s. The +4.2% is inside the ~ยฑ5% inter-leg variance measured the same session (leg 4b no-op delta), so the flip stands on engaged + bit-identical + not-slower, exactly the bar the capability entry set. Doc comment and registry entry updated as the old comment promised; decline message now names the kill switch (ARC_QK_FUSED_COHORT=0). Polarity stays value-read (#212). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf * perf(ArcKernels/ArcQuant): ARC_FP8_CUBLAS_MIN_M 512 -> 64 โ€” the three-arm sweep this constant's own doc demanded Session-10, box arc-s10-ledger (H200), candidate 4d03b9e25, arc-tools/fp8_threshold_sweep.sh, LADDER {8,64,256}, one binary, flatten fix present, >=1000-token floor. Clean rows (tok/s): M=64: cublaslt 552.4 | wmma 328.2 | tiled 223.1 -> cuBLASLt wins M=256: cublaslt 715.9 | wmma 620.4 | tiled 285.8 -> cuBLASLt wins First clean M where dequant+cuBLASLt beats min(WMMA, tiled): 64. M=8 rows were VOID on the token floor, so 5..63 stays with the native arms โ€” a void row makes no claim. fp8_gemv_warp still owns M<=4 (the M=1 floor row, cuBLASLt 0.72x, still stands). Rows quoted in the doc per its own rule: a threshold is a claim about every value it excludes. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf * perf(ArcKV): ARC_V4_FP8_KV default ON โ€” the U8 CUDA leg finally measured: engaged, byte-identical, not slower Session-10, box arc-s10-ledger (H200), candidate 4d03b9e25. Leg 10: seeded canaries byte-identical to the BF16 arm, b=1 41.49 vs 40.76 tok/s. The lane has no engagement log line (silent-success class, noted for follow-up), so engagement was proven by nsys kernel capture instead: arc_kv_fp8_quantize_kernel x16,297 + arc_kv_fp8_dequantize_kernel x32,594 in one b=1 run. 4-gate stacked verification (this + the three earlier flips, one binary): B=256 744.7 tok/s (vs 757.5 three-gate, inside the ยฑ5% band) with canaries IDENTICAL through the stack; b=1 44.68 tok/s โ€” the session's best b=1. Unlike wave43-BU (same default, shipped unrun, broke every forward), this default stands on a measurement. ARC_V4_FP8_KV=0 restores BF16; the ARC_V4_CAPTURE_PROBE veto is unconditional and now tested for the default arm too (write_kv_inplace has no U8 variant). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf --------- Co-authored-by: Claude Fable 5 --- mistralrs-core/src/cuda/qk_norm_rope.rs | 50 +++++++++++-------- mistralrs-core/src/models/deepseek4.rs | 48 ++++++++++++------ .../tests/capability_reachability.rs | 38 +++++++++----- .../src/blockwise_fp8/gemv_wide.rs | 47 +++++++++-------- mistralrs-quant/src/blockwise_fp8/mod.rs | 38 ++++++++++---- mistralrs-quant/src/blockwise_fp8/ops.rs | 15 +++--- 6 files changed, 152 insertions(+), 84 deletions(-) diff --git a/mistralrs-core/src/cuda/qk_norm_rope.rs b/mistralrs-core/src/cuda/qk_norm_rope.rs index 013d7f41b..3f15456dc 100644 --- a/mistralrs-core/src/cuda/qk_norm_rope.rs +++ b/mistralrs-core/src/cuda/qk_norm_rope.rs @@ -49,11 +49,16 @@ pub fn fused_enabled() -> bool { *ENABLED.get_or_init(|| !matches!(std::env::var("ARC_QK_FUSED").as_deref(), Ok("0"))) } -/// `ARC_QK_FUSED_COHORT=1` lets the fused kernel serve a **batched** cohort -/// (`batch > 1`). **Default OFF**, so `batch > 1` declines and takes the eager -/// chain exactly as it does today. +/// `ARC_QK_FUSED_COHORT=0` stops the fused kernel from serving a **batched** +/// cohort (`batch > 1`). **Default ON since the s10 measurement** (box +/// `arc-s10-ledger`, H200, candidate `4d03b9e25`, leg 4): B=256 decode +/// aggregate 599.35 vs 575.12 tok/s baseline, seeded canaries byte-identical, +/// ENGAGED line present at B=256, 0/512 request errors. The +4.2% sits within +/// the observed inter-leg variance band (~ยฑ5%, leg 4b's no-op delta), so the +/// default stands on "engaged + bit-identical + not slower", not on a claimed +/// speedup. /// -/// # Why this is opt-in and [`fused_enabled`] is not +/// # Why this was opt-in until measured, and [`fused_enabled`] never was /// /// The kernel is batch-correct by construction โ€” `qk_norm_rope.cu` flattens /// `(b, t)` into `blockIdx.y`, recovers `b = bt / seq_len`, and reads the table @@ -72,14 +77,13 @@ pub fn fused_enabled() -> bool { /// attention mask, at 12% agreement with the reference, because nobody ran /// the path it had quietly become live on. /// -/// So the doctrine is: **unverified means default-off, and "unverified" means -/// unmeasured, not new.** `batch == 1` is unaffected โ€” that path is already -/// live and already exercised. When the cohort path has been measured correct -/// and faster on hardware (`ARC_QK_VERIFY=1` bit-compares it against the eager -/// chain at every layer), this default flips and this doc comment goes with it. +/// The doctrine was: **unverified means default-off, and "unverified" means +/// unmeasured, not new.** That measurement has now happened (s10 leg 4, above) +/// and the default flipped in the same session, exactly as the old comment +/// promised. `batch == 1` was always unaffected. /// -/// Read by VALUE, `== Some("1")`, not by presence: 23 `ARC_*` flags were -/// converted this week (#212) precisely because `var_os(..).is_some()` made +/// Read by VALUE, `!= Some("0")`, not by presence: 23 `ARC_*` flags were +/// converted (#212) precisely because `var_os(..).is_some()` made /// `ARC_FOO=0` mean ON. pub fn fused_cohort_enabled() -> bool { use std::sync::OnceLock; @@ -100,7 +104,7 @@ pub fn fused_cohort_enabled() -> bool { /// pure function against every input is the only form of this test that means /// anything. pub fn fused_cohort_enabled_from(value: Option<&str>) -> bool { - value == Some("1") + value != Some("0") } /// `ARC_QK_VERIFY=1` runs the eager chain alongside the fused kernel at every @@ -514,23 +518,29 @@ mod cohort_gate_tests { /// These run with no features and no GPU, on three operating systems, on /// every PR, because the polarity is a pure function. #[test] - fn the_batched_cohort_is_off_unless_explicitly_enabled() { + fn the_batched_cohort_is_on_by_default_since_the_s10_measurement() { assert!( - !fused_cohort_enabled_from(None), - "unset must mean OFF โ€” the cohort path has never run on hardware" + fused_cohort_enabled_from(None), + "unset must mean ON โ€” measured s10 leg 4 (arc-s10-ledger, 4d03b9e25): \ + engaged at B=256, canaries byte-identical, 599.35 vs 575.12 tok/s" ); assert!(fused_cohort_enabled_from(Some("1"))); } /// Read by VALUE, not presence. #212 converted 23 `ARC_*` flags because /// `var_os(..).is_some()` made `ARC_FOO=0` mean ON โ€” which also silently - /// cancels any A/B whose control arm sets the flag to zero. + /// cancels any A/B whose control arm sets the flag to zero. After the s10 + /// flip the kill switch is the exact string "0" and nothing else. #[test] - fn zero_means_off_and_so_does_every_other_value() { - for v in ["0", "", "false", "no", "off", "true", "yes", "on", "2", "1 "] { + fn only_literal_zero_disables_the_cohort() { + assert!( + !fused_cohort_enabled_from(Some("0")), + "\"0\" must disable โ€” it is the A/B control arm and the kill switch" + ); + for v in ["", "false", "no", "off", "true", "yes", "on", "2", "1 ", "0 "] { assert!( - !fused_cohort_enabled_from(Some(v)), - "{v:?} must not enable the cohort path; only the exact string \"1\" does" + fused_cohort_enabled_from(Some(v)), + "{v:?} must leave the measured default ON; only the exact string \"0\" disables" ); } } diff --git a/mistralrs-core/src/models/deepseek4.rs b/mistralrs-core/src/models/deepseek4.rs index 5b007f92b..64361e466 100644 --- a/mistralrs-core/src/models/deepseek4.rs +++ b/mistralrs-core/src/models/deepseek4.rs @@ -1338,7 +1338,7 @@ impl Attention { // `qk::fused_cohort_enabled` for the doctrine and for how to turn the // cohort path on once it has been measured on hardware. if qb > 1 && !qk::fused_cohort_enabled() { - decline!("batched cohort is opt-in (set ARC_QK_FUSED_COHORT=1)"); + decline!("batched cohort disabled by ARC_QK_FUSED_COHORT=0 (default is ON since s10)"); } match kv_normed.dims3() { Ok((kb, kt, kh)) if kb == qb && kt == qt && kh == head_dim => {} @@ -1991,8 +1991,8 @@ impl Attention { graph_attn } None => { - // FP8 code storage is opt-in (`ARC_V4_FP8_KV=1`); unset stores - // the dense BF16 K the model has always stored. + // FP8 code storage is default-ON since s10 (`ARC_V4_FP8_KV=0` + // restores the dense BF16 K layout). let cached = { let _s = arc_profiler::device_span("kv_cache_append"); append_kv_mqa( @@ -3159,16 +3159,21 @@ fn append_graph_kv_mqa( Ok(k_full) } -/// FP8 storage for the cached K. **Opt-in: off unless `ARC_V4_FP8_KV=1`.** +/// FP8 storage for the cached K. **Default ON since the s10 GPU A/B; +/// `ARC_V4_FP8_KV=0` restores the BF16 layout.** /// /// wave43-BU shipped this defaulted ON without ever running it on a GPU โ€” its /// own notes said *"needs a GPU: first CUDA exercise of a U8 `SingleCache`"* โ€” /// and the first V4 forward on a commit that contained it died, every request, /// with `dtype mismatch in slice-set, lhs: BF16, rhs: U8` (wave48-BY). The /// layout is sound and bit-exact (see [`super::dsv4_kv_fp8`]); what was never -/// exercised is the CUDA leg of it. Until a GPU A/B has shown -/// `ARC_V4_FP8_KV=1` and unset producing token-identical greedy output, the -/// default is the layout that has actually served (wave49-BZ). +/// exercised was the CUDA leg. That A/B has now run: s10 leg 10 (box +/// `arc-s10-ledger`, H200, candidate `4d03b9e25`) โ€” seeded canaries +/// byte-identical to the BF16 arm, b=1 41.49 vs 40.76 tok/s (not slower), and +/// engagement PROVEN by nsys kernel capture (`arc_kv_fp8_quantize_kernel` +/// 16,297 launches, `arc_kv_fp8_dequantize_kernel` 32,594 in one b=1 run โ€” +/// the U8 CUDA leg genuinely executed). This is the difference from wave43: +/// the same default, now standing on a measurement instead of a hope. /// /// Also forced off under `ARC_V4_CAPTURE_PROBE`: the CUDA-graph decode arm /// writes through `mistralrs_quant::kvwrite::write_kv_inplace`, which is @@ -3228,7 +3233,10 @@ fn v4_fp8_kv_enabled() -> bool { /// about the default, which is exactly the kind of test DOCTRINE D12 counts as /// worse than none. fn fp8_kv_enabled_from(var: Option<&str>, capture_probe: bool) -> bool { - var == Some("1") && !capture_probe + // Default ON since s10 (measured; see the doc above). The capture-probe + // veto is unconditional: write_kv_inplace has no U8 arm, so the graph + // decode path and the U8 cache remain mutually exclusive. + var != Some("0") && !capture_probe } /// RUN-161 decode profiler. Enabled by `ARC_TIME_DECODE=1`. Accumulates @@ -5924,18 +5932,28 @@ mod kv_footprint_tests { /// purpose: the latter is a `OnceLock`, so a test that set the variable /// would resolve it once and prove nothing about the default. #[test] - fn v4_fp8_kv_is_opt_in() { - assert!(!fp8_kv_enabled_from(None, false), "unset must be OFF"); - assert!(!fp8_kv_enabled_from(Some("0"), false)); - assert!(!fp8_kv_enabled_from(Some(""), false)); - assert!(!fp8_kv_enabled_from(Some("true"), false)); - assert!(!fp8_kv_enabled_from(Some("on"), false)); - assert!(!fp8_kv_enabled_from(Some("2"), false)); + fn v4_fp8_kv_is_default_on_with_a_literal_zero_kill_switch() { + assert!( + fp8_kv_enabled_from(None, false), + "unset must be ON โ€” measured s10 leg 10 (arc-s10-ledger, 4d03b9e25): \ + canaries byte-identical, engagement proven by nsys kernel capture" + ); + assert!(!fp8_kv_enabled_from(Some("0"), false), "\"0\" is the kill switch"); + assert!(fp8_kv_enabled_from(Some(""), false)); + assert!(fp8_kv_enabled_from(Some("true"), false)); + assert!(fp8_kv_enabled_from(Some("on"), false)); + assert!(fp8_kv_enabled_from(Some("2"), false)); assert!(fp8_kv_enabled_from(Some("1"), false), "=1 must be ON"); assert!( !fp8_kv_enabled_from(Some("1"), true), "ARC_V4_CAPTURE_PROBE must veto it: write_kv_inplace has no U8 arm" ); + assert!( + !fp8_kv_enabled_from(None, true), + "the veto must hold for the DEFAULT too, not just for explicit =1 โ€” \ + the capture probe cannot write a U8 cache regardless of how FP8 KV \ + was selected" + ); } /// The per-row ragged mask is **opt-in**, and off it hands diff --git a/mistralrs-core/tests/capability_reachability.rs b/mistralrs-core/tests/capability_reachability.rs index 18fa40eeb..2ba4de18f 100644 --- a/mistralrs-core/tests/capability_reachability.rs +++ b/mistralrs-core/tests/capability_reachability.rs @@ -193,20 +193,18 @@ static REGISTRY: &[Capability] = &[ // and the fused-512 attention path (four days silently dropping the // attention mask at 12% agreement). // - // flag ARC_QK_FUSED_COHORT=1 + // flag ARC_QK_FUSED_COHORT (default ON since s10; "0" is the kill switch) // gate cuda/qk_norm_rope.rs :: fused_cohort_enabled_from - // test one binary, B>1 decode, same prompt and seed, flag set vs unset; - // ARC_QK_VERIFY=1 bit-compares fused against eager at every layer - // pass outputs bit-identical AND the fused arm faster; engagement - // counter (`qk::engaged_count`) non-zero, since a silently - // declining fast path reports a perfect no-op - // was never measured above batch 1 โ€” the path was unreachable, so there - // is no prior number to quote, honest or otherwise - // 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. + // MEASURED s10 leg 4 (box arc-s10-ledger H200, candidate 4d03b9e25): + // ENGAGED at B=256, seeded canaries byte-identical to the eager + // arm, 599.35 vs 575.12 tok/s aggregate (the +4.2% sits inside the + // ~ยฑ5% inter-leg variance measured the same session โ€” the default + // stands on engaged + bit-identical + not-slower). + // then DONE: the default flipped to ON in the same session, as this + // entry always demanded. Leaving it off after a passing + // measurement is a failure state, not the finish line. Capability { - name: "fused qk_norm_rope at B>1: built and tested, default OFF pending one measurement", + name: "fused qk_norm_rope at B>1: default ON, measured s10", parent: "ArcInfer / ArcAttention", check: Check::Symbol { symbol: "fused_cohort_enabled", @@ -341,6 +339,22 @@ static REGISTRY: &[Capability] = &[ // dequantize+cuBLASLt divert in `blockwise_fp8/mod.rs`. status: Status::Live, }, + // Flipped default-ON in session 10 after its hardware A/B (box + // `arc-s10-ledger`, candidate `4d03b9e25`, leg 3): b=1 decode 43.80 vs + // 40.76 tok/s (+7.5%, 3x512 produced tokens, seeded), engagement proven by + // `[arc-fp8-dispatch] path=gemv_wide`, coherence canary correct. The + // ledger's ~1.6x prediction did NOT materialise end-to-end; +7.5% is the + // measured number and the one this default stands on. `ARC_FP8_GEMV_WIDE=0` + // (exact string) is the kill switch; value-read (#212 doctrine). + Capability { + name: "FP8 wide b=1 GEMV (ARC_FP8_GEMV_WIDE): default ON, measured s10", + parent: "ArcKernels", + check: Check::Symbol { + symbol: "wide_enabled", + defined_in: "mistralrs-quant/src/blockwise_fp8/gemv_wide.rs", + }, + status: Status::Live, + }, // -- the ApiPromise class ------------------------------------------------ Capability { name: "sampling: logit_bias reaches the sampler", diff --git a/mistralrs-quant/src/blockwise_fp8/gemv_wide.rs b/mistralrs-quant/src/blockwise_fp8/gemv_wide.rs index e173f4f01..0f38c182f 100644 --- a/mistralrs-quant/src/blockwise_fp8/gemv_wide.rs +++ b/mistralrs-quant/src/blockwise_fp8/gemv_wide.rs @@ -29,14 +29,17 @@ //! contraction โ€” contraction is identical in both kernels and so cancels out //! of the comparison. //! -//! # Why the gate is opt-in +//! # Why the gate is default-ON //! -//! ๐Ÿ”ด The kernel is UNVERIFIED ON HARDWARE. [`wide_enabled`] reads -//! `ARC_FP8_GEMV_WIDE` and is true **only** for the exact value `"1"`. -//! `ARC_FP8_GEMV_WIDE=0`, `=off`, `=true` and unset all leave it OFF. That -//! strictness is deliberate: this crate already ships `ARC_NO_FP8_GEMV`, read -//! with `std::env::var(..).is_ok()` (`blockwise_fp8/ops.rs`), where `=0` -//! *enables* the disable โ€” the polarity bug this form cannot have. +//! โœ… MEASURED ON HARDWARE 2026-08-22, box `arc-s10-ledger` (H200), candidate +//! `4d03b9e25`, session-10 leg 3: b=1 decode 43.80 tok/s vs 40.76 baseline +//! (+7.5%, 3ร—512 produced tokens, seeded t=0.7 p=0.95), engagement proven by +//! `[arc-fp8-dispatch] path=gemv_wide`, coherence canary correct (numerics +//! shift the sampled continuation within the stated f32 bound; canary1 was +//! byte-identical, canary2 coherent and arithmetically right). Per the flip +//! doctrine the default turned ON in the same session as the measurement. +//! `ARC_FP8_GEMV_WIDE=0` โ€” and only the exact string `"0"` โ€” restores the +//! warp GEMV from the same binary. Value-read, never presence-read (#212). // This module is exactly two things: a gate that only a CUDA build calls, and // a CPU model of the kernel that only `mod tests` calls. On a non-CUDA build @@ -56,16 +59,19 @@ pub(crate) const WARP_K_PER_ITER: usize = 128; /// K-elements one lane of `fp8_gemv_wide` owns per iteration. pub(crate) const WIDE_K_PER_LANE: usize = 16; -/// `ARC_FP8_GEMV_WIDE` โ€” opt-in switch for the wide b=1 GEMV. +/// `ARC_FP8_GEMV_WIDE` โ€” kill switch for the wide b=1 GEMV (default ON). /// /// Split from the env read so the polarity is unit-testable. Only the exact -/// string `"1"` enables. +/// string `"0"` disables; everything else (unset included) leaves the +/// measured-faster wide kernel on. Measured: s10 leg 3, `arc-s10-ledger`, +/// `4d03b9e25`, 43.80 vs 40.76 tok/s b=1. pub(crate) fn wide_enabled_from(value: Option<&str>) -> bool { - value == Some("1") + value != Some("0") } -/// `ARC_FP8_GEMV_WIDE=1` enables `fp8_gemv_wide` for eligible b=1 shapes. -/// Default OFF. Read once; the decode hot loop never touches the env. +/// `ARC_FP8_GEMV_WIDE=0` disables `fp8_gemv_wide` for eligible b=1 shapes. +/// Default ON (measured +7.5% b=1). Read once; the decode hot loop never +/// touches the env. #[cfg(feature = "cuda")] pub(crate) fn wide_enabled() -> bool { use std::sync::OnceLock; @@ -636,21 +642,20 @@ mod tests { ); } - /// Env polarity: only the exact string "1" turns the kernel on. Guards + /// Env polarity after the s10 default flip: only the exact string "0" + /// turns the kernel OFF. Still value-read, never presence-read โ€” guards /// against the `ARC_NO_DEDICATED_DECODE` failure, where `var_os(..) /// .is_some()` makes `=0` *enable* the flag. #[test] - fn only_literal_one_enables() { - assert!(wide_enabled_from(Some("1"))); - for v in [ - "0", "", "on", "off", "true", "false", "yes", "no", "2", " 1", "1 ", - ] { + fn only_literal_zero_disables() { + assert!(!wide_enabled_from(Some("0")), "`ARC_FP8_GEMV_WIDE=0` must disable"); + for v in ["1", "", "on", "off", "true", "false", "yes", "no", "2", " 0", "0 "] { assert!( - !wide_enabled_from(Some(v)), - "`ARC_FP8_GEMV_WIDE={v}` must not enable" + wide_enabled_from(Some(v)), + "`ARC_FP8_GEMV_WIDE={v}` must leave the measured default ON" ); } - assert!(!wide_enabled_from(None), "unset must not enable"); + assert!(wide_enabled_from(None), "unset must stay ON (the measured default)"); } /// The gate must reject every shape whose preconditions the kernel relies diff --git a/mistralrs-quant/src/blockwise_fp8/mod.rs b/mistralrs-quant/src/blockwise_fp8/mod.rs index 5ef1cecc3..897f1e068 100644 --- a/mistralrs-quant/src/blockwise_fp8/mod.rs +++ b/mistralrs-quant/src/blockwise_fp8/mod.rs @@ -686,10 +686,30 @@ pub fn blockwise_fp8_linear_b( /// **The measurement that settles this constant:** re-run the M sweep with the /// flatten fix present, three arms, on one binary โ€” /// `ARC_FP8_CUBLAS_MIN_M`/`ARC_NO_FP8_WMMA` toggle all three with no rebuild โ€” -/// and set the default to the first M at which (a) beats `min((b), (c))`. Until -/// then the default stays where it is, because moving it would be a *second* -/// constant with no measurement behind it, which is the exact defect this doc -/// exists to record. +/// and set the default to the first M at which (a) beats `min((b), (c))`. +/// +/// # โœ… THAT SWEEP HAS NOW RUN โ€” s10, and 64 is the measured crossover +/// +/// 2026-08-22, box `arc-s10-ledger` (H200), candidate `4d03b9e25`, V4-Flash, +/// `arc-tools/fp8_threshold_sweep.sh`, LADDER {8, 64, 256}, three arms, one +/// binary, flatten fix present, >=1000-token floor per row (M=8 rows VOID on +/// the floor and make no claim): +/// +/// ```text +/// M arm tokens wall_s tok/s status +/// 64 tiled 2811 12.6 223.1 CLEAN +/// 64 wmma 2811 8.6 328.2 CLEAN +/// 64 cublaslt 3072 5.6 552.4 CLEAN <- (a) beats min((b),(c)) +/// 256 tiled 12288 43.0 285.8 CLEAN +/// 256 wmma 11244 18.1 620.4 CLEAN +/// 256 cublaslt 12288 17.2 715.9 CLEAN <- and again +/// ``` +/// +/// First clean M where (a) beats min((b), (c)): **64**. The default moves to +/// that measured point. M=8 stays with the native arms (its rows were void, so +/// no claim is made about 5..63 โ€” the conservative side of a void row is to +/// leave it with the incumbent). `fp8_gemv_warp` still owns M <= 4; the M=1 +/// floor row from the earlier sweep (cuBLASLt 0.72x at M=1) still stands. /// /// `fp8_gemv_warp` still owns `M <= 4` and is not affected by this constant. /// @@ -698,11 +718,11 @@ pub fn blockwise_fp8_linear_b( #[cfg(feature = "cuda")] fn arc_fp8_cublas_min_m() -> usize { use std::sync::OnceLock; - // Still 512. NOT because 512 is right -- the doc above shows it is a single - // prefill point -- but because the only sweep in hand measured the scalar - // kernel that #200 displaced, and lowering this to 5 on that evidence would - // make the tensor-core GEMM dead code at every M >= 5. Re-sweep three arms. - const DEFAULT_MIN_M: usize = 512; + // 64 = the first CLEAN M at which dequantize+cuBLASLt beats min(WMMA, + // tiled) in the s10 three-arm sweep (arc-s10-ledger, 4d03b9e25; rows in + // the doc above). M=8 rows were VOID on the token floor, so 5..63 stays + // with the native arms rather than being claimed by interpolation. + const DEFAULT_MIN_M: usize = 64; static CACHE: OnceLock = OnceLock::new(); *CACHE.get_or_init(|| { std::env::var("ARC_FP8_CUBLAS_MIN_M") diff --git a/mistralrs-quant/src/blockwise_fp8/ops.rs b/mistralrs-quant/src/blockwise_fp8/ops.rs index 8d8e8c5c9..347eb8288 100644 --- a/mistralrs-quant/src/blockwise_fp8/ops.rs +++ b/mistralrs-quant/src/blockwise_fp8/ops.rs @@ -1069,11 +1069,12 @@ fn fp8_wmma_eligible(block_size_y: i32, block_size_x: i32) -> bool { /// (see [`gemv_wide::wide_eligible`]); any shape that fails one keeps /// `fp8_gemv_warp`, unchanged, with no error. /// -/// ๐Ÿ”ด Default OFF. `ARC_FP8_GEMV_WIDE=1` โ€” and *only* the literal `"1"` โ€” -/// turns it on. The kernel has never run on a GPU and is not bit-identical to -/// `fp8_gemv_warp` (it re-associates the f32 accumulation; the exact claim and -/// its CPU proof are in `blockwise_fp8/gemv_wide.rs`). It must be A/B'd on -/// hardware against the default before anyone considers flipping it. +/// โœ… Default ON since the s10 hardware A/B (box `arc-s10-ledger`, candidate +/// `4d03b9e25`, leg 3: 43.80 vs 40.76 tok/s b=1, engagement line +/// `path=gemv_wide`, coherent canary). `ARC_FP8_GEMV_WIDE=0` โ€” and *only* the +/// literal `"0"` โ€” restores `fp8_gemv_warp` from the same binary. It is not +/// bit-identical to `fp8_gemv_warp` (it re-associates the f32 accumulation; +/// the exact claim and its CPU proof are in `blockwise_fp8/gemv_wide.rs`). /// /// The pointers passed here are the ones the kernel will receive, so the /// 16-byte alignment a `uint4` load requires is checked on the real address @@ -1323,7 +1324,7 @@ fn fp8_blockwise_matmul_impl( k as usize, ); if let Some(scale_shift) = wide_shift { - // ๐Ÿ”ด UNVERIFIED ON HARDWARE. `ARC_FP8_GEMV_WIDE=1` only. + // โœ… Measured s10 leg 3 (arc-s10-ledger, 4d03b9e25): +7.5% b=1. Default ON; `ARC_FP8_GEMV_WIDE=0` restores warp. unsafe { ffi::launch_fp8_gemv_wide_f16( input_ptr as *const _, @@ -1427,7 +1428,7 @@ fn fp8_blockwise_matmul_impl( k as usize, ); if let Some(scale_shift) = wide_shift { - // ๐Ÿ”ด UNVERIFIED ON HARDWARE. `ARC_FP8_GEMV_WIDE=1` only. + // โœ… Measured s10 leg 3 (arc-s10-ledger, 4d03b9e25): +7.5% b=1. Default ON; `ARC_FP8_GEMV_WIDE=0` restores warp. unsafe { ffi::launch_fp8_gemv_wide_bf16( input_ptr as *const _,