diff --git a/src/lib/mlxcel-core/src/cache.rs b/src/lib/mlxcel-core/src/cache.rs index 7728df347..991d10471 100644 --- a/src/lib/mlxcel-core/src/cache.rs +++ b/src/lib/mlxcel-core/src/cache.rs @@ -2615,6 +2615,258 @@ impl KVCache { /// the INT8 scale sidecars (their last dim is `1`); the head/tail bounds /// read every non-sequence dim from the array's own shape. Mirrors mlx-lm /// `RotatingKVCache._trim`. + /// Keep only the live-window slots named by `positions`, in that order, + /// compacted to the front of the buffer. + /// + /// The generalization of [`Self::trim`] that tree drafting needs. A linear + /// speculative accept set is a prefix, so rolling back is a tail trim; a + /// draft tree's accepted path is scattered across the verified positions, + /// so the surviving slots have to be selected and closed up (issue #1204). + /// + /// Supported for every dense mode, because every one of them stores its + /// per-token state on the same sequence axis: `Fp16` and `Int8` in + /// `keys` / `values` (plus Int8's scale sidecars), and the Turbo modes in + /// `v_packed` / `v_norms` / `v_rescale`, with `k_packed` / `k_norms` + /// alongside them in symmetric `Turbo4`. `trim` already slices all of + /// those on axis 2, so selecting a subset of that axis is the same + /// operation with an index list instead of a bound. + /// + /// Two cases refuse. `Turbo4Delegated` splits its V between packed cold + /// storage and a hot ring at `cold_offset`, so a selection spanning the + /// boundary has to rebuild it, which this does not attempt. Pool-backed + /// caches hold no dense buffer at all. A caller that gets an error is + /// expected to decline tree drafting rather than proceed. + /// + /// Note for anyone extending this: the `Fp16 | Int8` restriction on + /// [`Self::trim_front_keep`] does **not** apply here, and inheriting it + /// would be wrong. That one exists because a head trim advances + /// `live_start` instead of moving data, and the Turbo fetch paths slice + /// from `0`. This reindexes the buffers physically and leaves + /// `live_start` at `0`, so the hazard is not present. + /// + /// `positions` must be strictly increasing and inside the live window. A + /// tree path is increasing by construction (`parents[i] < i`), so a + /// violation is a caller bug rather than a shape this should tolerate. + pub fn gather_positions(&mut self, positions: &[i32]) -> Result<(), String> { + if self.paged_backing.is_some() { + return Err( + "KVCache::gather_positions: pool-backed caches have no dense buffer to reindex" + .to_string(), + ); + } + if self.mode == KVCacheMode::Turbo4Delegated { + return Err( + "KVCache::gather_positions: Turbo4Delegated splits V across packed cold storage \ + and a hot ring, and rebuilding cold_offset for a selection spanning both is not \ + implemented" + .to_string(), + ); + } + if self.live_start != 0 { + return Err(format!( + "KVCache::gather_positions: requires an untrimmed head (live_start={})", + self.live_start + )); + } + let live_len = self.live_len(); + if positions.is_empty() { + return Err("KVCache::gather_positions: positions must be non-empty".to_string()); + } + for pair in positions.windows(2) { + if pair[1] <= pair[0] { + return Err(format!( + "KVCache::gather_positions: positions must be strictly increasing, got {pair:?}" + )); + } + } + if let Some(&last) = positions.last() + && (positions[0] < 0 || last >= live_len) + { + return Err(format!( + "KVCache::gather_positions: positions {:?} outside the live window [0, {live_len})", + positions + )); + } + if positions.len() as i32 == live_len + && positions.iter().enumerate().all(|(i, &p)| p == i as i32) + { + // The identity selection. Skipping it keeps a full-accept round + // free of a buffer rewrite, mirroring the `trim` fast path. + return Ok(()); + } + + let index = ffi::from_slice_i32(positions, &[positions.len() as i32]); + let pick = |arr: &UniquePtr| -> UniquePtr { ffi::take(arr, &index, 2) }; + // Every buffer that carries a sequence axis is reindexed, whether or + // not the current mode populates it. Reindexing a subset would leave + // the survivors decoded against another token's scale, norm or sign + // vector, which is silent rather than loud. + self.keys = self.keys.as_ref().map(&pick); + self.values = self.values.as_ref().map(&pick); + self.key_scales = self.key_scales.as_ref().map(&pick); + self.val_scales = self.val_scales.as_ref().map(&pick); + self.v_packed = self.v_packed.as_ref().map(&pick); + self.v_norms = self.v_norms.as_ref().map(&pick); + self.v_rescale = self.v_rescale.as_ref().map(&pick); + self.k_packed = self.k_packed.as_ref().map(&pick); + self.k_norms = self.k_norms.as_ref().map(&pick); + self.offset = positions.len() as i32; + Ok(()) + } + + /// Keep only the tail slots named by `kept`, compacting them to the front + /// of the tail region and leaving everything before it alone. + /// + /// The scoped form of [`Self::gather_positions`], and the one a + /// speculative rollback actually wants. `gather_positions` names the + /// survivors of the whole live window, so a tree rollback has to re-list + /// the entire history every round and pays a full buffer rewrite for a + /// change confined to the last `tail_len` slots. This touches only that + /// region, through the same `slice_update` the append path uses, so the + /// cost scales with the verify block instead of with the context + /// (issue #1209). + /// + /// It also keeps the allocated capacity. `gather_positions` replaces each + /// buffer with a `take` result sized exactly to the survivors, so the next + /// append has to grow it again; writing back in place does not. + /// + /// `kept` are offsets **within** the tail, so `0` is the first slot of the + /// verified block and `tail_len - 1` the last. They must be strictly + /// increasing and inside `[0, tail_len)`. A draft tree's accepted path is + /// increasing by construction (`parents[i] < i`), so a violation is a + /// caller bug rather than a shape to tolerate. + /// + /// Unlike [`Self::gather_positions`] this places no requirement on + /// `live_start`: the tail is located in buffer coordinates through + /// `buffer_idx()`, which already accounts for a trimmed head. + /// + /// Refuses the two shapes `gather_positions` refuses, for the same + /// reasons: `Turbo4Delegated`, whose V is split between packed cold + /// storage and a hot ring, and pool-backed caches, which hold no dense + /// buffer. + pub fn gather_within_tail(&mut self, tail_len: i32, kept: &[i32]) -> Result<(), String> { + const WHO: &str = "KVCache::gather_within_tail"; + if self.paged_backing.is_some() { + return Err(format!( + "{WHO}: pool-backed caches have no dense buffer to reindex" + )); + } + if self.mode == KVCacheMode::Turbo4Delegated { + return Err(format!( + "{WHO}: Turbo4Delegated splits V across packed cold storage and a hot ring, and \ + rebuilding cold_offset for a selection spanning both is not implemented" + )); + } + Self::validate_tail_selection(WHO, tail_len, kept, self.live_len())?; + if Self::is_identity_selection(tail_len, kept) { + // Full accept. Mirrors the `trim` and `gather_positions` fast + // paths: no buffer touched and no offset moved. + return Ok(()); + } + + let tail_end = self.buffer_idx(); + let tail_start = tail_end - tail_len; + let kept_len = kept.len() as i32; + let index = ffi::from_slice_i32(kept, &[kept_len]); + let compact = |arr: &UniquePtr| -> UniquePtr { + Self::compact_tail_in_place(arr, &index, tail_start, tail_end, kept_len) + }; + // Every buffer carrying a sequence axis is rewritten, whether or not + // the current mode populates it. Rewriting a subset would leave the + // survivors decoded against another token's scale, norm or sign + // vector, which is silent rather than loud. + self.keys = self.keys.as_ref().map(&compact); + self.values = self.values.as_ref().map(&compact); + self.key_scales = self.key_scales.as_ref().map(&compact); + self.val_scales = self.val_scales.as_ref().map(&compact); + self.v_packed = self.v_packed.as_ref().map(&compact); + self.v_norms = self.v_norms.as_ref().map(&compact); + self.v_rescale = self.v_rescale.as_ref().map(&compact); + self.k_packed = self.k_packed.as_ref().map(&compact); + self.k_norms = self.k_norms.as_ref().map(&compact); + self.offset -= tail_len - kept_len; + Ok(()) + } + + /// Whether [`Self::gather_within_tail`] can serve this cache as it stands. + /// + /// Static here, unlike the rotating counterpart, but exposed with the same + /// name so a caller holding either kind asks one question. The point of + /// asking at all is the rotating case, where the answer changes with the + /// cache's state (issue #1204). + pub fn can_gather_within_tail(&self) -> bool { + self.paged_backing.is_none() && self.mode != KVCacheMode::Turbo4Delegated + } + + /// Shared validation for the tail-scoped selection, used by both dense + /// caches so the two cannot drift into disagreeing about what a legal + /// selection is. + fn validate_tail_selection( + who: &str, + tail_len: i32, + kept: &[i32], + live_len: i32, + ) -> Result<(), String> { + if tail_len <= 0 { + return Err(format!("{who}: tail_len must be positive (got {tail_len})")); + } + if tail_len > live_len { + return Err(format!( + "{who}: tail_len {tail_len} exceeds the live window [0, {live_len})" + )); + } + if kept.is_empty() { + return Err(format!( + "{who}: kept must be non-empty; a round that accepts nothing still keeps its root" + )); + } + for pair in kept.windows(2) { + if pair[1] <= pair[0] { + return Err(format!( + "{who}: kept must be strictly increasing, got {pair:?}" + )); + } + } + if let Some(&last) = kept.last() + && (kept[0] < 0 || last >= tail_len) + { + return Err(format!( + "{who}: kept {kept:?} outside the tail [0, {tail_len})" + )); + } + Ok(()) + } + + /// Whether the selection keeps the whole tail in order, i.e. the round was + /// a full accept and there is nothing to compact. + fn is_identity_selection(tail_len: i32, kept: &[i32]) -> bool { + kept.len() as i32 == tail_len && kept.iter().enumerate().all(|(i, &p)| p == i as i32) + } + + /// Rewrite `[tail_start, tail_start + kept_len)` with the selected slots of + /// `[tail_start, tail_end)`, leaving the rest of the buffer as it is. + /// + /// The slots past the new end keep stale contents, exactly as they do + /// after a `trim`: they are outside the live window and the next append + /// overwrites them. + fn compact_tail_in_place( + arr: &UniquePtr, + index: &UniquePtr, + tail_start: i32, + tail_end: i32, + kept_len: i32, + ) -> UniquePtr { + let s = ffi::array_shape(arr); + let tail = ffi::slice(arr, &[0, 0, tail_start, 0], &[s[0], s[1], tail_end, s[3]]); + let picked = ffi::take(&tail, index, 2); + ffi::slice_update( + arr, + &picked, + &[0, 0, tail_start, 0], + &[s[0], s[1], tail_start + kept_len, s[3]], + ) + } + fn slice_keep_sink(arr: &MlxArray, keep: i32, n: i32, live_len: i32) -> UniquePtr { let s = ffi::array_shape(arr); let head = ffi::slice(arr, &[0, 0, 0, 0], &[s[0], s[1], keep, s[3]]); @@ -5118,6 +5370,96 @@ impl RotatingKVCache { self.idx -= n; n } + + /// Keep only the tail slots named by `kept`, compacting them to the front + /// of the tail region. + /// + /// The rotating counterpart of [`KVCache::gather_within_tail`], and what a + /// draft tree needs from Gemma 4: its sliding-window layers hold this + /// cache, and [`Self::trim`] cannot serve a scattered accept set because + /// it moves no data at all, it only rewinds `offset` and `idx` and lets + /// the next append overwrite. A tail that has to be compacted has to be + /// written (issue #1204). + /// + /// Requires the ring to still be linear ([`Self::is_trimmable`]). That is + /// the precondition the speculative rollback path already runs under: + /// `enable_speculative_buffer` exists so the verify block append plus its + /// rollback stay inside the unwrapped region. Once the ring has wrapped, + /// physical slot order is not logical order, the tail is not contiguous, + /// and a selection expressed as tail offsets does not mean what it says. + /// + /// `start_position` is untouched, because compacting a tail does not move + /// slot `0`. + pub fn gather_within_tail(&mut self, tail_len: i32, kept: &[i32]) -> Result<(), String> { + const WHO: &str = "RotatingKVCache::gather_within_tail"; + if !self.tail_is_in_logical_order() { + return Err(format!( + "{WHO}: the ring has wrapped, so the tail is not contiguous and a selection \ + expressed as tail offsets cannot be compacted in place" + )); + } + if self.mode == KVCacheMode::Turbo4Delegated { + return Err(format!( + "{WHO}: Turbo4Delegated splits V across packed cold storage and a hot ring, and \ + rebuilding cold_offset for a selection spanning both is not implemented" + )); + } + KVCache::validate_tail_selection(WHO, tail_len, kept, self.idx)?; + if KVCache::is_identity_selection(tail_len, kept) { + return Ok(()); + } + + let tail_end = self.idx; + let tail_start = tail_end - tail_len; + let kept_len = kept.len() as i32; + let index = ffi::from_slice_i32(kept, &[kept_len]); + let compact = |arr: &UniquePtr| -> UniquePtr { + KVCache::compact_tail_in_place(arr, &index, tail_start, tail_end, kept_len) + }; + self.keys = self.keys.as_ref().map(&compact); + self.values = self.values.as_ref().map(&compact); + self.key_scales = self.key_scales.as_ref().map(&compact); + self.val_scales = self.val_scales.as_ref().map(&compact); + self.v_packed = self.v_packed.as_ref().map(&compact); + self.v_norms = self.v_norms.as_ref().map(&compact); + self.v_rescale = self.v_rescale.as_ref().map(&compact); + + let dropped = tail_len - kept_len; + self.offset -= dropped; + self.idx -= dropped; + Ok(()) + } + + /// Whether the last `n` physical slots are the last `n` logical tokens, in + /// order, for any `n` within the live window. + /// + /// Two regimes qualify, for different reasons. Buffered speculative mode + /// (`buffer_size > 0`) compacts from the front rather than wrapping, so + /// slot order is logical order by construction and `idx` is the live + /// length. Without the buffer the ring wraps in place, and only an + /// unwrapped one qualifies. + /// + /// [`Self::is_trimmable`] is **not** the right question here, despite + /// looking like it. It excludes buffered mode outright, deliberately, for + /// a snapshot-restore reason that has nothing to do with whether a tail is + /// contiguous. Gating on it refuses every MTP round, because + /// `enable_speculative_buffer` is what puts the cache in buffered mode in + /// the first place, and tree drafting runs in exactly that regime + /// (issue #1204). + pub fn tail_is_in_logical_order(&self) -> bool { + self.buffer_size > 0 || self.is_trimmable() + } + + /// Whether [`Self::gather_within_tail`] can serve this cache **right now**. + /// + /// State-dependent, and that is the property a caller has to respect: a + /// tree rollback available this round can be unavailable the next one, as + /// soon as an unbuffered ring wraps. Ask before appending a verify block, + /// not after. A refusal once the block is on the cache leaves a state only + /// a tree-aware rollback could undo (issue #1204). + pub fn can_gather_within_tail(&self) -> bool { + self.tail_is_in_logical_order() && self.mode != KVCacheMode::Turbo4Delegated + } } impl Default for RotatingKVCache { diff --git a/src/lib/mlxcel-core/src/cache/detach_tests.rs b/src/lib/mlxcel-core/src/cache/detach_tests.rs index 407dfa6b4..2b1746567 100644 --- a/src/lib/mlxcel-core/src/cache/detach_tests.rs +++ b/src/lib/mlxcel-core/src/cache/detach_tests.rs @@ -1046,3 +1046,484 @@ fn detached_cache_set_truncate_to_int8_preserves_dequantization() { "INT8 truncate_to+adopt+decode values diverge from fresh prefill" ); } + +// --------------------------------------------------------------------------- +// KVCache::gather_positions (issue #1204) +// --------------------------------------------------------------------------- + +#[test] +fn gather_positions_keeps_a_scattered_subset_and_compacts_it() { + // The case a linear trim cannot express: a draft tree's accepted path is + // scattered across the verified positions, so the survivors have to be + // selected and closed up rather than truncated from the tail. + let mut cache = KVCache::new(); + cache.update( + fp32_tokens(&[1.0, 2.0, 3.0, 4.0, 5.0]), + fp32_tokens(&[10.0, 20.0, 30.0, 40.0, 50.0]), + ); + + cache + .gather_positions(&[0, 2, 3]) + .expect("gather must succeed"); + + assert_eq!(cache.seq_len(), 3, "the cache holds exactly the survivors"); + assert_eq!( + visible_keys_as_fp32(&cache), + vec![1.0, 3.0, 4.0], + "survivors keep their order and are compacted to the front" + ); +} + +#[test] +fn gather_positions_of_a_prefix_matches_trim_to() { + // A linear accepted set is a prefix, so the tree path must agree with the + // chain path there. This is the equivalence that lets one rollback serve + // both once the tree path is wired up. + let mut gathered = KVCache::new(); + let mut trimmed = KVCache::new(); + for cache in [&mut gathered, &mut trimmed] { + cache.update( + fp32_tokens(&[1.0, 2.0, 3.0, 4.0]), + fp32_tokens(&[10.0, 20.0, 30.0, 40.0]), + ); + } + + gathered + .gather_positions(&[0, 1]) + .expect("gather must succeed"); + trimmed.trim_to(2).expect("trim_to must succeed"); + + assert_eq!(gathered.seq_len(), trimmed.seq_len()); + assert_eq!( + visible_keys_as_fp32(&gathered), + visible_keys_as_fp32(&trimmed) + ); +} + +#[test] +fn gather_positions_of_everything_is_a_no_op() { + let mut cache = KVCache::new(); + cache.update( + fp32_tokens(&[1.0, 2.0, 3.0]), + fp32_tokens(&[10.0, 20.0, 30.0]), + ); + + cache + .gather_positions(&[0, 1, 2]) + .expect("identity gather must succeed"); + + assert_eq!(cache.seq_len(), 3); + assert_eq!(visible_keys_as_fp32(&cache), vec![1.0, 2.0, 3.0]); +} + +#[test] +fn gather_positions_carries_the_int8_scale_sidecars() { + // Reindexing K and V while leaving their per-token scales behind would + // decode every surviving token against the wrong scale, which is exactly + // the class of silent corruption this cache's Turbo refusal exists to + // prevent. + let mut cache = KVCache::new_with_mode(KVCacheMode::Int8); + cache.update( + fp32_tokens(&[1.0, 2.0, 3.0, 4.0]), + fp32_tokens(&[10.0, 20.0, 30.0, 40.0]), + ); + + cache + .gather_positions(&[1, 3]) + .expect("gather must succeed"); + + assert_eq!(cache.seq_len(), 2); + // The raw INT8 codes carry no evidence here: a per-token scale normalizes + // every single-valued token to the same code. The scales are where the + // surviving identity lives, so that is what has to be reindexed with them. + for (name, sidecar) in [ + ("key_scales", cache.key_scales.as_ref()), + ("val_scales", cache.val_scales.as_ref()), + ] { + let sidecar = sidecar.unwrap_or_else(|| panic!("{name} must survive the gather")); + let shape = crate::ffi::array_shape(sidecar); + assert_eq!( + shape[2], 2, + "{name} must be reindexed alongside the tokens, got shape {shape:?}" + ); + } +} + +#[test] +fn gather_positions_rejects_shapes_it_cannot_serve() { + let mut cache = KVCache::new(); + cache.update( + fp32_tokens(&[1.0, 2.0, 3.0]), + fp32_tokens(&[10.0, 20.0, 30.0]), + ); + + assert!( + cache.gather_positions(&[]).is_err(), + "an empty selection would leave no context to decode from" + ); + assert!( + cache.gather_positions(&[1, 1]).is_err(), + "a repeated position is a caller bug, not a shape to tolerate" + ); + assert!( + cache.gather_positions(&[2, 0]).is_err(), + "a tree path is increasing by construction" + ); + assert!( + cache.gather_positions(&[0, 5]).is_err(), + "a position past the live window must not be reindexed" + ); + // The cache is untouched by every rejection above. + assert_eq!(cache.seq_len(), 3); + assert_eq!(visible_keys_as_fp32(&cache), vec![1.0, 2.0, 3.0]); +} + +#[test] +fn gather_positions_reindexes_every_turbo_sidecar() { + // The Turbo modes keep their per-token state on the same sequence axis + // `trim` already slices, so a selection along it is the same operation + // with an index list instead of a bound. Every sidecar the mode populates + // has to come along: a survivor decoded against another token's norm or + // sign vector fails silently rather than loudly. + // head_dim 8 rather than the 1 the plain fixture uses: Turbo3Asym packs + // 24 bits across the feature axis and rejects anything narrower, so a + // head_dim of 1 would test the fixture rather than the gather. + let wide = |tokens: usize, base: f32| -> UniquePtr { + let values: Vec = (0..tokens) + .flat_map(|t| (0..8).map(move |d| base + t as f32 + d as f32 * 0.01)) + .collect(); + crate::ffi::from_slice_f32(&values, &[1, 1, tokens as i32, 8]) + }; + + for mode in [ + KVCacheMode::Turbo4, + KVCacheMode::Turbo4Asym, + KVCacheMode::Turbo3Asym, + ] { + let mut cache = KVCache::new_with_mode(mode); + cache.update(wide(4, 1.0), wide(4, 10.0)); + let before = cache.seq_len(); + + cache + .gather_positions(&[0, 2]) + .unwrap_or_else(|e| panic!("{mode:?} must support gather, got: {e}")); + + assert_eq!(cache.seq_len(), 2, "{mode:?} keeps exactly the survivors"); + assert!(before > 2, "{mode:?} fixture must start wider than it ends"); + for (name, sidecar) in [ + ("v_packed", cache.v_packed.as_ref()), + ("v_norms", cache.v_norms.as_ref()), + ("v_rescale", cache.v_rescale.as_ref()), + ("k_packed", cache.k_packed.as_ref()), + ("k_norms", cache.k_norms.as_ref()), + ] { + let Some(sidecar) = sidecar else { continue }; + let shape = crate::ffi::array_shape(sidecar); + assert_eq!( + shape[2], 2, + "{mode:?} left {name} at shape {shape:?}, so it no longer describes the tokens" + ); + } + } +} + +#[test] +fn gather_positions_refuses_the_split_delegated_cache() { + // Turbo4Delegated is the one dense mode that refuses, and for a reason + // that is about its layout rather than about quantization: V lives partly + // in packed cold storage and partly in a hot ring, so a selection + // spanning `cold_offset` has to rebuild that boundary. + let mut cache = KVCache::new_with_mode(KVCacheMode::Turbo4Delegated); + cache.update( + fp32_tokens(&[1.0, 2.0, 3.0]), + fp32_tokens(&[10.0, 20.0, 30.0]), + ); + + let err = cache + .gather_positions(&[0, 2]) + .expect_err("the split layout must refuse rather than reindex one half"); + assert!( + err.contains("cold_offset"), + "the refusal should name the boundary it cannot rebuild, got: {err}" + ); +} + +// --------------------------------------------------------------------------- +// Tail-scoped gather (issue #1204, cost side of issue #1209) +// --------------------------------------------------------------------------- + +/// Visible keys of a rotating cache while its ring is still linear. +/// +/// Only valid unwrapped, which is the regime the speculative path keeps itself +/// in; once the ring wraps, slot order stops being logical order. +fn visible_rotating_keys_as_fp32(cache: &RotatingKVCache) -> Vec { + let keys = cache.keys.as_ref().expect("cache should have keys"); + let shape = crate::ffi::array_shape(keys); + let sliced = crate::ffi::slice( + keys, + &[0, 0, 0, 0], + &[shape[0], shape[1], cache.offset, shape[3]], + ); + flatten_fp32(&sliced) +} + +#[test] +fn gather_within_tail_keeps_a_scattered_subset_of_the_block() { + // Two tokens of history followed by a four-slot verify block. The accepted + // path takes the block's slots 0, 2 and 3, which no trim can express. + let mut cache = KVCache::new(); + cache.update( + fp32_tokens(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]), + fp32_tokens(&[10.0, 20.0, 30.0, 40.0, 50.0, 60.0]), + ); + + cache + .gather_within_tail(4, &[0, 2, 3]) + .expect("gather must succeed"); + + assert_eq!( + cache.seq_len(), + 5, + "two history tokens plus three survivors" + ); + assert_eq!( + visible_keys_as_fp32(&cache), + vec![1.0, 2.0, 3.0, 5.0, 6.0], + "the history is untouched and the block's survivors are compacted" + ); +} + +#[test] +fn gather_within_tail_agrees_with_the_whole_window_form() { + // The scoped form exists to avoid re-listing the history every round, not + // to mean something different. Naming the same survivors either way has to + // land in the same place, or the cheap path is not the same rollback. + let mut scoped = KVCache::new(); + let mut whole = KVCache::new(); + for cache in [&mut scoped, &mut whole] { + cache.update( + fp32_tokens(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]), + fp32_tokens(&[10.0, 20.0, 30.0, 40.0, 50.0, 60.0]), + ); + } + + scoped + .gather_within_tail(4, &[0, 2, 3]) + .expect("scoped gather must succeed"); + // The same selection said in absolute terms: keep the history, then the + // block slots at 2 + {0, 2, 3}. + whole + .gather_positions(&[0, 1, 2, 4, 5]) + .expect("whole-window gather must succeed"); + + assert_eq!(scoped.seq_len(), whole.seq_len()); + assert_eq!( + visible_keys_as_fp32(&scoped), + visible_keys_as_fp32(&whole), + "the scoped and whole-window forms must agree on the surviving tokens" + ); +} + +#[test] +fn gather_within_tail_of_a_prefix_matches_trim() { + // A linear accept set is a prefix of the block, so the tree rollback has + // to agree with the chain rollback there. This is the equivalence that + // lets one path serve both once the round loop is wired up. + let mut gathered = KVCache::new(); + let mut trimmed = KVCache::new(); + for cache in [&mut gathered, &mut trimmed] { + cache.update( + fp32_tokens(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]), + fp32_tokens(&[10.0, 20.0, 30.0, 40.0, 50.0, 60.0]), + ); + } + + gathered + .gather_within_tail(4, &[0, 1]) + .expect("gather must succeed"); + trimmed.trim(2); + + assert_eq!(gathered.seq_len(), trimmed.seq_len()); + assert_eq!( + visible_keys_as_fp32(&gathered), + visible_keys_as_fp32(&trimmed) + ); +} + +#[test] +fn gather_within_tail_of_the_whole_block_is_a_no_op() { + // Full accept. The round pays no buffer rewrite, mirroring the `trim` and + // `gather_positions` fast paths. + let mut cache = KVCache::new(); + cache.update( + fp32_tokens(&[1.0, 2.0, 3.0, 4.0]), + fp32_tokens(&[10.0, 20.0, 30.0, 40.0]), + ); + + cache + .gather_within_tail(3, &[0, 1, 2]) + .expect("gather must succeed"); + + assert_eq!(cache.seq_len(), 4); + assert_eq!(visible_keys_as_fp32(&cache), vec![1.0, 2.0, 3.0, 4.0]); +} + +#[test] +fn gather_within_tail_rejects_selections_it_cannot_serve() { + let mut cache = KVCache::new(); + cache.update( + fp32_tokens(&[1.0, 2.0, 3.0, 4.0]), + fp32_tokens(&[10.0, 20.0, 30.0, 40.0]), + ); + + assert!( + cache.gather_within_tail(0, &[0]).is_err(), + "a zero-length tail names no block" + ); + assert!( + cache.gather_within_tail(5, &[0]).is_err(), + "a tail longer than the live window is a caller bug" + ); + assert!( + cache.gather_within_tail(3, &[]).is_err(), + "a round that accepts nothing still keeps its root" + ); + assert!( + cache.gather_within_tail(3, &[1, 1]).is_err(), + "a tree path is strictly increasing" + ); + assert!( + cache.gather_within_tail(3, &[0, 3]).is_err(), + "offsets are within the tail, not within the cache" + ); + // None of the refusals may have moved the cache. + assert_eq!(cache.seq_len(), 4); + assert_eq!(visible_keys_as_fp32(&cache), vec![1.0, 2.0, 3.0, 4.0]); +} + +#[test] +fn rotating_gather_within_tail_compacts_the_block() { + // The case that decides whether Gemma 4 can roll a tree back at all: its + // sliding-window layers hold this cache, and `trim` moves no data. + let mut cache = RotatingKVCache::new(64); + let _ = cache.update_and_fetch( + fp32_tokens(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]), + fp32_tokens(&[10.0, 20.0, 30.0, 40.0, 50.0, 60.0]), + ); + + cache + .gather_within_tail(4, &[0, 2, 3]) + .expect("gather must succeed on an unwrapped ring"); + + assert_eq!(cache.offset, 5); + assert_eq!( + visible_rotating_keys_as_fp32(&cache), + vec![1.0, 2.0, 3.0, 5.0, 6.0], + "the history is untouched and the block's survivors are compacted" + ); +} + +#[test] +fn rotating_gather_within_tail_of_a_prefix_matches_trim() { + let mut gathered = RotatingKVCache::new(64); + let mut trimmed = RotatingKVCache::new(64); + for cache in [&mut gathered, &mut trimmed] { + let _ = cache.update_and_fetch( + fp32_tokens(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]), + fp32_tokens(&[10.0, 20.0, 30.0, 40.0, 50.0, 60.0]), + ); + } + + gathered + .gather_within_tail(4, &[0, 1]) + .expect("gather must succeed"); + trimmed.trim(2); + + assert_eq!(gathered.offset, trimmed.offset); + assert_eq!( + visible_rotating_keys_as_fp32(&gathered), + visible_rotating_keys_as_fp32(&trimmed) + ); +} + +#[test] +fn rotating_gather_within_tail_refuses_a_wrapped_ring() { + // Once the ring has wrapped, physical slot order is not logical order, so + // a selection expressed as tail offsets does not mean what it says. The + // speculative path avoids this with `enable_speculative_buffer`; a caller + // that has not gets a refusal rather than a silently wrong window. + let mut cache = RotatingKVCache::new(4); + for value in [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0] { + let _ = cache.update_and_fetch(fp32_tokens(&[value]), fp32_tokens(&[value * 10.0])); + } + assert!(!cache.is_trimmable(), "the ring must have wrapped"); + + let refusal = cache + .gather_within_tail(2, &[0]) + .expect_err("a wrapped ring must refuse rather than compact the wrong slots"); + assert!(refusal.contains("wrapped"), "the refusal names its cause"); +} + +#[test] +fn rotating_gather_within_tail_serves_the_buffered_speculative_regime() { + // The regime that actually matters, and the one an `is_trimmable` gate + // silently refused: Gemma 4's MTP path calls `enable_speculative_buffer` + // so a verify block can append and roll back without destroying window + // entries, and `is_unwrapped` excludes buffered mode outright. Gating the + // gather on it made every tree round decline while the flag looked on + // (issue #1204). + let mut cache = RotatingKVCache::new(64); + cache + .enable_speculative_buffer(16) + .expect("speculative buffer must enable"); + assert!( + !cache.is_trimmable(), + "buffered mode is deliberately not `is_trimmable`; that is the trap" + ); + assert!( + cache.tail_is_in_logical_order(), + "buffered mode compacts from the front, so the tail is still contiguous" + ); + + for value in [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0] { + let _ = cache.update_and_fetch(fp32_tokens(&[value]), fp32_tokens(&[value * 10.0])); + } + + cache + .gather_within_tail(4, &[0, 2, 3]) + .expect("a buffered rotating cache must serve a tail selection"); + + assert_eq!(cache.offset, 5); + assert_eq!( + visible_rotating_keys_as_fp32(&cache), + vec![1.0, 2.0, 3.0, 5.0, 6.0], + "history untouched, block survivors compacted" + ); +} + +#[test] +fn rotating_gather_within_tail_matches_trim_in_the_buffered_regime() { + // The prefix equivalence again, in the regime the speculative path runs + // in rather than only in the plain one. + let mut gathered = RotatingKVCache::new(64); + let mut trimmed = RotatingKVCache::new(64); + for cache in [&mut gathered, &mut trimmed] { + cache + .enable_speculative_buffer(16) + .expect("speculative buffer must enable"); + for value in [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0] { + let _ = cache.update_and_fetch(fp32_tokens(&[value]), fp32_tokens(&[value * 10.0])); + } + } + + gathered + .gather_within_tail(4, &[0, 1]) + .expect("gather must succeed"); + trimmed.trim(2); + + assert_eq!(gathered.offset, trimmed.offset); + assert_eq!( + visible_rotating_keys_as_fp32(&gathered), + visible_rotating_keys_as_fp32(&trimmed) + ); +} diff --git a/src/lib/mlxcel-core/src/drafter/mod.rs b/src/lib/mlxcel-core/src/drafter/mod.rs index 3ebd4a1e4..f88fdc2ba 100644 --- a/src/lib/mlxcel-core/src/drafter/mod.rs +++ b/src/lib/mlxcel-core/src/drafter/mod.rs @@ -963,6 +963,33 @@ pub trait Drafter { sampler: &crate::generate::SamplingConfig, ) -> Result, DrafterError>; + /// Propose a draft **tree** rather than a chain. + /// + /// The default wraps [`Self::draft_block`] into a linear tree, so every + /// drafter can serve the tree path from day one and a drafter that + /// branches overrides this. That degenerate case is not just a + /// convenience: it is what makes the tree wiring testable, because a + /// linear tree must produce byte-identical output to the chain path it + /// subsumes, and any divergence is the wiring rather than the topology + /// (issue #1204). + /// + /// Branching needs a confidence signal the drafter already computes and + /// discards: its per-step logits are reduced to an argmax today, and the + /// runner-up at an uncertain step is exactly the alternative worth + /// spending a verify position on. + fn draft_tree( + &mut self, + last_bonus: i32, + hidden: Option<&MlxArray>, + block_size: usize, + sampler: &crate::generate::SamplingConfig, + ) -> Result { + let drafts = self.draft_block(last_bonus, hidden, block_size, sampler)?; + Ok(crate::speculative::mtp::tree::DraftTree::linear( + last_bonus, &drafts, + )) + } + /// Produce a draft block as a device-side token array. /// /// DFlash can feed proposal tokens directly into the target verify graph, diff --git a/src/lib/mlxcel-core/src/speculative/mtp/generator.rs b/src/lib/mlxcel-core/src/speculative/mtp/generator.rs index 916eb79af..d92bef337 100644 --- a/src/lib/mlxcel-core/src/speculative/mtp/generator.rs +++ b/src/lib/mlxcel-core/src/speculative/mtp/generator.rs @@ -48,7 +48,8 @@ use std::time::{Duration, Instant}; use super::adaptive::effective_mtp_block_size; use super::target::{MtpTarget, MtpVerifyOutput}; -use super::walk::speculative_walk; +use super::tree::DraftTree; +use super::walk::{WalkResult, speculative_walk}; #[derive(Debug, Default)] struct MtpRoundDiagnostics { @@ -360,6 +361,14 @@ pub struct MtpSessionState { /// Once per session: the condition holds every round it holds at all, and /// a per-round warning would bury the run's own diagnostics. warned_block_override: bool, + /// Whether this session has already reported taking the tree path. + /// + /// Logged once, at a level an operator sees, because "the flag is set" and + /// "the tree path ran" are different claims: the per-round capability + /// question can answer no every round and leave the session silently on + /// the chain path. A linear-tree equivalence check that cannot tell those + /// apart proves nothing (issue #1204). + announced_tree_round: bool, /// Classic-step probe rounds still to run (issue #736). probes_remaining: usize, /// Round diagnostics, accumulated across steps exactly as the @@ -481,6 +490,30 @@ pub struct MtpGenerator { /// measured-cost estimator has a classic single-token step time to divide /// by. See [`Self::with_profile_probe_rounds`]. profile_probe_rounds: usize, + /// Whether this process may propose a draft tree instead of a chain. + /// + /// Read once at construction from `MLXCEL_MTP_TREE`, off unless set. Being + /// off is not the same as the target refusing: this is the operator switch, + /// and [`MtpTarget::tree_round_is_available`] is the per-round capability + /// question asked on top of it (issue #1204). + tree_drafting: bool, +} + +/// Whether tree drafting is switched on for this process. +/// +/// Off by default, because the tree path is new and the chain path it subsumes +/// is what every measurement so far rests on. The default `Drafter::draft_tree` +/// builds a linear tree, so switching this on without a branching drafter must +/// produce byte-identical output to leaving it off; that equivalence is the +/// first thing to check, and a divergence is the wiring rather than the +/// topology (issue #1204). +fn tree_drafting_enabled() -> bool { + std::env::var("MLXCEL_MTP_TREE") + .map(|value| { + let value = value.trim().to_ascii_lowercase(); + value == "1" || value == "true" || value == "on" || value == "yes" + }) + .unwrap_or(false) } impl MtpGenerator { @@ -505,6 +538,7 @@ impl MtpGenerator { prefer_requested_block_size, last_acceptance: None, profile_probe_rounds: 0, + tree_drafting: tree_drafting_enabled(), } } @@ -523,6 +557,18 @@ impl MtpGenerator { self } + /// Override the `MLXCEL_MTP_TREE` switch for this generator. + /// + /// The environment is read once at construction, which is right for a + /// process-wide operator switch and useless for a test that has to run + /// both arms in one process. Tests use this to pin the arm instead of + /// mutating the environment, which is not safe to do from a parallel + /// suite (issue #1204). + pub fn with_tree_drafting(mut self, enabled: bool) -> Self { + self.tree_drafting = enabled; + self + } + /// Block size (K). Test/diagnostic accessor. pub fn block_size(&self) -> usize { self.block_size @@ -758,6 +804,7 @@ impl MtpGenerator { emitted_count: 1, accept_lens: Vec::new(), warned_block_override: false, + announced_tree_round: false, probes_remaining: self.profile_probe_rounds, diagnostics, prefill_time, @@ -929,9 +976,35 @@ impl MtpGenerator { // classic decode. Probes are recorded separately from the // speculative aggregates below. let is_probe = state.probes_remaining > 0; - let draft_tokens: Vec = if is_probe { + // A round proposes either a chain or a tree. The capability question + // is asked here, before anything is drafted, because the target's + // rollback moves with its cache state and the verify forward has + // already appended the block by the time a refusal could surface + // (issue #1204). + let tree_available = + !is_probe && self.tree_drafting && self.target.tree_round_is_available(); + let use_tree = tree_available; + // Report the outcome once per session, at a level an operator sees. + // "The flag is set" and "the tree path ran" are different claims: the + // per-round capability question can answer no every round and leave + // the session silently on the chain path, which would make a + // linear-tree equivalence check pass for the wrong reason. + if self.tree_drafting && !is_probe && !state.announced_tree_round { + state.announced_tree_round = true; + if tree_available { + tracing::info!( + block_size = self.block_size, + "MTP is drafting trees this session (MLXCEL_MTP_TREE)" + ); + } else { + tracing::warn!( + "MLXCEL_MTP_TREE is set but the target declined a tree round; staying on the chain path" + ); + } + } + let (draft_tokens, draft_tree): (Vec, Option) = if is_probe { state.probes_remaining -= 1; - Vec::new() + (Vec::new(), None) } else { // Bound the block size by the remaining budget. When the // operator requested a block larger than the drafter's @@ -996,13 +1069,21 @@ impl MtpGenerator { // seed verify on the first round). let hidden = state.verify_out.next_hidden.as_ref(); let draft_start = Instant::now(); - match self.drafter.draft_block(state.bonus, hidden, bs, sampling) { - Ok(t) => { - state.diagnostics.draft_ms += duration_ms(draft_start.elapsed()); - t - } + // Both arms charge `draft_ms` on success and on failure, because a + // failing drafter still spent the time. + let outcome = if use_tree { + self.drafter + .draft_tree(state.bonus, hidden, bs, sampling) + .map(|tree| (Vec::new(), Some(tree))) + } else { + self.drafter + .draft_block(state.bonus, hidden, bs, sampling) + .map(|tokens| (tokens, None)) + }; + state.diagnostics.draft_ms += duration_ms(draft_start.elapsed()); + match outcome { + Ok(proposal) => proposal, Err(e) => { - state.diagnostics.draft_ms += duration_ms(draft_start.elapsed()); // Drafter failed; bail out cleanly. We've already // emitted at least the seed bonus, so finish with what // we have rather than panicking. Future hardening @@ -1018,25 +1099,76 @@ impl MtpGenerator { // If the drafter returned fewer than bs-1 proposals (e.g. // it short-circuited on a non-greedy path), `bs` shrinks // accordingly so the verify shape stays consistent. - let actual_bs = draft_tokens.len() + 1; - let mut verify_input = Vec::with_capacity(actual_bs); - verify_input.push(state.bonus); - verify_input.extend_from_slice(&draft_tokens); + let actual_bs = match draft_tree.as_ref() { + // A tree's node count already includes the root, which is the + // bonus the chain path prepends by hand, so the two agree on what + // the verify block length means. + Some(tree) => tree.len(), + None => draft_tokens.len() + 1, + }; + let verify_input: Vec = match draft_tree.as_ref() { + Some(tree) => tree.tokens().to_vec(), + None => { + let mut input = Vec::with_capacity(actual_bs); + input.push(state.bonus); + input.extend_from_slice(&draft_tokens); + input + } + }; // Phase 1: sink-aware forward. The target's KV cache is now // `bs` longer than before the call. We have target_tokens // for the walk; the captured state holds hidden + shared // K/V slabs for the finalize step. let verify_forward_start = Instant::now(); - let forward_out = self - .target - .verify_forward(&verify_input, sampling, logprobs_config); + let forward_out = match draft_tree.as_ref() { + Some(tree) => match self + .target + .verify_forward_tree(tree, sampling, logprobs_config) + { + Ok(out) => out, + Err(refusal) => { + // `tree_round_is_available` said yes above, so this is a + // target bug rather than a configuration. Nothing has been + // appended yet, so finishing is clean. + tracing::warn!( + reason = refusal.reason, + "MTP tree verify refused after the target reported a tree round available" + ); + state.finished = true; + return MtpStepOutput::finished_empty(); + } + }, + None => self + .target + .verify_forward(&verify_input, sampling, logprobs_config), + }; let verify_elapsed_ms = duration_ms(verify_forward_start.elapsed()); // Walk the draft against the target's argmax tokens. let budget = state.max_tokens - state.emitted_count; let walk_start = Instant::now(); - let walk = speculative_walk(&draft_tokens, &forward_out.target_tokens, budget); + // The tree walk carries one extra thing the chain walk has no need + // for: the accepted node path. For a linear tree that path is + // `[0, 1, ..., accepted]` and the rollback it drives reduces to the + // tail trim the chain path does, which is the equivalence to check + // first when this is switched on. + let (walk, accepted_path) = match draft_tree.as_ref() { + Some(tree) => { + let tree_walk = tree.walk(&forward_out.target_tokens, budget); + ( + WalkResult { + accepted: tree_walk.accepted, + new_tokens: tree_walk.new_tokens, + }, + Some(tree_walk.path), + ) + } + None => ( + speculative_walk(&draft_tokens, &forward_out.target_tokens, budget), + None, + ), + }; if !is_probe { state.diagnostics.speculative_walk_ms += duration_ms(walk_start.elapsed()); } @@ -1049,11 +1181,12 @@ impl MtpGenerator { state.diagnostics.record_probe(verify_elapsed_ms); } else { state.diagnostics.verify_forward_ms += verify_elapsed_ms; - state.diagnostics.record_round( - draft_tokens.len(), - walk.accepted, - walk.new_tokens.len(), - ); + // `actual_bs - 1` rather than `draft_tokens.len()`: identical on + // the chain path by construction, and the only form that is right + // for a tree, whose proposals are not a flat vector. + state + .diagnostics + .record_round(actual_bs - 1, walk.accepted, walk.new_tokens.len()); state.accept_lens.push(walk.accepted as f64); } @@ -1063,9 +1196,32 @@ impl MtpGenerator { // *before* `forward_out` is moved into `verify_finalize`. let target_logprobs = forward_out.target_logprobs; let verify_finalize_start = Instant::now(); - state.verify_out = - self.target - .verify_finalize(walk.accepted, actual_bs, forward_out.captured); + state.verify_out = match accepted_path.as_ref() { + Some(path) => { + match self + .target + .verify_finalize_tree(path, actual_bs, forward_out.captured) + { + Ok(out) => out, + Err(refusal) => { + // No fallback exists here. The verify block is on the + // cache and only a tree-aware rollback takes it off, so + // continuing would decode against a tail that no longer + // means what the drafter thinks it means. End the session + // instead. + tracing::error!( + reason = refusal.reason, + "MTP tree rollback refused after its block was appended; ending the speculative session" + ); + state.finished = true; + return MtpStepOutput::finished_empty(); + } + } + } + None => self + .target + .verify_finalize(walk.accepted, actual_bs, forward_out.captured), + }; if !is_probe { state.diagnostics.verify_finalize_ms += duration_ms(verify_finalize_start.elapsed()); } @@ -1111,7 +1267,18 @@ impl MtpGenerator { // Early-finish rounds returned above without this call — the drafter // is reset before its next session, so nothing is lost. Failures // degrade draft quality only and are logged, not fatal. - if let Some(hidden_full) = state.verify_out.verify_hidden_full.as_ref() { + if accepted_path.is_some() && state.verify_out.verify_hidden_full.is_some() { + // A stateful drafter trims its in-round tail against the chain it + // proposed, and a tree's accepted set is a path through a + // branching proposal rather than a prefix of one chain, so the + // hook has no tree form. Unreachable today rather than merely + // unimplemented: the one tree-capable target leaves + // `verify_hidden_full` as `None`. Warned rather than guessed + // (issue #1204). + tracing::warn!( + "MTP tree round on a stateful drafter: skipping accept_verified_tokens, which has no tree form yet" + ); + } else if let Some(hidden_full) = state.verify_out.verify_hidden_full.as_ref() { let accept_started = Instant::now(); let outcome = self.drafter.accept_verified_tokens( hidden_full, diff --git a/src/lib/mlxcel-core/src/speculative/mtp/mod.rs b/src/lib/mlxcel-core/src/speculative/mtp/mod.rs index 04015e5ae..51b152188 100644 --- a/src/lib/mlxcel-core/src/speculative/mtp/mod.rs +++ b/src/lib/mlxcel-core/src/speculative/mtp/mod.rs @@ -54,6 +54,7 @@ pub(crate) mod adaptive; pub mod generator; pub mod round_loop_batched; pub mod target; +pub mod tree; pub mod walk; #[cfg(test)] diff --git a/src/lib/mlxcel-core/src/speculative/mtp/target.rs b/src/lib/mlxcel-core/src/speculative/mtp/target.rs index cade7f142..7de17545a 100644 --- a/src/lib/mlxcel-core/src/speculative/mtp/target.rs +++ b/src/lib/mlxcel-core/src/speculative/mtp/target.rs @@ -203,6 +203,15 @@ impl MtpVerifyOutput { .collect() } } +/// Why a target declined to verify a draft tree. +/// +/// Carried rather than returned as a bare `None` so the round loop can log +/// the cause once and fall back to the linear path, instead of leaving an +/// operator to work out why tree drafting silently did nothing. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct TreeVerifyUnsupported { + pub reason: &'static str, +} /// Trait for MTP-compatible target models. /// @@ -298,6 +307,52 @@ pub trait MtpTarget { logprobs_config: &LogprobsConfig, ) -> VerifyForwardOutput; + /// Whether a tree round can complete on this target **right now**, asked + /// before anything is drafted. + /// + /// The two tree methods below are one capability and the round loop cannot + /// discover that by trying. [`Self::verify_forward_tree`] appends the + /// tree's block to the cache; if [`Self::verify_finalize_tree`] then + /// refuses, the round is stuck, because the block is on the cache and only + /// a tree-aware rollback can take it off. There is no falling back to the + /// linear path from there. + /// + /// The answer is per round rather than per model, which is why it is a + /// method and not a constant. Gemma 4 can verify a tree at any time but + /// can only roll one back while its sliding-window rings are unwrapped, so + /// its answer changes with cache state. + /// + /// The default is `false`, so a target opts in by answering, the same way + /// it opts into the two methods themselves. + fn tree_round_is_available(&self) -> bool { + false + } + + /// Verify a draft **tree** in one forward, or say why this target cannot. + /// + /// The default refuses, and that default is the safety property rather + /// than a placeholder: a target that has not implemented a tree-aware + /// forward must not receive one. Flattening a tree into the sequence a + /// linear verify expects would let a node's state reach its sibling, + /// which for an attention-only target the tree mask prevents and for a + /// recurrent one nothing does (issue #1204). + /// + /// Qwen 3.5 is the case that makes this matter. Its 64 layers are 48 + /// GatedDeltaNet and 16 attention, and a recurrence has no mask: a + /// flattened tree folds state through siblings and produces a wrong + /// answer silently. Its adapter overrides this to say so, rather than + /// relying on the default, so the refusal names its cause. + fn verify_forward_tree( + &self, + _tree: &super::tree::DraftTree, + _sampler: &SamplingConfig, + _logprobs_config: &LogprobsConfig, + ) -> Result { + Err(TreeVerifyUnsupported { + reason: "this target has no tree-aware verify path", + }) + } + /// Second-phase verify: apply the per-row tail-zero rollback to the /// target's KV cache based on the walk's `accepted` count, slice the /// captured shared K/V to the post-rollback length, and return the @@ -324,6 +379,57 @@ pub trait MtpTarget { captured: VerifyCaptured, ) -> MtpVerifyOutput; + /// Second-phase verify for a draft **tree**, or say why this target + /// cannot. + /// + /// [`Self::verify_finalize`] rolls the cache back by trimming a tail, + /// because a linear accept set is a prefix of the block that was + /// verified. A tree's accept set is a path, so the survivors are + /// scattered through the block and rolling back means selecting them and + /// closing the gap. That is a different operation on the cache rather + /// than a different argument to the same one, which is why it is its own + /// method instead of a `path` parameter on the linear finalize + /// (issue #1204). + /// + /// `path` is [`super::tree::TreeWalk::path`]: node indices from the root + /// through the last accepted node, strictly increasing because + /// `parents[i] < i`. `path[0]` is the root, so `path.len() - 1` is the + /// accept count the linear finalize would have been given. + /// + /// A linear tree must land exactly where the chain lands. For a + /// [`super::tree::DraftTree::linear`] the path is `[0, 1, ..., accepted]`, + /// the selection is a prefix, and the result must be byte-identical to + /// `verify_finalize(accepted, block_size, captured)`. That equivalence is + /// what lets the tree path be switched on for the configurations the + /// chain path already covers, so it is the first thing to test and the + /// first thing to suspect: a divergence there is wiring, not topology. + /// + /// The default refuses, for the reason [`Self::verify_forward_tree`] + /// refuses: a target that cannot select a scattered accept set out of its + /// cache must not be handed one. + /// + /// Accepting the forward and refusing the finalize is a real combination + /// and not a contradiction. A target can verify a tree under a mask while + /// its cache has no way to keep the path afterwards; Gemma 4's rotating + /// sliding-window caches are that case today, since `RotatingKVCache` + /// rolls back by rewinding `offset` and `idx` without moving data, which + /// a scattered selection cannot use. The pair therefore has to be probed + /// as one capability **before** a tree is drafted: once a tree forward has + /// appended its block, a refusal here leaves the cache in a state only a + /// tree-aware rollback can undo, which is also why `captured` is consumed + /// rather than handed back for a linear retry. + #[allow(unused_variables)] + fn verify_finalize_tree( + &self, + path: &[usize], + block_size: usize, + captured: VerifyCaptured, + ) -> Result { + Err(TreeVerifyUnsupported { + reason: "this target has no tree-aware cache rollback path", + }) + } + /// Number of decoder layers in the target. Used by the round-loop /// only for diagnostic logging. fn num_layers(&self) -> usize; diff --git a/src/lib/mlxcel-core/src/speculative/mtp/tests.rs b/src/lib/mlxcel-core/src/speculative/mtp/tests.rs index c18f333ac..e7ad1c99a 100644 --- a/src/lib/mlxcel-core/src/speculative/mtp/tests.rs +++ b/src/lib/mlxcel-core/src/speculative/mtp/tests.rs @@ -62,6 +62,10 @@ struct MockMtpTarget { /// before `prefill_and_seed`, then advances by `accepted + 1` per /// round. cumulative_offset: RefCell, + /// Whether this mock claims it can round-trip a draft tree. Off by + /// default so the refusal tests keep testing the default, and turned on + /// by the equivalence test that drives both arms. + tree_capable: bool, } impl MockMtpTarget { @@ -73,9 +77,24 @@ impl MockMtpTarget { call_count: RefCell::new(0), verify_log: RefCell::new(Vec::new()), cumulative_offset: RefCell::new(0), + tree_capable: false, } } + /// Claim the tree capability, and serve it by delegating to the linear + /// path. + /// + /// The delegation is the point rather than a shortcut. A linear tree's + /// node list *is* the chain's verify input, and its accepted path *is* + /// `[0, 1, ..., accepted]`, so a mock that forwards both to the linear + /// methods is a faithful tree target for the degenerate case, and any + /// difference the generator produces between the two arms is the + /// generator's own (issue #1204). + fn tree_capable(mut self) -> Self { + self.tree_capable = true; + self + } + /// Tiny dummy `[1, 1, 4]` FP32 tensor used wherever the target /// needs to surface a "hidden" or a "shared K/V slab". The /// downstream drafter's `set_shared_kv` only inspects the tensor @@ -197,6 +216,43 @@ impl MtpTarget for MockMtpTarget { } } + fn tree_round_is_available(&self) -> bool { + self.tree_capable + } + + fn verify_forward_tree( + &self, + tree: &super::tree::DraftTree, + sampler: &SamplingConfig, + logprobs_config: &LogprobsConfig, + ) -> Result { + if !self.tree_capable { + return Err(super::target::TreeVerifyUnsupported { + reason: "this mock was not built tree-capable", + }); + } + // A linear tree's node list is exactly the chain's verify input. + Ok(self.verify_forward(tree.tokens(), sampler, logprobs_config)) + } + + fn verify_finalize_tree( + &self, + path: &[usize], + block_size: usize, + captured: VerifyCaptured, + ) -> Result { + if !self.tree_capable { + return Err(super::target::TreeVerifyUnsupported { + reason: "this mock was not built tree-capable", + }); + } + // A linear tree's accepted path is `[0, 1, ..., accepted]`, so its + // length minus the root is the chain's accept count. Logging through + // the same `verify_log` is what lets the equivalence test compare the + // two arms' rollbacks and not just their tokens. + Ok(self.verify_finalize(path.len() - 1, block_size, captured)) + } + fn verify_finalize( &self, accepted: usize, @@ -1003,3 +1059,139 @@ fn an_overridden_block_request_is_not_reported_as_the_width_used() { "the recorded range must be ordered" ); } + +/// A target that has not implemented a tree-aware verify must refuse one. +/// +/// The default is the safety property, not a placeholder. Flattening a tree +/// into the sequence a linear verify expects lets a node's state reach its +/// sibling, which the tree mask prevents for attention and nothing prevents +/// for a recurrence, so the failure would be silent (issue #1204). A target +/// opts in by implementing the method; omission has to mean no. +#[test] +fn a_target_without_a_tree_path_refuses_a_tree() { + use super::tree::DraftTree; + + let target = MockMtpTarget::new(vec![vec![10, 11, 12]], vec![]); + let tree = DraftTree::linear(7, &[10, 11]); + + let refusal = target + .verify_forward_tree(&tree, &SamplingConfig::greedy(), &LogprobsConfig::default()) + .expect_err("a target with no tree path must refuse rather than flatten"); + assert!( + !refusal.reason.is_empty(), + "a refusal must carry a reason the round loop can log" + ); +} + +/// A target says nothing about trees unless it says yes. +/// +/// The round loop asks this before it drafts, because the two tree methods are +/// one capability: the forward appends a block to the cache and only a +/// tree-aware rollback can take it off, so a refusal discovered afterwards has +/// no fallback. A default of `true` would make every target that has not +/// thought about trees claim it can round-trip one (issue #1204). +#[test] +fn a_target_claims_no_tree_round_until_it_says_so() { + let target = MockMtpTarget::new(vec![vec![10, 11, 12]], vec![]); + assert!( + !target.tree_round_is_available(), + "the default has to be no, so opting in is an act rather than an omission" + ); +} + +/// The rollback half of the tree capability refuses by default too. +/// +/// Verifying a tree and keeping its accepted path are separate abilities, and +/// a target can have the first without the second: the mask that makes the +/// forward correct does nothing for a cache whose rollback is a tail trim. +/// If this half inherited a permissive default, a target that opted into the +/// forward would silently get a linear rollback applied to a scattered accept +/// set, which keeps the wrong entries rather than failing (issue #1204). +#[test] +fn a_target_without_a_tree_rollback_refuses_one() { + let target = MockMtpTarget::new(vec![vec![10, 11, 12]], vec![]); + let captured = VerifyCaptured { + tensors: Vec::new(), + scalars: Vec::new(), + }; + + // The path a full linear accept of two drafts produces: root plus both + // nodes. Even this, the shape the chain path handles, has to be refused + // by a target that has not implemented the selection. + let refusal = target + .verify_finalize_tree(&[0, 1, 2], 3, captured) + .expect_err("a target with no tree rollback must refuse rather than trim a tail"); + assert!( + !refusal.reason.is_empty(), + "a refusal must carry a reason the round loop can log" + ); +} + +/// A linear tree round must land exactly where the chain round lands, token +/// for token and rollback for rollback. +/// +/// This is the property the whole tree path rests on. `Drafter::draft_tree` +/// defaults to wrapping `draft_block` in a linear tree, so switching tree +/// drafting on without a branching drafter changes the plumbing and nothing +/// else. A divergence here is the wiring rather than the topology, which is +/// why this is the first thing to run after the round loop is wired and the +/// first thing to suspect when a branching drafter later misbehaves +/// (issue #1204). +#[test] +fn a_linear_tree_round_lands_where_the_chain_round_lands() { + // One script covering the three shapes a round takes: full accept, a + // partial accept that diverges mid-block, and an immediate rejection. + let scripted_targets = || { + vec![ + vec![10, 11, 12, 13], + vec![20, 99, 22, 23], + vec![77, 31, 32, 33], + ] + }; + let scripted_drafts = || vec![vec![10, 11, 12], vec![20, 21, 22], vec![30, 31, 32]]; + + let run = |tree: bool| { + let target = MockMtpTarget::new(scripted_targets(), vec![]); + let target = if tree { target.tree_capable() } else { target }; + let drafter = MockMtpDrafter::new(scripted_drafts()); + let mut gen_ = MtpGenerator::new(target, Box::new(drafter), 4).with_tree_drafting(tree); + let (tokens, _logprobs, stats) = gen_.generate( + &[1, 2, 3], + 13, + &SamplingConfig::greedy(), + &[], + &AtomicBool::new(false), + &LogprobsConfig::default(), + ); + let log = gen_.target().verify_log.borrow().clone(); + (tokens, stats.generated_tokens, log) + }; + + let (chain_tokens, chain_generated, chain_log) = run(false); + let (tree_tokens, tree_generated, tree_log) = run(true); + + assert_eq!( + tree_tokens, chain_tokens, + "a linear tree must emit exactly what the chain emits" + ); + assert_eq!( + tree_generated, chain_generated, + "the token counts must agree, not only the token values" + ); + assert_eq!( + tree_log, chain_log, + "the rollbacks must agree too: same accept count, same block size, same order" + ); + + // Guard against a vacuous comparison. If every round were a full accept + // the two paths would agree for a reason that has nothing to do with the + // wiring being right. + assert!( + chain_log.iter().any(|(accepted, _)| *accepted < 3), + "the script must exercise a partial accept, or this proves nothing" + ); + assert!( + chain_log.iter().any(|(accepted, _)| *accepted == 0), + "the script must exercise an immediate rejection too" + ); +} diff --git a/src/lib/mlxcel-core/src/speculative/mtp/tree.rs b/src/lib/mlxcel-core/src/speculative/mtp/tree.rs new file mode 100644 index 000000000..50d35a9c4 --- /dev/null +++ b/src/lib/mlxcel-core/src/speculative/mtp/tree.rs @@ -0,0 +1,293 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Draft-tree topology, its attention mask, and the accept walk over it. +//! +//! ## Why a tree +//! +//! A linear draft block spends its whole verify budget on one guess about the +//! future, so each extra position multiplies the chance that the prefix has +//! already diverged. Measured on M5 Max with `qwen3.8-27b-4bit` and the +//! `qwen3_5_mtp` drafter, widening the block from 3 to 12 drops acceptance +//! from 0.65 to 0.17 while emitted-per-verify moves only 2.3 to 2.8: the +//! verify cost rises with the block and the yield does not. Branching spends +//! the same extra positions on alternatives at uncertain steps instead +//! (issue #1204). +//! +//! ## What is here, and what is not +//! +//! Everything in this module is pure: a topology, a mask derived from it, and +//! a walk over target argmaxes. None of it touches a device, a cache, or a +//! model, which is what makes it testable without a checkpoint and reviewable +//! before the plumbing exists. +//! +//! The pieces that are *not* here are the ones with the real risk. Verifying a +//! tree needs the mask below handed to the target's forward instead of the +//! causal mask it builds itself, and rolling back needs the accepted path's KV +//! entries compacted, which today's caches cannot express: `KVCache` offers +//! `trim_to(new_len)` and nothing else, and tail truncation only works because +//! a linear accept set is a prefix. See the module notes on [`TreeWalk::path`] +//! for what a correct rollback has to do. + +/// A draft tree in topological order. +/// +/// Node 0 is the root: the bonus token the target already committed to this +/// round. Every other node is a proposal whose parent is the context it was +/// drafted from. `parents[i] < i` holds for every `i > 0`, which is what makes +/// the mask lower-triangular and the walk a simple descent. +/// +/// A linear draft block is the degenerate case where `parents[i] == i - 1`, +/// and it must behave exactly as the chain path does today. That equivalence +/// is the property the tests pin, because it is what lets the tree path be +/// switched on without changing output on the configurations it subsumes. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DraftTree { + tokens: Vec, + parents: Vec, +} + +impl DraftTree { + /// A tree holding only the root bonus token. + pub fn root(bonus: i32) -> Self { + Self { + tokens: vec![bonus], + parents: vec![0], + } + } + + /// The linear chain `[bonus, draft_0, ..., draft_{k-1}]`. + /// + /// The shape the round loop drafts today, expressed as a tree so both + /// paths can share one walk and one mask builder. + pub fn linear(bonus: i32, drafts: &[i32]) -> Self { + let mut tree = Self::root(bonus); + let mut parent = 0; + for &token in drafts { + parent = tree.push_child(parent, token); + } + tree + } + + /// Attach `token` under `parent` and return the new node's index. + /// + /// # Panics + /// + /// If `parent` is not an existing node. A parent index at or past the + /// current length would break the `parents[i] < i` invariant that the + /// mask and the walk both rely on, and silently accepting it would + /// produce a mask that lets a node attend to its own future. + pub fn push_child(&mut self, parent: usize, token: i32) -> usize { + assert!( + parent < self.tokens.len(), + "parent {parent} is not a node (tree has {} nodes)", + self.tokens.len() + ); + self.tokens.push(token); + self.parents.push(parent); + self.tokens.len() - 1 + } + + /// Number of nodes, including the root. + pub fn len(&self) -> usize { + self.tokens.len() + } + + pub fn is_empty(&self) -> bool { + self.tokens.is_empty() + } + + /// Tokens in node order: the verify input, flattened. + pub fn tokens(&self) -> &[i32] { + &self.tokens + } + + /// The parent of `node`. The root is its own parent. + pub fn parent(&self, node: usize) -> usize { + self.parents[node] + } + + /// Children of `node`, in insertion order. + /// + /// Linear in the tree size. Draft trees are tens of nodes at most, so + /// this stays cheaper than carrying a child list that has to be kept + /// consistent with `parents`. + pub fn children(&self, node: usize) -> Vec { + (1..self.tokens.len()) + .filter(|&i| self.parents[i] == node) + .collect() + } + + /// Root-to-node path, inclusive at both ends. + pub fn path_to(&self, node: usize) -> Vec { + let mut path = vec![node]; + let mut cursor = node; + while cursor != 0 { + cursor = self.parents[cursor]; + path.push(cursor); + } + path.reverse(); + path + } + + /// Whether `ancestor` is on the root-to-`node` path, `node` included. + pub fn is_ancestor_or_self(&self, ancestor: usize, node: usize) -> bool { + let mut cursor = node; + loop { + if cursor == ancestor { + return true; + } + if cursor == 0 { + return false; + } + cursor = self.parents[cursor]; + } + } + + /// Row-major `[n, n]` additive attention mask. + /// + /// `mask[q * n + k]` is `0.0` when node `k` is on the root-to-`q` path and + /// `masked` otherwise, so a node attends to its ancestors and itself and + /// to nothing else. Siblings are mutually invisible, which is the whole + /// point: two branches must not see each other's tokens. + /// + /// `masked` is taken rather than hard-coded because the value that means + /// "masked" is a property of the attention kernel this is handed to, and + /// picking `f32::NEG_INFINITY` here would produce NaN in a kernel that + /// adds the mask before a softmax over an all-masked row. + pub fn additive_mask(&self, masked: f32) -> Vec { + let n = self.tokens.len(); + let mut mask = vec![masked; n * n]; + for q in 0..n { + for k in self.path_to(q) { + mask[q * n + k] = 0.0; + } + } + mask + } + + /// Row-major `[n, offset + n]` additive mask, matching the convention + /// [`crate::utils::create_causal_mask`] uses: query row `q` sits at + /// logical position `q + offset`, `0.0` means attend and `masked` means + /// block. + /// + /// The first `offset` columns are the KV history already in the cache, + /// which every node attends in full: the tree branches only within the + /// block being verified, and everything before it is shared context. The + /// remaining `n` columns follow the topology, so a node sees its + /// ancestors and itself and no sibling. + /// + /// This is the shape a target forward actually takes. [`Self::additive_mask`] + /// is the `offset == 0` case and stays because it is what the topology + /// tests compare against a plain causal mask. + pub fn additive_mask_with_history(&self, offset: usize, masked: f32) -> Vec { + let n = self.tokens.len(); + let total = offset + n; + let mut mask = vec![masked; n * total]; + for q in 0..n { + let row = q * total; + for k in 0..offset { + mask[row + k] = 0.0; + } + for k in self.path_to(q) { + mask[row + offset + k] = 0.0; + } + } + mask + } + + /// Descend the branch the target agreed with. + /// + /// `target_tokens[i]` is the target's greedy choice conditional on the + /// context ending at node `i`, which is the same contract + /// [`super::walk::speculative_walk`] uses: one logit per verify position. + /// + /// Starting at the root, a child whose token equals the target's choice at + /// the current node is accepted and becomes the new current node. The walk + /// stops at the first node where no child matches, and the target's choice + /// there is the bonus that seeds the next round. A node with several + /// children can match at most one of them, since the target emits a single + /// argmax per position, so the accepted set is always a path and never a + /// subtree. + /// + /// `budget` caps the emitted tokens the way the linear walk does, so a + /// generation does not run past its declared `max_tokens`. + pub fn walk(&self, target_tokens: &[i32], budget: usize) -> TreeWalk { + debug_assert_eq!( + target_tokens.len(), + self.tokens.len(), + "the verify pass produces one target choice per tree node" + ); + debug_assert!(budget >= 1, "the round loop gates on budget before walking"); + + let mut path = vec![0usize]; + let mut new_tokens = Vec::new(); + let mut current = 0usize; + + // A short `target_tokens` is a shape bug upstream. Release builds + // stop rather than index out of bounds; the debug assertion above is + // what catches it in tests. + while let Some(&next) = target_tokens.get(current) { + let matched = self + .children(current) + .into_iter() + .find(|&child| self.tokens[child] == next); + match matched { + Some(child) if new_tokens.len() < budget => { + new_tokens.push(self.tokens[child]); + path.push(child); + current = child; + } + _ => { + if new_tokens.len() < budget { + new_tokens.push(next); + } + break; + } + } + } + + TreeWalk { + accepted: path.len() - 1, + path, + new_tokens, + } + } +} + +/// Result of [`DraftTree::walk`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TreeWalk { + /// Number of accepted draft nodes, excluding the root. + pub accepted: usize, + /// Node indices from the root through the last accepted node. + /// + /// **This is what makes tree rollback harder than chain rollback.** For a + /// linear draft the path is `[0, 1, ..., accepted]`, so the surviving KV + /// entries are a prefix and `trim_to(accepted + 1)` is correct. For a tree + /// the path is a scattered subset of the verified positions, so the cache + /// has to keep those entries and compact them, which `KVCache` cannot + /// express today. The recurrent side is better off than it looks: Qwen 3.5 + /// already rebuilds GatedDeltaNet state by replaying the captured block + /// over the accepted prefix rather than trimming it, and replaying a + /// gathered path instead of a prefix is an indexing change, not a new + /// mechanism. + pub path: Vec, + /// User-visible tokens for this round: the accepted drafts followed by + /// the target's choice at the stopping node, clamped to budget. + pub new_tokens: Vec, +} + +#[cfg(test)] +#[path = "tree_tests.rs"] +mod tree_tests; diff --git a/src/lib/mlxcel-core/src/speculative/mtp/tree_tests.rs b/src/lib/mlxcel-core/src/speculative/mtp/tree_tests.rs new file mode 100644 index 000000000..cffac13ec --- /dev/null +++ b/src/lib/mlxcel-core/src/speculative/mtp/tree_tests.rs @@ -0,0 +1,282 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Topology, mask and walk properties for [`DraftTree`]. +//! +//! The load-bearing tests are the two equivalences: a linear tree must walk +//! exactly as the chain path walks and mask exactly as a causal mask masks. +//! Everything the tree path subsumes has to keep behaving identically, or +//! turning it on changes output on configurations that were already correct. + +use super::super::walk::speculative_walk; +use super::{DraftTree, TreeWalk}; + +const MASKED: f32 = -1.0e9; + +fn causal_mask(n: usize) -> Vec { + let mut mask = vec![MASKED; n * n]; + for q in 0..n { + for k in 0..=q { + mask[q * n + k] = 0.0; + } + } + mask +} + +#[test] +fn a_linear_tree_masks_exactly_as_a_causal_mask() { + for k in 0..6 { + let drafts: Vec = (0..k).map(|i| 100 + i).collect(); + let tree = DraftTree::linear(7, &drafts); + assert_eq!( + tree.additive_mask(MASKED), + causal_mask(k as usize + 1), + "linear tree of {k} drafts must be indistinguishable from causal" + ); + } +} + +#[test] +fn a_linear_tree_walks_exactly_as_the_chain_walk() { + // Every accept length from none to all, plus the full-accept case where + // the target's last choice becomes the next round's bonus. + let drafts = vec![11, 12, 13]; + for accepted in 0..=drafts.len() { + // The target agrees for `accepted` positions and disagrees after, + // and its choice at the stopping position is the next bonus. + let mut target_tokens: Vec = (0..drafts.len()) + .map(|i| if i < accepted { drafts[i] } else { 999 }) + .collect(); + target_tokens.push(1000); + + let tree = DraftTree::linear(7, &drafts); + let tree_walk = tree.walk(&target_tokens, usize::MAX); + let chain_walk = speculative_walk(&drafts, &target_tokens, usize::MAX); + + assert_eq!( + tree_walk.accepted, chain_walk.accepted, + "accept count must match the chain walk at accepted={accepted}" + ); + assert_eq!( + tree_walk.new_tokens, chain_walk.new_tokens, + "emitted tokens must match the chain walk at accepted={accepted}" + ); + assert_eq!( + tree_walk.path, + (0..=tree_walk.accepted).collect::>(), + "a linear accept set is a prefix, which is what makes trim_to correct" + ); + } +} + +#[test] +fn siblings_cannot_see_each_other() { + // root -> {a, b}; a -> c + let mut tree = DraftTree::root(1); + let a = tree.push_child(0, 10); + let b = tree.push_child(0, 20); + let c = tree.push_child(a, 30); + let n = tree.len(); + let mask = tree.additive_mask(MASKED); + let visible = |q: usize, k: usize| mask[q * n + k] == 0.0; + + assert!(visible(a, 0) && visible(a, a), "a sees the root and itself"); + assert!(!visible(a, b), "a must not see its sibling"); + assert!(!visible(b, a), "b must not see its sibling"); + assert!( + visible(c, 0) && visible(c, a) && visible(c, c), + "c sees its whole ancestry" + ); + assert!(!visible(c, b), "c must not see the other branch"); + assert!(!visible(0, a), "the root sees nothing below it"); +} + +#[test] +fn the_walk_descends_the_branch_the_target_agreed_with() { + // root -> {a=10, b=20}; b -> d=40. The target picks b, then d. + let mut tree = DraftTree::root(1); + let _a = tree.push_child(0, 10); + let b = tree.push_child(0, 20); + let d = tree.push_child(b, 40); + + // target[root] = 20 picks b; target[b] = 40 picks d; target[d] = 50 stops. + let mut target = vec![0; tree.len()]; + target[0] = 20; + target[b] = 40; + target[d] = 50; + + let walk = tree.walk(&target, usize::MAX); + assert_eq!( + walk, + TreeWalk { + accepted: 2, + path: vec![0, b, d], + new_tokens: vec![20, 40, 50], + } + ); +} + +#[test] +fn a_branch_the_target_rejected_contributes_nothing() { + // Same tree, but the target disagrees with both children of the root. + let mut tree = DraftTree::root(1); + let _a = tree.push_child(0, 10); + let b = tree.push_child(0, 20); + let d = tree.push_child(b, 40); + + let mut target = vec![0; tree.len()]; + target[0] = 77; // matches neither 10 nor 20 + target[b] = 40; + target[d] = 50; + + let walk = tree.walk(&target, usize::MAX); + assert_eq!(walk.accepted, 0); + assert_eq!(walk.path, vec![0]); + assert_eq!( + walk.new_tokens, + vec![77], + "a fully rejected round still emits the target's own choice" + ); +} + +#[test] +fn the_accepted_set_is_a_path_even_when_two_branches_could_match() { + // Two children carrying the SAME token. The target emits one argmax per + // position, so at most one of them can be taken, and the walk must not + // fan out. This is what keeps rollback a path-compaction problem rather + // than a subtree-compaction one. + let mut tree = DraftTree::root(1); + let first = tree.push_child(0, 10); + let second = tree.push_child(0, 10); + tree.push_child(first, 30); + tree.push_child(second, 40); + + let mut target = vec![0; tree.len()]; + target[0] = 10; + target[first] = 30; + + let walk = tree.walk(&target, usize::MAX); + assert_eq!(walk.accepted, 2); + assert_eq!( + walk.path, + vec![0, first, first + 2], + "insertion order breaks the tie, and the result is still one path" + ); +} + +#[test] +fn budget_clamps_the_emission_without_changing_the_path() { + let drafts = vec![11, 12, 13]; + let tree = DraftTree::linear(7, &drafts); + let target = vec![11, 12, 13, 14]; + + let full = tree.walk(&target, usize::MAX); + assert_eq!(full.new_tokens, vec![11, 12, 13, 14]); + + let clamped = tree.walk(&target, 2); + assert_eq!(clamped.new_tokens, vec![11, 12]); + assert!( + clamped.new_tokens.len() <= 2, + "a walk must never emit past the caller's budget" + ); +} + +#[test] +fn path_and_ancestry_agree() { + let mut tree = DraftTree::root(1); + let a = tree.push_child(0, 10); + let b = tree.push_child(a, 20); + let c = tree.push_child(0, 30); + + assert_eq!(tree.path_to(b), vec![0, a, b]); + assert_eq!(tree.path_to(c), vec![0, c]); + for node in 0..tree.len() { + for candidate in 0..tree.len() { + assert_eq!( + tree.is_ancestor_or_self(candidate, node), + tree.path_to(node).contains(&candidate), + "ancestry and path must not disagree for ({candidate}, {node})" + ); + } + } +} + +#[test] +#[should_panic(expected = "is not a node")] +fn attaching_under_a_nonexistent_parent_is_rejected() { + let mut tree = DraftTree::root(1); + tree.push_child(5, 10); +} + +#[test] +fn a_history_mask_opens_the_cache_and_keeps_the_tree_shape() { + // root -> {a, b}; a -> c, with four tokens already in the KV cache. + let mut tree = DraftTree::root(1); + let a = tree.push_child(0, 10); + let b = tree.push_child(0, 20); + let c = tree.push_child(a, 30); + let offset = 4; + let n = tree.len(); + let total = offset + n; + let mask = tree.additive_mask_with_history(offset, MASKED); + let visible = |q: usize, k: usize| mask[q * total + k] == 0.0; + + assert_eq!(mask.len(), n * total); + for q in 0..n { + for k in 0..offset { + assert!(visible(q, k), "node {q} must see cached history column {k}"); + } + } + assert!(visible(c, offset + a), "c still sees its ancestor"); + assert!( + !visible(c, offset + b), + "c still cannot see the other branch" + ); + assert!( + !visible(a, offset + b), + "siblings stay invisible past the history" + ); +} + +#[test] +fn a_history_mask_at_zero_offset_is_the_plain_mask() { + let mut tree = DraftTree::root(1); + let a = tree.push_child(0, 10); + tree.push_child(0, 20); + tree.push_child(a, 30); + assert_eq!( + tree.additive_mask_with_history(0, MASKED), + tree.additive_mask(MASKED), + "the history form must degenerate to the plain one" + ); +} + +#[test] +fn a_linear_history_mask_matches_a_causal_mask_with_the_same_offset() { + // The equivalence that lets the tree path replace create_causal_mask on a + // linear block without changing what the attention sees. + let offset = 3usize; + let drafts = vec![11, 12, 13]; + let tree = DraftTree::linear(7, &drafts); + let n = tree.len(); + let total = offset + n; + + let mut expected = vec![MASKED; n * total]; + for q in 0..n { + for k in 0..=(q + offset) { + expected[q * total + k] = 0.0; + } + } + assert_eq!(tree.additive_mask_with_history(offset, MASKED), expected); +} diff --git a/src/models/gemma4.rs b/src/models/gemma4.rs index 9fd130de3..f2508ef79 100644 --- a/src/models/gemma4.rs +++ b/src/models/gemma4.rs @@ -1487,6 +1487,43 @@ impl Cache { } } + /// Whether this cache can roll back a **scattered** accept set right now. + /// + /// The tree counterpart of the implicit "trim always works" that + /// [`Self::trim_speculative`] rests on. It does not always work here: the + /// rotating variant needs its ring unwrapped, so the answer is a function + /// of state rather than of type, and a caller must ask each round before + /// it appends a verify block (issue #1204). + /// + /// Used by: Gemma 4 MTP tree rollback. + pub(crate) fn can_gather_speculative(&self) -> bool { + match self { + Self::Standard(cache) => cache.can_gather_within_tail(), + Self::Rotating(cache) => cache.can_gather_within_tail(), + } + } + + /// Keep only the verify block's `kept` slots, compacting them in place. + /// + /// The tree analogue of [`Self::trim_speculative`]. A linear accept set is + /// a prefix of the block, so rolling it back is a tail trim; a tree's is a + /// path, so the survivors are scattered and have to be selected and closed + /// up. `kept` are offsets within the block, which is exactly the node + /// index of each accepted node because the block is appended in the tree's + /// topological order. + /// + /// Used by: Gemma 4 MTP tree rollback. + pub(crate) fn gather_speculative( + &mut self, + block_size: i32, + kept: &[i32], + ) -> Result<(), String> { + match self { + Self::Standard(cache) => cache.gather_within_tail(block_size, kept), + Self::Rotating(cache) => cache.gather_within_tail(block_size, kept), + } + } + /// Add upstream-style rollback slack to sliding-window target caches used /// by Gemma 4 MTP. The logical attention window is unchanged; only the /// rotating cache's temporary storage grows so verify-block append + @@ -2503,6 +2540,54 @@ impl Attention { } } +/// The mask a sliding-window family must obey when the caller supplied one. +/// +/// A caller mask states which keys are structurally allowed. It cannot also +/// state the sliding-window band, because the same mask is handed to the +/// full-attention family, which must not have one. Below the window the two +/// constraints agree and this returns the caller mask unchanged. Past it they +/// stop agreeing, and a sliding layer has to satisfy both, so this returns +/// their intersection. +/// +/// The MTP draft-tree mask is what exposed the gap. Its history columns are +/// all-attend by construction, and [`trim_mask_to_keys`] crops from the right +/// without adding a band, so past the window the sliding layers were attending +/// outside their window and a linear-tree verify stopped reproducing the chain +/// path it is defined to match (issue #1204). +/// +/// Two properties this rests on. Additive masks intersect by `minimum`, so a +/// caller mask that already carries the band is unchanged by intersecting it +/// again, and every existing caller that builds a window-aware mask keeps its +/// exact behavior. And the intersection always keeps each query's own +/// position, which both constraints allow, so no row can go all-masked and +/// soften to NaN. +/// +/// The crop width is `min(sliding_live_len, window - 1) + l`, which is the +/// same quantity the caller of this function computes as `sliding_returned`, +/// so the result's key axis matches what the rotating cache actually returns. +/// +/// Used by: Gemma 4 caller-mask forward (speculative verify, batched and +/// adopted-prefix prefill). +fn sliding_mask_from_caller( + mask: &MlxArray, + l: i32, + sliding_live_len: i32, + window: i32, +) -> UniquePtr { + if window <= 0 || sliding_live_len + l <= window { + return mlxcel_core::copy(mask); + } + let band = create_sliding_window_prefill_mask(l, sliding_live_len, window); + let band_len = *mlxcel_core::array_shape(&band).last().unwrap_or(&0); + let mask_key_len = *mlxcel_core::array_shape(mask).last().unwrap_or(&0); + if band_len > 0 && mask_key_len >= band_len { + let cropped = slice_axis(mask, -1, mask_key_len - band_len, mask_key_len); + mlxcel_core::minimum(&cropped, &band) + } else { + band + } +} + fn trim_mask_to_keys( mask: Option<&MlxArray>, keys: &MlxArray, @@ -3255,7 +3340,10 @@ impl Gemma4TextModel { )), ) } else { - (Some(mlxcel_core::copy(mask)), Some(mlxcel_core::copy(mask))) + ( + Some(mlxcel_core::copy(mask)), + Some(sliding_mask_from_caller(mask, l, sliding_live_len, window)), + ) } } else if has_padding { // Resident prompt padding present (ragged burst). Build per-row @@ -5118,6 +5206,27 @@ impl Gemma4Wrapper { /// verification. /// /// Used by: [`crate::models::gemma4_mtp_target::Gemma4MtpTargetAdapter`]. + /// KV length the next forward on `seq_id` will treat as history. + /// + /// The same value `forward_with_speculative_sinks` derives internally + /// when it builds its own causal mask (`first_present_cache_offset`), so + /// a caller that supplies a mask instead can size it to match. Anything + /// else silently misaligns the mask against the cache. + /// + /// Exposed for draft-tree verification (issue #1204): a tree mask is + /// `[nodes, offset + nodes]` and the caller has to know `offset` to build + /// it. Reading the cache is preferable to threading the number down from + /// the round loop, which would leave two sources for one fact. + /// + /// Used by: `Gemma4MtpTargetAdapter::verify_forward_tree`. + pub(crate) fn sequence_kv_offset(&self, seq_id: Option) -> i32 { + self.sequence_state.with_or_create_sequence_state( + seq_id, + || self.model.make_caches(), + |caches| first_present_cache_offset(caches), + ) + } + pub(crate) fn forward_hidden_with_speculative_sinks( &self, input_ids: &MlxArray, @@ -5266,6 +5375,71 @@ impl Gemma4Wrapper { ) } + /// Whether every cache of this sequence can roll back a scattered accept + /// set right now. + /// + /// Asked **before** a tree is drafted, not after it is verified. The + /// rotating sliding-window caches answer from their state, so the tree + /// path is available per round rather than per model, and a round that + /// appends a tree block it cannot roll back has no recovery: the linear + /// rollback would keep the wrong entries rather than fail (issue #1204). + /// + /// Used by: [`crate::models::gemma4_mtp_target::Gemma4MtpTargetAdapter`]. + pub fn can_gather_speculative_cache(&self, seq_id: Option) -> bool { + self.sequence_state.with_or_create_sequence_state( + seq_id, + || self.model.make_caches(), + |sequence_caches| { + sequence_caches + .iter() + .all(|cache| cache.can_gather_speculative()) + }, + ) + } + + /// Keep only the verify block slots named by `kept`, on every cache of + /// this sequence. + /// + /// The tree counterpart of [`Self::rollback_speculative_cache`]. That one + /// takes an accept count and trims `block_size - (max(accepted) + 1)` from + /// the tail, which is correct exactly when the accept set is a prefix. + /// This one takes the accepted path's node indices, which are also the + /// block offsets because the verify block is appended in the tree's + /// topological order. + /// + /// A linear tree gives `kept == [0, 1, ..., accepted]`, and the two agree + /// there by construction: `KVCache::gather_within_tail` on a prefix is + /// asserted equal to `trim`. + /// + /// There is no per-row tail-zero here because there is no batched tree + /// path yet. Adding one has to revisit this: a per-row *path* is not a + /// per-row count, so the rows cannot share one global trim the way + /// `rollback_speculative_cache` lets them. + /// + /// Used by: [`crate::models::gemma4_mtp_target::Gemma4MtpTargetAdapter`]. + pub fn gather_speculative_cache( + &self, + seq_id: Option, + block_size: i32, + kept: &[i32], + ) -> Result<(), String> { + if block_size <= 0 { + return Err(format!( + "gather_speculative_cache: block_size must be positive (got {block_size})" + )); + } + self.sequence_state.with_or_create_sequence_state( + seq_id, + || self.model.make_caches(), + |sequence_caches| -> Result<(), String> { + for cache in sequence_caches.iter_mut() { + cache.gather_speculative(block_size, kept)?; + } + Ok(()) + }, + ) + } + /// Per-row tail-zero rollback against a **caller-owned** `[B, ...]` /// cache vector (batched MTP dispatch). /// @@ -5942,6 +6116,72 @@ mod gemma4_unified_mask_tests { assert!(trim_mask_to_keys(Some(&fixed), &k, m).is_some()); } + /// Below the window the caller mask is returned untouched. + /// + /// This is what keeps the fix scoped: every context that fits inside the + /// sliding window takes exactly the path it took before, so no existing + /// caller-mask path changes behavior (issue #1204). + #[test] + fn sliding_mask_from_caller_is_identity_below_the_window() { + let mask = create_causal_mask(4, 8); // [4, 12] + let out = sliding_mask_from_caller(&mask, 4, 8, 64); + mlxcel_core::eval(&out); + assert!( + mlxcel_core::item_bool(&mlxcel_core::allclose(&out, &mask, 0.0, 0.0)), + "a context inside the window must not be re-masked" + ); + } + + /// Past the window, a caller mask that lets a sliding layer see its whole + /// history gets the band added. + /// + /// This is the draft-tree shape: `DraftTree::additive_mask_with_history` + /// fills every history column with 0.0, so without this the sliding layers + /// attend outside their window and a linear-tree verify stops matching the + /// chain path it is defined to reproduce. + #[test] + fn sliding_mask_from_caller_adds_the_band_past_the_window() { + const L: i32 = 4; + const WINDOW: i32 = 8; + const LIVE: i32 = 16; + // All-attend over history + a causal block, the tree mask's shape. + let all_attend = mlxcel_core::zeros(&[L, LIVE + L], mlxcel_core::dtype::FLOAT32); + let out = sliding_mask_from_caller(&all_attend, L, LIVE, WINDOW); + mlxcel_core::eval(&out); + + let band = create_sliding_window_prefill_mask(L, LIVE, WINDOW); + mlxcel_core::eval(&band); + assert_eq!( + mlxcel_core::array_shape(&out), + mlxcel_core::array_shape(&band), + "the result must be sized to what the rotating cache returns" + ); + assert!( + mlxcel_core::item_bool(&mlxcel_core::allclose(&out, &band, 0.0, 0.0)), + "an all-attend caller mask past the window must reduce to the band" + ); + } + + /// Past the window, a caller mask that already carries the band is + /// unchanged by intersecting it again. + /// + /// The property every existing caller relies on: the batched and + /// adopted-prefix prefill paths build window-aware masks, so applying the + /// band a second time has to be a no-op rather than a tightening. + #[test] + fn sliding_mask_from_caller_leaves_an_already_banded_mask_alone() { + const L: i32 = 4; + const WINDOW: i32 = 8; + const LIVE: i32 = 16; + let banded = create_sliding_window_prefill_mask(L, LIVE, WINDOW); + let out = sliding_mask_from_caller(&banded, L, LIVE, WINDOW); + mlxcel_core::eval(&out); + assert!( + mlxcel_core::item_bool(&mlxcel_core::allclose(&out, &banded, 0.0, 0.0)), + "intersecting a banded mask with its own band must change nothing" + ); + } + /// A mask LONGER than the keys (correctly-offset caller mask on the sliding /// family) must still be cropped, not discarded: the pre-#885 crop path is /// preserved for `query_len > 1`. diff --git a/src/models/gemma4_mtp_target.rs b/src/models/gemma4_mtp_target.rs index 6696ff8cd..bbd6e40bd 100644 --- a/src/models/gemma4_mtp_target.rs +++ b/src/models/gemma4_mtp_target.rs @@ -190,6 +190,144 @@ pub struct Gemma4MtpTargetAdapter<'a> { } impl<'a> Gemma4MtpTargetAdapter<'a> { + /// Shared body of the linear and tree verify paths. + /// + /// `mask` is `None` for a linear block, where the model builds its own + /// causal mask, and `Some(tree_mask)` for a draft tree, whose topology no + /// causal mask can express (issue #1204). Everything after the forward is + /// identical, so the two paths differ by this argument and nothing else. + fn verify_forward_with_mask( + &self, + verify_input: &[i32], + mask: Option<&mlxcel_core::MlxArray>, + sampler: &SamplingConfig, + logprobs_config: &mlxcel_core::sampling::LogprobsConfig, + ) -> VerifyForwardOutput { + // Sink-aware forward over `[bonus, draft_0, …, draft_{K-2}]`. + let verify_arr = mlxcel_core::from_slice_i32(verify_input, &[1, verify_input.len() as i32]); + let mut sinks = Gemma4SpeculativeSinks::with_hidden_and_shared_kv(); + // Greedy/no-logprobs can use the latest upstream deferred path: run + // the transformer once with `skip_final_norm=True`, capture pre-norm + // hidden/shared K/V, then project hidden positions to logits only as + // needed. Keep the full-logits path for non-greedy or logprob + // requests so existing sampler/logprob semantics stay unchanged. + // The upstream Python reference uses deferred greedy hidden→logits + // projection by default. In Rust/MLX today that path projects one + // position at a time across the cxx bridge and is slower than the + // batched `[K, vocab]` LM-head projection for Gemma 4 31B on local + // Apple Silicon runs. Keep it available for parity experiments, but + // leave the faster full-logits verifier as the default until we have + // a fused/graph-side deferred walk. + let use_deferred_greedy = std::env::var("MLXCEL_ENABLE_MTP_DEFERRED").ok().as_deref() + == Some("1") + && sampler.temperature == 0.0 + && !logprobs_config.enabled; + let logits = if use_deferred_greedy { + let _ = self.wrapper.forward_hidden_with_speculative_sinks( + &verify_arr, + None, + None, + mask, + self.seq_id, + None, + Some(&mut sinks), + true, + ); + None + } else { + let raw_logits = self.wrapper.forward_with_speculative_sinks( + &verify_arr, + None, + None, + mask, + self.seq_id, + None, + Some(&mut sinks), + None, + ); + // issue #350: mask the model's output-illegal placeholder ids on the + // verify logits BEFORE the per-position argmax AND the logprob + // extraction, so a suppressed id can never win a near-tie verify + // position (the leak) and the emitted token plus its logprob stay + // consistent. The bias is the model's suppressed-id set, forwarded + // from the server `enqueue_request` through `MtpGenerator` into + // `sampler`. An empty map (every non-multimodal model) short-circuits + // here to the raw logits, so the baseline stays bit-exact. + if sampler.token_bias.is_empty() { + Some(raw_logits) + } else { + Some(mlxcel_core::sampling::apply_token_bias( + &raw_logits, + &sampler.token_bias, + )) + } + }; + + // Greedy-parity gate: pull the per-position argmax tokens from + // the verify logits. At temperature == 0 this is byte-identical + // to the drafter-less target's own argmax extension. + // + // At temperature > 0 a future enhancement plumbs the sampler + // through per-position; the round-loop driver's perf-sensitive + // path is greedy, so we keep argmax-only for now. + let target_tokens = if let Some(logits) = logits.as_ref() { + Self::argmax_per_position(logits) + } else { + // The hidden sink is still owned by `sinks`; pull it below but + // compute after extraction so the hidden handle is available for + // both token projection and `VerifyCaptured`. + Vec::new() + }; + + // Per-position log-probability data, aligned 1:1 with + // `target_tokens`. `None` (zero-overhead) when logprobs are + // disabled; the round loop forwards the entries for accepted + // positions on to `finalize_burst_success`. + let target_logprobs = logits.as_ref().and_then(|logits| { + Self::per_position_logprobs(logits, &target_tokens, logprobs_config) + }); + + // Capture the hidden + pre-slice shared K/V for the finalize step. + let hidden_sink = sinks + .hidden_sink + .expect("with_hidden_and_shared_kv installs the hidden sink"); + let hidden_full = hidden_sink + .into_iter() + .next_back() + .expect("hidden sink must carry at least one entry"); + let target_tokens = if target_tokens.is_empty() && use_deferred_greedy { + self.argmax_from_hidden_positions(hidden_full.as_ref().unwrap(), &sampler.token_bias) + } else { + target_tokens + }; + + let shared_kv_map = sinks + .shared_kv_sink + .expect("with_hidden_and_shared_kv installs the shared K/V sink"); + let mut captured_tensors: Vec> = Vec::with_capacity(5); + captured_tensors.push(hidden_full); + if let Some((k_full, v_full)) = shared_kv_map.get("full_attention") { + captured_tensors.push(mlxcel_core::copy(k_full.as_ref().unwrap())); + captured_tensors.push(mlxcel_core::copy(v_full.as_ref().unwrap())); + } + if let Some((k_swa, v_swa)) = shared_kv_map.get("sliding_attention") { + captured_tensors.push(mlxcel_core::copy(k_swa.as_ref().unwrap())); + captured_tensors.push(mlxcel_core::copy(v_swa.as_ref().unwrap())); + } + + VerifyForwardOutput { + target_tokens, + target_logprobs, + captured: VerifyCaptured { + tensors: captured_tensors, + // No per-step scalar metadata used today; the per-row + // batched path will populate these with kv_offset_pre / + // bonus_position_pre. + scalars: Vec::new(), + }, + } + } + /// Construct an adapter that routes every trait call through the /// per-sequence cache slot at `seq_id`. /// @@ -368,6 +506,66 @@ impl<'a> Gemma4MtpTargetAdapter<'a> { .collect() } + /// Keep the captured shared K/V history plus the verify block's `kept` + /// slots, dropping the rest. + /// + /// The tree counterpart of [`Self::slice_shared_kv`]. That one crops + /// `rejected` entries off the end, which is correct exactly when the + /// accept set is a prefix of the block. A tree's accepted path is + /// scattered, so the block region has to be selected and closed up while + /// the history in front of it is left alone. + /// + /// `kept` are block offsets, which equal the accepted nodes' indices + /// because the block was appended in the tree's topological order. Because + /// they are strictly increasing and inside `[0, block_size)`, a full-length + /// selection is the identity and returns the slabs untouched, matching + /// `slice_shared_kv`'s `rejected == 0` fast path. + pub(crate) fn gather_shared_kv( + tensors: Vec>, + block_size: usize, + kept: &[i32], + ) -> Vec> { + if kept.len() == block_size { + return tensors; + } + let index = mlxcel_core::from_slice_i32(kept, &[kept.len() as i32]); + tensors + .into_iter() + .map(|ptr| { + let array = ptr.as_ref().expect("shared K/V slab must be non-null"); + let shape = mlxcel_core::array_shape(array); + debug_assert!( + shape.len() == 4, + "shared K/V slab must be 4-D, got shape {:?}", + shape + ); + let kv_len = shape[2]; + let block_start = kv_len - block_size as i32; + debug_assert!( + block_start >= 0, + "gather_shared_kv: block_size ({block_size}) exceeds kv_len ({kv_len})", + ); + let block = mlxcel_core::slice( + array, + &[0, 0, block_start, 0], + &[shape[0], shape[1], kv_len, shape[3]], + ); + let picked = mlxcel_core::take(&block, &index, 2); + if block_start == 0 { + // No history in front of the block; concatenating an empty + // head would be a no-op with an extra graph node. + return picked; + } + let head = mlxcel_core::slice( + array, + &[0, 0, 0, 0], + &[shape[0], shape[1], block_start, shape[3]], + ); + mlxcel_core::concatenate(&head, &picked, 2) + }) + .collect() + } + /// Compute the argmax along the last axis for each row of a /// `[1, block_size, vocab]` logits tensor and return the /// per-position token ids. @@ -589,129 +787,137 @@ impl<'a> MtpTarget for Gemma4MtpTargetAdapter<'a> { sampler: &SamplingConfig, logprobs_config: &mlxcel_core::sampling::LogprobsConfig, ) -> VerifyForwardOutput { - // Sink-aware forward over `[bonus, draft_0, …, draft_{K-2}]`. - let verify_arr = mlxcel_core::from_slice_i32(verify_input, &[1, verify_input.len() as i32]); - let mut sinks = Gemma4SpeculativeSinks::with_hidden_and_shared_kv(); - // Greedy/no-logprobs can use the latest upstream deferred path: run - // the transformer once with `skip_final_norm=True`, capture pre-norm - // hidden/shared K/V, then project hidden positions to logits only as - // needed. Keep the full-logits path for non-greedy or logprob - // requests so existing sampler/logprob semantics stay unchanged. - // The upstream Python reference uses deferred greedy hidden→logits - // projection by default. In Rust/MLX today that path projects one - // position at a time across the cxx bridge and is slower than the - // batched `[K, vocab]` LM-head projection for Gemma 4 31B on local - // Apple Silicon runs. Keep it available for parity experiments, but - // leave the faster full-logits verifier as the default until we have - // a fused/graph-side deferred walk. - let use_deferred_greedy = std::env::var("MLXCEL_ENABLE_MTP_DEFERRED").ok().as_deref() - == Some("1") - && sampler.temperature == 0.0 - && !logprobs_config.enabled; - let logits = if use_deferred_greedy { - let _ = self.wrapper.forward_hidden_with_speculative_sinks( - &verify_arr, - None, - None, - None, - self.seq_id, - None, - Some(&mut sinks), - true, - ); - None - } else { - let raw_logits = self.wrapper.forward_with_speculative_sinks( - &verify_arr, - None, - None, - None, - self.seq_id, - None, - Some(&mut sinks), - None, - ); - // issue #350: mask the model's output-illegal placeholder ids on the - // verify logits BEFORE the per-position argmax AND the logprob - // extraction, so a suppressed id can never win a near-tie verify - // position (the leak) and the emitted token plus its logprob stay - // consistent. The bias is the model's suppressed-id set, forwarded - // from the server `enqueue_request` through `MtpGenerator` into - // `sampler`. An empty map (every non-multimodal model) short-circuits - // here to the raw logits, so the baseline stays bit-exact. - if sampler.token_bias.is_empty() { - Some(raw_logits) - } else { - Some(mlxcel_core::sampling::apply_token_bias( - &raw_logits, - &sampler.token_bias, - )) - } - }; - - // Greedy-parity gate: pull the per-position argmax tokens from - // the verify logits. At temperature == 0 this is byte-identical - // to the drafter-less target's own argmax extension. - // - // At temperature > 0 a future enhancement plumbs the sampler - // through per-position; the round-loop driver's perf-sensitive - // path is greedy, so we keep argmax-only for now. - let target_tokens = if let Some(logits) = logits.as_ref() { - Self::argmax_per_position(logits) - } else { - // The hidden sink is still owned by `sinks`; pull it below but - // compute after extraction so the hidden handle is available for - // both token projection and `VerifyCaptured`. - Vec::new() - }; + self.verify_forward_with_mask(verify_input, None, sampler, logprobs_config) + } - // Per-position log-probability data, aligned 1:1 with - // `target_tokens`. `None` (zero-overhead) when logprobs are - // disabled; the round loop forwards the entries for accepted - // positions on to `finalize_burst_success`. - let target_logprobs = logits.as_ref().and_then(|logits| { - Self::per_position_logprobs(logits, &target_tokens, logprobs_config) - }); + /// Verify a draft tree in one forward. + /// + /// Gemma 4 can do this because it has no recurrent state: the tree mask + /// makes siblings invisible to each other and every layer is attention, + /// so one flattened forward over the nodes is exact (issue #1204). A + /// hybrid target cannot, which is why the trait default refuses and why + /// Qwen 3.5 overrides it to say so. + /// + /// The mask is sized against the cache the next forward will read, taken + /// from `sequence_kv_offset` rather than threaded down from the round + /// loop, so there is one source for the offset instead of two that can + /// drift apart. + fn verify_forward_tree( + &self, + tree: &mlxcel_core::speculative::mtp::tree::DraftTree, + sampler: &SamplingConfig, + logprobs_config: &mlxcel_core::sampling::LogprobsConfig, + ) -> Result + { + let offset = self.wrapper.sequence_kv_offset(self.seq_id).max(0) as usize; + let nodes = tree.len(); + // MLX additive masks are f32 with 0 to attend; the blocking value is + // a large negative rather than -inf so an all-masked row cannot turn + // into NaN through the softmax. + let flat = tree.additive_mask_with_history(offset, -1.0e9); + let mask = mlxcel_core::from_slice_f32(&flat, &[nodes as i32, (offset + nodes) as i32]); + Ok(self.verify_forward_with_mask(tree.tokens(), Some(&mask), sampler, logprobs_config)) + } + + /// Gemma 4 can round-trip a tree whenever every cache of this sequence can + /// still address its verify block as a contiguous tail. + /// + /// The verify half is always available: the tree mask solves it and every + /// layer here is attention. The rollback half is not, because the + /// sliding-window layers hold a rotating cache and a wrapped ring has no + /// contiguous tail to select from. So this reads cache state rather than + /// returning a constant, and the round loop has to ask each round. + fn tree_round_is_available(&self) -> bool { + self.wrapper.can_gather_speculative_cache(self.seq_id) + } - // Capture the hidden + pre-slice shared K/V for the finalize step. - let hidden_sink = sinks - .hidden_sink - .expect("with_hidden_and_shared_kv installs the hidden sink"); - let hidden_full = hidden_sink - .into_iter() - .next_back() - .expect("hidden sink must carry at least one entry"); - let target_tokens = if target_tokens.is_empty() && use_deferred_greedy { - self.argmax_from_hidden_positions(hidden_full.as_ref().unwrap(), &sampler.token_bias) - } else { - target_tokens - }; + /// Roll the target cache back to a draft tree's accepted **path**. + /// + /// [`Self::verify_finalize`] trims `block_size - accepted - 1` entries off + /// the tail, which is right exactly when the accept set is a prefix. The + /// tree's is a path, so the survivors are scattered through the block and + /// the rollback selects them and closes the gap instead (issue #1204). + /// + /// `path` is the walk's node indices from the root through the last + /// accepted node. They double as block offsets, because the verify block + /// was appended in the tree's topological order, and `path[0]` is the root + /// the target already committed to, so it is always kept. + /// + /// For a linear tree the path is `[0, 1, ..., accepted]`, the selection is + /// a prefix, and every step below reduces to the linear one: the gather + /// equals a `trim`, the shared-K/V gather equals the tail slice, and the + /// hidden is read at the same position. That is the equivalence to test + /// first when the round loop is wired. + /// + /// The capability check is a guard, not the contract. It is here so a + /// wrapped sliding-window ring produces a named refusal rather than a + /// silently wrong window, but by this point the block is already on the + /// cache and a refusal is not recoverable. The round loop is what has to + /// ask `can_gather_speculative_cache` before it drafts a tree. + fn verify_finalize_tree( + &self, + path: &[usize], + block_size: usize, + captured: VerifyCaptured, + ) -> Result { + use mlxcel_core::speculative::mtp::target::TreeVerifyUnsupported; - let shared_kv_map = sinks - .shared_kv_sink - .expect("with_hidden_and_shared_kv installs the shared K/V sink"); - let mut captured_tensors: Vec> = Vec::with_capacity(5); - captured_tensors.push(hidden_full); - if let Some((k_full, v_full)) = shared_kv_map.get("full_attention") { - captured_tensors.push(mlxcel_core::copy(k_full.as_ref().unwrap())); - captured_tensors.push(mlxcel_core::copy(v_full.as_ref().unwrap())); + if path.is_empty() { + return Err(TreeVerifyUnsupported { + reason: "an accepted path always contains the root, so an empty one is a bug", + }); } - if let Some((k_swa, v_swa)) = shared_kv_map.get("sliding_attention") { - captured_tensors.push(mlxcel_core::copy(k_swa.as_ref().unwrap())); - captured_tensors.push(mlxcel_core::copy(v_swa.as_ref().unwrap())); + if !self.wrapper.can_gather_speculative_cache(self.seq_id) { + return Err(TreeVerifyUnsupported { + reason: "a Gemma 4 sliding-window ring has wrapped, so its verify block is no \ + longer a contiguous tail that a path selection can address", + }); } - VerifyForwardOutput { - target_tokens, - target_logprobs, - captured: VerifyCaptured { - tensors: captured_tensors, - // No per-step scalar metadata used today; the per-row - // batched path will populate these with kv_offset_pre / - // bonus_position_pre. - scalars: Vec::new(), - }, + let mut tensors = captured.tensors; + if tensors.is_empty() { + return Err(TreeVerifyUnsupported { + reason: "VerifyCaptured must carry at least the hidden tensor at index 0", + }); } + let hidden_full = tensors.remove(0); + + // The drafter's next hidden comes from the last accepted node. For a + // tree that is a node index rather than an accept count; the two + // coincide on a linear tree, which is the equivalence. + let last_accepted = *path.last().expect("checked non-empty above"); + let next_hidden = + self.draft_hidden_at_position(hidden_full.as_ref().unwrap(), last_accepted); + + let kept: Vec = path.iter().map(|&node| node as i32).collect(); + if self + .wrapper + .gather_speculative_cache(self.seq_id, block_size as i32, &kept) + .is_err() + { + return Err(TreeVerifyUnsupported { + reason: "the Gemma 4 target cache refused the accepted path selection", + }); + } + + let next_shared_kv = Self::gather_shared_kv(tensors, block_size, &kept); + + // Same rebind as the linear finalize: the absolute post-rollback cache + // offset, read from the wrapper rather than derived from the accept + // count, so the drafter's RoPE anchor cannot drift. + let kv_offset = self + .wrapper + .speculative_cache_offset(self.seq_id, "full_attention") + .max(0) as usize; + let bonus_position = mtp_draft_position(kv_offset); + + Ok(MtpVerifyOutput { + next_hidden, + next_shared_kv, + kv_offset, + bonus_position, + verify_hidden_full: None, + }) } fn verify_finalize( @@ -822,6 +1028,36 @@ impl<'a> Gemma4VLMtpTargetAdapter<'a> { } impl<'a> MtpTarget for Gemma4VLMtpTargetAdapter<'a> { + /// Delegated to the inner text adapter, which is where the caches a tree + /// rollback selects from actually live. Inheriting the refusing defaults + /// here made the whole tree path dead for every multimodal Gemma 4 + /// checkpoint while the text-only one worked, which is not a distinction + /// anything about trees justifies: the vision tower is not in the verify + /// forward at all (issue #1204). + fn tree_round_is_available(&self) -> bool { + self.inner.tree_round_is_available() + } + + fn verify_forward_tree( + &self, + tree: &mlxcel_core::speculative::mtp::tree::DraftTree, + sampler: &SamplingConfig, + logprobs_config: &mlxcel_core::sampling::LogprobsConfig, + ) -> Result + { + self.inner + .verify_forward_tree(tree, sampler, logprobs_config) + } + + fn verify_finalize_tree( + &self, + path: &[usize], + block_size: usize, + captured: VerifyCaptured, + ) -> Result { + self.inner.verify_finalize_tree(path, block_size, captured) + } + fn prefill_and_seed( &self, prompt_tokens: &[i32], @@ -916,6 +1152,36 @@ impl<'a> Gemma4UnifiedMtpTargetAdapter<'a> { } impl<'a> MtpTarget for Gemma4UnifiedMtpTargetAdapter<'a> { + /// Delegated to the inner text adapter, which is where the caches a tree + /// rollback selects from actually live. Inheriting the refusing defaults + /// here made the whole tree path dead for every multimodal Gemma 4 + /// checkpoint while the text-only one worked, which is not a distinction + /// anything about trees justifies: the vision tower is not in the verify + /// forward at all (issue #1204). + fn tree_round_is_available(&self) -> bool { + self.inner.tree_round_is_available() + } + + fn verify_forward_tree( + &self, + tree: &mlxcel_core::speculative::mtp::tree::DraftTree, + sampler: &SamplingConfig, + logprobs_config: &mlxcel_core::sampling::LogprobsConfig, + ) -> Result + { + self.inner + .verify_forward_tree(tree, sampler, logprobs_config) + } + + fn verify_finalize_tree( + &self, + path: &[usize], + block_size: usize, + captured: VerifyCaptured, + ) -> Result { + self.inner.verify_finalize_tree(path, block_size, captured) + } + fn prefill_and_seed( &self, prompt_tokens: &[i32], diff --git a/src/models/qwen3_5_mtp_target.rs b/src/models/qwen3_5_mtp_target.rs index 25b0709a0..9161baa37 100644 --- a/src/models/qwen3_5_mtp_target.rs +++ b/src/models/qwen3_5_mtp_target.rs @@ -268,6 +268,34 @@ impl<'a> MtpTarget for Qwen35MtpTargetAdapter<'a> { .expect("Qwen35Model exposes its embed_tokens table") } + /// Refuse a draft tree, naming the reason rather than inheriting the + /// default. + /// + /// Qwen 3.5 is a hybrid: with `full_attention_interval` 4 over 64 layers, + /// 48 are GatedDeltaNet and 16 are attention. The tree mask solves the + /// attention layers and does nothing for the other 48, because a + /// recurrence is a fold with no mask: flattening a tree through it carries + /// state from a node into its sibling and produces a wrong answer with no + /// signal. Verifying a tree here needs a tree-aware scan that forks state + /// at branch points, which does not exist (issue #1204). + /// + /// Overridden rather than left to the trait default so the refusal states + /// its cause at the place a reader looks for it. + fn verify_forward_tree( + &self, + _tree: &mlxcel_core::speculative::mtp::tree::DraftTree, + _sampler: &SamplingConfig, + _logprobs_config: &LogprobsConfig, + ) -> Result + { + Err( + mlxcel_core::speculative::mtp::target::TreeVerifyUnsupported { + reason: "Qwen 3.5 is a GatedDeltaNet hybrid; a recurrence has no attention mask, \ + so a flattened tree would fold state through siblings", + }, + ) + } + fn verify_forward( &self, verify_input: &[i32], @@ -312,6 +340,34 @@ impl<'a> MtpTarget for Qwen35MtpTargetAdapter<'a> { } } + /// Refuse a tree rollback, naming the reason rather than inheriting the + /// default. + /// + /// Unreachable in practice, because [`Self::verify_forward_tree`] refuses + /// first and no tree block ever reaches this cache. It is here because the + /// two halves are one capability: a reader checking whether Qwen 3.5 can + /// keep a tree path should find the answer at the rollback too, and a + /// future forward that stops refusing must not silently inherit a rollback + /// that cannot serve it. + /// + /// The cause is the same one. 48 of 64 layers are GatedDeltaNet, and their + /// rollback replays the captured block over the accepted prefix. Replaying + /// a gathered path instead of a prefix is an indexing change rather than a + /// new mechanism, so this is reachable work; it is simply not done + /// (issue #1204). + fn verify_finalize_tree( + &self, + _path: &[usize], + _block_size: usize, + _captured: VerifyCaptured, + ) -> Result { + Err( + mlxcel_core::speculative::mtp::target::TreeVerifyUnsupported { + reason: "Qwen 3.5 rolls GatedDeltaNet state back by replaying the accepted prefix, and replaying a scattered path is not implemented", + }, + ) + } + fn verify_finalize( &self, accepted: usize,