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 + } } } }