fix(ArcGraph): Tensor::arange in the captured V4 forward is a dangling host pointer (4 call sites)

.typos.toml

@@ -69,6 +69,11 @@ extend-ignore-re = [
 "ue" = "ue"
 "ba" = "ba"
 "ratatui" = "ratatui"
+# candle's (and numpy's) `Tensor::arange`. Already exempt as a bare Rust
+# identifier below, but it also appears as a *word inside* longer identifiers
+# — e.g. test names describing the arange expression they replaced — and
+# extend-identifiers only matches whole identifiers.
+"arange" = "arange"
 "optin" = "optin"
 
 [type.md.extend-identifiers]

mistralrs-core/src/layers.rs

@@ -2253,6 +2253,126 @@ 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<Vec<(usize, Tensor)>> =
+        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<Tensor> {
+    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)
+    })
+}
+
+// Cached `[0.0, 1.0, 2.0, …]` F32 ramps, one per device. Serves every absolute
+// position vector the V4 attention masks need. See `positions_f32`.
+std::thread_local! {
+    static IOTA_F32: std::cell::RefCell<Vec<Tensor>> =
+        const { std::cell::RefCell::new(Vec::new()) };
+}
+
+/// Smallest ramp worth allocating: 32 KB of F32. Large enough that a typical
+/// run allocates once, at warmup, and never again — which is the point, since a
+/// rebuild *during* capture is the bug this exists to avoid.
+const IOTA_MIN_LEN: usize = 8192;
+
+/// `[len]` F32 tensor holding `[start, start+1, …, start+len-1]`, served as a
+/// view into a cached device ramp.
+///
+/// Replaces `Tensor::arange(start, start+len, dev)?.to_dtype(F32)` on the V4
+/// attention mask path. `arange` is a host→device copy, and inside CUDA-graph
+/// capture it is worse than slow: the recorded `cuMemcpyHtoDAsync` keeps the
+/// HOST pointer of a `Vec` that is freed the moment the expression returns, so
+/// every launch and replay fills the tensor from freed memory. See
+/// [`compress_positions`] for the full mechanism and the crash it produced.
+///
+/// The ramp grows by powers of two so rebuilds are rare and land during warmup.
+pub fn positions_f32(start: usize, len: usize, device: &Device) -> Result<Tensor> {
+    let end = start + len;
+    IOTA_F32.with(|c| {
+        let mut cache = c.borrow_mut();
+        let hit = cache
+            .iter()
+            .position(|t| t.device().same_device(device) && t.dim(0).unwrap_or(0) >= end);
+        let idx = match hit {
+            Some(i) => i,
+            None => {
+                cache.retain(|t| !t.device().same_device(device));
+                let cap = end.max(IOTA_MIN_LEN).next_power_of_two();
+                // The one `arange` — deliberately here and not in the forward.
+                let ramp = Tensor::arange(0u32, cap as u32, device)?.to_dtype(DType::F32)?;
+                cache.push(ramp);
+                cache.len() - 1
+            }
+        };
+        cache[idx].narrow(0, start, len)
+    })
+}
+
 /// 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.
