diff --git a/examples/backends/vllm/launch/agg_multimodal_router.sh b/examples/backends/vllm/launch/agg_multimodal_router.sh index 267b712125ef..8739f0e65cc1 100755 --- a/examples/backends/vllm/launch/agg_multimodal_router.sh +++ b/examples/backends/vllm/launch/agg_multimodal_router.sh @@ -49,7 +49,7 @@ NATS_SERVER="${NATS_SERVER:-nats://127.0.0.1:4222}" ETCD_ENDPOINTS="${ETCD_ENDPOINTS:-http://127.0.0.1:2379}" VLLM_SYSTEM_PORT_BASE="${VLLM_SYSTEM_PORT_BASE:-18081}" KV_EVENTS_PORT_BASE="${KV_EVENTS_PORT_BASE:-5557}" -DYN_LOG_VAL="${DYN_LOG:-info,lightseek_mm=debug,dynamo_kv_router::scheduling=debug}" +DYN_LOG_VAL="${DYN_LOG:-info,lightseek_mm=debug,dynamo_kv_router::scheduling=debug,dynamo_llm::kv_router=debug}" # Pass-through extra args for `python -m dynamo.vllm`. VLLM_EXTRA_ARGS="${VLLM_EXTRA_ARGS:-}" diff --git a/lib/llm/src/model_card.rs b/lib/llm/src/model_card.rs index a7e98a81b1ef..edba1efb9f37 100644 --- a/lib/llm/src/model_card.rs +++ b/lib/llm/src/model_card.rs @@ -969,7 +969,7 @@ impl ModelDeploymentCard { // extracting special-token strings. `FastTokenizer` does not re-expose // `get_added_tokens_decoder`, so we must capture specials from the raw // HF tokenizer before any swap. - let hf = HfTokenizer::from_file(p) + let mut hf = HfTokenizer::from_file(p) .inspect_err(|err| { if let Some(serde_err) = err.downcast_ref::() && let Ok(contents) = std::fs::read_to_string(p) @@ -979,7 +979,16 @@ impl ModelDeploymentCard { }) .map_err(anyhow::Error::msg) .with_context(|| p.display().to_string())?; - + // Apply the tokenizer_config.json special-token merge eagerly so + // `extract_hf_special_tokens` below sees the same specials the + // wrapped tokenizer will use. Without this the L1 prefix cache's + // boundary list would diverge from the actual tokenizer + // (e.g. Qwen2-VL's `<|image_pad|>` would be in the tokenizer + // but missing from the cache specials), letting chat prefixes + // straddle a special-token boundary and reducing hit rate. + if let Some(model_dir) = p.parent() { + crate::tokenizers::hf::merge_special_tokens_from_config(&mut hf, model_dir); + } // Hold onto specials before any move of `hf`. let specials: Vec = if cache_enabled { extract_hf_special_tokens(&hf) @@ -987,6 +996,10 @@ impl ModelDeploymentCard { Vec::new() }; + // Merge already applied above; just wrap. + let wrap_hf = + |hf: HfTokenizer| crate::tokenizers::HuggingFaceTokenizer::from_tokenizer(hf); + // Pick the inner backend. let raw: Arc = if use_fast { if let Some(path_str) = p.to_str() { @@ -1000,9 +1013,7 @@ impl ModelDeploymentCard { %e, "Failed to load fastokens, falling back to HuggingFace" ); - Arc::new(crate::tokenizers::HuggingFaceTokenizer::from_tokenizer( - hf, - )) + Arc::new(wrap_hf(hf)) } } } else { @@ -1010,10 +1021,10 @@ impl ModelDeploymentCard { path = %p.display(), "Tokenizer path contains non-UTF-8 characters, skipping fastokens; falling back to HuggingFace" ); - Arc::new(crate::tokenizers::HuggingFaceTokenizer::from_tokenizer(hf)) + Arc::new(wrap_hf(hf)) } } else { - Arc::new(crate::tokenizers::HuggingFaceTokenizer::from_tokenizer(hf)) + Arc::new(wrap_hf(hf)) }; if cache_enabled { diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index 8ac3721ec78c..7ad4f0b3b8b5 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -185,19 +185,6 @@ fn mdc_model_dir(mdc: &ModelDeploymentCard) -> Option { cf.path()?.parent().map(std::path::PathBuf::from) } -/// Find the first occurrence of `needle` in `haystack`. Linear scan; the -/// needles here are tokenized chat-template placeholders (≤ 10 tokens for -/// Phi-3-style `<|image_N|>`), so the naive O(n·m) cost is fine. -#[cfg(feature = "lightseek-mm")] -fn find_subseq(haystack: &[T], needle: &[T]) -> Option { - if needle.is_empty() || needle.len() > haystack.len() { - return None; - } - haystack - .windows(needle.len()) - .position(|window| window == needle) -} - /// Shared SSRF-aware `MediaFetcher` + `reqwest::Client` for the dim-fetch /// path used by MM-aware routing. Inherits the same policy contract as the /// frontend-decode path (`MediaLoader`): blocklist DNS resolver, redirect @@ -248,11 +235,17 @@ pub struct OpenAIPreprocessor { /// unreadable. #[cfg(feature = "lightseek-mm")] image_token_counter: Option, - /// Image-placeholder token id resolved from the model's HF JSON configs. + /// Image-placeholder token id the routing-side sequence fills per image. + /// Resolved from `config.json`'s `image_token_id` field when present, + /// otherwise falls back to lightseek's `ModelProcessorSpec` value. This + /// is the id the backend's HF processor emits in the expanded sequence + /// (per-patch token for Qwen-VL families, the single placeholder for + /// LLaVA/Phi-3), so block hashes align bit-for-bit with the worker. + /// /// `None` disables MM-aware routing for this model and the router falls /// back to text-prefix routing. #[cfg(feature = "lightseek-mm")] - image_token_id: Option, + routing_image_token_id: Option, /// Per-family flatten-time image placeholder template (e.g. /// `"<|image_{n}|>"` for Phi-3, `""` for LLaVA-1.5). Threaded /// through from the formatter so the routing path can reverse the @@ -260,6 +253,13 @@ pub struct OpenAIPreprocessor { /// tokens when the chat template uses numbered markers. #[cfg(feature = "lightseek-mm")] image_placeholder_template: Option<&'static str>, + /// BOS token id to prepend to the routing-side sequence so per-block + /// hashes match the backend's HF processor output on models with + /// `add_bos_token: true` (Phi-3-vision and other `LlamaTokenizer` + /// families). `None` when the model doesn't need it or `bos_token` + /// doesn't round-trip to a single id. + #[cfg(feature = "lightseek-mm")] + routing_prepend_bos: Option, } impl OpenAIPreprocessor { @@ -388,8 +388,31 @@ impl OpenAIPreprocessor { // lightseek registry resolve fine-tunes loaded from custom-named // directories where the family substring isn't in the path. #[cfg(feature = "lightseek-mm")] - let image_token_inputs: Option<(String, String, std::path::PathBuf)> = mdc_model_dir(&mdc) - .map(|p| (mdc.source_path().to_string(), model_info.model_type(), p)); + let model_dir_for_routing: Option = mdc_model_dir(&mdc); + // TODO(mm-routing): fastokens lacks a special-token mutator, so it + // can't merge tokenizer_config.json specials and would BPE-shatter + // placeholders (e.g. Qwen2-VL `<|image_pad|>`). Disable MM-routing + // here; remove once fastokens upstream exposes the mutator. + #[cfg(feature = "lightseek-mm")] + let image_token_inputs: Option<(String, String, std::path::PathBuf)> = { + let fastokens_active = std::env::var("DYN_TOKENIZER").as_deref() == Ok("fastokens"); + if fastokens_active && model_dir_for_routing.is_some() { + tracing::warn!( + target: "mm_routing", + "DYN_TOKENIZER=fastokens is set; MM-aware KV routing disabled. \ + Unset DYN_TOKENIZER (or set it to 'default') to re-enable." + ); + None + } else { + model_dir_for_routing.as_ref().map(|p| { + ( + mdc.source_path().to_string(), + model_info.model_type(), + p.clone(), + ) + }) + } + }; let media_loader = match mdc.media_decoder { Some(media_decoder) => Some(MediaLoader::new(media_decoder, mdc.media_fetcher)?), @@ -399,73 +422,80 @@ impl OpenAIPreprocessor { let context_length = mdc.context_length; #[cfg(feature = "lightseek-mm")] - let (image_token_counter, image_token_id) = match image_token_inputs { - Some((model_id, model_type, model_dir)) => { - // Try counter init and image-token resolution independently. - // Each carries its own reason for failure; the summary log - // below names whichever pieces are missing so operators can - // tell at a glance whether the model needs a lightseek - // upstream PR (registry miss) or a non-standard placeholder - // location (resolver miss). - let (counter, counter_err): ( - Option, - Option, - ) = match lightseek_mm::LightseekMmCounter::try_new( - &model_id, - Some(&model_type), - &model_dir, - ) { - Ok(c) => (Some(c), None), - Err(e) => (None, Some(e.to_string())), - }; - let img_tok = lightseek_mm::resolve_image_token_id(&model_id, &model_dir); - - match (counter.is_some(), img_tok.is_some()) { - (true, true) => tracing::info!( - target: "mm_routing", - model = %model_id, - model_dir = %model_dir.display(), - "MM-aware KV routing enabled (lightseek)" - ), - (counter_ok, img_ok) => { - let mut reasons: Vec = Vec::new(); - if !counter_ok { - reasons.push(format!( - "model not supported by the lightseek registry ({})", - counter_err.as_deref().unwrap_or("unknown error") - )); - } - if !img_ok { - reasons.push( - "image-placeholder token unresolvable from \ + let (image_token_counter, routing_image_token_id, bos_token_string) = + match image_token_inputs { + Some((model_id, model_type, model_dir)) => { + // Resolve counter + image-token id independently so the + // summary log can name which piece is missing. + let (counter, counter_err): ( + Option, + Option, + ) = match lightseek_mm::LightseekMmCounter::try_new( + &model_id, + Some(&model_type), + &model_dir, + ) { + Ok(c) => (Some(c), None), + Err(e) => (None, Some(e.to_string())), + }; + // One-shot config/tokenizer_config read for all + // routing-side token info. Parsing lives in + // `lightseek_mm`, next to the spec resolution. + let routing_tokens = + lightseek_mm::resolve_routing_tokens(&model_id, &model_dir); + // `chat_placeholder_token_id` already prefers config.json's + // explicit field and falls back to the spec value, so it's + // the single id used both for the engagement gate and the + // routing-fill below. + let img_tok = routing_tokens.chat_placeholder_token_id; + let bos_tok_string = routing_tokens.bos_token_string; + + match (counter.is_some(), img_tok.is_some()) { + (true, true) => tracing::info!( + target: "mm_routing", + model = %model_id, + model_dir = %model_dir.display(), + "MM-aware KV routing enabled" + ), + (counter_ok, img_ok) => { + let mut reasons: Vec = Vec::new(); + if !counter_ok { + reasons.push(format!( + "model not supported by the MM-routing registry ({})", + counter_err.as_deref().unwrap_or("unknown error") + )); + } + if !img_ok { + reasons.push( + "image-placeholder token unresolvable from \ config.json / processor_config.json / \ tokenizer_config.json / vocab probe" - .to_string(), + .to_string(), + ); + } + tracing::warn!( + target: "mm_routing", + model = %model_id, + reasons = %reasons.join("; "), + "{} is not supported for MM-aware KV routing ({}). \ + Falling back to KV routing without MM awareness — \ + text-prefix overlap still works but the router \ + cannot distinguish requests by image content.", + model_id, + reasons.join("; ") ); } - tracing::warn!( - target: "mm_routing", - model = %model_id, - reasons = %reasons.join("; "), - "{} is not supported for MM-aware KV routing ({}). \ - Falling back to KV routing without MM awareness — \ - text-prefix overlap still works but the router \ - cannot distinguish requests by image content.", - model_id, - reasons.join("; ") - ); } + (counter, img_tok, bos_tok_string) } - (counter, img_tok) - } - None => { - tracing::debug!( - target: "mm_routing", - "model directory not derivable from MDC; MM-aware routing disabled" - ); - (None, None) - } - }; + None => { + tracing::debug!( + target: "mm_routing", + "model directory not derivable from MDC; MM-aware routing disabled" + ); + (None, None, None) + } + }; #[cfg(feature = "lightseek-mm")] let image_placeholder_template = formatter.image_placeholder_template(); @@ -477,11 +507,51 @@ impl OpenAIPreprocessor { // force (both lightseek hooks resolved to `None`) — no point // building a client they'll never use. #[cfg(feature = "lightseek-mm")] - if image_token_counter.is_some() || image_token_id.is_some() { + if image_token_counter.is_some() || routing_image_token_id.is_some() { std::sync::LazyLock::force(&DIM_FETCH_MEDIA_FETCHER); std::sync::LazyLock::force(&DIM_FETCH_HTTP_CLIENT); } + // Resolve the routing-side BOS prepend for models with + // `add_bos_token: true` (see `routing_prepend_bos` doc). Only kept + // when the configured `bos_token` round-trips to a single id. The + // BOS string was harvested above by `resolve_routing_tokens` from + // the same `tokenizer_config.json` pass. + #[cfg(feature = "lightseek-mm")] + let routing_prepend_bos = match bos_token_string { + Some(bos_text) => match tokenizer.encode(&bos_text) { + Ok(enc) if enc.token_ids().len() == 1 => { + let id = enc.token_ids()[0]; + tracing::debug!( + target: "mm_routing", + bos_token = %bos_text, + bos_token_id = id, + "routing-side BOS prepend enabled (tokenizer_config.json add_bos_token=true)" + ); + Some(id) + } + Ok(enc) => { + tracing::debug!( + target: "mm_routing", + bos_token = %bos_text, + round_trip_ids = ?enc.token_ids(), + "BOS token does not round-trip to a single id; routing-side prepend disabled" + ); + None + } + Err(e) => { + tracing::debug!( + target: "mm_routing", + bos_token = %bos_text, + error = %e, + "BOS token failed to re-encode; routing-side prepend disabled" + ); + None + } + }, + None => None, + }; + Ok(Arc::new(Self { formatter, tokenizer, @@ -496,9 +566,11 @@ impl OpenAIPreprocessor { #[cfg(feature = "lightseek-mm")] image_token_counter, #[cfg(feature = "lightseek-mm")] - image_token_id, + routing_image_token_id, #[cfg(feature = "lightseek-mm")] image_placeholder_template, + #[cfg(feature = "lightseek-mm")] + routing_prepend_bos, })) } @@ -569,13 +641,13 @@ impl OpenAIPreprocessor { let tokenize_start = Instant::now(); let (token_ids, annotations) = { let _nvtx = dynamo_nvtx_range!("preprocess.tokenize"); - self.gather_tokens(request, formatted_prompt.clone(), tracker) + self.gather_tokens(request, formatted_prompt.as_deref(), tracker) .with_context(|| "Failed to gather tokens")? }; TOKENIZE_SECONDS.observe(tokenize_start.elapsed().as_secs_f64()); let _mm_image_entries = self - .gather_multi_modal_data(request, &mut builder, formatted_prompt) + .gather_multi_modal_data(request, &mut builder, formatted_prompt.as_deref()) .await .with_context(|| "Failed to gather multimodal data")?; @@ -583,8 +655,13 @@ impl OpenAIPreprocessor { // mm_hashes) for the KV router. No-op when no images are present or // the model has no resolved image-placeholder. #[cfg(feature = "lightseek-mm")] - self.gather_mm_exact_routing_info(&mut builder, &_mm_image_entries, &token_ids) - .with_context(|| "Failed to build MM routing info")?; + self.gather_mm_exact_routing_info( + &mut builder, + &_mm_image_entries, + &token_ids, + formatted_prompt.as_deref(), + ) + .with_context(|| "Failed to build MM routing info")?; // Install tokens on the builder. Done after MM routing built its // view so the routing-side borrow stays cheap and builder ownership @@ -785,7 +862,7 @@ impl OpenAIPreprocessor { &self, request: &R, builder: &mut PreprocessedRequestBuilder, - formatted_prompt: Option, + formatted_prompt: Option<&str>, ) -> Result> { let mut media_map: MultimodalDataMap = HashMap::new(); let mut fetch_tasks: Vec<(String, &ChatCompletionRequestUserMessageContentPart)> = @@ -940,8 +1017,8 @@ impl OpenAIPreprocessor { // URL-passthrough path (media_loader is None): fetch image headers in // parallel to get (W, H) per image without downloading the full bytes. - // This is what enables MM-aware routing for vLLM-backed VLMs that - // register `media_decoder: null` and let the worker do its own decode. + // Enables MM-aware routing for backends that register + // `media_decoder: null` and decode images on the worker. #[cfg(feature = "lightseek-mm")] if !url_passthrough_images.is_empty() { let dim_results = futures::future::join_all( @@ -1013,8 +1090,12 @@ impl OpenAIPreprocessor { Self::strip_inline_data_urls(&mut extra_args["messages"]); } - if let Some(ref prompt) = formatted_prompt { - extra_args["formatted_prompt"] = serde_json::Value::String(prompt.clone()); + if let Some(prompt) = formatted_prompt { + // Clone here is the single owned allocation we actually need: + // the prompt is inserted into the request's `extra_args` JSON. + // The caller still holds the original `String`; passing + // `Option<&str>` keeps text-only requests (no MM) clone-free. + extra_args["formatted_prompt"] = serde_json::Value::String(prompt.to_string()); } if let Some(serde_json::Value::Object(backend_extra_args)) = @@ -1067,32 +1148,19 @@ impl OpenAIPreprocessor { Ok(Vec::new()) } - /// Build `MmRoutingInfo` for exact MM-aware KV routing. - /// - /// Computes per-image token counts via lightseek, expands the placeholder - /// tokens, builds per-block `BlockMmObjectInfo`, and writes the result to - /// `builder.mm_routing_info`. The worker-bound `token_ids` are left - /// unchanged — only the routing-side view is expanded. - /// - /// `token_ids` is the tokenized formatted prompt (one entry per - /// placeholder per image, before expansion); the caller threads it in - /// from `gather_tokens` to avoid a second tokenizer pass. - /// - /// Returns `Ok(())` with no work performed when: - /// - no images in the request, - /// - `image_token_id` was not resolved at startup, - /// - `image_token_counter` is unavailable, - /// - `kv_cache_block_size` is 0 (worker didn't advertise one), or - /// - the count of placeholder tokens in `token_ids` doesn't match - /// `mm_image_entries.len()` (mismatched expansion would misalign - /// offsets; falling back to text-prefix routing is safer than - /// producing incorrect block hashes). + /// Build `MmRoutingInfo` for exact MM-aware KV routing. The worker-bound + /// `token_ids` are unchanged — only the routing-side view is expanded. + /// `formatted_prompt` is only consumed for Phi-3-style numbered placeholder + /// templates; single-special-token families (Qwen-VL, LLaVA) ignore it. + /// Returns `Ok(())` with no work performed on any precondition miss + /// (caller falls back to text-prefix routing). #[cfg(feature = "lightseek-mm")] pub fn gather_mm_exact_routing_info( &self, builder: &mut PreprocessedRequestBuilder, mm_image_entries: &[MmImageEntry], token_ids: &[crate::protocols::TokenIdType], + formatted_prompt: Option<&str>, ) -> Result<()> { use crate::protocols::common::preprocessor::MmRoutingInfo; use dynamo_kv_router::protocols::{RequestExtraInfo, RequestMmObjectInfo}; @@ -1100,10 +1168,10 @@ impl OpenAIPreprocessor { if mm_image_entries.is_empty() { return Ok(()); } - let Some(image_token_id) = self.image_token_id else { + let Some(find_token_id) = self.routing_image_token_id else { tracing::debug!( target: "mm_routing", - "image_token_id unresolved; skipping MM routing info" + "routing_image_token_id unresolved; skipping MM routing info" ); return Ok(()); }; @@ -1123,40 +1191,34 @@ impl OpenAIPreprocessor { return Ok(()); } - // Sanity: number of placeholder tokens in the tokenized prompt must - // match the number of images in the request. If they disagree, the - // expansion would misplace ranges; better to skip MM routing entirely - // and fall back to text-prefix routing for this request. - // - // Families like Phi-3-vision use numbered placeholder text - // (`<|image_1|>`) that BPE-decomposes into multiple sub-tokens — - // `image_token_id` (the single `<|image|>` special token) never - // appears post-tokenization. For those we run a substring-match - // pass first that rewrites each numbered placeholder's BPE - // sub-sequence back to a single `image_token_id`, then proceed - // with the standard expansion below. - let placeholder_count = token_ids.iter().filter(|&&t| t == image_token_id).count(); + // Single-special-token placeholders (Qwen-VL `<|image_pad|>`, LLaVA + // ``) emit one `find_token_id` per image in the tokenized + // prompt and hit the fast path below. Numbered-text placeholders + // (Phi-3 `<|image_N|>`) BPE-shatter and need the splice helper to + // mirror what the worker hashes on. + let placeholder_count = token_ids.iter().filter(|&&t| t == find_token_id).count(); let normalized_token_ids: std::borrow::Cow<'_, [crate::protocols::TokenIdType]> = if placeholder_count == mm_image_entries.len() { std::borrow::Cow::Borrowed(token_ids) } else if let Some(tpl) = self.image_placeholder_template && tpl.contains("{n}") + && let Some(prompt) = formatted_prompt { - match self.normalize_numbered_placeholders( - token_ids, - image_token_id, + match self.splice_phi3_numbered_placeholders_at_token_level( + prompt, tpl, + find_token_id, mm_image_entries.len(), ) { - Some(v) => std::borrow::Cow::Owned(v), + Some(spliced) => std::borrow::Cow::Owned(spliced), None => { tracing::warn!( target: "mm_routing", placeholder_count, image_count = mm_image_entries.len(), - image_token_id = image_token_id, + routing_image_token_id = find_token_id, placeholder_template = tpl, - "numbered placeholder BPE rewrite failed; \ + "splice failed for numbered placeholder; \ skipping MM routing info (text-prefix routing only)" ); return Ok(()); @@ -1167,7 +1229,7 @@ impl OpenAIPreprocessor { target: "mm_routing", placeholder_count, image_count = mm_image_entries.len(), - image_token_id = image_token_id, + routing_image_token_id = find_token_id, "placeholder token count in tokenized prompt does not match image count; \ skipping MM routing info (text-prefix routing only)" ); @@ -1181,14 +1243,32 @@ impl OpenAIPreprocessor { .collect(); let n_total: usize = n_tokens.iter().sum(); + // Replace each placeholder occurrence with N copies of the same + // `find_token_id` (config.json's `image_token_id` when present). + // That id is what the backend's HF processor emits in the expanded + // sequence for every supported family; filling with lightseek-mm's + // per-spec value previously left Qwen2-VL / Qwen2.5-VL stuck at + // `effective_cached_blocks=1` because their per-patch id never + // appears in stored block hashes. + // BOS ownership: when the Phi-3 splice helper produced + // `normalized_token_ids` (Cow::Owned), it has already emitted the + // leading BOS as part of its complete sequence. Only the non-splice + // path (Cow::Borrowed = raw tokenized prompt) needs the caller to + // prepend BOS here. Avoids the prior split-brain where the leading + // push was here and the mid-segment pushes lived inside the helper. + let splice_owns_bos = matches!(normalized_token_ids, std::borrow::Cow::Owned(_)); + let bos_extra = (!splice_owns_bos && self.routing_prepend_bos.is_some()) as usize; let mut expanded: Vec = - Vec::with_capacity(normalized_token_ids.len() + n_total); + Vec::with_capacity(normalized_token_ids.len() + n_total + bos_extra); + if !splice_owns_bos && let Some(bos) = self.routing_prepend_bos { + expanded.push(bos); + } let mut img_ranges: Vec<(usize, usize)> = Vec::with_capacity(mm_image_entries.len()); let mut i = 0usize; for &t in normalized_token_ids.iter() { - if t == image_token_id && i < mm_image_entries.len() { + if t == find_token_id && i < mm_image_entries.len() { let start = expanded.len(); - expanded.extend(std::iter::repeat_n(image_token_id, n_tokens[i])); + expanded.extend(std::iter::repeat_n(find_token_id, n_tokens[i])); img_ranges.push((start, start + n_tokens[i])); i += 1; } else { @@ -1235,48 +1315,89 @@ impl OpenAIPreprocessor { Ok(()) } - /// Rewrites BPE-decomposed numbered image placeholders back into single - /// `image_token_id` tokens so the standard expansion can proceed. + /// Build routing-side tokens for Phi-3-vision's `<|image_N|>` numbered + /// placeholder template so the router's per-block hashes match the + /// worker's BlockStored events byte-for-byte. + /// + /// **Family contract — Phi-3-only.** The encode-decode roundtrip + + /// per-segment BOS pattern here is specific to Phi-3's HF processor + /// (and any future `LlamaTokenizer`-family with `<|image_{n}|>` + /// placeholders): the processor splits the prompt at `<|image_N|>` and + /// tokenizes each segment with `add_special_tokens=true`, prefixing + /// every segment (including the first) with a fresh BOS, and decodes- + /// then-re-encodes across special-token boundaries, which inserts + /// whitespace after each special token and bumps the SentencePiece + /// prefix token (e.g. 29871 `▁` -> 259 `▁▁` for `<|user|>\n`). Both + /// effects are reproduced here. /// - /// Used for Phi-3-vision-style templates whose flatten-time placeholder - /// is `<|image_{n}|>` (not a tokenizer special token, BPE-encodes into - /// ~7 sub-tokens) while the model's actual image token is `<|image|>` - /// (single special token = `image_token_id`). The backend's HF - /// processor recognises `<|image_{n}|>` in the prompt and replaces - /// each with N copies of `image_token_id` post-tokenization — we - /// replicate the routing-side equivalent here. + /// Caller is `gather_mm_exact_routing_info`, which gates on + /// `placeholder_tpl.contains("{n}")`. The `routing_prepend_bos` + /// requirement below acts as the family guard: a future model with a + /// `{n}`-style placeholder but no BOS prepend (i.e. not a + /// LlamaTokenizer-family) would not benefit from this roundtrip and + /// could silently produce wrong block hashes, so we bail and let the + /// caller fall back to text-prefix routing. /// - /// For each image index `i` in `1..=expected_count`, encodes the - /// substituted placeholder string and scans `token_ids` for the - /// resulting BPE sub-sequence. Each match collapses to a single - /// `image_token_id` in the returned vector, preserving every - /// surrounding token. Returns `None` if any expected placeholder is - /// missing or if scans go out of order — the caller falls back to - /// text-prefix routing in that case. + /// Returns one `find_token_id` per image; the caller's expansion loop + /// multiplies them to the per-image patch count. Leading BOS is + /// emitted here as the first element of the returned vector — the + /// caller does NOT prepend its own BOS when this helper succeeds. + /// Returns `None` (caller falls back to text-prefix routing) when the + /// family guard trips or on any tokenize/decode failure. #[cfg(feature = "lightseek-mm")] - fn normalize_numbered_placeholders( + fn splice_phi3_numbered_placeholders_at_token_level( &self, - token_ids: &[crate::protocols::TokenIdType], - image_token_id: crate::protocols::TokenIdType, + prompt: &str, placeholder_tpl: &str, + find_token_id: crate::protocols::TokenIdType, expected_count: usize, ) -> Option> { - let mut out: Vec = Vec::with_capacity(token_ids.len()); - let mut cursor = 0usize; + // Family guard: the encode-decode roundtrip + per-segment BOS + // semantics are LlamaTokenizer-family-specific. `routing_prepend_bos` + // is set iff the model declares `add_bos_token: true` in + // tokenizer_config.json, which is the marker for that family. + let bos = self.routing_prepend_bos?; + + // Encode-decode roundtrip mirrors vLLM's text-substitute fallback, + // which is what the Phi-3 worker actually hashes on. + let prompt_owned: String = self + .tokenizer + .encode(prompt) + .ok() + .and_then(|enc| self.tokenizer.decode(enc.token_ids(), false).ok()) + .map(Into::into)?; + let prompt: &str = &prompt_owned; + + let mut byte_ranges: Vec<(usize, usize)> = Vec::with_capacity(expected_count); + let mut search_from = 0usize; for idx in 1..=expected_count { - let placeholder_text = placeholder_tpl.replace("{n}", &idx.to_string()); - let encoding = self.tokenizer.encode(&placeholder_text).ok()?; - let sub_ids = encoding.token_ids(); - if sub_ids.is_empty() { - return None; - } - let pos = find_subseq(&token_ids[cursor..], sub_ids)? + cursor; - out.extend_from_slice(&token_ids[cursor..pos]); - out.push(image_token_id); - cursor = pos + sub_ids.len(); + let pattern = placeholder_tpl.replace("{n}", &idx.to_string()); + let rel = prompt[search_from..].find(&pattern)?; + let start = search_from + rel; + let end = start + pattern.len(); + byte_ranges.push((start, end)); + search_from = end; + } + + // Leading BOS is emitted here so the caller's general-purpose + // BOS-prepend path can skip when this helper produced the + // normalized tokens (avoids the prior split-brain where leading + // BOS was pushed by the caller and mid/suffix BOS pushed here). + let mut result: Vec = vec![bos]; + let mut prev_end = 0usize; + for &(ph_start, ph_end) in byte_ranges.iter() { + let seg = &prompt[prev_end..ph_start]; + let seg_enc = self.tokenizer.encode(seg).ok()?; + result.extend_from_slice(seg_enc.token_ids()); + result.push(find_token_id); + // Mid-prompt segments get a fresh BOS — Phi-3's HF processor + // tokenizes each segment with add_special_tokens=true. + result.push(bos); + prev_end = ph_end; } - out.extend_from_slice(&token_ids[cursor..]); - Some(out) + let suffix_enc = self.tokenizer.encode(&prompt[prev_end..]).ok()?; + result.extend_from_slice(suffix_enc.token_ids()); + Some(result) } /// xxh3-64 of the raw URL bytes. Used as the routing `mm_hash` in the @@ -1455,7 +1576,7 @@ impl OpenAIPreprocessor { >( &self, request: &R, - formatted_prompt: Option, + formatted_prompt: Option<&str>, tracker: Option<&RequestTracker>, ) -> Result<(Vec, HashMap)> { let mut annotations = HashMap::new(); @@ -1488,15 +1609,17 @@ impl OpenAIPreprocessor { if let Some(text_input) = request.extract_text() { match text_input { TextInput::Single(raw_prompt) => { - if let Some(f) = formatted_prompt.as_ref() + if let Some(f) = formatted_prompt && request.has_annotation(ANNOTATION_FORMATTED_PROMPT) { annotations .insert(ANNOTATION_FORMATTED_PROMPT.to_string(), f.to_string()); } - // Completions will use raw_prompt, no template - let prompt = formatted_prompt.unwrap_or(raw_prompt); + // Completions will use raw_prompt, no template. + // Borrow either input — no allocation needed; the + // tokenizer accepts `&str`. + let prompt: &str = formatted_prompt.unwrap_or(raw_prompt.as_str()); // If nvext.token_data is present, use the pre-computed tokens // directly and skip tokenization. This avoids redundant @@ -1526,10 +1649,10 @@ impl OpenAIPreprocessor { tracing::warn!( "backend_instance_id provided but no token_data; tokenizing prompt" ); - let encoding = self.encode_with_timing(&prompt, tracker)?; + let encoding = self.encode_with_timing(prompt, tracker)?; (encoding.token_ids().to_vec(), false) } else { - let encoding = self.encode_with_timing(&prompt, tracker)?; + let encoding = self.encode_with_timing(prompt, tracker)?; (encoding.token_ids().to_vec(), false) }; diff --git a/lib/llm/src/preprocessor/lightseek_mm.rs b/lib/llm/src/preprocessor/lightseek_mm.rs index 50ec7fbf205d..9d1a47d84c4a 100644 --- a/lib/llm/src/preprocessor/lightseek_mm.rs +++ b/lib/llm/src/preprocessor/lightseek_mm.rs @@ -146,7 +146,20 @@ impl LightseekMmCounter { /// - `tokenizer.json` or `config.json` is missing or unparseable, or /// - no `ModelProcessorSpec` matches the model (caller should fall back to /// text-prefix routing). +/// +/// Standalone wrapper around [`resolve_image_token_id_with_config`]. Prefer +/// [`resolve_routing_tokens`] when also fetching the chat-template placeholder +/// or BOS token (one config-parse pass instead of two). pub fn resolve_image_token_id(model_id: &str, model_dir: &Path) -> Option { + let config = read_json(model_dir, "config.json")?; + resolve_image_token_id_with_config(model_id, model_dir, &config) +} + +fn resolve_image_token_id_with_config( + model_id: &str, + model_dir: &Path, + config: &serde_json::Value, +) -> Option { // Try the HuggingFace fast tokenizer first; fall back to a no-op // tokenizer when `tokenizer.json` is missing (Kimi-K2.5 ships only // `tiktoken.model`, for example). Specs that read the placeholder @@ -174,34 +187,10 @@ pub fn resolve_image_token_id(model_id: &str, model_dir: &Path) -> Option &null_tokenizer, }; - let config_path = model_dir.join("config.json"); - let config_json = std::fs::read_to_string(&config_path) - .map_err(|e| { - tracing::warn!( - target: "mm_routing", - config = %config_path.display(), - err = %e, - "lightseek: failed to read config.json" - ); - e - }) - .ok()?; - let config: serde_json::Value = serde_json::from_str(&config_json) - .map_err(|e| { - tracing::warn!( - target: "mm_routing", - config = %config_path.display(), - err = %e, - "lightseek: failed to parse config.json" - ); - e - }) - .ok()?; - let metadata = ModelMetadata { model_id, tokenizer, - config: &config, + config, }; let spec = MODEL_REGISTRY.lookup(&metadata)?; @@ -227,6 +216,122 @@ pub fn resolve_image_token_id(model_id: &str, model_dir: &Path) -> Option, + /// Token id the chat template emits per image. Read from `config.json`'s + /// literal `image_token_id` field, falling back to `image_token_id` + /// above. Equals `image_token_id` for most VLMs; Qwen2-VL / Qwen2.5-VL + /// emit `<|image_pad|>` here while the per-patch id is `<|vision_pad|>`. + pub chat_placeholder_token_id: Option, + /// `bos_token` string from `tokenizer_config.json` when + /// `add_bos_token: true`. Caller encodes via its model tokenizer to + /// produce the routing-side prepend id. `None` for models that don't + /// prepend BOS. + pub bos_token_string: Option, +} + +/// Resolve all routing-side token info from a model directory in a single +/// pass. Reads `config.json` once for the per-spec image id + chat-template +/// placeholder, and `tokenizer_config.json` once for BOS. Replaces the +/// in-`preprocessor.rs` `read_image_token_id_from_config` / +/// `read_bos_token_from_config` helpers so config parsing lives next to +/// the rest of the MM-routing token resolution. +pub fn resolve_routing_tokens(model_id: &str, model_dir: &Path) -> RoutingTokens { + let config = read_json(model_dir, "config.json"); + let tokenizer_config = read_json(model_dir, "tokenizer_config.json"); + + let image_token_id = config + .as_ref() + .and_then(|c| resolve_image_token_id_with_config(model_id, model_dir, c)); + let chat_placeholder_token_id = config + .as_ref() + .and_then(extract_chat_placeholder_from_config) + .or(image_token_id); + let bos_token_string = tokenizer_config + .as_ref() + .and_then(extract_bos_token_from_tokenizer_config); + + RoutingTokens { + image_token_id, + chat_placeholder_token_id, + bos_token_string, + } +} + +/// Read + parse a JSON file under `model_dir`. Warns on read or parse +/// failure (missing files are silent — many models legitimately lack +/// `tokenizer_config.json`). Returns `None` on any error. +fn read_json(model_dir: &Path, filename: &str) -> Option { + let path = model_dir.join(filename); + let raw = match std::fs::read_to_string(&path) { + Ok(s) => s, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => return None, + Err(e) => { + tracing::warn!( + target: "mm_routing", + path = %path.display(), + err = %e, + "lightseek: failed to read {filename}" + ); + return None; + } + }; + match serde_json::from_str(&raw) { + Ok(v) => Some(v), + Err(e) => { + tracing::warn!( + target: "mm_routing", + path = %path.display(), + err = %e, + "lightseek: failed to parse {filename}" + ); + None + } + } +} + +/// Read the literal `image_token_id` field from a pre-parsed `config.json`. +/// Used by Qwen2-VL / Qwen2.5-VL where the chat-template-emitted placeholder +/// differs from the per-patch expansion token returned by the spec. +fn extract_chat_placeholder_from_config(config: &serde_json::Value) -> Option { + config + .get("image_token_id") + .and_then(|x| x.as_u64()) + .and_then(|id| u32::try_from(id).ok()) +} + +/// Return the `bos_token` string from a pre-parsed `tokenizer_config.json` +/// when `add_bos_token: true`. The routing-side sequence must prepend it to +/// match the backend's HF-processor output (Phi-3-vision and other +/// `LlamaTokenizer`-family models). Returns `None` otherwise. +fn extract_bos_token_from_tokenizer_config(cfg: &serde_json::Value) -> Option { + if !cfg + .get("add_bos_token") + .and_then(|x| x.as_bool()) + .unwrap_or(false) + { + return None; + } + // `bos_token` is usually a plain string ("") but the HF schema also + // allows it to be an `AddedToken` dict — handle both. + cfg.get("bos_token").and_then(|x| match x { + serde_json::Value::String(s) => Some(s.clone()), + serde_json::Value::Object(o) => o + .get("content") + .and_then(|c| c.as_str()) + .map(|s| s.to_owned()), + _ => None, + }) +} + #[cfg(test)] mod tests { //! Contract tests against the upstream lightseek registry. Pin the diff --git a/lib/tokenizers/src/hf.rs b/lib/tokenizers/src/hf.rs index 080a775719fe..585944df0249 100644 --- a/lib/tokenizers/src/hf.rs +++ b/lib/tokenizers/src/hf.rs @@ -1,7 +1,9 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use tokenizers::tokenizer::Tokenizer as HfTokenizer; +use std::path::Path; + +use tokenizers::tokenizer::{AddedToken, Tokenizer as HfTokenizer}; use super::{ Encoding, Error, Result, TokenIdType, @@ -13,16 +15,120 @@ pub struct HuggingFaceTokenizer { } impl HuggingFaceTokenizer { + /// Load from `tokenizer.json`, merging in special tokens declared only in + /// a sibling `tokenizer_config.json`'s `added_tokens_decoder`. Without + /// this, some releases (Qwen2-VL-2B's `<|image_pad|>`) BPE-shatter and + /// silently break MM-aware routing. The merge is idempotent. pub fn from_file(model_name: &str) -> Result { - let tokenizer = HfTokenizer::from_file(model_name) + let mut tokenizer = HfTokenizer::from_file(model_name) .map_err(|err| Error::msg(format!("Error loading tokenizer: {}", err)))?; + if let Some(parent) = Path::new(model_name).parent() { + merge_special_tokens_from_config(&mut tokenizer, parent); + } + Ok(HuggingFaceTokenizer { tokenizer }) } pub fn from_tokenizer(tokenizer: HfTokenizer) -> Self { HuggingFaceTokenizer { tokenizer } } + + /// Wrap an already-loaded `HfTokenizer` and merge in the sibling + /// `tokenizer_config.json` special tokens; see [`Self::from_file`]. + pub fn from_tokenizer_with_model_dir(tokenizer: HfTokenizer, model_dir: &Path) -> Self { + let mut tokenizer = tokenizer; + merge_special_tokens_from_config(&mut tokenizer, model_dir); + HuggingFaceTokenizer { tokenizer } + } +} + +/// Promote `tokenizer_config.json`'s `special: true` `added_tokens_decoder` +/// entries onto `tokenizer`. Missing-file / parse errors are swallowed since +/// the file is optional. See [`HuggingFaceTokenizer::from_file`]. +/// +/// `pub` so downstream crates (e.g. `dynamo-llm`'s model_card) can apply +/// the same promotion before extracting the special-token boundary list +/// for the L1 prefix cache — otherwise the cache and the wrapped +/// tokenizer would disagree on which strings are atomic specials. +pub fn merge_special_tokens_from_config(tokenizer: &mut HfTokenizer, model_dir: &Path) { + let cfg_path = model_dir.join("tokenizer_config.json"); + let Ok(raw) = std::fs::read_to_string(&cfg_path) else { + return; + }; + let cfg: serde_json::Value = match serde_json::from_str(&raw) { + Ok(v) => v, + Err(e) => { + tracing::debug!( + target: "tokenizer", + path = %cfg_path.display(), + error = %e, + "tokenizer_config.json parse failed; skipping special-token merge" + ); + return; + } + }; + let Some(decoder) = cfg.get("added_tokens_decoder").and_then(|v| v.as_object()) else { + return; + }; + + let mut to_add: Vec = Vec::new(); + for (_id, spec) in decoder { + let obj = match spec.as_object() { + Some(o) => o, + None => continue, + }; + // The id is informational — `add_special_tokens` reuses the existing + // vocab id via `Model::token_to_id` on the content string. + if obj.get("special").and_then(|v| v.as_bool()) != Some(true) { + continue; + } + let Some(content) = obj.get("content").and_then(|v| v.as_str()) else { + continue; + }; + if content.is_empty() { + continue; + } + let single_word = obj + .get("single_word") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + let lstrip = obj.get("lstrip").and_then(|v| v.as_bool()).unwrap_or(false); + let rstrip = obj.get("rstrip").and_then(|v| v.as_bool()).unwrap_or(false); + let normalized = obj + .get("normalized") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + let token = AddedToken::from(content.to_string(), true) + .single_word(single_word) + .lstrip(lstrip) + .rstrip(rstrip) + .normalized(normalized); + to_add.push(token); + } + + if to_add.is_empty() { + return; + } + // Dedups against existing added-tokens, so this is a no-op when + // tokenizer.json already had them. Return value = net-new count. + let added = tokenizer.add_special_tokens(&to_add); + if added > 0 { + // Warn (not debug) when the merge actually promotes anything — + // intentionally loud so accidental promotion of a previously- + // non-special token shows up immediately in worker logs. Lists + // the literal token strings so debugging doesn't require a + // second pass through `added_tokens_decoder`. + let promoted: Vec<&str> = to_add.iter().map(|t| t.content.as_str()).collect(); + tracing::warn!( + target: "tokenizer", + path = %cfg_path.display(), + added, + candidates = to_add.len(), + promoted = ?promoted, + "merged additional special tokens from tokenizer_config.json" + ); + } } impl Encoder for HuggingFaceTokenizer { @@ -70,3 +176,104 @@ impl From for HuggingFaceTokenizer { HuggingFaceTokenizer { tokenizer } } } + +#[cfg(test)] +mod tests { + //! The existing tool-calling / reasoning parser tests inject already- + //! decoded text and never run `tokenizer.decode`, so they cannot + //! catch the class of bug where a non-special marker gets accidentally + //! promoted to "special" and silently disappears under + //! `skip_special_tokens=True`. One unit test pins both halves of the + //! gate (promote special:true / skip special:false) end-to-end through + //! the actual encode -> decode round trip — the layer above which the + //! parser tests cannot reach. + use super::*; + use std::fs; + use tempfile::TempDir; + + #[test] + fn merge_gate_round_trips_through_decode() { + // Minimal WordLevel `tokenizer.json`. `<|special_kept|>` and + // `<|special_dropped|>` are deliberately NOT pre-declared as + // special in `added_tokens` here — that's exactly the shape + // `merge_special_tokens_from_config` is supposed to fix from + // tokenizer_config.json. + const TOKENIZER_JSON: &str = r#"{ + "version": "1.0", + "truncation": null, + "padding": null, + "added_tokens": [ + {"id": 0, "content": "", "special": true, "single_word": false, "lstrip": false, "rstrip": false, "normalized": false} + ], + "normalizer": null, + "pre_tokenizer": null, + "post_processor": null, + "decoder": null, + "model": { + "type": "WordLevel", + "vocab": {"": 0, "hello": 1, "world": 2, "<|special_kept|>": 3, "<|special_dropped|>": 4}, + "unk_token": "" + } + }"#; + + // `<|special_kept|>` is `special: true` → must be promoted, and + // therefore stripped under skip_special_tokens=true. + // `<|special_dropped|>` is `special: false` → must be skipped, + // and therefore survive skip_special_tokens=true. The latter is + // the Ryan/Keiven concern: tool-call / reasoning markers are + // universally declared `special: false` precisely so the parser + // still sees them; a regression here would silently break + // parsing without any of the parser unit tests turning red. + const TOKENIZER_CONFIG_JSON: &str = r#"{ + "added_tokens_decoder": { + "3": {"content": "<|special_kept|>", "special": true, "single_word": false, "lstrip": false, "rstrip": false, "normalized": false}, + "4": {"content": "<|special_dropped|>", "special": false, "single_word": false, "lstrip": false, "rstrip": false, "normalized": false} + } + }"#; + + let dir = TempDir::new().unwrap(); + fs::write(dir.path().join("tokenizer.json"), TOKENIZER_JSON).unwrap(); + fs::write( + dir.path().join("tokenizer_config.json"), + TOKENIZER_CONFIG_JSON, + ) + .unwrap(); + + let mut tokenizer = HfTokenizer::from_file(dir.path().join("tokenizer.json")).unwrap(); + merge_special_tokens_from_config(&mut tokenizer, dir.path()); + + // Registry assertion: only the special:true entry was promoted. + let specials: Vec = { + let mut v: Vec = tokenizer + .get_added_tokens_decoder() + .values() + .filter(|t| t.special) + .map(|t| t.content.clone()) + .collect(); + v.sort(); + v + }; + assert_eq!( + specials, + vec!["".to_string(), "<|special_kept|>".to_string()], + "<|special_kept|> promoted; <|special_dropped|> stayed non-special" + ); + + // Decode round-trip assertion (the layer parser tests cannot + // reach): the registry mutation actually changes downstream + // decode behavior in the expected direction. + let enc_kept = tokenizer.encode("<|special_kept|>", false).unwrap(); + let decoded_strip = tokenizer.decode(enc_kept.get_ids(), true).unwrap(); + assert!( + !decoded_strip.contains("<|special_kept|>"), + "promoted special:true token must be stripped under skip_special_tokens=true; got {decoded_strip:?}" + ); + + let enc_drop = tokenizer.encode("<|special_dropped|>", false).unwrap(); + let decoded_keep = tokenizer.decode(enc_drop.get_ids(), true).unwrap(); + assert!( + decoded_keep.contains("<|special_dropped|>"), + "non-promoted special:false token must survive skip_special_tokens=true; got {decoded_keep:?}" + ); + } +} diff --git a/tests/mm_router/test_router_rust_mm_router_e2e.py b/tests/mm_router/test_router_rust_mm_router_e2e.py index 74636b1a7ee9..f9af7fd4f6d6 100644 --- a/tests/mm_router/test_router_rust_mm_router_e2e.py +++ b/tests/mm_router/test_router_rust_mm_router_e2e.py @@ -532,8 +532,8 @@ def test_router_rust_mm_logs_lightseek_initialization( ) log_text = router_proc.read_logs() assert ( - "MM-aware KV routing enabled (lightseek)" in log_text - ), "frontend should emit lightseek init log line on model registration" + "MM-aware KV routing enabled" in log_text + ), "frontend should emit MM-routing init log line on model registration" assert ( "resolved image-placeholder token id" in log_text ), "image_token resolver should log which tier produced the hit" diff --git a/tests/serve/multimodal_profiles/vllm.py b/tests/serve/multimodal_profiles/vllm.py index 92b29e149932..eec4e1edbb71 100644 --- a/tests/serve/multimodal_profiles/vllm.py +++ b/tests/serve/multimodal_profiles/vllm.py @@ -109,7 +109,15 @@ profiled_vram_gib=13.0, requested_vllm_kv_cache_bytes=536_870_912, env={"SINGLE_GPU": "true"}, - tests=[MmCase(payload=make_image_payload_cached_tokens(["green"]))], + tests=[ + MmCase( + payload=make_image_payload_cached_tokens( + ["green"], + require_rust_processor_init=True, + min_avg_kv_hit_rate=0.9, + ) + ) + ], ), # The chat-processor variant of the MM-aware router: same routing # architecture, but the frontend uses --dyn-chat-processor=vllm @@ -200,7 +208,22 @@ profiled_vram_gib=19.0, requested_vllm_kv_cache_bytes=1_719_075_000, env={"SINGLE_GPU": "true"}, - tests=[MmCase(payload=make_image_payload(["green"]))], + # Qwen2-VL / Qwen2.5-VL: chat template emits `<|image_pad|>` + # (151655) and vLLM's HF processor expands the same id N + # times in the prompt sequence — routing-side fills with + # this id so block hashes align with what the worker + # stores. (lightseek's per-spec id is `<|vision_pad|>` + # 151654; the routing path now uses config.json's + # `image_token_id` instead, see preprocessor.rs splice.) + tests=[ + MmCase( + payload=make_image_payload_cached_tokens( + ["green"], + require_rust_processor_init=True, + min_avg_kv_hit_rate=0.9, + ) + ) + ], ), }, ), @@ -214,7 +237,16 @@ profiled_vram_gib=16.0, requested_vllm_kv_cache_bytes=1_719_075_000, env={"SINGLE_GPU": "true"}, - tests=[MmCase(payload=make_image_payload(["green"]))], + # Dual-token routing path — see qwen2.5-vl-3b above. + tests=[ + MmCase( + payload=make_image_payload_cached_tokens( + ["green"], + require_rust_processor_init=True, + min_avg_kv_hit_rate=0.9, + ) + ) + ], ), }, ), @@ -237,7 +269,15 @@ # engaged (2nd identical request hits the warm worker's KV # cache); a silent regression to text-prefix-only routing # would still return "green" but 0 cached tokens. - tests=[MmCase(payload=make_image_payload_cached_tokens(["green"]))], + tests=[ + MmCase( + payload=make_image_payload_cached_tokens( + ["green"], + require_rust_processor_init=True, + min_avg_kv_hit_rate=0.9, + ) + ) + ], ), }, # Phi-3-vision uses --trust-remote-code for its custom processor. @@ -360,7 +400,15 @@ requested_vllm_kv_cache_bytes=4_318_854_000, # cached_tokens-asserting payload proves MM-aware routing # engaged for LLaVA-1.5 (placeholder-template `` path). - tests=[MmCase(payload=make_image_payload_cached_tokens(["green"]))], + tests=[ + MmCase( + payload=make_image_payload_cached_tokens( + ["green"], + require_rust_processor_init=True, + min_avg_kv_hit_rate=0.9, + ) + ) + ], ), "agg": TopologyConfig( # nightly-only: 7B 1-GPU footprint is tight (vram=19.2 GiB). @@ -478,7 +526,15 @@ requested_vllm_kv_cache_bytes=4_318_854_000, # cached_tokens-asserting payload proves MM-aware routing # engaged for LLaVA-NeXT (anyres multi-crop processor). - tests=[MmCase(payload=make_image_payload_cached_tokens(["green"]))], + tests=[ + MmCase( + payload=make_image_payload_cached_tokens( + ["green"], + require_rust_processor_init=True, + min_avg_kv_hit_rate=0.9, + ) + ) + ], ), }, ), diff --git a/tests/utils/multimodal.py b/tests/utils/multimodal.py index dcdd8676f4be..ba15bc135b33 100644 --- a/tests/utils/multimodal.py +++ b/tests/utils/multimodal.py @@ -66,21 +66,18 @@ def make_image_payload( def make_image_payload_cached_tokens( expected_response: list[str], *, - repeat_count: int = 2, + repeat_count: int = 3, min_cached_tokens: int = 1, + require_rust_processor_init: bool = False, + require_vllm_mm_processor_init: bool = False, + min_avg_kv_hit_rate: float = 0.0, ) -> CachedTokensChatPayload: - """Image payload that also asserts MM-aware KV cache reuse on repeats. - - Same body shape as :func:`make_image_payload`, but wrapped in a - :class:`CachedTokensChatPayload` so the 2nd+ request validates that - ``usage.prompt_tokens_details.cached_tokens >= min_cached_tokens``. - Two identical MM requests through an MM-routing-aware frontend must - land on the same warm worker and reuse the prefix cache; if routing - silently regresses to text-only the second request will report 0 - cached tokens and this payload fails. - - Used to harden the ``agg_router`` pre_merge smoke against silent - regressions in the Rust+lightseek routing path. + """Image payload that asserts MM-aware KV cache reuse on repeats. + + ``require_rust_processor_init`` / ``require_vllm_mm_processor_init`` assert + the MM-routing init log fired. ``min_avg_kv_hit_rate`` asserts the + post-R1 mean of router_kv_hit_rate >= threshold (fails closed when + router-side hashes diverge from the worker). """ return CachedTokensChatPayload( body={ @@ -103,6 +100,9 @@ def make_image_payload_cached_tokens( repeat_count=repeat_count, expected_response=expected_response, min_cached_tokens=min_cached_tokens, + require_rust_processor_init=require_rust_processor_init, + require_vllm_mm_processor_init=require_vllm_mm_processor_init, + min_avg_kv_hit_rate=min_avg_kv_hit_rate, ) diff --git a/tests/utils/payloads.py b/tests/utils/payloads.py index 2b20fa3d2ab8..6204dfbf3651 100644 --- a/tests/utils/payloads.py +++ b/tests/utils/payloads.py @@ -405,18 +405,38 @@ def __init__( timeout: int = 60, min_cached_tokens: int = 1, router_nvext_expectation: RouterNvextExpectation | None = None, + require_rust_processor_init: bool = False, + require_vllm_mm_processor_init: bool = False, + min_avg_kv_hit_rate: float = 0.0, ): + log_patterns: List[str] = list(expected_log or []) + if require_rust_processor_init: + log_patterns.append(r"MM-aware KV routing enabled") + if require_vllm_mm_processor_init: + log_patterns.append(r"\[mm-routing\] Transfer mode:") super().__init__( body=body, repeat_count=repeat_count, expected_response=expected_response or [], - expected_log=expected_log or [], + expected_log=log_patterns, timeout=timeout, ) self.min_cached_tokens = min_cached_tokens self._request_count = 0 self._cached_tokens_found = False self.router_nvext_expectation = router_nvext_expectation + # Asserts the post-R1 mean of router_kv_hit_rate >= threshold. Catches + # router/worker hash divergence (overlap=0) that cached_tokens alone + # can miss via load-balance luck on vLLM's per-worker prefix cache. + # + # TODO(mm-routing): this field + _metrics_baseline + the kv_hit_rate + # delta assertion in final_validation() are router-metric-specific + # logic accreting onto a general-purpose Cached*Payload base. If + # more strong-gate metrics get added (decode-imbalance, routing- + # block-count, etc.), move into a RouterMetricsAssertion + # mixin/subclass. + self.min_avg_kv_hit_rate = min_avg_kv_hit_rate + self._metrics_baseline: Optional[tuple[float, float]] = None def validate(self, response: Any, content: str) -> None: """Validate response and check for cached tokens on repeated requests.""" @@ -457,10 +477,47 @@ def validate(self, response: Any, content: str) -> None: f"(expected >= {self.min_cached_tokens})" ) - def final_validation(self) -> None: - """Called after all requests are processed to ensure we saw cached tokens. + # Snapshot after R1 so the delta in final_validation isolates R2+. + if ( + self._metrics_baseline is None + and self._request_count == 1 + and self.min_avg_kv_hit_rate > 0 + ): + self._metrics_baseline = self._scrape_router_kv_hit_rate() - Raises AssertionError if cached tokens were not found on any repeated request. + def _scrape_router_kv_hit_rate(self) -> Optional[tuple[float, float]]: + """Return ``(sum, count)`` for ``router_kv_hit_rate`` from the + frontend /metrics endpoint, or ``None`` if the endpoint is + unreachable. The component MetricsHierarchy auto-prepends + ``dynamo_component_`` to the exported name. + """ + url = f"http://localhost:{self.port}/metrics" + try: + text = requests.get(url, timeout=5).text + except requests.RequestException as e: + # Narrow to HTTP/network errors per .ai/python-guidelines.md: + # we expect transient endpoint flakes here (timeout, connection + # refused while the frontend is still binding /metrics) and + # the strong gate has its own `is None` guard. Programming + # errors propagate so they surface at test-time instead of + # being swallowed. + logger.warning("Failed to scrape %s: %s", url, e) + return None + # Compose from canonical constants so a metric rename in + # prometheus_names cascades here instead of silently breaking + # the kv_hit_rate strong gate. + full = ( + f"{prometheus_names.name_prefix.COMPONENT}_" + f"{prometheus_names.router.KV_HIT_RATE}" + ) + return ( + sum_metric_samples(text, f"{full}_sum"), + sum_metric_samples(text, f"{full}_count"), + ) + + def final_validation(self) -> None: + """Assert cached_tokens >= min_cached_tokens on at least one repeat, + and (if set) router_kv_hit_rate post-R1 mean >= min_avg_kv_hit_rate. """ if self.repeat_count > 1 and not self._cached_tokens_found: raise AssertionError( @@ -473,6 +530,41 @@ def final_validation(self) -> None: "✓ Final validation PASSED: cached_tokens found in repeated requests" ) + if self.min_avg_kv_hit_rate <= 0: + return + if self._metrics_baseline is None: + raise AssertionError( + "min_avg_kv_hit_rate set but no metrics baseline captured " + "(R1 validate() didn't run or /metrics was unreachable)." + ) + after = self._scrape_router_kv_hit_rate() + if after is None: + raise AssertionError( + "router_kv_hit_rate scrape failed at final_validation; " + "/metrics endpoint unreachable from test." + ) + bsum, bcount = self._metrics_baseline + asum, acount = after + d_sum, d_count = asum - bsum, acount - bcount + if d_count <= 0: + raise AssertionError( + f"router_kv_hit_rate: no new observations between R1 and final " + f"(baseline_count={bcount}, after_count={acount}); " + f"MM-routing likely not engaging on repeat requests." + ) + avg = d_sum / d_count + if avg < self.min_avg_kv_hit_rate: + raise AssertionError( + f"router_kv_hit_rate: mean over R2+ ({avg:.3f}) below required " + f"min ({self.min_avg_kv_hit_rate}). delta_n={d_count}, " + f"delta_sum={d_sum:.3f}. Router-side block hashes did not " + f"match the worker — MM-aware routing degraded silently." + ) + logger.info( + f"✓ router_kv_hit_rate: mean over R2+ = {avg:.3f} " + f"(>= {self.min_avg_kv_hit_rate})" + ) + @dataclass class LoraTestChatPayload(ChatPayload):