diff --git a/experimental/sgl-router/Cargo.toml b/experimental/sgl-router/Cargo.toml index de419ab8b..80416b29d 100644 --- a/experimental/sgl-router/Cargo.toml +++ b/experimental/sgl-router/Cargo.toml @@ -43,6 +43,16 @@ serde_json = { version = "1", features = ["preserve_order"] } # this pulls no openssl/native-tls (matching reqwest's rustls-tls above). hf-hub = { version = "0.4", default-features = false, features = ["ureq"] } +# Chat-template rendering for cache-aware routing: the engine caches tokens +# AFTER applying the model's chat template, so the router renders the same +# template (from tokenizer_config.json) before hashing — otherwise its query +# token_ids diverge from the engine's stored blocks. `pycompat` supplies the +# Python str/dict methods HF chat templates rely on (.startswith, .items, ...). +minijinja = { version = "2", features = ["loop_controls", "json"] } +minijinja-contrib = { version = "2", features = ["pycompat"] } +# `strftime_now` chat-template helper (some templates inject the current date). +chrono = { version = "0.4", default-features = false, features = ["clock"] } + # Utilities anyhow = "1" thiserror = "2" diff --git a/experimental/sgl-router/src/policies/cache_aware_zmq.rs b/experimental/sgl-router/src/policies/cache_aware_zmq.rs index e9e38ea01..e23f9fa9c 100644 --- a/experimental/sgl-router/src/policies/cache_aware_zmq.rs +++ b/experimental/sgl-router/src/policies/cache_aware_zmq.rs @@ -17,10 +17,12 @@ //! balance_rel_threshold`, skip the cache lookup and pick the //! lowest-load worker. This prevents one hot worker from dominating //! cache-aware selection while every other worker idles. -//! 2. **Tokenize.** Pull the prompt text out of the JSON body (`messages` or -//! `prompt` field), run it through the per-model tokenizer. On any -//! failure (no body, no tokenizer, encode error, empty tokens), fall -//! through to step 4 (min-load fallback). +//! 2. **Tokenize.** For chat requests (`messages`) on a model with a chat +//! encoder (a Jinja template, or a built-in encoder like DeepSeek-V4's), +//! render it and tokenize the result so the query tokens match what the +//! engine cached (BOS + role markers + content); otherwise tokenize the raw +//! `prompt`/`text`. On any failure (no body, no tokenizer, encode error, +//! empty tokens), fall through to step 4 (min-load fallback). //! 3. **Hash + match.** Compute block hashes via //! [`super::kv_events::compute_block_hashes`], query the shared hash tree //! for the longest matching prefix. If `match_rate > cache_threshold`, @@ -130,9 +132,20 @@ impl CacheAwareZmqPolicy { abs_diff > self.config.balance_abs_threshold && max_load > rel_threshold } - /// Extract a prompt-text candidate from a JSON request body. Returns - /// `None` if the body isn't valid JSON or doesn't contain a routable - /// text field; the caller falls back to non-cache-aware routing. + /// Byte-slice convenience wrapper over [`Self::extract_prompt_text_from_value`]. + /// Only the parsed-value form is on the hot path ([`Self::tokens_for_request`]); + /// this wrapper exists for the extraction unit tests. + #[cfg(test)] + fn extract_prompt_text(body: &[u8]) -> Option { + let v: serde_json::Value = serde_json::from_slice(body).ok()?; + Self::extract_prompt_text_from_value(&v) + } + + /// Extract a raw prompt-text candidate from an already-parsed JSON request + /// body (the body is parsed once in [`Self::tokens_for_request`]). Returns + /// `None` when there's no routable text field; the caller falls back to + /// non-cache-aware routing. This is the raw path — chat requests on a model + /// with a chat template are tokenized via the template instead. /// /// Supported shapes (in priority order): /// 1. `"prompt": "..."` — `/v1/completions`-style. @@ -145,8 +158,7 @@ impl CacheAwareZmqPolicy { /// 5. `"text": "..."` — SGLang `/generate` native form. /// /// Anything else yields `None`. - fn extract_prompt_text(body: &[u8]) -> Option { - let v: serde_json::Value = serde_json::from_slice(body).ok()?; + fn extract_prompt_text_from_value(v: &serde_json::Value) -> Option { if let Some(s) = v.get("prompt").and_then(|p| p.as_str()) { return Some(s.to_string()); } @@ -189,6 +201,28 @@ impl CacheAwareZmqPolicy { None } + /// Produce the token sequence to hash for this request. + /// + /// Chat requests (`messages`) on a model that has a chat template are + /// rendered through that template and tokenized the way the engine does, so + /// the query hashes match the engine's cached blocks (which are keyed on + /// chat-templated tokens). Everything else — `/v1/completions` (`prompt`), + /// `/generate` (`text`), or a chat model without a template — tokenizes the + /// raw extracted prompt text, unchanged. A failed template render/encode + /// falls through to the raw path rather than failing the request. + fn tokens_for_request(&self, model_id: &ModelId, body: &[u8]) -> Option> { + let value: serde_json::Value = serde_json::from_slice(body).ok()?; + if self.tokenizers.has_chat_encoder(&model_id.0) { + if let Some(messages) = value.get("messages").filter(|m| m.is_array()) { + if let Some(tokens) = self.tokenizers.encode_chat(&model_id.0, messages) { + return Some(tokens); + } + } + } + let text = Self::extract_prompt_text_from_value(&value)?; + self.tokenize(model_id, &text) + } + /// Tokenize `text` for `model_id`. Returns `None` if no tokenizer is /// loaded (the model_id may be misconfigured) or if encoding fails. /// Errors log at debug — they degrade routing but are not fatal. @@ -221,19 +255,17 @@ impl Policy for CacheAwareZmqPolicy { return Self::pick_min_load(workers); } - // 2. Extract the prompt text. + // 2. Tokenize the request (chat-template-aware for chat traffic on + // models that ship a template; raw prompt text otherwise). let body = match ctx.request_body() { Some(b) if !b.is_empty() => b, _ => return Self::pick_min_load(workers), }; - let Some(text) = Self::extract_prompt_text(body) else { + let Some(tokens) = self.tokens_for_request(ctx.model(), body) else { return Self::pick_min_load(workers); }; - // 3. Tokenize + hash + match. - let Some(tokens) = self.tokenize(ctx.model(), &text) else { - return Self::pick_min_load(workers); - }; + // 3. Hash + match. // Source block_size from the worker — the router can only hash // prompts at the block size the workers publish at. If no worker // has registered yet (oracle empty), cache-aware routing has no @@ -713,6 +745,237 @@ mod tests { } } + /// A chat-completions request on a model with a chat template must route by + /// the **chat-templated** tokens (BOS + role markers + content) — the tokens + /// the engine actually cached — not by the raw joined content. Worker w0 + /// published its blocks under the templated tokens; only a router that + /// renders the same template hashes a matching query. Hashing the raw + /// content instead would match nothing, leaving live `overlap_blocks_sum` + /// at 0 for chat traffic. + #[test] + fn chat_request_routes_by_templated_tokens() { + let registry = tokenizer_registry_with_tiny(); + let template = serde_json::json!({ + "chat_template": "{{ bos_token }}{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] }}{% endfor %}<|assistant|>", + "bos_token": "", + }); + registry.attach_chat_template_for_test("tiny", &template); + + let messages = serde_json::json!([{"role":"user","content":"hello world hello world"}]); + // Engine-side blocks are keyed on tokenize(render(messages)). + let templated_tokens = registry.encode_chat("tiny", &messages).unwrap(); + let block_size = 4u32; + let templated_hashes = compute_block_hashes(&templated_tokens, block_size as usize); + assert!( + !templated_hashes.is_empty(), + "templated prompt must produce at least one block" + ); + + let tree = Arc::new(HashTree::new()); + tree.insert( + &KvWorkerId::new("http://w0:30000".into(), 0), + None, + &templated_hashes, + ); + + let policy = CacheAwareZmqPolicy::new( + CacheAwareConfig { + cache_threshold: 0.0, + balance_abs_threshold: 32, + balance_rel_threshold: 1.1, + }, + tree, + registry, + oracle_for_tests(block_size), + ); + let w0 = worker("http://w0:30000", "tiny"); + let w1 = worker("http://w1:30000", "tiny"); + let workers = vec![Arc::clone(&w0), Arc::clone(&w1)]; + let model = ModelId("tiny".into()); + let body = serde_json::to_vec(&serde_json::json!({ + "model": "tiny", + "messages": messages, + })) + .unwrap(); + let ctx = SelectionContext::new(&model, Some(&body)); + let chosen = policy.select(&workers, &ctx).expect("must pick"); + assert_eq!( + chosen.url, "http://w0:30000", + "chat request must route by chat-templated tokens to the worker holding that prefix" + ); + } + + /// Templated and raw-content hashings must genuinely differ, confirming + /// the chat-template path does real work (a no-op template would make this + /// assertion fail, and raw-content hashes would miss the engine's + /// templated blocks). + #[test] + fn chat_templated_hashes_differ_from_raw_content_hashes() { + let registry = tokenizer_registry_with_tiny(); + let template = serde_json::json!({ + "chat_template": "{{ bos_token }}{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] }}{% endfor %}<|assistant|>", + "bos_token": "", + }); + registry.attach_chat_template_for_test("tiny", &template); + let content = "hello world hello world"; + let messages = serde_json::json!([{"role":"user","content":content}]); + + let templated = registry.encode_chat("tiny", &messages).unwrap(); + let raw = adapter::encode(®istry.get("tiny").unwrap(), content).unwrap(); + assert_ne!( + compute_block_hashes(&templated, 4), + compute_block_hashes(&raw, 4), + "templated and raw-content block hashes must differ" + ); + } + + /// The DeepSeek-V4 built-in encoder is dispatched for chat requests when a + /// model has it (no Jinja template). The query tokens come from the V4 + /// encoder, so a worker holding that encoded prefix is matched. (The V4 + /// markers aren't special tokens in the tiny fixture, but the dispatch + + /// routing wiring is what's under test; byte-exact V4 token parity is pinned + /// by `dsv4`'s string goldens and validated live.) + #[test] + fn chat_request_routes_via_dsv4_encoder() { + let registry = tokenizer_registry_with_tiny(); + registry.attach_chat_encoder_for_test("tiny", crate::tokenizer::ChatEncoder::DeepSeekV4); + assert!(registry.has_chat_encoder("tiny")); + + let messages = + serde_json::json!([{"role":"user","content":"hello world hello world hello world"}]); + let encoded = registry.encode_chat("tiny", &messages).unwrap(); + let block_size = 4u32; + let hashes = compute_block_hashes(&encoded, block_size as usize); + assert!(!hashes.is_empty()); + + let tree = Arc::new(HashTree::new()); + tree.insert(&KvWorkerId::new("http://w0:30000".into(), 0), None, &hashes); + let policy = CacheAwareZmqPolicy::new( + CacheAwareConfig { + cache_threshold: 0.0, + balance_abs_threshold: 32, + balance_rel_threshold: 1.1, + }, + tree, + registry, + oracle_for_tests(block_size), + ); + let w0 = worker("http://w0:30000", "tiny"); + let w1 = worker("http://w1:30000", "tiny"); + let workers = vec![Arc::clone(&w0), Arc::clone(&w1)]; + let model = ModelId("tiny".into()); + let body = serde_json::to_vec(&serde_json::json!({ "messages": messages })).unwrap(); + let ctx = SelectionContext::new(&model, Some(&body)); + let chosen = policy.select(&workers, &ctx).expect("must pick"); + assert_eq!( + chosen.url, "http://w0:30000", + "dsv4 chat request must route by the V4-encoded prefix" + ); + } + + /// Helper: a tree holding `content`'s RAW-tokenized block hashes on w0, the + /// two workers, and a policy — the fixture the raw-fallback routing tests + /// share. Returns (policy, workers, model). + fn raw_prefix_fixture( + registry: Arc, + content: &str, + ) -> (CacheAwareZmqPolicy, Vec>, ModelId) { + let raw_tokens = adapter::encode(®istry.get("tiny").unwrap(), content).unwrap(); + let hashes = compute_block_hashes(&raw_tokens, 4); + assert!( + !hashes.is_empty(), + "raw content must produce at least one block" + ); + let tree = Arc::new(HashTree::new()); + tree.insert(&KvWorkerId::new("http://w0:30000".into(), 0), None, &hashes); + let policy = CacheAwareZmqPolicy::new( + CacheAwareConfig { + cache_threshold: 0.0, + balance_abs_threshold: 32, + balance_rel_threshold: 1.1, + }, + tree, + registry, + oracle_for_tests(4), + ); + let workers = vec![ + worker("http://w0:30000", "tiny"), + worker("http://w1:30000", "tiny"), + ]; + (policy, workers, ModelId("tiny".into())) + } + + /// Graceful degradation: a model that HAS a chat template whose render fails + /// (here it always raises) must fall back to hashing the RAW content and + /// still route by prefix — not error, not blindly min-load. Exercises the + /// `tokens_for_request` fall-through that the leaf `encode_chat`-returns-None + /// tests don't reach at the routing level. + #[test] + fn chat_render_failure_falls_back_to_raw_routing() { + let registry = tokenizer_registry_with_tiny(); + registry.attach_chat_template_for_test( + "tiny", + &serde_json::json!({ + "chat_template": "{{ raise_exception('boom') }}", + "bos_token": "", + }), + ); + let content = "hello world hello world hello world"; + let (policy, workers, model) = raw_prefix_fixture(registry, content); + let body = serde_json::to_vec(&serde_json::json!({ + "messages": [{"role": "user", "content": content}], + })) + .unwrap(); + let ctx = SelectionContext::new(&model, Some(&body)); + let chosen = policy.select(&workers, &ctx).expect("must pick"); + assert_eq!( + chosen.url, "http://w0:30000", + "a failed template render must degrade to raw-content routing" + ); + } + + /// A chat request on a model WITHOUT a chat template routes by the raw + /// joined `messages[*].content` — the common config where the model ships + /// no `chat_template`. Covers the `tokens_for_request` path that skips the + /// template block entirely for a `messages` body. + #[test] + fn chat_on_template_less_model_routes_by_raw_content() { + let registry = tokenizer_registry_with_tiny(); // no template attached + assert!(!registry.has_chat_encoder("tiny")); + let content = "hello world hello world hello world"; + let (policy, workers, model) = raw_prefix_fixture(registry, content); + let body = serde_json::to_vec(&serde_json::json!({ + "messages": [{"role": "user", "content": content}], + })) + .unwrap(); + let ctx = SelectionContext::new(&model, Some(&body)); + let chosen = policy.select(&workers, &ctx).expect("must pick"); + assert_eq!(chosen.url, "http://w0:30000"); + } + + /// A `/v1/completions` (`prompt`) request on a model that DOES have a chat + /// template must still use the raw path — the template applies only to + /// `messages` traffic. Guards the `messages`-presence gate in + /// `tokens_for_request`. + #[test] + fn completions_prompt_on_templated_model_uses_raw_path() { + let registry = tokenizer_registry_with_tiny(); + registry.attach_chat_template_for_test( + "tiny", + &serde_json::json!({ + "chat_template": "{{ bos_token }}{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] }}{% endfor %}", + "bos_token": "", + }), + ); + let content = "hello world hello world hello world"; + let (policy, workers, model) = raw_prefix_fixture(registry, content); + // `prompt` body (no `messages`) -> raw path, so it matches the raw tree. + let body = serde_json::to_vec(&serde_json::json!({ "prompt": content })).unwrap(); + let ctx = SelectionContext::new(&model, Some(&body)); + let chosen = policy.select(&workers, &ctx).expect("must pick"); + assert_eq!(chosen.url, "http://w0:30000"); + } + /// Two workers both hold the prefix; the lower-load one wins. #[test] fn tie_break_by_lowest_active_load() { diff --git a/experimental/sgl-router/src/tokenizer/adapter.rs b/experimental/sgl-router/src/tokenizer/adapter.rs index efec68053..ca9007467 100644 --- a/experimental/sgl-router/src/tokenizer/adapter.rs +++ b/experimental/sgl-router/src/tokenizer/adapter.rs @@ -44,23 +44,68 @@ fn looks_like_path(source: &str) -> bool { } /// Download `tokenizer.json` for a HuggingFace repo id and return the cached -/// local path. Uses the blocking `ureq` API (this runs once at startup, -/// before the server begins serving) and `from_env` so `HF_TOKEN` / -/// `HF_HOME` / endpoint overrides are honored. +/// local path, adding an actionable error context. The actual fetch (blocking +/// `ureq`, `from_env` so `HF_TOKEN` / `HF_HOME` / endpoint overrides apply) +/// lives in [`download_repo_file`]. fn download_tokenizer_json(repo_id: &str) -> Result { + download_repo_file(repo_id, "tokenizer.json").with_context(|| { + format!( + "download tokenizer.json for HuggingFace repo {repo_id:?} \ + (pass --tokenizer-path with a local tokenizer.json, or set HF_TOKEN \ + for a gated/private repo)" + ) + }) +} + +/// Download `file` from a HuggingFace repo id and return the cached local path. +/// Shared by `tokenizer.json` (required) and `tokenizer_config.json` (optional). +fn download_repo_file(repo_id: &str, file: &str) -> Result { use hf_hub::api::sync::ApiBuilder; let api = ApiBuilder::from_env() .build() .context("initialize HuggingFace Hub client")?; api.model(repo_id.to_string()) - .get("tokenizer.json") - .with_context(|| { - format!( - "download tokenizer.json for HuggingFace repo {repo_id:?} \ - (pass --tokenizer-path with a local tokenizer.json, or set HF_TOKEN \ - for a gated/private repo)" - ) - }) + .get(file) + .with_context(|| format!("download {file} for HuggingFace repo {repo_id:?}")) +} + +/// Load the `tokenizer_config.json` co-located with the tokenizer named by +/// `source` (the same value passed to [`load`]). For a local +/// `.../tokenizer.json` path this is the sibling file; for an HF repo id it is +/// downloaded from the same repo. +/// +/// Returns `Ok(None)` when the model ships no `tokenizer_config.json` (rare but +/// valid) — the caller then has no chat template and routes via raw prompt text. +pub fn load_tokenizer_config(source: &str) -> Result> { + let path = if Path::new(source).is_file() || looks_like_path(source) { + match Path::new(source).parent() { + Some(dir) => dir.join("tokenizer_config.json"), + None => return Ok(None), + } + } else { + // HF repo id. The download error type doesn't distinguish a genuine + // 404 (repo ships no tokenizer_config.json — benign) from auth/network + // failures (wrong/expired HF_TOKEN, gated repo, timeout), so warn with + // the cause rather than asserting the benign case at debug: a swallowed + // auth error here silently disables chat-template routing. + match download_repo_file(source, "tokenizer_config.json") { + Ok(p) => p, + Err(e) => { + tracing::warn!(repo = %source, error = %e, + "could not download tokenizer_config.json; chat-template routing disabled for this model \ + (expected if the repo ships none — otherwise check HF_TOKEN / network for a gated or private repo)"); + return Ok(None); + } + } + }; + if !path.is_file() { + return Ok(None); + } + let bytes = std::fs::read(&path) + .with_context(|| format!("read tokenizer_config.json at {}", path.display()))?; + let value = serde_json::from_slice(&bytes) + .with_context(|| format!("parse tokenizer_config.json at {}", path.display()))?; + Ok(Some(value)) } pub fn encode(t: &Tokenizer, text: &str) -> Result> { diff --git a/experimental/sgl-router/src/tokenizer/chat_template.rs b/experimental/sgl-router/src/tokenizer/chat_template.rs new file mode 100644 index 000000000..2e1968630 --- /dev/null +++ b/experimental/sgl-router/src/tokenizer/chat_template.rs @@ -0,0 +1,366 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Chat-template rendering for cache-aware routing. +//! +//! The engine caches KV blocks keyed on tokens it produces *after* applying the +//! model's chat template (BOS + role/special markers + content). The router's +//! cache-aware selection must hash the same token sequence, so it renders the +//! same template before tokenizing — otherwise its query hashes never match the +//! engine's stored blocks and cache-aware routing silently degrades to min-load +//! (`sgl_router_overlap_blocks_sum` stuck at 0). +//! +//! The template and its special-token strings come from the model's +//! `tokenizer_config.json` — the HuggingFace built-in template, which is what +//! the engine uses unless launched with an explicit chat-template override. +//! +//! Tokenization does not auto-prepend special tokens (the `dynamo_tokenizers` +//! HF wrapper hardcodes `add_special_tokens = false`; [`super::adapter::encode`] +//! adds none of its own), so the rendered text must already contain `bos_token` +//! and the role markers as literal text. That matches HuggingFace +//! `apply_chat_template(tokenize=True)` semantics, where the template — not the +//! tokenizer's special-token insertion — is the single source of the leading +//! specials. + +use anyhow::{Context, Result}; +use minijinja::{ + value::Value as JinjaValue, Environment, Error as JinjaError, ErrorKind as JinjaErrorKind, + UndefinedBehavior, +}; +use std::collections::BTreeMap; + +/// Template registered under a fixed name in the per-model environment. +const TEMPLATE_NAME: &str = "chat"; + +/// The named special tokens HuggingFace injects into the template context via +/// `special_tokens_map`. Each is supplied from `tokenizer_config.json`, or as +/// the empty string when absent — jinja2 renders an undefined name as `""`, so +/// an absent token must not surface as anything else (minijinja would otherwise +/// print a `none` value as the literal string "none", silently diverging every +/// block hash from the engine's). +const SPECIAL_TOKEN_KEYS: [&str; 7] = [ + "bos_token", + "eos_token", + "unk_token", + "sep_token", + "pad_token", + "cls_token", + "mask_token", +]; + +/// A compiled chat template plus the special-token strings it references. +/// +/// One per model, built once at startup from `tokenizer_config.json` and held +/// in the [`super::TokenizerRegistry`]. Rendering is read-only and thread-safe. +pub struct ChatTemplate { + env: Environment<'static>, + /// `(name, token)` pairs for [`SPECIAL_TOKEN_KEYS`]; absent tokens are `""`. + special_tokens: Vec<(&'static str, String)>, +} + +impl std::fmt::Debug for ChatTemplate { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ChatTemplate") + .field("special_tokens", &self.special_tokens) + .finish() + } +} + +impl ChatTemplate { + /// Build from a parsed `tokenizer_config.json`. Returns `Ok(None)` when the + /// config carries no `chat_template` (the model is then routed via the raw + /// prompt-text path, unchanged). + pub fn from_tokenizer_config(cfg: &serde_json::Value) -> Result> { + let Some(template_src) = extract_chat_template(cfg) else { + return Ok(None); + }; + let special_tokens = SPECIAL_TOKEN_KEYS + .iter() + .map(|&key| (key, extract_token_str(cfg, key).unwrap_or_default())) + .collect(); + + let mut env = Environment::new(); + // HuggingFace compiles chat templates with trim_blocks + lstrip_blocks; + // mirror that or rendered whitespace (and thus tokens) diverge. + env.set_trim_blocks(true); + env.set_lstrip_blocks(true); + // Printing a variable the router didn't supply (a custom + // `chat_template_kwargs` entry, a date var, ...) must be a render + // error so the caller falls back to raw-text hashing — under the + // default lenient behavior it would render as `""` and produce a + // plausible-but-divergent prompt whose hashes silently never match + // the engine's. If-tests and iteration over undefined stay permitted + // (`{% if enable_thinking is defined %}`-style guards are common). + env.set_undefined_behavior(UndefinedBehavior::SemiStrict); + // Python str/dict methods used by real templates (.startswith, .items, + // .strip, ...) that minijinja doesn't implement natively. + env.set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback); + env.add_function("raise_exception", raise_exception); + env.add_function("strftime_now", strftime_now); + env.add_template_owned(TEMPLATE_NAME, template_src) + .context("compile chat template from tokenizer_config.json")?; + + Ok(Some(Self { + env, + special_tokens, + })) + } + + /// Render `messages` (the request's `messages` array) into the prompt text + /// the engine would tokenize, with `add_generation_prompt = true`. + /// + /// `messages` is passed through as-is; templates expect string `content`. + /// Multimodal content arrays are out of scope (text-only routing): a + /// template may stringify the array (divergent hashes → min-load) or error + /// (raw prompt-text fallback); neither fails the request. + /// + /// `tools` and `documents` are supplied as `none` — the context HuggingFace + /// renders with when a request carries neither, so tools-branching + /// templates take the no-tools path. A request that does carry them renders + /// the no-tools form, so its hashes won't match the engine and it routes by + /// min-load — no worse than before this path existed. Any other variable + /// the template prints is a render error (semi-strict undefined), falling + /// back to raw rather than hashing a silently divergent prompt. + pub fn render(&self, messages: &serde_json::Value) -> Result { + let tmpl = self + .env + .get_template(TEMPLATE_NAME) + .context("chat template not registered")?; + let mut ctx: BTreeMap<&str, JinjaValue> = BTreeMap::new(); + ctx.insert("messages", JinjaValue::from_serialize(messages)); + ctx.insert("add_generation_prompt", JinjaValue::from(true)); + ctx.insert("tools", JinjaValue::from(())); + ctx.insert("documents", JinjaValue::from(())); + for (name, token) in &self.special_tokens { + ctx.insert(name, JinjaValue::from(token.clone())); + } + tmpl.render(ctx).context("render chat template") + } +} + +/// Pull the chat-template source out of `tokenizer_config.json`. +/// +/// Accepts both shapes HuggingFace ships: +/// - `"chat_template": ""` — the common single-template case. +/// - `"chat_template": [{"name": "default", "template": ""}, ...]` — +/// multi-template models; we take the entry named `default`, else the first. +fn extract_chat_template(cfg: &serde_json::Value) -> Option { + match cfg.get("chat_template")? { + serde_json::Value::String(s) => Some(s.clone()), + serde_json::Value::Array(arr) => arr + .iter() + .find(|e| e.get("name").and_then(|n| n.as_str()) == Some("default")) + .or_else(|| arr.first()) + .and_then(|e| e.get("template").and_then(|t| t.as_str())) + .map(str::to_owned), + _ => None, + } +} + +/// Read a special-token string, accepting both the plain-string form and the +/// `AddedToken` object form (`{"content": "", ...}`) HuggingFace uses. +fn extract_token_str(cfg: &serde_json::Value, key: &str) -> Option { + match cfg.get(key)? { + serde_json::Value::String(s) => Some(s.clone()), + serde_json::Value::Object(o) => { + o.get("content").and_then(|c| c.as_str()).map(str::to_owned) + } + _ => None, + } +} + +/// `raise_exception(msg)` — templates call this to reject malformed message +/// sequences (e.g. a non-alternating role order). Surfaces as a render error. +fn raise_exception(msg: String) -> std::result::Result { + Err(JinjaError::new(JinjaErrorKind::InvalidOperation, msg)) +} + +/// `strftime_now(format)` — current local time, matching the helper HuggingFace +/// injects so templates can stamp the date. Both engine and router render +/// within the same day, so the date prefix is stable enough to share a cache +/// block. +fn strftime_now(format: String) -> String { + chrono::Local::now().format(&format).to_string() +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + /// A small but representative instruct template: emits `bos_token`, wraps + /// each turn in role markers, and appends a generation prompt. Exercises the + /// variables the renderer must supply (`messages`, `bos_token`, + /// `add_generation_prompt`). + const SIMPLE_TEMPLATE: &str = "{{ bos_token }}{% for m in messages %}<|{{ m['role'] }}|>\n{{ m['content'] }}<|end|>\n{% endfor %}{% if add_generation_prompt %}<|assistant|>\n{% endif %}"; + + fn messages() -> serde_json::Value { + json!([ + {"role": "system", "content": "be brief"}, + {"role": "user", "content": "hi"} + ]) + } + + #[test] + fn no_chat_template_returns_none() { + let cfg = json!({"bos_token": "", "eos_token": ""}); + assert!(ChatTemplate::from_tokenizer_config(&cfg).unwrap().is_none()); + } + + #[test] + fn renders_roles_bos_and_generation_prompt() { + let cfg = json!({ + "chat_template": SIMPLE_TEMPLATE, + "bos_token": "", + "eos_token": "", + }); + let tmpl = ChatTemplate::from_tokenizer_config(&cfg).unwrap().unwrap(); + let out = tmpl.render(&messages()).unwrap(); + assert_eq!( + out, + "<|system|>\nbe brief<|end|>\n<|user|>\nhi<|end|>\n<|assistant|>\n" + ); + } + + /// `add_generation_prompt` is always true on the routing side (we hash the + /// prompt the engine will prefill, which includes the assistant header). + #[test] + fn generation_prompt_is_always_appended() { + let cfg = json!({ "chat_template": SIMPLE_TEMPLATE, "bos_token": "" }); + let tmpl = ChatTemplate::from_tokenizer_config(&cfg).unwrap().unwrap(); + assert!(tmpl + .render(&messages()) + .unwrap() + .ends_with("<|assistant|>\n")); + } + + /// The list form `[{name, template}, ...]` selects the `default` entry. + #[test] + fn list_form_selects_default_template() { + let cfg = json!({ + "chat_template": [ + {"name": "tool_use", "template": "TOOLS"}, + {"name": "default", "template": SIMPLE_TEMPLATE}, + ], + "bos_token": "", + }); + let tmpl = ChatTemplate::from_tokenizer_config(&cfg).unwrap().unwrap(); + assert!(tmpl + .render(&messages()) + .unwrap() + .starts_with("<|system|>")); + } + + /// `bos_token` in the `AddedToken` object form is read from `.content`. + #[test] + fn bos_token_object_form_is_extracted() { + let cfg = json!({ + "chat_template": "{{ bos_token }}X", + "bos_token": {"content": "<|begin|>", "lstrip": false}, + }); + let tmpl = ChatTemplate::from_tokenizer_config(&cfg).unwrap().unwrap(); + assert_eq!(tmpl.render(&json!([])).unwrap(), "<|begin|>X"); + } + + /// `raise_exception` surfaces as a render error (caller then falls back to + /// the raw prompt-text path rather than failing the request). + #[test] + fn raise_exception_surfaces_as_error() { + let cfg = json!({ + "chat_template": "{{ raise_exception('bad messages') }}", + "bos_token": "", + }); + let tmpl = ChatTemplate::from_tokenizer_config(&cfg).unwrap().unwrap(); + let err = tmpl.render(&messages()).unwrap_err(); + // The minijinja message is the cause; check the full anyhow chain. + assert!(format!("{err:#}").contains("bad messages"), "got: {err:#}"); + } + + /// pycompat exposes Python str methods (`.startswith`, `.upper`, ...) that + /// real HuggingFace templates lean on; without the callback these error. + #[test] + fn pycompat_string_methods_available() { + let cfg = json!({ + "chat_template": "{% for m in messages %}{% if m['role'].startswith('sys') %}{{ m['content'].upper() }}{% endif %}{% endfor %}", + "bos_token": "", + }); + let tmpl = ChatTemplate::from_tokenizer_config(&cfg).unwrap().unwrap(); + assert_eq!(tmpl.render(&messages()).unwrap(), "BE BRIEF"); + } + + /// An absent special token renders as `""` exactly like an undefined name + /// under HuggingFace's jinja2 — never as minijinja's literal `"none"`, + /// which would corrupt block 0 (and thus every chained block hash). + #[test] + fn absent_special_tokens_render_empty() { + let cfg = json!({"chat_template": "A{{ bos_token }}{{ pad_token }}B"}); + let tmpl = ChatTemplate::from_tokenizer_config(&cfg).unwrap().unwrap(); + assert_eq!(tmpl.render(&json!([])).unwrap(), "AB"); + } + + /// Every name in HuggingFace's `special_tokens_map` is threaded from + /// `tokenizer_config.json`, not just `bos_token`/`eos_token`. + #[test] + fn named_special_tokens_from_config_are_supplied() { + let cfg = json!({ + "chat_template": "{{ pad_token }}|{{ unk_token }}", + "pad_token": "", + "unk_token": {"content": ""}, + }); + let tmpl = ChatTemplate::from_tokenizer_config(&cfg).unwrap().unwrap(); + assert_eq!(tmpl.render(&json!([])).unwrap(), "|"); + } + + /// Printing a variable the router doesn't supply is a render error + /// (semi-strict undefined) so the caller falls back to raw-text hashing, + /// instead of rendering a plausible-but-divergent prompt. + #[test] + fn printing_unsupplied_variable_fails_render() { + let cfg = json!({ + "chat_template": "{{ custom_kwarg }}", + "bos_token": "", + }); + let tmpl = ChatTemplate::from_tokenizer_config(&cfg).unwrap().unwrap(); + tmpl.render(&messages()).unwrap_err(); + } + + /// Undefined names stay usable in if-tests (semi-strict only rejects + /// printing them); common `{% if enable_thinking is defined %}`-style + /// guards must keep rendering. + #[test] + fn undefined_in_if_test_is_permitted() { + let cfg = json!({ + "chat_template": "{% if enable_thinking is defined and enable_thinking %}T{% endif %}X", + }); + let tmpl = ChatTemplate::from_tokenizer_config(&cfg).unwrap().unwrap(); + assert_eq!(tmpl.render(&messages()).unwrap(), "X"); + } + + /// `tools` is `none` in the render context — the same context HuggingFace + /// renders with for a request that carries no tools — so tools-branching + /// templates take the no-tools path instead of erroring or mis-branching. + #[test] + fn tools_supplied_as_none_takes_no_tools_branch() { + let cfg = json!({ + "chat_template": "{% if tools is not none %}TOOLS{% endif %}X", + }); + let tmpl = ChatTemplate::from_tokenizer_config(&cfg).unwrap().unwrap(); + assert_eq!(tmpl.render(&messages()).unwrap(), "X"); + } + + /// trim_blocks + lstrip_blocks match HuggingFace's compilation: the newline + /// after a block tag and leading whitespace before one are stripped, so a + /// block-per-line template renders without spurious blank lines. + #[test] + fn trim_and_lstrip_blocks_match_huggingface() { + let cfg = json!({ + "chat_template": "{% for m in messages %}\n {% if true %}\n{{ m['role'] }}\n {% endif %}\n{% endfor %}", + "bos_token": "", + }); + let tmpl = ChatTemplate::from_tokenizer_config(&cfg).unwrap().unwrap(); + // Each iteration emits just "\n"; lstrip removes the two leading + // spaces before the `{% if %}`/`{% endif %}`, trim removes the newline + // immediately after each block tag. + assert_eq!(tmpl.render(&messages()).unwrap(), "system\nuser\n"); + } +} diff --git a/experimental/sgl-router/src/tokenizer/dsv4.rs b/experimental/sgl-router/src/tokenizer/dsv4.rs new file mode 100644 index 000000000..07baa3292 --- /dev/null +++ b/experimental/sgl-router/src/tokenizer/dsv4.rs @@ -0,0 +1,274 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! DeepSeek-V4 prompt encoder for cache-aware routing. +//! +//! DeepSeek-V4 ships no Jinja chat template; the engine builds the prompt in +//! code (`python/sglang/srt/entrypoints/openai/encoding_dsv4.py`, selected for +//! the `DeepseekV4` architecture). So to make the router's query tokens match +//! the engine's cached blocks, this reproduces that encoder's output for the +//! routing-relevant subset. +//! +//! # Scope +//! +//! Text content, chat (non-thinking) mode — the engine default +//! (`SGLANG_DEFAULT_THINKING=false`). For a user turn the engine emits +//! `BOS <|User|> content <|Assistant|> `. Two preprocessing steps mirror +//! `serving_chat.py`/`encode_messages`: an empty system message is inserted when +//! the first message isn't a system message, and consecutive user turns are +//! merged into one (joined with `\n\n`). Tools, tasks, and per-turn reasoning +//! content are out of scope: the engine renders tools immediately after the +//! system content at the front of the prompt, so a tools-carrying request +//! diverges from the first block and routes by min-load; tasks alter only the +//! trailing turn transition; reasoning content is never emitted in chat mode, +//! so it causes no divergence. +//! +//! Tokenization does not auto-prepend special tokens (the `dynamo_tokenizers` +//! HF wrapper hardcodes `add_special_tokens = false`; +//! [`super::adapter::encode`] adds none of its own), so the literal marker text +//! below is what maps to the special token ids. Pinned byte-exact against the +//! live engine's `/tokenize` (DeepSeek-V4-Flash, snapshot `6976c7ff`): +//! `[{user:"ABCD"}]` → `[0, 128803, 51453, 128804, 128822]`. + +/// Beginning-of-sequence marker (token id 0). +const BOS: &str = "<|begin▁of▁sentence|>"; +/// End-of-sequence marker, closing each prior assistant turn (token id 1). +const EOS: &str = "<|end▁of▁sentence|>"; +/// User-turn marker (token id 128803). +const USER: &str = "<|User|>"; +/// Assistant-turn marker, opening the generation prompt (token id 128804). +const ASSISTANT: &str = "<|Assistant|>"; +/// Thinking-end marker; the chat-mode generation prompt ends with it (128822). +const THINK_END: &str = ""; + +/// Render `messages` into the DeepSeek-V4 chat prompt for routing. +/// +/// Mirrors `encoding_dsv4.encode_messages` for the routing subset (chat mode, +/// text content, no tools/tasks). `messages` is the request's `messages` array; +/// non-array input renders to just the BOS marker (the caller then tokenizes it +/// and, finding no useful prefix, degrades to min-load like any short prompt). +pub fn render_messages(messages: &serde_json::Value) -> String { + let mut msgs: Vec<(String, String)> = messages + .as_array() + .map(|arr| { + arr.iter() + .map(|m| { + let role = m + .get("role") + .and_then(|r| r.as_str()) + .unwrap_or("") + .to_string(); + (role, content_to_string(m.get("content"))) + }) + .collect() + }) + .unwrap_or_default(); + + // The engine inserts an empty system message when the first message isn't a + // system message; it renders to nothing but keeps the index logic aligned. + if msgs.first().map(|(r, _)| r != "system").unwrap_or(true) { + msgs.insert(0, ("system".to_string(), String::new())); + } + + merge_consecutive_user_turns(&mut msgs); + + let mut out = String::from(BOS); + for i in 0..msgs.len() { + render_one(i, &msgs, &mut out); + } + out +} + +/// Collapse runs of consecutive `user` turns into one, joining their content +/// with `\n\n` — the engine merges them (`merge_tool_messages`) before encoding, +/// so back-to-back user messages must hash like the single turn it builds. +/// `developer` and other roles break a run and are left as-is. +fn merge_consecutive_user_turns(msgs: &mut Vec<(String, String)>) { + let mut merged: Vec<(String, String)> = Vec::with_capacity(msgs.len()); + for (role, content) in msgs.drain(..) { + match merged.last_mut() { + Some((last_role, last_content)) if last_role == "user" && role == "user" => { + last_content.push_str("\n\n"); + last_content.push_str(&content); + } + _ => merged.push((role, content)), + } + } + *msgs = merged; +} + +/// Append message `i`'s encoded form to `out`. +fn render_one(i: usize, msgs: &[(String, String)], out: &mut String) { + let (role, content) = &msgs[i]; + match role.as_str() { + "system" => out.push_str(content), + "user" | "developer" => { + out.push_str(USER); + out.push_str(content); + } + "assistant" => { + // Chat mode emits no reasoning block, so a prior assistant turn is + // just its content closed by EOS. + out.push_str(content); + out.push_str(EOS); + } + // Unknown roles aren't part of routing traffic; emit the content so a + // stray role still contributes something rather than vanishing. + _ => out.push_str(content), + } + + // Generation-prompt transition. The engine appends it only when this is the + // last message OR the next message is an assistant/reminder turn, and only + // for user/developer messages. + let next_takes_transition = match msgs.get(i + 1) { + Some((next_role, _)) => next_role == "assistant" || next_role == "latest_reminder", + None => true, + }; + if next_takes_transition && (role == "user" || role == "developer") { + out.push_str(ASSISTANT); + out.push_str(THINK_END); + } +} + +/// Flatten a message `content` field to a string: a plain string as-is, an +/// OpenAI parts array to its concatenated `text` parts, anything else to empty. +fn content_to_string(content: Option<&serde_json::Value>) -> String { + match content { + Some(serde_json::Value::String(s)) => s.clone(), + Some(serde_json::Value::Array(parts)) => parts + .iter() + .filter_map(|p| p.get("text").and_then(|t| t.as_str())) + .collect(), + _ => String::new(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + /// Byte-exact against the engine's `/tokenize`: a single user turn renders + /// `BOS <|User|> content <|Assistant|> `. + #[test] + fn single_user_turn() { + let out = render_messages(&json!([{"role":"user","content":"ABCD"}])); + assert_eq!( + out, + "<|begin▁of▁sentence|><|User|>ABCD<|Assistant|>" + ); + } + + /// A leading system message renders as bare content (no marker), before the + /// user turn. + #[test] + fn system_then_user() { + let out = render_messages(&json!([ + {"role":"system","content":"SYS"}, + {"role":"user","content":"ABCD"} + ])); + assert_eq!( + out, + "<|begin▁of▁sentence|>SYS<|User|>ABCD<|Assistant|>" + ); + } + + /// Multi-turn: each prior user turn gets the generation prompt, the prior + /// assistant turn is closed by EOS. Matches the engine token stream + /// `[0,128803,55,19,128804,128822,35,19,1,128803,55,20,128804,128822]`. + #[test] + fn multi_turn() { + let out = render_messages(&json!([ + {"role":"user","content":"U1"}, + {"role":"assistant","content":"A1"}, + {"role":"user","content":"U2"} + ])); + assert_eq!( + out, + "<|begin▁of▁sentence|><|User|>U1<|Assistant|>A1<|end▁of▁sentence|><|User|>U2<|Assistant|>" + ); + } + + /// An empty leading system message (already present) is not duplicated and + /// renders to nothing — same result as a bare user turn. + #[test] + fn explicit_empty_system_is_not_duplicated() { + let out = render_messages(&json!([ + {"role":"system","content":""}, + {"role":"user","content":"ABCD"} + ])); + assert_eq!( + out, + "<|begin▁of▁sentence|><|User|>ABCD<|Assistant|>" + ); + } + + /// Array (multimodal) content flattens to its text parts. + #[test] + fn array_content_flattens_text_parts() { + let out = render_messages(&json!([ + {"role":"user","content":[{"type":"text","text":"AB"},{"type":"text","text":"CD"}]} + ])); + assert_eq!( + out, + "<|begin▁of▁sentence|><|User|>ABCD<|Assistant|>" + ); + } + + /// Consecutive user turns merge into one `<|User|>` turn joined with `\n\n` + /// (the engine's `merge_tool_messages`), so only one user marker and one + /// generation prompt are emitted — not a marker per message. + #[test] + fn consecutive_user_turns_merge() { + let out = render_messages(&json!([ + {"role":"user","content":"U1"}, + {"role":"user","content":"U2"} + ])); + assert_eq!( + out, + "<|begin▁of▁sentence|><|User|>U1\n\nU2<|Assistant|>" + ); + } + + /// A run of user turns split by an assistant turn does NOT merge across the + /// assistant: each side is its own user turn. + #[test] + fn user_runs_do_not_merge_across_assistant() { + let out = render_messages(&json!([ + {"role":"user","content":"U1"}, + {"role":"user","content":"U2"}, + {"role":"assistant","content":"A1"}, + {"role":"user","content":"U3"} + ])); + assert_eq!( + out, + "<|begin▁of▁sentence|><|User|>U1\n\nU2<|Assistant|>A1<|end▁of▁sentence|><|User|>U3<|Assistant|>" + ); + } + + /// A `developer` turn renders identically to a user turn for text content + /// (the engine nests the same `<|User|>` marker) and takes the generation + /// prompt. Developer turns are not merged (only `user` runs merge), so two + /// developers emit two markers. + #[test] + fn developer_role_renders_like_user_without_merging() { + assert_eq!( + render_messages(&json!([{"role":"developer","content":"D1"}])), + "<|begin▁of▁sentence|><|User|>D1<|Assistant|>" + ); + assert_eq!( + render_messages(&json!([ + {"role":"developer","content":"D1"}, + {"role":"developer","content":"D2"} + ])), + "<|begin▁of▁sentence|><|User|>D1<|User|>D2<|Assistant|>" + ); + } + + /// An empty messages list renders to just the BOS marker — the documented + /// degrade path (the caller then routes by min-load on the empty prefix). + #[test] + fn empty_messages_renders_bos_only() { + assert_eq!(render_messages(&json!([])), "<|begin▁of▁sentence|>"); + } +} diff --git a/experimental/sgl-router/src/tokenizer/mod.rs b/experimental/sgl-router/src/tokenizer/mod.rs index 6b950d748..8df65fd57 100644 --- a/experimental/sgl-router/src/tokenizer/mod.rs +++ b/experimental/sgl-router/src/tokenizer/mod.rs @@ -2,15 +2,78 @@ // SPDX-License-Identifier: Apache-2.0 pub mod adapter; +pub mod chat_template; +pub mod dsv4; use anyhow::Result; +use chat_template::ChatTemplate; use dashmap::DashMap; use dynamo_tokenizers::Tokenizer; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; +/// How to turn a chat request's `messages` into the prompt the engine tokenizes +/// and caches. Cache-aware routing renders this before hashing so its query +/// tokens match the engine's stored blocks. +pub enum ChatEncoder { + /// HuggingFace Jinja chat template from `tokenizer_config.json` (most + /// models). Boxed: it holds a minijinja `Environment`, far larger than the + /// other variants. + Jinja(Box), + /// DeepSeek-V4 ships no template; the engine encodes in code. See [`dsv4`]. + DeepSeekV4, +} + +impl ChatEncoder { + /// Render `messages` into the engine-equivalent prompt text. + fn render(&self, messages: &serde_json::Value) -> Result { + match self { + ChatEncoder::Jinja(t) => t.render(messages), + ChatEncoder::DeepSeekV4 => Ok(dsv4::render_messages(messages)), + } + } +} + +/// A model's chat encoder plus its fallback-logging state. +struct ChatEncoderEntry { + encoder: ChatEncoder, + fallback_warned: AtomicBool, +} + +impl ChatEncoderEntry { + fn new(encoder: ChatEncoder) -> Self { + Self { + encoder, + fallback_warned: AtomicBool::new(false), + } + } + + /// Log a per-request fallback to raw prompt-text hashing. "Enabled but + /// failing every request" must be distinguishable from "healthy" at the + /// default (info) log level — otherwise cache-aware overlap silently + /// degrades to 0 with no signal — so the first failure for a model logs at + /// warn; subsequent ones at debug to avoid a per-request log flood. + fn log_fallback(&self, model_id: &str, cause: &str) { + if !self.fallback_warned.swap(true, Ordering::Relaxed) { + tracing::warn!(model = %model_id, %cause, + "chat-encoder failed; falling back to raw prompt-text hashing \ + (cache-aware overlap degrades for this model; further failures log at debug)"); + } else { + tracing::debug!(model = %model_id, %cause, + "chat-encoder failed; falling back to raw prompt-text hashing"); + } + } +} + #[derive(Default)] pub struct TokenizerRegistry { inner: DashMap>, + /// Per-model chat encoder, present only when the model's prompt format is + /// known (a `tokenizer_config.json` chat template, or a built-in encoder + /// like DeepSeek-V4's). Cache-aware routing uses it to tokenize chat + /// requests the way the engine does; models without one fall back to raw + /// prompt-text tokenization. + encoders: DashMap>, } impl std::fmt::Debug for TokenizerRegistry { @@ -27,16 +90,128 @@ impl TokenizerRegistry { let m = &cfg.model; let t = adapter::load(&m.tokenizer_path)?; me.inner.insert(m.id.clone(), t); + // Resolve the chat encoder, best-effort: a Jinja template from + // tokenizer_config.json, else a built-in encoder for a recognized model + // (DeepSeek-V4), else none (chat traffic routes via raw text). Every + // path logs its outcome — whether chat-aware routing is live for this + // model is the single most useful signal for diagnosing "cache-aware + // routing degraded to overlap=0 on chat traffic", so it must never be + // silent. + if let Some(encoder) = me.resolve_chat_encoder(&m.id, &m.tokenizer_path) { + me.encoders + .insert(m.id.clone(), Arc::new(ChatEncoderEntry::new(encoder))); + } Ok(me) } + /// Pick the chat encoder for a model, logging the outcome on every branch. + fn resolve_chat_encoder(&self, model_id: &str, tokenizer_path: &str) -> Option { + match adapter::load_tokenizer_config(tokenizer_path) { + Ok(Some(cfg_json)) => match ChatTemplate::from_tokenizer_config(&cfg_json) { + Ok(Some(tmpl)) => { + tracing::info!(model = %model_id, + "chat-template routing enabled; chat requests route by templated tokens"); + return Some(ChatEncoder::Jinja(Box::new(tmpl))); + } + Ok(None) => {} // no template — fall through to built-in detection + Err(e) => tracing::warn!(model = %model_id, error = %e, + "failed to compile chat template; falling back to built-in detection"), + }, + Ok(None) => {} + Err(e) => tracing::warn!(model = %model_id, error = %e, + "failed to load tokenizer_config.json; falling back to built-in detection"), + } + if is_deepseek_v4(model_id) { + tracing::info!(model = %model_id, + "DeepSeek-V4 routing enabled; chat requests route via the built-in V4 encoder"); + return Some(ChatEncoder::DeepSeekV4); + } + tracing::info!(model = %model_id, + "no chat template or built-in encoder; chat traffic routes via raw prompt text"); + None + } + pub fn get(&self, model_id: &str) -> Option> { self.inner.get(model_id).map(|r| Arc::clone(&*r)) } + /// Whether this model has a chat encoder (and thus the chat-aware + /// tokenization path is available for it). + pub fn has_chat_encoder(&self, model_id: &str) -> bool { + self.encoders.contains_key(model_id) + } + + /// Render `messages` through the model's chat encoder, then tokenize the + /// result the same way the engine does (`add_special_tokens = false`, so the + /// encoder's literal `bos_token`/role markers carry the specials). Returns + /// `None` — caller falls back to raw routing — when the model has no + /// encoder, no tokenizer, or rendering/encoding fails or yields no tokens. + pub fn encode_chat(&self, model_id: &str, messages: &serde_json::Value) -> Option> { + // Clone the Arc and drop the DashMap guard before the CPU-bound + // render+encode (mirrors `get`), so no shard read-lock is held across it. + let entry = Arc::clone(&*self.encoders.get(model_id)?); + let tokenizer = self.get(model_id)?; + let rendered = entry + .encoder + .render(messages) + .inspect_err(|e| { + // `{e:#}` prints the full anyhow chain, so the underlying + // minijinja cause (e.g. a `raise_exception` message) is + // visible, not just the "render chat template" context. + entry.log_fallback(model_id, &format!("render failed: {e:#}")) + }) + .ok()?; + match adapter::encode(&tokenizer, &rendered) { + Ok(ids) if !ids.is_empty() => Some(ids), + Ok(_) => { + entry.log_fallback(model_id, "rendered prompt tokenized to zero tokens"); + None + } + Err(e) => { + entry.log_fallback(model_id, &format!("tokenize failed: {e:#}")); + None + } + } + } + pub fn ids(&self) -> Vec { self.inner.iter().map(|kv| kv.key().clone()).collect() } + + /// Attach a chat encoder to an already-loaded model. Lets policy tests in + /// other modules exercise the chat-aware routing path without a co-located + /// fixture. + #[cfg(test)] + pub(crate) fn attach_chat_encoder_for_test(&self, model_id: &str, encoder: ChatEncoder) { + self.encoders.insert( + model_id.to_string(), + Arc::new(ChatEncoderEntry::new(encoder)), + ); + } + + /// Convenience: attach a Jinja chat encoder built from an inline + /// `tokenizer_config.json` value. + #[cfg(test)] + pub(crate) fn attach_chat_template_for_test( + &self, + model_id: &str, + tokenizer_config: &serde_json::Value, + ) { + let template = ChatTemplate::from_tokenizer_config(tokenizer_config) + .expect("valid test chat template") + .expect("test tokenizer_config has a chat_template"); + self.attach_chat_encoder_for_test(model_id, ChatEncoder::Jinja(Box::new(template))); + } +} + +/// Whether `model_id` denotes a DeepSeek-V4 model, which the engine encodes via +/// the built-in [`dsv4`] encoder rather than a Jinja template. Heuristic on the +/// served model id (the router has no model architecture from `/server_info`); +/// scoped to "deepseek" + "v4" so it doesn't claim V3-family models, whose +/// encoding differs. +fn is_deepseek_v4(model_id: &str) -> bool { + let id = model_id.to_ascii_lowercase(); + id.contains("deepseek") && id.contains("v4") } #[cfg(test)] @@ -196,4 +371,116 @@ mod tests { let err = TokenizerRegistry::load_from_config(&c).unwrap_err(); assert!(err.to_string().to_lowercase().contains("tokenizer")); } + + #[test] + fn load_tokenizer_config_reads_sibling() { + let dir = tempfile::tempdir().unwrap(); + let tok = dir.path().join("tokenizer.json"); + std::fs::write(&tok, "{}").unwrap(); + std::fs::write( + dir.path().join("tokenizer_config.json"), + r#"{"chat_template":"X","bos_token":""}"#, + ) + .unwrap(); + let cfg = adapter::load_tokenizer_config(tok.to_str().unwrap()) + .unwrap() + .expect("sibling tokenizer_config.json is loaded"); + assert_eq!(cfg["chat_template"], "X"); + } + + #[test] + fn load_tokenizer_config_absent_returns_none() { + let dir = tempfile::tempdir().unwrap(); + let tok = dir.path().join("tokenizer.json"); + std::fs::write(&tok, "{}").unwrap(); + assert!(adapter::load_tokenizer_config(tok.to_str().unwrap()) + .unwrap() + .is_none()); + } + + /// `encode_chat` renders the template then tokenizes the result — and that + /// token sequence differs from tokenizing the raw message content (the very + /// reason raw-content hashing missed the engine's chat-templated blocks). + #[test] + fn encode_chat_renders_then_tokenizes() { + let reg = TokenizerRegistry::default(); + reg.inner.insert( + "tiny".into(), + adapter::load("tests/fixtures/tiny_tokenizer.json").unwrap(), + ); + let cfg = serde_json::json!({ + "chat_template": "{{ bos_token }}{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] }}{% endfor %}", + "bos_token": "", + }); + reg.attach_chat_template_for_test("tiny", &cfg); + assert!(reg.has_chat_encoder("tiny")); + + let messages = serde_json::json!([{"role":"user","content":"hi"}]); + let chat_ids = reg.encode_chat("tiny", &messages).expect("encode_chat"); + assert!(!chat_ids.is_empty()); + + let tok = reg.get("tiny").unwrap(); + let raw_ids = adapter::encode(&tok, "hi").unwrap(); + assert_ne!( + chat_ids, raw_ids, + "chat-templated tokens must differ from raw-content tokens" + ); + + // encode_chat is exactly tokenize(render(messages)). + let rendered = reg + .encoders + .get("tiny") + .unwrap() + .encoder + .render(&messages) + .unwrap(); + assert_eq!(chat_ids, adapter::encode(&tok, &rendered).unwrap()); + } + + #[test] + fn encode_chat_none_without_template() { + let reg = TokenizerRegistry::default(); + reg.inner.insert( + "tiny".into(), + adapter::load("tests/fixtures/tiny_tokenizer.json").unwrap(), + ); + assert!(!reg.has_chat_encoder("tiny")); + let messages = serde_json::json!([{"role":"user","content":"hi"}]); + assert!(reg.encode_chat("tiny", &messages).is_none()); + } + + /// A template that fails to render (here, one that calls `raise_exception`) + /// makes `encode_chat` return `None`, so the policy falls back to the raw + /// prompt-text path rather than failing the request. + #[test] + fn encode_chat_none_on_render_failure() { + let reg = TokenizerRegistry::default(); + reg.inner.insert( + "tiny".into(), + adapter::load("tests/fixtures/tiny_tokenizer.json").unwrap(), + ); + reg.attach_chat_template_for_test( + "tiny", + &serde_json::json!({ + "chat_template": "{{ raise_exception('nope') }}", + "bos_token": "", + }), + ); + assert!(reg.has_chat_encoder("tiny")); + let messages = serde_json::json!([{"role":"user","content":"hi"}]); + assert!( + reg.encode_chat("tiny", &messages).is_none(), + "a failing render must yield None so routing falls back to raw text" + ); + } + + #[test] + fn is_deepseek_v4_matches_v4_only() { + assert!(is_deepseek_v4("deepseek-ai/DeepSeek-V4-Flash")); + assert!(is_deepseek_v4("DeepSeek-V4-Pro")); + // Not V4-family models. + assert!(!is_deepseek_v4("deepseek-ai/DeepSeek-V3.2")); + assert!(!is_deepseek_v4("Qwen/Qwen3-0.6B")); + assert!(!is_deepseek_v4("tiny")); + } } diff --git a/experimental/sgl-router/tests/e2e/chat_completions/test_two_router_convergence.py b/experimental/sgl-router/tests/e2e/chat_completions/test_two_router_convergence.py index 34fe97075..efcd35f61 100644 --- a/experimental/sgl-router/tests/e2e/chat_completions/test_two_router_convergence.py +++ b/experimental/sgl-router/tests/e2e/chat_completions/test_two_router_convergence.py @@ -43,12 +43,15 @@ import time import httpx import pytest from infra.gateway import Gateway -from infra.model_pool import PASSTHROUGH_CHAT_TEMPLATE_PATH, spawn_worker +from infra.model_pool import spawn_worker from infra.model_specs import get_model_spec -# Disjoint prefixes — share no common opening text, so block 0 hashes -# differ from the first block onward and each worker's HashTree -# contribution is uniquely identifying. +# Disjoint prefixes — share no common content. Under the chat template both +# render with the same leading role header (``<|im_start|>user`` ...; Qwen3 has +# no BOS token), so the first block(s) may hash identically; the disjoint +# content then diverges +# well within the matched region, making each worker's HashTree contribution +# uniquely identifying. # # Length matters: each prefix must span ≥2 SGLang blocks at the default # block_size of 64 tokens so the worker actually emits BlockStored @@ -127,21 +130,12 @@ def _direct_warm(worker_url: str, model_id: str, prefix: str) -> None: worker to populate, so the two workers' HashTree state would no longer be uniquely identifying. - Token alignment with the router — ``cache_aware_zmq`` hashes - ``messages[*].content`` RAW (``cache_aware_zmq.rs::extract_prompt_text``) - using ``add_special_tokens=false``. By default SGLang's chat - endpoint would wrap ``prefix`` in the model's chat template before - tokenizing — adding role tags, end-of-turn markers, and a - generation prompt — and the resulting block hashes would never - match what the router computes from raw content. - - The test launches each worker with ``--chat-template - ``: a Jinja template that emits - only ``messages[*].content`` (the same shape the router extracts), - and which combines with Transformers' ``apply_chat_template( - tokenize=True, add_special_tokens=False)`` to produce the same - token stream the router will compute. So warm and route hash the - same blocks via the same endpoint. + Token alignment with the router — the workers run with the model's + real chat template (no override), so the engine caches blocks keyed + on chat-templated tokens (role markers + content + generation prompt). + ``cache_aware_zmq`` mirrors this: for a chat request on a model that + ships a chat template, it renders the same template and tokenizes the + result before hashing, so warm and route hash the same blocks. """ r = httpx.post( f"{worker_url}/v1/chat/completions", @@ -194,24 +188,22 @@ def test_two_routers_route_by_prefix_content( """ spec = get_model_spec("qwen3-0.6b") gpus = gpu_allocator.acquire(2) - # Passthrough chat template — see _direct_warm for the rationale. Both - # workers must run with the same template; otherwise their KV blocks - # would hash template-wrapped tokens while the router hashes raw - # content, and every lookup would miss the tree. - worker_chat_template_args = ["--chat-template", PASSTHROUGH_CHAT_TEMPLATE_PATH] + # Workers run with the model's REAL chat template (no override): the engine + # caches chat-templated tokens, and the router renders the same template + # (loaded from the model's tokenizer_config.json) before hashing. This + # exercises the production chat-template tokenization path, which aligns + # router query hashes with the engine's templated blocks. try: with ( spawn_worker( "qwen3-0.6b", gpu_ids=[gpus[0]], enable_kv_events=True, - extra_args=worker_chat_template_args, ) as worker_x, spawn_worker( "qwen3-0.6b", gpu_ids=[gpus[1]], enable_kv_events=True, - extra_args=worker_chat_template_args, ) as worker_y, Gateway() as router_a, Gateway() as router_b, diff --git a/experimental/sgl-router/tests/e2e/infra/model_pool.py b/experimental/sgl-router/tests/e2e/infra/model_pool.py index 2386bcbbc..c60bbee64 100644 --- a/experimental/sgl-router/tests/e2e/infra/model_pool.py +++ b/experimental/sgl-router/tests/e2e/infra/model_pool.py @@ -26,7 +26,6 @@ import socket import subprocess import time from dataclasses import dataclass, field -from pathlib import Path import httpx @@ -34,16 +33,6 @@ from .model_specs import get_model_spec logger = logging.getLogger(__name__) -# Passthrough Jinja chat template that emits ONLY `messages[*].content` -# joined with `\n` — matching the router's cache_aware_zmq prompt -# extraction. A worker launched with -# ``--chat-template `` tokenizes the -# raw content string, so its KV-block hashes align with what the -# router computes from the same chat-completions request. Test-only. -PASSTHROUGH_CHAT_TEMPLATE_PATH = str( - Path(__file__).parent / "passthrough_chat_template.jinja" -) - def _get_open_port() -> int: """Allocate an ephemeral TCP port in the range [20000, 55535]. diff --git a/experimental/sgl-router/tests/e2e/infra/passthrough_chat_template.jinja b/experimental/sgl-router/tests/e2e/infra/passthrough_chat_template.jinja deleted file mode 100644 index 8d542dd66..000000000 --- a/experimental/sgl-router/tests/e2e/infra/passthrough_chat_template.jinja +++ /dev/null @@ -1,13 +0,0 @@ -{#- - Passthrough chat template for cache-aware-zmq e2e tests. - - Emits ONLY `messages[*].content` joined with `\n` — no role markers, - no special tokens, no generation prompt. This is the SAME shape the - router's cache_aware_zmq policy produces in `extract_prompt_text`, - so a worker launched with `--chat-template ` tokenizes the - same string the router will tokenize for routing — making block - hashes align across worker KV cache and router HashTree. - - Use only for tests; not appropriate for any real chat workload. --#} -{{- messages | map(attribute='content') | join('\n') -}}