From 1e6cff1161c040259d07be7b9191b3f354e1f05c Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick Date: Fri, 21 Aug 2026 01:41:59 +0100 Subject: [PATCH 1/6] =?UTF-8?q?perf(ArcAttention):=20batch=20V4=20RoPE=20o?= =?UTF-8?q?ver=20the=20cohort=20=E2=80=94=2099,072=20launches/step=20->=20?= =?UTF-8?q?301?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `DeepSeekV2RotaryEmbedding::{forward,forward_inverse_tail}` gated the batched path on `seqlen_offsets.len() == 1` — the LENGTH of the offset vector, where the thing that matters is the DISTINCTNESS of its values. `len()` is the batch size, so the batched path was unreachable at every batch size above one. The offsets are uniform anyway. `scheduler::default_scheduler`'s `select_running_bucket` admits a forward pass only when cache lengths are exactly equal, and the one producer of a ragged dense batch is gated behind `ARC_MTP_PER_SEQ_KV`, which defaults off (`deepseek4::ragged_row_q0`). So the loop was performing B bit-identical recomputations and concatenating them back. Measured, B=256 pure decode, steady state (nsys, 204.80 tok/s): kernel share launches/step per (layer, seq) ucopy_bf16 10.2% 22,794 2 (2*43*256 = 22,016) copy2d_bf16 5.0% 33,610 3 (3*43*256 = 33,024) rope_i_bf16 3.9% 33,054 3 (3*43*256 = 33,024) uneg_bf16 0.8% 11,008 1 (1*43*256 = 11,008, exact) 19.9% of step time, and the only cost in the profile that gets WORSE as batch grows. `copy2d` is the 256-argument `Tensor::cat` that exists solely to undo the split the same function just made; it goes to zero. `uneg` is 11,008 launches to negate a [1,32] slice of a CONSTANT table, so `neg_sin` is now materialised once at construction — not a kernel to fuse, a kernel to delete. Cohort form is 43 layers * 7 = 301 launches (129 of them rope_i), which lands RoPE at the same order as fp8_matmul_tiled (301) and qtip2b_grouped_gemm (129) instead of dwarfing them. PREDICTION, not a result: this is a launch-count fix, the bytes moved are unchanged, and how much of the 19.9% actually returns must be measured on hardware. Ragged offsets still take the loop verbatim, so enabling ARC_MTP_PER_SEQ_KV behaves exactly as today. No ragged index_select path is added — nothing ships dark. Bit-identical by construction: with 2-D [T, D/2] cos/sin the kernel's stride_b is 0 (candle-nn/src/rotary_emb.rs), so every batch row reads the same cos/sin row and each output element is an independent two-multiply-one-add. Seven tests hold the pre-fix loop as a verbatim oracle and assert raw IEEE-754 bit equality (not a tolerance) for uniform AND ragged offsets, at seq_len 1 and 2, with a one-ULP negative control proving the comparator can fail. Engagement is counted (`rope_cohort_stats::counts`): a fast path that silently declined would pass every equality assertion trivially. The counter already earned itself — it caught a missing lock in one of these tests. Co-Authored-By: Claude Opus 5 (1M context) --- mistralrs-core/src/layers.rs | 351 +++++++++++++++++++++++++++++++++-- 1 file changed, 337 insertions(+), 14 deletions(-) diff --git a/mistralrs-core/src/layers.rs b/mistralrs-core/src/layers.rs index 71701e852..82202d508 100644 --- a/mistralrs-core/src/layers.rs +++ b/mistralrs-core/src/layers.rs @@ -375,12 +375,27 @@ impl QRmsNorm { /// `None` keeps the old loop verbatim, so ragged cohorts behave exactly as they /// do today. This is a dispatch fix, not a new ragged path. /// -/// # Relationship to `DeepSeekV2RotaryEmbedding::uniform_offset` +/// # The V4 rotary, collapsed onto this predicate /// -/// The V4 rotary got this same correction separately (PR #198, still unmerged -/// at the time of writing) as a private associated function with an identical -/// body. When these two lines of work meet, collapse that one onto this one -- -/// there must not be two spellings of the same predicate. +/// The V4 rotary (`DeepSeekV2RotaryEmbedding`) got this same correction +/// separately, in #198, as a private associated function `uniform_offset` with +/// a byte-identical body. #198 is now merged and that second spelling is gone: +/// `DeepSeekV2RotaryEmbedding::{forward, forward_inverse_tail}` call this +/// function directly, so the two cannot drift apart. +/// +/// Two details #198 established that this doc did not previously carry: +/// +/// * The uniformity is not an accident of the current scheduler, it is +/// *enforced* by it. `scheduler::default_scheduler::select_running_bucket` +/// admits a forward pass only when cache lengths are exactly equal +/// ("Sequences may only share a forward pass when their cache lengths are +/// exactly equal"), and the single producer of a ragged dense batch is gated +/// behind `ARC_MTP_PER_SEQ_KV`, which defaults off +/// (`models::deepseek4::ragged_row_q0`). +/// * Measured at B=256 on a real decode profile, the per-sequence loop cost +/// ~100,000 GPU launches per step, 19.9% of step time -- of which the +/// `copy2d` half existed only to undo a split the same function had just +/// made. That number was measured on #198's branch, not on this tree. #[inline] pub(crate) fn uniform_seqlen_offset(seqlen_offsets: &[usize]) -> Option { match seqlen_offsets { @@ -1540,8 +1555,28 @@ impl Qwen2_5VLRotaryEmbedding { pub struct DeepSeekV2RotaryEmbedding { sin: Tensor, cos: Tensor, + /// `-sin`, materialised once at construction. + /// + /// [`Self::forward_inverse_tail`] rotates by the conjugate `(cos, -sin)`. + /// It used to spell that as `self.sin.narrow(..)?.neg()?` *inside* the + /// per-sequence loop, which at B=256 × 43 layers launched `uneg_bf16` + /// **11,008 times per decode step** to negate a `[1, 32]` slice of a + /// **constant** table — 0.8% of a measured B=256 step spent re-deriving a + /// value that never changes. + /// + /// Narrowing a pre-negated table is bit-identical to negating a narrowed + /// one: negation is a sign-bit flip, exact in every float format and + /// elementwise, so it commutes with a view. See `neg_sin_commutes_with_narrow`. + neg_sin: Tensor, } +// The engagement counters for this path are the ones declared beside +// `uniform_seqlen_offset` (`rope_cohort_stats`, above). #198 carried a second +// module of the same name, with the same two atomics and a `counts()` of its +// own; two `pub mod rope_cohort_stats` in one file is a hard name collision, +// and two counters for one claim would let a harness read `cohort > 0` off the +// Phi path while this one never engaged. Collapsed onto the shared module, as +// `uniform_seqlen_offset`'s doc comment asks. #[derive(Debug, Clone, Deserialize, Serialize)] #[serde(untagged)] pub enum DeepSeekV2RopeScaling { @@ -1583,6 +1618,19 @@ pub struct DeepSeekV2RopeConfig { } impl DeepSeekV2RotaryEmbedding { + /// The single place the tables are stored, so `neg_sin` cannot be derived + /// in one constructor and forgotten in the other. + fn from_tables(sin: Tensor, cos: Tensor) -> Result { + let neg_sin = sin.neg()?; + Ok(Self { sin, cos, neg_sin }) + } + + // The cohort predicate for this rotary is the free function + // `uniform_seqlen_offset`, not a method here. #198 introduced + // `Self::uniform_offset` with a byte-identical body; per that function's + // doc comment ("there must not be two spellings of the same predicate") + // the duplicate is collapsed away, and the reasoning it carried has moved + // onto the survivor. fn new_unscaled(cfg: &DeepSeekV2RopeConfig, dtype: DType, dev: &Device) -> Result { let max_seq_len = cfg.max_position_embeddings; let dim = cfg.qk_rope_head_dim; @@ -1601,7 +1649,7 @@ impl DeepSeekV2RotaryEmbedding { let sin = freqs.sin()?.to_dtype(dtype)?; let cos = freqs.cos()?.to_dtype(dtype)?; - Ok(Self { sin, cos }) + Self::from_tables(sin, cos) } fn yarn_find_correction_dim( @@ -1693,7 +1741,7 @@ impl DeepSeekV2RotaryEmbedding { let sin = (freqs.sin()? * mscale as f64)?.to_dtype(dtype)?; let cos = (freqs.cos()? * mscale as f64)?.to_dtype(dtype)?; - Ok(Self { sin, cos }) + Self::from_tables(sin, cos) } pub fn new(cfg: &DeepSeekV2RopeConfig, dtype: DType, dev: &Device) -> Result { @@ -1746,13 +1794,26 @@ impl DeepSeekV2RotaryEmbedding { ) -> Result<(Tensor, Tensor)> { let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; - if seqlen_offsets.len() == 1 { - let cos = self.cos.narrow(0, seqlen_offsets[0], seq_len)?; - let sin = self.sin.narrow(0, seqlen_offsets[0], seq_len)?; + // One `rope_i` over the whole cohort whenever every row rotates at the + // same position — which is every request today. See + // [`uniform_seqlen_offset`] for why this used to read `len() == 1`, and + // what that cost. + // + // Bit-identical to the loop below, not merely close: with a 2-D + // `[T, D/2]` cos/sin the kernel's `stride_b` is 0 + // (`candle-nn/src/rotary_emb.rs`), so every batch row reads the same + // cos/sin row, and each output element is an independent + // two-multiply-one-add of its own inputs. Batching changes which + // launch computes an element, never the arithmetic that produces it. + if let Some(offset) = uniform_seqlen_offset(seqlen_offsets) { + rope_cohort_stats::record_cohort(); + let cos = self.cos.narrow(0, offset, seq_len)?; + let sin = self.sin.narrow(0, offset, seq_len)?; let q_embed = candle_nn::rotary_emb::rope_i(&q.contiguous()?, &cos, &sin)?; let k_embed = candle_nn::rotary_emb::rope_i(&k.contiguous()?, &cos, &sin)?; Ok((q_embed, k_embed)) } else { + rope_cohort_stats::record_per_sequence(); let mut q_embeds = Vec::new(); let mut k_embeds = Vec::new(); for (i, offset) in seqlen_offsets.iter().enumerate() { @@ -1789,16 +1850,21 @@ impl DeepSeekV2RotaryEmbedding { let (_b, _h, seq_len, head_dim) = x.dims4()?; let nope = head_dim - rope_dim; let x_nope = x.narrow(3, 0, nope)?; - let rotated = if seqlen_offsets.len() == 1 { - let cos = self.cos.narrow(0, seqlen_offsets[0], seq_len)?; - let sin = self.sin.narrow(0, seqlen_offsets[0], seq_len)?.neg()?; + // Same cohort dispatch as `forward`, and the conjugate `sin` now comes + // from the pre-negated table (see the `neg_sin` field) rather than a + // `neg` launch per sequence per layer. + let rotated = if let Some(offset) = uniform_seqlen_offset(seqlen_offsets) { + rope_cohort_stats::record_cohort(); + let cos = self.cos.narrow(0, offset, seq_len)?; + let sin = self.neg_sin.narrow(0, offset, seq_len)?; let x_pe = x.narrow(3, nope, rope_dim)?.contiguous()?; candle_nn::rotary_emb::rope_i(&x_pe, &cos, &sin)? } else { + rope_cohort_stats::record_per_sequence(); let mut outs = Vec::new(); for (i, offset) in seqlen_offsets.iter().enumerate() { let cos = self.cos.narrow(0, *offset, seq_len)?; - let sin = self.sin.narrow(0, *offset, seq_len)?.neg()?; + let sin = self.neg_sin.narrow(0, *offset, seq_len)?; let x_pe = x .i(i)? .unsqueeze(0)? @@ -4243,3 +4309,260 @@ mod rope_cohort_tests { } } } + +/// Bit-exactness and engagement tests for the DeepSeek V2/V4 RoPE cohort path. +/// +/// The change under test replaced `seqlen_offsets.len() == 1` — a test of the +/// *length* of the offset vector — with a test of the *distinctness of its +/// values*. At B=256 the length is 256, so the batched path was unreachable at +/// every batch size above one, while the values were uniform anyway; the +/// per-sequence loop was therefore performing B bit-identical recomputations +/// and concatenating them back together (~100,000 GPU launches per decode step +/// on a measured profile, 19.9% of step time). +/// +/// Each test below holds the old loop as a verbatim oracle and asserts the new +/// dispatch reproduces it **bit for bit**, not within a tolerance — a +/// tolerance-based assertion here would be vacuous, since the claim is that no +/// arithmetic changed at all. +#[cfg(test)] +mod deepseek_rope_cohort_tests { + use super::*; + use std::sync::{Mutex, MutexGuard, OnceLock}; + + const MAX_POS: usize = 32; + const ROPE_DIM: usize = 8; + const HEAD_DIM: usize = 12; + const N_HEADS: usize = 3; + const BATCH: usize = 4; + + /// The engagement counters are process-global, so every test that drives a + /// forward takes this lock and the deltas it observes are its own. + fn guard() -> MutexGuard<'static, ()> { + static LOCK: OnceLock> = OnceLock::new(); + LOCK.get_or_init(|| Mutex::new(())) + .lock() + .unwrap_or_else(|e| e.into_inner()) + } + + fn rope() -> DeepSeekV2RotaryEmbedding { + let cfg = DeepSeekV2RopeConfig { + rope_scaling: None, + max_position_embeddings: MAX_POS, + rope_theta: 10000.0, + qk_rope_head_dim: ROPE_DIM, + }; + DeepSeekV2RotaryEmbedding::new(&cfg, DType::F32, &Device::Cpu).unwrap() + } + + /// Deterministic, non-degenerate input. Both paths consume the same tensor, + /// so randomness would be sound too, but a fixed pattern makes a failure + /// reproducible. + fn xs(b: usize, h: usize, t: usize, d: usize) -> Tensor { + let n = b * h * t * d; + let v: Vec = (0..n) + .map(|i| ((i as f32) * 0.37).sin() * 1.7 + 0.11) + .collect(); + Tensor::from_vec(v, (b, h, t, d), &Device::Cpu).unwrap() + } + + /// Raw IEEE-754 bit patterns. Float equality would accept `-0.0 == 0.0` and + /// silently pass a sign error in the conjugate `sin`; bit patterns do not. + fn bits(t: &Tensor) -> Vec { + t.flatten_all() + .unwrap() + .to_vec1::() + .unwrap() + .into_iter() + .map(f32::to_bits) + .collect() + } + + /// The pre-fix `forward`, verbatim, as the oracle. Always loops, whatever + /// the offsets are. + fn reference_forward( + re: &DeepSeekV2RotaryEmbedding, + q: &Tensor, + k: &Tensor, + offsets: &[usize], + ) -> Result<(Tensor, Tensor)> { + let (_b, _h, seq_len, _n) = q.dims4()?; + let mut q_embeds = Vec::new(); + let mut k_embeds = Vec::new(); + for (i, offset) in offsets.iter().enumerate() { + let cos = re.cos.narrow(0, *offset, seq_len)?; + let sin = re.sin.narrow(0, *offset, seq_len)?; + q_embeds.push(candle_nn::rotary_emb::rope_i( + &q.i(i)?.unsqueeze(0)?.contiguous()?, + &cos, + &sin, + )?); + k_embeds.push(candle_nn::rotary_emb::rope_i( + &k.i(i)?.unsqueeze(0)?.contiguous()?, + &cos, + &sin, + )?); + } + Ok((Tensor::cat(&q_embeds, 0)?, Tensor::cat(&k_embeds, 0)?)) + } + + /// The pre-fix `forward_inverse_tail`, verbatim — including the old + /// `sin.narrow(..).neg()` spelling, so this oracle also pins the `neg_sin` + /// precomputation. + fn reference_inverse_tail( + re: &DeepSeekV2RotaryEmbedding, + x: &Tensor, + rope_dim: usize, + offsets: &[usize], + ) -> Result { + let (_b, _h, seq_len, head_dim) = x.dims4()?; + let nope = head_dim - rope_dim; + let x_nope = x.narrow(3, 0, nope)?; + let mut outs = Vec::new(); + for (i, offset) in offsets.iter().enumerate() { + let cos = re.cos.narrow(0, *offset, seq_len)?; + let sin = re.sin.narrow(0, *offset, seq_len)?.neg()?; + let x_pe = x + .i(i)? + .unsqueeze(0)? + .narrow(3, nope, rope_dim)? + .contiguous()?; + outs.push(candle_nn::rotary_emb::rope_i(&x_pe, &cos, &sin)?); + } + let rotated = Tensor::cat(&outs, 0)?; + Tensor::cat(&[&x_nope, &rotated], 3)?.contiguous() + } + + // The predicate itself is tested once, in `rope_cohort_tests` + // (`uniform_offset_tests_values_not_length`), against the same six cases. + // #198 carried a byte-identical copy of that test aimed at + // `DeepSeekV2RotaryEmbedding::uniform_offset`; with the two spellings + // collapsed onto `uniform_seqlen_offset` there is one predicate, so a + // second copy would assert nothing the first does not. What is NOT + // duplicated, and stays below, is the V4-specific behaviour: `neg_sin`, + // `forward_inverse_tail`, and the bit-exactness of this rotary's own + // cohort dispatch. + #[test] + fn neg_sin_commutes_with_narrow() { + let re = rope(); + for offset in [0usize, 1, 17] { + let precomputed = re.neg_sin.narrow(0, offset, 3).unwrap(); + let on_the_fly = re.sin.narrow(0, offset, 3).unwrap().neg().unwrap(); + assert_eq!( + bits(&precomputed), + bits(&on_the_fly), + "pre-negating the table must be bit-identical at offset {offset}" + ); + } + } + + #[test] + fn cohort_forward_is_bit_identical_to_the_per_sequence_loop() { + let _g = guard(); + let re = rope(); + // seq_len 1 is decode; seq_len 2 covers uniform-offset chunked prefill. + for seq_len in [1usize, 2] { + let q = xs(BATCH, N_HEADS, seq_len, ROPE_DIM); + let k = xs(BATCH, 1, seq_len, ROPE_DIM); + let offsets = vec![5usize; BATCH]; + + let before = rope_cohort_stats::counts(); + let (q_new, k_new) = re.forward(&q, &k, &offsets).unwrap(); + let after = rope_cohort_stats::counts(); + let (q_ref, k_ref) = reference_forward(&re, &q, &k, &offsets).unwrap(); + + assert_eq!(q_new.dims(), q_ref.dims()); + assert_eq!(bits(&q_new), bits(&q_ref), "Q differs at seq_len={seq_len}"); + assert_eq!(bits(&k_new), bits(&k_ref), "K differs at seq_len={seq_len}"); + // Prove the batched branch is the one that ran: a fast path that + // silently declined would pass the equality above trivially. + assert_eq!( + (after.0 - before.0, after.1 - before.1), + (1, 0), + "expected exactly one cohort call and no per-sequence loop" + ); + } + } + + #[test] + fn ragged_forward_still_takes_the_loop_and_still_matches() { + let _g = guard(); + let re = rope(); + let q = xs(BATCH, N_HEADS, 1, ROPE_DIM); + let k = xs(BATCH, 1, 1, ROPE_DIM); + let offsets = vec![5usize, 6, 5, 9]; + + let before = rope_cohort_stats::counts(); + let (q_new, k_new) = re.forward(&q, &k, &offsets).unwrap(); + let after = rope_cohort_stats::counts(); + let (q_ref, k_ref) = reference_forward(&re, &q, &k, &offsets).unwrap(); + + assert_eq!(bits(&q_new), bits(&q_ref)); + assert_eq!(bits(&k_new), bits(&k_ref)); + assert_eq!( + (after.0 - before.0, after.1 - before.1), + (0, 1), + "ragged rows must still take the per-sequence loop" + ); + } + + #[test] + fn cohort_inverse_tail_is_bit_identical_to_the_per_sequence_loop() { + let _g = guard(); + let re = rope(); + for seq_len in [1usize, 2] { + let x = xs(BATCH, N_HEADS, seq_len, HEAD_DIM); + let offsets = vec![5usize; BATCH]; + + let before = rope_cohort_stats::counts(); + let new = re.forward_inverse_tail(&x, ROPE_DIM, &offsets).unwrap(); + let after = rope_cohort_stats::counts(); + let reference = reference_inverse_tail(&re, &x, ROPE_DIM, &offsets).unwrap(); + + assert_eq!(new.dims(), reference.dims()); + assert_eq!( + bits(&new), + bits(&reference), + "inverse tail differs at seq_len={seq_len}" + ); + assert_eq!((after.0 - before.0, after.1 - before.1), (1, 0)); + } + } + + #[test] + fn ragged_inverse_tail_still_takes_the_loop_and_still_matches() { + let _g = guard(); + let re = rope(); + let x = xs(BATCH, N_HEADS, 1, HEAD_DIM); + let offsets = vec![5usize, 6, 5, 9]; + + let before = rope_cohort_stats::counts(); + let new = re.forward_inverse_tail(&x, ROPE_DIM, &offsets).unwrap(); + let after = rope_cohort_stats::counts(); + let reference = reference_inverse_tail(&re, &x, ROPE_DIM, &offsets).unwrap(); + + assert_eq!(bits(&new), bits(&reference)); + assert_eq!((after.0 - before.0, after.1 - before.1), (0, 1)); + } + + /// The comparator must be able to fail. Every bit-equality assertion above + /// is worthless if `bits` had degenerated into something that always + /// matches, so poison one element by a single ULP and require it flagged. + #[test] + fn the_bit_comparator_is_live() { + // Drives a forward, so it must serialise with the counter assertions. + let _g = guard(); + let re = rope(); + let x = xs(BATCH, N_HEADS, 1, HEAD_DIM); + let offsets = vec![5usize; BATCH]; + let good = re.forward_inverse_tail(&x, ROPE_DIM, &offsets).unwrap(); + + let mut poisoned = bits(&good); + poisoned[0] += 1; // one ULP + assert_ne!( + bits(&good), + poisoned, + "the comparator accepted a one-ULP difference; every other \ + assertion in this module would be vacuous" + ); + } +} From 44f3c52b0364e53550ef908e2c0a23a0b4b98486 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick Date: Fri, 21 Aug 2026 01:42:11 +0100 Subject: [PATCH 2/6] perf(ArcAttention): un-deaden the fused qk_norm_rope kernel above batch 1 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `fused_qk_gate` declined with "per-sequence position offsets" whenever `seqlen_offsets.len() != 1` — the same length-vs-values confusion as the RoPE dispatch in the parent commit, and with the same consequence: the fused one-launch kernel was dead at EVERY batch size above one, i.e. in exactly the serving regime it was written for. Its own doc comment advertises "replacing 16 candle launches per layer with 1"; at B=256 it replaced nothing. The kernel is already batch-correct and needed no change. `qk_norm_rope.cu` flattens (b, t) into blockIdx.y, recovers `b = bt / seq_len`, and reads the table at `pos_offset + t` — a position independent of `b`, which is exactly right when the rows share an offset. Its grid is (n_heads + 1, batch * seq_len) and the wrapper sizes buffers as `batch * ...`. Nothing assumed batch == 1. Gate now takes the uniform offset via the shared `DeepSeekV2RotaryEmbedding::uniform_offset` helper, so both call sites cannot drift apart. Genuinely ragged rows still decline and still take the eager chain. Split from the parent commit so it can be reverted independently: this one swaps in a CUDA kernel that has never run above batch 1, and unlike the parent it is not bit-identical-by-construction — it is bit-identical by the kernel's own design, which `ARC_QK_VERIFY=1` checks against the eager chain at every layer. Verify before trusting; `ARC_QK_FUSED=0` restores the eager path from the same binary. Co-Authored-By: Claude Opus 5 (1M context) --- mistralrs-core/src/models/deepseek4.rs | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/mistralrs-core/src/models/deepseek4.rs b/mistralrs-core/src/models/deepseek4.rs index 766a5f59c..cceec48ae 100644 --- a/mistralrs-core/src/models/deepseek4.rs +++ b/mistralrs-core/src/models/deepseek4.rs @@ -1287,9 +1287,21 @@ impl Attention { if cos.dtype() != candle_core::DType::BF16 || sin.dtype() != candle_core::DType::BF16 { decline!("RoPE tables are not BF16"); } - if seqlen_offsets.len() != 1 { + // The kernel needs ONE `pos_offset`, not a batch of one row. It reads + // the table at `pos_offset + t` where `t = blockIdx.y % seq_len` + // (`qk_norm_rope.cu`), a position that does not depend on `b`, and its + // grid is `(n_heads + 1, batch * seq_len)` with buffers sized + // `batch * ...` — so it is already correct at any batch size whose rows + // share an offset. This used to read `seqlen_offsets.len() != 1`, which + // tested the length of the vector rather than the distinctness of its + // values and therefore declined at EVERY batch size above one, leaving + // the fused path dead in exactly the serving regime it was written for. + // Rows that genuinely differ still decline, and still take the eager + // chain. Prove the swap with `ARC_QK_VERIFY=1`, which bit-compares + // against the eager output at every layer. + let Some(uniform_offset) = crate::layers::uniform_seqlen_offset(seqlen_offsets) else { decline!("per-sequence position offsets"); - } + }; if self.num_kv_heads != 1 { decline!("more than one KV head"); } @@ -1322,7 +1334,7 @@ impl Attention { Ok((kb, kt, kh)) if kb == qb && kt == qt && kh == head_dim => {} _ => decline!("kv_norm output shape does not match [B, T, head_dim]"), } - let pos_offset = seqlen_offsets[0]; + let pos_offset = uniform_offset; match cos.dims2() { Ok((rows, cols)) if cols == rope_dim / 2 && rows >= pos_offset + qt => {} _ => decline!("RoPE table does not cover [pos_offset, pos_offset + T)"), From e6d28562fc315938801e84a3cd45eec0cc903fe5 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 17:21:07 +0100 Subject: [PATCH 3/6] fix(ArcAttention): one cohort predicate, one counter pair, per-thread test counts Integration fix for #198 landing after #204. Both PRs made the same correction -- `seqlen_offsets.len() == 1` tests the length of the vector where the property needed is the distinctness of its values -- and each shipped its own spelling of the predicate and its own engagement counters. `uniform_seqlen_offset`'s doc comment on the integration branch asked for exactly this collapse ("there must not be two spellings of the same predicate"); this is that collapse. 1. `pub mod rope_cohort_stats` was declared TWICE in one file (layers.rs:400 from #204, :1568 from #198). That is a hard name collision -- the file did not compile. #204's is the superset (it has `record_cohort` / `record_per_sequence` as well as `counts`), so #198's is the one that goes. 2. `DeepSeekV2RotaryEmbedding::uniform_offset` had a byte-identical body to `uniform_seqlen_offset`. Removed; `forward`, `forward_inverse_tail` and `deepseek4::fused_qk_gate` now call the free function, so the two cannot drift apart. The reasoning #198's doc comment carried and #204's did not -- that uniformity is *enforced* by `select_running_bucket` rather than merely observed, and the measured B=256 cost -- moves onto the survivor. 3. Collapsing the counters made them shared, and shared process-global counters cannot be delta-asserted from a parallel test binary. `deepseek_rope_cohort_ tests` held a private mutex, `rope_cohort_tests` held none (until now it was the only writer), and three tests duly failed with `(2, 1)` where they wanted `(1, 0)`. A shared mutex was the first fix and it is also wrong: it obliges every present and future test that touches ANY rotary to know about a lock in another module. Two such tests already exist and did not -- `deepseek4::tests::{standard_layer_rope_is_unscaled_and_compressed_is_not, mtp_block_takes_the_standard_unscaled_rope_table}` both drive `DeepSeekV2RotaryEmbedding::forward` -- and they race the cohort tests. That was caught by merging #198 and #121 together and running the full suite: `--test-threads=1` passes 707, the parallel run fails one. So the counters now also keep a `#[cfg(test)]` per-thread mirror, and the tests assert on `local_counts()`. Each test's delta is its own by construction, no cooperation from other modules is required, and the serving instrument -- the global `counts()` -- is byte-for-byte unchanged. Also drops #198's copy of `uniform_offset_tests_values_not_length`: with one predicate it asserted nothing the surviving copy does not, against the same six cases. `cargo test -p mistralrs-core --lib`: 699 passed, 0 failed; 707 passed on the merged #198+#181+#130+#121 tree. The 99,072 -> 301 launch figure remains #198's own branch measurement and is NOT re-measured here. Co-Authored-By: Claude Opus 5 (1M context) --- mistralrs-core/src/layers.rs | 100 ++++++++++++++++++++++++++--------- 1 file changed, 74 insertions(+), 26 deletions(-) diff --git a/mistralrs-core/src/layers.rs b/mistralrs-core/src/layers.rs index 82202d508..628e7ee20 100644 --- a/mistralrs-core/src/layers.rs +++ b/mistralrs-core/src/layers.rs @@ -433,11 +433,64 @@ pub mod rope_cohort_stats { #[inline] pub(super) fn record_cohort() { COHORT.fetch_add(1, Ordering::Relaxed); + #[cfg(test)] + bump_local(1, 0); } #[inline] pub(super) fn record_per_sequence() { PER_SEQUENCE.fetch_add(1, Ordering::Relaxed); + #[cfg(test)] + bump_local(0, 1); + } + + // ---- Test-only per-thread mirror of the two counters above. ----------- + // + // The atomics are process-global and, since #198 collapsed its own private + // copy of this module onto this one, they are written by EVERY rotary in + // this file -- `RotaryEmbedding`, `Llama3RotaryEmbedding`, + // `Phi4MMRotaryEmbedding` and `DeepSeekV2RotaryEmbedding` alike. A test + // that brackets one call with `counts()` is therefore not measuring its own + // call: `cargo test` runs the suite on many threads, and any other test + // that drives any rotary lands between the `before` and the `after`. + // + // That is not hypothetical. Two rotary forwards in `deepseek4.rs` + // (`standard_layer_rope_is_unscaled_and_compressed_is_not`, + // `mtp_block_takes_the_standard_unscaled_rope_table`) race the cohort tests + // exactly this way -- caught by merging #198 and #121 together and running + // the full suite, where `--test-threads=1` passes 707 and the parallel run + // fails one. + // + // A mutex was the first fix and it is the wrong one: it obliges every + // present and future test that touches any rotary to know about a lock in + // another module, which nothing enforces and the next such test will + // forget. Counting per-thread instead makes each test's delta its own by + // construction, needs no cooperation from anyone else, and leaves the + // serving instrument -- the global `counts()` -- byte-for-byte unchanged. + + #[cfg(test)] + thread_local! { + static LOCAL: std::cell::Cell<(u64, u64)> = const { std::cell::Cell::new((0, 0)) }; + } + + #[cfg(test)] + #[inline] + fn bump_local(cohort: u64, per_sequence: u64) { + LOCAL.with(|c| { + let (a, b) = c.get(); + c.set((a + cohort, b + per_sequence)); + }); + } + + /// `(cohort, per_sequence)` for THIS thread only — the test-facing + /// counterpart of [`counts`]. + /// + /// Same monotonic-delta contract, same meaning, but immune to whatever the + /// rest of the suite is doing on other threads. Tests assert on deltas + /// around their own call; nothing resets it. + #[cfg(test)] + pub fn local_counts() -> (u64, u64) { + LOCAL.with(|c| c.get()) } } @@ -4176,9 +4229,9 @@ mod rope_cohort_tests { let k = xs(&[B, H, T, D_HEAD], 0.019); let offsets = vec![5usize; B]; - let before = rope_cohort_stats::counts(); + let before = rope_cohort_stats::local_counts(); let (qb, kb) = re.forward(&q, &k, &offsets).unwrap(); - let after = rope_cohort_stats::counts(); + let after = rope_cohort_stats::local_counts(); // A green result must prove work happened: the batched arm has to have // been the one that ran. @@ -4209,9 +4262,9 @@ mod rope_cohort_tests { let k = xs(&[3, H, T, D_HEAD], 0.019); let offsets = vec![0usize, 4, 9]; - let before = rope_cohort_stats::counts(); + let before = rope_cohort_stats::local_counts(); let (qb, kb) = re.forward(&q, &k, &offsets).unwrap(); - let after = rope_cohort_stats::counts(); + let after = rope_cohort_stats::local_counts(); assert!( after.1 > before.1, "ragged cohort must NOT take the batched path" @@ -4327,7 +4380,6 @@ mod rope_cohort_tests { #[cfg(test)] mod deepseek_rope_cohort_tests { use super::*; - use std::sync::{Mutex, MutexGuard, OnceLock}; const MAX_POS: usize = 32; const ROPE_DIM: usize = 8; @@ -4335,14 +4387,15 @@ mod deepseek_rope_cohort_tests { const N_HEADS: usize = 3; const BATCH: usize = 4; - /// The engagement counters are process-global, so every test that drives a - /// forward takes this lock and the deltas it observes are its own. - fn guard() -> MutexGuard<'static, ()> { - static LOCK: OnceLock> = OnceLock::new(); - LOCK.get_or_init(|| Mutex::new(())) - .lock() - .unwrap_or_else(|e| e.into_inner()) - } + // Engagement is asserted through `rope_cohort_stats::local_counts()`, the + // per-thread mirror of the global counters. #198 originally serialised + // these tests on a private mutex, which worked only while this module also + // carried its own private copy of `rope_cohort_stats`; once the two counter + // modules were collapsed the atomics became shared and a private lock + // stopped serialising anything that mattered. A shared lock was the next + // attempt and it is still wrong -- it obliges unrelated tests in other + // modules to know about it. See `rope_cohort_stats` for why per-thread + // counting is the fix. fn rope() -> DeepSeekV2RotaryEmbedding { let cfg = DeepSeekV2RopeConfig { @@ -4457,7 +4510,6 @@ mod deepseek_rope_cohort_tests { #[test] fn cohort_forward_is_bit_identical_to_the_per_sequence_loop() { - let _g = guard(); let re = rope(); // seq_len 1 is decode; seq_len 2 covers uniform-offset chunked prefill. for seq_len in [1usize, 2] { @@ -4465,9 +4517,9 @@ mod deepseek_rope_cohort_tests { let k = xs(BATCH, 1, seq_len, ROPE_DIM); let offsets = vec![5usize; BATCH]; - let before = rope_cohort_stats::counts(); + let before = rope_cohort_stats::local_counts(); let (q_new, k_new) = re.forward(&q, &k, &offsets).unwrap(); - let after = rope_cohort_stats::counts(); + let after = rope_cohort_stats::local_counts(); let (q_ref, k_ref) = reference_forward(&re, &q, &k, &offsets).unwrap(); assert_eq!(q_new.dims(), q_ref.dims()); @@ -4485,15 +4537,14 @@ mod deepseek_rope_cohort_tests { #[test] fn ragged_forward_still_takes_the_loop_and_still_matches() { - let _g = guard(); let re = rope(); let q = xs(BATCH, N_HEADS, 1, ROPE_DIM); let k = xs(BATCH, 1, 1, ROPE_DIM); let offsets = vec![5usize, 6, 5, 9]; - let before = rope_cohort_stats::counts(); + let before = rope_cohort_stats::local_counts(); let (q_new, k_new) = re.forward(&q, &k, &offsets).unwrap(); - let after = rope_cohort_stats::counts(); + let after = rope_cohort_stats::local_counts(); let (q_ref, k_ref) = reference_forward(&re, &q, &k, &offsets).unwrap(); assert_eq!(bits(&q_new), bits(&q_ref)); @@ -4507,15 +4558,14 @@ mod deepseek_rope_cohort_tests { #[test] fn cohort_inverse_tail_is_bit_identical_to_the_per_sequence_loop() { - let _g = guard(); let re = rope(); for seq_len in [1usize, 2] { let x = xs(BATCH, N_HEADS, seq_len, HEAD_DIM); let offsets = vec![5usize; BATCH]; - let before = rope_cohort_stats::counts(); + let before = rope_cohort_stats::local_counts(); let new = re.forward_inverse_tail(&x, ROPE_DIM, &offsets).unwrap(); - let after = rope_cohort_stats::counts(); + let after = rope_cohort_stats::local_counts(); let reference = reference_inverse_tail(&re, &x, ROPE_DIM, &offsets).unwrap(); assert_eq!(new.dims(), reference.dims()); @@ -4530,14 +4580,13 @@ mod deepseek_rope_cohort_tests { #[test] fn ragged_inverse_tail_still_takes_the_loop_and_still_matches() { - let _g = guard(); let re = rope(); let x = xs(BATCH, N_HEADS, 1, HEAD_DIM); let offsets = vec![5usize, 6, 5, 9]; - let before = rope_cohort_stats::counts(); + let before = rope_cohort_stats::local_counts(); let new = re.forward_inverse_tail(&x, ROPE_DIM, &offsets).unwrap(); - let after = rope_cohort_stats::counts(); + let after = rope_cohort_stats::local_counts(); let reference = reference_inverse_tail(&re, &x, ROPE_DIM, &offsets).unwrap(); assert_eq!(bits(&new), bits(&reference)); @@ -4550,7 +4599,6 @@ mod deepseek_rope_cohort_tests { #[test] fn the_bit_comparator_is_live() { // Drives a forward, so it must serialise with the counter assertions. - let _g = guard(); let re = rope(); let x = xs(BATCH, N_HEADS, 1, HEAD_DIM); let offsets = vec![5usize; BATCH]; From eea014cb4155433f998ece2cff9840d8f5b3e515 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 18:17:53 +0100 Subject: [PATCH 4/6] fix(ArcAttention): the batched fused qk_norm_rope path is OPT-IN MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The parent commit un-deadens the fused Q/K kernel above batch 1 by fixing a length-vs-values gate. That is the right fix to the gate and the wrong default: it makes a CUDA kernel LIVE on a path it has never run on. This repo has paid for that twice in one week. * TCFRAG shipped default-ON carrying "UNVERIFIED ON HARDWARE — NEVER RUN" in its own header. It held 63 GB and permanently broke a layer through a poisoned `OnceLock`. #209 retired it to opt-in. * The fused-512 attention path ran for four days silently dropping the attention mask, at 12% agreement with the reference, because nobody ran the path it had quietly become live on. So the doctrine, and this commit: **unverified means default-off, and "unverified" means unmeasured, not new.** `ARC_QK_FUSED_COHORT` gates the batched path, default OFF. With it unset, `batch > 1` declines and takes the eager chain exactly as master does today — so the parent commit ships its gate fix with no behaviour change, and the cohort path is available to be measured rather than discovered in production. `batch == 1` is untouched: that path is already live and already exercised. Read by VALUE, `== Some("1")`, split into the pure `fused_cohort_enabled_from` so the polarity is testable without mutating the environment (racy across `cargo test`'s threads, `unsafe` since the 2024 edition, and meaningless against a `OnceLock` that latches the first read). Two tests pin it, and they run with no features and no GPU on every PR: unset is OFF, and `0`/`false`/ `off`/`true`/`on`/`2`/`"1 "` are all OFF. That second one is not pedantry — #212 converted 23 `ARC_*` flags this week because `var_os(..).is_some()` made `ARC_FOO=0` mean ON, which silently cancels any A/B whose control sets zero. To turn it on, once `ARC_QK_VERIFY=1` has bit-compared it against the eager chain on hardware and it is measured faster: `ARC_QK_FUSED_COHORT=1`. When that measurement exists, this default flips and the doc comment goes with it. `cargo test -p mistralrs-core --lib`: 701 passed, 0 failed. Co-Authored-By: Claude Opus 5 (1M context) --- mistralrs-core/src/cuda/qk_norm_rope.rs | 92 +++++++++++++++++++ mistralrs-core/src/models/deepseek4.rs | 10 ++ .../tests/capability_reachability.rs | 44 +++++++++ 3 files changed, 146 insertions(+) diff --git a/mistralrs-core/src/cuda/qk_norm_rope.rs b/mistralrs-core/src/cuda/qk_norm_rope.rs index 68314ec23..b276020b5 100644 --- a/mistralrs-core/src/cuda/qk_norm_rope.rs +++ b/mistralrs-core/src/cuda/qk_norm_rope.rs @@ -49,6 +49,60 @@ pub fn fused_enabled() -> bool { *ENABLED.get_or_init(|| !matches!(std::env::var("ARC_QK_FUSED").as_deref(), Ok("0"))) } +/// `ARC_QK_FUSED_COHORT=1` lets the fused kernel serve a **batched** cohort +/// (`batch > 1`). **Default OFF**, so `batch > 1` declines and takes the eager +/// chain exactly as it does today. +/// +/// # Why this is opt-in and [`fused_enabled`] is not +/// +/// The kernel is batch-correct by construction — `qk_norm_rope.cu` flattens +/// `(b, t)` into `blockIdx.y`, recovers `b = bt / seq_len`, and reads the table +/// at `pos_offset + t`, a position independent of `b` — and until #198 the +/// dispatch gate tested `seqlen_offsets.len() != 1`, the LENGTH of the offset +/// vector where the property needed is the DISTINCTNESS of its values. So the +/// fused path was dead at every batch size above one. #198 fixes that gate. +/// +/// Fixing the gate, though, makes a kernel **live on a path it has never run +/// on**, and this repo has paid for that twice in one week: +/// +/// * TCFRAG shipped default-ON with "UNVERIFIED ON HARDWARE — NEVER RUN" in its +/// own header. It held 63 GB and permanently broke a layer through a poisoned +/// `OnceLock`; #209 made it opt-in. +/// * The fused-512 attention path ran for four days silently dropping the +/// attention mask, at 12% agreement with the reference, because nobody ran +/// the path it had quietly become live on. +/// +/// So the doctrine is: **unverified means default-off, and "unverified" means +/// unmeasured, not new.** `batch == 1` is unaffected — that path is already +/// live and already exercised. When the cohort path has been measured correct +/// and faster on hardware (`ARC_QK_VERIFY=1` bit-compares it against the eager +/// chain at every layer), this default flips and this doc comment goes with it. +/// +/// Read by VALUE, `== Some("1")`, not by presence: 23 `ARC_*` flags were +/// converted this week (#212) precisely because `var_os(..).is_some()` made +/// `ARC_FOO=0` mean ON. +pub fn fused_cohort_enabled() -> bool { + use std::sync::OnceLock; + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| { + fused_cohort_enabled_from(std::env::var("ARC_QK_FUSED_COHORT").ok().as_deref()) + }) +} + +/// Pure half of [`fused_cohort_enabled`], so the polarity is testable without +/// mutating the process environment. +/// +/// Same reasoning as `mistralrs_quant::env_flag_value` and +/// `qtip::tcfrag2b_enabled_from`: mutating the environment from a test is racy +/// across `cargo test`'s threads and, since the 2024 edition, `unsafe` — and +/// the caller latches its answer in a `OnceLock`, so such a test would prove +/// only whatever the first reader in the process happened to see. Testing the +/// pure function against every input is the only form of this test that means +/// anything. +pub fn fused_cohort_enabled_from(value: Option<&str>) -> bool { + value == Some("1") +} + /// `ARC_QK_VERIFY=1` runs the eager chain alongside the fused kernel at every /// layer and compares the raw output bytes. /// @@ -443,3 +497,41 @@ static VERIFIED: AtomicU64 = AtomicU64::new(0); pub fn verified_elements() -> u64 { VERIFIED.load(Ordering::Relaxed) } + +#[cfg(test)] +mod cohort_gate_tests { + use super::*; + + /// The batched fused path must be OFF unless explicitly asked for. + /// + /// This is the assertion that stops a repeat of TCFRAG — default-ON with + /// "UNVERIFIED ON HARDWARE — NEVER RUN" in its own header, 63 GB held, a + /// layer permanently broken through a poisoned `OnceLock`, retired by #209 + /// — and of the fused-512 attention path, which ran for four days silently + /// dropping the attention mask at 12% agreement with the reference. Both + /// were live on a path nobody had run. + /// + /// These run with no features and no GPU, on three operating systems, on + /// every PR, because the polarity is a pure function. + #[test] + fn the_batched_cohort_is_off_unless_explicitly_enabled() { + assert!( + !fused_cohort_enabled_from(None), + "unset must mean OFF — the cohort path has never run on hardware" + ); + assert!(fused_cohort_enabled_from(Some("1"))); + } + + /// Read by VALUE, not presence. #212 converted 23 `ARC_*` flags because + /// `var_os(..).is_some()` made `ARC_FOO=0` mean ON — which also silently + /// cancels any A/B whose control arm sets the flag to zero. + #[test] + fn zero_means_off_and_so_does_every_other_value() { + for v in ["0", "", "false", "no", "off", "true", "yes", "on", "2", "1 "] { + assert!( + !fused_cohort_enabled_from(Some(v)), + "{v:?} must not enable the cohort path; only the exact string \"1\" does" + ); + } + } +} diff --git a/mistralrs-core/src/models/deepseek4.rs b/mistralrs-core/src/models/deepseek4.rs index cceec48ae..5b007f92b 100644 --- a/mistralrs-core/src/models/deepseek4.rs +++ b/mistralrs-core/src/models/deepseek4.rs @@ -1330,6 +1330,16 @@ impl Attention { if qh != self.num_attention_heads * head_dim { decline!("q_proj output width is not n_heads * head_dim"); } + // The batched cohort is OPT-IN. Fixing the length-vs-values gate above + // makes this kernel live on a path it has never run on, and this repo + // has paid for exactly that twice this week (TCFRAG, and the fused-512 + // attention path that dropped the mask for four days). `batch == 1` is + // untouched — already live, already exercised. See + // `qk::fused_cohort_enabled` for the doctrine and for how to turn the + // cohort path on once it has been measured on hardware. + if qb > 1 && !qk::fused_cohort_enabled() { + decline!("batched cohort is opt-in (set ARC_QK_FUSED_COHORT=1)"); + } match kv_normed.dims3() { Ok((kb, kt, kh)) if kb == qb && kt == qt && kh == head_dim => {} _ => decline!("kv_norm output shape does not match [B, T, head_dim]"), diff --git a/mistralrs-core/tests/capability_reachability.rs b/mistralrs-core/tests/capability_reachability.rs index c203fd761..8b929cb3c 100644 --- a/mistralrs-core/tests/capability_reachability.rs +++ b/mistralrs-core/tests/capability_reachability.rs @@ -170,6 +170,50 @@ static REGISTRY: &[Capability] = &[ }, status: Status::Live, }, + // ── SHIPPED OFF, WITH A NAMED EXPERIMENT ──────────────────────────────── + // + // This entry is `Live` because the GATE is reached on every layer of every + // step — that is what this file checks, and it must not go dark. What is + // switched off is the gate's DEFAULT ANSWER at `batch > 1`, which no + // reachability check can see. + // + // Recorded here anyway, because "built, correct, tested, and switched off" + // is this repo's most expensive failure mode and this file is where someone + // goes looking for it. 26 pieces of finished work are parked that way right + // now, including a complete GPU sampler with zero callers; we shipped + // 34 tok/s of landed work while master served 17, for exactly this reason. + // + // The kernel is batch-correct by construction and the dispatch bug that + // kept it dead above batch 1 (`seqlen_offsets.len() != 1` — the LENGTH of + // the vector where the property needed is the DISTINCTNESS of its values) + // is fixed. It is off only because fixing that gate makes a CUDA kernel + // live on a path it has never run on, and this repo has paid for that twice + // in one week: TCFRAG (default-ON, "UNVERIFIED ON HARDWARE — NEVER RUN" in + // its own header, 63 GB held, a layer permanently broken, retired by #209) + // and the fused-512 attention path (four days silently dropping the + // attention mask at 12% agreement). + // + // flag ARC_QK_FUSED_COHORT=1 + // gate cuda/qk_norm_rope.rs :: fused_cohort_enabled_from + // test one binary, B>1 decode, same prompt and seed, flag set vs unset; + // ARC_QK_VERIFY=1 bit-compares fused against eager at every layer + // pass outputs bit-identical AND the fused arm faster; engagement + // counter (`qk::engaged_count`) non-zero, since a silently + // declining fast path reports a perfect no-op + // was never measured above batch 1 — the path was unreachable, so there + // is no prior number to quote, honest or otherwise + // then the default flips to ON in the same change that records the + // number. Leaving it off after a passing measurement is a failure + // state, not the finish line. + Capability { + name: "fused qk_norm_rope at B>1: built and tested, default OFF pending one measurement", + parent: "ArcInfer / ArcAttention", + check: Check::Symbol { + symbol: "fused_cohort_enabled", + defined_in: "mistralrs-core/src/cuda/qk_norm_rope.rs", + }, + status: Status::Live, + }, Capability { name: "prefill admission cap (--prefill-max-seqs)", parent: "ArcInfer / ArcSched", From 192f8ffd6b5246f8bc5566f1a69264decbf0aa23 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 20:12:07 +0100 Subject: [PATCH 5/6] =?UTF-8?q?perf(ArcAttention):=20gather=20the=20RAGGED?= =?UTF-8?q?=20cohort's=20RoPE=20=E2=80=94=20one=20rank-3=20rope=5Fi,=20opt?= =?UTF-8?q?-in=20ARC=5FROPE=5FCOHORT?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The uniform-offset fix (#198, and the shared predicate it now rides on) left one arm untouched: rows at genuinely DIFFERENT positions still take the per-sequence loop — ~9 ops x 43 layers x B per step on the measured B=256 profile (~99,000 launches, 19.9% of step time), of which the copy2d half only undoes a split the same function just made. Nothing structural required the loop: candle's rope_i accepts a rank-3 [B, T, D/2] cos/sin whenever B matches the input batch (rope_check_cs, candle-nn/src/rotary_emb.rs, pinned rev 89ab14e; the CUDA wrapper sets stride_b for rank-3, the CPU path indexes i + b_i*t*d/2). So the ragged arm becomes: gather the per-row table rows once — [B*T] u32 indices, host-built because the offsets ARE host data, memoised per instance so the H2D and the three index_selects run once per step per device (the model shares one rotary per device per table kind), CLAUDE.md pitfall #5 — then ONE rope_i per projection. ~5 launches per layer where the loop paid ~9 x B. DEFAULT OFF behind ARC_ROPE_COHORT, read by VALUE via env_flag_is_set (=0 means OFF; #212), latched once. Unset keeps the loop byte-for-byte. Registered in capability_reachability.rs beside the fused-cohort entry, same doctrine: unverified means default-off, and unverified means unmeasured. Both ragged arms are named methods driven directly by deepseek_ragged_cohort_tests — bit-pattern parity (not tolerance) for mixed offsets at B in {1,3}, T in {1,2}, ascending and descending, forward and inverse tail, plus a stale-memo test (the cache must rebuild when the cohort moves) and a pinned unset-default. engaged_count in cuda/qk_norm_rope.rs also gains its first reader: a test that proves the counters advance, closing the silent-success hole where an A/B could read a dead path as live. Parent system: ArcInfer / ArcAttention. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf --- mistralrs-core/src/cuda/qk_norm_rope.rs | 24 + mistralrs-core/src/layers.rs | 454 ++++++++++++++++-- .../tests/capability_reachability.rs | 24 + 3 files changed, 470 insertions(+), 32 deletions(-) diff --git a/mistralrs-core/src/cuda/qk_norm_rope.rs b/mistralrs-core/src/cuda/qk_norm_rope.rs index b276020b5..013d7f41b 100644 --- a/mistralrs-core/src/cuda/qk_norm_rope.rs +++ b/mistralrs-core/src/cuda/qk_norm_rope.rs @@ -534,4 +534,28 @@ mod cohort_gate_tests { ); } } + + /// [`engaged_count`] is the instrument the A/B harness reads to tell "the + /// kernel ran" apart from "the benchmark got faster" — and until this + /// test, nothing in the tree read it, so it could have silently stopped + /// advancing and every claim built on it would still look green. The + /// counters are process-global and other tests may bump them + /// concurrently, so the assertion is on a lower bound of the delta, which + /// only this thread's calls can guarantee. + #[test] + fn engaged_count_advances_when_the_paths_are_taken() { + let (e0, d0) = engaged_count(); + note_engaged(128, 64, 16, "bf16"); + note_declined("test: exercising the decline counter"); + let (e1, d1) = engaged_count(); + assert!( + e1 >= e0 + 1, + "note_engaged did not advance the engaged counter ({e0} -> {e1}); \ + an A/B reading engaged_count() would call a dead path live" + ); + assert!( + d1 >= d0 + 1, + "note_declined did not advance the declined counter ({d0} -> {d1})" + ); + } } diff --git a/mistralrs-core/src/layers.rs b/mistralrs-core/src/layers.rs index 628e7ee20..b8f6a06ae 100644 --- a/mistralrs-core/src/layers.rs +++ b/mistralrs-core/src/layers.rs @@ -494,6 +494,44 @@ pub mod rope_cohort_stats { } } +/// `ARC_ROPE_COHORT=1` routes a **ragged** decode cohort — rows at genuinely +/// different positions — through ONE gathered `rope_i` per projection instead +/// of the per-sequence loop. **Default OFF**; unset and every off-spelling +/// keep the loop byte-for-byte as it is today. +/// +/// Uniform cohorts are NOT behind this flag: when every row shares an offset, +/// [`uniform_seqlen_offset`] already dispatches one batched `rope_i` with a +/// 2-D table, which is bit-identical to the loop and needs no gather. This +/// flag governs only the arm the uniform fix could not reach — offsets that +/// genuinely differ, where the loop was the only spelling. On the measured +/// B=256 profile that loop is ~9 ops × 43 layers × B ≈ 99,000 launches per +/// step (19.9% of step time); the gathered path is ~5 launches per layer. +/// +/// candle's `rope_i` accepts a rank-3 `[B, T, D/2]` cos/sin when `B` matches +/// the input batch (`rope_check_cs`, `candle-nn/src/rotary_emb.rs`; the CUDA +/// wrapper sets `stride_b`, the CPU path indexes `i + b_i * t * d / 2`), so a +/// per-row table is expressible in one launch — the loop was never structural. +/// +/// Why opt-in rather than default-ON: the same doctrine as +/// `cuda::qk_norm_rope::fused_cohort_enabled` — unverified means default-off, +/// and unverified means unmeasured on hardware, not new. The A/B box flips +/// this flag; when the gathered arm has been measured correct and faster, the +/// default flips in the same change that records the number. +/// +/// Read by VALUE (`mistralrs_quant::env_flag_is_set`), never by presence: +/// `ARC_ROPE_COHORT=0` means OFF (#212 converted 23 flags because +/// `var_os(..).is_some()` made `=0` mean ON). Latched once per process — the +/// dispatch runs on every layer of every step and `std::env::var` takes a +/// global lock. The polarity table lives with `env_flag_value`'s own tests; +/// the two ragged arms themselves are compared bit-for-bit in +/// `deepseek_ragged_cohort_tests`, which drives them directly so the latch +/// cannot make the parity claim untestable. +pub(crate) fn rope_cohort_enabled() -> bool { + use std::sync::OnceLock; + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| mistralrs_quant::env_flag_is_set("ARC_ROPE_COHORT")) +} + /// Tile a `[max_seq_len, rot_dim]` rotary table into the `[b * seq_len, rot_dim]` /// layout `mistralrs_quant::rotary::apply_rotary_inplace` requires. /// @@ -1621,6 +1659,30 @@ pub struct DeepSeekV2RotaryEmbedding { /// one: negation is a sign-bit flip, exact in every float format and /// elementwise, so it commutes with a view. See `neg_sin_commutes_with_narrow`. neg_sin: Tensor, + /// Memoised per-row table gather for the ragged-cohort path + /// ([`rope_cohort_enabled`]). Behind an `Arc` so the per-device sharing in + /// `deepseek4.rs` (one rotary per device per table kind, cloned into every + /// layer on that device) shares ONE cache: the offsets are host data, so + /// one small H2D per step per device is the floor, and this cache holds it + /// there instead of paying it again on all 43 layers × 3 rope calls + /// (CLAUDE.md pitfall #5 — no `Tensor::from_vec` per layer on the hot + /// path). `None` until the gated path first runs; the flag-off paths never + /// touch it. + cohort_gather: Arc>>, +} + +/// One step's gathered `[B, T, rope_dim / 2]` tables for a ragged cohort, +/// keyed by the exact `(seqlen_offsets, seq_len)` that built them. All three +/// tables are gathered together: `forward` reads (cos, sin), +/// `forward_inverse_tail` reads (cos, neg_sin), and they share the step. +#[derive(Debug)] +struct CohortGather { + offsets: Vec, + seq_len: usize, + /// `[B, T, rope_dim / 2]` each, contiguous, on the tables' device. + cos: Tensor, + sin: Tensor, + neg_sin: Tensor, } // The engagement counters for this path are the ones declared beside @@ -1675,7 +1737,156 @@ impl DeepSeekV2RotaryEmbedding { /// in one constructor and forgotten in the other. fn from_tables(sin: Tensor, cos: Tensor) -> Result { let neg_sin = sin.neg()?; - Ok(Self { sin, cos, neg_sin }) + Ok(Self { + sin, + cos, + neg_sin, + cohort_gather: Arc::new(std::sync::Mutex::new(None)), + }) + } + + /// The step's `(cos, sin, neg_sin)` — each `[B, T, rope_dim / 2]` — for a + /// ragged cohort, one table row per `(b, t)` at absolute position + /// `seqlen_offsets[b] + t`. + /// + /// The row indices are host data (`seqlen_offsets` is a host slice), so + /// one H2D copy per step is the floor; everything after it — the three + /// `index_select`s included — happens once per step per device and is + /// served from [`Self::cohort_gather`] to the other rope calls of the same + /// step (all layers on this device share one instance of `self`; see the + /// field's doc). `index_select` output is contiguous and `reshape` of a + /// contiguous tensor is a view, so the returned tables satisfy `rope_i`'s + /// contiguity requirement with no extra copy. + fn cohort_tables( + &self, + seqlen_offsets: &[usize], + seq_len: usize, + ) -> Result<(Tensor, Tensor, Tensor)> { + let mut cache = self + .cohort_gather + .lock() + .expect("cohort gather cache poisoned"); + if let Some(entry) = cache.as_ref() { + if entry.seq_len == seq_len && entry.offsets == seqlen_offsets { + return Ok((entry.cos.clone(), entry.sin.clone(), entry.neg_sin.clone())); + } + } + let b = seqlen_offsets.len(); + let mut rows: Vec = Vec::with_capacity(b * seq_len); + for &off in seqlen_offsets { + for t in 0..seq_len { + rows.push((off + t) as u32); + } + } + let idx = Tensor::from_vec(rows, b * seq_len, self.cos.device())?; + let half = self.cos.dim(D::Minus1)?; + let cos = self + .cos + .index_select(&idx, 0)? + .reshape((b, seq_len, half))?; + let sin = self + .sin + .index_select(&idx, 0)? + .reshape((b, seq_len, half))?; + let neg_sin = self + .neg_sin + .index_select(&idx, 0)? + .reshape((b, seq_len, half))?; + *cache = Some(CohortGather { + offsets: seqlen_offsets.to_vec(), + seq_len, + cos: cos.clone(), + sin: sin.clone(), + neg_sin: neg_sin.clone(), + }); + Ok((cos, sin, neg_sin)) + } + + /// The gated ragged arm of [`Self::forward`]: ONE `rope_i` per projection + /// over the whole cohort, with the per-row rotation expressed through a + /// rank-3 `[B, T, D/2]` cos/sin (see [`rope_cohort_enabled`] for why that + /// is expressible at all). Bit-identical to [`Self::forward_ragged_loop`]: + /// the gather copies table rows verbatim, and each output element is the + /// same two-multiply-one-add of the same inputs — batching changes which + /// launch computes an element, never the arithmetic. Asserted bit-for-bit + /// in `deepseek_ragged_cohort_tests`. + fn forward_ragged_cohort( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offsets: &[usize], + seq_len: usize, + ) -> Result<(Tensor, Tensor)> { + let (cos, sin, _neg_sin) = self.cohort_tables(seqlen_offsets, seq_len)?; + let q_embed = candle_nn::rotary_emb::rope_i(&q.contiguous()?, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope_i(&k.contiguous()?, &cos, &sin)?; + Ok((q_embed, k_embed)) + } + + /// The per-sequence loop, extracted verbatim so both [`Self::forward`] + /// arms name the thing they dispatch between and the parity tests can + /// drive each directly whatever `ARC_ROPE_COHORT` latched to. + fn forward_ragged_loop( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offsets: &[usize], + seq_len: usize, + ) -> Result<(Tensor, Tensor)> { + let mut q_embeds = Vec::new(); + let mut k_embeds = Vec::new(); + for (i, offset) in seqlen_offsets.iter().enumerate() { + let cos = self.cos.narrow(0, *offset, seq_len)?; + let sin = self.sin.narrow(0, *offset, seq_len)?; + let q_embed = + candle_nn::rotary_emb::rope_i(&q.i(i)?.unsqueeze(0)?.contiguous()?, &cos, &sin)?; + let k_embed = + candle_nn::rotary_emb::rope_i(&k.i(i)?.unsqueeze(0)?.contiguous()?, &cos, &sin)?; + q_embeds.push(q_embed); + k_embeds.push(k_embed); + } + Ok((Tensor::cat(&q_embeds, 0)?, Tensor::cat(&k_embeds, 0)?)) + } + + /// Gated ragged arm of [`Self::forward_inverse_tail`]: one `rope_i` over + /// the whole cohort, conjugate rotation via the pre-negated table gathered + /// per row. Returns only the rotated tail; the caller re-attaches the NoPE + /// dims (shared with the loop arm, so the `cat` is spelled once). + fn inverse_tail_ragged_cohort( + &self, + x: &Tensor, + rope_dim: usize, + nope: usize, + seqlen_offsets: &[usize], + seq_len: usize, + ) -> Result { + let (cos, _sin, neg_sin) = self.cohort_tables(seqlen_offsets, seq_len)?; + let x_pe = x.narrow(3, nope, rope_dim)?.contiguous()?; + candle_nn::rotary_emb::rope_i(&x_pe, &cos, &neg_sin) + } + + /// Per-sequence loop arm of [`Self::forward_inverse_tail`], extracted + /// verbatim (minus the trailing NoPE `cat`, which the caller owns). + fn inverse_tail_ragged_loop( + &self, + x: &Tensor, + rope_dim: usize, + nope: usize, + seqlen_offsets: &[usize], + seq_len: usize, + ) -> Result { + let mut outs = Vec::new(); + for (i, offset) in seqlen_offsets.iter().enumerate() { + let cos = self.cos.narrow(0, *offset, seq_len)?; + let sin = self.neg_sin.narrow(0, *offset, seq_len)?; + let x_pe = x + .i(i)? + .unsqueeze(0)? + .narrow(3, nope, rope_dim)? + .contiguous()?; + outs.push(candle_nn::rotary_emb::rope_i(&x_pe, &cos, &sin)?); + } + Tensor::cat(&outs, 0) } // The cohort predicate for this rotary is the free function @@ -1865,27 +2076,14 @@ impl DeepSeekV2RotaryEmbedding { let q_embed = candle_nn::rotary_emb::rope_i(&q.contiguous()?, &cos, &sin)?; let k_embed = candle_nn::rotary_emb::rope_i(&k.contiguous()?, &cos, &sin)?; Ok((q_embed, k_embed)) + } else if rope_cohort_enabled() { + // Ragged rows, one gathered rope_i per projection — opt-in, see + // [`rope_cohort_enabled`]. + rope_cohort_stats::record_cohort(); + self.forward_ragged_cohort(q, k, seqlen_offsets, seq_len) } else { rope_cohort_stats::record_per_sequence(); - let mut q_embeds = Vec::new(); - let mut k_embeds = Vec::new(); - for (i, offset) in seqlen_offsets.iter().enumerate() { - let cos = self.cos.narrow(0, *offset, seq_len)?; - let sin = self.sin.narrow(0, *offset, seq_len)?; - let q_embed = candle_nn::rotary_emb::rope_i( - &q.i(i)?.unsqueeze(0)?.contiguous()?, - &cos, - &sin, - )?; - let k_embed = candle_nn::rotary_emb::rope_i( - &k.i(i)?.unsqueeze(0)?.contiguous()?, - &cos, - &sin, - )?; - q_embeds.push(q_embed); - k_embeds.push(k_embed); - } - Ok((Tensor::cat(&q_embeds, 0)?, Tensor::cat(&k_embeds, 0)?)) + self.forward_ragged_loop(q, k, seqlen_offsets, seq_len) } } @@ -1912,20 +2110,14 @@ impl DeepSeekV2RotaryEmbedding { let sin = self.neg_sin.narrow(0, offset, seq_len)?; let x_pe = x.narrow(3, nope, rope_dim)?.contiguous()?; candle_nn::rotary_emb::rope_i(&x_pe, &cos, &sin)? + } else if rope_cohort_enabled() { + // Ragged rows, one gathered rope_i — opt-in, see + // [`rope_cohort_enabled`]. + rope_cohort_stats::record_cohort(); + self.inverse_tail_ragged_cohort(x, rope_dim, nope, seqlen_offsets, seq_len)? } else { rope_cohort_stats::record_per_sequence(); - let mut outs = Vec::new(); - for (i, offset) in seqlen_offsets.iter().enumerate() { - let cos = self.cos.narrow(0, *offset, seq_len)?; - let sin = self.neg_sin.narrow(0, *offset, seq_len)?; - let x_pe = x - .i(i)? - .unsqueeze(0)? - .narrow(3, nope, rope_dim)? - .contiguous()?; - outs.push(candle_nn::rotary_emb::rope_i(&x_pe, &cos, &sin)?); - } - Tensor::cat(&outs, 0)? + self.inverse_tail_ragged_loop(x, rope_dim, nope, seqlen_offsets, seq_len)? }; Tensor::cat(&[&x_nope, &rotated], 3)?.contiguous() } @@ -4614,3 +4806,201 @@ mod deepseek_rope_cohort_tests { ); } } + +/// Parity for the **gated** ragged-cohort arms ([`rope_cohort_enabled`], +/// `ARC_ROPE_COHORT`): one gathered rank-3 `rope_i` against the per-sequence +/// loop it replaces, bit for bit. +/// +/// The arms are driven **directly**, not through the public dispatch: the flag +/// latches once per process (`OnceLock`), so a test that reached the gathered +/// arm by mutating the environment would prove only whatever the first reader +/// happened to see — the same reasoning as `env_flag_value`'s own tests. The +/// dispatch's default (flag unset ⇒ ragged rows keep the loop) is pinned by +/// `deepseek_rope_cohort_tests::ragged_forward_still_takes_the_loop_and_still_matches`, +/// which runs the public `forward` in exactly that environment. +#[cfg(test)] +mod deepseek_ragged_cohort_tests { + use super::*; + + const MAX_POS: usize = 32; + const ROPE_DIM: usize = 8; + const HEAD_DIM: usize = 12; + const N_HEADS: usize = 3; + + fn rope() -> DeepSeekV2RotaryEmbedding { + let cfg = DeepSeekV2RopeConfig { + rope_scaling: None, + max_position_embeddings: MAX_POS, + rope_theta: 10000.0, + qk_rope_head_dim: ROPE_DIM, + }; + DeepSeekV2RotaryEmbedding::new(&cfg, DType::F32, &Device::Cpu).unwrap() + } + + fn xs(b: usize, h: usize, t: usize, d: usize) -> Tensor { + let n = b * h * t * d; + let v: Vec = (0..n) + .map(|i| ((i as f32) * 0.29).sin() * 1.3 + 0.07) + .collect(); + Tensor::from_vec(v, (b, h, t, d), &Device::Cpu).unwrap() + } + + /// Raw IEEE-754 bit patterns — the claim is that no arithmetic changed at + /// all, so a tolerance would be vacuous and `-0.0 == 0.0` would slip a + /// sign error through. + fn bits(t: &Tensor) -> Vec { + t.flatten_all() + .unwrap() + .to_vec1::() + .unwrap() + .into_iter() + .map(f32::to_bits) + .collect() + } + + /// Mixed offsets, B ∈ {1, 3}, decode (T=1) and uniform-length chunk (T=2): + /// the gathered arm must reproduce the loop bit for bit. + #[test] + fn ragged_cohort_forward_matches_the_loop_bitwise() { + let re = rope(); + for (offsets, t) in [ + (vec![0usize, 4, 9], 1usize), + (vec![0usize, 4, 9], 2), + (vec![9usize, 4, 0], 1), // descending — order must be per-row + (vec![7usize], 1), // B = 1 through the same arms + ] { + let b = offsets.len(); + let q = xs(b, N_HEADS, t, ROPE_DIM); + let k = xs(b, 1, t, ROPE_DIM); + + let (q_c, k_c) = re.forward_ragged_cohort(&q, &k, &offsets, t).unwrap(); + let (q_l, k_l) = re.forward_ragged_loop(&q, &k, &offsets, t).unwrap(); + + assert_eq!(q_c.dims(), q_l.dims()); + assert_eq!( + bits(&q_c), + bits(&q_l), + "gathered Q differs from the loop for offsets {offsets:?}, T={t}" + ); + assert_eq!( + bits(&k_c), + bits(&k_l), + "gathered K differs from the loop for offsets {offsets:?}, T={t}" + ); + + // Non-vacuity: rows at different offsets must actually rotate + // differently, or the equality above would pass under a + // broadcast-one-row bug on BOTH sides. + if b > 1 { + let d01 = (q_c.i(0).unwrap() - q_c.i(1).unwrap()) + .unwrap() + .abs() + .unwrap() + .max_all() + .unwrap() + .to_scalar::() + .unwrap(); + assert!(d01 > 0.0, "rows at different offsets rotated identically"); + } + } + } + + /// Same parity for the conjugate (inverse-tail) rotation, and for the + /// fully assembled `[NoPE | rotated]` output. + #[test] + fn ragged_cohort_inverse_tail_matches_the_loop_bitwise() { + let re = rope(); + let nope = HEAD_DIM - ROPE_DIM; + for (offsets, t) in [ + (vec![0usize, 4, 9], 1usize), + (vec![0usize, 4, 9], 2), + (vec![7usize], 1), + ] { + let b = offsets.len(); + let x = xs(b, N_HEADS, t, HEAD_DIM); + + let tail_c = re + .inverse_tail_ragged_cohort(&x, ROPE_DIM, nope, &offsets, t) + .unwrap(); + let tail_l = re + .inverse_tail_ragged_loop(&x, ROPE_DIM, nope, &offsets, t) + .unwrap(); + assert_eq!(tail_c.dims(), tail_l.dims()); + assert_eq!( + bits(&tail_c), + bits(&tail_l), + "gathered inverse tail differs from the loop for offsets {offsets:?}, T={t}" + ); + + // The assembled output the model actually consumes. + let x_nope = x.narrow(3, 0, nope).unwrap(); + let full_c = Tensor::cat(&[&x_nope, &tail_c], 3) + .unwrap() + .contiguous() + .unwrap(); + let full_l = Tensor::cat(&[&x_nope, &tail_l], 3) + .unwrap() + .contiguous() + .unwrap(); + assert_eq!(bits(&full_c), bits(&full_l)); + } + } + + /// The gather memo must serve repeats AND rebuild when the cohort moves — + /// a stale entry would rotate every layer of the next step at the previous + /// step's positions, which is exactly the class of wrong-but-plausible + /// output that never crashes. + #[test] + fn cohort_gather_rebuilds_when_the_cohort_moves() { + let re = rope(); + let t = 1usize; + + for offsets in [vec![0usize, 4, 9], vec![1usize, 5, 10], vec![2usize, 3, 4]] { + // Twice with the same key: the second call is the memo hit. + for pass in 0..2 { + let (cos, sin, neg_sin) = re.cohort_tables(&offsets, t).unwrap(); + assert_eq!(cos.dims(), &[offsets.len(), t, ROPE_DIM / 2]); + for (b_i, &off) in offsets.iter().enumerate() { + let want_cos = re.cos.narrow(0, off, t).unwrap(); + let want_sin = re.sin.narrow(0, off, t).unwrap(); + let want_neg = re.neg_sin.narrow(0, off, t).unwrap(); + assert_eq!( + bits(&cos.i(b_i).unwrap()), + bits(&want_cos), + "cos row {b_i} wrong for offsets {offsets:?} (pass {pass})" + ); + assert_eq!(bits(&sin.i(b_i).unwrap()), bits(&want_sin)); + assert_eq!(bits(&neg_sin.i(b_i).unwrap()), bits(&want_neg)); + } + } + } + } + + /// `ARC_ROPE_COHORT` is read by VALUE through `env_flag_is_set`, whose + /// polarity table (`=0`/unset ⇒ OFF, only yes-spellings ⇒ ON) is pinned in + /// `mistralrs_quant::env_flag`. What this pins HERE is the dispatch + /// default: with the flag unset — the state of every CI environment this + /// runs in — ragged rows must take the per-sequence loop, so the gathered + /// arm cannot go live by accident of a latch. + #[test] + fn ragged_dispatch_defaults_to_the_loop() { + assert!( + std::env::var_os("ARC_ROPE_COHORT").is_none(), + "test environment unexpectedly sets ARC_ROPE_COHORT; this test \ + pins the UNSET default and cannot do so with the flag present" + ); + assert!(!rope_cohort_enabled(), "unset must mean OFF"); + + let re = rope(); + let q = xs(3, N_HEADS, 1, ROPE_DIM); + let k = xs(3, 1, 1, ROPE_DIM); + let before = rope_cohort_stats::local_counts(); + re.forward(&q, &k, &[0usize, 4, 9]).unwrap(); + let after = rope_cohort_stats::local_counts(); + assert_eq!( + (after.0 - before.0, after.1 - before.1), + (0, 1), + "with ARC_ROPE_COHORT unset a ragged cohort must take the loop" + ); + } +} diff --git a/mistralrs-core/tests/capability_reachability.rs b/mistralrs-core/tests/capability_reachability.rs index 8b929cb3c..35fad7178 100644 --- a/mistralrs-core/tests/capability_reachability.rs +++ b/mistralrs-core/tests/capability_reachability.rs @@ -214,6 +214,30 @@ static REGISTRY: &[Capability] = &[ }, status: Status::Live, }, + // Same shape as the entry above, same doctrine, one rung lower in the + // stack: the V4 rotary's RAGGED arm. Uniform cohorts already take one + // batched `rope_i` unconditionally (that dispatch is exact and shipped + // ON); this flag governs rows at genuinely different positions, which the + // per-sequence loop served at ~9 ops × 43 layers × B per step. The + // gathered rank-3 path is bit-identical to the loop + // (`deepseek_ragged_cohort_tests`) and default OFF until measured. + // + // flag ARC_ROPE_COHORT=1 (read by value via env_flag_is_set) + // gate layers.rs :: rope_cohort_enabled + // test one binary, ragged B>1 decode (ARC_MTP_PER_SEQ_KV on, or any + // producer of unequal offsets), flag set vs unset + // pass outputs bit-identical AND the gathered arm faster; + // rope_cohort_stats::counts() shows cohort > 0 on the flagged arm + // then the default flips in the same change that records the number + Capability { + name: "ragged-cohort RoPE: one gathered rope_i for unequal offsets, default OFF", + parent: "ArcInfer / ArcAttention", + check: Check::Symbol { + symbol: "rope_cohort_enabled", + defined_in: "mistralrs-core/src/layers.rs", + }, + status: Status::Live, + }, Capability { name: "prefill admission cap (--prefill-max-seqs)", parent: "ArcInfer / ArcSched", From 227c34f8ac23462d80377bb1c9321d5cf70a1ba2 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 20:25:33 +0100 Subject: [PATCH 6/6] =?UTF-8?q?fix(ArcAttention):=20bounds-guard=20the=20c?= =?UTF-8?q?ohort=20gather=20=E2=80=94=20both=20ragged=20arms=20must=20fail?= =?UTF-8?q?=20LOUD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The loop arm's narrow(0, offset, seq_len) hard-errors past the table end; the gathered arm fed the same out-of-range index to index_select, which is not reliably bounds-checked on the CUDA backend and can read garbage SILENTLY. A failure-mode asymmetry on exactly the arm the hardware A/B will exercise is the D18 house fault, so the offsets — host data, the check is free — are now guarded host-side in cohort_tables before any index reaches the device, with the offending offset, seq_len and table size named in the error. out_of_bounds_offset_fails_loud_on_both_arms pins: the loop arm errors, the cohort arm errors on forward AND inverse tail, the message is the GUARD's (a backend error cannot pass the test by coincidence), and offset 31 + T=1 on a 32-row table still works, so the bound is not off by one. Parent system: ArcInfer / ArcAttention. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf --- mistralrs-core/src/layers.rs | 69 ++++++++++++++++++++++++++++++++++++ 1 file changed, 69 insertions(+) diff --git a/mistralrs-core/src/layers.rs b/mistralrs-core/src/layers.rs index b8f6a06ae..818f235a1 100644 --- a/mistralrs-core/src/layers.rs +++ b/mistralrs-core/src/layers.rs @@ -1771,6 +1771,21 @@ impl DeepSeekV2RotaryEmbedding { return Ok((entry.cos.clone(), entry.sin.clone(), entry.neg_sin.clone())); } } + // Host-side bounds guard, BEFORE any index reaches `index_select`. The + // loop arm's `narrow(0, offset, seq_len)` hard-errors past the table + // end; `index_select` with an out-of-range index is not reliably + // bounds-checked on the CUDA backend and can read garbage SILENTLY. + // Both arms must fail loud, and identically-loud — the offsets are + // host data, so the check is free. + let table_rows = self.cos.dim(0)?; + if let Some(&worst) = seqlen_offsets.iter().max() { + if worst + seq_len > table_rows { + candle_core::bail!( + "rope cohort gather out of bounds: offset {worst} + seq_len {seq_len} \ + exceeds the {table_rows}-row cos/sin table (offsets {seqlen_offsets:?})" + ); + } + } let b = seqlen_offsets.len(); let mut rows: Vec = Vec::with_capacity(b * seq_len); for &off in seqlen_offsets { @@ -5003,4 +5018,58 @@ mod deepseek_ragged_cohort_tests { "with ARC_ROPE_COHORT unset a ragged cohort must take the loop" ); } + + /// An offset past the table end must ERROR on the cohort arm exactly as it + /// does on the loop arm — and it must be THIS module's host-side guard + /// that fires, not the backend. The loop's `narrow` fails loud everywhere, + /// but `index_select` with an out-of-range index is not reliably + /// bounds-checked on CUDA and can read garbage silently; the arm we A/B on + /// hardware must not be the one with the quiet failure mode (D18). The + /// message substring is pinned so a backend error passing this test by + /// coincidence is impossible. + #[test] + fn out_of_bounds_offset_fails_loud_on_both_arms() { + let re = rope(); + let t = 1usize; + // MAX_POS = 32 rows: offset 31 is the last valid row at T=1; 32 is + // past the end. The valid rows in the same cohort prove the guard + // reports the MAX, not merely offsets[0]. + let offsets = vec![5usize, 32]; + let q = xs(2, N_HEADS, t, ROPE_DIM); + let k = xs(2, 1, t, ROPE_DIM); + + // Loop arm: narrow(0, 32, 1) on a 32-row table hard-errors. + assert!( + re.forward_ragged_loop(&q, &k, &offsets, t).is_err(), + "the loop arm accepted an offset past the table end" + ); + + // Cohort arm: must also error, and with the guard's own message. + let err = re + .forward_ragged_cohort(&q, &k, &offsets, t) + .expect_err("the cohort arm accepted an offset past the table end"); + let msg = err.to_string(); + assert!( + msg.contains("rope cohort gather out of bounds"), + "the cohort arm rejected the offset, but not via the host-side \ + guard — got: {msg}" + ); + assert!( + msg.contains("offset 32") && msg.contains("32-row"), + "the guard's message must name the offending offset and the table \ + size — got: {msg}" + ); + + // Same guard on the inverse-tail arm (shared `cohort_tables`). + let x = xs(2, N_HEADS, t, HEAD_DIM); + let err = re + .inverse_tail_ragged_cohort(&x, ROPE_DIM, HEAD_DIM - ROPE_DIM, &offsets, t) + .expect_err("the inverse-tail cohort arm accepted an offset past the table end"); + assert!(err.to_string().contains("rope cohort gather out of bounds")); + + // The boundary itself is not off by one: the last valid row works. + let ok = vec![5usize, 31]; + re.forward_ragged_cohort(&q, &k, &ok, t) + .expect("offset 31 + seq_len 1 on a 32-row table is in bounds"); + } }