fix: MM-aware KV routing for Phi-3, Qwen2-VL, Qwen2.5-VL

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:-}"

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::<serde_json::Error>()
                             && let Ok(contents) = std::fs::read_to_string(p)
@@ -979,14 +979,27 @@ 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<String> = if cache_enabled {
                     extract_hf_special_tokens(&hf)
                 } else {
                     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<dyn crate::tokenizers::traits::Tokenizer> = if use_fast {
                     if let Some(path_str) = p.to_str() {
@@ -1000,20 +1013,18 @@ 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 {
                         tracing::warn!(
                             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 {

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<TokenIdType> {
+    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<TokenIdType> {
     // 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<TokenI
         None => &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<TokenI
     Some(id as TokenIdType)
 }
 
+/// Bundle of routing-side token info resolved from a model's HF JSON
+/// configs. All fields default to `None` when the corresponding lookup
+/// fails — callers disable the respective routing path without erroring.
+///
+/// Built by [`resolve_routing_tokens`]; reads `config.json` and
+/// `tokenizer_config.json` at most once each.
+pub struct RoutingTokens {
+    /// Image-placeholder token id resolved via `ModelProcessorSpec`
+    /// (per-family `config.json` field). `None` disables MM-aware routing.
+    pub image_token_id: Option<TokenIdType>,
+    /// 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<TokenIdType>,
+    /// `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<String>,
+}
+
+/// 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<serde_json::Value> {
+    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<TokenIdType> {
+    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<String> {
+    if !cfg
+        .get("add_bos_token")
+        .and_then(|x| x.as_bool())
+        .unwrap_or(false)
+    {
+        return None;
+    }
+    // `bos_token` is usually a plain string ("<s>") 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