ArcGraph: three replay-correctness blockers fixed; capture now blocked on host heap corruption, not on V4

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"

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::<u32>().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::<u32>().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 {