From 3ce7e2a29f01195fc4feceef00f62feadc37bfd9 Mon Sep 17 00:00:00 2001 From: Kan Wu Date: Thu, 17 Sep 2026 00:57:20 -0700 Subject: [PATCH] [sgl-router] Render chat prompts with dynamo-render (#38983) Co-authored-by: Claude Fable 5.1 --- experimental/sgl-router/Cargo.lock | 3 - experimental/sgl-router/Cargo.toml | 15 +- experimental/sgl-router/README.md | 42 ++ experimental/sgl-router/src/config/cli.rs | 19 + experimental/sgl-router/src/config/mod.rs | 1 + experimental/sgl-router/src/config/types.rs | 3 + .../sgl-router/src/policies/factory.rs | 1 + experimental/sgl-router/src/policies/mod.rs | 28 +- .../sgl-router/src/server/app_context.rs | 1 + .../sgl-router/src/server/routes/chat.rs | 205 +++++--- .../sgl-router/src/server/routes/models.rs | 1 + .../sgl-router/src/server/routes/tokenize.rs | 1 + .../sgl-router/src/tokenizer/adapter.rs | 5 - .../src/tokenizer/chat_formatter.rs | 488 ++++++++++++++++++ .../sgl-router/src/tokenizer/chat_template.rs | 365 ------------- experimental/sgl-router/src/tokenizer/dsv4.rs | 274 ---------- experimental/sgl-router/src/tokenizer/mod.rs | 238 +++++---- .../sgl-router/src/workers/manager.rs | 1 + .../tests/component/discovery/static_urls.rs | 1 + .../sgl-router/tests/proxy/bucket_routing.rs | 1 + .../sgl-router/tests/proxy/chat_routing.rs | 1 + .../tests/proxy/common/cache_aware_fixture.rs | 1 + .../sgl-router/tests/proxy/failover.rs | 1 + .../tests/proxy/graceful_shutdown.rs | 1 + .../tests/proxy/header_forwarding.rs | 1 + .../tests/proxy/pd_bootstrap_injection.rs | 1 + .../tests/proxy/pd_pool_isolation.rs | 1 + .../tests/proxy/pd_protocol_binding.rs | 1 + .../tests/proxy/roundrobin_input_ids.rs | 207 +++++++- .../tests/proxy/shared_prefill_admission.rs | 1 + .../tests/proxy/sticky_input_ids.rs | 1 + .../sgl-router/tests/proxy/sticky_routing.rs | 1 + .../sgl-router/tests/proxy/timeout.rs | 1 + 33 files changed, 1048 insertions(+), 864 deletions(-) create mode 100644 experimental/sgl-router/src/tokenizer/chat_formatter.rs delete mode 100644 experimental/sgl-router/src/tokenizer/chat_template.rs delete mode 100644 experimental/sgl-router/src/tokenizer/dsv4.rs diff --git a/experimental/sgl-router/Cargo.lock b/experimental/sgl-router/Cargo.lock index 592987e60..1e0d4e238 100644 --- a/experimental/sgl-router/Cargo.lock +++ b/experimental/sgl-router/Cargo.lock @@ -2072,7 +2072,6 @@ dependencies = [ "memo-map", "percent-encoding", "serde", - "serde_json", ] [[package]] @@ -3285,7 +3284,6 @@ dependencies = [ "anyhow", "axum", "bytes", - "chrono", "clap", "criterion", "dashmap", @@ -3301,7 +3299,6 @@ dependencies = [ "k8s-openapi", "kube", "minijinja", - "minijinja-contrib", "parking_lot", "rand 0.8.8", "reqwest", diff --git a/experimental/sgl-router/Cargo.toml b/experimental/sgl-router/Cargo.toml index 3f27aec9a..ff35f0266 100644 --- a/experimental/sgl-router/Cargo.toml +++ b/experimental/sgl-router/Cargo.toml @@ -24,17 +24,12 @@ unused_qualifications = "warn" [dependencies] # Pin Dynamo versions and commit Cargo.lock so builds are reproducible. dynamo-tokenizers = "=1.8.1" -# Prepare the Dynamo renderer dependency for the following migration. +# Renders the model's chat template (HF Jinja, or Dynamo's built-in encoder for +# template-less models like DeepSeek-V4) so cache-aware routing hashes the same +# tokens the engine caches. dynamo-renderer = "=5.1.2" -# Chat-template rendering for cache-aware routing, retained until the -# renderer migration replaces it: the engine caches tokens AFTER applying the -# model's chat template, so the router must render the same template before -# hashing or its token_ids diverge from the engine's stored blocks. `pycompat` -# supplies the Python str/dict methods HF templates call; `chrono` backs -# `strftime_now`. -minijinja = { version = "2.24", features = ["loop_controls", "json"] } -minijinja-contrib = { version = "2", features = ["pycompat"] } -chrono = { version = "0.4", default-features = false, features = ["clock"] } +# `OAIChatLikeRequest` speaks minijinja values. +minijinja = "2.24" # Async runtime + http tokio = { version = "1.42", features = ["full"] } diff --git a/experimental/sgl-router/README.md b/experimental/sgl-router/README.md index d1b8eeefe..c026dbecd 100644 --- a/experimental/sgl-router/README.md +++ b/experimental/sgl-router/README.md @@ -169,6 +169,48 @@ queue slot and no engine round-trip — the `sampling_contract_violation` code and per-parameter counter, bands, and one contract applied at a shared ingress across engines whose own flags the router operator may not control. +## Chat rendering + +The router renders chat requests with dynamo-render (`dynamo-renderer`): the model's +HF Jinja template from `tokenizer_config.json` or a sibling +`chat_template.jinja`, or dynamo-render's built-in DeepSeek encoder (V4 family, V3.2) +for template-less models. Cache-aware routing hashes the rendered tokens so its +prefix queries match the blocks the engine caches. Models the engine encodes in +code but dynamo-render cannot tokenize here (Inkling, Kimi K3) route via raw prompt +text, as does any model whose template fails to load or render. + +Plain text chat requests (string `content`, no tools, no template kwargs or +reasoning controls or historical `reasoning_content`, no assistant continuation, +no consecutive users or non-leading system turns) additionally forward the +rendered tokens to the engine as `input_ids`, retaining the original messages, +so the engine skips re-tokenizing. Every other request shape is rendered for +routing only: the router renders with dynamo-render and does not replicate +SGLang's request normalization, so forwarding is enabled shape by shape as +parity is verified. Use matching model files on the router and workers; worker +template overrides and default kwargs are not observable from the request. + +Set `--disable-input-ids-forwarding` for this router's model when worker-side +rendering has not been verified to match. This disables router-generated IDs +for every routing policy; cache-aware routing still renders and tokenizes +locally, and the original messages reach the workers for engine processing. +Caller-supplied `input_ids` remain caller-owned and pass through unchanged. + +Forwarding logs its assumptions at startup. In particular, disable it for +`SGLANG_DEFAULT_THINKING=true`, a non-default `SGLANG_DSV4_REASONING_EFFORT`, +worker parser overrides such as `--tool-call-parser deepseekv32` that select a +native encoder over a shipped template, or conversation templates with stop +strings (the engine's `input_ids` path skips those template stops). These worker +settings are not inferred from the router's environment. Disabling forwarding preserves +engine behavior but does not establish parity for local routing hashes. + +Also set `--disable-input-ids-forwarding` for array-only templates: Dynamo may wrap +string content into arrays differently from the worker. Dynamo 5.1.2 does not expose +its conversion flag, so the router cannot automatically block these templates. +Detailed content-format parity coverage follows in #39133. + +The Dynamo crates are pinned exactly and `Cargo.lock` is committed; CI builds +with `--locked`, so rendered bytes cannot change without a reviewed diff. + ## HTTP/2 There is nothing to configure. The router negotiates per connection inbound and diff --git a/experimental/sgl-router/src/config/cli.rs b/experimental/sgl-router/src/config/cli.rs index 66de72441..9c463d163 100644 --- a/experimental/sgl-router/src/config/cli.rs +++ b/experimental/sgl-router/src/config/cli.rs @@ -71,6 +71,11 @@ pub struct Cli { /// as the repo id (download honors `HF_TOKEN` / `HF_HOME`). #[arg(long)] pub tokenizer_path: Option, + /// Disable router-generated input_ids for this model. Workers tokenize messages + /// themselves; cache-aware routing still renders locally. Use for worker-only + /// thinking/effort defaults, parser/template overrides, or template stop strings. + #[arg(long)] + pub disable_input_ids_forwarding: bool, /// Routing policy. #[arg(long, value_enum, default_value = "round_robin")] pub policy: PolicyKind, @@ -702,6 +707,7 @@ impl Cli { // HuggingFace repo id) when --tokenizer-path is omitted. tokenizer_path: self.tokenizer_path.unwrap_or_else(|| self.model_id.clone()), id: self.model_id, + disable_input_ids_forwarding: self.disable_input_ids_forwarding, policy: self.policy, decode_policy: self.decode_policy, bucket_config, @@ -938,6 +944,19 @@ mod tests { assert_eq!(c.model.tokenizer_path, "/models/qwen3/tokenizer.json"); } + #[test] + fn input_ids_forwarding_can_be_disabled_for_the_model() { + let defaults = into_config_owned(with_model(&["--worker-urls", "http://x:30000"])).unwrap(); + assert!(!defaults.model.disable_input_ids_forwarding); + let disabled = into_config_owned(with_model(&[ + "--worker-urls", + "http://x:30000", + "--disable-input-ids-forwarding", + ])) + .unwrap(); + assert!(disabled.model.disable_input_ids_forwarding); + } + #[test] fn static_urls_backend() { let c = into_config_owned(with_model(&[ diff --git a/experimental/sgl-router/src/config/mod.rs b/experimental/sgl-router/src/config/mod.rs index fea8f1c2c..9d20f82e4 100644 --- a/experimental/sgl-router/src/config/mod.rs +++ b/experimental/sgl-router/src/config/mod.rs @@ -279,6 +279,7 @@ mod tests { model: ModelConfig { id: model_id.into(), tokenizer_path: "/tmp/tok.json".into(), + disable_input_ids_forwarding: false, policy: PolicyKind::RoundRobin, decode_policy: DecodePolicyKind::PowerOfTwo, bucket_config: None, diff --git a/experimental/sgl-router/src/config/types.rs b/experimental/sgl-router/src/config/types.rs index 5860cf87f..dfa02f28d 100644 --- a/experimental/sgl-router/src/config/types.rs +++ b/experimental/sgl-router/src/config/types.rs @@ -350,6 +350,9 @@ pub struct ModelConfig { /// id (downloaded on demand). Defaults to `id` when `--tokenizer-path` /// is omitted. Resolved by [`crate::tokenizer::adapter::load`]. pub tokenizer_path: String, + /// Disable router-generated input IDs for this model; keep routing tokenization. + /// Use when workers have rendering defaults or template stops the router cannot see. + pub disable_input_ids_forwarding: bool, pub policy: PolicyKind, /// Selection policy for the decode pool. pub decode_policy: DecodePolicyKind, diff --git a/experimental/sgl-router/src/policies/factory.rs b/experimental/sgl-router/src/policies/factory.rs index d78c3bb76..132755cb9 100644 --- a/experimental/sgl-router/src/policies/factory.rs +++ b/experimental/sgl-router/src/policies/factory.rs @@ -337,6 +337,7 @@ mod tests { model: ModelConfig { id: id.into(), tokenizer_path: "/tmp/x".into(), + disable_input_ids_forwarding: false, policy, decode_policy: Default::default(), bucket_config: None, diff --git a/experimental/sgl-router/src/policies/mod.rs b/experimental/sgl-router/src/policies/mod.rs index e473584b1..cf54c2862 100644 --- a/experimental/sgl-router/src/policies/mod.rs +++ b/experimental/sgl-router/src/policies/mod.rs @@ -52,14 +52,14 @@ pub fn request_tokens_for( model_id: &ModelId, value: &serde_json::Value, ) -> Option { - if tokenizers.has_chat_formatter(&model_id.0) { - if let Some(messages) = value.get("messages").filter(|m| m.is_array()) { - if let Some(ids) = tokenizers.encode_chat(&model_id.0, messages) { - return Some(RequestTokens { - ids, - rendered_from_chat: true, - }); - } + if tokenizers.has_chat_formatter(&model_id.0) + && value.get("messages").is_some_and(|m| m.is_array()) + { + if let Some(ids) = tokenizers.encode_chat(&model_id.0, value) { + return Some(RequestTokens { + ids, + rendered_from_chat: true, + }); } } let text = extract_prompt_text_from_value(value)?; @@ -510,14 +510,10 @@ pub trait Policy: Send + Sync + std::fmt::Debug { } /// Whether this policy's routing decision needs request tokens (i.e. - /// it routes by prompt prefix). Ingress tokenization itself is no longer - /// gated on this — that is a model property (`has_chat_formatter`) decided at - /// ingress via [`request_tokens_for`]. This flag is the EXTRA gate that - /// keeps the cache-aware policy's RAW-prompt routing path alive: a - /// cache-aware model with no chat formatter still wants its `/v1/completions` - /// /`text` prompt tokenized for tree matching, which `has_chat_formatter` - /// alone would not trigger. Default `false` for load-only and sticky - /// routes; only the cache-aware policy overrides it. + /// it routes by prompt prefix). This keeps routing tokenization active when + /// generated input-ID forwarding is disabled or no chat formatter exists. + /// Models with forwarding enabled also tokenize independently of this flag. + /// Default `false` for load-only and sticky routes. fn needs_request_tokens(&self) -> bool { false } diff --git a/experimental/sgl-router/src/server/app_context.rs b/experimental/sgl-router/src/server/app_context.rs index 952ed1a68..599d271ed 100644 --- a/experimental/sgl-router/src/server/app_context.rs +++ b/experimental/sgl-router/src/server/app_context.rs @@ -169,6 +169,7 @@ impl AppContext { model: crate::config::ModelConfig { id: "stub-model".into(), tokenizer_path: "stub".into(), + disable_input_ids_forwarding: false, policy: crate::config::PolicyKind::RoundRobin, decode_policy: Default::default(), bucket_config: None, diff --git a/experimental/sgl-router/src/server/routes/chat.rs b/experimental/sgl-router/src/server/routes/chat.rs index b05e04580..2ac5c052d 100644 --- a/experimental/sgl-router/src/server/routes/chat.rs +++ b/experimental/sgl-router/src/server/routes/chat.rs @@ -507,7 +507,7 @@ pub async fn chat_completions( // MODEL (does it have a chat formatter so the router can produce // engine-equivalent tokens?), not of how we pick the worker. Two gates: // - // * `has_chat_formatter` → a chat request on this model yields + // * Forwarding is enabled and `has_chat_formatter` -> a chat request yields // engine-equivalent ids we can forward as `input_ids` so the engine // skips re-tokenizing. This enables the offload for EVERY policy — // sticky and round-robin included — not just cache-aware. @@ -523,8 +523,10 @@ pub async fn chat_completions( // body. When parsed, this single value is reused for the routing // tokenization and the outgoing-body injection below (and PD bootstrap // injection). `parse_probe` already validated the object shape. + let can_forward_input_ids = !ctx.config.model.disable_input_ids_forwarding + && ctx.tokenizers.has_chat_formatter(&model_str); let want_tokens = should_tokenize_request( - ctx.tokenizers.has_chat_formatter(&model_str), + can_forward_input_ids, policy.needs_request_tokens(), ctx.bucket_selector.is_enabled(), ); @@ -798,8 +800,9 @@ pub async fn chat_completions( // Forward the router-computed tokens to the engine as `input_ids` so it // skips re-tokenizing the same prompt — but only when they are - // engine-equivalent (chat-formatter path) AND the request contains nothing - // the router's formatter didn't replicate (see `input_ids_safe_to_forward`). + // enabled for this model, engine-equivalent (chat-formatter path), and + // the request has no unreplicated rendering controls (see + // `input_ids_safe_to_forward`). // Otherwise omit them and the engine tokenizes from `messages` as usual — // a transparent, always-correct fallback (`messages` are always retained // in the forwarded body). `forward_input_ids` is `Some` only when @@ -807,7 +810,9 @@ pub async fn chat_completions( // predicate always has a parsed body to inspect. let forward_input_ids: Option<&[u32]> = match (request_tokens.as_ref(), request_value.as_ref()) { - (Some(t), Some(v)) if t.rendered_from_chat && input_ids_safe_to_forward(v) => { + (Some(t), Some(v)) + if can_forward_input_ids && t.rendered_from_chat && input_ids_safe_to_forward(v) => + { Some(t.ids.as_slice()) } _ => None, @@ -819,7 +824,7 @@ pub async fn chat_completions( // `ingress_tokenize_offload_failed`); successful forwards and expected // omissions are not problems. if ingress_tokenize_offload_failed( - ctx.tokenizers.has_chat_formatter(&model_str), + can_forward_input_ids, request_value.as_ref(), request_tokens.as_ref(), ) { @@ -1198,11 +1203,11 @@ fn parse_optional_positive_f64_header( } fn should_tokenize_request( - has_chat_formatter: bool, + can_forward_input_ids: bool, policy_needs_request_tokens: bool, bucket_enabled: bool, ) -> bool { - has_chat_formatter || policy_needs_request_tokens || bucket_enabled + can_forward_input_ids || policy_needs_request_tokens || bucket_enabled } /// Estimate prefill-token count from the raw request body for use as @@ -1401,51 +1406,30 @@ fn build_outgoing_body( Ok(Bytes::from(bytes)) } -/// Whether the router's `input_ids` may be forwarded for this request. +/// Forward generated IDs only for request shapes verified against the engine. +/// The engine uses `input_ids` verbatim, bypassing its chat-template processing. /// -/// We forward only when the engine, fed `input_ids`, would have produced the -/// SAME prompt the router tokenized. When `input_ids` is present the engine -/// uses it verbatim and ignores everything that would otherwise steer its -/// `messages`-side tokenization (only stop tokens / tool-call constraint are -/// still taken from `messages`). So any request field that changes that -/// tokenization but which the router's chat formatter does not replicate makes -/// the forwarded ids wrong. This predicate is conservative by construction — -/// any such signal returns `false` and the engine tokenizes from `messages` -/// (always correct). +/// Exclude requests that may render differently with dynamo-render: +/// - Non-leading system turns or consecutive users, which strict templates rewrite. +/// - Historical `reasoning_content`, which may be injected into message content. +/// - Tools and tool-call history, which the engine merges and normalizes +/// before rendering. +/// - Non-string or missing content, which the engine flattens or blanks. +/// - Template overrides, kwargs, reasoning controls, or task selection. +/// - Assistant continuations, whose final turn the engine handles separately. /// -/// Replicated-and-safe: plain text `messages` with a string `content`. -/// Not replicated → omit: -/// * `tools` / `functions` — the formatter doesn't render tool schemas. -/// * non-string or missing `content` (arrays, `null`): the engine normalizes -/// these before rendering; the router's formatter renders them verbatim. -/// * `chat_template` — an OpenAI-compatible per-request template override -/// (e.g. vLLM); the router renders with the model's default template, so a -/// custom one would diverge. (SGLang ignores it today, but block it so the -/// offload stays correct across engines / future versions.) -/// * `chat_template_kwargs` (carries `enable_thinking`/`thinking`), -/// `reasoning` / `reasoning_effort`, `task` — thinking/mode toggles the -/// formatter renders in the engine's default mode only. -/// * `continue_final_message: true`, or a trailing `assistant` message — the -/// engine rewrites/strips the final assistant turn; the formatter renders it -/// verbatim. -/// -/// NOTE: the router's chat formatter renders in the engine's default -/// (non-thinking) mode. Current sglang derives thinking from the request -/// (`chat_template_kwargs`), which this guard already omits, so a plain request -/// the router rendered matches the engine. The only way to diverge is an engine -/// build that applies a non-default thinking mode the router can't observe from -/// the request — the same router↔engine tokenization-parity assumption that -/// cache-aware routing already depends on. The same assumption covers -/// `add_special_tokens`: the router renders specials via the chat template, which -/// matches the engine on tokenizers that auto-add them (the common case); a -/// tokenizer that does not would diverge by a leading special, again undetectable -/// from the request. +/// Matching model files and engine defaults are still required. Worker template +/// overrides and default kwargs cannot be inferred from the request. +/// `--disable-input-ids-forwarding` gates forwarding separately for such fleets. fn input_ids_safe_to_forward(value: &serde_json::Value) -> bool { - if request_has_tools(value) || request_has_non_text_content(value) { + if request_has_tools(value) + || request_has_non_text_content(value) + || request_has_reasoning_content(value) + || request_has_role_rewrites(value) + { return false; } - // Fields that steer the engine's template tokenization but which the - // router's formatter does not thread through. + // Request controls whose rendering has not been verified against the engine. for key in [ "chat_template", "chat_template_kwargs", @@ -1469,15 +1453,15 @@ fn input_ids_safe_to_forward(value: &serde_json::Value) -> bool { /// Whether to increment `sgl_router_ingress_tokenize_errors_total`. /// -/// Count chats with a configured formatter that pass the forwarding guard +/// Count chats with forwarding enabled that pass the forwarding guard /// but lack chat-rendered tokens. Excluded requests are expected fallbacks, /// even when rendering fails. fn ingress_tokenize_offload_failed( - has_chat_formatter: bool, + can_forward_input_ids: bool, request_value: Option<&serde_json::Value>, request_tokens: Option<&RequestTokens>, ) -> bool { - if !has_chat_formatter { + if !can_forward_input_ids { return false; } let chat_request = request_value.is_some_and(|v| { @@ -1501,23 +1485,61 @@ fn last_message_is_assistant(value: &serde_json::Value) -> bool { == Some("assistant") } -/// Whether the request carries tool / function definitions. The router's chat -/// formatter renders only `messages`, so its `input_ids` would omit the tool -/// schemas the engine's template injects into the prompt — the caller must let -/// the engine tokenize these itself. +/// Tool schemas and tool-call history require engine normalization before +/// rendering: the engine merges message-level `tools` into the template's tools +/// and parses `tool_calls` arguments; dynamo-render does neither the same way. fn request_has_tools(value: &serde_json::Value) -> bool { - let nonempty = |key: &str| { - value.get(key).is_some_and(|v| match v { - serde_json::Value::Array(a) => !a.is_empty(), - serde_json::Value::Null => false, - _ => true, - }) + let nonempty = |v: &serde_json::Value| match v { + serde_json::Value::Array(a) => !a.is_empty(), + serde_json::Value::Null => false, + _ => true, }; - nonempty("tools") || nonempty("functions") + if ["tools", "functions"] + .iter() + .any(|key| value.get(key).is_some_and(nonempty)) + { + return true; + } + value + .get("messages") + .and_then(|m| m.as_array()) + .is_some_and(|messages| { + messages.iter().any(|message| { + message["role"] == "tool" + || ["tools", "tool_calls", "function_call"] + .iter() + .any(|key| message.get(key).is_some_and(nonempty)) + }) + }) +} + +/// dynamo-render may inject historical reasoning into content the engine leaves unchanged. +fn request_has_reasoning_content(value: &serde_json::Value) -> bool { + value + .get("messages") + .and_then(|messages| messages.as_array()) + .is_some_and(|messages| { + messages.iter().any(|message| { + message + .get("reasoning_content") + .is_some_and(|v| !v.is_null()) + }) + }) +} + +/// Message orders dynamo-render may rewrite for strict templates. +fn request_has_role_rewrites(value: &serde_json::Value) -> bool { + let Some(messages) = value.get("messages").and_then(|v| v.as_array()) else { + return false; + }; + messages.iter().skip(1).any(|m| m["role"] == "system") + || messages + .windows(2) + .any(|pair| pair[0]["role"] == "user" && pair[1]["role"] == "user") } /// Detect non-string or missing content, which requires engine tokenization: -/// the engine normalizes arrays and nulls differently from the router's formatter. +/// the engine normalizes arrays and nulls differently from dynamo-render. fn request_has_non_text_content(value: &serde_json::Value) -> bool { value .get("messages") @@ -1815,8 +1837,8 @@ mod tests { ); } - /// Tool / function requests are detected so the caller omits `input_ids` - /// (the router's formatter doesn't render tools). + /// Tool schemas and tool-call history are detected so the caller omits + /// `input_ids`; empty lists and nulls are not tools. #[test] fn request_has_tools_detects_tools_and_functions() { assert!(request_has_tools( @@ -1827,6 +1849,14 @@ mod tests { )); assert!(!request_has_tools(&serde_json::json!({"tools":[]}))); assert!(!request_has_tools(&serde_json::json!({"messages":[]}))); + for message in [ + serde_json::json!({"role":"system","content":"s","tools":[{"type":"function"}]}), + serde_json::json!({"role":"assistant","content":"","tool_calls":[{"function":{"name":"f","arguments":"{}"}}]}), + ] { + assert!(request_has_tools( + &serde_json::json!({"messages":[message]}) + )); + } } /// Arrays, nulls, and missing content block `input_ids` forwarding. @@ -1860,6 +1890,47 @@ mod tests { }))); } + #[test] + fn reasoning_history_is_an_expected_forwarding_omission() { + let mut value = serde_json::json!({"messages": [ + {"role":"user", "content":"hi"}, + {"role":"assistant", "content":"answer", "reasoning_content":"prior reasoning"}, + {"role":"user", "content":"next"} + ]}); + assert!(!input_ids_safe_to_forward(&value)); + assert!(!ingress_tokenize_offload_failed(true, Some(&value), None)); + value["messages"][1]["reasoning_content"] = serde_json::Value::Null; + assert!(input_ids_safe_to_forward(&value)); + value["messages"][1] + .as_object_mut() + .unwrap() + .remove("reasoning_content"); + assert!(input_ids_safe_to_forward(&value)); + } + + #[test] + fn role_rewrites_are_expected_forwarding_omissions() { + for roles in [ + vec!["user", "user"], + vec!["system", "system", "user"], + vec!["user", "assistant", "system", "user"], + ] { + let messages: Vec<_> = roles + .iter() + .map(|role| serde_json::json!({"role": role, "content": "text"})) + .collect(); + let value = serde_json::json!({"messages": messages}); + assert!(!input_ids_safe_to_forward(&value), "{roles:?}"); + assert!(!ingress_tokenize_offload_failed(true, Some(&value), None)); + } + assert!(input_ids_safe_to_forward(&serde_json::json!({"messages": [ + {"role": "system", "content": "instructions"}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"role": "user", "content": "next"} + ]}))); + } + /// Every field the engine honors on the `messages` path but which the /// router's formatter does not replicate must block forwarding — otherwise /// the engine uses the router's ids verbatim and silently runs a different @@ -1967,7 +2038,7 @@ mod tests { )); } - /// Non-chat-formatter models never expected the offload → not a failure even + /// Non-chat-formatter models never expected the offload -> not a failure even /// with no tokens. #[test] fn offload_failed_false_without_chat_formatter() { @@ -1975,7 +2046,7 @@ mod tests { assert!(!ingress_tokenize_offload_failed(false, Some(&value), None)); } - /// A non-chat (no `messages`) request on a chat-formatter model — e.g. + /// A non-chat (no `messages`) request on a chat-formatter model, e.g. /// `/v1/completions` `prompt` — never expected the chat-encode offload, so /// the absence of engine-equivalent ids is not a failure. #[test] diff --git a/experimental/sgl-router/src/server/routes/models.rs b/experimental/sgl-router/src/server/routes/models.rs index dedcdb79d..8b1a2ecf7 100644 --- a/experimental/sgl-router/src/server/routes/models.rs +++ b/experimental/sgl-router/src/server/routes/models.rs @@ -50,6 +50,7 @@ mod tests { ctx.config.model = crate::config::ModelConfig { id: "qwen3".into(), tokenizer_path: "x".into(), + disable_input_ids_forwarding: false, policy: PolicyKind::RoundRobin, decode_policy: Default::default(), bucket_config: None, diff --git a/experimental/sgl-router/src/server/routes/tokenize.rs b/experimental/sgl-router/src/server/routes/tokenize.rs index 9a27f8f72..46131aaf7 100644 --- a/experimental/sgl-router/src/server/routes/tokenize.rs +++ b/experimental/sgl-router/src/server/routes/tokenize.rs @@ -118,6 +118,7 @@ mod tests { model: crate::config::ModelConfig { id: "tiny".into(), tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(), + disable_input_ids_forwarding: false, policy: PolicyKind::RoundRobin, decode_policy: Default::default(), bucket_config: None, diff --git a/experimental/sgl-router/src/tokenizer/adapter.rs b/experimental/sgl-router/src/tokenizer/adapter.rs index 39800d3f0..ba2639831 100644 --- a/experimental/sgl-router/src/tokenizer/adapter.rs +++ b/experimental/sgl-router/src/tokenizer/adapter.rs @@ -161,11 +161,6 @@ impl ModelFiles { } } -/// Load the sibling tokenizer config for the current chat formatter. -pub fn load_tokenizer_config(source: &str) -> Result> { - ModelFiles::open(source).json("tokenizer_config.json") -} - pub fn encode(t: &Tokenizer, text: &str) -> Result> { let enc = t.encode(text).context("encode")?; Ok(enc.token_ids().to_vec()) diff --git a/experimental/sgl-router/src/tokenizer/chat_formatter.rs b/experimental/sgl-router/src/tokenizer/chat_formatter.rs new file mode 100644 index 000000000..e36b60da3 --- /dev/null +++ b/experimental/sgl-router/src/tokenizer/chat_formatter.rs @@ -0,0 +1,488 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Chat rendering via dynamo-render for cache-aware routing and input ID forwarding. +//! +//! Engine-specific request normalization is not replicated here. The forwarding +//! guard omits IDs for request shapes whose engine rendering has not been verified. + +use std::collections::HashMap; +use std::sync::Arc; + +use anyhow::{Context, Result}; +use dynamo_renderer::{ + deepseek_formatter_for, may_be_fix_tool_schema, ChatTemplate, ContextMixins, + OAIChatLikeRequest, OAIPromptFormatter, PromptFormatter, +}; +use minijinja::Value; +use serde_json::Value as JsonValue; + +pub type ChatTemplateKwargs = HashMap; + +// HF special_tokens_map names. dynamo-render supplies bos/eos/unk; the rest +// are passed as template context defaults. +const SPECIAL_TOKEN_KEYS: [&str; 7] = [ + "bos_token", + "eos_token", + "unk_token", + "sep_token", + "pad_token", + "cls_token", + "mask_token", +]; + +/// Renders and tokenizes chat requests through dynamo-render. +pub struct ChatFormatter { + formatter: Arc, + /// Template context defaults; request `chat_template_kwargs` override them. + defaults: ChatTemplateKwargs, +} + +impl ChatFormatter { + /// Load model files and select a template or native formatter from dynamo-render. + pub fn load(model_id: &str, tokenizer_path: &str) -> Result> { + let files = super::adapter::ModelFiles::open(tokenizer_path); + let model_type = files + .json("config.json")? + .and_then(|cfg| cfg["model_type"].as_str().map(str::to_owned)); + match model_type.as_deref() { + // These require tokenization paths not yet supported by this adapter. + Some("inkling_mm_model" | "kimi_k3") => return Ok(None), + Some(t) if t.starts_with("deepseek_v4") => { + return Ok(Self::deepseek_native(model_type.as_deref(), model_id)); + } + _ => {} + } + let cfg = files + .json("tokenizer_config.json")? + .unwrap_or_else(|| serde_json::json!({})); + let jinja = files.text("chat_template.jinja")?; + Ok(Self::from_tokenizer_config(cfg, jinja.as_deref())? + .or_else(|| Self::deepseek_native(model_type.as_deref(), model_id))) + } + + /// HF Jinja template from `tokenizer_config.json`, overridden by a sibling + /// `chat_template.jinja` when present (transformers' precedence); `Ok(None)` + /// when the model ships neither. + pub fn from_tokenizer_config( + mut cfg: JsonValue, + chat_template_jinja: Option<&str>, + ) -> Result> { + let config = cfg + .as_object_mut() + .context("tokenizer_config.json must be an object")?; + // Tokenizer-only settings have types the renderer's config does not support. + config.retain(|key, _| { + key == "chat_template" + || SPECIAL_TOKEN_KEYS.contains(&key.as_str()) + || key == "additional_special_tokens" + }); + if let Some(template) = chat_template_jinja { + cfg["chat_template"] = template.into(); + } + // HuggingFace supplies None when no retrieval documents are present. + let mut defaults = HashMap::from([("documents".into(), JsonValue::Null)]); + for key in SPECIAL_TOKEN_KEYS { + // Convert HF AddedToken objects to strings for dynamo-render. + if let Some(content) = added_token_content(&cfg[key]) { + cfg[key] = content.into(); + } + if matches!(key, "bos_token" | "eos_token" | "unk_token") { + // Missing tokens render as `None` in dynamo-render; HF uses "". + // These config values take precedence over kwargs. + if cfg[key].is_null() { + cfg[key] = "".into(); + } + } else { + // Never reach the template except through kwargs. + defaults.insert(key.to_owned(), cfg[key].as_str().unwrap_or_default().into()); + } + } + if let Some(extra) = cfg["additional_special_tokens"].as_array() { + let extra: Vec = extra + .iter() + .map(|t| added_token_content(t).map_or_else(|| t.clone(), Into::into)) + .collect(); + cfg["additional_special_tokens"] = extra.clone().into(); + defaults.insert("additional_special_tokens".into(), extra.into()); + } + // HF's `[{name, template}]` list form -> dynamo-render's `[{name: template}]`. + for entry in cfg["chat_template"].as_array_mut().into_iter().flatten() { + if let (Some(name), Some(template)) = + (entry["name"].as_str(), entry["template"].as_str()) + { + *entry = serde_json::json!({ name: template }); + } + } + let template: ChatTemplate = + serde_json::from_value(cfg).context("parse tokenizer_config.json")?; + if template.chat_template.is_none() { + return Ok(None); + } + let PromptFormatter::OAI(formatter) = + PromptFormatter::from_parts(template, ContextMixins::default(), true) + .context("compile chat template")?; + Ok(Some(Self { + formatter, + defaults, + })) + } + + /// dynamo-render's code-based DeepSeek encoders, for V4 (including variants + /// such as V4.1) and V3.2 non-Exp: the only built-in formatters verified + /// against the engine. `model_type` (from `config.json`) is authoritative; + /// the model id's last path segment is the fallback. + pub fn deepseek_native(model_type: Option<&str>, model_id: &str) -> Option { + let name = model_id + .rsplit('/') + .next() + .unwrap_or(model_id) + .to_lowercase(); + // The engine treats every `deepseek_v4*` variant (e.g. V4.1) as V4. + let model_type = model_type.map(str::to_lowercase).map(|t| { + if t.starts_with("deepseek_v4") { + "deepseek_v4".into() + } else { + t + } + }); + let PromptFormatter::OAI(formatter) = deepseek_formatter_for(&model_type, &name)?; + // Engine defaults: chat mode (`SGLANG_DEFAULT_THINKING=false`) and no + // reasoning-effort preamble; dynamo-render defaults to thinking at high effort. + let defaults = HashMap::from([ + ("thinking".into(), false.into()), + ("reasoning_effort".into(), "low".into()), + ]); + Some(Self { + formatter, + defaults, + }) + } + + fn template_kwargs(&self, request: &JsonValue) -> Result { + let mut kwargs: ChatTemplateKwargs = match request.get("chat_template_kwargs") { + None | Some(JsonValue::Null) => ChatTemplateKwargs::new(), + Some(v) => serde_json::from_value(v.clone()).context("chat_template_kwargs")?, + }; + for (key, value) in &self.defaults { + kwargs.entry(key.clone()).or_insert_with(|| value.clone()); + } + Ok(kwargs) + } + + /// Render the prompt text dynamo-render produces for `request`. + pub fn render(&self, request: &JsonValue) -> Result { + anyhow::ensure!(request["messages"].is_array(), "messages must be an array"); + let kwargs = self.template_kwargs(request)?; + self.formatter + .render(&ChatRequest { request, kwargs }) + .context("render chat template") + } + + pub fn encode( + &self, + tokenizer: &dynamo_tokenizers::Tokenizer, + request: &JsonValue, + ) -> Result> { + super::adapter::encode(tokenizer, &self.render(request)?) + } +} + +/// `content` of an HF `AddedToken` object (`{"content": "", "lstrip": ...}`). +fn added_token_content(token: &JsonValue) -> Option { + token + .as_object()? + .get("content")? + .as_str() + .map(str::to_owned) +} + +struct ChatRequest<'a> { + request: &'a JsonValue, + kwargs: ChatTemplateKwargs, +} + +/// Mirrors dynamo-render's own impl for its wire type: request fields pass through. +impl OAIChatLikeRequest for ChatRequest<'_> { + fn model(&self) -> String { + self.request["model"] + .as_str() + .unwrap_or_default() + .to_owned() + } + fn messages(&self) -> Value { + Value::from_serialize(&self.request["messages"]) + } + fn tools(&self) -> Option { + let tools = self.request.get("tools")?; + // HF and the engine treat an empty list as "no tools"; dynamo-render's schema + // fixer would hand the template `[]`, which tools-branching templates + // render as a tool preamble. + if tools.as_array().is_none_or(|t| t.is_empty()) { + return None; + } + may_be_fix_tool_schema(tools.clone()) + } + fn tool_choice(&self) -> Option { + self.request.get("tool_choice").map(Value::from_serialize) + } + fn reasoning_effort(&self) -> Option { + self.request + .get("reasoning_effort") + .map(Value::from_serialize) + } + /// Withheld: the engine enforces `response_format` by constrained decoding + /// and never renders it, while dynamo-render's DeepSeek formatters would + /// append a "## Response Format" schema preamble to the system turn. + fn response_format(&self) -> Option { + None + } + fn should_add_generation_prompt(&self) -> bool { + true + } + fn chat_template_args(&self) -> Option<&ChatTemplateKwargs> { + Some(&self.kwargs) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + 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 jinja(cfg: JsonValue) -> ChatFormatter { + ChatFormatter::from_tokenizer_config(cfg, None) + .unwrap() + .expect("config has a chat_template") + } + + fn request(messages: JsonValue) -> JsonValue { + json!({"model": "m", "messages": messages}) + } + + fn deepseek_v4() -> ChatFormatter { + ChatFormatter::deepseek_native(Some("deepseek_v4"), "any").unwrap() + } + + #[test] + fn no_chat_template_returns_none() { + let cfg = json!({"bos_token": "", "eos_token": ""}); + assert!(ChatFormatter::from_tokenizer_config(cfg, None) + .unwrap() + .is_none()); + } + + /// A sibling `chat_template.jinja` wins over `tokenizer_config.json`, as in + /// transformers, and suffices on its own. + #[test] + fn chat_template_jinja_file_takes_precedence() { + let cfg = json!({"chat_template": "CONFIG"}); + let enc = ChatFormatter::from_tokenizer_config(cfg, Some("FILE")) + .unwrap() + .unwrap(); + assert_eq!(enc.render(&request(json!([]))).unwrap(), "FILE"); + let enc = ChatFormatter::from_tokenizer_config(json!({}), Some("FILE")) + .unwrap() + .unwrap(); + assert_eq!(enc.render(&request(json!([]))).unwrap(), "FILE"); + } + + /// Both the string and the `AddedToken` object forms are accepted. + #[test] + fn additional_special_tokens_are_supplied() { + let enc = jinja(json!({ + "chat_template": "{{ additional_special_tokens | join(',') }}", + "additional_special_tokens": ["", {"content": "", "special": true}] + })); + assert_eq!(enc.render(&request(json!([]))).unwrap(), ","); + } + + #[test] + fn renders_roles_bos_and_generation_prompt() { + let enc = jinja(json!({"chat_template": SIMPLE_TEMPLATE, "bos_token": ""})); + let out = enc + .render(&request(json!([ + {"role": "system", "content": "be brief"}, + {"role": "user", "content": "hi"} + ]))) + .unwrap(); + assert_eq!( + out, + "<|system|>\nbe brief<|end|>\n<|user|>\nhi<|end|>\n<|assistant|>\n" + ); + } + + /// The list form `[{name, template}, ...]` and the `AddedToken` object form + /// of `bos_token` are both accepted. + #[test] + fn list_form_and_bos_object_form() { + let enc = jinja(json!({ + "chat_template": [ + {"name": "tool_use", "template": "TOOLS"}, + {"name": "default", "template": "{{ bos_token }}X"}, + ], + "bos_token": {"content": "<|begin|>", "lstrip": false, "normalized": false, + "rstrip": false, "single_word": false, "special": true}, + })); + assert_eq!(enc.render(&request(json!([]))).unwrap(), "<|begin|>X"); + } + + #[test] + fn absent_documents_match_huggingface_default() { + let enc = jinja(json!({ + "chat_template": "{% if documents is not none %}DOCS{% endif %}{% for m in messages %}{{ m.content }}{% endfor %}" + })); + let mut req = request(json!([{"role":"user","content":"hi"}])); + assert_eq!(enc.render(&req).unwrap(), "hi"); + req["chat_template_kwargs"] = json!({"documents": [{"text": "reference"}]}); + assert_eq!(enc.render(&req).unwrap(), "DOCShi"); + } + + #[test] + fn absent_special_tokens_render_empty() { + let enc = jinja(json!({ + "chat_template": "A{{ bos_token }}{{ eos_token }}{{ unk_token }}{{ sep_token }}{{ pad_token }}{{ cls_token }}{{ mask_token }}B" + })); + assert_eq!(enc.render(&request(json!([]))).unwrap(), "AB"); + } + + #[test] + fn all_special_tokens_from_config_are_supplied() { + let enc = jinja(json!({ + "chat_template": "{{ bos_token }}{{ eos_token }}{{ unk_token }}{{ sep_token }}{{ pad_token }}{{ cls_token }}{{ mask_token }}", + "bos_token": {"content": ""}, + "eos_token": "", + "unk_token": {"content": ""}, + "sep_token": "", + "pad_token": {"content": ""}, + "cls_token": "", + "mask_token": "" + })); + assert_eq!( + enc.render(&request(json!([]))).unwrap(), + "" + ); + } + + /// An unrecognized token shape is a load error, not a silent "". + #[test] + fn malformed_special_token_fails_to_load() { + let cfg = json!({"chat_template": "X", "eos_token": [""]}); + assert!(ChatFormatter::from_tokenizer_config(cfg, None).is_err()); + } + + #[test] + fn tokenizer_settings_do_not_affect_chat_rendering() { + let enc = jinja(json!({ + "chat_template": "{{ bos_token }}{% for m in messages %}{{ m.content }}{% endfor %}", + "bos_token": "", + "sp_model_kwargs": {"enable_sampling": false, "nbest_size": -1, "alpha": 0.1}, + "added_tokens_decoder": {"0": {"content": ""}}, + "truncation_size": 4096 + })); + assert_eq!( + enc.render(&request(json!([{"role": "user", "content": "hi"}]))) + .unwrap(), + "hi" + ); + } + + #[test] + fn chat_template_kwargs_reach_the_template() { + let enc = jinja(json!({"chat_template": "t={{ enable_thinking }}"})); + let mut req = request(json!([])); + assert_eq!(enc.render(&req).unwrap(), "t="); + req["chat_template_kwargs"] = json!({"enable_thinking": true}); + assert_eq!(enc.render(&req).unwrap(), "t=True"); + req["chat_template_kwargs"] = json!("not an object"); + assert!(enc.render(&req).is_err()); + } + + /// Tools pass through to the template; an empty list counts as no tools. + #[test] + fn tools_reach_the_template_and_empty_means_none() { + let enc = jinja(json!({ + "chat_template": "{% if tools is not none %}T:{{ tools | length }}{% endif %}X" + })); + let mut req = request(json!([])); + assert_eq!(enc.render(&req).unwrap(), "X"); + req["tools"] = json!([]); + assert_eq!(enc.render(&req).unwrap(), "X"); + req["tools"] = json!([{"type": "function", "function": {"name": "f"}}]); + assert_eq!(enc.render(&req).unwrap(), "T:1X"); + } + + #[test] + fn raise_exception_surfaces_as_error() { + let enc = jinja(json!({"chat_template": "{{ raise_exception('bad messages') }}"})); + let err = enc.render(&request(json!([]))).unwrap_err(); + assert!(format!("{err:#}").contains("bad messages"), "got: {err:#}"); + } + + #[test] + fn missing_messages_is_an_error() { + assert!(deepseek_v4().render(&json!({"model": "m"})).is_err()); + } + + #[test] + fn deepseek_native_detection() { + assert!(ChatFormatter::deepseek_native(None, "deepseek-ai/DeepSeek-V4-Flash").is_some()); + assert!(ChatFormatter::deepseek_native(None, "deepseek-v4-tiny").is_some()); + assert!(ChatFormatter::deepseek_native(Some("deepseek_v4"), "alias").is_some()); + assert!(ChatFormatter::deepseek_native(Some("deepseek_v41"), "alias").is_some()); + assert!(ChatFormatter::deepseek_native(Some("deepseek_v32"), "DeepSeek-V3.2").is_some()); + assert!(ChatFormatter::deepseek_native(Some("inkling_mm_model"), "inkling").is_none()); + assert!(ChatFormatter::deepseek_native(Some("llama"), "deepseek-v4").is_none()); + assert!(ChatFormatter::deepseek_native(None, "deepseek-ai/DeepSeek-V3.2-Exp").is_none()); + assert!(ChatFormatter::deepseek_native(None, "Qwen/Qwen3-0.6B").is_none()); + } + + /// Byte-exact against the engine's `/tokenize` in its default chat mode: + /// `[{user:"ABCD"}]` -> `[0, 128803, 51453, 128804, 128822]`. + #[test] + fn v4_single_user_turn() { + let out = deepseek_v4() + .render(&request(json!([{"role":"user","content":"ABCD"}]))) + .unwrap(); + assert_eq!( + out, + "<｜begin▁of▁sentence｜><｜User｜>ABCD<｜Assistant｜>" + ); + } + + #[test] + fn v4_system_then_multi_turn() { + let out = deepseek_v4() + .render(&request(json!([ + {"role":"system","content":"SYS"}, + {"role":"user","content":"U1"}, + {"role":"assistant","content":"A1"}, + {"role":"user","content":"U2"} + ]))) + .unwrap(); + assert_eq!( + out, + "<｜begin▁of▁sentence｜>SYS<｜User｜>U1<｜Assistant｜>A1<｜end▁of▁sentence｜><｜User｜>U2<｜Assistant｜>" + ); + } + + /// Request kwargs override the chat-mode default. + #[test] + fn v4_thinking_kwarg_overrides_chat_default() { + let mut req = request(json!([{"role":"user","content":"ABCD"}])); + req["chat_template_kwargs"] = json!({"thinking": true}); + let out = deepseek_v4().render(&req).unwrap(); + assert!(out.ends_with("<｜Assistant｜>"), "got: {out}"); + } + + /// The engine never renders `response_format` into the prompt. + #[test] + fn v4_ignores_response_format() { + let mut req = request(json!([{"role":"user","content":"ABCD"}])); + let plain = deepseek_v4().render(&req).unwrap(); + req["response_format"] = json!({"type": "json_object"}); + assert_eq!(deepseek_v4().render(&req).unwrap(), plain); + } +} diff --git a/experimental/sgl-router/src/tokenizer/chat_template.rs b/experimental/sgl-router/src/tokenizer/chat_template.rs deleted file mode 100644 index 613b5164e..000000000 --- a/experimental/sgl-router/src/tokenizer/chat_template.rs +++ /dev/null @@ -1,365 +0,0 @@ -// 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. -//! -//! 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 deleted file mode 100644 index 07baa3292..000000000 --- a/experimental/sgl-router/src/tokenizer/dsv4.rs +++ /dev/null @@ -1,274 +0,0 @@ -// 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 c246524e2..9dd594da8 100644 --- a/experimental/sgl-router/src/tokenizer/mod.rs +++ b/experimental/sgl-router/src/tokenizer/mod.rs @@ -2,38 +2,15 @@ // SPDX-License-Identifier: Apache-2.0 pub mod adapter; -pub mod chat_template; -pub mod dsv4; +pub mod chat_formatter; use anyhow::Result; -use chat_template::ChatTemplate; +use chat_formatter::ChatFormatter; 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 ChatFormatter { - /// 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 ChatFormatter { - /// Render `messages` into the engine-equivalent prompt text. - fn render(&self, messages: &serde_json::Value) -> Result { - match self { - ChatFormatter::Jinja(t) => t.render(messages), - ChatFormatter::DeepSeekV4 => Ok(dsv4::render_messages(messages)), - } - } -} - /// A model's chat formatter plus its fallback-logging state. struct ChatFormatterEntry { formatter: ChatFormatter, @@ -69,10 +46,7 @@ impl ChatFormatterEntry { pub struct TokenizerRegistry { inner: DashMap>, /// Per-model chat formatter, present only when the model's prompt format is - /// known (a `tokenizer_config.json` chat template, or a built-in formatter - /// 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. + /// known; models without one fall back to raw prompt-text tokenization. formatters: DashMap>, } @@ -90,51 +64,33 @@ impl TokenizerRegistry { let m = &cfg.model; let t = adapter::load(&m.tokenizer_path)?; me.inner.insert(m.id.clone(), t); - // Resolve the chat formatter, best-effort: a Jinja template from - // tokenizer_config.json, else a built-in formatter 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(formatter) = me.resolve_chat_formatter(&m.id, &m.tokenizer_path) { - me.formatters - .insert(m.id.clone(), Arc::new(ChatFormatterEntry::new(formatter))); + match ChatFormatter::load(&m.id, &m.tokenizer_path) { + Ok(Some(formatter)) => { + me.formatters + .insert(m.id.clone(), Arc::new(ChatFormatterEntry::new(formatter))); + tracing::info!(model = %m.id, "dynamo-render chat rendering enabled"); + } + Ok(None) => tracing::info!(model = %m.id, + "no supported chat formatter; chat traffic routes via raw prompt text"), + Err(e) => tracing::warn!(model = %m.id, error = %format!("{e:#}"), + "failed to load chat formatter; chat traffic routes via raw prompt text"), + } + if m.disable_input_ids_forwarding { + tracing::info!(model = %m.id, + "router-generated input_ids forwarding disabled; workers tokenize messages; \ + routing tokenization remains available"); + } else if me.has_chat_formatter(&m.id) { + tracing::warn!(model = %m.id, + "router-generated input_ids forwarding enabled: requires matching worker model \ + files and template defaults; native DeepSeek assumes SGLANG_DEFAULT_THINKING=false \ + and no SGLANG_DSV4_REASONING_EFFORT preamble; worker parser overrides \ + (including --tool-call-parser deepseekv32), content-format detection, and \ + conversation-template stop strings are not replicated. Use \ + --disable-input-ids-forwarding for array-only templates or when these assumptions do not hold"); } Ok(me) } - /// Pick the chat formatter for a model, logging the outcome on every branch. - fn resolve_chat_formatter( - &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(ChatFormatter::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 formatter"); - return Some(ChatFormatter::DeepSeekV4); - } - tracing::info!(model = %model_id, - "no chat template or built-in formatter; 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)) } @@ -145,34 +101,20 @@ impl TokenizerRegistry { self.formatters.contains_key(model_id) } - /// Render `messages` through the model's chat formatter, then tokenize the - /// result the same way the engine does (`add_special_tokens = false`, so the - /// formatter's literal `bos_token`/role markers carry the specials). Returns - /// `None` — caller falls back to raw routing — when the model has no - /// formatter, no tokenizer, or rendering/encoding fails or yields no tokens. - pub fn encode_chat(&self, model_id: &str, messages: &serde_json::Value) -> Option> { + /// Render with dynamo-render and tokenize; return `None` when unavailable or unsuccessful. + pub fn encode_chat(&self, model_id: &str, request: &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.formatters.get(model_id)?); let tokenizer = self.get(model_id)?; - let rendered = entry - .formatter - .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) { + match entry.formatter.encode(&tokenizer, request) { 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:#}")); + entry.log_fallback(model_id, &format!("render or tokenize failed: {e:#}")); None } } @@ -201,23 +143,13 @@ impl TokenizerRegistry { model_id: &str, tokenizer_config: &serde_json::Value, ) { - let template = ChatTemplate::from_tokenizer_config(tokenizer_config) + let formatter = ChatFormatter::from_tokenizer_config(tokenizer_config.clone(), None) .expect("valid test chat template") .expect("test tokenizer_config has a chat_template"); - self.attach_chat_formatter_for_test(model_id, ChatFormatter::Jinja(Box::new(template))); + self.attach_chat_formatter_for_test(model_id, formatter); } } -/// Whether `model_id` denotes a DeepSeek-V4 model, which the engine encodes via -/// the built-in [`dsv4`] formatter 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)] mod tests { use super::*; @@ -236,6 +168,7 @@ mod tests { model: crate::config::ModelConfig { id: "tiny".into(), tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(), + disable_input_ids_forwarding: false, policy: PolicyKind::RoundRobin, decode_policy: Default::default(), bucket_config: None, @@ -384,7 +317,7 @@ mod tests { } #[test] - fn load_tokenizer_config_reads_sibling() { + fn model_files_json_reads_sibling() { let dir = tempfile::tempdir().unwrap(); let tok = dir.path().join("tokenizer.json"); std::fs::write(&tok, "{}").unwrap(); @@ -393,18 +326,64 @@ mod tests { r#"{"chat_template":"X","bos_token":""}"#, ) .unwrap(); - let cfg = adapter::load_tokenizer_config(tok.to_str().unwrap()) + let cfg = adapter::ModelFiles::open(tok.to_str().unwrap()) + .json("tokenizer_config.json") .unwrap() .expect("sibling tokenizer_config.json is loaded"); assert_eq!(cfg["chat_template"], "X"); } + /// Families the engine encodes in code skip a shipped template; V4.1 counts as V4. #[test] - fn load_tokenizer_config_absent_returns_none() { + fn chat_formatter_load_preserves_native_precedence() { 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()) + std::fs::write(dir.path().join("chat_template.jinja"), "T").unwrap(); + let resolve = |model_type: &str| { + let cfg = serde_json::json!({ "model_type": model_type }).to_string(); + std::fs::write(dir.path().join("config.json"), cfg).unwrap(); + ChatFormatter::load("m", tok.to_str().unwrap()).unwrap() + }; + let request = serde_json::json!({"messages": [{"role": "user", "content": "hi"}]}); + for model_type in ["llama", "deepseek_v32"] { + assert_eq!(resolve(model_type).unwrap().render(&request).unwrap(), "T"); + } + assert!(resolve("inkling_mm_model").is_none()); + assert!(resolve("kimi_k3").is_none()); + assert_eq!( + resolve("deepseek_v41").unwrap().render(&request).unwrap(), + "<｜begin▁of▁sentence｜><｜User｜>hi<｜Assistant｜>" + ); + } + + #[test] + fn invalid_chat_template_keeps_tokenizer_available() { + let dir = tempfile::tempdir().unwrap(); + let tok = dir.path().join("tokenizer.json"); + std::fs::copy("tests/fixtures/tiny_tokenizer.json", &tok).unwrap(); + std::fs::write( + dir.path().join("config.json"), + r#"{"model_type":"deepseek_v32"}"#, + ) + .unwrap(); + std::fs::write(dir.path().join("chat_template.jinja"), "{% invalid %}").unwrap(); + let mut cfg = cfg(); + cfg.model.tokenizer_path = tok.to_str().unwrap().to_owned(); + + let reg = TokenizerRegistry::load_from_config(&cfg).unwrap(); + let tokenizer = reg.get(&cfg.model.id).unwrap(); + assert!(!adapter::encode(&tokenizer, "hello").unwrap().is_empty()); + assert!(!reg.has_chat_formatter(&cfg.model.id)); + } + + #[test] + fn model_files_json_absent_returns_none() { + let dir = tempfile::tempdir().unwrap(); + let tok = dir.path().join("tokenizer.json"); + std::fs::write(&tok, "{}").unwrap(); + assert!(adapter::ModelFiles::open(tok.to_str().unwrap()) + .json("tokenizer_config.json") .unwrap() .is_none()); } @@ -426,8 +405,8 @@ mod tests { reg.attach_chat_template_for_test("tiny", &cfg); assert!(reg.has_chat_formatter("tiny")); - let messages = serde_json::json!([{"role":"user","content":"hi"}]); - let chat_ids = reg.encode_chat("tiny", &messages).expect("encode_chat"); + let request = serde_json::json!({"messages": [{"role":"user","content":"hi"}]}); + let chat_ids = reg.encode_chat("tiny", &request).expect("encode_chat"); assert!(!chat_ids.is_empty()); let tok = reg.get("tiny").unwrap(); @@ -437,17 +416,42 @@ mod tests { "chat-templated tokens must differ from raw-content tokens" ); - // encode_chat is exactly tokenize(render(messages)). + // encode_chat is exactly tokenize(render(request)). let rendered = reg .formatters .get("tiny") .unwrap() .formatter - .render(&messages) + .render(&request) .unwrap(); assert_eq!(chat_ids, adapter::encode(&tok, &rendered).unwrap()); } + #[test] + fn routing_tokenization_receives_tools_and_template_kwargs() { + let reg = TokenizerRegistry::default(); + let tok = adapter::load("tests/fixtures/tiny_tokenizer.json").unwrap(); + reg.inner.insert("tiny".into(), Arc::clone(&tok)); + reg.attach_chat_template_for_test( + "tiny", + &serde_json::json!({ + "chat_template": "{{ tools[0].function.name }} {{ greeting }}" + }), + ); + let request = serde_json::json!({ + "messages": [{"role": "user", "content": "hi"}], + "tools": [{"type": "function", "function": {"name": "hello"}}], + "chat_template_kwargs": {"greeting": "world"} + }); + let tokens = crate::policies::request_tokens_for( + ®, + &crate::discovery::ModelId("tiny".into()), + &request, + ) + .expect("request tokenizes"); + assert_eq!(tokens.ids, adapter::encode(&tok, "hello world").unwrap()); + } + #[test] fn encode_chat_none_without_template() { let reg = TokenizerRegistry::default(); @@ -456,8 +460,8 @@ mod tests { adapter::load("tests/fixtures/tiny_tokenizer.json").unwrap(), ); assert!(!reg.has_chat_formatter("tiny")); - let messages = serde_json::json!([{"role":"user","content":"hi"}]); - assert!(reg.encode_chat("tiny", &messages).is_none()); + let request = serde_json::json!({"messages": [{"role":"user","content":"hi"}]}); + assert!(reg.encode_chat("tiny", &request).is_none()); } /// A template that fails to render (here, one that calls `raise_exception`) @@ -478,20 +482,10 @@ mod tests { }), ); assert!(reg.has_chat_formatter("tiny")); - let messages = serde_json::json!([{"role":"user","content":"hi"}]); + let request = serde_json::json!({"messages": [{"role":"user","content":"hi"}]}); assert!( - reg.encode_chat("tiny", &messages).is_none(), + reg.encode_chat("tiny", &request).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/src/workers/manager.rs b/experimental/sgl-router/src/workers/manager.rs index 4643e6bf8..58ac77bb9 100644 --- a/experimental/sgl-router/src/workers/manager.rs +++ b/experimental/sgl-router/src/workers/manager.rs @@ -591,6 +591,7 @@ mod tests { model: ModelConfig { id: id.into(), tokenizer_path: "/tmp/x".into(), + disable_input_ids_forwarding: false, policy: PolicyKind::RoundRobin, decode_policy: Default::default(), bucket_config: None, diff --git a/experimental/sgl-router/tests/component/discovery/static_urls.rs b/experimental/sgl-router/tests/component/discovery/static_urls.rs index d6dbfb780..57068a9e8 100644 --- a/experimental/sgl-router/tests/component/discovery/static_urls.rs +++ b/experimental/sgl-router/tests/component/discovery/static_urls.rs @@ -130,6 +130,7 @@ async fn static_urls_pd_role_resolved_end_to_end() { model: sgl_router::config::ModelConfig { id: "tiny".into(), tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(), + disable_input_ids_forwarding: false, policy: sgl_router::config::PolicyKind::RoundRobin, decode_policy: Default::default(), bucket_config: None, diff --git a/experimental/sgl-router/tests/proxy/bucket_routing.rs b/experimental/sgl-router/tests/proxy/bucket_routing.rs index 9f555c7d8..f35f5ef6e 100644 --- a/experimental/sgl-router/tests/proxy/bucket_routing.rs +++ b/experimental/sgl-router/tests/proxy/bucket_routing.rs @@ -61,6 +61,7 @@ fn build_app_context( model: ModelConfig { id: "tiny".into(), tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(), + disable_input_ids_forwarding: false, policy, decode_policy: Default::default(), bucket_config: Some(bucket_config), diff --git a/experimental/sgl-router/tests/proxy/chat_routing.rs b/experimental/sgl-router/tests/proxy/chat_routing.rs index 9ecf68fcc..7adb26f9f 100644 --- a/experimental/sgl-router/tests/proxy/chat_routing.rs +++ b/experimental/sgl-router/tests/proxy/chat_routing.rs @@ -34,6 +34,7 @@ fn config_for(_worker_url: &str) -> Config { model: ModelConfig { id: "tiny".into(), tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(), + disable_input_ids_forwarding: false, policy: PolicyKind::RoundRobin, decode_policy: Default::default(), bucket_config: None, diff --git a/experimental/sgl-router/tests/proxy/common/cache_aware_fixture.rs b/experimental/sgl-router/tests/proxy/common/cache_aware_fixture.rs index 72cdc4ab5..bfbe55e6e 100644 --- a/experimental/sgl-router/tests/proxy/common/cache_aware_fixture.rs +++ b/experimental/sgl-router/tests/proxy/common/cache_aware_fixture.rs @@ -26,6 +26,7 @@ pub fn config() -> Config { model: ModelConfig { id: MODEL.into(), tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(), + disable_input_ids_forwarding: false, policy: PolicyKind::CacheAware, decode_policy: Default::default(), bucket_config: None, diff --git a/experimental/sgl-router/tests/proxy/failover.rs b/experimental/sgl-router/tests/proxy/failover.rs index 6e4696ee9..9a3a97789 100644 --- a/experimental/sgl-router/tests/proxy/failover.rs +++ b/experimental/sgl-router/tests/proxy/failover.rs @@ -35,6 +35,7 @@ async fn failover_when_one_worker_dies() { model: ModelConfig { id: "tiny".into(), tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(), + disable_input_ids_forwarding: false, policy: PolicyKind::RoundRobin, decode_policy: Default::default(), bucket_config: None, diff --git a/experimental/sgl-router/tests/proxy/graceful_shutdown.rs b/experimental/sgl-router/tests/proxy/graceful_shutdown.rs index 5cf8d7a66..6eae3320c 100644 --- a/experimental/sgl-router/tests/proxy/graceful_shutdown.rs +++ b/experimental/sgl-router/tests/proxy/graceful_shutdown.rs @@ -53,6 +53,7 @@ fn build_ctx_with_worker(worker_url: &str) -> Arc { model: ModelConfig { id: "tiny".into(), tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(), + disable_input_ids_forwarding: false, policy: PolicyKind::RoundRobin, decode_policy: Default::default(), bucket_config: None, diff --git a/experimental/sgl-router/tests/proxy/header_forwarding.rs b/experimental/sgl-router/tests/proxy/header_forwarding.rs index 374817603..bad3360f4 100644 --- a/experimental/sgl-router/tests/proxy/header_forwarding.rs +++ b/experimental/sgl-router/tests/proxy/header_forwarding.rs @@ -31,6 +31,7 @@ async fn forwards_whitelisted_headers_strips_others() { model: ModelConfig { id: "tiny".into(), tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(), + disable_input_ids_forwarding: false, policy: PolicyKind::RoundRobin, decode_policy: Default::default(), bucket_config: None, diff --git a/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs b/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs index aa2ab1542..d0fd95422 100644 --- a/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs +++ b/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs @@ -46,6 +46,7 @@ fn config() -> Config { model: ModelConfig { id: "tiny".into(), tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(), + disable_input_ids_forwarding: false, policy: PolicyKind::RoundRobin, decode_policy: Default::default(), bucket_config: None, diff --git a/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs b/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs index 460f9923f..2d09f97f4 100644 --- a/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs +++ b/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs @@ -45,6 +45,7 @@ fn config() -> Config { model: ModelConfig { id: "tiny".into(), tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(), + disable_input_ids_forwarding: false, policy: PolicyKind::RoundRobin, decode_policy: Default::default(), bucket_config: None, diff --git a/experimental/sgl-router/tests/proxy/pd_protocol_binding.rs b/experimental/sgl-router/tests/proxy/pd_protocol_binding.rs index f5eccbceb..f310fb296 100644 --- a/experimental/sgl-router/tests/proxy/pd_protocol_binding.rs +++ b/experimental/sgl-router/tests/proxy/pd_protocol_binding.rs @@ -133,6 +133,7 @@ fn config() -> Config { model: ModelConfig { id: "tiny".into(), tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(), + disable_input_ids_forwarding: false, policy: PolicyKind::RoundRobin, decode_policy: Default::default(), bucket_config: None, diff --git a/experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs b/experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs index cbe84b272..fbd2ff073 100644 --- a/experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs +++ b/experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs @@ -16,11 +16,12 @@ use sgl_router::config::{ }; use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; use sgl_router::policies::factory::build_registry_with_defaults; +use sgl_router::policies::{Policy, SelectionContext}; use sgl_router::proxy::Proxy; use sgl_router::server::app::build_router; use sgl_router::server::app_context::AppContext; use sgl_router::tokenizer::TokenizerRegistry; -use sgl_router::workers::WorkerRegistry; +use sgl_router::workers::{Worker, WorkerRegistry}; use std::sync::Arc; use std::time::Duration; use tower::ServiceExt; @@ -42,6 +43,7 @@ fn config() -> Config { model: ModelConfig { id: MODEL.into(), tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(), + disable_input_ids_forwarding: false, policy: PolicyKind::RoundRobin, decode_policy: Default::default(), bucket_config: None, @@ -62,7 +64,10 @@ fn config() -> Config { } fn build_ctx(url: String) -> Arc { - let cfg = config(); + build_ctx_with_config(url, config()) +} + +fn build_ctx_with_config(url: String, cfg: Config) -> Arc { // The handler tokenizes via the AppContext's registry (which carries the V4 // encoder); the RoundRobin policy itself needs no tokenizer. let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap()); @@ -80,6 +85,36 @@ fn build_ctx(url: String) -> Arc { Arc::new(AppContext::new(cfg, tokenizers, proxy, registry, policies)) } +fn template_config(tokenizer_config: Value) -> (tempfile::TempDir, Config) { + let dir = tempfile::tempdir().unwrap(); + let tokenizer = dir.path().join("tokenizer.json"); + std::fs::copy("tests/fixtures/tiny_tokenizer.json", &tokenizer).unwrap(); + std::fs::write( + dir.path().join("tokenizer_config.json"), + tokenizer_config.to_string(), + ) + .unwrap(); + let mut cfg = config(); + cfg.model.tokenizer_path = tokenizer.to_str().unwrap().into(); + (dir, cfg) +} + +fn without_forwarding(mut cfg: Config, policy: PolicyKind) -> Config { + cfg.model.policy = policy; + cfg.model.cache_aware = (policy == PolicyKind::CacheAware).then(Default::default); + cfg.model.disable_input_ids_forwarding = true; + cfg +} + +async fn assert_forwarded_unchanged(ctx: &Arc, mock: &MockWorker, request: &Value) { + assert_eq!(send(Arc::clone(ctx), request.clone()).await, StatusCode::OK); + assert_eq!(captured(mock), *request); + assert!(!ctx + .metrics + .render() + .contains("sgl_router_ingress_tokenize_errors_total{")); +} + async fn send(ctx: Arc, body: Value) -> StatusCode { let app = build_router(ctx); let req = Request::builder() @@ -130,6 +165,98 @@ async fn round_robin_plain_chat_forwards_input_ids() { ); } +#[tokio::test] +async fn forwarding_opt_out_preserves_messages_and_caller_ids() { + for policy in [PolicyKind::RoundRobin, PolicyKind::CacheAware] { + let mock = MockWorker::start(vec![]).await; + let ctx = build_ctx_with_config(mock.url.clone(), without_forwarding(config(), policy)); + let mut request = + json!({"model": MODEL, "messages": [{"role": "user", "content": "hello"}]}); + assert_forwarded_unchanged(&ctx, &mock, &request).await; + request["input_ids"] = json!([42, 43]); + assert_forwarded_unchanged(&ctx, &mock, &request).await; + } +} + +#[tokio::test] +async fn forwarding_opt_out_keeps_ingress_tokens_for_routing() { + #[derive(Debug)] + struct ExpectTokens(Vec); + impl Policy for ExpectTokens { + fn needs_request_tokens(&self) -> bool { + true + } + + fn select( + &self, + workers: &[Arc], + ctx: &SelectionContext<'_>, + ) -> Option> { + assert_eq!(ctx.request_tokens(), Some(self.0.as_slice())); + workers.first().cloned() + } + } + + let mock = MockWorker::start(vec![]).await; + let cfg = without_forwarding(config(), PolicyKind::RoundRobin); + let ctx = build_ctx_with_config(mock.url.clone(), cfg); + let request = json!({"model": MODEL, "messages": [{"role": "user", "content": "hello"}]}); + let expected = ctx.tokenizers.encode_chat(MODEL, &request).unwrap(); + ctx.policies + .insert(ModelId(MODEL.into()), Arc::new(ExpectTokens(expected))); + assert_forwarded_unchanged(&ctx, &mock, &request).await; +} + +/// Array-only deployments must opt out until Dynamo exposes its conversion flag. +#[tokio::test] +async fn array_only_template_opt_out_preserves_engine_processing() { + let (_dir, cfg) = template_config(json!({ + "chat_template": "{% for m in messages %}{% for part in m.content %}{{ part.text }}{% endfor %}{% endfor %}" + })); + let mock = MockWorker::start(vec![]).await; + let ctx = build_ctx_with_config( + mock.url.clone(), + without_forwarding(cfg, PolicyKind::CacheAware), + ); + let request = json!({"model": MODEL, "messages": [{"role": "user", "content": "hello"}]}); + assert!(!ctx + .tokenizers + .encode_chat(MODEL, &request) + .unwrap() + .is_empty()); + assert_forwarded_unchanged(&ctx, &mock, &request).await; +} + +#[tokio::test] +async fn template_with_date_helper_forwards_input_ids() { + // GPT-OSS uses strftime_now; a bare Jinja probe incorrectly blocks it. + let (_dir, cfg) = template_config(json!({ + "chat_template": "{{ strftime_now('%Y-%m-%d') }}{% for m in messages %}{{ m.content }}{% endfor %}" + })); + let mock = MockWorker::start(vec![]).await; + let ctx = build_ctx_with_config(mock.url.clone(), cfg); + let request = json!({"model": MODEL, "messages": [{"role": "user", "content": "hello"}]}); + let expected = ctx.tokenizers.encode_chat(MODEL, &request).unwrap(); + assert_eq!(send(ctx, request.clone()).await, StatusCode::OK); + let body = captured(&mock); + assert_eq!(body["input_ids"], json!(expected)); + assert_eq!(body["messages"], request["messages"]); +} + +#[tokio::test] +async fn disabled_forwarding_does_not_count_routing_render_failures_as_offload_errors() { + let (_dir, cfg) = + template_config(json!({"chat_template": "{{ raise_exception('cannot render') }}"})); + let mock = MockWorker::start(vec![]).await; + let ctx = build_ctx_with_config( + mock.url.clone(), + without_forwarding(cfg, PolicyKind::CacheAware), + ); + let request = json!({"model": MODEL, "messages": [{"role": "user", "content": "hello"}]}); + assert!(ctx.tokenizers.encode_chat(MODEL, &request).is_none()); + assert_forwarded_unchanged(&ctx, &mock, &request).await; +} + /// Even under round-robin, a tool request omits `input_ids` (the safe predicate /// is policy-independent too). #[tokio::test] @@ -196,3 +323,79 @@ async fn successful_forward_does_not_emit_ingress_tokenize_error() { "healthy forwards (and expected omissions) must not emit the error counter; got:\n{m}", ); } + +/// History that dynamo-render rewrites stays intact for engine-side tokenization. +#[tokio::test] +async fn reasoning_history_preserves_messages_without_forwarding_ids() { + let (_dir, cfg) = template_config(json!({ + "chat_template": "{% for m in messages %}{{ m.role }}:{{ m.content }};{% endfor %}" + })); + let mock = MockWorker::start(vec![]).await; + let ctx = build_ctx_with_config(mock.url.clone(), cfg); + let mut request = json!({"model": MODEL, "messages": [ + {"role":"user", "content":"hi"}, + {"role":"assistant", "content":"answer", "reasoning_content":"prior reasoning"}, + {"role":"user", "content":"next"} + ]}); + assert!(!ctx + .tokenizers + .encode_chat(MODEL, &request) + .unwrap() + .is_empty()); + assert_forwarded_unchanged(&ctx, &mock, &request).await; + + request["messages"][1] + .as_object_mut() + .unwrap() + .remove("reasoning_content"); + assert_eq!(send(ctx, request).await, StatusCode::OK); + assert!(captured(&mock).get("input_ids").is_some()); +} + +/// Strict-template rewrites are used for routing only; the engine gets the original turns. +#[tokio::test] +async fn role_rewrites_preserve_messages_without_forwarding_ids() { + let template = concat!( + "{%- set ns = namespace(prev='') -%}", + "{%- for m in messages -%}", + "{%- if m.role == 'system' and not loop.first -%}", + "{{ raise_exception('System message must be first.') }}", + "{%- endif -%}", + "{%- if m.role == 'user' and ns.prev == 'user' -%}", + "{{ raise_exception('Conversation roles must alternate.') }}", + "{%- endif -%}", + "{{ m.role }}:{{ m.content }};", + "{%- set ns.prev = m.role -%}", + "{%- endfor -%}" + ); + let (_dir, cfg) = template_config(json!({ + "chat_template": template, "sp_model_kwargs": {"enable_sampling": false} + })); + let mock = MockWorker::start(vec![]).await; + let ctx = build_ctx_with_config(mock.url.clone(), cfg); + for roles in [ + vec!["user", "user"], + vec!["system", "system", "user"], + vec!["user", "assistant", "system", "user"], + ] { + let messages: Vec<_> = roles + .iter() + .map(|role| json!({"role": role, "content": "text"})) + .collect(); + let request = json!({"model": MODEL, "messages": messages}); + assert!(!ctx + .tokenizers + .encode_chat(MODEL, &request) + .unwrap() + .is_empty()); + assert_forwarded_unchanged(&ctx, &mock, &request).await; + } + let request = json!({"model": MODEL, "messages": [ + {"role": "system", "content": "instructions"}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"role": "user", "content": "next"} + ]}); + assert_eq!(send(ctx, request).await, StatusCode::OK); + assert!(captured(&mock).get("input_ids").is_some()); +} diff --git a/experimental/sgl-router/tests/proxy/shared_prefill_admission.rs b/experimental/sgl-router/tests/proxy/shared_prefill_admission.rs index 3eabdcb62..52414e1aa 100644 --- a/experimental/sgl-router/tests/proxy/shared_prefill_admission.rs +++ b/experimental/sgl-router/tests/proxy/shared_prefill_admission.rs @@ -149,6 +149,7 @@ fn config(policy: PolicyKind) -> Config { model: ModelConfig { id: "tiny".into(), tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(), + disable_input_ids_forwarding: false, policy, decode_policy: Default::default(), bucket_config: None, diff --git a/experimental/sgl-router/tests/proxy/sticky_input_ids.rs b/experimental/sgl-router/tests/proxy/sticky_input_ids.rs index 1f22f38cc..8c20010f7 100644 --- a/experimental/sgl-router/tests/proxy/sticky_input_ids.rs +++ b/experimental/sgl-router/tests/proxy/sticky_input_ids.rs @@ -54,6 +54,7 @@ fn config() -> Config { model: ModelConfig { id: MODEL.into(), tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(), + disable_input_ids_forwarding: false, policy: PolicyKind::Sticky, decode_policy: Default::default(), bucket_config: None, diff --git a/experimental/sgl-router/tests/proxy/sticky_routing.rs b/experimental/sgl-router/tests/proxy/sticky_routing.rs index 02fd1fcbc..96ddad1e0 100644 --- a/experimental/sgl-router/tests/proxy/sticky_routing.rs +++ b/experimental/sgl-router/tests/proxy/sticky_routing.rs @@ -42,6 +42,7 @@ fn build_sticky_ctx(header_name: &str, worker_urls: &[String]) -> Arc Config { model: ModelConfig { id: "tiny".into(), tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(), + disable_input_ids_forwarding: false, policy: PolicyKind::RoundRobin, decode_policy: Default::default(), bucket_config: None,