diff --git a/mistralrs-core/src/kv_cache/mod.rs b/mistralrs-core/src/kv_cache/mod.rs index e6a2fd378..cf3d0136d 100644 --- a/mistralrs-core/src/kv_cache/mod.rs +++ b/mistralrs-core/src/kv_cache/mod.rs @@ -14,6 +14,10 @@ mod rotating_cache; mod single_cache; pub mod turboquant_cache; mod xs_rolling; +/// Thread-local `ARC_V4_XS_PER_SEQ` override, so a test outside this module can +/// exercise both sides of the flag (the production read is a `OnceLock`). +#[cfg(test)] +pub(crate) use xs_rolling::test_override as xs_per_seq_test_override; pub use full_cache::{EitherCache, LayerCaches}; pub use hybrid_cache::{ diff --git a/mistralrs-core/src/models/deepseek4.rs b/mistralrs-core/src/models/deepseek4.rs index e11dc1400..1b1fcc9f2 100644 --- a/mistralrs-core/src/models/deepseek4.rs +++ b/mistralrs-core/src/models/deepseek4.rs @@ -1542,6 +1542,7 @@ impl Attention { compress_ratio: self.compress_ratio, sliding_window: self.sliding_window, raw_prefix: 0, + row_q0: None, }; // Faithful V4: the compressed (distant-context) KV, with compress-θ // RoPE at the strided compressed positions. `dsv4_attention` runs a @@ -1700,6 +1701,14 @@ impl Attention { &self.sdpa_params, super::dsv4_attention::Dsv4AttentionConfig { raw_prefix, + // The dense path is the only one whose K/V rows are a + // left-aligned ragged cohort (`front_pad_kv_cache`). + // PagedAttention keeps per-request block tables, and + // the graph-decode arm reads a FIXED window whose + // columns are device slots, not absolute positions — + // neither can be masked from `seqlen_offsets`, so both + // keep `row_q0: None` and behave exactly as before. + row_q0: ragged_row_q0(seqlen_offsets), ..dsv4_cfg }, )? @@ -2465,6 +2474,40 @@ fn append_graph_kv_mqa( /// /// Resolved once: `deepseek4` used to scan the environment per timer call and /// paid ~390 `getenv`s per forward for it (wave33). +/// The per-row absolute query positions [`super::dsv4_attention`] needs when +/// the dense batch is a **left-aligned ragged** cohort, or `None` when it is +/// not — which is every request today, because the flag defaults off. +/// +/// `seqlen_offsets[i]` is already row `i`'s own past length: the engine threads +/// it to every attention call so RoPE places the row's queries at their true +/// absolute positions. That makes it exactly the `q0` V4's *compressed* branch +/// needs and the `lead` its dead prefix needs, so no new value is plumbed — +/// what changes is that `dsv4_attention` stops collapsing the vector to the +/// single position it inferred from the cache width. +/// +/// Gated on `ARC_MTP_PER_SEQ_KV` rather than on a flag of its own: +/// [`crate::kv_cache::front_align_batch`] is the only producer of a ragged +/// dense batch and it runs exclusively under `KvAdvance::PerSequence`, which +/// that flag authorizes. With it off no batch is ever left-aligned, so the +/// offsets would be uniform and `dsv4_attention` would take its scalar path +/// anyway — the gate makes that a guarantee instead of a coincidence. +/// +/// Resolved once, for the reason [`v4_fp8_kv_enabled`] documents: this is on +/// the per-layer path, and a `getenv` per layer per step is ~43 per forward. +fn ragged_row_q0(seqlen_offsets: &[usize]) -> Option<&[usize]> { + static ENABLED: std::sync::OnceLock = std::sync::OnceLock::new(); + let enabled = *ENABLED.get_or_init(crate::pipeline::mtp_pipeline::per_sequence_kv_requested); + ragged_row_q0_from(enabled, seqlen_offsets) +} + +/// The decision itself, separated from the process-wide [`OnceLock`] so it can +/// be tested against both settings rather than against whatever the test +/// runner's environment happened to be (`ragged_row_q0_is_opt_in`). Same shape +/// as [`fp8_kv_enabled_from`], for the same reason. +fn ragged_row_q0_from(enabled: bool, seqlen_offsets: &[usize]) -> Option<&[usize]> { + enabled.then_some(seqlen_offsets) +} + fn v4_fp8_kv_enabled() -> bool { static ENABLED: std::sync::OnceLock = std::sync::OnceLock::new(); *ENABLED.get_or_init(|| { @@ -5036,6 +5079,28 @@ mod kv_footprint_tests { ); } + /// The per-row ragged mask is **opt-in**, and off it hands + /// `dsv4_attention` nothing at all — not a uniform vector it would have to + /// recognise, but `None`, so the scalar path is taken structurally. + /// + /// Tests [`ragged_row_q0_from`] rather than [`ragged_row_q0`] for the same + /// reason `v4_fp8_kv_is_opt_in` does: the latter is a `OnceLock`. + #[test] + fn ragged_row_q0_is_opt_in() { + let offsets = [10usize, 8, 5]; + assert!( + ragged_row_q0_from(false, &offsets).is_none(), + "with per-sequence KV advance off no batch is ever left-aligned, so the model must \ + not be handed per-row positions at all" + ); + assert_eq!( + ragged_row_q0_from(true, &offsets), + Some(&offsets[..]), + "on, the row positions are `seqlen_offsets` verbatim — the value the engine already \ + threads for RoPE, not a new one" + ); + } + /// A non-`Normal` KV slot must be refused loudly rather than silently /// quantizing or windowing a 1-wide marker. #[test] @@ -6070,6 +6135,7 @@ mod tests { compress_ratio: CompressRatio::Csa, sliding_window: 4, raw_prefix: 0, + row_q0: None, }; // Compressed KV (precomputed by `Attention::compressed_kv` in the real @@ -6197,6 +6263,7 @@ mod tests { compress_ratio: CompressRatio::Standard, sliding_window: 1024, raw_prefix: 0, + row_q0: None, }; // Correct: sequence 0 attends over its own 3 keys. @@ -6340,6 +6407,7 @@ mod tests { compress_ratio: CompressRatio::Hca, sliding_window: 8, raw_prefix: 0, + row_q0: None, }; // Compressed KV: T_c = t_k / ratio = 1 entry, shape [B, 1, T_c, D]. @@ -6426,6 +6494,7 @@ mod tests { compress_ratio: CompressRatio::Standard, sliding_window: 4, raw_prefix: 0, + row_q0: None, }; // Standard layers must bypass the compressed branch entirely: even if diff --git a/mistralrs-core/src/models/dsv4_attention.rs b/mistralrs-core/src/models/dsv4_attention.rs index f645b20ed..abf06abbe 100644 --- a/mistralrs-core/src/models/dsv4_attention.rs +++ b/mistralrs-core/src/models/dsv4_attention.rs @@ -76,7 +76,7 @@ //! `sinks_attn`'s varlen backend, which has no mask parameter (see the routing //! comment in `attention/backends/sinks.rs`). -use candle_core::{DType, Result, Tensor}; +use candle_core::{DType, Device, Result, Tensor}; use mistralrs_quant::MatMul; use super::deepseek4::CompressRatio; @@ -256,6 +256,131 @@ fn compose_caller_mask( local.broadcast_add(&caller.contiguous()?) } +/// Does this call need the per-row treatment, and is the vector it was handed +/// coherent with the layout that treatment assumes? +/// +/// Returning `None` for a **uniform** batch is the load-bearing part, and it is +/// a dispatch, not an optimisation: `None` sends [`dsv4_attention`] down the +/// pre-change scalar path verbatim, so every B=1 request, every prefill and +/// every batch built with per-sequence KV advance off runs the same characters +/// it always did. That is what makes flag-off byte-identity structural rather +/// than a number a test happened to observe — and it is why this is a separate +/// function: the two branches agree bit-for-bit on their common domain, so no +/// numeric test can tell you which one ran (a mutation that routed uniform +/// batches down the ragged path survived every other test in this module). +/// +/// The two refusals both describe the *same* invariant from opposite sides: +/// [`crate::kv_cache::front_align_batch`] pads every row up to the batch +/// maximum, so `q0` is that maximum — no row may exceed it, and at least one +/// row must reach it. A batch that is ragged in some other way is not the +/// layout this module knows how to mask, and guessing would be a wrong answer +/// nothing downstream catches. +fn resolve_ragged_rows(rows: Option<&[usize]>, q0: usize, b_sz: usize) -> Result> { + let Some(rows) = rows.filter(|rows| rows.iter().any(|&r| r != q0)) else { + return Ok(None); + }; + if rows.len() != b_sz { + return Err(candle_core::Error::Msg(format!( + "dsv4_attention: {} per-row query positions for a batch of {b_sz}", + rows.len() + ))); + } + if let Some(&ahead) = rows.iter().find(|&&r| r > q0) { + return Err(candle_core::Error::Msg(format!( + "dsv4_attention: row at absolute query position {ahead} is AHEAD of the batch's \ + shared end column {q0}. Under left-alignment the longest row defines that column, \ + so `q0` is the batch maximum by construction — a larger value means the caller's \ + per-row lengths and the cache layout disagree, and masking from them would be \ + silently wrong" + ))); + } + if rows.iter().copied().max() != Some(q0) { + return Err(candle_core::Error::Msg(format!( + "dsv4_attention: no row sits at the batch's shared end column {q0} (the furthest is \ + {:?}), so this batch is not the left-aligned cohort `front_pad_kv_cache` produces. \ + Per-row masking is only defined for that layout", + rows.iter().copied().max() + ))); + } + Ok(Some(rows)) +} + +/// The `[b, 1, t_q, t_k + t_c]` additive mask a **left-aligned ragged** cohort +/// needs, from the rank-2 uniform validity the batch leader would have used. +/// +/// `rows[i]` is row `i`'s own absolute position for query row 0 and `q0` is the +/// batch's (the maximum, see [`Dsv4AttentionConfig::row_q0`]), so +/// `lead_i = q0 - rows[i]` is exactly the dead prefix +/// [`crate::kv_cache::front_pad_kv_cache`] left ahead of that row. Two +/// corrections, one per time base: +/// +/// * **raw** (end-anchored): the window/causality comparison is already right +/// for every row — shifting a row's queries and keys together preserves their +/// difference — so all that is wrong is the dead prefix, and cache column `c` +/// is dead for row `i` iff `c < lead_i`. Row `i`'s own position sits at +/// column `q0 + r >= lead_i` for every `r`, so no row is ever fully masked +/// and softmax cannot produce a NaN row. +/// * **compressed** (start-anchored): column `b` is absolute block `b` for +/// every row, so the threshold is rebuilt from `rows[i] + r` rather than +/// `q0 + r`. The batch-wide value is an *over*-estimate, which is why this +/// correction only ever removes columns. +/// +/// Only `b` host values cross to the device (the leads and the positions); the +/// `[b, 1, t_q, ·]` masks themselves are built on-device by broadcast. +#[allow(clippy::too_many_arguments)] +fn ragged_union_mask( + rows: &[usize], + q0: usize, + valid: &Tensor, + raw_base: usize, + t_k: usize, + ratio: usize, + dtype: DType, + dev: &Device, +) -> Result { + let b_sz = rows.len(); + let (t_q, n_keys) = valid.dims2()?; + let t_c = n_keys - t_k; + + // Raw half: kill each row's dead prefix, keep the shared window/causality. + let lead = rows.iter().map(|&r| (q0 - r) as f32).collect::>(); + let lead = Tensor::from_vec(lead, (b_sz, 1, 1, 1), dev)?; + let kp = Tensor::arange(raw_base as u32, (raw_base + t_k) as u32, dev)? + .to_dtype(DType::F32)? + .reshape((1, 1, 1, t_k))?; + let live = kp.broadcast_ge(&lead)?; // [b, 1, 1, t_k] (u8) + let raw = valid + .narrow(1, 0, t_k)? + .reshape((1, 1, t_q, t_k))? + .broadcast_mul(&live)? + .contiguous()?; // [b, 1, t_q, t_k] + + let valid = if t_c == 0 { + raw + } else { + let qp = Tensor::from_vec( + rows.iter().map(|&r| r as f32).collect::>(), + (b_sz, 1, 1, 1), + dev, + )? + .broadcast_add( + &Tensor::arange(0u32, t_q as u32, dev)? + .to_dtype(DType::F32)? + .reshape((1, 1, t_q, 1))?, + )?; // [b, 1, t_q, 1] — row `i`'s query `r` at `rows[i] + r` + let threshold = ((qp + 1.0)? / ratio as f64)?.floor()?; + let bp = Tensor::arange(0u32, t_c as u32, dev)? + .to_dtype(DType::F32)? + .reshape((1, 1, 1, t_c))?; + let comp = bp.broadcast_lt(&threshold)?.contiguous()?; // [b, 1, t_q, t_c] + Tensor::cat(&[&raw, &comp], 3)? + }; + + let zeros = Tensor::zeros((b_sz, 1, t_q, n_keys), dtype, dev)?; + let neg_inf = Tensor::full(f32::NEG_INFINITY, (b_sz, 1, t_q, n_keys), dev)?.to_dtype(dtype)?; + valid.contiguous()?.where_cond(&zeros, &neg_inf) +} + /// The span of raw keys any query row in this block can reach, as /// `(base, len)` into the full raw cache `[0, t_k_full)`. /// @@ -310,7 +435,7 @@ pub(crate) fn sdpa_sliding_window( /// Per-call configuration for V4 hybrid attention. #[derive(Debug, Clone, Copy)] -pub struct Dsv4AttentionConfig { +pub struct Dsv4AttentionConfig<'a> { /// Per-layer compress ratio (0/4/128 → `Standard`/`Csa`/`Hca`). pub compress_ratio: CompressRatio, /// Sliding-window size for the always-on local branch (V4 default 128). @@ -334,6 +459,40 @@ pub struct Dsv4AttentionConfig { /// `raw_prefix` above [`raw_keep_span`]'s base is a hard error, not a /// silent truncation. pub raw_prefix: usize, + /// Per-batch-row absolute position of query row 0, when the batch is a + /// **left-aligned ragged** cohort (`front_pad_kv_cache`). `None` — every + /// path but the dense per-sequence-MTP decode — means every row shares the + /// single position this module has always inferred from the cache width, + /// and the code below is then character-for-character what it was. + /// + /// It exists because left-alignment preserves one of V4's two time bases + /// and breaks the other: + /// + /// * the **raw** window is end-anchored, so shifting a row's live run right + /// by `lead` shifts its queries *and* its keys by the same `lead`. Both + /// the causal and the sliding-window comparison are differences of the + /// two, so both are invariant — in cache-column space the batch-wide + /// position is exactly right, for every row. What left-alignment *does* + /// add is a dead prefix of `lead` zero columns; those are killed here too + /// (`row < lead` is not a live key), so this module needs no help from + /// the caller's mask to be correct. + /// * the **compressed** blocks are start-anchored — column `b` is absolute + /// block `b` for every row — so their threshold `b < (q_abs + 1) / ratio` + /// needs each row's *own* `q_abs`. Fed the batch-wide one it is too + /// **permissive**: a row that has advanced less would see compressed + /// blocks covering tokens it has not generated. That is a wrong answer, + /// not a slow one, which is why this is a hard blocker rather than a + /// pessimisation. + /// + /// The value is `seqlen_offsets[i]` — the per-row past length the engine + /// already threads to every attention call for RoPE. No new plumbing + /// carries it; the raggedness was already in the arguments. + /// + /// A row may not be *ahead* of the batch (`row_q0[i] > q0` is rejected): + /// under left-alignment the longest row defines the shared end column, so + /// `q0` is the maximum by construction and a larger value means the caller + /// and the cache disagree about the layout. + pub row_q0: Option<&'a [usize]>, } /// V4 hybrid attention dispatch — a single softmax over the union of the raw @@ -348,7 +507,7 @@ pub fn dsv4_attention( attention_mask: Option<&Tensor>, flash_params: &FlashParams, sdpa_params: &SdpaParams, - cfg: Dsv4AttentionConfig, + cfg: Dsv4AttentionConfig<'_>, ) -> Result { // ---- Standard layers: sliding-window + sink, same raw branch as CSA/HCA // (see module docs — the reference SWA-onlys ratio-0 layers; dense causal @@ -367,7 +526,7 @@ pub fn dsv4_attention( compressed_kv }; - let (_b, _h, t_q, _d) = q.dims4()?; + let (b_sz, _h, t_q, _d) = q.dims4()?; let t_k_given = k.dim(2)?; // `k` spans absolute positions `[raw_prefix, t_k_full)`; `raw_prefix == 0` // (every path but packed-KV decode) makes this the cache length, exactly as @@ -388,6 +547,16 @@ pub fn dsv4_attention( cfg.raw_prefix )))?; + // ---- Left-aligned ragged cohort: does any row sit behind the batch? ---- + // Dispatching on "every row is at `q0`" rather than on the flag is what + // makes flag-off byte-identity **structural**: a uniform batch — which is + // every B=1 request, every prefill, and every batch the engine builds with + // per-sequence KV advance off — takes `None` here and not one line below + // this point differs from the pre-change code. See + // [`Dsv4AttentionConfig::row_q0`] for why the two branches need it and the + // raw branch does not. + let row_q0 = resolve_ragged_rows(cfg.row_q0, q0, b_sz)?; + // ---- Raw working-set narrowing ---------------------------------------- // Only the trailing `t_q + window - 1` raw keys are reachable by ANY query // row in this block: row `r` sits at absolute position `q0 + r` and attends @@ -471,11 +640,29 @@ pub fn dsv4_attention( // Boolean validity → additive mask (0 / -inf), broadcast over batch+heads. let n_keys = valid.dim(1)?; - let zeros = Tensor::zeros((t_q, n_keys), q.dtype(), dev)?; - let neg_inf = Tensor::full(f32::NEG_INFINITY, (t_q, n_keys), dev)?.to_dtype(q.dtype())?; - let mask = valid - .where_cond(&zeros, &neg_inf)? - .reshape((1, 1, t_q, n_keys))?; + let mask = match row_q0 { + // Uniform batch: verbatim the pre-change code. + None => { + let zeros = Tensor::zeros((t_q, n_keys), q.dtype(), dev)?; + let neg_inf = + Tensor::full(f32::NEG_INFINITY, (t_q, n_keys), dev)?.to_dtype(q.dtype())?; + valid + .where_cond(&zeros, &neg_inf)? + .reshape((1, 1, t_q, n_keys))? + } + // Ragged cohort: `[b, 1, t_q, n_keys]`, because the compressed + // threshold and the dead prefix are both per-row. + Some(rows) => ragged_union_mask( + rows, + q0, + &valid, + raw_base, + t_k, + cfg.compress_ratio.ratio(), + q.dtype(), + dev, + )?, + }; // Fold in whatever the caller knows that this module cannot see: padding // columns in a ragged batch, the graph-decode length mask, custom bias. @@ -696,6 +883,7 @@ mod tests { compress_ratio: CompressRatio::Standard, sliding_window: window, raw_prefix: 0, + row_q0: None, }; let actual = dsv4_attention(&q, &k, &v, None, None, &flash, &sdpa, cfg)?; @@ -742,6 +930,7 @@ mod tests { compress_ratio: CompressRatio::Standard, sliding_window: t, raw_prefix: 0, + row_q0: None, }; let actual = dsv4_attention(&q, &k, &v, None, None, &flash, &sdpa, cfg)?; @@ -797,6 +986,7 @@ mod tests { compress_ratio: CompressRatio::Standard, sliding_window: window, raw_prefix: 0, + row_q0: None, }; let actual = dsv4_attention(&q, &k, &v, None, None, &flash, &sdpa, cfg)?; @@ -883,6 +1073,7 @@ mod tests { compress_ratio: ratio, sliding_window: window, raw_prefix: 0, + row_q0: None, }, )?; @@ -903,6 +1094,7 @@ mod tests { compress_ratio: ratio, sliding_window: window, raw_prefix: base, + row_q0: None, }, )?; @@ -951,6 +1143,7 @@ mod tests { compress_ratio: CompressRatio::Standard, sliding_window: 16, raw_prefix: 4, + row_q0: None, }, ) .unwrap_err(); @@ -1000,6 +1193,7 @@ mod tests { compress_ratio: CompressRatio::Csa, sliding_window: window, raw_prefix: 0, + row_q0: None, }; let out = dsv4_attention(&q, &k, &v, Some(&comp), None, &flash, &sdpa, cfg)?; let out_v: Vec = out.to_dtype(DType::F32)?.flatten_all()?.to_vec1()?; @@ -1123,6 +1317,7 @@ mod tests { compress_ratio: CompressRatio::Csa, sliding_window: t, // window covers everything → windowed == causal raw_prefix: 0, + row_q0: None, }; let actual = dsv4_attention(&q, &k, &v, None, None, &flash, &sdpa, cfg)?; @@ -1152,6 +1347,7 @@ mod tests { compress_ratio: CompressRatio::Csa, sliding_window: 4, raw_prefix: 0, + row_q0: None, }; let out = dsv4_attention(&q, &k, &v, Some(&comp), None, &flash, &sdpa, cfg)?; assert_eq!(out.dims(), &[b, h, t_q, d]); @@ -1221,6 +1417,10 @@ mod tests { compress_ratio: CompressRatio::Csa, sliding_window: window, raw_prefix: 0, + // Uniform batch: this fixture predates `row_q0` and builds every + // row at the same length, so `None` is the identity that preserves + // exactly what it asserted before the field existed. + row_q0: None, }; let with_comp = dsv4_attention(&q, &k, &v, Some(&comp), None, &flash, &sdpa, cfg)?; let without_comp = dsv4_attention(&q, &k, &v, None, None, &flash, &sdpa, cfg)?; @@ -1328,6 +1528,7 @@ mod tests { compress_ratio: CompressRatio::Hca, sliding_window: 128, raw_prefix: 0, + row_q0: None, }; let out = dsv4_attention(&q, &k, &v, Some(&comp), None, &flash, &sdpa, cfg)?; assert_eq!(out.dims(), &[b, h, t_q, d]); @@ -1371,6 +1572,7 @@ mod tests { compress_ratio: CompressRatio::Csa, sliding_window: window, raw_prefix: 0, + row_q0: None, }; // The engine pads every sequence to the batch max and reports PADDED @@ -1446,6 +1648,7 @@ mod tests { compress_ratio: CompressRatio::Csa, sliding_window: window, // window >= t, so ONLY causality is on trial raw_prefix: 0, + row_q0: None, }; let padded = vec![t as u32; b]; let flash = varlen_flash_params(&padded, &padded, &device); @@ -1534,6 +1737,7 @@ mod tests { compress_ratio: CompressRatio::Csa, sliding_window: window, raw_prefix: 0, + row_q0: None, }; let padded = vec![t as u32; b]; let flash = varlen_flash_params(&padded, &padded, &device); @@ -1612,6 +1816,7 @@ mod tests { // overlap, which is what makes the double counting observable. sliding_window: t, raw_prefix: 0, + row_q0: None, }; // Rank-4 causal mask (what `CausalMasker` hands the layer), raw width. let caller = left_padded_causal_mask(&[0], t, &device)?; @@ -1772,6 +1977,7 @@ mod tests { compress_ratio: CompressRatio::Standard, sliding_window: window, raw_prefix: 0, + row_q0: None, }; let (base, keep) = raw_keep_span(1, window, t_k); @@ -1816,6 +2022,7 @@ mod tests { compress_ratio: CompressRatio::Standard, sliding_window: 4, raw_prefix: 0, + row_q0: None, }; let bad = Tensor::zeros((1, 1, t, t + 3), DType::F32, &device)?; let err = dsv4_attention(&q, &k, &v, None, Some(&bad), &flash, &sdpa, cfg).unwrap_err(); @@ -1836,4 +2043,486 @@ mod tests { ); Ok(()) } + + // ======================================================================= + // LEFT-ALIGNED RAGGED COHORT (per-sequence MTP KV advance, PR #92/#95) + // ======================================================================= + + /// Lay `rows[i]` (`[1, kv, live_i, d]`) into a **left-aligned** batched + /// buffer of width `t_k_full`: `t_k_full - live_i` zero columns, then the + /// row's real content. This is what `front_pad_kv_cache` produces, zeros + /// and all — and a zero K row is not a masked row, which is the whole + /// reason the dead prefix has to be masked rather than tolerated. + fn left_align_rows(rows: &[Tensor], t_k_full: usize, dev: &Device) -> Result { + let mut padded = Vec::with_capacity(rows.len()); + for row in rows { + let (_, kv, live, d) = row.dims4()?; + let lead = t_k_full - live; + padded.push(if lead == 0 { + row.clone() + } else { + let zeros = Tensor::zeros((1, kv, lead, d), row.dtype(), dev)?; + Tensor::cat(&[&zeros, row], 2)? + }); + } + Tensor::cat(&padded.iter().collect::>(), 0) + } + + /// The shared ragged fixture: three rows with DIFFERENT committed pasts, + /// each with its own input stream, left-aligned into one buffer. + /// + /// The numbers are chosen so all three degeneracy traps are open — see the + /// assertions in [`ragged_cohort_matches_each_row_run_alone`], which run + /// before anything else and fail loudly if a later edit closes one. + struct RaggedFixture { + lens: Vec, + t_q: usize, + window: usize, + t_c: usize, + h: usize, + d: usize, + padded: usize, + t_k_full: usize, + q: Tensor, + k: Tensor, + v: Tensor, + comp: Tensor, + k_rows: Vec, + v_rows: Vec, + } + + fn ragged_fixture(dev: &Device) -> Result { + let (h, d, window, t_q, t_c) = (2usize, 16usize, 10usize, 2usize, 4usize); + let lens = vec![10usize, 8, 5]; + let b = lens.len(); + let padded = *lens.iter().max().unwrap(); + let t_k_full = padded + t_q; + + // Every row gets its OWN stream. One shared stream would make block `j` + // cover identical tokens on every row and hide any row mix-up — the + // exact fixture defect PR #95 found in its first three mutations. + let mut k_rows = Vec::with_capacity(b); + let mut v_rows = Vec::with_capacity(b); + for (i, &l) in lens.iter().enumerate() { + k_rows.push(mk(1, 1, l + t_q, d, 0.13 + 0.05 * i as f32, dev)?); + v_rows.push(mk(1, 1, l + t_q, d, 0.21 + 0.05 * i as f32, dev)?); + } + Ok(RaggedFixture { + q: mk(b, h, t_q, d, 0.09, dev)?, + comp: mk(b, 1, t_c, d, 0.04, dev)?, + k: left_align_rows(&k_rows, t_k_full, dev)?, + v: left_align_rows(&v_rows, t_k_full, dev)?, + k_rows, + v_rows, + lens, + t_q, + window, + t_c, + h, + d, + padded, + t_k_full, + }) + } + + impl RaggedFixture { + fn leads(&self) -> Vec { + self.lens.iter().map(|&l| self.padded - l).collect() + } + fn cfg(&self) -> Dsv4AttentionConfig<'_> { + Dsv4AttentionConfig { + compress_ratio: CompressRatio::Csa, + sliding_window: self.window, + raw_prefix: 0, + row_q0: None, + } + } + /// The same call the engine makes for a left-aligned cohort: row `i`'s + /// own absolute query position is its committed past `lens[i]`, which + /// is precisely what `seqlen_offsets[i]` already carries. + fn ragged_cfg(&self) -> Dsv4AttentionConfig<'_> { + Dsv4AttentionConfig { + row_q0: Some(&self.lens), + ..self.cfg() + } + } + } + + /// **The dispatch itself**, pinned directly — because no numeric test can + /// pin it. The two branches of [`dsv4_attention`]'s mask build agree + /// bit-for-bit on a uniform batch, so a mutation that routes uniform + /// batches down the ragged path is numerically invisible and survived + /// every other test in this module. It is still a real defect: it would + /// build a `[b, 1, t_q, n_keys]` mask and two extra host→device copies on + /// every layer of every B=1 decode step, and it would end the structural + /// flag-off guarantee. + #[test] + fn a_uniform_batch_never_reaches_the_ragged_path() -> Result<()> { + // Uniform, at the batch's shared end column: not ragged, whatever the + // value is. + for (q0, b) in [(0usize, 1usize), (7, 1), (7, 4), (128, 3)] { + let rows = vec![q0; b]; + assert!( + resolve_ragged_rows(Some(&rows), q0, b)?.is_none(), + "q0={q0} b={b}: a uniform vector must dispatch to the scalar path, not merely \ + compute the same numbers on the ragged one" + ); + } + // Absent entirely — every path but per-sequence MTP. + assert!(resolve_ragged_rows(None, 12, 3)?.is_none()); + // Genuinely ragged, and coherent: taken. + let ragged = [12usize, 10, 7]; + assert_eq!( + resolve_ragged_rows(Some(&ragged), 12, 3)?, + Some(&ragged[..]), + "a cohort with rows behind the leader is exactly what the ragged path is for" + ); + Ok(()) + } + + /// 🔑 **THE claim.** Every row of a left-aligned ragged cohort must produce + /// exactly what that row produces run ALONE with its dead prefix simply + /// absent — which is the definition of "the batch changed nothing". + /// + /// Both of V4's time bases are on trial and they fail in opposite + /// directions: + /// + /// * the **raw** window is end-anchored, so left-alignment shifts a row's + /// queries and keys together and the window/causality comparison is + /// already right — what it adds is `lead_i` zero columns that score + /// logit 0 and take real softmax weight; + /// * the **compressed** blocks are start-anchored, so the batch-wide + /// threshold `b < (q0 + r + 1) / ratio` is too **permissive** for a row + /// that has advanced less: it would attend blocks covering tokens that + /// row has not generated. + /// + /// The second is why a `[B, 1, t_q, k]` K/V mask over the raw columns — + /// everything PR #92 built — is not sufficient on its own. + #[test] + fn ragged_cohort_matches_each_row_run_alone() -> Result<()> { + let device = Device::Cpu; + let f = ragged_fixture(&device)?; + let ratio = CompressRatio::Csa.ratio(); + + // ---- Non-degeneracy FIRST (D12). A fixture where every row lands on + // the same compressed column proves nothing; PR #95 shipped three + // surviving mutations before noticing exactly that. + let thresholds: Vec> = f + .lens + .iter() + .map(|&l| (0..f.t_q).map(|r| (l + r + 1) / ratio).collect()) + .collect(); + assert!( + thresholds.iter().any(|t| *t != thresholds[0]), + "degenerate fixture: every row sees the same compressed blocks, so the per-row \ + threshold is never exercised and this test would pass with the fix deleted" + ); + assert!( + thresholds.iter().flatten().copied().min().unwrap() < f.t_c, + "degenerate fixture: no row is short enough for a compressed block to be masked" + ); + let (raw_base, _) = raw_keep_span(f.t_q, f.window, f.t_k_full); + assert!( + f.leads().iter().any(|&lead| lead > raw_base), + "degenerate fixture: every dead prefix falls outside the retained raw span, so the \ + raw-side correction is never exercised" + ); + + let (sdpa, _sinks) = sdpa_params_with_sinks(f.d, f.h, &device)?; + let flash = empty_flash_params(); + let batched = dsv4_attention( + &f.q, + &f.k, + &f.v, + Some(&f.comp), + None, + &flash, + &sdpa, + f.ragged_cfg(), + )?; + assert_eq!(batched.dims(), &[f.lens.len(), f.h, f.t_q, f.d]); + assert!( + flat(&batched).iter().all(|x| x.is_finite()), + "a row was fully masked — softmax produced a NaN row" + ); + + let mut solos: Vec> = Vec::new(); + for (i, &l) in f.lens.iter().enumerate() { + let qi = f.q.narrow(0, i, 1)?.contiguous()?; + let ci = f.comp.narrow(0, i, 1)?.contiguous()?; + // Run alone: no dead prefix exists, so `q0` is inferred as `l` and + // the compressed threshold is right by construction. + let solo = dsv4_attention( + &qi, + &f.k_rows[i], + &f.v_rows[i], + Some(&ci), + None, + &flash, + &sdpa, + f.cfg(), + )?; + let got = flat(&batched.narrow(0, i, 1)?.contiguous()?); + let want = flat(&solo); + // Batch-3 and batch-1 GEMMs tile — and therefore round — + // differently on `MatMul`'s F16 CPU path; same documented noise + // floor as `union_decode_matches_scalar_reference`. + let diff = max_abs_diff(&got, &want); + assert!( + diff < 1.5e-3, + "row {i} (committed past {l}, dead prefix {}) differs between the ragged cohort \ + and its solo run (max abs diff {diff})", + f.padded - l + ); + solos.push(want); + } + + // Teeth: the rows must be grossly distinguishable, or the agreement + // above is the vacuous kind. + assert!( + max_abs_diff(&solos[0], &solos[2]) > 1e-2, + "fixture rows are too similar for the equality assertions to have teeth" + ); + Ok(()) + } + + /// The negative control for the test above: fed the batch-wide position — + /// i.e. with the per-row `q0` deleted — the shorter rows must DISAGREE + /// with their solo runs. + /// + /// Without this, `ragged_cohort_matches_each_row_run_alone` cannot + /// distinguish "the fix works" from "the fixture never needed it". + #[test] + fn without_the_per_row_q0_the_shorter_rows_are_wrong() -> Result<()> { + let device = Device::Cpu; + let f = ragged_fixture(&device)?; + let (sdpa, _sinks) = sdpa_params_with_sinks(f.d, f.h, &device)?; + let flash = empty_flash_params(); + + // `row_q0: None` is exactly the pre-change behaviour: one scalar `q0` + // inferred from the cache width, applied to every row. + let batched = dsv4_attention( + &f.q, + &f.k, + &f.v, + Some(&f.comp), + None, + &flash, + &sdpa, + f.cfg(), + )?; + let mut worst = 0f32; + for i in 0..f.lens.len() { + let qi = f.q.narrow(0, i, 1)?.contiguous()?; + let ci = f.comp.narrow(0, i, 1)?.contiguous()?; + let solo = dsv4_attention( + &qi, + &f.k_rows[i], + &f.v_rows[i], + Some(&ci), + None, + &flash, + &sdpa, + f.cfg(), + )?; + worst = worst.max(max_abs_diff( + &flat(&batched.narrow(0, i, 1)?.contiguous()?), + &flat(&solo), + )); + } + assert!( + worst > 1e-2, + "the batch-wide position produced the right answer anyway (max abs diff {worst}), so \ + this fixture cannot detect the bug the per-row `q0` fixes" + ); + Ok(()) + } + + /// **Flag-off byte-identity, asserted as well as dispatched.** A uniform + /// batch — every B=1 request, every prefill, every batch the engine builds + /// with per-sequence KV advance off — must produce bit-identical output + /// whether or not the per-row vector is supplied. + /// + /// The guarantee is structural (`dsv4_attention` dispatches on "does any + /// row differ from `q0`" and the uniform branch is the pre-change code + /// verbatim), so this asserts EXACT equality, not a tolerance. B=1 is + /// included on purpose: it is the only regime with a solid measurement and + /// it must not move. + #[test] + fn a_uniform_row_q0_is_bit_identical_to_none() -> Result<()> { + let device = Device::Cpu; + let (h, d) = (2usize, 16usize); + let (sdpa, _sinks) = sdpa_params_with_sinks(d, h, &device)?; + let flash = empty_flash_params(); + let mut checked = 0usize; + for (ratio, t_c) in [ + (CompressRatio::Standard, 0usize), + (CompressRatio::Csa, 3), + (CompressRatio::Hca, 2), + ] { + for b in [1usize, 3] { + for (t_k, t_q) in [(12usize, 12usize), (12, 1), (12, 4)] { + for window in [4usize, 10, 32] { + let q = mk(b, h, t_q, d, 0.09, &device)?; + let k = mk(b, 1, t_k, d, 0.19, &device)?; + let v = mk(b, 1, t_k, d, 0.29, &device)?; + let comp = (t_c > 0) + .then(|| mk(b, 1, t_c, d, 0.04, &device)) + .transpose()?; + let cfg = Dsv4AttentionConfig { + compress_ratio: ratio, + sliding_window: window, + raw_prefix: 0, + row_q0: None, + }; + let rows = vec![t_k - t_q; b]; + let with = dsv4_attention( + &q, + &k, + &v, + comp.as_ref(), + None, + &flash, + &sdpa, + Dsv4AttentionConfig { + row_q0: Some(&rows), + ..cfg + }, + )?; + let without = + dsv4_attention(&q, &k, &v, comp.as_ref(), None, &flash, &sdpa, cfg)?; + assert_eq!( + max_abs_diff(&flat(&with), &flat(&without)), + 0.0, + "{ratio:?} b={b} t_k={t_k} t_q={t_q} window={window}: a uniform \ + per-row vector changed the result, so flag-off is NOT byte-identical" + ); + checked += 1; + } + } + } + } + assert!(checked >= 54, "only {checked} configurations exercised"); + Ok(()) + } + + /// **The dead prefix cannot vote.** Fill each row's zero-filled prefix with + /// loud, distinctive keys and values; no row's output may move by a single + /// bit. A `-inf` logit gives an exactly-zero softmax weight, so this is an + /// exact-equality assertion, not a tolerance. + /// + /// This is the FP8-KV failure mode restated: serving from a zero-filled + /// prefix is a wrong answer nothing downstream catches. + #[test] + fn the_dead_prefix_cannot_influence_its_row() -> Result<()> { + let device = Device::Cpu; + let f = ragged_fixture(&device)?; + let leads = f.leads(); + let (sdpa, _sinks) = sdpa_params_with_sinks(f.d, f.h, &device)?; + let flash = empty_flash_params(); + + let clean = dsv4_attention( + &f.q, + &f.k, + &f.v, + Some(&f.comp), + None, + &flash, + &sdpa, + f.ragged_cfg(), + )?; + + let (kp, vp) = perturb_keys(&f.k, &f.v, |i, j| j < leads[i])?; + let perturbed = dsv4_attention( + &f.q, + &kp, + &vp, + Some(&f.comp), + None, + &flash, + &sdpa, + f.ragged_cfg(), + )?; + assert_eq!( + max_abs_diff(&flat(&perturbed), &flat(&clean)), + 0.0, + "the zero-filled dead prefix took softmax weight — a zero K row is not a masked row" + ); + + // Negative control: the SAME perturbation one column later — the first + // genuinely live column of the longest-padded row — must move the + // output by O(1), or the probe above is vacuous. + let last = leads.len() - 1; + let (kc, vc) = perturb_keys(&f.k, &f.v, |i, j| i == last && j == leads[last])?; + let control = dsv4_attention( + &f.q, + &kc, + &vc, + Some(&f.comp), + None, + &flash, + &sdpa, + f.ragged_cfg(), + )?; + assert!( + max_abs_diff(&flat(&control), &flat(&clean)) > 1e-2, + "perturbing a LIVE key did not move the output, so the dead-prefix probe proves nothing" + ); + Ok(()) + } + + /// A row cannot be **ahead** of the cohort, and the vector must describe + /// the batch it is handed. Both are caller/cache disagreements that would + /// otherwise mask from the wrong geometry silently. + #[test] + fn an_inconsistent_row_q0_vector_is_refused() -> Result<()> { + let device = Device::Cpu; + let f = ragged_fixture(&device)?; + let (sdpa, _sinks) = sdpa_params_with_sinks(f.d, f.h, &device)?; + let flash = empty_flash_params(); + let call = |rows: &[usize]| { + dsv4_attention( + &f.q, + &f.k, + &f.v, + Some(&f.comp), + None, + &flash, + &sdpa, + Dsv4AttentionConfig { + row_q0: Some(rows), + ..f.cfg() + }, + ) + }; + + // `q0` is the batch maximum by construction (left-alignment ends every + // row at the same column), so a larger value is incoherent. + let ahead = vec![f.padded + 1, f.lens[1], f.lens[2]]; + let err = call(&ahead).unwrap_err().to_string(); + assert!( + err.contains("AHEAD of the batch"), + "expected an ahead-of-batch refusal, got: {err}" + ); + + // Wrong width: silently zipping a short vector against the batch would + // mask row 2 from row 0's geometry. + let short = vec![f.lens[0], f.lens[1]]; + let err = call(&short).unwrap_err().to_string(); + assert!( + err.contains("per-row query positions for a batch of"), + "expected a batch-width refusal, got: {err}" + ); + + // Ragged, but with NO row at the shared end column: not the layout + // `front_pad_kv_cache` produces, so the geometry this module would mask + // from does not hold. + let floating: Vec = f.lens.iter().map(|l| l - 1).collect(); + let err = call(&floating).unwrap_err().to_string(); + assert!( + err.contains("no row sits at the batch's shared end column"), + "expected a not-left-aligned refusal, got: {err}" + ); + Ok(()) + } } diff --git a/mistralrs-core/src/pipeline/mtp_pipeline.rs b/mistralrs-core/src/pipeline/mtp_pipeline.rs index cb7e2723b..4fc90f201 100644 --- a/mistralrs-core/src/pipeline/mtp_pipeline.rs +++ b/mistralrs-core/src/pipeline/mtp_pipeline.rs @@ -1300,30 +1300,40 @@ impl MtpSpeculativePipeline { /// How this pipeline advances the KV cache after a verify, decided once. /// /// Default [`KvAdvance::Cohort`]. `ARC_MTP_PER_SEQ_KV=1` asks for - /// [`KvAdvance::PerSequence`]; it is granted only when **both** of these - /// hold, and the reason it was refused is logged by name so a run that - /// asked for it never silently gets the ceiling instead: + /// [`KvAdvance::PerSequence`]; it is granted only when **all three** of + /// these hold, and the reason it was refused is logged by name so a run + /// that asked for it never silently gets the ceiling instead: /// /// 1. every cache slot can carry its own length - /// ([`cache_supports_per_sequence_advance`]); - /// 2. the target model masks a left-aligned ragged batch's dead prefix - /// ([`Self::target_masks_ragged_batches`]). + /// ([`cache_supports_per_sequence_advance`]) — cleared for V4 by + /// wave63-CO under `ARC_V4_XS_PER_SEQ`; + /// 2. the target model masks a left-aligned ragged batch's dead prefix AND + /// its start-anchored compressed columns + /// ([`Self::target_masks_ragged_batches`]) — cleared for V4 by the + /// per-row `row_q0` in `dsv4_attention`; + /// 3. this pipeline's own step arithmetic can express a batch whose rows + /// hold different lengths ([`Self::step_carries_per_sequence_lengths`]) + /// — **still refused**, and now the only one. /// - /// (2) is not decoration. Per-sequence lengths in one dense buffer mean - /// `front_pad_kv_cache` zero-fills ahead of the shorter rows, and a zero K - /// row is **not** a masked row — it scores logit 0 and takes real softmax - /// weight. Serving from those columns would be a wrong answer nothing - /// downstream catches, which is the exact failure this project already paid - /// for once with FP8 KV. Refusing is the only correct behaviour until a - /// model threads [`crate::layers_masker::RaggedKvLens`]. + /// None of these is decoration. Per-sequence lengths in one dense buffer + /// mean `front_pad_kv_cache` zero-fills ahead of the shorter rows, and a + /// zero K row is **not** a masked row — it scores logit 0 and takes real + /// softmax weight. Serving from those columns would be a wrong answer + /// nothing downstream catches, which is the exact failure this project + /// already paid for once with FP8 KV. + /// + /// (3) is worth stating explicitly because granting the mode on (1) and (2) + /// alone would be **actively harmful**, not merely useless: `window_ok` + /// would fail on every ragged step, the fused path would defer to the + /// target's own decode, and that path builds no mask at all at `t_q == 1` + /// (`layers_masker.rs:285`) — so the dead prefix would be attended as real + /// keys by the very fallback that was supposed to be the safe one. pub(crate) fn kv_advance(&self) -> KvAdvance { *self.kv_advance.get_or_init(|| { if !per_sequence_kv_requested() { return KvAdvance::Cohort; } - let refusal = cache_supports_per_sequence_advance(&self.target_cache) - .err() - .or_else(|| self.target_masks_ragged_batches().err()); + let refusal = per_sequence_refusal(&self.target_cache); match refusal { None => { tracing::info!( @@ -1345,31 +1355,6 @@ impl MtpSpeculativePipeline { }) } - /// Whether the target model applies a per-row mask over a left-aligned - /// ragged batch's dead prefix. - /// - /// **No model in the tree does yet**, and this says so rather than - /// pretending. The plumbing that would flip it is exactly one wire: - /// `clone_in_cache` already computes `lead_pad` per sequence - /// ([`crate::kv_cache::front_align_batch`]) and - /// [`crate::layers_masker::CausalMasker::make_left_padded_causal_mask`] - /// already turns that into the `[B, 1, t_q, k]` additive mask a model wants - /// — what is missing is carrying it from the pipeline through - /// `inputs_processor` into the model forward, where today the mask is - /// rebuilt from `seqlen_offsets` alone. - /// - /// It is stated as a capability rather than hardcoded at the call site so - /// that the model that adds the wire overrides one thing, here. - fn target_masks_ragged_batches(&self) -> std::result::Result<(), String> { - Err( - "no model in this tree threads a ragged-batch mask into its forward yet, so the \ - zero-filled dead prefix a left-aligned per-sequence cache leaves would be attended \ - as real keys (a zero K row scores logit 0 and takes softmax weight). See \ - `MtpSpeculativePipeline::target_masks_ragged_batches`." - .to_string(), - ) - } - /// Snapshot of this pipeline's MTP acceptance counters. pub fn acceptance(&self) -> MtpAcceptance { self.acceptance.snapshot() @@ -1929,6 +1914,111 @@ pub(crate) fn cache_supports_per_sequence_advance( Ok(()) } +/// The whole of [`MtpSpeculativePipeline::kv_advance`]'s decision: the three +/// capability checks, composed in order, or `None` when all three pass. +/// +/// Composed here rather than inline in `kv_advance` so that "which checks are +/// required" is itself testable — dropping one is otherwise invisible, because +/// `kv_advance` needs a whole target pipeline to call. +/// +/// Order matters only for which reason gets reported. Each check is +/// independent, and each has been the *last* one at some point: (1) fell to +/// wave63-CO, (2) to the per-row `row_q0`, and (3) is where it stands now. +pub(crate) fn per_sequence_refusal(cache: &EitherCache) -> Option { + cache_supports_per_sequence_advance(cache) + .err() + .or_else(|| model_masks_ragged_batches(cache).err()) + .or_else(|| step_carries_per_sequence_lengths().err()) +} + +/// Whether the model behind this cache masks a left-aligned ragged batch, or +/// why it does not. The cache-inspecting half of +/// [`MtpSpeculativePipeline::kv_advance`]'s second check, split out so it is +/// testable without standing up a whole target pipeline — the same shape +/// [`cache_supports_per_sequence_advance`] has. +/// +/// The discriminator is the `XsRolling` compressor slot. DeepSeek V4 is the +/// only architecture in the tree that allocates any, and it is exactly the +/// architecture whose attention now takes a per-row `row_q0`. +pub(crate) fn model_masks_ragged_batches(cache: &EitherCache) -> std::result::Result<(), String> { + let EitherCache::Normal(normal) = cache else { + return Err( + "the target is on PagedAttention; this dense-cache mask is not what it needs" + .to_string(), + ); + }; + let is_v4 = normal + .lock() + .unwrap() + .0 + .iter() + .any(|slot| matches!(slot, crate::kv_cache::KvCache::XsRolling(_))); + if is_v4 { + return Ok(()); + } + Err( + "the target has no `XsRolling` compressor slots, so it is not DeepSeek V4 and no other \ + model in this tree threads a ragged-batch mask into its forward. The zero-filled dead \ + prefix a left-aligned per-sequence cache leaves would be attended as real keys (a zero \ + K row scores logit 0 and takes softmax weight). See \ + `MtpSpeculativePipeline::target_masks_ragged_batches`." + .to_string(), + ) +} + +/// Whether the **fused step's own bookkeeping** can express a batch whose +/// rows hold different committed lengths. +/// +/// It cannot, and this is the blocker that remains after the mask. It is +/// stated separately, and named, because the two previous waves each +/// believed the mask was the last one. +/// +/// [`MtpSpeculativePipeline::step_full`] keys the entire step on ONE scalar +/// `cache_len = current_normal_cache_len(self)` — the *batched* buffer's +/// width, which under left-alignment is `max_j L_j`, not any row's length. +/// Four of its uses are arithmetically wrong the moment the rows diverge, +/// and they fail in different ways: +/// +/// * `uncached[i] = toks_i.len() - cache_len` saturates to **0** for every +/// row but the longest, which fails `window_ok`'s `(1..=w)` check and +/// sends every step down the deferral path — per-sequence advance would +/// be *reachable but inert*, and the deferral is itself unsafe here +/// because the target's own decode builds no mask at `t_q == 1`; +/// * `toks[cache_len..]` (the verify window) is a **panic** site once +/// `cache_len > toks_i.len()`, reachable only because `window_ok` bails +/// first; +/// * `run_target_forward_batch(.., Some((w, cache_len)))` becomes +/// `seqlen_offsets[i] = cache_len` for every row +/// (`inputs_processor.rs:244` — `last_n_context_len.1` is a shared +/// scalar), so RoPE would place a shorter row's queries `lead_i` past +/// their true positions. That same vector is what feeds the model's +/// `row_q0`, so the mask this PR built is correct but would be handed +/// uniform values; +/// * `set_len(cache_len + committed_i)` counts the dead prefix **inside** +/// the recorded length (`front_pad_single` sets `current_seq_len = +/// target_len`, and `clone_out_cache` copies the batched length verbatim +/// into every row). So the next `front_align_batch` pads relative to +/// already-padded lengths and the dead prefix **accumulates linearly**, +/// not as the `sqrt(steps)` random walk `front_pad_single`'s doc claims. +/// Nothing in `Sequence` or `SingleCache` records a row's true live +/// length; the only honest source is `seq.get_toks().len()`. +/// +/// Fixing it is a change to this pipeline's step arithmetic plus a per-row +/// route through `inputs_processor` — not a model change, which is why it +/// is not in the PR that added the mask. +pub(crate) fn step_carries_per_sequence_lengths() -> std::result::Result<(), String> { + Err( + "`step_full` keys the whole fused step on one shared `cache_len` (the BATCHED buffer \ + width, which left-alignment makes `max_j L_j`): `uncached[i]` saturates to 0 for \ + every row but the longest and trips `window_ok`, `toks[cache_len..]` would panic, \ + `prefill_window` collapses `seqlen_offsets` to one scalar so the model's per-row \ + `row_q0` is fed uniform values, and `set_len(cache_len + committed_i)` counts the \ + dead prefix inside the length so it accumulates. See \ + `MtpSpeculativePipeline::step_carries_per_sequence_lengths`." + .to_string(), + ) +} + /// Greedy argmax over a 1-D logits tensor — returns the token ID. fn argmax_token(logits: &Tensor) -> Result { // logits shape [vocab] @@ -4253,6 +4343,91 @@ mod tests { ); } + /// The **mask** capability, which this wave moves. A V4 target passes it + /// now (`dsv4_attention` takes a per-row `row_q0`); a plain-`Normal` target + /// still does not, and must say so rather than inherit V4's answer. + /// + /// Paired with the refusal that remains, so the pair reads as one + /// statement: the mask is no longer the blocker, and something else is. + #[test] + fn the_mask_capability_passes_for_v4_and_still_refuses_for_everyone_else() { + use crate::kv_cache::{EitherCache, KvCache, NormalCache, XsRollingCache}; + use std::sync::{Arc, Mutex}; + + let v4_shaped = EitherCache::Normal(Arc::new(Mutex::new(NormalCache(vec![ + KvCache::new_normal(2, 4096, 512), + KvCache::XsRolling(Box::new(XsRollingCache::new(4, 2, 64, 2048))), + ])))); + assert!( + model_masks_ragged_batches(&v4_shaped).is_ok(), + "a V4 target threads the per-row `row_q0` into `dsv4_attention`, so the ragged mask \ + is no longer what blocks it" + ); + + let plain = EitherCache::Normal(Arc::new(Mutex::new(NormalCache(vec![ + KvCache::new_normal(2, 4096, 512), + KvCache::new_normal(2, 4096, 512), + ])))); + let err = model_masks_ragged_batches(&plain).unwrap_err(); + assert!( + err.contains("XsRolling") && err.contains("dead"), + "a non-V4 target must still refuse, and name why; got {err}" + ); + } + + /// 🔴 **The refusal moved; it did not disappear.** With both of the + /// previously-named blockers cleared — the cache carries per-row lengths + /// (wave63-CO) and V4's attention masks a ragged cohort (this wave) — a V4 + /// target must STILL be refused, and the reason must now name the fused + /// step's shared `cache_len`. + /// + /// This is the test that stops the third check from being quietly dropped. + /// Granting the mode on the first two alone is not merely useless, it is + /// harmful: `window_ok` fails on every ragged step, the fused path defers + /// to the target's own decode, and that path builds no mask at all at + /// `t_q == 1` — so the dead prefix would be attended as real keys by the + /// fallback that was supposed to be the safe one. + #[test] + fn a_fully_capable_v4_cache_is_still_refused_and_now_names_the_step() { + use crate::kv_cache::{EitherCache, KvCache, NormalCache, XsRollingCache}; + use std::sync::{Arc, Mutex}; + + let v4 = EitherCache::Normal(Arc::new(Mutex::new(NormalCache(vec![ + KvCache::new_normal(2, 4096, 512), + KvCache::XsRolling(Box::new(XsRollingCache::new(4, 2, 64, 2048))), + ])))); + + // Force the cache capability ON, thread-locally, so checks (1) and (2) + // both pass and check (3) is the only thing left standing. + let refusal = crate::kv_cache::xs_per_seq_test_override::with(true, || { + assert!( + cache_supports_per_sequence_advance(&v4).is_ok(), + "precondition: with ARC_V4_XS_PER_SEQ on, a V4 cache carries per-row lengths" + ); + assert!( + model_masks_ragged_batches(&v4).is_ok(), + "precondition: V4's attention now masks a ragged cohort" + ); + per_sequence_refusal(&v4) + }); + + let why = refusal.expect( + "per-sequence advance must STILL be refused: the fused step keys on one shared \ + `cache_len`, and granting the mode without fixing that defers every ragged step \ + into a decode path that builds no mask at all", + ); + assert!( + why.contains("cache_len") && why.contains("window_ok"), + "the refusal must name the fused step's shared `cache_len`, not a blocker that is \ + already cleared; got {why}" + ); + assert!( + !why.contains("XsRolling"), + "the refusal still names the compressor slot, so the flag is not reaching the \ + cache and this wave changed nothing; got {why}" + ); + } + /// The plan functions agree on the one thing they must: the slots a /// sequence keeps. They differ only in what happens to the surplus — the /// cohort rule rolls it back and re-feeds it as a longer tail next step,