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] diff --git a/mistralrs-core/src/layers.rs b/mistralrs-core/src/layers.rs index 94fd20401..9f941b5e8 100644 --- a/mistralrs-core/src/layers.rs +++ b/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> = + 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) + }) +} + +// 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> = + 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. @@ -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::()?; + 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()); + } + + /// `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)] 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) } 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)?;