From 2743a440020c9939a52dad3de0a14c944fa23dd1 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Wed, 19 Aug 2026 21:30:19 +0100 Subject: [PATCH 1/4] fix(ArcGraph): Tensor::arange in the captured forward is a dangling HOST pointer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit V4 decode capture reached "forward RECORDED" and then SIGSEGV'd on its very first cuGraphLaunch, with a device-side assert from candle's index_select: ids[id_i] < src_dim_size (T = __nv_bfloat16, I = unsigned int) That dtype pair is the bf16 cos table indexed by U32 positions, i.e. `self.cos.index_select(&positions, 0)` in `DeepSeekV2RotaryEmbedding::forward_at_positions` — NOT the TD-MoE `tid2eid.index_select` (tid2eid is i64, so it cannot produce T=bf16), and not the embedding (its ids come from the already-address-stable graph-mode buffer). `positions` was built one line earlier by `compressed_kv_from_rows` with `Tensor::arange(0u32, t_c, dev)`. `arange` materialises a transient host `Vec` and uploads it via `CudaDevice::clone_htod` -> an ASYNC `cuMemcpyHtoDAsync`. Under `cuStreamBeginCapture` that 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 when the expression returns, so the graph copied freed host memory into `positions` and handed garbage indices to index_select. This is not an address-staleness bug — it fires on the first launch, before any replay — which is why the static input_ids buffer did not prevent it. candle already documents this exact mechanism and works around it in `CudaDevice::htod_info` (it leaks the host source while capturing); nothing protects `clone_htod`, which is the path `arange` takes. Fix: serve the strided compressed positions from a per-(ratio, device) table built once, outside capture, and take a zero-copy `narrow` view each step. The values are a pure function of (ratio, capacity) and never change. This also removes a host round trip from every compressor layer of every decode step. Keyed per ratio because V4 interleaves CSA and HCA compression ratios; a single-slot cache would thrash and reintroduce the copy on every layer. --- mistralrs-core/src/layers.rs | 128 +++++++++++++++++++++++++ mistralrs-core/src/models/deepseek4.rs | 17 ++-- 2 files changed, 139 insertions(+), 6 deletions(-) diff --git a/mistralrs-core/src/layers.rs b/mistralrs-core/src/layers.rs index 94fd20401..ab494f04d 100644 --- a/mistralrs-core/src/layers.rs +++ b/mistralrs-core/src/layers.rs @@ -2253,6 +2253,81 @@ 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. @@ -3519,6 +3594,59 @@ 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::()?; + 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()); + } +} + #[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 00907a056..b9009c29c 100644 --- a/mistralrs-core/src/models/deepseek4.rs +++ b/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) } From fb58fc63c7f9b48daae4f563374c49905c5a0acf Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Wed, 19 Aug 2026 21:37:10 +0100 Subject: [PATCH 2/4] fix(ArcGraph): the other three arange calls in the captured V4 forward MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A sweep of every host->device constructor reachable from the B=1 decode forward found three more `Tensor::arange` calls with the same defect as the compressor positions, all in `dsv4_attention`, all unconditional on every layer of every step: :610 kp = arange(raw_base, raw_base + t_k) absolute key positions :613 qp = arange(q0, q0 + t_q) query positions (ONE element at decode -- a whole clone_htod for 4 bytes, 43x per token) :628 bp = arange(0, t_c) compressed-block indices Each is a contiguous integer ramp cast to F32, so all three are views of a single cached device ramp: `layers::positions_f32(start, len, dev)` narrows `[0.0, 1.0, 2.0, …]`. Capacity grows by powers of two from 8192, so a typical run allocates the ramp once during warmup and never again -- rebuilding inside capture is precisely the hazard being removed. Fixing only the compressor site would not have been enough: any one of these would have re-armed the same dangling-host-pointer crash on the first launch. Also removes three H2D copies per layer per token from the decode hot path (129 per token at 43 layers), independent of graph capture. --- mistralrs-core/src/layers.rs | 77 +++++++++++++++++++++ mistralrs-core/src/models/dsv4_attention.rs | 17 +++-- 2 files changed, 85 insertions(+), 9 deletions(-) diff --git a/mistralrs-core/src/layers.rs b/mistralrs-core/src/layers.rs index ab494f04d..9f941b5e8 100644 --- a/mistralrs-core/src/layers.rs +++ b/mistralrs-core/src/layers.rs @@ -2328,6 +2328,51 @@ pub fn compress_positions(t_c: usize, ratio: usize, device: &Device) -> Result> = + 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 { + 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. @@ -3645,6 +3690,38 @@ mod compress_positions_tests { 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::()?; + let got = positions_f32(start, len, &dev)?.to_vec1::()?; + 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::()?, 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)] diff --git a/mistralrs-core/src/models/dsv4_attention.rs b/mistralrs-core/src/models/dsv4_attention.rs index f8ae30fb2..50189f2e7 100644 --- a/mistralrs-core/src/models/dsv4_attention.rs +++ b/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)?; From 489134b91c629765480bee9d94a163e036db4061 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Thu, 20 Aug 2026 16:22:54 +0100 Subject: [PATCH 3/4] fix(ci): allow 'arange' as a word inside longer identifiers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Same root cause as #181, and the identical fix: `.typos.toml` already exempted the bare Rust identifier `arange` (candle's/numpy's `Tensor::arange`) under [type.rust.extend-identifiers], but that only matches WHOLE identifiers — the word inside `matches_the_arange_expression_it_replaced` still tripped the gate. Add the word-level exemption rather than renaming, because `arange` is a real API name, not a misspelling of 'arrange'. This branch and #181 both carry those test functions, so both hit it. Whichever lands first makes the other's copy redundant. Co-Authored-By: Claude Opus 5 (1M context) --- .typos.toml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.typos.toml b/.typos.toml index 2932a4ac7..1842403a3 100644 --- a/.typos.toml +++ b/.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] From 240982c8cf32fe246a4706fd6a62b498d6daf891 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Thu, 20 Aug 2026 17:58:07 +0100 Subject: [PATCH 4/4] chore: force CI dispatch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two prior attempts (a push to the branch and a close/reopen of the PR) fired only the pull_request_target Analysis workflow; Continuous integration and CUDA compile check never dispatched for 489134b91, leaving the PR displaying only 'comment' — an absence that reads as a pass. Empty commit to force a synchronize on a fresh sha. Co-Authored-By: Claude Opus 5 (1M context)