diff --git a/crates/onnx-genai-engine/src/pipeline/flat_autoregressive.rs b/crates/onnx-genai-engine/src/pipeline/flat_autoregressive.rs index e7c1882179..a50d590e2b 100644 --- a/crates/onnx-genai-engine/src/pipeline/flat_autoregressive.rs +++ b/crates/onnx-genai-engine/src/pipeline/flat_autoregressive.rs @@ -4,7 +4,70 @@ //! decode driver and its paged-sequence and prefix-reuse helpers. use super::paged_decode::{PagedMirror, PipelineDecodeLoopBackend}; +use super::prefix_reuse::{KvPrefixStore, ReusedPrefix, apply_prefix_reuse}; use super::*; + +/// A native decoder's session-resident KV, as a [`KvPrefixStore`]. +struct NativeKvStore<'a>(&'a mut dyn PipelineDecoderComponent); + +impl KvPrefixStore for NativeKvStore<'_> { + fn current_kv_len(&self) -> usize { + // `None` means the backend keeps nothing across turns, which is the + // same starting point as an empty cache. + self.0.current_kv_len().unwrap_or(0) + } + + fn use_kv(&self) -> bool { + self.0.use_kv() + } + + fn rewind_to(&mut self, target: usize) -> anyhow::Result { + self.0.rewind_kv(target) + } + + fn reset(&mut self) -> anyhow::Result<()> { + // A decoder that cannot rewind to empty must not keep a cache the next + // turn would read as its prefix, so a decline is an error here rather + // than a silent carry-over. + if self.0.rewind_kv(0)? { + return Ok(()); + } + anyhow::bail!("native decoder declined to reset its KV cache") + } +} + +/// ORT's past tensors, held by the pipeline rather than the decoder, as a +/// [`KvPrefixStore`]. `current` is tracked outside the state, so it is passed +/// in rather than read back. +struct OrtKvStore<'a> { + state: &'a mut Option, + current: usize, +} + +impl KvPrefixStore for OrtKvStore<'_> { + fn current_kv_len(&self) -> usize { + self.current + } + + fn use_kv(&self) -> bool { + self.state.as_ref().is_some_and(|state| state.use_kv) + } + + fn rewind_to(&mut self, target: usize) -> anyhow::Result { + let Some(state) = self.state.as_mut() else { + return Ok(false); + }; + state.truncate_past(self.current, target) + } + + fn reset(&mut self) -> anyhow::Result<()> { + // Dropping the state is how ORT empties its cache: a fresh one is built + // below for every turn that reuses nothing. + *self.state = None; + Ok(()) + } +} + impl PipelineEngine { /// Resolve which components a flat autoregressive decode drives natively, /// unifying the pure-native backend and the hybrid env-flag injection so both @@ -122,30 +185,36 @@ impl PipelineEngine { let paged_enabled = self.paged.is_some() && inputs_digest.is_some() && (!use_native_decoder || native_supports_paging); - let reused = if use_native_decoder && !paged_enabled { - let shared = match (inputs_digest, self.retained.as_ref()) { - (Some(inputs), Some(retained)) => retained.reusable_prefix(inputs, &prompt_tokens), - _ => 0, - }; - let current = native_decoder_component - .as_ref() - .and_then(|decoder| decoder.current_kv_len()) - .unwrap_or(0); - let reusable = shared.min(current); - let must_reset = reusable == 0 || (reusable < current && !self.positions_are_linear()); - let target = if must_reset { 0 } else { reusable }; - if let Some(decoder) = native_decoder_component.as_mut() - && target != current - && !decoder.rewind_kv(target)? - { - 0 - } else { - target + // One prefix-reuse policy for every backend (`prefix_reuse`): the + // shared-prefix question and the decision it feeds are identical, so + // only the cache mechanism is chosen here. A native decoder rewinds its + // session-resident KV; ORT slices pipeline-held past tensors. + let shared = match (inputs_digest, self.retained.as_ref()) { + (Some(inputs), Some(retained)) => retained.reusable_prefix(inputs, &prompt_tokens), + _ => 0, + }; + let retained_len = self.retained.as_ref().map_or(0, |r| r.tokens.len()); + let positions_are_linear = self.positions_are_linear(); + let reused = if paged_enabled { + ReusedPrefix::NONE + } else if use_native_decoder { + match native_decoder_component.as_mut() { + Some(decoder) => apply_prefix_reuse( + &mut NativeKvStore(decoder.as_mut()), + shared, + positions_are_linear, + )?, + None => ReusedPrefix::NONE, } - } else if paged_enabled || use_native_decoder { - 0 } else { - self.reusable_prefix_len(inputs_digest, &prompt_tokens) + apply_prefix_reuse( + &mut OrtKvStore { + state: &mut self.decoder_state, + current: retained_len, + }, + shared, + positions_are_linear, + )? }; // Any failure below leaves the decoder KV in an unknown state, so the // retention is dropped now and only re-established on success. @@ -183,8 +252,8 @@ impl PipelineEngine { // tensor covering the whole prompt. Prefilling only a suffix would hand // it positions for tokens it is not being given, so such a pipeline // recomputes rather than reuses. - let reused = if reused > 0 && self.decoder_positions_are_routed(&decoder_in_edges) { - 0 + let reused = if !reused.is_empty() && self.decoder_positions_are_routed(&decoder_in_edges) { + ReusedPrefix::NONE } else { reused }; @@ -226,11 +295,11 @@ impl PipelineEngine { Self::admit_paged_sequence(paged, state, decoder, inputs, &prompt_tokens)? }; paged_session = Some((seq, inputs)); - reused = shared; + reused = ReusedPrefix::from_paged_admission(shared); } let chain = build_processor_chain(&options, Some(self.tokenizer()?))?; - if reused == 0 && paged_session.is_none() { + if reused.is_empty() && paged_session.is_none() { self.decoder_state = Some(Self::new_decoder_state( &self.models, &ar.decoder, @@ -329,16 +398,16 @@ impl PipelineEngine { decoder_in_edges, static_cross_kv, context_tokens: prompt_tokens, - retained_len: reused, + retained_len: reused.len(), prompt_len: 0, generated_count: 0, - kv_len: reused, + kv_len: reused.len(), prefill_chunk_size: self.prefill_chunk_size, }; // Prefill only what the retained KV does not already cover. backend.prompt_len = backend.context_tokens.len() - backend.retained_len; let prefilled = backend.prompt_len; - let mut loop_state = DecodeLoopState::new(reused, options.seed, options.top_logprobs); + let mut loop_state = DecodeLoopState::new(reused.len(), options.seed, options.top_logprobs); // Taken without `?` so a failed generation still releases its sequence // below: an abandoned sequence holds its pages out of the pool for the // life of the process. @@ -388,7 +457,7 @@ impl PipelineEngine { self.component_cache .borrow_mut() - .note_prefix_reuse(reused, prefilled); + .note_prefix_reuse(reused.len(), prefilled); match paged_session { // Publish what this generation computed so the next request can // attach to it, then let go of the sequence. @@ -623,52 +692,4 @@ impl PipelineEngine { .iter() .any(|(_, port)| port == position_input) } - - /// How many leading prompt tokens the retained decoder KV can serve. - /// - /// Zero whenever anything about the request's identity changed, whenever - /// the prompt shares no leading token with the retained context, or - /// whenever the attachments could not be digested. - /// - /// When the prompt diverges part-way, the retained KV is first truncated to - /// the shared head. Truncation can decline — an opaque past with no - /// identifiable sequence axis, or fixed loop-carried state — in which case - /// nothing is reused and the turn is recomputed. - fn reusable_prefix_len(&mut self, inputs: Option, prompt_tokens: &[TokenId]) -> usize { - let (Some(inputs), Some(retained)) = (inputs, self.retained.as_ref()) else { - return 0; - }; - let retained_len = retained.tokens.len(); - let shared = retained.reusable_prefix(inputs, prompt_tokens); - let Some(state) = self.decoder_state.as_mut() else { - return 0; - }; - if !state.use_kv || shared == 0 { - return 0; - } - if shared == retained_len { - return shared; - } - // Extending keeps the carried position state valid; truncating does not, - // and only a linear continuation can be rebuilt from the absolute past - // length alone. A model that carries or is handed its coordinates would - // resume from positions describing tokens that no longer exist. - if !self.positions_are_linear() { - self.decoder_state = None; - return 0; - } - let Some(state) = self.decoder_state.as_mut() else { - return 0; - }; - match state.truncate_past(retained_len, shared) { - Ok(true) => shared, - // Declining to truncate is a normal outcome, and so is a failed - // slice: either way the KV is no longer trustworthy for reuse, so - // the state is dropped and the turn recomputes from scratch. - _ => { - self.decoder_state = None; - 0 - } - } - } } diff --git a/crates/onnx-genai-engine/src/pipeline/mod.rs b/crates/onnx-genai-engine/src/pipeline/mod.rs index ef25bebab9..f7e57dc64b 100644 --- a/crates/onnx-genai-engine/src/pipeline/mod.rs +++ b/crates/onnx-genai-engine/src/pipeline/mod.rs @@ -45,6 +45,7 @@ mod flat_autoregressive; mod iterative; mod nested_autoregressive; mod paged_decode; +mod prefix_reuse; mod routing; mod schedulers; #[cfg(feature = "native-backend")] diff --git a/crates/onnx-genai-engine/src/pipeline/prefix_reuse.rs b/crates/onnx-genai-engine/src/pipeline/prefix_reuse.rs new file mode 100644 index 0000000000..cbc0d4aed7 --- /dev/null +++ b/crates/onnx-genai-engine/src/pipeline/prefix_reuse.rs @@ -0,0 +1,363 @@ +//! The single prefix-reuse policy, shared by every decoder backend. +//! +//! Both backends answer the same question at the top of a turn: the retained +//! context says some prompt prefix is reusable, the KV holds some number of +//! tokens — how much of it survives, and what has to happen to the cache +//! first? Only the *mechanism* differs. ORT owns its past as pipeline-held +//! tensors it slices; a native decoder owns a session-resident cache it +//! rewinds. That difference belongs in [`KvPrefixStore`], and nowhere else. +//! +//! Keeping the policy in one place is not tidiness. When it lived in two, the +//! copies drifted: the native copy had no `use_kv` gate and the ORT copy had no +//! clamp against the KV's real length, so each was missing a guard the other +//! had. Neither omission could show up as a crash — a prefix-reuse bug produces +//! fluent text conditioned on the wrong context. + +/// A prefix length that came from the shared policy. +/// +/// This is the DRY mechanism, not decoration. The field is private and the +/// only constructors live in this module, so a new backend **cannot produce a +/// reuse length at all** without going through [`apply_prefix_reuse`] — the +/// exact mistake that produced two divergent copies of this policy in the +/// first place is now a compile error rather than a review catch. +/// +/// If you find yourself wanting a fourth constructor, that is the signal to +/// stop: you are about to add a second policy. Extend [`plan_prefix_reuse`] +/// instead, so every backend gets the change. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct ReusedPrefix(usize); + +impl ReusedPrefix { + /// Nothing carried over; the turn prefills the whole prompt. + pub(crate) const NONE: Self = Self(0); + + /// The paged cache admits its own shared prefix through `claim_paged_prefix`, + /// which is already a single implementation both backends share. It is a + /// different *source* of a prefix, not a different policy. + pub(crate) fn from_paged_admission(len: usize) -> Self { + Self(len) + } + + /// Tokens the turn may skip re-prefilling. + pub(crate) fn len(self) -> usize { + self.0 + } + + pub(crate) fn is_empty(self) -> bool { + self.0 == 0 + } +} + +/// What the KV cache must become before this turn can reuse a prefix. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum PrefixDecision { + /// Nothing is reusable; the cache starts empty. + Reset, + /// The cache already holds exactly the reusable prefix. No work. + Keep(usize), + /// The cache holds more than the prompt shares; drop the tail first. + Rewind(usize), +} + +impl PrefixDecision { + /// Tokens the turn may skip re-prefilling, before any rewind is attempted. + #[cfg(test)] + pub(crate) fn reused_len(self) -> usize { + match self { + Self::Reset => 0, + Self::Keep(len) | Self::Rewind(len) => len, + } + } +} + +/// Decide how much of `current` KV a prompt sharing `shared` leading tokens may +/// reuse. Pure, so the policy is testable without a model or a device. +/// +/// `shared` comes from the retained context and describes *tokens*; `current` +/// is what the cache actually holds. They are maintained separately, so the +/// smaller governs — trusting `shared` alone would reuse tokens the cache never +/// had. +pub(crate) fn plan_prefix_reuse( + shared: usize, + current: usize, + use_kv: bool, + positions_are_linear: bool, +) -> PrefixDecision { + if !use_kv || shared == 0 { + return PrefixDecision::Reset; + } + let reusable = shared.min(current); + if reusable == 0 { + return PrefixDecision::Reset; + } + if reusable == current { + // Pure extension: the cache is a prefix of this prompt, so the carried + // position state stays valid and nothing has to be rewound. + return PrefixDecision::Keep(reusable); + } + // Truncating invalidates carried positions, and only a linear continuation + // can be rebuilt from the absolute past length alone. A model that carries + // or is handed its coordinates would resume from positions describing + // tokens that no longer exist. + if !positions_are_linear { + return PrefixDecision::Reset; + } + PrefixDecision::Rewind(reusable) +} + +/// A rewindable KV cache. The only thing the two backends do differently. +pub(crate) trait KvPrefixStore { + /// Tokens currently held. `0` when the backend keeps no reusable cache. + fn current_kv_len(&self) -> usize; + + /// Whether this decoder carries a KV cache at all. + fn use_kv(&self) -> bool; + + /// Drop everything after `target`. `Ok(false)` declines — an opaque past + /// with no identifiable sequence axis, or fixed loop-carried state. + fn rewind_to(&mut self, target: usize) -> anyhow::Result; + + /// Return to an empty cache. Must not leave stale entries behind. + fn reset(&mut self) -> anyhow::Result<()>; +} + +/// Apply [`plan_prefix_reuse`] to `store`, returning the reusable prefix length. +/// +/// A declined rewind resets. This is the whole reason application is shared +/// rather than left to each call site: "I could not rewind" and "reuse nothing" +/// are only consistent if the cache is actually emptied. Reporting `0` while +/// leaving the old tokens in place would prefill the new prompt *on top of* a +/// previous turn's KV. +pub(crate) fn apply_prefix_reuse( + store: &mut S, + shared: usize, + positions_are_linear: bool, +) -> anyhow::Result { + let decision = plan_prefix_reuse( + shared, + store.current_kv_len(), + store.use_kv(), + positions_are_linear, + ); + match decision { + PrefixDecision::Reset => { + store.reset()?; + Ok(ReusedPrefix::NONE) + } + PrefixDecision::Keep(len) => Ok(ReusedPrefix(len)), + PrefixDecision::Rewind(len) => { + if store.rewind_to(len)? { + Ok(ReusedPrefix(len)) + } else { + store.reset()?; + Ok(ReusedPrefix::NONE) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_decoder_without_kv_reuses_nothing() { + assert_eq!(plan_prefix_reuse(8, 8, false, true), PrefixDecision::Reset); + assert_eq!(plan_prefix_reuse(8, 8, false, true).reused_len(), 0); + } + + #[test] + fn no_shared_prompt_prefix_reuses_nothing() { + assert_eq!(plan_prefix_reuse(0, 8, true, true), PrefixDecision::Reset); + } + + #[test] + fn an_empty_cache_reuses_nothing_however_much_the_prompt_shares() { + assert_eq!(plan_prefix_reuse(8, 0, true, true), PrefixDecision::Reset); + } + + #[test] + fn a_cache_that_is_exactly_the_shared_prefix_is_kept_untouched() { + assert_eq!(plan_prefix_reuse(8, 8, true, true), PrefixDecision::Keep(8)); + } + + #[test] + fn the_cache_length_bounds_the_reuse_when_the_context_claims_more() { + // The retained context and the cache are maintained separately, so a + // context claiming more than the cache holds must not be believed. + assert_eq!( + plan_prefix_reuse(12, 8, true, true), + PrefixDecision::Keep(8) + ); + } + + #[test] + fn a_diverging_prompt_rewinds_the_cache_to_the_shared_head() { + assert_eq!( + plan_prefix_reuse(5, 8, true, true), + PrefixDecision::Rewind(5) + ); + } + + #[test] + fn a_diverging_prompt_cannot_rewind_nonlinear_positions() { + assert_eq!(plan_prefix_reuse(5, 8, true, false), PrefixDecision::Reset); + } + + #[test] + fn nonlinear_positions_still_allow_a_pure_extension() { + // Nothing is truncated, so no carried position is invalidated. + assert_eq!( + plan_prefix_reuse(8, 8, true, false), + PrefixDecision::Keep(8) + ); + } + + /// Records what a backend was actually asked to do. + #[derive(Default)] + struct FakeStore { + len: usize, + use_kv: bool, + rewind_declines: bool, + rewound_to: Option, + was_reset: bool, + } + + impl KvPrefixStore for FakeStore { + fn current_kv_len(&self) -> usize { + self.len + } + fn use_kv(&self) -> bool { + self.use_kv + } + fn rewind_to(&mut self, target: usize) -> anyhow::Result { + self.rewound_to = Some(target); + Ok(!self.rewind_declines) + } + fn reset(&mut self) -> anyhow::Result<()> { + self.was_reset = true; + Ok(()) + } + } + + #[test] + fn a_kept_prefix_touches_the_cache_at_all() -> anyhow::Result<()> { + let mut store = FakeStore { + len: 8, + use_kv: true, + ..FakeStore::default() + }; + assert_eq!(apply_prefix_reuse(&mut store, 8, true)?.len(), 8); + assert_eq!(store.rewound_to, None, "a kept prefix needs no rewind"); + assert!(!store.was_reset); + Ok(()) + } + + #[test] + fn a_declined_rewind_empties_the_cache_rather_than_leaving_its_tail() -> anyhow::Result<()> { + // Reporting zero reuse while the previous turn's tokens stay in the + // cache would prefill the new prompt on top of them. + let mut store = FakeStore { + len: 8, + use_kv: true, + rewind_declines: true, + ..FakeStore::default() + }; + assert!(apply_prefix_reuse(&mut store, 5, true)?.is_empty()); + assert_eq!(store.rewound_to, Some(5)); + assert!(store.was_reset, "a declined rewind must reset"); + Ok(()) + } + + #[test] + fn reusing_nothing_empties_the_cache() -> anyhow::Result<()> { + let mut store = FakeStore { + len: 8, + use_kv: true, + ..FakeStore::default() + }; + assert!(apply_prefix_reuse(&mut store, 0, true)?.is_empty()); + assert!(store.was_reset); + Ok(()) + } + + /// Files allowed to drive the KV-rewind primitives directly. + /// + /// `flat_autoregressive.rs` holds the two `KvPrefixStore` adapters, which + /// are the sanctioned mechanism half of this module. `state.rs` defines + /// `truncate_past` and exercises it in its own unit tests. + const KV_REWIND_CALLERS: [&str; 3] = [ + "pipeline\\flat_autoregressive.rs", + "pipeline\\prefix_reuse.rs", + "decode\\state.rs", + ]; + + /// The primitives that mutate or inspect a KV cache's length. Reaching for + /// one of these outside the adapters means reimplementing prefix reuse. + const KV_REWIND_PRIMITIVES: [&str; 3] = [".rewind_kv(", ".truncate_past(", ".current_kv_len("]; + + fn rust_sources(dir: &std::path::Path, out: &mut Vec) { + for entry in std::fs::read_dir(dir).expect("crate src is readable") { + let path = entry.expect("readable entry").path(); + if path.is_dir() { + rust_sources(&path, out); + } else if path.extension().is_some_and(|ext| ext == "rs") { + out.push(path); + } + } + } + + /// The DRY tripwire. + /// + /// A second copy of this policy cannot appear without calling one of the + /// KV primitives, so confining those calls to the adapters confines the + /// policy to this module. This runs under `cargo test`, so it cannot be + /// missed by not looking at CI. + /// + /// If this fails, the fix is to route the new backend through + /// [`apply_prefix_reuse`] and add a [`KvPrefixStore`] adapter -- **not** to + /// add the offending file to `KV_REWIND_CALLERS`. Widening the allowlist is + /// how the duplication this module exists to prevent gets back in. + #[test] + fn the_kv_rewind_primitives_are_only_driven_by_the_shared_policy() { + let src = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("src"); + let mut sources = Vec::new(); + rust_sources(&src, &mut sources); + assert!( + sources.len() > 10, + "the source scan found almost nothing, so it cannot be trusted to \ + have looked: {} files under {}", + sources.len(), + src.display() + ); + + let mut offenders = Vec::new(); + for path in sources { + let relative = path.strip_prefix(&src).expect("scanned under src"); + let display = relative.display().to_string().replace('/', "\\"); + if KV_REWIND_CALLERS + .iter() + .any(|allowed| display.ends_with(allowed)) + { + continue; + } + let text = std::fs::read_to_string(&path).expect("source is readable"); + for (line_number, line) in text.lines().enumerate() { + for primitive in KV_REWIND_PRIMITIVES { + if line.contains(primitive) { + offenders.push(format!("{display}:{} {}", line_number + 1, line.trim())); + } + } + } + } + + assert!( + offenders.is_empty(), + "these call a KV-rewind primitive outside the `KvPrefixStore` \ + adapters, which is how prefix reuse gets reimplemented per \ + backend. Route them through `apply_prefix_reuse` instead of \ + widening `KV_REWIND_CALLERS`:\n{}", + offenders.join("\n") + ); + } +}