From c19289cc45454e93ff31baeeb4f499b882bc8835 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Tue, 6 Oct 2026 18:07:24 +0900 Subject: [PATCH 1/2] fix(prompt-cache): stop a whole-prompt hit from duplicating the last token Replaying an identical prompt hits the prompt cache over every prompt token. The adopt restored all of them, and admission then backed only the prefill cursor off to `len - 1`, so prefill forwarded the last prompt token onto a cache that already held it. Measured on CUDA with per-layer offsets after the clamp: the cache held `len` tokens (Gemma 3 115/115, Gemma 4 119/119, Qwen3 128/128, Llama 3.2 160/160), which is hypothesis (a) of the issue, the same broadcast abort #2096 reported on M5. `try_adopt_cached_prefix` now caps every hit at `len - 1` (`whole_prompt_reuse_cap`). On the snapshot branch the cap turns the restore into a truncating one; a model that refuses that truncation (recurrent state, wrapped sliding ring) falls back to cold prefill and records a `layout_constraints` reject. On the KV branch the dense truncate and paged block floors install the capped length. The admission back-off stays as a warn-level backstop. `usage.cached_tokens` now reports the restored length instead of over-counting by one. Tests: three tiny-Gemma 3 scheduler tests read the restored per-layer offset from the model's own snapshot, including the refusal fallback, and an ignored e2e test in tests/prompt_cache_e2e.rs asserts `cached_tokens == prompt_tokens - 1` and byte-identical cold and warm replies (it fails on the pre-fix binary with 20 vs 19). Closes #1760 --- src/server/batch/scheduler/admission.rs | 20 +- src/server/batch/scheduler/mod.rs | 4 + src/server/batch/scheduler/prompt_cache.rs | 96 ++++++++- .../scheduler_model_owned_cache_tests.rs | 39 ++-- .../batch/scheduler_whole_prompt_hit_tests.rs | 189 ++++++++++++++++++ src/server/prompt_cache/metrics.rs | 5 +- tests/prompt_cache_e2e.rs | 112 +++++++++++ 7 files changed, 438 insertions(+), 27 deletions(-) create mode 100644 src/server/batch/scheduler_whole_prompt_hit_tests.rs diff --git a/src/server/batch/scheduler/admission.rs b/src/server/batch/scheduler/admission.rs index 85cf5442f..15089439a 100644 --- a/src/server/batch/scheduler/admission.rs +++ b/src/server/batch/scheduler/admission.rs @@ -494,16 +494,22 @@ impl BatchScheduler { self.prompt_cache_seq_ctx.insert(seq_id, ctx); } - // Guard against a degenerate cache hit where the adopted prefix - // covers the entire tokenized prompt. This can legitimately happen - // when a client replays an identical prompt. Back off one token so - // the prefill path still runs and the sampler sees fresh logits. + // A cache hit never adopts the whole prompt: `try_adopt_cached_prefix` + // stops one token short (`whole_prompt_reuse_cap`, #1760) so the + // prefill's forward over the last prompt token lands on a cache that + // does not already hold it. Backing the prefill cursor off here + // instead, as this guard used to, left the restored cache holding all + // `len` tokens and then appended the last one a second time. The + // guard stays as a backstop for the sampler (prefill must forward at + // least one token), and a hit reaching it is a broken invariant. let prefill_start_offset = if prefill_start_offset >= prompt_tokens.len() && !prompt_tokens.is_empty() { - tracing::debug!( + tracing::warn!( seq_id = %seq_id, - "prompt-cache hit covered the entire prompt; re-running the \ - last token through prefill to produce a sampling logit" + adopted = prefill_start_offset, + prompt_len = prompt_tokens.len(), + "prompt-cache adopt covered the entire prompt, which the reuse cap \ + should prevent; re-running the last token on top of the adopted cache" ); prompt_tokens.len() - 1 } else { diff --git a/src/server/batch/scheduler/mod.rs b/src/server/batch/scheduler/mod.rs index 634a14c5a..371f7c0d5 100644 --- a/src/server/batch/scheduler/mod.rs +++ b/src/server/batch/scheduler/mod.rs @@ -1024,6 +1024,10 @@ mod scheduler_muse_glimmer_parallel_tests; #[path = "../scheduler_model_owned_cache_tests.rs"] mod scheduler_model_owned_cache_tests; +#[cfg(test)] +#[path = "../scheduler_whole_prompt_hit_tests.rs"] +mod scheduler_whole_prompt_hit_tests; + /// Resolve a request's context-retention values against the server policy /// (#1472), pure so the arithmetic is unit-testable without a scheduler. /// diff --git a/src/server/batch/scheduler/prompt_cache.rs b/src/server/batch/scheduler/prompt_cache.rs index 69c05a62f..43eaf6676 100644 --- a/src/server/batch/scheduler/prompt_cache.rs +++ b/src/server/batch/scheduler/prompt_cache.rs @@ -171,12 +171,55 @@ impl BatchScheduler { } if let SnapshotLookupOutcome::Hit { entry: snapshot_entry, - matched_len, + matched_len: store_matched_len, } = snapshot_outcome { - // A `matched_len` shorter than the stored entry means the store - // adopted at the longest common prefix, which only happens when - // the model agreed it could truncate there (#1145). + // The store reports a match that covers the whole stored entry + // (`whole_entry`) or one it can only adopt by truncating (#1145). + // The multimodal gate below keys on that, not on the reuse cap. + let whole_entry = store_matched_len >= snapshot_entry.tokens.len(); + // Whole-prompt hit (#1760): never restore the last prompt token. + // The sampler needs logits from a forward over at least one token, + // so prefill always re-runs `tokens[len - 1]`; restoring it as well + // would put that token in the cache twice. + let matched_len = whole_prompt_reuse_cap(store_matched_len, tokens.len()); + if matched_len < store_matched_len { + let truncatable = matched_len >= store.min_prefix_tokens().max(1) + && snapshot_entry.with_snapshot(|snapshot| { + self.model.snapshot_truncatable_to(snapshot, matched_len) + }); + if !truncatable { + // Recurrent families, and rotating families whose ring + // has wrapped, cannot drop the last token. A cold prefill + // is the only correct answer; a partial or duplicated + // restore would silently change the output. + tracing::debug!( + matched = store_matched_len, + reuse_cap = matched_len, + stored = snapshot_entry.tokens.len(), + "prompt-cache hit covered the entire prompt but the snapshot cannot \ + be truncated to leave the last token for prefill; falling back to \ + cold prefill" + ); + self.batch_observability + .record_prompt_cache_reject_detailed( + PromptCacheRejectReason::LayoutConstraints, + None, + matched_len, + Some(snapshot_entry.tokens.len()), + ); + return None; + } + tracing::debug!( + matched = store_matched_len, + restored = matched_len, + "prompt-cache hit covered the entire prompt; restoring all but the last \ + token so prefill re-runs it to produce a sampling logit" + ); + } + // A `matched_len` shorter than the stored entry means the restore + // truncates: the store adopted at the longest common prefix + // (#1145), or the whole-prompt cap above dropped the last token. let partial = matched_len < snapshot_entry.tokens.len(); // #124 step c, on the snapshot branch. Same rule and same reason as // the KV branch below: a partial match can leave image or audio @@ -187,7 +230,11 @@ impl BatchScheduler { // 4 VL and Unified wrappers, which already forwarded // `snapshot_truncatable_to`; an exact-prefix restore is unaffected // because it leaves `partial` false. - if require_whole_entry && partial { + // A whole-entry match that only the reuse cap shortened leaves the + // last prompt token as the suffix, which is exactly what this + // request forwarded through the token path before #1760, so it + // stays adoptable here. + if require_whole_entry && !whole_entry { tracing::debug!( matched = matched_len, stored = snapshot_entry.tokens.len(), @@ -276,21 +323,43 @@ impl BatchScheduler { if self.model.sequence_state_layout().backend == SequenceStateBackend::ModelOwned { return None; } - let (entry, matched_len) = store.lookup_longest_prefix(&key, tokens)?; + let (entry, store_matched_len) = store.lookup_longest_prefix(&key, tokens)?; // #124 step c: multimodal sharing requires the matched prefix to cover // the ENTIRE stored entry. A partial (e.g. APC block-clamped) match // could leave image/audio placeholder tokens in the suffix, which the // token-path suffix prefill would mis-handle. Decline here (falling // back to a cold prefill) before consuming anything; the entry stays // available for a later exact match. - if require_whole_entry && matched_len < entry.tokens.len() { + if require_whole_entry && store_matched_len < entry.tokens.len() { self.batch_observability.record_prompt_cache_reject( PromptCacheRejectReason::ModeMismatch, None, + store_matched_len, + ); + return None; + } + // Whole-prompt hit (#1760), same rule as the snapshot branch: adopt + // at most `len - 1` tokens so the prefill's forward over the last + // prompt token lands on a cache that does not already hold it. The + // dense truncate and the paged block floors below then install + // exactly this many (or fewer) tokens. + let matched_len = whole_prompt_reuse_cap(store_matched_len, tokens.len()); + if matched_len < store.min_prefix_tokens().max(1) { + self.batch_observability.record_prompt_cache_reject( + PromptCacheRejectReason::PrefixTooShort, + None, matched_len, ); return None; } + if matched_len < store_matched_len { + tracing::debug!( + matched = store_matched_len, + adopted = matched_len, + "prompt-cache hit covered the entire prompt; adopting all but the last token \ + so prefill re-runs it to produce a sampling logit" + ); + } // Length the adopted cache actually covers. The dense path truncates // to exactly `matched_len`; the paged paths floor to the pool block // boundary (#225), so they report their own value. @@ -1373,3 +1442,16 @@ impl BatchScheduler { self.drain_store_paged_releases(); } } + +/// Longest prefix a request of `prompt_len` tokens may adopt from a cache hit +/// that matched `matched_len` of them (#1760). +/// +/// Prefill must forward at least the last prompt token so the sampler has a +/// logit row to draw the first generated token from. When a hit covers the +/// whole prompt (a client replaying an identical prompt), the adopt therefore +/// stops one token short, and the prefill's forward over that last token lands +/// on the position it belongs to. Adopting all `prompt_len` tokens and then +/// re-running the last one would leave it in the cache twice. +pub(super) fn whole_prompt_reuse_cap(matched_len: usize, prompt_len: usize) -> usize { + matched_len.min(prompt_len.saturating_sub(1)) +} diff --git a/src/server/batch/scheduler_model_owned_cache_tests.rs b/src/server/batch/scheduler_model_owned_cache_tests.rs index 6cdb4a864..007878a11 100644 --- a/src/server/batch/scheduler_model_owned_cache_tests.rs +++ b/src/server/batch/scheduler_model_owned_cache_tests.rs @@ -58,7 +58,7 @@ fn put(weights: &mut WeightMap, key: &str, shape: &[i32], value: f32) { weights.insert(key.to_string(), tensor(shape, value)); } -fn tiny_gemma3_args() -> Gemma3ModelArgs { +pub(super) fn tiny_gemma3_args() -> Gemma3ModelArgs { Gemma3ModelArgs { model_type: "gemma3_text".to_string(), hidden_size: HIDDEN as usize, @@ -80,7 +80,7 @@ fn tiny_gemma3_args() -> Gemma3ModelArgs { } } -fn tiny_gemma3_weights() -> WeightMap { +pub(super) fn tiny_gemma3_weights() -> WeightMap { let mut w: WeightMap = WeightMap::new(); put(&mut w, "model.embed_tokens.weight", &[VOCAB, HIDDEN], 0.05); put( @@ -150,13 +150,13 @@ fn tiny_gemma3_weights() -> WeightMap { w } -fn tiny_gemma3() -> Gemma3Wrapper { +pub(super) fn tiny_gemma3() -> Gemma3Wrapper { let args = tiny_gemma3_args(); let weights = tiny_gemma3_weights(); Gemma3Wrapper::new(Gemma3Model::from_weights(&weights, &args).expect("tiny gemma3 loads")) } -fn test_store() -> Arc { +pub(super) fn test_store() -> Arc { Arc::new(PromptCacheStore::with_config(PromptCacheConfig::new( true, 1 << 20, @@ -168,10 +168,19 @@ fn test_store() -> Arc { /// A scheduler shaped like the default server: batching on, paged decode /// storage, prompt cache installed. -fn scheduler(store: Arc) -> BatchScheduler { +pub(super) fn scheduler(store: Arc) -> BatchScheduler { + scheduler_with_model(tiny_gemma3(), store) +} + +/// [`scheduler`] around a caller-built Gemma 3 (for example one whose only +/// layer is a sliding one). +pub(super) fn scheduler_with_model( + model: Gemma3Wrapper, + store: Arc, +) -> BatchScheduler { let (_tx, rx) = mpsc::channel(); let sched = BatchScheduler::with_config( - LoadedModel::Gemma3(tiny_gemma3()), + LoadedModel::Gemma3(model), MlxcelTokenizer::stub(), vec![EOS_ID], rx, @@ -190,7 +199,7 @@ fn scheduler(store: Arc) -> BatchScheduler { sched } -fn cache_ctx() -> PromptCacheRequestContext { +pub(super) fn cache_ctx() -> PromptCacheRequestContext { PromptCacheRequestContext { model_id: "tiny-gemma3".to_string(), lora_id: None, @@ -202,7 +211,7 @@ fn cache_ctx() -> PromptCacheRequestContext { } } -fn options() -> ServerGenerateOptions { +pub(super) fn options() -> ServerGenerateOptions { ServerGenerateOptions { n_indent: 0, t_max_predict_ms: None, @@ -230,7 +239,10 @@ fn options() -> ServerGenerateOptions { } } -fn enqueue(sched: &mut BatchScheduler, prompt_tokens: Vec) -> mpsc::Receiver { +pub(super) fn enqueue( + sched: &mut BatchScheduler, + prompt_tokens: Vec, +) -> mpsc::Receiver { let (tx, rx) = mpsc::channel(); sched.enqueue_request( "prompt".to_string(), @@ -246,7 +258,10 @@ fn enqueue(sched: &mut BatchScheduler, prompt_tokens: Vec) -> mpsc::Receive rx } -fn run_to_completion(sched: &mut BatchScheduler, rx: &mpsc::Receiver) { +pub(super) fn run_to_completion( + sched: &mut BatchScheduler, + rx: &mpsc::Receiver, +) -> crate::server::model_provider::GenerationResult { for _ in 0..8 { match sched.decide_action() { BatchSchedulerAction::Prefill(id) => sched.execute_prefill(id), @@ -263,7 +278,7 @@ fn run_to_completion(sched: &mut BatchScheduler, rx: &mpsc::Receiver {} - Ok(GenerateEvent::Done(_)) => return, + Ok(GenerateEvent::Done(result)) => return result, Ok(GenerateEvent::Error(err)) => panic!("unexpected generation error: {err}"), Err(err) => panic!("generation did not finish: {err}"), } @@ -271,7 +286,7 @@ fn run_to_completion(sched: &mut BatchScheduler, rx: &mpsc::Receiver Vec { +pub(super) fn prompt(n: usize) -> Vec { (0..n).map(|i| (i % 6) as i32).collect() } diff --git a/src/server/batch/scheduler_whole_prompt_hit_tests.rs b/src/server/batch/scheduler_whole_prompt_hit_tests.rs new file mode 100644 index 000000000..50e764c35 --- /dev/null +++ b/src/server/batch/scheduler_whole_prompt_hit_tests.rs @@ -0,0 +1,189 @@ +// Copyright 2025-2026 Lablup Inc. +// +// 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. + +//! Whole-prompt prompt-cache hits (issue #1760). +//! +//! A client that replays an identical prompt gets a cache hit covering every +//! prompt token. Prefill still has to forward the last one to produce a +//! sampling logit, so the adopt must install `len - 1` tokens and leave the +//! last for prefill. Before #1760 the adopt installed all `len` and admission +//! backed only the prefill cursor off, so the last token landed in the cache +//! twice and the warm reply diverged from the cold one. +//! +//! These tests read the restored per-layer cache offset itself, through the +//! model's own snapshot, rather than only the offset the scheduler reports. + +use super::scheduler_model_owned_cache_tests::{ + cache_ctx, enqueue, options, prompt, run_to_completion, scheduler, scheduler_with_model, + test_store, tiny_gemma3_args, tiny_gemma3_weights, +}; +use super::*; +use mlxcel_core::generate::LanguageModel; +use std::sync::atomic::AtomicBool; +use std::sync::{Arc, mpsc}; + +use crate::models::{Gemma3Model, Gemma3Wrapper}; +use crate::server::prompt_cache::PromptCacheRejectReason; + +/// Per-layer cache offset of `seq_id`'s live model-owned state, read from the +/// snapshot the model itself would donate. The tiny fixture has one layer. +fn restored_offset(sched: &BatchScheduler, seq_id: SequenceId, kind: &str) -> i32 { + let snapshot = sched + .model + .snapshot_sequence_state(seq_id, 0) + .expect("restored sequence has model-owned state"); + let name = format!("layer0.{kind}.offset"); + let offset = snapshot + .tensor(&name) + .unwrap_or_else(|| panic!("snapshot carries {name}")); + mlxcel_core::item_i32(offset) +} + +/// The fixture with its only layer made a sliding one (window 8), so any +/// prompt longer than the window wraps the ring and the model refuses every +/// truncating restore. +fn tiny_sliding_gemma3() -> Gemma3Wrapper { + let mut args = tiny_gemma3_args(); + args.sliding_window_pattern = 2; + Gemma3Wrapper::new( + Gemma3Model::from_weights(&tiny_gemma3_weights(), &args).expect("tiny gemma3 loads"), + ) +} + +/// Turn 1 on `first`, returning the full stored conversation +/// (`first ++ generated`), which is the completion snapshot's token vector. +/// +/// The tiny fixture samples EOS straight away, so `ignore_eos` with +/// `max_tokens == 1` is what makes the stored entry one token longer than the +/// prompt; without it the stored entry is the prompt itself. +fn donate_turn(sched: &mut BatchScheduler, first: &[i32], ignore_eos: bool) -> Vec { + let (tx, rx) = mpsc::channel(); + let mut opts = options(); + opts.ignore_eos = ignore_eos; + sched.enqueue_request( + "prompt".to_string(), + Some(first.to_vec()), + opts, + Vec::new(), + Vec::new(), + Vec::new(), + tx, + Arc::new(AtomicBool::new(false)), + true, + ); + let result = run_to_completion(sched, &rx); + let mut stored = first.to_vec(); + stored.extend_from_slice(&result.generated_token_ids); + stored +} + +/// An identical replay of turn 1's prompt: the stored completion entry is +/// longer than the prompt, so the store adopts by truncating to the prompt +/// length. The adopt must go one further and leave the last prompt token for +/// prefill, and the restored cache must hold exactly that many tokens. +#[test] +fn identical_prompt_replay_restores_all_but_the_last_token() { + let mut sched = scheduler(test_store()); + let first = prompt(40); + let stored = donate_turn(&mut sched, &first, true); + assert_eq!( + stored.len(), + first.len() + 1, + "turn 1 stored one generated token" + ); + + let (seq_id, adopted) = sched + .try_adopt_cached_prefix(&cache_ctx(), &first, false) + .expect("an identical replay hits the cache"); + assert_eq!(adopted, first.len() - 1, "adopt stops one token short"); + assert_eq!( + restored_offset(&sched, seq_id, "standard"), + (first.len() - 1) as i32, + "the restored cache must not already hold the last prompt token" + ); +} + +/// The exact-entry shape: the request is the stored conversation itself, so +/// the store matches the whole entry without truncating. The cap must turn +/// that into a truncating restore of `len - 1`, and the admission path must +/// then prefill exactly the one remaining token. +#[test] +fn whole_entry_hit_restores_len_minus_one_and_prefills_one_token() { + let store = test_store(); + let mut sched = scheduler(store.clone()); + let replay = donate_turn(&mut sched, &prompt(40), false); + + let (seq_id, adopted) = sched + .try_adopt_cached_prefix(&cache_ctx(), &replay, false) + .expect("the stored conversation hits its own entry"); + assert_eq!(adopted, replay.len() - 1); + assert_eq!( + restored_offset(&sched, seq_id, "standard"), + (replay.len() - 1) as i32 + ); + sched.release_sequence_caches(seq_id); + + // Through admission: the queued sequence starts prefill at `len - 1` and + // reports the same count as cached, so `usage.cached_tokens` never claims + // the token prefill re-ran. + let _rx = enqueue(&mut sched, replay.clone()); + let queued = sched + .prefill_queue + .dequeue() + .expect("the replay is queued for prefill"); + assert_eq!(queued.prefill_start_offset, replay.len() - 1); + assert_eq!(queued.already_cached_tokens, replay.len() - 1); + assert_eq!( + restored_offset(&sched, queued.seq_id, "standard"), + (replay.len() - 1) as i32 + ); +} + +/// A family that cannot truncate (here a wrapped sliding ring) must fall back +/// to a cold prefill on a whole-prompt hit, record the decline, and never +/// install the full entry for prefill to append the last token onto. +#[test] +fn untruncatable_whole_entry_hit_falls_back_to_cold_prefill() { + let store = test_store(); + let mut sched = scheduler_with_model(tiny_sliding_gemma3(), store.clone()); + let replay = donate_turn(&mut sched, &prompt(40), false); + assert_eq!(store.stats().snapshot_entries, 1, "turn 1 donated"); + + let before = sched + .batch_observability + .prompt_cache_reject_reasons + .count(PromptCacheRejectReason::LayoutConstraints); + let adopted = sched.try_adopt_cached_prefix(&cache_ctx(), &replay, false); + assert!( + adopted.is_none(), + "a wrapped ring cannot drop the last token, so the hit must decline" + ); + assert_eq!( + sched + .batch_observability + .prompt_cache_reject_reasons + .count(PromptCacheRejectReason::LayoutConstraints), + before + 1, + "the decline is recorded" + ); + + // Admission then prefills the whole prompt cold. + let _rx = enqueue(&mut sched, replay.clone()); + let queued = sched + .prefill_queue + .dequeue() + .expect("the replay is queued for prefill"); + assert_eq!(queued.prefill_start_offset, 0); + assert_eq!(queued.already_cached_tokens, 0); +} diff --git a/src/server/prompt_cache/metrics.rs b/src/server/prompt_cache/metrics.rs index 8e7c4276a..bc80d2994 100644 --- a/src/server/prompt_cache/metrics.rs +++ b/src/server/prompt_cache/metrics.rs @@ -59,7 +59,10 @@ pub enum PromptCacheRejectReason { /// capturable state). EmptySet, /// A pool-level operation failed: paged clone, dense truncate, paged - /// trim, snapshot restore, or the final `adopt`/`adopt_paged` call. + /// trim, snapshot restore, or the final `adopt`/`adopt_paged` call. Also + /// recorded when a snapshot hit covers the whole prompt and the model + /// refuses to truncate it by the one token prefill must re-run (#1760), + /// so the request falls back to a cold prefill. LayoutConstraints, /// The paged pool's block-size floor pushed the adoptable/donatable /// length below `min_prefix_tokens`. diff --git a/tests/prompt_cache_e2e.rs b/tests/prompt_cache_e2e.rs index aa41ba5ac..378867e8f 100644 --- a/tests/prompt_cache_e2e.rs +++ b/tests/prompt_cache_e2e.rs @@ -527,3 +527,115 @@ async fn multi_turn_chat_with_cache_disabled_never_reports_cached_tokens() { stop_server(&mut child); } + +/// One greedy, seeded chat completion of a single user turn. Returns the whole +/// assistant message (reasoning plus content, so a thinking model's output is +/// compared too), `usage.prompt_tokens` and `usage.prompt_tokens_details.cached_tokens`. +async fn replay_turn( + client: &reqwest::Client, + base_url: &str, + model_alias: &str, + prompt: &str, +) -> (String, u64, Option) { + let body = serde_json::json!({ + "model": model_alias, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": 48, + "temperature": 0.0, + "seed": 0, + "user": "prompt-cache-e2e-replay", + }); + let resp = client + .post(format!("{base_url}/v1/chat/completions")) + .json(&body) + .send() + .await + .expect("send chat request"); + let status = resp.status(); + let value: serde_json::Value = resp.json().await.expect("parse chat response JSON"); + assert!(status.is_success(), "status={status} body={value}"); + let message = &value["choices"][0]["message"]; + let text = format!( + "{}\u{0}{}", + message["reasoning_content"].as_str().unwrap_or_default(), + message["content"].as_str().unwrap_or_default() + ); + let prompt_tokens = value["usage"]["prompt_tokens"].as_u64().unwrap_or(0); + let cached = value["usage"]["prompt_tokens_details"]["cached_tokens"].as_u64(); + (text, prompt_tokens, cached) +} + +/// Issue #1760: replaying an identical prompt hits the cache over the whole +/// prompt. The adopt must restore `prompt_tokens - 1` tokens and let prefill +/// re-run the last one, so the warm reply is byte-identical to the cold one. +/// Before the fix the restore installed all `prompt_tokens` and the last token +/// landed in the cache twice, and the replies diverged. +/// +/// `--apc-block-size 1` (dense backend via `--batch-size 1`) lets the KV +/// lookup match the full prompt at any length instead of flooring it to a +/// 16-token block, so the whole-prompt hit is reached on every run. +/// `cached_tokens` is the restored length the scheduler installed, which is +/// the cache offset prefill starts from. +#[tokio::test] +#[ignore = "requires local model weights (qwen3-0.6b-4bit) and the mlxcel-server binary"] +async fn identical_prompt_replay_restores_all_but_the_last_token() { + let model_dir = repo_model_dir(QWEN3_MODEL); + let binary = repo_binary_path("mlxcel-server"); + if !model_dir.exists() || !binary.exists() { + eprintln!("Skipping: model or mlxcel-server binary not present"); + return; + } + let port = reserve_port(); + let base_url = format!("http://127.0.0.1:{port}"); + let port_str = port.to_string(); + let model_str = model_dir.to_string_lossy().to_string(); + let model_alias = "qwen3-replay"; + let mut child = spawn_server(&[ + "--model", + &model_str, + "--alias", + model_alias, + "--host", + "127.0.0.1", + "--port", + &port_str, + "--parallel", + "1", + "--batch-size", + "1", + "--no-warmup", + "--prompt-cache-enabled=true", + "--prompt-cache-min-prefix", + "4", + "--apc-block-size", + "1", + ]); + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(120)) + .build() + .expect("build reqwest client"); + if !wait_for_health_soft(&client, &base_url, Duration::from_secs(90)).await { + eprintln!("Skipping: mlxcel-server did not become healthy at {base_url}"); + stop_server(&mut child); + return; + } + + let prompt = "List three facts about the Moon, one short sentence each."; + let (cold, cold_prompt_tokens, cold_cached) = + replay_turn(&client, &base_url, model_alias, prompt).await; + let (warm, warm_prompt_tokens, warm_cached) = + replay_turn(&client, &base_url, model_alias, prompt).await; + stop_server(&mut child); + + assert_eq!(cold_cached, Some(0), "the first request is a cold prefill"); + assert_eq!(cold_prompt_tokens, warm_prompt_tokens); + assert_eq!( + warm_cached, + Some(warm_prompt_tokens - 1), + "a whole-prompt hit restores every prompt token but the last" + ); + assert_eq!( + warm, cold, + "the warm replay must match the cold request byte for byte" + ); +} From 1381d5968d3a7ef9d1725d2dbb5901eeb868344e Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Tue, 6 Oct 2026 18:08:53 +0900 Subject: [PATCH 2/2] docs: add technical report for PR #2136 --- .../pr-2136-whole-prompt-hit-reuse-cap.en.md | 65 +++++++++++++++++++ .../pr-2136-whole-prompt-hit-reuse-cap.ko.md | 65 +++++++++++++++++++ 2 files changed, 130 insertions(+) create mode 100644 TECHNICAL_REPORTS/pr-2136-whole-prompt-hit-reuse-cap.en.md create mode 100644 TECHNICAL_REPORTS/pr-2136-whole-prompt-hit-reuse-cap.ko.md diff --git a/TECHNICAL_REPORTS/pr-2136-whole-prompt-hit-reuse-cap.en.md b/TECHNICAL_REPORTS/pr-2136-whole-prompt-hit-reuse-cap.en.md new file mode 100644 index 000000000..0e8acf138 --- /dev/null +++ b/TECHNICAL_REPORTS/pr-2136-whole-prompt-hit-reuse-cap.en.md @@ -0,0 +1,65 @@ +# PR #2136: A whole-prompt cache hit no longer duplicates the last prompt token + +**Date**: 2026-10-06 +**Status**: Verified on CUDA (GB10); the M5 tile-padded abort path is not reproducible on this host +**Risk**: Low + +## Summary + +When a client replays an identical prompt, the prompt cache matches every prompt token. `try_adopt_cached_prefix` restored all `len` tokens, and admission then moved only the prefill cursor back to `len - 1` so the sampler would get a logit row. Prefill therefore forwarded the last prompt token onto a cache that already held it, at position `len` instead of `len - 1`. The PR caps every adopt at `len - 1` on both the snapshot and KV branches, so the re-run token lands where it belongs. + +Closes #1760. + +## 1. Problem + +Issue #1760 left two explanations open: (a) a real duplicate in the cache, or (b) only the forward-width rounding effect, because warm forwards one token where cold forwards the whole prompt. Per-layer cache offsets read right after the admission clamp settle it as (a): + +| Checkpoint | Path | prompt_len | offset after clamp | +|---|---|---|---| +| gemma-3-4b-it-4bit | snapshot (truncating restore) | 115 | 115 | +| gemma-4-12b-it-4bit | snapshot (truncating restore) | 119 | 119 | +| qwen3-4b-4bit | paged KV | 128 | 128 | +| llama-3.2-1b-4bit | paged KV | 160 | 160 | + +On the KV path a whole-prompt hit needs a prompt length that is a multiple of the paged block (32 by default), because the clone path floors to whole blocks. That matches the 32- and 48-token replays that #2096 saw abort on M5, where the padded mask was built for `len - 1` cached rows over a cache holding `len`. + +## 2. Change + +- `whole_prompt_reuse_cap(matched_len, prompt_len) = min(matched_len, prompt_len - 1)` is applied right after each store lookup in `src/server/batch/scheduler/prompt_cache.rs`. +- Snapshot branch: when the cap shortens the match, the model is asked `snapshot_truncatable_to(snapshot, len - 1)`. If it agrees, the restore goes through `restore_sequence_state_truncated`. If it refuses (recurrent state, or a sliding ring that has wrapped), the request falls back to cold prefill and records a `layout_constraints` reject with the stored entry length. A partial restore is never installed. +- KV branch: the cap runs before the dense `truncate_to` and the paged block floors, which then install `len - 1` or fewer tokens. A cap below `min_prefix_tokens` declines as `prefix_too_short`. +- The multimodal whole-entry gate (#124 step c) keys on the store's match, not on the capped value. A whole-entry VLM replay therefore still adopts, and its one-token suffix is exactly what the old back-off already forwarded through the token path. +- The admission back-off in `admission.rs` stays as a backstop and logs at warn, since reaching it now means the invariant broke. +- Side effect: `usage.prompt_tokens_details.cached_tokens` and the shared-budget charge report the length actually restored, one less than before for a whole-prompt hit. + +## 3. Decision: cap in the adopt, not in admission + +The issue offered the same direction. Moving only the prefill cursor cannot work, because the restore has already installed the token. Restoring `len - 1` makes the forward land on the correct slot for every family, and the truncation predicate the store already uses decides whether that is possible. Recording the refusal under the existing `layout_constraints` reason avoids widening the fixed reject enum and the `/v1/cache/stats` schema for one decline site. + +## 4. Validation (GB10, CUDA, `mlxcel-server`, prompt cache on, temperature 0, seed 0, 200 tokens) + +| Checkpoint | Before | After | +|---|---|---| +| gemma-3-4b-it-4bit | warm diverged at char 226 | 3/3 byte-identical to cold | +| llama-3.2-1b-4bit | diverged at char 230 | 3/3 byte-identical | +| gemma-4-12b-it-4bit | diverged at token 71 | diverges at token 27, a bf16 tie (warm top-2 gap 0.0) | +| qwen3-4b-4bit | diverged at token 76 | diverges at token 82, a bf16 tie (warm top-2 gap 0.0) | +| qwen3.5-0.8b-4bit (recurrent) | identical (boundary-entry hit) | identical | + +The Gemma 4 and Qwen3 residuals are the forward-width effect (b): the warm request forwards 1 token (Gemma 4) or 32 tokens (Qwen3, after the paged block floor) where cold forwards a wider chunk. A control that removes the width difference (`--prefill-chunk-size 1` in both arms, boundary snapshot disabled for Gemma 4) separates the two causes. After the fix, both checkpoints are token-identical with bitwise-equal logprobs. Before the fix, both still diverge, and Gemma 4 already differs at the first token's logprob (-0.5 vs -0.625). + +Tests: + +- `src/server/batch/scheduler_whole_prompt_hit_tests.rs` has three tests on the tiny real Gemma 3. They read the restored offset from the model's own snapshot for an identical replay (truncating), a whole-entry hit through admission (`prefill_start_offset == already_cached_tokens == len - 1`), and a sliding-window refusal that falls back to cold with the reject counted. +- `tests/prompt_cache_e2e.rs::identical_prompt_replay_restores_all_but_the_last_token` is ignored and runs qwen3-0.6b-4bit with `--apc-block-size 1`. It asserts `cached_tokens == prompt_tokens - 1` and byte-identical replies. It fails on the pre-fix binary (20 vs 19). + +## 5. Not verified here + +- `trinity-nano-preview-4bit` (AFMoE) is not present locally. Qwen3 and Llama 3.2 cover the KV branch instead. +- The M5 broadcast abort needs Metal hardware. +- A real-checkpoint refusal was not reachable through chat requests, because recurrent families hit their shorter boundary entry. The unit test covers that path. + +## 6. Related + +- #1754: completion-origin snapshot token counts versus the offsets the caches hold. The measurements above also show `offset == token_len - 1` for `length` finishes on Gemma 4 and the KV path. +- #2096 / PR #2097: the pool-backed trim fix that reported the M5 abort. diff --git a/TECHNICAL_REPORTS/pr-2136-whole-prompt-hit-reuse-cap.ko.md b/TECHNICAL_REPORTS/pr-2136-whole-prompt-hit-reuse-cap.ko.md new file mode 100644 index 000000000..d0450fb66 --- /dev/null +++ b/TECHNICAL_REPORTS/pr-2136-whole-prompt-hit-reuse-cap.ko.md @@ -0,0 +1,65 @@ +# PR #2136: 프롬프트 전체 캐시 히트가 마지막 프롬프트 토큰을 중복시키지 않도록 수정 + +**작성일**: 2026-10-06 +**상태**: CUDA(GB10)에서 검증 완료. M5 tile-padded abort 경로는 이 호스트에서 재현 불가 +**위험도**: 낮음 + +## 요약 + +클라이언트가 동일한 프롬프트를 다시 보내면 prompt cache가 프롬프트의 모든 토큰과 일치한다. `try_adopt_cached_prefix`는 `len`개 토큰을 모두 복원했고, admission은 sampler가 logit을 얻도록 prefill 커서만 `len - 1`로 되돌렸다. 그 결과 prefill은 이미 캐시에 들어 있는 마지막 프롬프트 토큰을 `len - 1`이 아닌 `len` 위치에 한 번 더 넣었다. 이 PR은 snapshot 경로와 KV 경로 모두에서 adopt 길이를 `len - 1`로 제한해, 다시 실행되는 토큰이 제자리에 들어가게 한다. + +Closes #1760. + +## 1. 문제 정의 + +이슈 #1760은 두 가설을 열어 두었다. (a) 캐시에 실제 중복이 있다. (b) warm 요청은 토큰 1개만, cold 요청은 프롬프트 전체를 forward하는 데서 오는 forward 폭 반올림 효과일 뿐이다. admission clamp 직후 측정한 layer별 캐시 offset으로 (a)임을 확인했다. + +| 체크포인트 | 경로 | prompt_len | clamp 직후 offset | +|---|---|---|---| +| gemma-3-4b-it-4bit | snapshot (truncating restore) | 115 | 115 | +| gemma-4-12b-it-4bit | snapshot (truncating restore) | 119 | 119 | +| qwen3-4b-4bit | paged KV | 128 | 128 | +| llama-3.2-1b-4bit | paged KV | 160 | 160 | + +KV 경로는 clone 경로가 블록 단위로 내림하기 때문에, 프롬프트 길이가 paged 블록(기본 32)의 배수일 때만 프롬프트 전체 히트가 된다. #2096이 M5에서 32/48 토큰 replay에서만 abort를 본 것과 맞는다. 그 abort는 padded mask를 `len - 1`개의 캐시 행 기준으로 만들었는데 캐시에는 `len`개가 있어서 발생했다. + +## 2. 변경 내용 + +- `src/server/batch/scheduler/prompt_cache.rs`에서 store lookup 직후 `whole_prompt_reuse_cap(matched_len, prompt_len) = min(matched_len, prompt_len - 1)`을 적용한다. +- Snapshot 경로: 상한 때문에 일치 길이가 줄어들면 모델에 `snapshot_truncatable_to(snapshot, len - 1)`을 묻는다. 모델이 동의하면 `restore_sequence_state_truncated`로 복원한다. 거부하면(recurrent state, 또는 wrap된 sliding ring) cold prefill로 되돌아가고 저장된 엔트리 길이와 함께 `layout_constraints` reject를 기록한다. 부분 복원은 설치하지 않는다. +- KV 경로: 상한을 dense `truncate_to`와 paged 블록 내림보다 먼저 적용하므로, 이후 단계는 `len - 1`개 이하만 설치한다. 상한이 `min_prefix_tokens`보다 작으면 `prefix_too_short`로 거절한다. +- 멀티모달 whole-entry 게이트(#124 step c)는 상한 적용 값이 아니라 store가 보고한 일치 길이를 기준으로 판단한다. 그래서 whole-entry VLM replay는 계속 adopt되며, 남는 토큰 1개짜리 suffix는 이전 back-off가 이미 token 경로로 forward하던 것과 같다. +- `admission.rs`의 back-off는 안전장치로 남기고 warn으로 로그한다. 이 경로에 도달하면 불변식이 깨졌다는 뜻이다. +- 부수 효과: 프롬프트 전체 히트에서 `usage.prompt_tokens_details.cached_tokens`와 shared budget 계산이 실제로 복원한 길이, 즉 이전보다 1 작은 값을 보고한다. + +## 3. 기술적 선택: admission이 아니라 adopt에서 제한 + +이슈가 제안한 방향과 같다. 복원이 이미 토큰을 설치한 뒤라서 prefill 커서만 옮겨서는 해결되지 않는다. `len - 1`까지만 복원하면 모든 family에서 forward가 올바른 슬롯에 들어가며, 그것이 가능한지는 store가 이미 쓰는 truncation predicate가 판단한다. 거부 사유는 기존 `layout_constraints`로 기록해, 거절 지점 하나 때문에 고정된 reject enum과 `/v1/cache/stats` 스키마를 넓히지 않았다. + +## 4. 검증 (GB10, CUDA, `mlxcel-server`, prompt cache on, temperature 0, seed 0, 200 토큰) + +| 체크포인트 | 수정 전 | 수정 후 | +|---|---|---| +| gemma-3-4b-it-4bit | 226번째 문자에서 분기 | 3/3 cold와 바이트 단위 동일 | +| llama-3.2-1b-4bit | 230번째 문자에서 분기 | 3/3 바이트 단위 동일 | +| gemma-4-12b-it-4bit | 71번째 토큰에서 분기 | 27번째 토큰에서 분기, bf16 tie (warm top-2 gap 0.0) | +| qwen3-4b-4bit | 76번째 토큰에서 분기 | 82번째 토큰에서 분기, bf16 tie (warm top-2 gap 0.0) | +| qwen3.5-0.8b-4bit (recurrent) | 동일 (boundary 엔트리 히트) | 동일 | + +Gemma 4와 Qwen3에 남은 분기는 forward 폭 효과 (b)다. warm 요청은 토큰 1개(Gemma 4) 또는 32개(Qwen3, paged 블록 내림 이후)를 forward하고 cold는 더 넓은 chunk를 forward한다. 폭 차이를 없앤 대조 실험(양쪽 모두 `--prefill-chunk-size 1`, Gemma 4는 boundary snapshot 비활성화)으로 두 원인을 분리했다. 수정 후에는 두 체크포인트 모두 토큰이 동일하고 logprob도 비트 단위로 같다. 수정 전에는 둘 다 여전히 분기하며, Gemma 4는 첫 토큰의 logprob부터 다르다(-0.5 vs -0.625). + +테스트: + +- `src/server/batch/scheduler_whole_prompt_hit_tests.rs`: tiny 실제 Gemma 3로 세 테스트를 돌린다. 모델 자체 snapshot에서 복원된 offset을 읽어 동일 replay(truncating), admission을 거친 whole-entry 히트(`prefill_start_offset == already_cached_tokens == len - 1`), sliding window 거부 시 cold fallback과 reject 집계를 확인한다. +- `tests/prompt_cache_e2e.rs::identical_prompt_replay_restores_all_but_the_last_token`: ignored 테스트로, qwen3-0.6b-4bit와 `--apc-block-size 1`로 실행한다. `cached_tokens == prompt_tokens - 1`과 바이트 단위로 동일한 응답을 확인하며, 수정 전 바이너리에서는 실패한다(20 vs 19). + +## 5. 이 호스트에서 검증하지 못한 항목 + +- `trinity-nano-preview-4bit`(AFMoE)가 로컬에 없다. 대신 Qwen3와 Llama 3.2로 KV 경로를 검증했다. +- M5 broadcast abort는 Metal 하드웨어가 필요하다. +- 실제 체크포인트에서 거부 경로는 chat 요청으로 도달하지 못했다. recurrent family가 더 짧은 boundary 엔트리에서 히트하기 때문이며, 이 경로는 단위 테스트가 다룬다. + +## 6. 관련 항목 + +- #1754: completion-origin snapshot의 token 수와 캐시가 실제로 가진 offset의 차이. 위 측정에서도 Gemma 4와 KV 경로의 `length` 종료는 `offset == token_len - 1`이었다. +- #2096 / PR #2097: M5 abort를 보고한 pool-backed trim 수정.