@@ -3519,6 +3639,91 @@ 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.
+    #[test]
+    fn matches_the_arange_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::<u32>()?;
+                let got = compress_positions(t_c, ratio, &dev)?.to_vec1::<u32>()?;
+                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());
+    }
+
+    /// `positions_f32` replaced `arange(start, start+len).to_f32()` at three
+    /// V4 attention-mask sites. A ramp served at the wrong offset silently
+    /// shifts the sliding-window mask — wrong tokens attended, no error.
+    #[test]
+    fn positions_f32_matches_the_arange_expression_it_replaced() -> Result<()> {
+        let dev = Device::Cpu;
+        for (start, len) in [(0usize, 1usize), (0, 7), (5, 1), (13, 128), (1000, 64)] {
+            let want = Tensor::arange(start as u32, (start + len) as u32, &dev)?
+                .to_dtype(DType::F32)?
+                .to_vec1::<f32>()?;
+            let got = positions_f32(start, len, &dev)?.to_vec1::<f32>()?;
+            assert_eq!(got, want, "start={start} len={len}");
+        }
+        Ok(())
+    }
+
+    /// The ramp must survive growth past its initial capacity, and a view of it
+    /// must still reshape (the three call sites all reshape immediately).
+    #[test]
+    fn positions_f32_grows_and_its_views_reshape() -> Result<()> {
+        let dev = Device::Cpu;
+        let short = positions_f32(0, 4, &dev)?;
+        assert_eq!(short.reshape((1, 4))?.dims(), &[1, 4]);
+        let past_min = IOTA_MIN_LEN + 3;
+        let long = positions_f32(past_min, 2, &dev)?;
+        assert_eq!(long.to_vec1::<f32>()?, vec![past_min as f32, past_min as f32 + 1.0]);
+        assert_eq!(long.reshape((2, 1))?.dims(), &[2, 1]);
+        // One ramp per device, not one per request.
+        assert_eq!(IOTA_F32.with(|c| c.borrow().len()), 1);
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod graph_mode_mask_tests {
     use super::*;

mistralrs-core/src/models/deepseek4.rs

@@ -1319,13 +1319,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 = (Tensor::arange(0u32, t_c as u32, dev)?.to_dtype(DType::F32)?
-            * (ratio as f64))?
-        .to_dtype(DType::U32)?;
+        let positions = crate::layers::compress_positions(t_c, ratio, dev)?;
         self.rotary_emb
             .forward_at_positions(&comp, self.cfg.qk_rope_head_dim, &positions)
     }

mistralrs-core/src/models/dsv4_attention.rs

@@ -607,12 +607,13 @@ pub fn dsv4_attention(
     // valid, so no query row is fully masked (no softmax NaN). `kp` carries the
     // ABSOLUTE position of each retained key, so the comparison against `qp` is
     // unchanged by the narrowing above.
-    let kp = Tensor::arange(raw_base as u32, (raw_base + t_k) as u32, dev)?
-        .to_dtype(DType::F32)?
-        .reshape((1, t_k))?;
-    let qp = Tensor::arange(q0 as u32, (q0 + t_q) as u32, dev)?
-        .to_dtype(DType::F32)?
-        .reshape((t_q, 1))?;
+    //
+    // These are views into a cached device ramp, not `Tensor::arange`. `arange`
+    // is a host→device copy, and inside CUDA-graph capture the recorded memcpy
+    // holds the host pointer of a `Vec` that is already freed at launch time —
+    // see `layers::positions_f32`.
+    let kp = crate::layers::positions_f32(raw_base, t_k, dev)?.reshape((1, t_k))?;
+    let qp = crate::layers::positions_f32(q0, t_q, dev)?.reshape((t_q, 1))?;
     let causal = kp.broadcast_le(&qp)?;
     let lower = (&qp - window as f64)?;
     let in_window = kp.broadcast_gt(&lower)?;
@@ -625,9 +626,7 @@ pub fn dsv4_attention(
             // Compressed branch mask: query r attends compressed block b iff
             // b < (q0+r+1)/ratio (causal over fully-completed ratio-blocks).
             let ratio = cfg.compress_ratio.ratio();
-            let bp = Tensor::arange(0u32, t_c as u32, dev)?
-                .to_dtype(DType::F32)?
-                .reshape((1, t_c))?;
+            let bp = crate::layers::positions_f32(0, t_c, dev)?.reshape((1, t_c))?;
             let threshold = ((&qp + 1.0)? / ratio as f64)?.floor()?; // [t_q, 1]
             let comp_valid = bp.broadcast_lt(&threshold)?; // [t_q, t_c] (u8)
             let valid = Tensor::cat(&[&raw_valid, &comp_valid], 1)?;