From 003ca8d774a391146759cf715374f9a0e5e79171 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 03:19:07 +0100 Subject: [PATCH 1/9] fix(ArcModels): Gemma reported a 4096 context because it read the serde default `Model::new` set `max_seq_len: default_max_position_embeddings()` -- the literal `4096` fallback used when a config file is silent -- in a struct literal whose two adjacent lines read `cfg.max_position_embeddings` correctly. Every Gemma checkpoint ships an explicit `max_position_embeddings` (8192 for gemma-2b and gemma-7b), so the usable context of every Gemma was halved. `NormalModel::max_seq_len` is not cosmetic: the engine consults it for prompt rejection, prompt truncation, `is_done`, CUDA-graph block sizing, and the context length reported to API clients. All five were being told 4096 while the KV cache next to them was sized for 8192. Gemma was the only model in `models/` that did this -- deepseek2/3/4, gemma2, glm4*, gpt_oss, llama, mistral, mixtral, phi2, phi3 and the rest all read `cfg.max_position_embeddings` in the same position. UNVERIFIED ON HARDWARE: no GPU was used. The change is a one-line substitution of a config-derived value for a constant; there is no arithmetic to measure. Parent system: ArcModels. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01UMmjFy8TvsgypxVWNVhhC7 --- mistralrs-core/src/models/gemma.rs | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/mistralrs-core/src/models/gemma.rs b/mistralrs-core/src/models/gemma.rs index 6be961476b..6e25256ff6 100644 --- a/mistralrs-core/src/models/gemma.rs +++ b/mistralrs-core/src/models/gemma.rs @@ -456,7 +456,19 @@ impl Model { cfg.num_hidden_layers, cfg.max_position_embeddings, )), - max_seq_len: default_max_position_embeddings(), + // 🔴 This read `default_max_position_embeddings()` -- the literal + // 4096 serde fallback -- while the two lines that bracket it read + // `cfg.max_position_embeddings` correctly. `default_*` is the value + // to use when the config file is SILENT, never the value to use when + // it has spoken; every Gemma checkpoint ships an explicit + // `max_position_embeddings` (8192 for gemma-2b/7b, 8192 for + // codegemma), so this silently halved the usable context of every + // Gemma. `NormalModel::max_seq_len` is what the engine consults for + // prompt rejection, truncation, `is_done`, CUDA-graph block sizing + // and the context length reported to API clients -- all five were + // being told 4096. Every sibling model in `models/` reads + // `cfg.max_position_embeddings` here; gemma was the only exception. + max_seq_len: cfg.max_position_embeddings, cfg: ModelConfigMetadata { max_seq_len: cfg.max_position_embeddings, num_layers: cfg.num_hidden_layers, From 7ccb189d71fe0d98191d280c7c567f1c3852bb8e Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 03:19:30 +0100 Subject: [PATCH 2/9] fix(ArcInfer/ArcKV): Gemma-2's global layers were given a sliding-window ring `Model::new` built its cache with `NormalCache::new_sliding(cfg.num_hidden_layers, _, Some(cfg.sliding_window))`, which fills `num_hidden_layers` slots with the SAME `KvCache::Rotating` prototype. Gemma-2 alternates SWA and global attention -- `Attention::new` sets `use_sliding_window: layer_idx.is_multiple_of(2)`, "Order is SWA, global, SWA" -- so the odd, global-attention layers were handed a `sliding_window`-entry ring buffer as well. Below `sliding_window` tokens that is only wasteful. At and beyond it the ring wraps, so a global layer sees the last `sliding_window` keys in RING order rather than positional order, while `forward` pairs it with the non-windowed `attention_mask`. The mask cannot compensate -- it has no way to know the rows were rotated -- so the model degrades silently rather than erroring. gemma-2-9b's `max_position_embeddings` is 8192 against a 4096 window, so the broken regime is inside the supported context, not past it. `NormalCache::from_types` already exists for exactly this and `qwen3.rs` already uses it; gemma2 predates it. The per-layer kinds are now built by `cache_types`, whose predicate is pinned against `Attention::new`'s by test. Tests (CPU, no GPU required): - `cache_types_match_attention_layers` walks all 42 layers of a gemma-2-9b shaped config and asserts SWA layers get a window and global layers get a full-length cache, with the right window and max_seq_len values. - `global_layers_get_a_full_length_cache_not_a_ring` asserts on the cache actually constructed: 21 of 42 slots rotating, not 42. Before this change all 42 were rotating, so that count is what fails. Parent system: ArcInfer / ArcKV. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01UMmjFy8TvsgypxVWNVhhC7 --- mistralrs-core/src/models/gemma2.rs | 139 ++++++++++++++++++++++++++-- 1 file changed, 133 insertions(+), 6 deletions(-) diff --git a/mistralrs-core/src/models/gemma2.rs b/mistralrs-core/src/models/gemma2.rs index c41c46ee45..241039ac68 100644 --- a/mistralrs-core/src/models/gemma2.rs +++ b/mistralrs-core/src/models/gemma2.rs @@ -23,7 +23,8 @@ use crate::{ pipeline::{ extract_logits, text_models_inputs_processor::{FlashParams, PagedAttentionInputMetadata}, - EitherCache, IsqModel, KvCache, NormalCache, NormalLoadingMetadata, NormalModel, + EitherCache, IsqModel, KvCache, NormalCache, NormalCacheType, NormalLoadingMetadata, + NormalModel, }, utils::{progress::NiceProgressBar, unvarbuilder::UnVarBuilder}, }; @@ -392,6 +393,43 @@ pub struct Model { cfg: ModelConfigMetadata, } +/// One [`NormalCacheType`] per decoder layer, matching Gemma-2's alternation. +/// +/// 🔴 This exists because the constructor used to call +/// `NormalCache::new_sliding(cfg.num_hidden_layers, .., Some(cfg.sliding_window))`, +/// which hands **every** layer a `KvCache::Rotating` ring of `sliding_window` +/// entries. Gemma-2 alternates SWA / global (`Attention::new`: +/// `use_sliding_window: layer_idx.is_multiple_of(2)`, "Order is SWA, global, +/// SWA"), so the odd, **global-attention** layers were given a 4096-entry ring +/// buffer too. +/// +/// Below 4096 tokens that is merely wasteful. At and beyond 4096 it is wrong: +/// the ring wraps, so a global layer sees only the last 4096 keys, in *ring* +/// order rather than positional order, while `forward` masks it with the +/// non-windowed `attention_mask`. The mask cannot correct for it -- it does not +/// know the rows have been rotated -- so the model degrades silently instead of +/// erroring. Gemma-2's own `max_position_embeddings` is 8192, i.e. exactly +/// twice the window, so the broken regime is inside the supported context. +/// +/// The predicate here must stay identical to `Attention::new`'s; the pairing is +/// pinned by `gemma2_cache_type_tests::cache_types_match_attention_layers`. +pub(crate) fn cache_types(cfg: &Config) -> Vec { + (0..cfg.num_hidden_layers) + .map(|layer_idx| { + if layer_idx.is_multiple_of(2) { + // Order is SWA, global, SWA -- same predicate as `Attention::new`. + NormalCacheType::SlidingWindow { + window: cfg.sliding_window, + } + } else { + NormalCacheType::Normal { + max_seq_len: cfg.max_position_embeddings, + } + } + }) + .collect() +} + impl Model { pub fn new( cfg: &Config, @@ -483,11 +521,10 @@ impl Model { ))?), device: normal_loading_metadata.real_device, hidden_size: cfg.hidden_size, - cache: EitherCache::Normal(NormalCache::new_sliding( - cfg.num_hidden_layers, - cfg.max_position_embeddings, - Some(cfg.sliding_window), - )), + // Per-layer cache kinds -- see `cache_types`. This used to be + // `new_sliding(.., Some(window))`, which gave the global layers a + // ring buffer as well. + cache: EitherCache::Normal(NormalCache::from_types(cache_types(cfg))), max_seq_len: cfg.max_position_embeddings, sliding_window: cfg.sliding_window, final_logit_softcapping: cfg.final_logit_softcapping, @@ -826,3 +863,93 @@ impl AnyMoeBaseModelMixin for Model { true } } + +/// Gemma-2 gives every *other* layer a sliding window; the KV cache must agree. +/// +/// The bug these pin: `NormalCache::new_sliding(n, _, Some(window))` builds `n` +/// **identical** `Rotating` slots, so the global-attention layers were handed a +/// `sliding_window`-entry ring buffer. Past `sliding_window` tokens the ring +/// wraps and a global layer can no longer see the whole prefix -- silently, +/// because the non-windowed mask it is paired with has no way to notice. +#[cfg(test)] +mod gemma2_cache_type_tests { + use super::*; + + /// google/gemma-2-9b-it: 42 layers, window 4096, context 8192 -- i.e. the + /// context is twice the window, so the broken regime is *inside* the + /// supported range, not beyond it. + fn cfg() -> Config { + Config { + num_hidden_layers: 42, + sliding_window: 4096, + max_position_embeddings: 8192, + ..Default::default() + } + } + + /// The layer kinds must alternate SWA / global starting at SWA, which is + /// the same predicate `Attention::new` uses for `use_sliding_window`. + #[test] + fn cache_types_match_attention_layers() { + let cfg = cfg(); + let types = cache_types(&cfg); + assert_eq!(types.len(), cfg.num_hidden_layers); + for (layer_idx, ty) in types.iter().enumerate() { + // Mirrors `Attention::new`: `use_sliding_window: layer_idx.is_multiple_of(2)`. + let attention_layer_slides = layer_idx.is_multiple_of(2); + match ty { + NormalCacheType::SlidingWindow { window } => { + assert!( + attention_layer_slides, + "layer {layer_idx} attends globally but was given a ring buffer" + ); + assert_eq!(*window, cfg.sliding_window); + } + NormalCacheType::Normal { max_seq_len } => { + assert!( + !attention_layer_slides, + "layer {layer_idx} uses SWA but was given a full-length cache" + ); + assert_eq!(*max_seq_len, cfg.max_position_embeddings); + } + } + } + } + + /// The regression itself, at the level of the cache actually constructed: + /// before the fix every slot was `Rotating`, so the `Normal` assertion on + /// odd layers is exactly what used to fail. + #[test] + fn global_layers_get_a_full_length_cache_not_a_ring() { + let cfg = cfg(); + let cache = NormalCache::from_types(cache_types(&cfg)); + let cache = cache.lock().unwrap(); + assert_eq!(cache.0.len(), cfg.num_hidden_layers); + + let rotating = cache + .0 + .iter() + .filter(|c| matches!(c, KvCache::Rotating { .. })) + .count(); + // 42 layers -> 21 SWA, 21 global. Not 42 rotating, which is what + // `new_sliding` produced. + assert_eq!(rotating, 21, "expected half the layers to slide, got {rotating}"); + + for (layer_idx, slot) in cache.0.iter().enumerate() { + if layer_idx.is_multiple_of(2) { + assert!( + matches!(slot, KvCache::Rotating { .. }), + "SWA layer {layer_idx} must use a rotating cache" + ); + } else { + assert!( + !matches!(slot, KvCache::Rotating { .. }), + "global layer {layer_idx} must NOT use a rotating cache -- \ + it wraps at {} tokens inside a {}-token context", + cfg.sliding_window, + cfg.max_position_embeddings + ); + } + } + } +} From 7b69e42ec3f04cd66e33076439dd2260841a0ce1 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 03:19:30 +0100 Subject: [PATCH 3/9] fix(ArcMoE): Phi-3.5-MoE's router clamped the wrong way -- min where HF has max HF PhiMoE's `sparsemixer` computes its routing mask as factor = scores.abs().clamp(min=mask_logits_threshold) mask = ((mask_logits_threshold - scores) / factor) > (2 * jitter_eps) `torch.clamp(min=x)` is a LOWER BOUND -- elementwise `max(input, x)`. Both call sites in this port spelled it `broadcast_minimum`, i.e. `min(|s|, t)`, the exact opposite. Verified against `transformers/src/transformers/models/phimoe/modeling_phimoe.py`, where the top-1 and top-2 blocks are identical in this respect. Consequence: `min <= max`, and the numerator `t - s` is non-negative for the top-1 threshold, so the ratio was systematically too LARGE. The `> 2 * jitter_eps` mask therefore over-fired, sending extra experts to `-inf` before the softmax and changing the gate multiplier on every token of every forward. The model still produces plausible text, which is why this survived. There is also a structural difference, not just a numeric one. With `clamp(min=t)` and `t > 0` the denominator is floored at `t`, so a score of exactly zero divides by `t`. Under `minimum` the denominator became `min(0, t) == 0`, producing `inf` and an unconditional mask -- an exposure HF does not have. The duplicated expression is now one function, `sparsemixer_jitter_mask`, so the two call sites cannot drift apart again. Tests (CPU, no GPU required), against a scalar reference written straight from the HF formula so it cannot drift with the code it checks: - `matches_hf_clamp_min_semantics` - `minimum_and_maximum_actually_disagree_here` -- the vacuity guard. It fails if the chosen input cannot separate the bug from the fix, which it did on the first input tried; the mask only notices inside the band where the two ratios straddle `2 * jitter_eps` (`s in [0.98, 1/1.02)` for `t = 1`, `eps = 0.01`). - `zero_score_does_not_divide_by_zero` -- pins the `inf` exposure above. Parent system: ArcInfer / ArcMoE. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01UMmjFy8TvsgypxVWNVhhC7 --- mistralrs-core/src/models/phi3_5_moe.rs | 179 ++++++++++++++++++++++-- 1 file changed, 169 insertions(+), 10 deletions(-) diff --git a/mistralrs-core/src/models/phi3_5_moe.rs b/mistralrs-core/src/models/phi3_5_moe.rs index 07db3b7a0d..a6847d3cb1 100644 --- a/mistralrs-core/src/models/phi3_5_moe.rs +++ b/mistralrs-core/src/models/phi3_5_moe.rs @@ -344,6 +344,44 @@ struct MoeMlp { num_experts: usize, } +/// HF PhiMoE's `sparsemixer` routing mask: +/// +/// ```text +/// factor = scores.abs().clamp(min=mask_logits_threshold) +/// mask = ((mask_logits_threshold - scores) / factor) > (2 * jitter_eps) +/// ``` +/// +/// 🔴 Both call sites used to spell `clamp(min=…)` as `broadcast_minimum`. +/// `torch.clamp(min=x)` is a **lower bound** -- it is `max(input, x)`, +/// elementwise -- so `broadcast_minimum` computes the exact opposite: +/// `min(|s|, threshold)` where HF computes `max(|s|, threshold)`. The +/// denominator was therefore too small on every element where +/// `|s| < threshold`, which for the top-1 threshold is *every non-argmax +/// element*, so the ratio was inflated, the `> 2ε` mask over-fired, and extra +/// experts were masked to `-inf` before the softmax. That changes the gate +/// multiplier on every token of every forward -- silently, since the model +/// still produces plausible text. +/// +/// Signs matter for the direction of the error but not for its existence: +/// `factor` is only ever compared as a divisor, and the numerator +/// `threshold - scores` is non-negative for the top-1 case, so an +/// under-estimated `factor` can only push the ratio *up*, never down. +/// +/// Source, verified against +/// `transformers/src/transformers/models/phimoe/modeling_phimoe.py` +/// (both the top-1 and top-2 blocks, which are identical in this respect). +fn sparsemixer_jitter_mask( + scores: &Tensor, + mask_logits_threshold: &Tensor, + jitter_eps: f64, +) -> Result { + let factor = scores.abs()?.broadcast_maximum(mask_logits_threshold)?; + mask_logits_threshold + .broadcast_sub(scores)? + .broadcast_div(&factor)? + .gt(2. * jitter_eps) +} + impl MoeMlp { fn new( cfg: &Config, @@ -376,11 +414,8 @@ impl MoeMlp { // Compute mask for sparsity let selected_experts = scores.argmax_keepdim(D::Minus1)?; let mask_logits_threshold = scores.gather(&selected_experts, D::Minus1)?; - let factor = scores.abs()?.broadcast_minimum(&mask_logits_threshold)?; - let mask_logits_threshold = mask_logits_threshold - .broadcast_sub(scores)? - .broadcast_div(&factor)? - .gt(2. * jitter_eps)?; + let mask_logits_threshold = + sparsemixer_jitter_mask(scores, &mask_logits_threshold, jitter_eps)?; // Apply mask let masked_gates = masked_fill(scores, &mask_logits_threshold, f64::NEG_INFINITY)?; @@ -401,11 +436,11 @@ impl MoeMlp { // Compute mask for sparsity let selected_experts_top2 = masked_scores.argmax_keepdim(D::Minus1)?; let mask_logits_threshold = masked_scores.gather(&selected_experts_top2, D::Minus1)?; - let factor = scores.abs()?.broadcast_minimum(&mask_logits_threshold)?; - let mask_logits_threshold = mask_logits_threshold - .broadcast_sub(scores)? - .broadcast_div(&factor)? - .gt(2. * jitter_eps)?; + // NB: the numerator threshold comes from `masked_scores` but both the + // subtrahend and the `abs()` are taken from the *unmasked* `scores` -- + // that asymmetry is HF's, verbatim, not a typo here. + let mask_logits_threshold = + sparsemixer_jitter_mask(scores, &mask_logits_threshold, jitter_eps)?; // Apply mask let masked_gates_top2 = @@ -888,3 +923,127 @@ impl NormalModel for Model { } impl AnyMoeBaseModelMixin for Model {} + +/// `sparsemixer`'s routing mask against a hand-evaluated reference. +/// +/// The bug pinned here: `scores.abs().clamp(min=t)` is `max(|s|, t)`, and the +/// port used `broadcast_minimum` -- `min(|s|, t)`. The reference below is +/// written out in scalar Rust straight from the HF formula, so it cannot drift +/// with the implementation it is checking. +#[cfg(test)] +mod sparsemixer_jitter_tests { + use super::*; + use candle_core::{Device, Tensor}; + + const JITTER_EPS: f64 = 0.01; + + /// HF: `((t - s) / max(|s|, t)) > 2 * jitter_eps`, evaluated elementwise. + fn reference(scores: &[f32], threshold: f32, jitter_eps: f64) -> Vec { + scores + .iter() + .map(|&s| { + let factor = s.abs().max(threshold); + let ratio = (threshold - s) / factor; + u8::from(ratio > (2.0 * jitter_eps) as f32) + }) + .collect() + } + + /// The `broadcast_minimum` the fix replaced, so the test can show the two + /// formulas genuinely disagree on this input rather than merely agreeing + /// with itself. + fn reference_with_minimum(scores: &[f32], threshold: f32, jitter_eps: f64) -> Vec { + scores + .iter() + .map(|&s| { + let factor = s.abs().min(threshold); + let ratio = (threshold - s) / factor; + u8::from(ratio > (2.0 * jitter_eps) as f32) + }) + .collect() + } + + fn run(scores: &[f32], threshold: f32) -> Vec { + let dev = Device::Cpu; + let n = scores.len(); + let s = Tensor::from_vec(scores.to_vec(), (1, n), &dev).unwrap(); + let t = Tensor::from_vec(vec![threshold], (1, 1), &dev).unwrap(); + sparsemixer_jitter_mask(&s, &t, JITTER_EPS) + .unwrap() + .flatten_all() + .unwrap() + .to_vec1::() + .unwrap() + } + + /// Router logits with the top-1 threshold taken from the max, which is the + /// shape `sparsemixer` actually calls this with. + /// + /// `0.9802` is deliberate. `min` and `max` always disagree on the *value* + /// of `factor`, but the mask only notices inside the band where the two + /// ratios straddle `2 * jitter_eps`. For `t = 1` and `2ε = 0.02` that band + /// is `s ∈ [0.98, 1/1.02)`: below it both ratios clear the threshold, above + /// it neither does. A score of `0.98` alone would pass under either + /// formula, which is exactly the trap the vacuity test below guards. + #[test] + fn matches_hf_clamp_min_semantics() { + let scores = [1.0f32, 0.9802, 0.5, 0.05, -0.3, -2.0]; + let threshold = 1.0f32; // = max(scores), as `argmax`/`gather` produce. + assert_eq!( + run(&scores, threshold), + reference(&scores, threshold, JITTER_EPS) + ); + } + + /// A green result must prove work happened: if `minimum` and `maximum` + /// agreed on this input the test above would pass either way. + #[test] + fn minimum_and_maximum_actually_disagree_here() { + let scores = [1.0f32, 0.9802, 0.5, 0.05, -0.3, -2.0]; + let threshold = 1.0f32; + + // The denominators differ on every element where |s| != t. This part + // holds for any input and is the structural claim. + let differing_factors = scores + .iter() + .filter(|&&s| s.abs().max(threshold) != s.abs().min(threshold)) + .count(); + assert!( + differing_factors > 0, + "no element distinguishes max(|s|,t) from min(|s|,t)" + ); + + // And the difference reaches the mask, which is what the model consumes. + let with_max = reference(&scores, threshold, JITTER_EPS); + let with_min = reference_with_minimum(&scores, threshold, JITTER_EPS); + assert_ne!( + with_max, with_min, + "input is degenerate -- it cannot distinguish the bug from the fix" + ); + // The live implementation is on the `maximum` side of that split. + assert_eq!(run(&scores, threshold), with_max); + assert_ne!(run(&scores, threshold), with_min); + } + + /// `clamp(min=t)` with `t > 0` makes the denominator at least `t`, so a + /// score of exactly zero divides by `t`, not by zero. Under `minimum` the + /// denominator was `min(0, t) == 0`, producing `inf` and an unconditional + /// mask. This is the structural difference, not just a numeric one. + #[test] + fn zero_score_does_not_divide_by_zero() { + let scores = [1.5f32, 0.0, -0.75]; + let threshold = 1.5f32; + let got = run(&scores, threshold); + assert_eq!(got, reference(&scores, threshold, JITTER_EPS)); + + let dev = Device::Cpu; + let s = Tensor::from_vec(scores.to_vec(), (1, 3), &dev).unwrap(); + let t = Tensor::from_vec(vec![threshold], (1, 1), &dev).unwrap(); + let factor = s.abs().unwrap().broadcast_maximum(&t).unwrap(); + let factor: Vec = factor.flatten_all().unwrap().to_vec1().unwrap(); + assert!( + factor.iter().all(|f| f.is_finite() && *f >= threshold), + "clamp(min=t) must floor the denominator at t, got {factor:?}" + ); + } +} From 875b5e45e3104b83a133eefd4265262fed60b09a Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 03:20:01 +0100 Subject: [PATCH 4/9] perf(ArcModels): Qwen3-Next synced the device 36 times per token for one vector `Model::forward`'s layer loop called `indices.to_vec1()?` on the recurrent state indices inside each linear-attention arm. `to_vec1` is a device-to-host copy, i.e. a full stall of the compute stream, and `indices` is bound once ABOVE the loop and never written to -- the loop hands it to `pool.gather_*` and `pool.scatter_*`, which read it and mutate the pool, never the index tensor. Qwen3-Next has `full_attention_interval == 4`, so 36 of its 48 layers are linear attention: 36 identical syncs per forward where the data supports one. Hoisted to a single read before the loop. The two `bail!`s that lived next to it move with it. The emptiness check is now gated on the same `has_linear_layers` predicate the existing "indices are required" bail uses, so a model with no linear-attention layers sees no new failure mode. The per-layer offset-divergence check stays in the loop -- `pool` is per-layer, so `first_offset` is genuinely not loop-invariant. UNVERIFIED ON HARDWARE: no GPU was used, and a host sync costs nothing measurable on a CPU device. The claim is 36 device syncs become 1, which is a count of `to_vec1` call sites, not a timing. Parent system: ArcModels. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01UMmjFy8TvsgypxVWNVhhC7 --- mistralrs-core/src/models/qwen3_next.rs | 36 ++++++++++++++++++------- 1 file changed, 27 insertions(+), 9 deletions(-) diff --git a/mistralrs-core/src/models/qwen3_next.rs b/mistralrs-core/src/models/qwen3_next.rs index 74d1aee8d9..010b32e75b 100644 --- a/mistralrs-core/src/models/qwen3_next.rs +++ b/mistralrs-core/src/models/qwen3_next.rs @@ -913,17 +913,35 @@ impl Model { let mut hybrid_cache = self.kv_cache.hybrid(); let state_indices = hybrid_cache.state_indices().cloned(); - if self + let has_linear_layers = self .layer_types .iter() - .any(|lt| matches!(lt, LayerType::LinearAttention)) - && state_indices.is_none() - { + .any(|lt| matches!(lt, LayerType::LinearAttention)); + if has_linear_layers && state_indices.is_none() { candle_core::bail!( "Hybrid recurrent state indices are required for linear-attention layers." ); } + // 🔴 `indices.to_vec1()` is a device->host copy, i.e. a full sync of the + // compute stream. It used to sit *inside* the layer loop, on a tensor + // bound once above and never written to -- so every linear-attention + // layer re-read the same bytes and stalled the pipeline again. Qwen3-Next + // has `full_attention_interval == 4`, so 36 of its 48 layers are linear: + // 36 syncs per forward where the data supports one. + // + // Hoisting is sound because `state_indices` is loop-invariant by + // construction: the loop passes `indices` to `pool.gather_*` / + // `pool.scatter_*`, which read it and mutate the *pool*, never the index + // tensor itself. + let state_indices_host: Option> = state_indices + .as_ref() + .map(|indices| indices.to_vec1::()) + .transpose()?; + if has_linear_layers && state_indices_host.as_ref().is_some_and(|v| v.is_empty()) { + candle_core::bail!("Hybrid recurrent state indices are empty."); + } + let mask = CausalMasker.make_causal_mask_matrix( input_ids, metadata @@ -968,10 +986,10 @@ impl Model { let indices = state_indices.as_ref().expect( "checked above: linear-attention layers require recurrent indices", ); - let indices_vec: Vec = indices.to_vec1()?; - if indices_vec.is_empty() { - candle_core::bail!("Hybrid recurrent state indices are empty."); - } + // Read once before the loop -- see `state_indices_host`. + let indices_vec = state_indices_host.as_ref().expect( + "checked above: linear-attention layers require recurrent indices", + ); let first_offset = pool.get_seqlen_offset(indices_vec[0] as usize); if indices_vec @@ -998,7 +1016,7 @@ impl Model { pool.scatter_recurrent_state(indices, &gdn_cache.recurrent_state)?; let delta = gdn_cache.seqlen_offset.saturating_sub(first_offset); - for &idx in &indices_vec { + for &idx in indices_vec { let updated = pool.get_seqlen_offset(idx as usize) + delta; pool.set_seqlen_offset(idx as usize, updated); } From baaf5f40f4d16fc952e93710d95a5477781a9093 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 03:20:01 +0100 Subject: [PATCH 5/9] perf(ArcAttention): RoPE tested the LENGTH of seqlen_offsets at seven more sites Every rotary implementation in `layers.rs` gated its batched path on `seqlen_offsets.len() == 1` -- the number of ENTRIES -- when the property it needs is whether those entries are EQUAL. `len()` is the batch size, so the batched path was unreachable at every batch size above one, while the values were uniform anyway: the scheduler admits a forward pass only when the participating cache lengths are equal. The per-sequence loop was therefore performing B bit-identical recomputations of the same rotation and concatenating them back together. This is the same defect PR #198 fixed for `DeepSeekV2RotaryEmbedding`. Payoff there is measured, not projected: the `rope` node alone spent 728.6 ms/step of HOST time against 163.5 ms/step of device time at B=256 on an H200, and the fix moved a step from 1082.3 ms to 794.0 ms (1.36x). Same shape, five more model families. NOT A CLAIM ABOUT THESE SITES -- no GPU was used here, and no speedup is asserted for them. Seven sites, all `candle_nn::rotary_emb::{rope, rope_i}` or the fused CUDA kernel: - `PhiRotaryEmbedding::forward`, both arms (partial-rotary and full) - `Phi4MMRotaryEmbedding::forward` - `RotaryEmbedding::forward`, CUDA and candle arms - `GptOssRotaryEmbedding::forward`, CUDA and candle arms `DeepSeekV2RotaryEmbedding`'s two sites are deliberately untouched -- they belong to PR #198, which is still open against `release/openrouter-ready`. The predicate here is `uniform_seqlen_offset`, a free function with a body identical to that PR's `DeepSeekV2RotaryEmbedding::uniform_offset`; the doc comment says so and asks for them to be collapsed when the branches meet. The `rope_cohort_stats` counters keep that PR's names. Second, independent defect in the same hunks: `RotaryEmbedding::forward` gated the fused CUDA kernel on `qh == kh` and `GptOssRotaryEmbedding::forward` on `qh == k.dim(1)?`. That is false for every GQA model -- qwen2, qwen3, mistral, gemma2, phi3, and GPT-OSS itself (64 query heads, 8 KV heads) -- so the kernel was dead for the only architectures that reach it. The equality was never a kernel requirement: `mistralrs_quant::rotary::apply_rotary_inplace` passes `num_heads` and `num_kv_heads` as separate arguments and `kernels/rotary/rotary.cu` walks them as two independent loops (`nq = num_heads * rot_dim`, `nk = num_kv_heads * rot_dim`). The only shape it enforces is `(num_tokens, head_size)` matching between q and k, which GQA preserves. Read from the kernel source; UNVERIFIED ON HARDWARE. The CUDA arm also gains an `is_cuda()` device test. It previously checked only the compile-time feature, so a cuda-enabled binary running on CPU or Metal fell into `apply_rotary_inplace`'s "expects a cuda tensor" bail. Widening the head-count gate would have widened that exposure too. The CUDA cos/sin tables cannot simply narrow for a uniform cohort the way the candle paths can: the kernel indexes `cos_cache + token_idx * rot_dim` over the b-major flattened token axis and hard-bails unless the cache is exactly `(b * seq_len, rot_dim)`. `cohort_rotary_table` therefore tiles -- one broadcast plus one copy for a uniform cohort, and the original `b`-way `cat` for a ragged one, which is the only thing that can express it. Tests (CPU, no GPU required), asserting on raw IEEE-754 bit patterns rather than a tolerance, since the claim is that no arithmetic changed at all: - `uniform_offset_tests_values_not_length` -- `[7; 256]` is uniform. - `uniform_cohort_is_bit_identical_to_the_loop` and `..._for_gpt_j_style` -- both rope styles, against the pre-fix loop held as a verbatim oracle. Each asserts the `rope_cohort_stats` COHORT counter advanced, so the test cannot pass while the fast path sits dark. - `ragged_cohort_keeps_the_loop_and_stays_per_row` -- ragged cohorts must still loop, must still be per-row correct, and must genuinely differ between rows. - `cohort_rotary_table_matches_the_cat_it_replaces` -- shape AND bit equality against the `cat` for uniform, single-row and two ragged cases. This is the only CPU-reachable assertion on the CUDA table layout. Parent system: ArcInfer / ArcAttention. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01UMmjFy8TvsgypxVWNVhhC7 --- mistralrs-core/src/layers.rs | 449 +++++++++++++++++++++++++++++++---- 1 file changed, 401 insertions(+), 48 deletions(-) diff --git a/mistralrs-core/src/layers.rs b/mistralrs-core/src/layers.rs index 3a07b09fba..10ddf1f3e6 100644 --- a/mistralrs-core/src/layers.rs +++ b/mistralrs-core/src/layers.rs @@ -350,6 +350,122 @@ impl QRmsNorm { } } +/// The offset shared by every row of a decode cohort, or `None` when the rows +/// genuinely differ. +/// +/// # Why this is a *value* test and not a length test +/// +/// Every rotary implementation in this file gated its batched path on +/// `seqlen_offsets.len() == 1` -- the **length of the vector**, where the +/// property it needs is the **distinctness of the values**. `len()` is the +/// batch size, so the batched path was unreachable at every batch size above +/// one, while the values were uniform anyway: `crate::scheduler` admits a +/// forward pass only when the participating cache lengths are equal. The +/// per-sequence loop was therefore performing B **bit-identical** +/// recomputations of the same `cos`/`sin` rotation and concatenating them back +/// together. +/// +/// Returning `Some` is exact rather than approximate. `candle_nn`'s rope +/// kernels take a 2-D `[T, D/2]` cos/sin whose batch stride is 0, so every +/// batch row already reads the same table 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. That is +/// asserted bit-for-bit, not within a tolerance, in `rope_cohort_tests`. +/// +/// `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 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. +#[inline] +pub(crate) fn uniform_seqlen_offset(seqlen_offsets: &[usize]) -> Option { + match seqlen_offsets { + [] => None, + [first, rest @ ..] => rest.iter().all(|o| o == first).then_some(*first), + } +} + +/// Engagement counters for the uniform-offset cohort path -- see +/// [`uniform_seqlen_offset`]. +/// +/// House rule: "a green result must prove work happened". A fast path that +/// silently never engages produces a perfectly green no-op, so the batched +/// branch counts itself and the per-sequence branch counts itself, and a +/// harness that sees `cohort == 0` must treat that as an environment failure +/// rather than a result. +pub mod rope_cohort_stats { + use std::sync::atomic::{AtomicU64, Ordering}; + + pub(super) static COHORT: AtomicU64 = AtomicU64::new(0); + pub(super) static PER_SEQUENCE: AtomicU64 = AtomicU64::new(0); + + /// `(cohort, per_sequence)` -- batched RoPE calls vs per-sequence loops. + /// + /// Tests assert on *deltas* around a call rather than absolute values, so + /// no reset hook is needed and the counters stay monotonic for a serving + /// harness reading them at the end of a run. + pub fn counts() -> (u64, u64) { + ( + COHORT.load(Ordering::Relaxed), + PER_SEQUENCE.load(Ordering::Relaxed), + ) + } + + #[inline] + pub(super) fn record_cohort() { + COHORT.fetch_add(1, Ordering::Relaxed); + } + + #[inline] + pub(super) fn record_per_sequence() { + PER_SEQUENCE.fetch_add(1, Ordering::Relaxed); + } +} + +/// Tile a `[max_seq_len, rot_dim]` rotary table into the `[b * seq_len, rot_dim]` +/// layout `mistralrs_quant::rotary::apply_rotary_inplace` requires. +/// +/// The CUDA kernel indexes `cos_cache + token_idx * rot_dim` with one block per +/// token of the *flattened* `[b * seq_len, h, d]` tensor, and +/// `mistralrs_quant::rotary` hard-bails unless the cache is exactly +/// `(num_tokens, rot_dim)`. So -- unlike the `candle_nn` paths -- a uniform +/// cohort cannot simply narrow to `[seq_len, rot_dim]`; the rows must actually +/// be repeated, b-major, to match `q.transpose(1, 2)?.flatten(0, 1)?`. +/// +/// Uniform cohorts get one broadcast + one copy instead of `b` narrows and a +/// `b`-way `cat`; ragged cohorts keep the `cat`, which is the only thing that +/// can express them. +fn cohort_rotary_table( + table: &Tensor, + seqlen_offsets: &[usize], + seq_len: usize, +) -> Result { + match uniform_seqlen_offset(seqlen_offsets) { + Some(offset) => { + let rows = table.narrow(0, offset, seq_len)?; + if seqlen_offsets.len() == 1 { + return rows.contiguous(); + } + let rot_dim = rows.dim(D::Minus1)?; + rows.unsqueeze(0)? + .broadcast_as((seqlen_offsets.len(), seq_len, rot_dim))? + .contiguous()? + .reshape((seqlen_offsets.len() * seq_len, rot_dim)) + } + None => { + let mut rows = Vec::with_capacity(seqlen_offsets.len()); + for offset in seqlen_offsets { + rows.push(table.narrow(0, *offset, seq_len)?); + } + Tensor::cat(&rows, 0) + } + } +} + /// RoPE supporting LongRope #[derive(Debug, Clone)] pub struct PhiRotaryEmbedding { @@ -656,13 +772,17 @@ impl PhiRotaryEmbedding { let k_rot = k.narrow(D::Minus1, 0, rot_dim)?; let k_pass = k.narrow(D::Minus1, rot_dim, k.dim(D::Minus1)? - rot_dim)?; - let (q_rot, k_rot) = if seqlen_offsets.len() == 1 { - let cos = cos.narrow(0, seqlen_offsets[0], seq_len)?; - let sin = sin.narrow(0, seqlen_offsets[0], seq_len)?; + // Batched whenever every row shares an offset -- see + // `uniform_seqlen_offset`. This used to require `len() == 1`. + let (q_rot, k_rot) = if let Some(offset) = uniform_seqlen_offset(seqlen_offsets) { + rope_cohort_stats::record_cohort(); + let cos = cos.narrow(0, offset, seq_len)?; + let sin = sin.narrow(0, offset, seq_len)?; let q_embed = candle_nn::rotary_emb::rope(&q_rot.contiguous()?, &cos, &sin)?; let k_embed = candle_nn::rotary_emb::rope(&k_rot.contiguous()?, &cos, &sin)?; (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() { @@ -690,13 +810,15 @@ impl PhiRotaryEmbedding { Tensor::cat(&[q_rot, q_pass], D::Minus1)?.contiguous()?, Tensor::cat(&[k_rot, k_pass], D::Minus1)?.contiguous()?, )) - } else if seqlen_offsets.len() == 1 { - let cos = cos.narrow(0, seqlen_offsets[0], seq_len)?; - let sin = sin.narrow(0, seqlen_offsets[0], seq_len)?; + } else if let Some(offset) = uniform_seqlen_offset(seqlen_offsets) { + rope_cohort_stats::record_cohort(); + let cos = cos.narrow(0, offset, seq_len)?; + let sin = sin.narrow(0, offset, seq_len)?; let q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?; let k_embed = candle_nn::rotary_emb::rope(&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() { @@ -1884,13 +2006,17 @@ impl Phi4MMRotaryEmbedding { let k_rot = k.narrow(D::Minus1, 0, rot_dim)?; let k_pass = k.narrow(D::Minus1, rot_dim, k.dim(D::Minus1)? - rot_dim)?; - let (q_rot, k_rot) = if seqlen_offsets.len() == 1 { - let cos = cos.narrow(0, seqlen_offsets[0], seq_len)?; - let sin = sin.narrow(0, seqlen_offsets[0], seq_len)?; + // Batched whenever every row shares an offset -- see + // `uniform_seqlen_offset`. This used to require `len() == 1`. + let (q_rot, k_rot) = if let Some(offset) = uniform_seqlen_offset(seqlen_offsets) { + rope_cohort_stats::record_cohort(); + let cos = cos.narrow(0, offset, seq_len)?; + let sin = sin.narrow(0, offset, seq_len)?; let q_embed = candle_nn::rotary_emb::rope(&q_rot.contiguous()?, &cos, &sin)?; let k_embed = candle_nn::rotary_emb::rope(&k_rot.contiguous()?, &cos, &sin)?; (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() { @@ -2652,21 +2778,28 @@ impl RotaryEmbedding { } } - if cfg!(feature = "cuda") && qh == kh { - let (cos, sin) = if seqlen_offsets.len() == 1 { - ( - self.cos.narrow(0, seqlen_offsets[0], seq_len)?, - self.sin.narrow(0, seqlen_offsets[0], seq_len)?, - ) - } else { - let mut cos_s = Vec::new(); - let mut sin_s = Vec::new(); - for offset in seqlen_offsets { - cos_s.push(self.cos.narrow(0, *offset, seq_len)?); - sin_s.push(self.sin.narrow(0, *offset, seq_len)?); - } - (Tensor::cat(&cos_s, 0)?, Tensor::cat(&sin_s, 0)?) - }; + // 🔴 This used to read `cfg!(feature = "cuda") && qh == kh`, which + // switched the fused kernel off for **every GQA model** -- qwen2, qwen3, + // mistral, gemma2, phi3 and the rest all have `qh != kh` by definition, + // so they fell through to the `candle_nn` path below. + // + // The head-count equality was never a kernel requirement. + // `mistralrs_quant::rotary::apply_rotary_inplace` takes `num_heads` and + // `num_kv_heads` as separate arguments and the kernel walks them as two + // independent loops (`nq = num_heads * rot_dim`, `nk = num_kv_heads * + // rot_dim`, `kernels/rotary/rotary.cu`); the only shape it insists on is + // `(num_tokens, head_size)` matching between q and k, which GQA + // preserves. UNVERIFIED ON HARDWARE -- read from the kernel source, not + // measured. + // + // The `is_cuda()` term is new as well: the old condition tested only the + // compile-time feature, so a cuda-enabled binary running a model on CPU + // or Metal walked into `apply_rotary_inplace` and hit its + // "expects a cuda tensor" bail. That was reachable before only for MHA + // models; widening the head-count gate would have widened that too. + if cfg!(feature = "cuda") && q.device().is_cuda() { + let cos = cohort_rotary_table(&self.cos, seqlen_offsets, seq_len)?; + let sin = cohort_rotary_table(&self.sin, seqlen_offsets, seq_len)?; let q_embed = q.transpose(1, 2)?.flatten(0, 1)?; let k_embed = k.transpose(1, 2)?.flatten(0, 1)?; @@ -2688,13 +2821,15 @@ impl RotaryEmbedding { k = k.contiguous()?; } Ok((q, k)) - } else 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)?; + } else 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 = rope(&q.contiguous()?, &cos, &sin)?; let k_embed = rope(&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() { @@ -2830,24 +2965,18 @@ impl GptOssRotaryEmbedding { #[allow(unused_variables)] let (_b_sz, kh, _seq_len, _n_embd) = k.dims4()?; - // Use CUDA optimized kernel when available and q/k have same number of heads - // The CUDA kernel uses is_neox=true for chunked/GPT-NeoX style rotary + // Use the CUDA kernel when available. The kernel uses is_neox=true for + // chunked/GPT-NeoX style rotary. + // + // 🔴 This used to also require `qh == k.dim(1)?`, i.e. MHA. GPT-OSS is + // GQA (64 query heads, 8 KV heads), so that condition is false for the + // only architecture that reaches this code and the kernel was dead. + // `apply_rotary_inplace` takes `num_kv_heads` separately -- see the + // longer note on `RotaryEmbedding::forward`. UNVERIFIED ON HARDWARE. #[cfg(feature = "cuda")] - if q.device().is_cuda() && qh == k.dim(1)? { - let (cos, sin) = if seqlen_offsets.len() == 1 { - ( - self.cos.narrow(0, seqlen_offsets[0], seq_len)?, - self.sin.narrow(0, seqlen_offsets[0], seq_len)?, - ) - } else { - let mut cos_s = Vec::new(); - let mut sin_s = Vec::new(); - for offset in seqlen_offsets { - cos_s.push(self.cos.narrow(0, *offset, seq_len)?); - sin_s.push(self.sin.narrow(0, *offset, seq_len)?); - } - (Tensor::cat(&cos_s, 0)?, Tensor::cat(&sin_s, 0)?) - }; + if q.device().is_cuda() { + let cos = cohort_rotary_table(&self.cos, seqlen_offsets, seq_len)?; + let sin = cohort_rotary_table(&self.sin, seqlen_offsets, seq_len)?; // Reshape for CUDA kernel: [b, h, seq, dim] -> [b*seq, h, dim] let q_embed = q.transpose(1, 2)?.flatten(0, 1)?; @@ -2871,14 +3000,18 @@ impl GptOssRotaryEmbedding { return Ok((q, k)); } - // CPU fallback using candle_nn's rope (GPT-NeoX/chunked style) - 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)?; + // CPU fallback using candle_nn's rope (GPT-NeoX/chunked style). + // Batched whenever every row shares an offset -- see + // `uniform_seqlen_offset`. This used to require `len() == 1`. + 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(&q.contiguous()?, &cos, &sin)?; let k_embed = candle_nn::rotary_emb::rope(&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() { @@ -3620,3 +3753,223 @@ mod graph_mode_mask_tests { assert!(graph_mode_length_mask(&positions, 0, DType::F32).is_err()); } } + +/// Bit-exactness and engagement tests for the uniform-offset RoPE cohort path. +/// +/// The change under test replaced `seqlen_offsets.len() == 1` -- a test of the +/// *length* of the offset vector -- with [`uniform_seqlen_offset`], a test of +/// the *distinctness of its values*, at every rotary implementation in this +/// file. 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. +/// +/// Each test 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 would be vacuous, since the claim is that no +/// arithmetic changed at all. +#[cfg(test)] +mod rope_cohort_tests { + use super::*; + use candle_core::{DType, Device}; + + const B: usize = 4; + const H: usize = 2; + const T: usize = 3; + const D_HEAD: usize = 8; + const MAX_POS: usize = 64; + + fn xs(dims: &[usize], seed: f32) -> Tensor { + let n: usize = dims.iter().product(); + let v: Vec = (0..n) + .map(|i| ((i as f32) * seed).sin() * 1.7 + 0.11) + .collect(); + Tensor::from_vec(v, dims, &Device::Cpu).unwrap() + } + + /// Raw IEEE-754 bit patterns. Float equality would accept `-0.0 == 0.0` and + /// silently pass a sign error; bit patterns do not. + fn bits(t: &Tensor) -> Vec { + t.flatten_all() + .unwrap() + .to_vec1::() + .unwrap() + .into_iter() + .map(f32::to_bits) + .collect() + } + + fn rope() -> RotaryEmbedding { + RotaryEmbedding::new(10000.0, D_HEAD, MAX_POS, &Device::Cpu, true, DType::F32).unwrap() + } + + /// The pre-fix per-sequence loop, verbatim, as the oracle. Always loops, + /// whatever the offsets are. + fn reference(re: &RotaryEmbedding, q: &Tensor, k: &Tensor, offsets: &[usize]) -> (Tensor, Tensor) { + let seq_len = q.dim(2).unwrap(); + 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).unwrap(); + let sin = re.sin.narrow(0, *offset, seq_len).unwrap(); + q_embeds.push( + candle_nn::rotary_emb::rope(&q.i(i).unwrap().unsqueeze(0).unwrap().contiguous().unwrap(), &cos, &sin) + .unwrap(), + ); + k_embeds.push( + candle_nn::rotary_emb::rope(&k.i(i).unwrap().unsqueeze(0).unwrap().contiguous().unwrap(), &cos, &sin) + .unwrap(), + ); + } + ( + Tensor::cat(&q_embeds, 0).unwrap(), + Tensor::cat(&k_embeds, 0).unwrap(), + ) + } + + /// The regression this whole change is about: a B-long vector of one + /// repeated value IS uniform, and `len() == 1` said it was not. + #[test] + fn uniform_offset_tests_values_not_length() { + assert_eq!(uniform_seqlen_offset(&[7; 256]), Some(7)); + assert_eq!(uniform_seqlen_offset(&[7]), Some(7)); + assert_eq!(uniform_seqlen_offset(&[0, 0, 0]), Some(0)); + // Genuinely ragged rows must still decline, so the loop is preserved. + assert_eq!(uniform_seqlen_offset(&[7, 7, 8]), None); + assert_eq!(uniform_seqlen_offset(&[8, 7, 7]), None); + assert_eq!(uniform_seqlen_offset(&[]), None); + } + + /// Batching a uniform cohort must be bit-identical to the loop it replaces. + #[test] + fn uniform_cohort_is_bit_identical_to_the_loop() { + let re = rope(); + let q = xs(&[B, H, T, D_HEAD], 0.013); + let k = xs(&[B, H, T, D_HEAD], 0.019); + let offsets = vec![5usize; B]; + + let before = rope_cohort_stats::counts(); + let (qb, kb) = re.forward(&q, &k, &offsets).unwrap(); + let after = rope_cohort_stats::counts(); + + // A green result must prove work happened: the batched arm has to have + // been the one that ran. + assert!( + after.0 > before.0, + "cohort path did not engage -- test would pass vacuously" + ); + + let (qref, kref) = reference(&re, &q, &k, &offsets); + assert_eq!(bits(&qb), bits(&qref), "batched q differs from per-sequence"); + assert_eq!(bits(&kb), bits(&kref), "batched k differs from per-sequence"); + } + + /// Ragged offsets must still take the per-sequence loop and still be right. + /// A uniform rotation applied to a ragged cohort would be silently wrong. + #[test] + fn ragged_cohort_keeps_the_loop_and_stays_per_row() { + let re = rope(); + let q = xs(&[3, H, T, D_HEAD], 0.013); + let k = xs(&[3, H, T, D_HEAD], 0.019); + let offsets = vec![0usize, 4, 9]; + + let before = rope_cohort_stats::counts(); + let (qb, kb) = re.forward(&q, &k, &offsets).unwrap(); + let after = rope_cohort_stats::counts(); + assert!( + after.1 > before.1, + "ragged cohort must NOT take the batched path" + ); + + let (qref, kref) = reference(&re, &q, &k, &offsets); + assert_eq!(bits(&qb), bits(&qref)); + assert_eq!(bits(&kb), bits(&kref)); + + // And the rotation must genuinely differ between rows, or the test is + // vacuous -- it would pass even if every row got offset 0. + let d01 = (qb.i(0).unwrap() - qb.i(1).unwrap()) + .unwrap() + .abs() + .unwrap() + .max_all() + .unwrap() + .to_scalar::() + .unwrap(); + assert!(d01 > 0.0, "rows at different offsets rotated identically"); + } + + /// GPT-J style (`is_gpt_neox == false`, i.e. `rope_i`) takes the same + /// dispatch, so it gets the same guarantee. + #[test] + fn uniform_cohort_is_bit_identical_for_gpt_j_style() { + let re = + RotaryEmbedding::new(10000.0, D_HEAD, MAX_POS, &Device::Cpu, false, DType::F32).unwrap(); + let q = xs(&[B, H, T, D_HEAD], 0.023); + let k = xs(&[B, H, T, D_HEAD], 0.029); + let offsets = vec![11usize; B]; + + let (qb, kb) = re.forward(&q, &k, &offsets).unwrap(); + + let seq_len = T; + 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).unwrap(); + let sin = re.sin.narrow(0, *offset, seq_len).unwrap(); + q_embeds.push( + candle_nn::rotary_emb::rope_i( + &q.i(i).unwrap().unsqueeze(0).unwrap().contiguous().unwrap(), + &cos, + &sin, + ) + .unwrap(), + ); + k_embeds.push( + candle_nn::rotary_emb::rope_i( + &k.i(i).unwrap().unsqueeze(0).unwrap().contiguous().unwrap(), + &cos, + &sin, + ) + .unwrap(), + ); + } + assert_eq!(bits(&qb), bits(&Tensor::cat(&q_embeds, 0).unwrap())); + assert_eq!(bits(&kb), bits(&Tensor::cat(&k_embeds, 0).unwrap())); + } + + /// `cohort_rotary_table` feeds the CUDA kernel, which indexes + /// `cos_cache + token_idx * rot_dim` over the b-major flattened token axis + /// and hard-bails unless the cache is exactly `(b * seq_len, rot_dim)`. + /// The uniform fast path must therefore *tile*, not narrow -- this is the + /// assertion that would catch a shape or ordering mistake on a path CPU + /// tests cannot otherwise reach. + #[test] + fn cohort_rotary_table_matches_the_cat_it_replaces() { + let table = xs(&[MAX_POS, D_HEAD / 2], 0.031); + + for offsets in [ + vec![5usize; B], // uniform, B > 1 -> must tile + vec![5usize], // single row + vec![0usize, 4, 9], // ragged -> must keep the cat + vec![3usize, 3, 3, 7], // trailing mismatch -> ragged + ] { + let mut rows = Vec::new(); + for offset in &offsets { + rows.push(table.narrow(0, *offset, T).unwrap()); + } + let expected = Tensor::cat(&rows, 0).unwrap(); + + let got = cohort_rotary_table(&table, &offsets, T).unwrap(); + assert_eq!( + got.dims(), + &[offsets.len() * T, D_HEAD / 2], + "wrong shape for offsets {offsets:?} -- the kernel bails on this" + ); + assert_eq!( + bits(&got), + bits(&expected), + "tiled table differs from the per-offset cat for offsets {offsets:?}" + ); + } + } +} From d7f1482d6e2108690d7854298574735d461bd7de Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 03:23:19 +0100 Subject: [PATCH 6/9] fix(ArcModels): the phi2 GGUF loader was unreachable, and hid a missing GELU Two defects, one masking the other. 1. `PropsGGUF::try_from` listed `attention.layer_norm_rms_epsilon` in `has_required_keys`, so it demanded `phi2.attention.layer_norm_rms_epsilon`. Phi-2 is a LayerNorm architecture -- this very file builds a `candle_nn::LayerNorm` from that epsilon, complete with a bias -- and no phi2 GGUF carries an RMS-eps key. llama.cpp loads phi2's epsilon with `ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps)` and builds the graph with `LLM_NORM` (`src/models/phi2.cpp`); `gguf-py` spells that key `{arch}.attention.layer_norm_epsilon` (`gguf/constants.py`). So the loader rejected every real file. `quantized_starcoder2.rs`, the other LayerNorm-based GGUF loader in this tree, already reads the non-RMS key. The epsilon is now resolved by `layer_norm_eps`, which accepts either spelling and errors naming both. Nothing that loaded before stops loading. 2. `Mlp::forward` had no activation: `ffn_up` straight into `ffn_down`. Two stacked linear maps compose to a single linear map, so the whole FFN contributed no nonlinearity. `git show 1269bd8ab` ("Implement GPTQ quantization", Aug 2024) is the deletion: - xs.apply(&self.ffn_up)?.gelu()?.apply(&self.ffn_down) + MatMul.qmethod_matmul(&MatMul.qmethod_matmul(xs, &*self.ffn_up)?, ...) The sibling `quantized_phi3.rs`, rewritten in the same commit, kept its activation -- a transcription slip, not a decision. GELU, not SiLU, and specifically the tanh approximation: `models/phi2.rs` (the unquantized path) applies `cfg.hidden_act`, which is `gelu_new` in microsoft/phi-2's config.json, and llama.cpp builds phi2's FFN with `LLM_FFN_GELU`. Candle's `Tensor::gelu` is the tanh form; `gelu_erf` is not. Defect 2 alone would have been unobservable, since defect 1 meant the code never ran. Fixing 1 without 2 would have shipped a phi2 that loads and answers with a linear FFN. Tests (CPU, no GPU required): - `accepts_the_layer_norm_key_llama_cpp_writes` -- the file shape that used to be rejected outright. - `still_accepts_the_rms_spelling` -- the fix cannot regress anything that worked before. - `missing_epsilon_names_both_keys` - `gelu_is_not_a_linear_map` -- asserts GELU breaks additivity, i.e. that the activation is what stops the FFN collapsing to one matmul. - `gelu_is_the_tanh_approximation_not_erf` -- pins which GELU, since the two forms are separate functions in candle. UNVERIFIED ON HARDWARE: no GPU was used, and no phi2 GGUF was loaded end to end. The metadata-key claim is read from llama.cpp and gguf-py source; the activation claim is read from this repo's own git history and its unquantized sibling. Parent system: ArcModels. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01UMmjFy8TvsgypxVWNVhhC7 --- mistralrs-core/src/models/quantized_phi2.rs | 165 +++++++++++++++++++- 1 file changed, 162 insertions(+), 3 deletions(-) diff --git a/mistralrs-core/src/models/quantized_phi2.rs b/mistralrs-core/src/models/quantized_phi2.rs index 55e5c9cacf..5593de56a0 100644 --- a/mistralrs-core/src/models/quantized_phi2.rs +++ b/mistralrs-core/src/models/quantized_phi2.rs @@ -38,7 +38,25 @@ struct Mlp { impl Module for Mlp { fn forward(&self, xs: &Tensor) -> Result { - MatMul.qmethod_matmul(&MatMul.qmethod_matmul(xs, &*self.ffn_up)?, &*self.ffn_down) + // 🔴 The `.gelu()` between the two projections was deleted by + // `1269bd8ab` ("Implement GPTQ quantization", Aug 2024), which rewrote + // xs.apply(&self.ffn_up)?.gelu()?.apply(&self.ffn_down) + // into a bare `qmethod_matmul` composition and dropped the activation on + // the way through. The sibling `quantized_phi3.rs`, rewritten in the + // same commit, kept its activation -- this was a transcription slip, not + // a decision. + // + // Without it the FFN is two stacked linear maps, which collapse to a + // single linear map: the block contributes no nonlinearity at all. + // + // Phi-2 uses GELU, not SiLU: `models/phi2.rs` (the unquantized path) + // applies `cfg.hidden_act`, which is `gelu_new` in + // microsoft/phi-2's config.json, and llama.cpp builds the phi2 graph + // with `LLM_FFN_GELU` (`src/models/phi2.cpp`). Candle's `Tensor::gelu` + // is the tanh approximation, which is what `gelu_new` names; `gelu_erf` + // is the exact form and would NOT match. + let up = MatMul.qmethod_matmul(xs, &*self.ffn_up)?.gelu()?; + MatMul.qmethod_matmul(&up, &*self.ffn_down) } } @@ -184,19 +202,55 @@ struct PropsGGUF { max_seq_len: usize, } +/// Phi-2's LayerNorm epsilon, under whichever key the writer used. +/// +/// llama.cpp and `gguf-py` write `{arch}.attention.layer_norm_epsilon` for +/// phi2 (it is a LayerNorm architecture). The RMS spelling is accepted as a +/// fallback so a file produced by some other converter still loads, and an +/// error names both keys rather than only the one that happened to be checked. +fn layer_norm_eps(c: &ContentMetadata<'_>) -> std::result::Result { + if let Some(eps) = c.get_option_value::("attention.layer_norm_epsilon")? { + return Ok(eps as f64); + } + if let Some(eps) = c.get_option_value::("attention.layer_norm_rms_epsilon")? { + return Ok(eps as f64); + } + anyhow::bail!( + "phi2 GGUF is missing its LayerNorm epsilon: neither \ + `phi2.attention.layer_norm_epsilon` nor \ + `phi2.attention.layer_norm_rms_epsilon` is present" + ) +} + impl TryFrom> for PropsGGUF { type Error = anyhow::Error; fn try_from(c: ContentMetadata) -> std::result::Result { c.verify_arch("phi2")?; + // 🔴 `attention.layer_norm_rms_epsilon` used to be in this list, and it + // is a key phi2 GGUFs do not carry -- so `has_required_keys` failed on + // every real file and this loader was unreachable. That is what masked + // the missing FFN activation above. + // + // Phi-2 normalises with LayerNorm, not RMSNorm -- this very file builds + // a `candle_nn::LayerNorm` from `ln_eps`, complete with a bias. Writers + // agree: llama.cpp loads phi2's epsilon with + // `ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps)` and + // builds the graph with `LLM_NORM` (`src/models/phi2.cpp`), and + // `gguf-py` spells that key `{arch}.attention.layer_norm_epsilon` + // (`gguf/constants.py`). `quantized_starcoder2.rs`, the other + // LayerNorm-based GGUF loader here, already reads the non-RMS key. + // + // The epsilon is resolved below rather than listed here so that a file + // carrying either spelling loads, and a file carrying neither gets an + // error naming both. let required = [ "attention.head_count", "attention.head_count_kv", "block_count", "embedding_length", "rope.dimension_count", - "attention.layer_norm_rms_epsilon", "context_length", ]; c.has_required_keys(&required)?; @@ -209,7 +263,7 @@ impl TryFrom> for PropsGGUF { block_count: c.get_value::("block_count")? as usize, embedding_length: c.get_value::("embedding_length")? as usize, rope_dim: c.get_value::("rope.dimension_count")? as usize, - ln_eps: c.get_value::("attention.layer_norm_rms_epsilon")? as f64, + ln_eps: layer_norm_eps(&c)?, max_seq_len: c .get_value::("context_length") .ok() @@ -389,3 +443,108 @@ impl ModelWeights { self.output.forward(&xs) } } + +/// The two defects in the phi2 GGUF path, and the interaction between them. +/// +/// `PropsGGUF` demanded `phi2.attention.layer_norm_rms_epsilon`, a key phi2 +/// GGUFs do not carry -- Phi-2 is a LayerNorm architecture, and llama.cpp reads +/// `LLM_KV_ATTENTION_LAYERNORM_EPS` for it. So the loader rejected every real +/// file, which is why nobody noticed that `Mlp::forward` had lost its GELU. +#[cfg(test)] +mod quantized_phi2_tests { + use super::*; + use candle_core::quantized::gguf_file::Value; + + fn metadata(pairs: &[(&str, f32)]) -> std::collections::HashMap { + pairs + .iter() + .map(|(k, v)| (format!("phi2.{k}"), Value::F32(*v))) + .collect() + } + + /// The key llama.cpp actually writes must be accepted. Before the fix this + /// file was rejected outright. + #[test] + fn accepts_the_layer_norm_key_llama_cpp_writes() { + let md = metadata(&[("attention.layer_norm_epsilon", 1e-5)]); + let c = ContentMetadata { + path_prefix: "phi2", + metadata: &md, + }; + assert_eq!(layer_norm_eps(&c).unwrap(), 1e-5f32 as f64); + } + + /// A file carrying the RMS spelling still loads, so the fix cannot regress + /// anything that happened to work before. + #[test] + fn still_accepts_the_rms_spelling() { + let md = metadata(&[("attention.layer_norm_rms_epsilon", 1e-5)]); + let c = ContentMetadata { + path_prefix: "phi2", + metadata: &md, + }; + assert_eq!(layer_norm_eps(&c).unwrap(), 1e-5f32 as f64); + } + + /// Missing entirely is still an error, and the message names both keys + /// rather than only the one that used to be checked. + #[test] + fn missing_epsilon_names_both_keys() { + let md = metadata(&[("attention.head_count", 32.0)]); + let c = ContentMetadata { + path_prefix: "phi2", + metadata: &md, + }; + let err = layer_norm_eps(&c).unwrap_err().to_string(); + assert!(err.contains("layer_norm_epsilon"), "got: {err}"); + assert!(err.contains("layer_norm_rms_epsilon"), "got: {err}"); + } + + /// The FFN must be nonlinear. Two stacked linear maps compose to a single + /// linear map, so `Mlp` without its activation is provably indistinguishable + /// from one matmul -- which is what the missing `.gelu()` made it. + /// + /// Checked here on the activation itself rather than through `Mlp`, whose + /// fields are `Arc` and need a GGUF to build: GELU must + /// break additivity, `f(a + b) != f(a) + f(b)`. + #[test] + fn gelu_is_not_a_linear_map() { + let dev = Device::Cpu; + let a = Tensor::from_vec(vec![-2.0f32, -0.5, 0.5, 2.0], (1, 4), &dev).unwrap(); + let b = Tensor::from_vec(vec![1.0f32, 0.25, -0.75, 3.0], (1, 4), &dev).unwrap(); + + let lhs = (&a + &b).unwrap().gelu().unwrap(); + let rhs = (a.gelu().unwrap() + b.gelu().unwrap()).unwrap(); + let gap = (lhs - rhs) + .unwrap() + .abs() + .unwrap() + .max_all() + .unwrap() + .to_scalar::() + .unwrap(); + assert!( + gap > 1e-3, + "GELU behaved additively (gap {gap}) -- the FFN would collapse to one linear map" + ); + } + + /// Candle's `gelu` is the tanh approximation (`gelu_new`), which is what + /// Phi-2's `hidden_act` names and what llama.cpp's `LLM_FFN_GELU` uses. + /// `gelu_erf` is the exact form; picking it would be a different function. + /// This pins which one `Mlp::forward` must call. + #[test] + fn gelu_is_the_tanh_approximation_not_erf() { + let dev = Device::Cpu; + let x = Tensor::from_vec(vec![-3.0f32, -1.0, 0.5, 3.0], (1, 4), &dev).unwrap(); + let tanh: Vec = x.gelu().unwrap().flatten_all().unwrap().to_vec1().unwrap(); + let erf: Vec = x + .gelu_erf() + .unwrap() + .flatten_all() + .unwrap() + .to_vec1() + .unwrap(); + assert_ne!(tanh, erf, "the two GELU forms are indistinguishable here"); + } +} From 15fad01edf582c69c9df4bd5bd3eee10665297ed Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 03:25:24 +0100 Subject: [PATCH 7/9] fix(ArcModels): phi3 and starcoder2 GGUF read their logits from the padding MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both models ended `forward` with .i((.., seq_len - 1, ..)) `seq_len` is the batch MAXIMUM, not any particular sequence's length. On a ragged batch that takes the last row of the right-hand padding for every short sequence -- and the same wrong row for all of them -- so every sequence but the longest was sampled from a padded position. The repo already documents this as banned, at `pipeline/inputs_processor.rs`: 🔑 `new_len`, NOT `padded.len()`. `padded.len()` is the batch max, so on a ragged batch every short row would read its logits out of the right-hand PADDING — same wrong row for every sequence. `context_lens` carries the per-row `(start, len)` and `extract_logits` narrows row by row. It was simply never plumbed into these two `forward`s: `pipeline/gguf.rs` passed `context_lens` to Llama, Phi2, Qwen, Qwen3, Qwen3MoE and both X-LoRA arms, and skipped exactly these two. Extraction still happens BEFORE the output projection, so the "only the kept rows go through the vocab matmul" property of the original code is preserved. The result shape becomes `[b, len, vocab]`, matching every other GGUF model here. Uniform batches are unaffected: there every row's `start + len` is `seq_len`, so the change is the identity -- which is why this survived. Tests (CPU, no GPU required), in `pipeline/mod.rs` where `extract_logits` lives, on hidden states whose values encode `(batch, token)` so a wrong row is identifiable rather than merely unequal: - `ragged_batch_reads_each_sequences_own_last_row` - `last_row_indexing_reads_padding_on_a_ragged_batch` -- states the bug as a disagreement and asserts the two disagree, so it cannot pass vacuously. - `uniform_batch_is_unchanged` -- pins the identity case. UNVERIFIED ON HARDWARE: no GPU was used and no GGUF was loaded end to end. Note that the phi3 path is additionally gated by `quantized_phi3.rs`'s own metadata requirements, which were not audited here. Parent system: ArcModels. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01UMmjFy8TvsgypxVWNVhhC7 --- mistralrs-core/src/models/quantized_phi3.rs | 29 +++++-- .../src/models/quantized_starcoder2.rs | 29 +++++-- mistralrs-core/src/pipeline/gguf.rs | 4 +- mistralrs-core/src/pipeline/mod.rs | 76 +++++++++++++++++++ 4 files changed, 124 insertions(+), 14 deletions(-) diff --git a/mistralrs-core/src/models/quantized_phi3.rs b/mistralrs-core/src/models/quantized_phi3.rs index 7c3a3e2a08..f1d363a86d 100644 --- a/mistralrs-core/src/models/quantized_phi3.rs +++ b/mistralrs-core/src/models/quantized_phi3.rs @@ -9,7 +9,7 @@ use crate::layers::{CausalMasker, MatMul, RmsNorm, Sdpa}; use crate::layers_masker::PastKvLenCache; use crate::paged_attention::{AttentionImplementation, PagedAttention}; use crate::pipeline::text_models_inputs_processor::PagedAttentionInputMetadata; -use crate::pipeline::{EitherCache, KvCache, NormalCache}; +use crate::pipeline::{extract_logits, EitherCache, KvCache, NormalCache}; use crate::utils::gguf_metadata::ContentMetadata; use crate::utils::model_config as ModelConfig; use crate::utils::progress::{new_multi_progress, NiceProgressBar}; @@ -366,6 +366,7 @@ impl ModelWeights { &self, input_ids: &Tensor, seqlen_offsets: &[usize], + context_lens: Vec<(usize, usize)>, metadata: Option<(Vec<(Tensor, Tensor, Option, Option)>, &PagedAttentionInputMetadata)>, ) -> Result { let (_b_sz, seq_len) = input_ids.dims2()?; @@ -409,10 +410,26 @@ impl ModelWeights { let ys = layer.mlp.forward(&ys)?; xs = (ys + residual)? } - let xs = xs - .apply(&self.output_norm)? - .i((.., seq_len - 1, ..))? - .contiguous()?; - MatMul.qmatmul(&xs, &self.output) + // 🔴 This used to be `.i((.., seq_len - 1, ..))`: the LAST row of the + // padded batch, taken for every sequence. `seq_len` is the batch + // maximum, so on a ragged batch every short row read its logits out of + // the right-hand PADDING -- and the same wrong row for all of them. + // + // The repo already documents this as banned, at + // `pipeline/inputs_processor.rs`: "`new_len`, NOT `padded.len()`. + // `padded.len()` is the batch max, so on a ragged batch every short row + // would read its logits out of the right-hand PADDING". `context_lens` + // carries the per-row `(start, len)` that says where each sequence + // actually ends, and `extract_logits` narrows row by row. + // + // `context_lens` was simply never plumbed into this `forward` -- + // `pipeline/gguf.rs` passed it to every other GGUF architecture and + // skipped these two. Uniform batches are unaffected: there every row's + // `start + len` is `seq_len`, so this is the identity. + // + // Extracting BEFORE the output projection is deliberate and preserved: + // only the kept rows go through the vocab matmul. + let xs = extract_logits(&xs.apply(&self.output_norm)?, context_lens)?; + MatMul.qmatmul(&xs.contiguous()?, &self.output) } } diff --git a/mistralrs-core/src/models/quantized_starcoder2.rs b/mistralrs-core/src/models/quantized_starcoder2.rs index 21e2c2e169..8bea8fc857 100644 --- a/mistralrs-core/src/models/quantized_starcoder2.rs +++ b/mistralrs-core/src/models/quantized_starcoder2.rs @@ -10,7 +10,7 @@ use crate::layers::{CausalMasker, MatMul, QLinear, RotaryEmbedding, Sdpa}; use crate::layers_masker::PastKvLenCache; use crate::paged_attention::{AttentionImplementation, PagedAttention}; use crate::pipeline::text_models_inputs_processor::PagedAttentionInputMetadata; -use crate::pipeline::{EitherCache, KvCache, NormalCache}; +use crate::pipeline::{extract_logits, EitherCache, KvCache, NormalCache}; use crate::utils::gguf_metadata::ContentMetadata; use crate::utils::model_config as ModelConfig; use crate::utils::progress::{new_multi_progress, NiceProgressBar}; @@ -356,6 +356,7 @@ impl ModelWeights { &self, input_ids: &Tensor, seqlen_offsets: &[usize], + context_lens: Vec<(usize, usize)>, metadata: Option<(Vec<(Tensor, Tensor, Option, Option)>, &PagedAttentionInputMetadata)>, ) -> Result { let (_b_sz, seq_len) = input_ids.dims2()?; @@ -398,10 +399,26 @@ impl ModelWeights { let ys = layer.mlp.forward(&ys)?; xs = (ys + residual)? } - let xs = xs - .apply(&self.output_norm)? - .i((.., seq_len - 1, ..))? - .contiguous()?; - MatMul.qmatmul(&xs, &self.output) + // 🔴 This used to be `.i((.., seq_len - 1, ..))`: the LAST row of the + // padded batch, taken for every sequence. `seq_len` is the batch + // maximum, so on a ragged batch every short row read its logits out of + // the right-hand PADDING -- and the same wrong row for all of them. + // + // The repo already documents this as banned, at + // `pipeline/inputs_processor.rs`: "`new_len`, NOT `padded.len()`. + // `padded.len()` is the batch max, so on a ragged batch every short row + // would read its logits out of the right-hand PADDING". `context_lens` + // carries the per-row `(start, len)` that says where each sequence + // actually ends, and `extract_logits` narrows row by row. + // + // `context_lens` was simply never plumbed into this `forward` -- + // `pipeline/gguf.rs` passed it to every other GGUF architecture and + // skipped these two. Uniform batches are unaffected: there every row's + // `start + len` is `seq_len`, so this is the identity. + // + // Extracting BEFORE the output projection is deliberate and preserved: + // only the kept rows go through the vocab matmul. + let xs = extract_logits(&xs.apply(&self.output_norm)?, context_lens)?; + MatMul.qmatmul(&xs.contiguous()?, &self.output) } } diff --git a/mistralrs-core/src/pipeline/gguf.rs b/mistralrs-core/src/pipeline/gguf.rs index 1d91b8885e..cba17223de 100644 --- a/mistralrs-core/src/pipeline/gguf.rs +++ b/mistralrs-core/src/pipeline/gguf.rs @@ -813,7 +813,7 @@ impl Pipeline for GGUFPipeline { flash_meta_full.as_ref().unwrap_or(&flash_meta), )?, Model::Phi3(ref model) => { - model.forward(&input_ids, &seqlen_offsets, paged_attn_meta)? + model.forward(&input_ids, &seqlen_offsets, context_lens, paged_attn_meta)? } Model::XLoraPhi3(ref model) => model.forward( &input_ids, @@ -827,7 +827,7 @@ impl Pipeline for GGUFPipeline { flash_meta_full.as_ref().unwrap_or(&flash_meta), )?, Model::Starcoder2(ref model) => { - model.forward(&input_ids, &seqlen_offsets, paged_attn_meta)? + model.forward(&input_ids, &seqlen_offsets, context_lens, paged_attn_meta)? } Model::Qwen(ref model) => { model.forward(&input_ids, &seqlen_offsets, context_lens, paged_attn_meta)? diff --git a/mistralrs-core/src/pipeline/mod.rs b/mistralrs-core/src/pipeline/mod.rs index 8244606347..37110ed39e 100644 --- a/mistralrs-core/src/pipeline/mod.rs +++ b/mistralrs-core/src/pipeline/mod.rs @@ -1924,3 +1924,79 @@ mod host_copy_tests { } } } + +/// `extract_logits` vs the `.i((.., seq_len - 1, ..))` it replaced in the phi3 +/// and starcoder2 GGUF models. +/// +/// `seq_len` is the batch MAXIMUM. On a ragged batch, indexing it takes the +/// last row of the padding for every short sequence -- and the same row for all +/// of them. `pipeline/inputs_processor.rs` already documents this as banned; +/// those two models simply never had `context_lens` plumbed into their +/// `forward`, so they could not honour it. +#[cfg(test)] +mod extract_logits_tests { + use super::extract_logits; + use candle_core::{Device, IndexOp, Tensor}; + + /// `[batch, seq, hidden]` where element `(b, t, 0)` encodes `b * 100 + t`, + /// so a wrong row is identifiable rather than merely unequal. + fn hidden(batch: usize, seq: usize, hidden: usize) -> Tensor { + let mut v = Vec::with_capacity(batch * seq * hidden); + for b in 0..batch { + for t in 0..seq { + for h in 0..hidden { + v.push((b * 100 + t) as f32 + h as f32 / 1000.0); + } + } + } + Tensor::from_vec(v, (batch, seq, hidden), &Device::Cpu).unwrap() + } + + /// Three sequences of true lengths 4, 2 and 3 padded to 4. + fn ragged_context_lens() -> Vec<(usize, usize)> { + vec![(3, 1), (1, 1), (2, 1)] + } + + #[test] + fn ragged_batch_reads_each_sequences_own_last_row() { + let xs = hidden(3, 4, 2); + let got = extract_logits(&xs, ragged_context_lens()).unwrap(); + assert_eq!(got.dims(), &[3, 1, 2]); + + // Row b must be the token at that sequence's real end: 3, 1, 2. + let markers: Vec = got.i((.., 0, 0)).unwrap().to_vec1().unwrap(); + assert_eq!(markers, vec![3.0, 101.0, 202.0]); + } + + /// The bug, stated as a disagreement: `seq_len - 1` takes token 3 from every + /// row, which is padding for sequences 1 and 2. + #[test] + fn last_row_indexing_reads_padding_on_a_ragged_batch() { + let xs = hidden(3, 4, 2); + let seq_len = 4; + + let padded: Vec = xs.i((.., seq_len - 1, 0)).unwrap().to_vec1().unwrap(); + assert_eq!(padded, vec![3.0, 103.0, 203.0]); + + let correct = extract_logits(&xs, ragged_context_lens()).unwrap(); + let correct: Vec = correct.i((.., 0, 0)).unwrap().to_vec1().unwrap(); + assert_ne!( + padded, correct, + "input is degenerate -- it cannot distinguish the bug from the fix" + ); + // Only sequence 0 actually ends at the batch maximum, so only it agrees. + assert_eq!(padded[0], correct[0]); + assert_ne!(padded[1], correct[1]); + assert_ne!(padded[2], correct[2]); + } + + /// Uniform batches must be untouched -- this is where the old code was + /// accidentally right, and where the fix must be the identity. + #[test] + fn uniform_batch_is_unchanged() { + let xs = hidden(3, 4, 2); + let got = extract_logits(&xs, vec![(3, 1); 3]).unwrap(); + let markers: Vec = got.i((.., 0, 0)).unwrap().to_vec1().unwrap(); + assert_eq!(markers, vec![3.0, 103.0, 203.0]); + } +} From 6e682b77b33c387752509445a1d67a683489be9a Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 03:30:45 +0100 Subject: [PATCH 8/9] perf(ArcInfer/ArcAttention): V4 copied its whole KV cache twice, for one tensor Under V4's fused MQA the compressor stores ONE tensor and attention reads it as both K and V: three of `deepseek4.rs`'s four `dsv4_attention` call sites pass the same binding twice (`&k, &k` at :1836, `&k_full, &k_full` at :1885, `&k_cached, &k_cached` at :1977). `dsv4_attention` then materialised it twice, at two separate points: 1. the raw working-set narrowing -- `Some(v.narrow(2, rel_base, keep)?.contiguous()?)` next to the identical call for K; 2. the union build -- `Tensor::cat(&[v, comp], 2)?.contiguous()?` next to the identical call for K. Both allocate and copy the full retained key span, once per layer per step. The existing comment above (1) already names the cost -- "the `Tensor::cat` below, which copies the whole raw cache (twice, once for K and once for V) on every decode step" -- without noticing that when V is K the second copy is the same bytes as the first. The two are coupled, which is why this is not a one-line change at the `cat`: two independent `narrow(..).contiguous()` calls produce two distinct handles even when their input was one tensor, so an identity test placed after (1) always says "not aliased". The flag is computed once, before the narrowing, and threaded down. `v_aliases_k` is a HANDLE test, not a value test. Candle's `Tensor` is a refcounted handle carrying a `TensorId` that survives `clone`, so `k.id() == v.id()` holds exactly when the two arguments name one tensor with one layout -- the only condition under which the duplicated work is guaranteed bit-identical. Distinct tensors holding equal values decline and keep both copies, so this can only be conservative, never wrong. `deepseek4.rs:1821` passes genuinely distinct `&k_full, &v_full` and is unaffected. `fused_v_stats::counts()` reports `(aliased, materialized)`: skipping a copy leaves no trace in the output, so a harness that sees `aliased == 0` on a V4 decode must treat that as an environment failure rather than a result. Tests (CPU, no GPU required): - `alias_predicate_is_handle_identity` -- the decision, with no shared state: same binding and `clone` alias; `copy()` and a different tensor do not. - `aliased_v_is_bit_identical_and_engages` -- passing K twice must equal passing a distinct `copy()` of K, compared on raw IEEE-754 bit patterns, and must advance the counters (lower bounds; the counters are process-global and this file's other 60 tests drive them concurrently). - `distinct_v_is_not_aliased` -- negative control against "always reuse k_cat", which would silently compute attention over keys. UNVERIFIED ON HARDWARE: no GPU was used. The claim is that two of the copies of a tensor are provably redundant and are no longer made -- a count of allocations, not a measured time or a measured memory figure. No byte or millisecond number is asserted. Parent system: ArcInfer / ArcAttention. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01UMmjFy8TvsgypxVWNVhhC7 --- mistralrs-core/src/models/dsv4_attention.rs | 258 +++++++++++++++++++- 1 file changed, 253 insertions(+), 5 deletions(-) diff --git a/mistralrs-core/src/models/dsv4_attention.rs b/mistralrs-core/src/models/dsv4_attention.rs index 6d0e5bf72c..90e18bbfc9 100644 --- a/mistralrs-core/src/models/dsv4_attention.rs +++ b/mistralrs-core/src/models/dsv4_attention.rs @@ -644,6 +644,52 @@ pub struct Dsv4AttentionConfig<'a> { pub graph_positions: Option<&'a Tensor>, } +/// Engagement counters for the aliased-V fast path in [`dsv4_attention`]. +/// +/// House rule: "a green result must prove work happened". Skipping a `cat` +/// leaves no trace in the output -- the results are bit-identical either way -- +/// so a harness that sees `aliased == 0` on a V4 decode must treat that as an +/// environment failure rather than a result. +pub mod fused_v_stats { + use std::sync::atomic::{AtomicU64, Ordering}; + + static ALIASED: AtomicU64 = AtomicU64::new(0); + static MATERIALIZED: AtomicU64 = AtomicU64::new(0); + + /// `(aliased, materialized)` — union-key concatenations skipped because + /// V is K, vs genuinely built. + pub fn counts() -> (u64, u64) { + ( + ALIASED.load(Ordering::Relaxed), + MATERIALIZED.load(Ordering::Relaxed), + ) + } + + #[inline] + pub(super) fn record_aliased() { + ALIASED.fetch_add(1, Ordering::Relaxed); + } + + #[inline] + pub(super) fn record_materialized() { + MATERIALIZED.fetch_add(1, Ordering::Relaxed); + } +} + +/// Whether `v` is literally the same tensor as `k`, so every materialisation of +/// V duplicates one already performed for K. +/// +/// This is a HANDLE test, not a value test. Candle's `Tensor` is a reference- +/// counted handle carrying a unique `TensorId` that survives `clone`, so +/// `k.id() == v.id()` holds exactly when the two arguments name one tensor with +/// one layout -- the only condition under which duplicated work is guaranteed +/// bit-identical. Two distinct tensors that happen to hold equal values return +/// `false` and keep the copying path, so this can only be conservative. +#[inline] +pub(crate) fn v_aliases_k(k: &Tensor, v: &Tensor) -> bool { + k.id() == v.id() +} + /// V4 hybrid attention dispatch — a single softmax over the union of the raw /// sliding-window KV and the compressed KV, with the per-head `attn_sink` /// (carried in `sdpa_params.sinks`) as an extra softmax-denominator column. @@ -738,13 +784,42 @@ pub fn dsv4_attention( cfg.raw_prefix )) })?; + // 🔴 Under V4's fused MQA, V *is* K: three of `deepseek4.rs`'s four call + // sites pass the same binding twice (`&k, &k`; `&k_full, &k_full`; + // `&k_cached, &k_cached`), because the compressor stores one tensor and the + // attention reads it as both. The two full-cache materialisations below -- + // this `narrow(..).contiguous()` and the `Tensor::cat` further down -- each + // built a byte-identical duplicate of it, once per layer per step. The + // comment above already names the cost ("copies the whole raw cache, twice, + // once for K and once for V") without noticing that the second copy is the + // same bytes as the first. + // + // Identity is tested with `Tensor::id`, not by comparing values: candle's + // `Tensor` is a handle whose id survives `clone`, so `k.id() == v.id()` + // means the two arguments are literally the same tensor with the same + // layout -- the only condition under which the duplicated work is + // guaranteed to agree bit for bit. Distinct tensors that merely happen to + // hold equal values take the old path, so this can only be conservative, + // never wrong. `deepseek4.rs:1821` passes genuinely distinct + // `&k_full, &v_full` and still gets both copies. + // + // The flag is computed HERE, before the narrowing, and threaded down -- + // re-deriving it after this point would always say "not aliased", because + // two independent `narrow(..).contiguous()` calls produce two distinct + // handles even when their input was one tensor. + let v_aliases_k = v_aliases_k(k, v); let (k_owned, v_owned) = if rel_base == 0 && keep == t_k_given { (None, None) } else { - ( - Some(k.narrow(2, rel_base, keep)?.contiguous()?), - Some(v.narrow(2, rel_base, keep)?.contiguous()?), - ) + let k_narrowed = k.narrow(2, rel_base, keep)?.contiguous()?; + let v_narrowed = if v_aliases_k { + fused_v_stats::record_aliased(); + k_narrowed.clone() + } else { + fused_v_stats::record_materialized(); + v.narrow(2, rel_base, keep)?.contiguous()? + }; + (Some(k_narrowed), Some(v_narrowed)) }; let k = k_owned.as_ref().unwrap_or(k); let v = v_owned.as_ref().unwrap_or(v); @@ -792,7 +867,18 @@ pub fn dsv4_attention( }; let valid = Tensor::cat(&[&raw_valid, &comp_valid], 1)?; let k_cat = Tensor::cat(&[k, comp], 2)?.contiguous()?; - let v_cat = Tensor::cat(&[v, comp], 2)?.contiguous()?; + // The second half of the same duplication -- see `v_aliases_k` + // above. `v_cat` was an unconditional + // `Tensor::cat(&[v, comp], 2)?.contiguous()?`, allocating and + // copying a byte-identical duplicate of the entire union key + // tensor whenever V aliased K. + let v_cat = if v_aliases_k { + fused_v_stats::record_aliased(); + k_cat.clone() + } else { + fused_v_stats::record_materialized(); + Tensor::cat(&[v, comp], 2)?.contiguous()? + }; (k_cat, v_cat, valid) } None => (k.clone(), v.clone(), raw_valid), @@ -3127,3 +3213,165 @@ mod tests { Ok(()) } } + +/// The aliased-V fast path: under V4's fused MQA the caller passes the same +/// tensor as K and as V, and the second `Tensor::cat` that built `v_cat` was a +/// byte-identical duplicate of `k_cat`. +/// +/// Because the two paths agree exactly, skipping the cat is invisible in the +/// output -- which is why these tests assert on the `fused_v_stats` counters as +/// well as on the values. A test that only compared outputs would pass whether +/// or not the fast path ever engaged. +#[cfg(test)] +mod fused_v_alias_tests { + use super::*; + use candle_core::{DType, Device, Result, Tensor}; + use std::collections::HashMap; + use std::sync::{Mutex, MutexGuard, OnceLock}; + + /// The engagement counters are process-global and `cargo test` runs these + /// in parallel, 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 mk(b: usize, h: usize, t: usize, d: usize, seed: f32, dev: &Device) -> Result { + let data: Vec = (0..(b * h * t * d)) + .map(|i| ((i as f32) * seed).sin()) + .collect(); + Tensor::from_vec(data, (b, h, t, d), dev) + } + + fn flash() -> FlashParams { + FlashParams { + max_q: 0, + max_k: 0, + cumulative_seqlens_q: HashMap::new(), + cumulative_seqlens_k: HashMap::new(), + causal: false, + } + } + + fn cfg<'a>() -> Dsv4AttentionConfig<'a> { + Dsv4AttentionConfig { + compress_ratio: CompressRatio::Csa, + sliding_window: 128, + raw_prefix: 0, + row_q0: None, + graph_positions: None, + } + } + + fn sdpa(h: usize, d: usize, dev: &Device) -> Result { + let sinks: Vec = (0..h).map(|i| 0.5 + 0.25 * i as f32).collect(); + Ok(SdpaParams { + n_kv_groups: h, + softcap: None, + softmax_scale: 1.0 / (d as f32).sqrt(), + sliding_window: None, + sinks: Some(Tensor::from_vec(sinks, (1, h, 1, 1), dev)?), + }) + } + + /// Passing K twice (what `deepseek4.rs` does at three of its four call + /// sites) must produce exactly what passing two equal-valued but distinct + /// tensors produces -- and must take the fast path while doing it. + #[test] + fn aliased_v_is_bit_identical_and_engages() -> Result<()> { + let _g = guard(); + let dev = Device::Cpu; + let (b, h, d, t_k) = (1usize, 2usize, 8usize, 256usize); + let t_c = t_k / 4; + + let q = mk(b, h, 1, d, 0.03, &dev)?; + let k = mk(b, 1, t_k, d, 0.11, &dev)?; + let comp = mk(b, 1, t_c, d, 0.05, &dev)?; + let sdpa = sdpa(h, d, &dev)?; + + // A distinct tensor holding the same values: this is the slow path. + let v_copy = k.copy()?; + assert_ne!(k.id(), v_copy.id(), "copy() must yield a distinct handle"); + + let before = fused_v_stats::counts(); + let aliased = dsv4_attention(&q, &k, &k, Some(&comp), None, &flash(), &sdpa, cfg())?; + let mid = fused_v_stats::counts(); + let separate = dsv4_attention(&q, &k, &v_copy, Some(&comp), None, &flash(), &sdpa, cfg())?; + let after = fused_v_stats::counts(); + + // Two duplications are avoided per call, not one: the raw + // `narrow(..).contiguous()` and the union `cat`. The counters are + // process-global and the rest of this file's tests drive + // `dsv4_attention` too, so these are lower bounds -- the exact-delta + // claim is made by `alias_predicate_is_handle_identity` below, which + // has no shared state. + assert!( + mid.0 >= before.0 + 2, + "aliased call did not take both fast paths ({before:?} -> {mid:?})" + ); + assert!( + after.1 >= mid.1 + 2, + "distinct-tensor call must still materialize both copies of V" + ); + + let a: Vec = aliased.to_dtype(DType::F32)?.flatten_all()?.to_vec1()?; + let s: Vec = separate.to_dtype(DType::F32)?.flatten_all()?.to_vec1()?; + assert_eq!( + a.iter().map(|x| x.to_bits()).collect::>(), + s.iter().map(|x| x.to_bits()).collect::>(), + "skipping the duplicate cat changed the result" + ); + Ok(()) + } + + /// The decision itself, with no shared state: handle identity, not value + /// equality. This is the assertion the counter deltas cannot make, because + /// `dsv4_attention`'s counters are process-global and every other test in + /// this file drives them concurrently. + #[test] + fn alias_predicate_is_handle_identity() -> Result<()> { + let dev = Device::Cpu; + let k = mk(1, 1, 8, 4, 0.11, &dev)?; + + // What `deepseek4.rs` passes: one binding, twice. + assert!(v_aliases_k(&k, &k)); + // A clone is the same handle -- this is what the fix relies on when it + // threads `k_narrowed.clone()` through as V. + assert!(v_aliases_k(&k, &k.clone())); + // Equal values, different handle: must decline and keep copying. + let same_values = k.copy()?; + assert!(!v_aliases_k(&k, &same_values)); + // Different values: must decline. + let other = mk(1, 1, 8, 4, 0.17, &dev)?; + assert!(!v_aliases_k(&k, &other)); + Ok(()) + } + + /// A genuinely different V must NOT be aliased onto K. Guards against + /// "always reuse k_cat", which would silently compute self-attention over + /// keys. + #[test] + fn distinct_v_is_not_aliased() -> Result<()> { + let _g = guard(); + let dev = Device::Cpu; + let (b, h, d, t_k) = (1usize, 2usize, 8usize, 256usize); + let t_c = t_k / 4; + + let q = mk(b, h, 1, d, 0.03, &dev)?; + let k = mk(b, 1, t_k, d, 0.11, &dev)?; + let v = mk(b, 1, t_k, d, 0.17, &dev)?; + let comp = mk(b, 1, t_c, d, 0.05, &dev)?; + let sdpa = sdpa(h, d, &dev)?; + + let with_v = dsv4_attention(&q, &k, &v, Some(&comp), None, &flash(), &sdpa, cfg())?; + let with_k = dsv4_attention(&q, &k, &k, Some(&comp), None, &flash(), &sdpa, cfg())?; + + let a: Vec = with_v.to_dtype(DType::F32)?.flatten_all()?.to_vec1()?; + let bb: Vec = with_k.to_dtype(DType::F32)?.flatten_all()?.to_vec1()?; + assert_ne!(a, bb, "a different V produced the same output as V == K"); + Ok(()) + } +} From 9c51f50ab91379e4e4987794aab73d70e0712118 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 03:38:13 +0100 Subject: [PATCH 9/9] chore: rustfmt the code this wave added, and clear the warnings it introduced Scoped cleanup for the eight fixes in this branch. No behaviour change. - `layers.rs`, `gemma2.rs`: rustfmt applied to the lines this wave added. `cargo fmt --all` is NOT used -- it reformats ~90 upstream files, which the fork policy forbids. Both files were checked against `3c22cb5b5` first: `layers.rs` was rustfmt-clean at baseline and every deviation was in new code; `gemma2.rs` had 5 pre-existing own-file deviations and now has 4, because replacing the multi-line `NormalCache::new_sliding(..)` call happened to remove one. Every other file this branch touches was already unformatted upstream and its deviation count is unchanged. - `quantized_phi3.rs`, `quantized_starcoder2.rs`: `seq_len` and starcoder2's `IndexOp` import became unused when `.i((.., seq_len - 1, ..))` was replaced by `extract_logits`. - `lib.rs`: re-export `models::dsv4_attention::fused_v_stats`. `models` is a private module, so without this the engagement counters are dead code -- and shipping a counter nothing can read is the same failure the counter exists to prevent. Placed alongside the existing `pub use models::deepseek4::{..}` escape hatch, with a `//` comment rather than `///` so it does not split the sorted `use` group. Verified: `cargo test -p mistralrs-core --lib` 668 passed / 0 failed; `cargo check -p mistralrs-core --lib` at 24 warnings, the same count as `3c22cb5b5`; the scoped clippy lane (`--no-deps -p arc-bench -p arc-engine -p arc-cuda-graph -p arc-cli -p mistralrs-quant --tests --examples -- -D warnings`) green. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01UMmjFy8TvsgypxVWNVhhC7 --- mistralrs-core/src/layers.rs | 45 +++++++++++++------ mistralrs-core/src/lib.rs | 5 +++ mistralrs-core/src/models/gemma2.rs | 5 ++- mistralrs-core/src/models/quantized_phi3.rs | 2 +- .../src/models/quantized_starcoder2.rs | 4 +- 5 files changed, 43 insertions(+), 18 deletions(-) diff --git a/mistralrs-core/src/layers.rs b/mistralrs-core/src/layers.rs index 10ddf1f3e6..af68e3aa18 100644 --- a/mistralrs-core/src/layers.rs +++ b/mistralrs-core/src/layers.rs @@ -439,11 +439,7 @@ pub mod rope_cohort_stats { /// Uniform cohorts get one broadcast + one copy instead of `b` narrows and a /// `b`-way `cat`; ragged cohorts keep the `cat`, which is the only thing that /// can express them. -fn cohort_rotary_table( - table: &Tensor, - seqlen_offsets: &[usize], - seq_len: usize, -) -> Result { +fn cohort_rotary_table(table: &Tensor, seqlen_offsets: &[usize], seq_len: usize) -> Result { match uniform_seqlen_offset(seqlen_offsets) { Some(offset) => { let rows = table.narrow(0, offset, seq_len)?; @@ -3805,7 +3801,12 @@ mod rope_cohort_tests { /// The pre-fix per-sequence loop, verbatim, as the oracle. Always loops, /// whatever the offsets are. - fn reference(re: &RotaryEmbedding, q: &Tensor, k: &Tensor, offsets: &[usize]) -> (Tensor, Tensor) { + fn reference( + re: &RotaryEmbedding, + q: &Tensor, + k: &Tensor, + offsets: &[usize], + ) -> (Tensor, Tensor) { let seq_len = q.dim(2).unwrap(); let mut q_embeds = Vec::new(); let mut k_embeds = Vec::new(); @@ -3813,12 +3814,20 @@ mod rope_cohort_tests { let cos = re.cos.narrow(0, *offset, seq_len).unwrap(); let sin = re.sin.narrow(0, *offset, seq_len).unwrap(); q_embeds.push( - candle_nn::rotary_emb::rope(&q.i(i).unwrap().unsqueeze(0).unwrap().contiguous().unwrap(), &cos, &sin) - .unwrap(), + candle_nn::rotary_emb::rope( + &q.i(i).unwrap().unsqueeze(0).unwrap().contiguous().unwrap(), + &cos, + &sin, + ) + .unwrap(), ); k_embeds.push( - candle_nn::rotary_emb::rope(&k.i(i).unwrap().unsqueeze(0).unwrap().contiguous().unwrap(), &cos, &sin) - .unwrap(), + candle_nn::rotary_emb::rope( + &k.i(i).unwrap().unsqueeze(0).unwrap().contiguous().unwrap(), + &cos, + &sin, + ) + .unwrap(), ); } ( @@ -3860,8 +3869,16 @@ mod rope_cohort_tests { ); let (qref, kref) = reference(&re, &q, &k, &offsets); - assert_eq!(bits(&qb), bits(&qref), "batched q differs from per-sequence"); - assert_eq!(bits(&kb), bits(&kref), "batched k differs from per-sequence"); + assert_eq!( + bits(&qb), + bits(&qref), + "batched q differs from per-sequence" + ); + assert_eq!( + bits(&kb), + bits(&kref), + "batched k differs from per-sequence" + ); } /// Ragged offsets must still take the per-sequence loop and still be right. @@ -3902,8 +3919,8 @@ mod rope_cohort_tests { /// dispatch, so it gets the same guarantee. #[test] fn uniform_cohort_is_bit_identical_for_gpt_j_style() { - let re = - RotaryEmbedding::new(10000.0, D_HEAD, MAX_POS, &Device::Cpu, false, DType::F32).unwrap(); + let re = RotaryEmbedding::new(10000.0, D_HEAD, MAX_POS, &Device::Cpu, false, DType::F32) + .unwrap(); let q = xs(&[B, H, T, D_HEAD], 0.023); let k = xs(&[B, H, T, D_HEAD], 0.029); let offsets = vec![11usize; B]; diff --git a/mistralrs-core/src/lib.rs b/mistralrs-core/src/lib.rs index bb6c5fd8cb..26b8b14d96 100644 --- a/mistralrs-core/src/lib.rs +++ b/mistralrs-core/src/lib.rs @@ -160,6 +160,11 @@ pub use kv_cache::{ /// same reason `set_mtp_load_depth` is: the identity guard has to drive the /// same model both ways from one test binary. pub use models::deepseek4::{set_v4_kv_window, v4_kv_window_enabled}; +// `fused_v_stats`: engagement counters for V4's aliased-V fast path, reachable +// so a serving or benchmark harness can read them. Skipping a redundant KV copy +// leaves no trace in the output, so `aliased == 0` after a V4 decode is an +// environment failure, not a result. See the module's own docs. +pub use models::dsv4_attention::fused_v_stats; pub use request::{ ApproximateUserLocation, Constraint, DetokenizationRequest, ImageGenerationResponseFormat, LlguidanceGrammar, MessageContent, NormalRequest, ReasoningEffort, Request, RequestMessage, diff --git a/mistralrs-core/src/models/gemma2.rs b/mistralrs-core/src/models/gemma2.rs index 241039ac68..f753ad317e 100644 --- a/mistralrs-core/src/models/gemma2.rs +++ b/mistralrs-core/src/models/gemma2.rs @@ -933,7 +933,10 @@ mod gemma2_cache_type_tests { .count(); // 42 layers -> 21 SWA, 21 global. Not 42 rotating, which is what // `new_sliding` produced. - assert_eq!(rotating, 21, "expected half the layers to slide, got {rotating}"); + assert_eq!( + rotating, 21, + "expected half the layers to slide, got {rotating}" + ); for (layer_idx, slot) in cache.0.iter().enumerate() { if layer_idx.is_multiple_of(2) { diff --git a/mistralrs-core/src/models/quantized_phi3.rs b/mistralrs-core/src/models/quantized_phi3.rs index f1d363a86d..a521f19a59 100644 --- a/mistralrs-core/src/models/quantized_phi3.rs +++ b/mistralrs-core/src/models/quantized_phi3.rs @@ -369,7 +369,7 @@ impl ModelWeights { context_lens: Vec<(usize, usize)>, metadata: Option<(Vec<(Tensor, Tensor, Option, Option)>, &PagedAttentionInputMetadata)>, ) -> Result { - let (_b_sz, seq_len) = input_ids.dims2()?; + let (_b_sz, _seq_len) = input_ids.dims2()?; let mut xs = self.tok_embeddings.forward(input_ids)?; let cache = &mut self.cache.normal().0; let mask = CausalMasker.make_sliding_window_causal_mask_matrix( diff --git a/mistralrs-core/src/models/quantized_starcoder2.rs b/mistralrs-core/src/models/quantized_starcoder2.rs index 8bea8fc857..7c1a75ba8d 100644 --- a/mistralrs-core/src/models/quantized_starcoder2.rs +++ b/mistralrs-core/src/models/quantized_starcoder2.rs @@ -16,7 +16,7 @@ use crate::utils::model_config as ModelConfig; use crate::utils::progress::{new_multi_progress, NiceProgressBar}; use candle_core::quantized::QMatMul; use candle_core::quantized::QTensor; -use candle_core::{DType, Device, IndexOp, Module, Result, Tensor}; +use candle_core::{DType, Device, Module, Result, Tensor}; use candle_nn::{Embedding, LayerNorm}; use mistralrs_quant::{GgufMatMul, QuantMethod, QuantMethodConfig}; @@ -359,7 +359,7 @@ impl ModelWeights { context_lens: Vec<(usize, usize)>, metadata: Option<(Vec<(Tensor, Tensor, Option, Option)>, &PagedAttentionInputMetadata)>, ) -> Result { - let (_b_sz, seq_len) = input_ids.dims2()?; + let (_b_sz, _seq_len) = input_ids.dims2()?; let mut xs = self.tok_embeddings.forward(input_ids)?; let cache = &mut self.cache.normal().0; let mask = CausalMasker.make_causal_mask_matrix(