diff --git a/.claude/run_capture.sh b/.claude/run_capture.sh new file mode 100644 index 0000000000..dd2ffb3c6c --- /dev/null +++ b/.claude/run_capture.sh @@ -0,0 +1,74 @@ +#!/bin/bash +# ArcGraph replay-correctness measurement leg (bench mode). +# +# Bench mode, NOT server mode: capture requires ARC_CANDLE_ALLOC_CACHE, whose +# server-mode leak (132 allocs/step, zero frees) OOMs a 143 GB card on the first +# completion. Another agent owns that fix; this leg does not touch it. +# +# Usage: run_capture.sh +set -u +export CUDA_HOME=/usr/local/cuda-12.4 +export PATH=$CUDA_HOME/bin:$HOME/.cargo/bin:$PATH +export LD_LIBRARY_PATH=$CUDA_HOME/lib64:${LD_LIBRARY_PATH:-} + +NAME="$1"; shift +GENLEN="$1"; shift +PROMPTLEN="$1"; shift +LEG_ENV="$*" + +LOGDIR=/root/logs/capture; mkdir -p "$LOGDIR" +LOG="$LOGDIR/${NAME}.log" +# Run under a DIFFERENT process name. The other agent's harness on this box +# reaps with `pkill -x mistralrs`, which matches on the exact process name and +# therefore kills our legs too (it killed two: SIGTERM, BENCH_EXIT=143, one of +# them 4 s into the run). Renaming the argv[0] the kernel sees makes our process +# invisible to that pattern without changing what either side measures β€” the +# file is a copy of the same binary, and its sha256 is recorded below. +SRC=/root/arc-rc/target/release/mistralrs +BIN=/root/arc-rc/target/release/arcgraphbench +cp -f "$SRC" "$BIN" 2>/dev/null || true +cd /root/arc-rc || exit 2 + +# Kill by exact PID, and ONLY our own binary. `pkill -f mistralrs` matches the +# ssh command line that launched this script and has killed agents through SSH; +# `pgrep -x mistralrs` is safe from that but still kills every other agent's +# run on this shared box (it killed one of ours, and we killed one of theirs). +# Resolving /proc/PID/exe scopes the kill to processes started from THIS +# checkout, so concurrent legs from /root/arc or /root/target-alloc survive. +for p in $(pgrep -x arcgraphbench); do + case "$(readlink -f "/proc/$p/exe" 2>/dev/null)" in + /root/arc-rc/*) kill -9 "$p" 2>/dev/null ;; + esac +done +sleep 3 + +# Refuse to start while another agent holds the GPU: two V4 instances do not fit +# in 143 GB, and a leg that OOMs mid-run is a void measurement, not a slow one. +for _ in $(seq 1 120); do + used=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits) + [ "${used:-99999}" -lt 2000 ] && break + sleep 15 +done +used=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits) +if [ "${used:-99999}" -ge 2000 ]; then + echo "VOID: GPU still busy (${used} MiB) after waiting; not measuring." + exit 5 +fi + +# Provenance: a number without the binary that produced it is not evidence. +{ + echo "=== LEG $NAME ===" + echo "git_sha=$(git rev-parse HEAD)" + echo "bin_sha256=$(sha256sum "$BIN" | cut -d' ' -f1)" + echo "env=$LEG_ENV gen_len=$GENLEN prompt_len=$PROMPTLEN" + echo "gpu_free_before=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader)" + echo "date=$(date -Is)" +} > "$LOG" + +# shellcheck disable=SC2086 +env $LEG_ENV "$BIN" bench \ + --prompt-len "$PROMPTLEN" --gen-len "$GENLEN" --iterations 1 \ + auto -m /root/models/v4-src --from-uqff /root/models/v4-uqff/qtip2-0.uqff \ + >> "$LOG" 2>&1 +echo "BENCH_EXIT=$?" >> "$LOG" +echo "--- LOG: $LOG" diff --git a/arc-cuda-graph/src/graph.rs b/arc-cuda-graph/src/graph.rs index d90677216f..c288998f74 100644 --- a/arc-cuda-graph/src/graph.rs +++ b/arc-cuda-graph/src/graph.rs @@ -367,11 +367,32 @@ output_trusted={} verify_remaining={} verify_failed={}", batch_size={batch_size}): cudaError {sync}" ); } + // πŸ”΄ BOOBY TRAP β€” this is a CLONE OF THE CAPTURED OUTPUT TENSOR, i.e. a + // handle to storage the graph OVERWRITES on its next launch. It is not + // a snapshot. Two consequences the caller owns: + // + // * every value this has ever returned aliases every other one, so + // holding an "old" result and comparing it to a "new" one compares a + // tensor with itself; + // * a verification path that launches the graph again β€” even for a + // diagnostic β€” silently rewrites the bytes it was about to grade. + // `normal.rs` therefore takes its verdict from the tensor BEFORE any + // extra launch. This nearly voided the probe that found it. + // + // Copy it if you need it to outlive the next launch. let out = captured .output .clone() .ok_or_else(|| candle_core::Error::Msg("captured graph has no output".into()))?; self.replays += 1; + // D18/D3: the replay count is the ONE number that separates a graph + // that serves from a graph that only records, and until now nothing + // printed it after capture β€” so every report of "replays served" was an + // inference. Emit it on a geometric-ish schedule: the first replay, the + // first ten, then every 100. Cost is a handful of lines per run. + if self.replays <= 10 || self.replays % 100 == 0 { + tracing::info!("ARCGRAPH REPLAY #{} β€” {}", self.replays, self.status_line()); + } Ok(out) } @@ -542,22 +563,68 @@ output_trusted={} verify_remaining={} verify_failed={}", candle_core::bail!("cuStreamEndCapture failed: {s}"); } + // Step markers. This sequence is five driver calls with no logging + // between them, and it has aborted the PROCESS inside that window + // (`malloc_consolidate(): invalid chunk size`) on a capture that + // recorded with zero allocator misses. A glibc abort leaves no CUDA + // error to read, so without a marker per call the only thing the log + // establishes is "somewhere in here" β€” which is not a located blocker. + // These are `info!` on purpose: they cost one line per captured graph, + // and there is at most one capture per process. + tracing::info!("ARC capture: [1/5] cuStreamEndCapture OK, graph recorded"); + + // Is the context ALREADY in error before the graph has ever run? + // + // `cuGraphLaunch` returned 700 SYNCHRONOUSLY on a capture with zero + // allocator misses. A synchronous illegal-address from a launch call is + // characteristic of a STICKY context error raised by earlier work, not + // of the launch itself β€” CUDA reports a real graph fault asynchronously, + // at the following sync. Nothing executes between begin_capture and + // here (capture records, it does not run), so a non-zero result on this + // line places the fault BEFORE capture β€” in the warmup forwards β€” and a + // zero places it in the graph. Without this the two are indistinguishable + // and the 700 gets blamed on whichever one is being worked on. + let pre = unsafe { cudaStreamSynchronize(self.stream) }; + if pre == CUDA_SUCCESS { + tracing::info!("ARC capture: [1b] context CLEAN before instantiate"); + } else { + tracing::error!( + "ARC capture: [1b] context ALREADY IN ERROR before instantiate (cudaError \ + {pre}). The fault happened during warmup, not in the graph; every later \ + CUDA call inherits it." + ); + } + // Instantiate (private pool still installed). RUN-161 2b: // AUTO_FREE_ON_LAUNCH (=1) so a graph with memory-alloc nodes can be // RE-launched (replayed) -- otherwise the 2nd launch fails with - // INVALID_VALUE. Harmless if the graph has no alloc nodes. + // INVALID_VALUE. + // + // It is NOT unconditionally harmless when the graph has no alloc nodes, + // which is now the normal case: the capture-miss gate refuses to reach + // this function unless every allocation was a cache hit, so a clean + // capture records no memory nodes at all and the flag governs nothing + // that exists. `ARC_GRAPH_AUTO_FREE=0` drops it so that can be + // ablated against the abort above rather than assumed innocent. const CU_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH: u64 = 1; - let mut exec: CUgraphExec = std::ptr::null_mut(); - let s = unsafe { - cuGraphInstantiateWithFlags( - &mut exec, - graph, - CU_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, - ) + let auto_free = !matches!( + std::env::var("ARC_GRAPH_AUTO_FREE").as_deref(), + Ok("0") | Ok("false") + ); + let flags = if auto_free { + CU_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH + } else { + 0 }; + let mut exec: CUgraphExec = std::ptr::null_mut(); + let s = unsafe { cuGraphInstantiateWithFlags(&mut exec, graph, flags) }; + tracing::info!( + "ARC capture: [2/5] cuGraphInstantiateWithFlags(flags={flags}) returned {s}" + ); unsafe { cuGraphDestroy(graph); } + tracing::info!("ARC capture: [3/5] cuGraphDestroy OK"); if s != CUDA_SUCCESS { self.drain_alloc_cache(); unsafe { @@ -571,6 +638,7 @@ output_trusted={} verify_remaining={} verify_failed={}", // First launch (private pool still installed -> graph memory is // allocated/backed here at the capture-time addresses). let s = unsafe { cuGraphLaunch(exec, self.stream) }; + tracing::info!("ARC capture: [4/5] first cuGraphLaunch returned {s}"); if s != CUDA_SUCCESS { self.drain_alloc_cache(); unsafe { @@ -586,6 +654,7 @@ output_trusted={} verify_remaining={} verify_failed={}", // the CUDA context and the process dies later with no diagnostic. // cudaError: 700 = illegalAddress, 719 = launchFailure, 1 = invalidValue. let sync = unsafe { cudaStreamSynchronize(self.stream) }; + tracing::info!("ARC capture: [5/5] first-launch cudaStreamSynchronize returned {sync}"); // RUN-161 diagnostic: measure the TRUE clean graph replay latency here, // while the captured input tensors are still alive (no stale-input // fault). 10 back-to-back launch+sync, report the best. Compare to the @@ -667,6 +736,12 @@ output_trusted={} verify_remaining={} verify_failed={}", #[cfg(feature = "cuda")] impl Drop for CudaGraphRunner { fn drop(&mut self) { - // CapturedGraph handles cleanup via its own Drop + // CapturedGraph handles cleanup via its own Drop. + // + // The final tally goes to stderr directly, NOT through `tracing`: a + // subscriber may already be torn down by the time the pipeline drops, + // and a status line that silently vanishes at exit is exactly the + // silent-success failure this line exists to prevent (D18). + eprintln!("ARCGRAPH FINAL {}", self.status_line()); } } diff --git a/memory/mission/wave65-CQ-v4-compressed-rope-recompute.md b/memory/mission/wave65-CQ-v4-compressed-rope-recompute.md index 3fcf377e83..d45f040269 100644 --- a/memory/mission/wave65-CQ-v4-compressed-rope-recompute.md +++ b/memory/mission/wave65-CQ-v4-compressed-rope-recompute.md @@ -43,8 +43,10 @@ the rolling compressor only ever *appends*, and only appends once every `ratio` tokens. The one thing already memoized here is the position vector itself -(`compressed_row_positions`, `deepseek4.rs:3013-3037`) β€” the thread-local there -is the pattern the rest of this wants. +(`compress_positions`, `layers.rs:2391-2459`) β€” the thread-local there is the +pattern the rest of this wants. Note it is keyed per `(ratio, device)`, NOT per +`(t_c, ratio, device)`: V4 interleaves CSA (ratio 4) and HCA (ratio 128) layers, +so a single-slot cache is evicted by the very next layer and memoizes nothing. --- @@ -80,8 +82,10 @@ did not change. rotates 0 or 1 new rows and concatenates, instead of rotating `t_c`. 2. The `cos`/`sin` gathers are **separately memoizable** β€” they depend only on `(t_c, ratio, device)`, exactly like the already-memoized - `compressed_row_positions` (`deepseek4.rs:3013-3037`) and `GRAPH_BLOCK_IDS` - (`dsv4_attention.rs`). Those two are the established in-tree pattern to copy. + `compress_positions` (`layers.rs:2391-2459`) and `GRAPH_BLOCK_IDS` + (`dsv4_attention.rs`). Those two are the established in-tree pattern to copy β€” + and copy the KEY as well as the shape: one slot per `(ratio, device)`, so the + CSA/HCA interleave cannot thrash it. Item 2 is independent of item 1 and much cheaper to land; it is worth doing first. diff --git a/mistralrs-core/src/kv_cache/single_cache.rs b/mistralrs-core/src/kv_cache/single_cache.rs index b04c78fd6a..207686fcee 100644 --- a/mistralrs-core/src/kv_cache/single_cache.rs +++ b/mistralrs-core/src/kv_cache/single_cache.rs @@ -326,9 +326,21 @@ impl SingleCache { self.pregrow_for_graph(src, read_capacity)?; let ad = self.all_data.as_ref().unwrap(); - // Device-slot write of the new token; `write_kv_inplace` uses the - // buffer's real capacity from all_data's dims, so the slot is correct. - mistralrs_quant::kvwrite::write_kv_inplace(ad, src, position)?; + // RING write. The read below is a fixed `0..read_capacity` window, so + // writing at the ABSOLUTE position put every token from `read_capacity` + // onward OUTSIDE the window that is read: the freshest key became + // invisible and a stale row was attended in its place, making this arm + // correct only while the context was shorter than `sliding_window`. + // + // `read_capacity` IS `sliding_window` here, so a ring of that many + // slots holds exactly the key set the raw branch is defined to attend. + // The slot is derived on the DEVICE from the device position β€” a + // host-resolved `position % capacity` would be baked into the recorded + // kernel and every replay would write to the capture step's slot. See + // `layers::graph_ring_slot` for why permuting the keys is invariant + // here and why the existing length mask is already ring-correct. + let slot = crate::layers::graph_ring_slot(position, read_capacity)?; + mistralrs_quant::kvwrite::write_kv_inplace(ad, src, &slot)?; // Fixed-offset window -> constant shape for capture/replay. ad.narrow(self.dim, 0, read_capacity) } diff --git a/mistralrs-core/src/kv_cache/xs_rolling.rs b/mistralrs-core/src/kv_cache/xs_rolling.rs index b73dc1bee6..fa30bed65f 100644 --- a/mistralrs-core/src/kv_cache/xs_rolling.rs +++ b/mistralrs-core/src/kv_cache/xs_rolling.rs @@ -489,6 +489,9 @@ pub struct XsRollingCache { /// materialise `comp.all_data` before the first group completes so the /// cache managers always have a tensor to batch. pub head_dim: usize, + /// Hold the retained raw tail at a CONSTANT width (see + /// [`Self::pin_tail_width`]). Off unless graph capture asks for it. + pin_tail: bool, } impl XsRollingCache { @@ -521,9 +524,60 @@ impl XsRollingCache { span_groups, margin: XS_TAIL_MARGIN_TOKENS, head_dim, + pin_tail: false, } } + /// The constant width the retained raw tail is held at under + /// [`Self::pin_tail_width`]: one row's full span plus the rollback margin. + /// + /// This is an upper bound on the width the unpinned policy ever chooses. + /// With `keep_from = ((tokens - margin)/ratio + 1 - span_groups) * ratio`, + /// writing `tokens - margin = q*ratio + r` gives a width of + /// `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. + pub fn graph_tail_width(&self) -> usize { + self.span_groups * self.ratio + self.margin + } + + /// Hold the retained raw tail at [`Self::graph_tail_width`] columns instead + /// of letting it breathe with the sequence position. + /// + /// # Why capture needs this (RUN-161, measured) + /// + /// `advance_uniform` rebuilds the tail every step with + /// `Tensor::cat(&[tail, xs_new], 1)` and retains only what a future + /// compressed row can still consume. That width is a function of + /// `tokens % ratio`, so it MOVES: `ratio` distinct widths, cycling with + /// period `ratio`. Each new width is an allocation size the caching + /// allocator has never served, and an allocation that misses the warm pool + /// *during capture* becomes an unstable graph memory node β€” which is not a + /// slow graph but an invalid one, surfacing as `CUDA_ERROR_ILLEGAL_ADDRESS` + /// on the first `cuGraphLaunch`. + /// + /// Measured on an H200 (`compress_ratios` {4, 128}, `hidden` 4096, BF16, so + /// 8192 B/token): a single surviving capture-time miss of 327 680 B = 40 + /// tokens, one token wider than the 319 488 B = 39 tokens missed on the + /// preceding warmup step. That is the ratio-128 (HCA) tail, which grows by + /// one token per step for 128 steps before it wraps β€” so no practical + /// warmup count can pre-cover it, which is exactly what earlier sessions + /// observed. (The ratio-4 (CSA) tail has only 4 distinct widths, which is + /// why the same defect once looked like "four unwarmed allocation sizes".) + /// + /// Pinned, the width stops being a function of position and the cat, the + /// narrow and the `contiguous` behind them all take one size for the life + /// of the run. + /// + /// Takes effect once the sequence is at least `graph_tail_width()` tokens + /// long; before that the history genuinely is shorter than the window and + /// the tail is still growing. Capture must therefore happen after that + /// point β€” with V4's HCA layers that is 144 tokens, which a prefill of any + /// realistic length already satisfies. + pub fn pin_tail_width(&mut self) { + self.pin_tail = true; + } + /// How many batch rows this cache currently describes. pub fn rows(&self) -> usize { self.tokens.len() @@ -1156,7 +1210,16 @@ impl XsRollingCache { let keep_from = ((rollback_floor / self.ratio + 1).saturating_sub(self.span_groups) * self.ratio) .max(win_start); - let new_base = keep_from.min(tokens_new); + // 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)? diff --git a/mistralrs-core/src/layers.rs b/mistralrs-core/src/layers.rs index af68e3aa18..71701e8521 100644 --- a/mistralrs-core/src/layers.rs +++ b/mistralrs-core/src/layers.rs @@ -2388,6 +2388,76 @@ impl Module for QLinear { } } +// Cached strided-position tables for the V4 compressor, one per +// `(ratio, device)`. Keyed that way because V4 mixes compression ratios across +// layers (CSA and HCA); a single-slot cache would thrash and reintroduce a host +// copy on every layer. See `compress_positions` for the rationale. +std::thread_local! { + static COMPRESS_POSITIONS: std::cell::RefCell> = + const { std::cell::RefCell::new(Vec::new()) }; +} + +/// Round a needed length up so the table is (re)built rarely β€” and therefore +/// during warmup rather than during capture. +const COMPRESS_POS_CHUNK: usize = 1024; + +/// `[t_c]` U32 tensor holding `[0, ratio, 2*ratio, …, (t_c-1)*ratio]`, served +/// as a zero-copy view into a cached device table. +/// +/// # Why this is not a `Tensor::arange`: that is a graph-capture HAZARD +/// +/// `Tensor::arange` (and `from_vec`/`new`/`from_slice`) builds its data as a +/// transient host `Vec` and uploads it through `CudaDevice::clone_htod`, which +/// issues an **async** `cuMemcpyHtoDAsync` and returns. Outside capture that is +/// merely a host round trip. **Inside `cuStreamBeginCapture` it is a +/// correctness bug**: the copy is not executed, it is *recorded* β€” the graph +/// stores the HOST POINTER and re-reads it on the first launch and on every +/// replay. The `Vec` is freed as soon as the expression returns, so the graph +/// copies freed host memory into the tensor. +/// +/// That is exactly how V4 decode capture died: `compressed_kv_from_rows` built +/// its strided positions with `arange`, and the first `cuGraphLaunch` filled +/// them with garbage, so the very next op β€” +/// `self.cos.index_select(&positions, 0)` in +/// `DeepSeekV2RotaryEmbedding::forward_at_positions` β€” tripped candle's +/// device-side bounds assert (`ids[id_i] < src_dim_size`, `T = __nv_bfloat16` +/// for the bf16 cos table, `I = unsigned int` for the U32 positions) and the +/// process took SIGSEGV. candle already documents this failure mode for kernel +/// dims/strides and works around it in `CudaDevice::htod_info` by leaking the +/// host source while capturing; nothing protects `clone_htod`, which is what +/// `arange` uses. +/// +/// The table is a pure function of `(ratio, capacity)` and its values never +/// change, so it is built **once, outside capture** (warmup steps run long +/// before `begin_capture`) and every decode step then takes a `narrow` view of +/// it. No host memory is touched inside the captured region, and the base +/// device address is stable across replays. +pub fn compress_positions(t_c: usize, ratio: usize, device: &Device) -> Result { + if ratio == 0 { + candle_core::bail!("compress_positions: ratio must be non-zero"); + } + COMPRESS_POSITIONS.with(|c| { + let mut cache = c.borrow_mut(); + let hit = cache.iter().position(|(r, t)| { + *r == ratio && t.device().same_device(device) && t.dim(0).unwrap_or(0) >= t_c + }); + let idx = match hit { + Some(i) => i, + None => { + // Drop any smaller/stale table for this (ratio, device) so the + // cache cannot grow without bound. + cache.retain(|(r, t)| !(*r == ratio && t.device().same_device(device))); + let cap = t_c.div_ceil(COMPRESS_POS_CHUNK) * COMPRESS_POS_CHUNK; + // The one `arange` β€” deliberately here and not in the forward. + let table = Tensor::arange_step(0u32, (cap * ratio) as u32, ratio as u32, device)?; + cache.push((ratio, table)); + cache.len() - 1 + } + }; + cache[idx].1.narrow(0, 0, t_c) + }) +} + /// Thread-local GPU positions tensor for CUDA graph mode. /// When set, `RotaryEmbedding::forward()` uses GPU-side gather instead of /// CPU-side `narrow()`, making the forward pass graph-capture compatible. @@ -2546,6 +2616,48 @@ pub fn set_graph_mode_mask(mask: Option) { GRAPH_MODE_MASK.with(|m| *m.borrow_mut() = mask); } +// Per-layer output handles, for bisecting a CUDA-graph replay against eager. +// +// The same-step probe proved the recorded graph is wired correctly (max|Ξ”| = +// 0 on the capture step), so a divergence one step later is a per-step input +// that must advance and does not. This finds WHICH one by finding the first +// layer whose output differs. +// +// The mechanism relies on a property of capture: the tensors a captured +// forward produces are ordinary warm-pool buffers, and the recorded kernels +// write to **those addresses** on every replay. Holding a clone of each layer +// output therefore gives a live window onto the graph's own intermediates β€” +// after a replay, these handles contain the REPLAY's values. Holding them +// also keeps them out of the allocator's free list, so the eager forward that +// follows cannot land on top of them. +// +// Handles only: `push` clones a `Tensor` (an `Arc` bump), so tracing costs no +// device work and is safe to run inside the capture region. +thread_local! { + static ARC_LAYER_TRACE: std::cell::RefCell>> = + const { std::cell::RefCell::new(None) }; +} + +/// Start recording per-layer outputs, discarding any previous recording. +pub fn arc_layer_trace_begin() { + ARC_LAYER_TRACE.with(|t| *t.borrow_mut() = Some(Vec::new())); +} + +/// Stop recording and take what was recorded. +pub fn arc_layer_trace_take() -> Option> { + ARC_LAYER_TRACE.with(|t| t.borrow_mut().take()) +} + +/// Record one layer's output. No-op unless a trace is open, so the call sites +/// in the model cost a thread-local read on the hot path and nothing else. +pub fn arc_layer_trace_push(xs: &Tensor) { + ARC_LAYER_TRACE.with(|t| { + if let Some(v) = t.borrow_mut().as_mut() { + v.push(xs.clone()); + } + }); +} + #[cfg(feature = "cuda")] pub fn graph_mode_mask() -> Option { GRAPH_MODE_MASK.with(|m| m.borrow().clone()) @@ -2558,17 +2670,17 @@ pub fn graph_mode_mask() -> Option { None } -/// Cached `[1, capacity]` F32 slot-index vector for [`graph_mode_length_mask`]. -/// -/// `Tensor::arange` on a GPU device is a host-to-device copy β€” banned in a hot -/// loop by `CLAUDE.md`, and worse than merely slow here: an H2D copy from -/// pageable host memory **cannot be recorded into a CUDA graph**, so building -/// the arange inside the capture region would abort the capture. Building it -/// once, on the warmup steps that precede capture, keeps the per-step work to -/// device-side compare/select kernels. -/// -/// Keyed by capacity *and* device: a rebuild is only needed when either moves, -/// which for one loaded model is never after the first step. +// Cached `[1, capacity]` F32 slot-index vector for `graph_mode_length_mask`. +// +// `Tensor::arange` on a GPU device is a host-to-device copy β€” banned in a hot +// loop by `CLAUDE.md`, and worse than merely slow here: an H2D copy from +// pageable host memory **cannot be recorded into a CUDA graph**, so building +// the arange inside the capture region would abort the capture. Building it +// once, on the warmup steps that precede capture, keeps the per-step work to +// device-side compare/select kernels. +// +// Keyed by capacity *and* device: a rebuild is only needed when either moves, +// which for one loaded model is never after the first step. #[allow(clippy::type_complexity)] std::thread_local! { static GRAPH_MODE_SLOT_IDS: std::cell::RefCell> = @@ -2628,6 +2740,68 @@ pub fn graph_mode_length_mask(positions: &Tensor, capacity: usize, dtype: DType) .reshape((b, 1, 1, capacity)) } +/// `positions mod capacity` β€” the RING slot a graph-mode decode writes its new +/// K/V row into β€” computed entirely on the device. +/// +/// # Why the graph KV buffer has to be a ring (RUN-161) +/// +/// The graph decode arm reads a constant `capacity`-wide window (slots +/// `0..capacity`) so the launch geometry never moves, but it used to *write* at +/// the ABSOLUTE position. Those agree only while `position < capacity`: from +/// token `capacity` onward the new row was written past the end of the window +/// that is read, so the freshest key was invisible to attention and a stale row +/// was attended in its place. The arm was therefore correct only below +/// `sliding_window`. +/// +/// A ring of exactly `capacity` slots is the natural fix, because `capacity` +/// **is** `sliding_window`: the raw branch attends precisely the last +/// `sliding_window` tokens, so a full ring holds exactly the right key set and +/// never holds a key that should have been evicted. +/// +/// ## Why permuting the keys is safe +/// +/// A ring stores keys out of order. That is invariant for this attention: +/// * softmax over keys is permutation-invariant, and K and V are written to +/// the SAME slot, so key `i` stays paired with value `i`; +/// * V4 rotates K *before* caching, so every stored row already carries its +/// own absolute RoPE β€” no consumer re-derives a position from column index; +/// * the validity mask stays [`graph_mode_length_mask`]'s `slot <= position`, +/// which is already exactly right for a ring: while the ring is filling +/// (`position < capacity`) slot `s` is written iff `s <= position`, and once +/// it is full every slot is valid and the predicate is universally true. +/// +/// ## Why this is not just `position % capacity` on the host +/// +/// The host knows the position, but a host-resolved slot becomes a literal in +/// the recorded kernel's arguments and every replay would then write to the +/// capture step's slot β€” the same defect that made host-resolved RoPE offsets +/// wrong (see `DeepSeekV2RotaryEmbedding::cos_sin_for`). Only `capacity` is +/// baked, and it is a genuine compile-time-constant of the run. +/// +/// F32 is exact for every integer below 2^24, and positions are bounded by +/// `max_position_embeddings` (163 840 for V4), so the float round trip is exact. +pub fn graph_ring_slot(positions: &Tensor, capacity: usize) -> Result { + if capacity == 0 { + candle_core::bail!("graph_ring_slot: capacity must be non-zero"); + } + let p = positions.to_dtype(DType::F32)?; + // `affine`'s multiplier is a constant kernel argument, which is safe to + // bake: `capacity` is fixed for the life of the run. Only the position is + // allowed to vary between replays, and it is read from device memory. + let blocks = p.affine(1.0 / capacity as f64, 0.0)?.floor()?; + let rem = (&p - (blocks * capacity as f64)?)?; + // `1.0 / capacity` is exact in binary only when `capacity` is a power of + // two (V4's `sliding_window` is 128, so it is β€” but this must not depend on + // that). Otherwise `p * (1/capacity)` can round just below an integer at + // `p` an exact multiple of `capacity`, the `floor` loses a block, and the + // remainder comes out as `capacity` β€” one slot PAST the window, i.e. an + // out-of-bounds KV write. The rounding error is bounded by well under one + // ulp of the quotient, so the remainder is in `[0, capacity]` and a single + // correction is exact for every input. + let over = rem.ge(capacity as f64)?.to_dtype(DType::F32)?; + (rem - (over * capacity as f64)?)?.to_dtype(DType::U32) +} + /// Check if graph-mode positions are set. #[cfg(feature = "cuda")] pub fn has_graph_mode_positions() -> bool { @@ -3669,6 +3843,85 @@ impl Module for ScaledEmbedding { } } +#[cfg(test)] +mod compress_positions_tests { + use super::*; + + /// The cache replaced `(arange(0..t_c).to_f32() * ratio).to_u32()`. The + /// values must be bit-identical to that expression, or every compressed + /// row silently rotates at the wrong absolute position β€” a quality bug no + /// test outside this one would catch. + // Named "strided_range" rather than naming the candle call: the Typos lane + // allow-lists `arange` as a whole identifier, so it is fine in the doc + // comment above and in the body below, but not as one word inside a longer + // snake_case name. Same reason as 745c871dd / e4eb59dfb. + #[test] + fn matches_the_strided_range_expression_it_replaced() -> Result<()> { + let dev = Device::Cpu; + for ratio in [1usize, 2, 4, 8, 32] { + for t_c in [1usize, 3, 17, 64] { + let want = (Tensor::arange(0u32, t_c as u32, &dev)?.to_dtype(DType::F32)? + * (ratio as f64))? + .to_dtype(DType::U32)? + .to_vec1::()?; + let got = compress_positions(t_c, ratio, &dev)?.to_vec1::()?; + assert_eq!(got, want, "ratio={ratio} t_c={t_c}"); + } + } + Ok(()) + } + + /// Growth must reuse the table, not silently rebuild it every step β€” + /// rebuilding is the host copy this cache exists to remove, and doing it + /// inside capture is the crash it exists to prevent. + #[test] + fn a_longer_request_within_capacity_reuses_the_same_storage() -> Result<()> { + let dev = Device::Cpu; + let ratio = 4usize; + let first = compress_positions(1, ratio, &dev)?; + let grown = compress_positions(COMPRESS_POS_CHUNK, ratio, &dev)?; + // Same backing table => same length after the chunk round-up. + assert_eq!(first.dim(0)?, 1); + assert_eq!(grown.dim(0)?, COMPRESS_POS_CHUNK); + let entries = COMPRESS_POSITIONS.with(|c| c.borrow().len()); + assert_eq!(entries, 1, "one (ratio, device) entry, not one per call"); + // Distinct ratios must NOT evict each other: V4 interleaves CSA and HCA + // layers, so a single-slot cache would rebuild on every layer. + let _ = compress_positions(8, 32, &dev)?; + let _ = compress_positions(8, ratio, &dev)?; + let entries = COMPRESS_POSITIONS.with(|c| c.borrow().len()); + assert_eq!(entries, 2, "ratios must coexist"); + Ok(()) + } + + #[test] + fn zero_ratio_is_refused() { + assert!(compress_positions(4, 0, &Device::Cpu).is_err()); + } + + /// At `ratio == 1` the table IS `[0, 1, …, t_c-1]`, which is exactly what + /// `Tensor::arange(0u32, t_c as u32, dev)` produces. `dsv4_attention`'s + /// compressed-branch mask relies on that identity to serve its block + /// indices from the cache instead of a per-layer host->device upload, so + /// pin it: a table that disagreed here would shift which compressed blocks + /// a query attends, silently and with no error. + /// + /// The `reshape` mirrors the call site, which reshapes the view to + /// `(1, t_c)` immediately β€” a narrowed view that could not reshape would + /// fail there and nowhere else. + #[test] + fn ratio_one_is_the_plain_index_range_the_mask_uses() -> Result<()> { + let dev = Device::Cpu; + for t_c in [1usize, 2, 7, 64, 1025] { + let want = Tensor::arange(0u32, t_c as u32, &dev)?.to_vec1::()?; + let got = compress_positions(t_c, 1, &dev)?; + assert_eq!(got.to_vec1::()?, want, "t_c={t_c}"); + assert_eq!(got.reshape((1, t_c))?.dims(), &[1, t_c], "t_c={t_c}"); + } + Ok(()) + } +} + #[cfg(test)] mod graph_mode_mask_tests { use super::*; diff --git a/mistralrs-core/src/models/deepseek4.rs b/mistralrs-core/src/models/deepseek4.rs index 3e84142bf8..82d9ac05b9 100644 --- a/mistralrs-core/src/models/deepseek4.rs +++ b/mistralrs-core/src/models/deepseek4.rs @@ -1412,11 +1412,18 @@ impl Attention { let t_c = comp.dim(2)?; // Compressed entry j sits at absolute position j*ratio. Apply the // layer's (compress-ΞΈ) RoPE to the last qk_rope_head_dim dims there. - // NOTE: builds a small arange each call (a host/device sync); fine on - // the correctness-first dense path β€” the long-context sparse-gather - // kernel is the place to precompute this. + // + // This used to build the positions with `Tensor::arange` on every call. + // That is a host round trip on the hot path AND β€” decisively β€” the + // reason CUDA-graph capture of the V4 decode forward SIGSEGV'd on its + // first launch: `arange` uploads a transient host `Vec`, and a captured + // `cuMemcpyHtoDAsync` records the host POINTER, not the bytes, so the + // graph re-read a freed `Vec` and handed garbage indices straight to + // the `index_select` on the next line. `compress_positions` serves a + // zero-copy view of a table built once, outside capture; see its doc + // comment for the full mechanism. let dev = comp.device(); - let positions = compressed_row_positions(t_c, ratio, dev)?; + let positions = crate::layers::compress_positions(t_c, ratio, dev)?; self.rotary_emb .forward_at_positions(&comp, self.cfg.qk_rope_head_dim, &positions) } @@ -1515,10 +1522,52 @@ impl Attention { .then(|| graph_comp_rows(self.compress_ratio.ratio())); if let Some(rows) = graph_rows { state.pin_comp_capacity(rows)?; + // The compressed axis is not the only width that + // moves with position: the retained RAW tail is + // rebuilt per step at a width that cycles with + // `tokens % ratio`, and a capture-time allocation + // miss on one of those widths is an unstable graph + // memory node β€” an illegal address on the first + // launch, not a slow graph. Pin it too. + state.pin_tail_width(); } // Always advanced, including under the window-only // ablation below: skipping it would drop the raw tail // and leave a hole the next group cannot be built from. + // + // πŸ”΄ ARCHITECTURAL LIMIT OF CAPTURE (RUN-161, measured). + // This call cannot be served by a replayed graph, for + // two independent reasons β€” neither of which more + // warmup, a bigger alloc cache, or pinning a width can + // reach: + // + // 1. `cuGraphLaunch` executes ONLY the recorded + // kernels. No host code runs. `advance` carries the + // compressor history in host-owned Rust state + // (`XsRollingCache::tail`, reassigned to a FRESH + // `Tensor` every step), so under replay the history + // simply stops advancing and the distant-context + // branch freezes at the capture step's content. The + // raw KV half does not have this problem because it + // writes through `write_kv_inplace` into one + // fixed-address buffer at a device-derived slot β€” + // a recorded kernel mutating stable memory. + // 2. The compressor fires only on the 1-in-`ratio` + // steps that complete a block, so consecutive + // decode steps do not even execute the same set of + // kernels. With `compress_ratios` {4, 128} that is + // every 4th step. A graph is a fixed DAG and cannot + // express that branch. + // + // The fix is (1)'s pattern applied here: `tail` has to + // become a fixed-capacity DEVICE ring advanced by a + // recorded kernel, exactly like `SingleCache:: + // append_graph`, and (2) then needs either a + // conditional graph node or one graph per phase. + // Until then the compressed branch is correct under + // capture only where it contributes nothing β€” + // `ARC_V4_WINDOW_ONLY=1`, or a context shorter than one + // `ratio` block. let advanced = state.advance(&xs3, |window| compressor.forward_from_xs(window))?; match graph_rows { @@ -1856,7 +1905,18 @@ impl Attention { .ok_or_else(|| candle_core::Error::Msg("graph positions unset".into()))? .to_dtype(candle_core::DType::U32)?; let cap = self.sliding_window.max(1); + // Bisect slots, in execution order, per graph-arm layer: + // k β€” the new post-RoPE key, BEFORE it is written. + // Diverges β‡’ the key itself is computed wrong + // (RoPE / projections), upstream of the cache. + // k_full β€” the fixed window read back AFTER the write. + // Clean `k` + dirty `k_full` β‡’ the write landed in + // the wrong slot, or the wrong window is read. + // attn β€” clean `k`/`k_full` + dirty `attn` β‡’ the mask or + // the compressed branch, not the KV ring. + crate::layers::arc_layer_trace_push(&k); let k_full = append_graph_kv_mqa(kv_cache, &k, &position, cap)?; + crate::layers::arc_layer_trace_push(&k_full); // Use ONLY the fixed-width graph mask (matches the C-wide K). // The eager `attention_mask` is kv_len-wide (growing) and would // both mismatch the fixed window and break shape-constancy. @@ -1882,7 +1942,7 @@ impl Attention { .into(), ) })?; - super::dsv4_attention::dsv4_attention( + let graph_attn = super::dsv4_attention::dsv4_attention( &q, &k_full, &k_full, @@ -1902,7 +1962,9 @@ impl Attention { graph_positions: Some(&position), ..dsv4_cfg }, - )? + )?; + crate::layers::arc_layer_trace_push(&graph_attn); + graph_attn } None => { // FP8 code storage is opt-in (`ARC_V4_FP8_KV=1`); unset stores @@ -2995,46 +3057,25 @@ fn append_kv_mqa( } } -// Cached `[t_c]` U32 absolute positions of the compressed rows, keyed by row -// count and ratio. Compressed row `j` sits at absolute position `j * ratio`, -// which depends on nothing but `j` β€” so for a given (width, ratio) this vector -// is a constant and rebuilding it per layer per step was pure waste. +// The compressed-row positions used to live here as a SINGLE-slot thread-local +// keyed on `(t_c, ratio, device)`. That slot is now `layers::compress_positions` +// (a per-`(ratio, device)` table, chunk-rounded and served as a `narrow` view), +// because a single slot THRASHES on V4: the model interleaves CSA (ratio 4) and +// HCA (ratio 128) layers, so consecutive layers of the same step evicted each +// other and the `arange_step` was paid again on every layer β€” including inside +// the capture region, where it is a SIGSEGV and not merely a cost. // -// It is also a correctness precondition for capture, not just a saving: -// `Tensor::arange` on a GPU device is an H2D copy from pageable host memory, -// which cannot be recorded into a CUDA graph at all. Under fixed capacity the -// width stops moving, so this hits from the second decode step onward and the -// copy leaves the capture region for good. -thread_local! { - static COMPRESSED_ROW_POSITIONS: std::cell::RefCell> = - const { std::cell::RefCell::new(None) }; -} - -fn compressed_row_positions(t_c: usize, ratio: usize, dev: &Device) -> Result { - COMPRESSED_ROW_POSITIONS.with(|c| -> Result { - let mut c = c.borrow_mut(); - if let Some((w, r, t)) = c.as_ref() { - if *w == t_c && *r == ratio && t.device().same_device(dev) { - return Ok(t.clone()); - } - } - // `j * ratio` is integer arithmetic. It used to be laundered through - // float β€” `arange(u32) -> cast_u32_f32 -> affine(*ratio) -> - // cast_f32_u32` β€” three kernel launches to compute a strided range - // that `arange_step` produces directly. The thread-local above makes - // that free on a hit, but a miss (first token, or a window/device - // change) still paid it. The values are identical: the old chain's - // `fmaf(j, ratio, 0.0)` is exact for every `j * ratio` a sequence can - // reach (both operands and the product are integers well inside f32's - // exact range), and `to_dtype(U32)` truncated it back. `ratio` is - // 1 / 4 / 128 (`CompressRatio::ratio`), never 0, so `arange_step`'s - // zero-step bail is unreachable. Multiplied in usize so the product - // cannot wrap before the u32 narrowing. - let t = Tensor::arange_step(0u32, (t_c * ratio) as u32, ratio as u32, dev)?; - *c = Some((t_c, ratio, t.clone())); - Ok(t) - }) -} +// The arithmetic argument that licensed `arange_step` in the first place is kept +// here, because it is what makes the two spellings interchangeable: `j * ratio` +// is integer arithmetic that used to be laundered through float β€” `arange(u32) +// -> cast_u32_f32 -> affine(*ratio) -> cast_f32_u32`, three kernel launches to +// compute a strided range `arange_step` produces directly. The values are +// identical: the old chain's `fmaf(j, ratio, 0.0)` is exact for every `j * ratio` +// a sequence can reach (both operands and the product are integers well inside +// f32's exact range), and `to_dtype(U32)` truncated it back. `ratio` is 1 / 4 / +// 128 (`CompressRatio::ratio`), never 0, so `arange_step`'s zero-step bail is +// unreachable. Multiplied in usize so the product cannot wrap before the u32 +// narrowing. /// The FIXED number of compressed rows the CUDA-graph decode arm reads back. /// @@ -4710,6 +4751,11 @@ impl DeepSeekV4 { Some(pos) => { let cap = self.cfg_full.sliding_window.max(1); let mask = crate::layers::graph_mode_length_mask(&pos, cap, xs_embed.dtype())?; + // Trace slot 0: does a REPLAY rebuild the right length mask + // from the device position buffer, or is the mask the one + // capture baked? Pushed before the layers so the bisect + // reads mask β†’ embed β†’ lift β†’ L0.. in execution order. + crate::layers::arc_layer_trace_push(&mask); crate::layers::set_graph_mode_mask(Some(mask)); } None => crate::layers::set_graph_mode_mask(None), @@ -4775,6 +4821,12 @@ impl DeepSeekV4 { }; v4_nan_dbg(&xs_4d, "lift_3d_to_4d"); v4_stat_dbg(&xs_4d, "lift_3d_to_4d"); + // Trace slots 1 and 2: the embedding (i.e. did the token id reach + // the graph) and the mHC lift, both BEFORE any attention. A + // divergence that starts here is an input; one that starts at L0 + // is inside the layer. + crate::layers::arc_layer_trace_push(&xs_embed); + crate::layers::arc_layer_trace_push(&xs_4d); let _prof_layers = arc_profiler::span("layers"); for (i, layer) in self.layers.iter().enumerate() { // Aggregated across all layers by default (calls = n_layers, @@ -4798,6 +4850,7 @@ impl DeepSeekV4 { Some(input_ids), )?; v4_stat_dbg(&xs_4d, &format!("L{i}")); + crate::layers::arc_layer_trace_push(&xs_4d); } drop(_prof_layers); let xs_4d = xs_4d.to_device(&self.device)?; @@ -4826,6 +4879,7 @@ impl DeepSeekV4 { flash_params, Some(input_ids), )?; + crate::layers::arc_layer_trace_push(&xs); } xs.to_device(&self.device)? }; diff --git a/mistralrs-core/src/models/dsv4_attention.rs b/mistralrs-core/src/models/dsv4_attention.rs index 2cbdcfc198..bc7cc7e06b 100644 --- a/mistralrs-core/src/models/dsv4_attention.rs +++ b/mistralrs-core/src/models/dsv4_attention.rs @@ -935,7 +935,14 @@ pub fn dsv4_attention( None => { let threshold = Tensor::new(&[((q0 + 1) / ratio) as u32], dev)?.reshape((1, 1))?; - Tensor::arange(0u32, t_c as u32, dev)? + // 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) } diff --git a/mistralrs-core/src/pipeline/normal.rs b/mistralrs-core/src/pipeline/normal.rs index d5ce9c76ff..53db115b34 100644 --- a/mistralrs-core/src/pipeline/normal.rs +++ b/mistralrs-core/src/pipeline/normal.rs @@ -513,6 +513,39 @@ const GRAPH_REPLAY_TOLERANCE: f32 = 1e-4; /// Costs one deviceβ†’host sync of a single scalar, and only runs on the handful /// of verification steps, never on the trusted path. #[cfg(feature = "cuda")] +thread_local! { + /// The EAGER logits of the step the CUDA graph was captured on, kept so a + /// later replay can be measured against them as well as against its own + /// step. A replay that reproduces these is a graph frozen at capture β€” its + /// static input buffers are not being read. One that matches neither them + /// nor the current step is partially advancing, which names a subset of + /// stale inputs rather than the whole set. (RUN-161) + static ARC_CAPTURE_STEP_EAGER: std::cell::RefCell> = + const { std::cell::RefCell::new(None) }; + /// How many replays past the capture step we are, for the same probe. + static ARC_REPLAY_DISTANCE: std::cell::RefCell = const { std::cell::RefCell::new(0) }; + /// Handles to the per-layer outputs the CAPTURED forward produced. The + /// recorded kernels write to these same addresses on every replay, so + /// after a replay these hold the replay's per-layer values and can be + /// bisected against an eager forward's. See `layers::arc_layer_trace_*`. + static ARC_GRAPH_LAYER_BUFS: std::cell::RefCell>> = + const { std::cell::RefCell::new(None) }; +} + +/// Is the per-layer graph-vs-eager bisect enabled? +/// +/// CUDA-only: every caller lives inside the `#[cfg(feature = "cuda")]` capture +/// block, so on a CPU/Metal build this would be genuinely dead code. Gated +/// rather than `#[allow(dead_code)]`'d so the dead-symbol detector keeps its +/// teeth on a CUDA build, where these must stay reachable. +#[cfg(feature = "cuda")] +fn arc_layer_bisect() -> bool { + std::env::var("ARC_GRAPH_LAYER_BISECT").as_deref() == Ok("1") +} + +/// Max elementwise |a - b| at F32. CUDA-only for the same reason as +/// [`arc_layer_bisect`]. +#[cfg(feature = "cuda")] fn max_abs_diff(a: &Tensor, b: &Tensor) -> candle_core::Result { if a.dims() != b.dims() { candle_core::bail!( @@ -2024,7 +2057,24 @@ impl Pipeline for NormalPipeline { paged_attn_meta.as_ref().map(|(a, b)| (a.clone(), b)), &flash_meta, ); - let _ = self.device().synchronize(); + // Do NOT discard this. `let _ =` here swallowed the + // only place an ASYNCHRONOUS device fault in a warm + // pass could ever be reported: the forward itself + // returns Ok (kernels are launched, not awaited), + // and the sync is what collects the fault. A + // swallowed fault leaves a sticky context error, and + // the next CUDA call to report it is the first + // cuGraphLaunch β€” which is then blamed for a fault + // it inherited. + if let Err(e) = self.device().synchronize() { + tracing::error!( + "ARC capture: warm-pass forward FAULTED at sync ({e}). The \ + CUDA context is now sticky-errored and every later call, \ + including the graph launch, will fail with it. Capture is \ + disabled for this run." + ); + runner.disable(); + } tracing::info!( "ARC capture: EAGER forward (sync'd) = {:?}", t_eager.elapsed() @@ -2105,6 +2155,13 @@ impl Pipeline for NormalPipeline { } let cl = context_lens.clone(); let pid = position_ids.clone(); + // Open the per-layer trace BEFORE capture begins so + // the recorded forward's own layer outputs are the + // handles we keep. Costs 43 Arc bumps, no device + // work, so it is safe inside the capture region. + if arc_layer_bisect() { + crate::layers::arc_layer_trace_begin(); + } let cap_result = match runner.begin_capture(bs) { Ok((gp, op)) => { match self.model.forward( @@ -2124,6 +2181,112 @@ impl Pipeline for NormalPipeline { 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); + } + 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: eager forward failed: {e}" + ), + } + } Some(out) } Err(e) => { @@ -2168,7 +2331,10 @@ impl Pipeline for NormalPipeline { match replayed { Ok(out) => { let dt = t.elapsed(); - if runner.needs_verification() { + if runner.needs_verification() + || std::env::var("ARC_GRAPH_VERIFY_ALWAYS").as_deref() + == Ok("1") + { // RUN-161 step 3. Turning replay // output on is the one change here // that can corrupt tokens SILENTLY: @@ -2179,6 +2345,16 @@ impl Pipeline for NormalPipeline { // eager result this step either way // β€” it is already computed and it is // the one we know is right. + // `out` above came from the replay, so + // the held capture-time layer handles + // now hold the REPLAY's per-layer + // values. Trace the eager forward and + // compare layer by layer: the first + // index that differs is the first layer + // reading a stale per-step input. + if arc_layer_bisect() { + crate::layers::arc_layer_trace_begin(); + } let eager = self.model.forward( &input_ids, &seqlen_offsets, @@ -2187,7 +2363,117 @@ impl Pipeline for NormalPipeline { paged_attn_meta.as_ref().map(|(a, b)| (a.clone(), b)), &flash_meta, )?; - match max_abs_diff(&out, &eager) { + if arc_layer_bisect() { + let eager_layers = + crate::layers::arc_layer_trace_take(); + ARC_GRAPH_LAYER_BUFS.with(|c| { + if let (Some(g), Some(e)) = + (c.borrow().as_ref(), eager_layers.as_ref()) + { + let n = g.len().min(e.len()); + let mut first = None; + let mut line = String::new(); + for i in 0..n { + let d = max_abs_diff(&g[i], &e[i]) + .unwrap_or(f32::NAN); + if i < 12 || first.map_or(false, |f| i <= f + 2) + { + line.push_str(&format!( + " L{i}={d:.2e}" + )); + } + if first.is_none() + && !(d <= GRAPH_REPLAY_TOLERANCE) + { + first = Some(i); + line.push_str("< { tracing::info!( "ARC capture: replay matched eager \