diff --git a/Cargo.toml b/Cargo.toml index b8dde1f07..b8852835d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -49,11 +49,11 @@ rust-version = "1.88" # 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 = "88d86a2019ee552923b217284e89847d2785dcdf" } -candle-nn = { git = "https://github.com/aeonmindai/candle.git", rev = "88d86a2019ee552923b217284e89847d2785dcdf" } -candle-flash-attn-v3 = { git = "https://github.com/aeonmindai/candle.git", rev = "88d86a2019ee552923b217284e89847d2785dcdf" } -candle-flash-attn = { git = "https://github.com/aeonmindai/candle.git", rev = "88d86a2019ee552923b217284e89847d2785dcdf" } -candle-metal-kernels = { git = "https://github.com/aeonmindai/candle.git", rev = "88d86a2019ee552923b217284e89847d2785dcdf" } +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" } axum = "0.8.8" anyhow = "1.0.100" diff --git a/mistralrs-core/src/pipeline/normal.rs b/mistralrs-core/src/pipeline/normal.rs index d412a7cd9..0bcc24e8a 100644 --- a/mistralrs-core/src/pipeline/normal.rs +++ b/mistralrs-core/src/pipeline/normal.rs @@ -98,6 +98,175 @@ pub struct NormalPipeline { autonomous_runner: Option, } +/// What a forward pass should do with candle's caching allocator before it +/// runs. See [`alloc_cache_action`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AllocCacheAction { + /// Turn the cache on. Frees start being recycled instead of returned to + /// the driver. + Enable, + /// Return every held buffer to the driver and turn the cache off. + DrainAndDisable, + /// Already in the right state; touch nothing. + Leave, +} + +/// Decide whether candle's caching allocator should be on for this forward. +/// +/// # Why this is not simply "on" +/// +/// The cache (fork `88d86a2`, `candle-core/src/cuda_backend/device.rs:42`) is +/// keyed on **exact byte count** — `free: HashMap>`, +/// looked up with `free.get_mut(&bytes)`. No bucketing, no smallest-fit, no +/// splitting. It also has **no capacity bound and no eviction**: the only ways +/// memory goes back to the driver are `set_alloc_cache_enabled(false)` and +/// `drain_alloc_cache_and_free()`. +/// +/// Those two facts together decide the policy: +/// +/// * **Decode wants it on.** A decode step allocates the same shapes 43 times +/// over (once per layer), and the shapes that do not depend on KV length — +/// hidden states, MLP and expert intermediates — repeat step after step. +/// Exact-size keying is a perfect fit for that traffic. +/// * **Prefill must drain it.** Prefill shapes scale with prompt length, are +/// large, and are hit once. Leaving the cache on across a prefill parks +/// those buffers for the process lifetime under a byte-size key nothing will +/// ever request again. That is the failure mode `ARC_NO_DEDICATED_DECODE` +/// already exists to work around, and it is why this returns +/// [`AllocCacheAction::DrainAndDisable`] rather than `Leave` on `seq_len != 1`. +/// +/// # The measurement this is waiting on +/// +/// Buffers whose size tracks KV length — the causal mask most obviously — +/// change size every decode step, so each step files one more never-reused +/// entry. That is bounded by `O(context^2)` bytes in the worst case and is +/// *not* bounded by this policy. It is small at short context (a `[1,1,1,kv]` +/// BF16 mask over 4k tokens sums to ~16 MB) and is not small at 128k. +/// +/// The fixed-capacity graph-mode path (`deepseek4.rs:4344-4353`, which swaps +/// the growing causal mask for `graph_mode_length_mask` at +/// `cfg_full.sliding_window`) removes that growth entirely, but it is reached +/// only under `ARC_V4_CAPTURE_PROBE`. Until the shape-invariance work lands, +/// `ARC_CANDLE_ALLOC_CACHE=0` is the kill switch, and the long-context +/// high-water mark is the number a GPU run should falsify this with. +/// +/// # Contract +/// +/// * `seq_len` — the forward's sequence length. `1` is decode. +/// * `enabled` — what candle reports *now* (`alloc_cache_enabled()`), so the +/// action is idempotent and we never drain a cache that is already off. +/// * `killed` — `ARC_CANDLE_ALLOC_CACHE=0` was set. +/// +/// 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 { + let want = !killed && seq_len == 1; + match (want, enabled) { + (true, false) => AllocCacheAction::Enable, + (false, true) => AllocCacheAction::DrainAndDisable, + _ => AllocCacheAction::Leave, + } +} + +/// `ARC_CANDLE_ALLOC_CACHE=0` turns the caching allocator off entirely. +/// +/// Any other value — and *unset* — leaves the default policy in force. The +/// variable used to be the on-switch, so the value `1` the ops scripts pass +/// (`arc-tools/arcgraph_heap_probe.sh:191`) still means what it always meant. +#[cfg_attr(not(feature = "cuda"), allow(dead_code))] +pub(crate) fn alloc_cache_killed() -> bool { + std::env::var("ARC_CANDLE_ALLOC_CACHE").is_ok_and(|v| v == "0") +} + +/// Retention cap for the caching allocator, in bytes, or `None` to leave +/// candle's own default (1 GiB) in force. +/// +/// `ARC_ALLOC_CACHE_MAX_MB=0` means unbounded — the pre-bounding behaviour, kept +/// reachable so the leak it causes can be A/B'd rather than argued about. +/// Measured unbounded on V4: **+6.04 MiB per decoded token with no plateau**, +/// against +0.057 MiB/token with the cache off entirely. +#[cfg_attr(not(feature = "cuda"), allow(dead_code))] +pub(crate) fn alloc_cache_capacity_bytes() -> Option { + parse_alloc_cache_capacity(std::env::var("ARC_ALLOC_CACHE_MAX_MB").ok().as_deref()) +} + +/// Pure half of [`alloc_cache_capacity_bytes`], so the parse is testable. +/// +/// `None` (unset, or unparsable) leaves candle's own default in force rather +/// than silently picking a different one — a typo in an ops script must not +/// quietly hand back the unbounded allocator this change exists to remove. +pub(crate) fn parse_alloc_cache_capacity(raw: Option<&str>) -> Option { + let mb: usize = raw?.trim().parse().ok()?; + Some(if mb == 0 { + usize::MAX + } else { + mb.saturating_mul(1024 * 1024) + }) +} + +/// Emit the allocator's counters every `ARC_ALLOC_CACHE_STATS` decode steps. +/// +/// The counters are what this cache has to be judged on. A green log — the +/// server did not crash, the output looked fine — is not evidence that the +/// cache is working: an earlier arena in this codebase reported "accounting OK" +/// and bit-identical output across 52 steps while silently bypassing itself for +/// every buffer under 128 bytes (`KERNEL_RULES.md:977-984`). What distinguishes +/// a working bounded cache from that is arithmetic, and it is printed here: +/// +/// * `alloc/step` — real `cuMemAllocAsync` calls per decode step. Low means the +/// cache is absorbing the ~11 k allocations a step makes. +/// * `free/step` — real `cuMemFreeAsync` calls per decode step. **Non-zero is +/// the point.** The unbounded cache's was exactly zero, forever, which is why +/// it grew without bound. +/// * `held` — bytes retained right now, against the cap. +#[cfg(feature = "cuda")] +fn report_alloc_cache_step(cd: &candle_core::CudaDevice, seq_len: usize) { + use std::sync::atomic::{AtomicU64, Ordering}; + static EVERY: std::sync::OnceLock = std::sync::OnceLock::new(); + let every = *EVERY.get_or_init(|| { + std::env::var("ARC_ALLOC_CACHE_STATS") + .ok() + .and_then(|v| v.trim().parse::().ok()) + .unwrap_or(0) + }); + if every == 0 || seq_len != 1 { + return; + } + static STEP: AtomicU64 = AtomicU64::new(0); + static LAST_ALLOC: AtomicU64 = AtomicU64::new(0); + static LAST_FREE: AtomicU64 = AtomicU64::new(0); + static LAST_STEP: AtomicU64 = AtomicU64::new(0); + let step = STEP.fetch_add(1, Ordering::Relaxed) + 1; + if step % every != 0 { + return; + } + let s = cd.alloc_cache_stats(); + let d_step = step - LAST_STEP.swap(step, Ordering::Relaxed); + let d_alloc = s.misses - LAST_ALLOC.swap(s.misses, Ordering::Relaxed); + let d_free = s.frees() - LAST_FREE.swap(s.frees(), Ordering::Relaxed); + let n = d_step.max(1) as f64; + tracing::info!( + "[alloc-cache] step {step} alloc/step {:.1} free/step {:.1} \ + hit-rate {:.4} held {:.1} MiB / cap {} high-water {:.1} MiB sizes {}", + d_alloc as f64 / n, + d_free as f64 / n, + s.hits as f64 / (s.hits + s.misses).max(1) as f64, + s.cached_bytes as f64 / (1024.0 * 1024.0), + if s.capacity_bytes == usize::MAX { + "unbounded".to_string() + } else { + format!("{} MiB", s.capacity_bytes / (1024 * 1024)) + }, + s.high_water_bytes as f64 / (1024.0 * 1024.0), + s.size_classes, + ); +} + /// A loader for a "normal" (non-quantized) model. pub struct NormalLoader { inner: Box, @@ -1702,15 +1871,37 @@ impl Pipeline for NormalPipeline { "normal.rs:1554", ); } - // Enable the candle caching allocator for graph-capture - // safety (RUN-161). Idempotent; gated so it's off during - // model load and only active for decode. Warmup decode - // forwards populate the cache before capture. - if probe && seq_len == 1 && std::env::var_os("ARC_CANDLE_ALLOC_CACHE").is_some() - { - if let candle_core::Device::Cuda(cd) = self.device() { - cd.set_alloc_cache_enabled(true); + // Candle's caching allocator, driven by `alloc_cache_action`. + // + // This used to be `probe && seq_len == 1 && env(...)`, which + // made a general-purpose allocator reachable only when the + // V4 capture probe was also on. The cache recycles frees for + // ANY decode step — it is not capture machinery — and the + // three stacked default-off gates meant the ~11k allocations + // per token were a disabled feature rather than a missing + // one. It is now on for decode by default and drained on + // prefill; see `alloc_cache_action` for why prefill must + // drain rather than coast. + if let candle_core::Device::Cuda(cd) = self.device() { + match alloc_cache_action( + seq_len, + cd.alloc_cache_enabled(), + alloc_cache_killed(), + ) { + AllocCacheAction::Enable => { + cd.set_alloc_cache_enabled(true); + if let Some(cap) = alloc_cache_capacity_bytes() { + cd.set_alloc_cache_capacity(cap); + } + } + AllocCacheAction::DrainAndDisable => { + cd.set_alloc_cache_enabled(false) + } + AllocCacheAction::Leave => {} } + report_alloc_cache_step(&cd, seq_len); + } + if probe && seq_len == 1 { // RUN-161 step 2b. Set the graph-mode device position: // drives RoPE + the fixed-capacity KV write slot, and // makes warmup forwards take the shape-constant path so @@ -2528,3 +2719,132 @@ impl AnyMoePipelineMixin for NormalPipeline { self.model.amoe_supported() } } + +/// The caching-allocator policy's contract. +/// +/// The allocator itself has **no tests at all** — not in this repo, and not in +/// the candle fork that implements it (`grep -rn "alloc_cache" --include="*.rs"` +/// over `candle-core` finds only `cuda_backend/device.rs` and +/// `cuda_backend/mod.rs`, with no `#[cfg(test)]` in either). Everything that +/// observes it needs a GPU, so nothing observes it in CI. +/// +/// The decision of *when* it is on does not need a GPU, so it is tested here. +#[cfg(test)] +mod alloc_cache_policy_tests { + use super::{alloc_cache_action, AllocCacheAction}; + + /// The change this policy exists to make: a decode step turns the cache on + /// without any capture probe being involved. + #[test] + fn decode_enables_the_cache() { + assert_eq!( + alloc_cache_action(1, false, false), + AllocCacheAction::Enable + ); + } + + /// Prefill must hand its buffers back. They are large, they are keyed by an + /// exact byte size no later request will ask for, and nothing evicts them. + #[test] + fn prefill_drains_the_cache() { + assert_eq!( + alloc_cache_action(512, true, false), + AllocCacheAction::DrainAndDisable + ); + } + + /// Called once per forward, so it must be a no-op in the steady state + /// rather than re-enabling (or re-draining) every step. + #[test] + fn steady_state_touches_nothing() { + assert_eq!(alloc_cache_action(1, true, false), AllocCacheAction::Leave); + assert_eq!( + alloc_cache_action(512, false, false), + AllocCacheAction::Leave + ); + } + + /// The kill switch has to work from either state, including turning off a + /// cache that a previous step already enabled — otherwise setting it + /// mid-run would leave the buffers stranded. + #[test] + fn kill_switch_disables_and_drains_from_either_state() { + assert_eq!( + alloc_cache_action(1, true, true), + AllocCacheAction::DrainAndDisable + ); + assert_eq!(alloc_cache_action(1, false, true), AllocCacheAction::Leave); + } + + /// A prompt that happens to be one token long is a prefill by every other + /// measure, but it allocates decode-shaped buffers, so the policy keys on + /// the shape it will actually see. Pinned so the equivalence is deliberate. + #[test] + fn one_token_prompt_is_treated_as_decode() { + assert_eq!( + alloc_cache_action(1, false, false), + AllocCacheAction::Enable + ); + } + + /// A zero-length forward is not decode. Guards against `seq_len == 0` + /// (which `dims2().unwrap_or((0, 0))` produces on a shape error) quietly + /// 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, true, false), + AllocCacheAction::DrainAndDisable + ); + } +} + +/// The retention cap's parse. +/// +/// The cap is the whole of the fix — candle's caching allocator had none, and +/// retained 6.04 MiB per decoded token forever — so the way it is configured has +/// to be unambiguous. In particular a typo must not fall back to "unbounded". +#[cfg(test)] +mod alloc_cache_capacity_tests { + use super::parse_alloc_cache_capacity; + + #[test] + fn unset_leaves_candles_default_in_force() { + assert_eq!(parse_alloc_cache_capacity(None), None); + } + + /// A typo is not a licence to run unbounded. `None` means "don't touch it", + /// and candle's default is bounded, so a bad value degrades to bounded. + #[test] + fn an_unparsable_value_does_not_become_unbounded() { + for bad in ["", " ", "lots", "1GiB", "-1", "1.5"] { + assert_eq!(parse_alloc_cache_capacity(Some(bad)), None, "{bad:?}"); + } + } + + /// `0` is the documented escape hatch back to the old unbounded allocator, + /// kept reachable only so its leak can be re-measured rather than argued + /// about. + #[test] + fn zero_is_the_explicit_unbounded_opt_in() { + assert_eq!(parse_alloc_cache_capacity(Some("0")), Some(usize::MAX)); + } + + #[test] + fn megabytes_convert_and_do_not_overflow() { + assert_eq!(parse_alloc_cache_capacity(Some("1")), Some(1024 * 1024)); + assert_eq!( + parse_alloc_cache_capacity(Some(" 512 ")), + Some(512 * 1024 * 1024) + ); + assert_eq!( + parse_alloc_cache_capacity(Some(&usize::MAX.to_string())), + Some(usize::MAX), + "saturates instead of wrapping to a tiny cap" + ); + } +}