[router] Apply chat template before cache-aware hashing (fix overlap=0 on chat traffic) (#27386)
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Kangyan-Zhou
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
276c98c6cf
commit
21647f1f5d
@@ -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"
|
||||
|
||||
@@ -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<String> {
|
||||
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<String> {
|
||||
let v: serde_json::Value = serde_json::from_slice(body).ok()?;
|
||||
fn extract_prompt_text_from_value(v: &serde_json::Value) -> Option<String> {
|
||||
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<Vec<u32>> {
|
||||
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": "<s>",
|
||||
});
|
||||
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": "<s>",
|
||||
});
|
||||
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<TokenizerRegistry>,
|
||||
content: &str,
|
||||
) -> (CacheAwareZmqPolicy, Vec<Arc<Worker>>, 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": "<s>",
|
||||
}),
|
||||
);
|
||||
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": "<s>",
|
||||
}),
|
||||
);
|
||||
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() {
|
||||
|
||||
@@ -44,17 +44,11 @@ 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<std::path::PathBuf> {
|
||||
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(|| {
|
||||
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 \
|
||||
@@ -63,6 +57,57 @@ fn download_tokenizer_json(repo_id: &str) -> Result<std::path::PathBuf> {
|
||||
})
|
||||
}
|
||||
|
||||
/// 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<std::path::PathBuf> {
|
||||
use hf_hub::api::sync::ApiBuilder;
|
||||
let api = ApiBuilder::from_env()
|
||||
.build()
|
||||
.context("initialize HuggingFace Hub client")?;
|
||||
api.model(repo_id.to_string())
|
||||
.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<Option<serde_json::Value>> {
|
||||
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<Vec<u32>> {
|
||||
let enc = t.encode(text).context("encode")?;
|
||||
Ok(enc.token_ids().to_vec())
|
||||
|
||||
@@ -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<Option<Self>> {
|
||||
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<String> {
|
||||
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": "<jinja>"` — the common single-template case.
|
||||
/// - `"chat_template": [{"name": "default", "template": "<jinja>"}, ...]` —
|
||||
/// multi-template models; we take the entry named `default`, else the first.
|
||||
fn extract_chat_template(cfg: &serde_json::Value) -> Option<String> {
|
||||
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": "<tok>", ...}`) HuggingFace uses.
|
||||
fn extract_token_str(cfg: &serde_json::Value, key: &str) -> Option<String> {
|
||||
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<String, JinjaError> {
|
||||
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": "<s>", "eos_token": "</s>"});
|
||||
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": "<s>",
|
||||
"eos_token": "</s>",
|
||||
});
|
||||
let tmpl = ChatTemplate::from_tokenizer_config(&cfg).unwrap().unwrap();
|
||||
let out = tmpl.render(&messages()).unwrap();
|
||||
assert_eq!(
|
||||
out,
|
||||
"<s><|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": "<s>" });
|
||||
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": "<s>",
|
||||
});
|
||||
let tmpl = ChatTemplate::from_tokenizer_config(&cfg).unwrap().unwrap();
|
||||
assert!(tmpl
|
||||
.render(&messages())
|
||||
.unwrap()
|
||||
.starts_with("<s><|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": "<s>",
|
||||
});
|
||||
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": "<s>",
|
||||
});
|
||||
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": "<pad>",
|
||||
"unk_token": {"content": "<unk>"},
|
||||
});
|
||||
let tmpl = ChatTemplate::from_tokenizer_config(&cfg).unwrap().unwrap();
|
||||
assert_eq!(tmpl.render(&json!([])).unwrap(), "<pad>|<unk>");
|
||||
}
|
||||
|
||||
/// 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": "<s>",
|
||||
});
|
||||
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": "<s>",
|
||||
});
|
||||
let tmpl = ChatTemplate::from_tokenizer_config(&cfg).unwrap().unwrap();
|
||||
// Each iteration emits just "<role>\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");
|
||||
}
|
||||
}
|
||||
@@ -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｜> </think>`. 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 = "</think>";
|
||||
|
||||
/// 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｜> </think>`.
|
||||
#[test]
|
||||
fn single_user_turn() {
|
||||
let out = render_messages(&json!([{"role":"user","content":"ABCD"}]));
|
||||
assert_eq!(
|
||||
out,
|
||||
"<｜begin▁of▁sentence｜><｜User｜>ABCD<｜Assistant｜></think>"
|
||||
);
|
||||
}
|
||||
|
||||
/// 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｜></think>"
|
||||
);
|
||||
}
|
||||
|
||||
/// 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｜></think>A1<｜end▁of▁sentence｜><｜User｜>U2<｜Assistant｜></think>"
|
||||
);
|
||||
}
|
||||
|
||||
/// 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｜></think>"
|
||||
);
|
||||
}
|
||||
|
||||
/// 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｜></think>"
|
||||
);
|
||||
}
|
||||
|
||||
/// 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｜></think>"
|
||||
);
|
||||
}
|
||||
|
||||
/// 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｜></think>A1<｜end▁of▁sentence｜><｜User｜>U3<｜Assistant｜></think>"
|
||||
);
|
||||
}
|
||||
|
||||
/// 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｜></think>"
|
||||
);
|
||||
assert_eq!(
|
||||
render_messages(&json!([
|
||||
{"role":"developer","content":"D1"},
|
||||
{"role":"developer","content":"D2"}
|
||||
])),
|
||||
"<｜begin▁of▁sentence｜><｜User｜>D1<｜User｜>D2<｜Assistant｜></think>"
|
||||
);
|
||||
}
|
||||
|
||||
/// 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｜>");
|
||||
}
|
||||
}
|
||||
@@ -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<ChatTemplate>),
|
||||
/// 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<String> {
|
||||
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<String, Arc<Tokenizer>>,
|
||||
/// 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<String, Arc<ChatEncoderEntry>>,
|
||||
}
|
||||
|
||||
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<ChatEncoder> {
|
||||
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<Arc<Tokenizer>> {
|
||||
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<Vec<u32>> {
|
||||
// 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<String> {
|
||||
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":"<s>"}"#,
|
||||
)
|
||||
.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": "<s>",
|
||||
});
|
||||
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": "<s>",
|
||||
}),
|
||||
);
|
||||
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"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
<PASSTHROUGH_CHAT_TEMPLATE_PATH>``: 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,
|
||||
|
||||
@@ -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 <PASSTHROUGH_CHAT_TEMPLATE_PATH>`` 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].
|
||||
|
||||
@@ -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 <this file>` 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') -}}
|
||||
Reference in New Issue
Block a user