From 90a7b3c8321979395871d9f54a62f358b377b4e5 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 3 Sep 2026 18:03:14 +0200 Subject: [PATCH 1/7] fix: cache aware routing with token ids Signed-off-by: simondanielsson --- src/protocols/spec.rs | 39 ++++++++++++++++++++++++++++----------- 1 file changed, 28 insertions(+), 11 deletions(-) diff --git a/src/protocols/spec.rs b/src/protocols/spec.rs index 3ed54dc8..108418ba 100644 --- a/src/protocols/spec.rs +++ b/src/protocols/spec.rs @@ -2336,6 +2336,15 @@ pub enum PromptInput { String(String), } +/// Serialize a token-id sequence into a stable, prefix-preserving routing key. +/// Matches the space-separated encoding used by InferenceGenerateRequest. +fn encode_token_ids(ids: &[i32]) -> String { + ids.iter() + .map(|id| id.to_string()) + .collect::>() + .join(" ") +} + impl PromptInput { /// Get the number of items in the PromptInput pub fn len(&self) -> usize { @@ -2358,21 +2367,19 @@ impl PromptInput { } /// Extract text representation for routing decisions - /// For token IDs, converts to a string representation + /// For token IDs, serialize the ids themselves so prefix-based cache-aware + /// routing can match shared prefixes (e.g. multiturn conversations); a + /// count-only key would collide every same-length prompt onto one node. pub fn extract_text_for_routing(&self) -> String { match self { PromptInput::String(s) => s.clone(), PromptInput::StringArray(arr) => arr.join(" "), - PromptInput::IntArray(ids) => { - // Convert token IDs to string representation for routing - // Format: "token_ids:" to indicate this is a token-based prompt - format!("token_ids:{}", ids.len()) - } - PromptInput::IntBatch(batches) => { - // For batches, use total token count - let total_tokens: usize = batches.iter().map(|b| b.len()).sum(); - format!("token_ids_batch:{}:{}", batches.len(), total_tokens) - } + PromptInput::IntArray(ids) => encode_token_ids(ids), + PromptInput::IntBatch(batches) => batches + .iter() + .map(|b| encode_token_ids(b)) + .collect::>() + .join(";"), } } @@ -2425,6 +2432,16 @@ mod tests { assert_eq!(req.extract_text_for_routing(), "151644 8948 198 2610"); } + #[test] + fn test_prompt_input_token_ids_routing_key() { + // Token-id prompts must serialize their ids (not a count) so cache-aware + // routing can prefix-match shared conversation prefixes. + let single = PromptInput::IntArray(vec![151644, 8948, 198, 2610]); + assert_eq!(single.extract_text_for_routing(), "151644 8948 198 2610"); + let batch = PromptInput::IntBatch(vec![vec![1, 2], vec![3, 4]]); + assert_eq!(batch.extract_text_for_routing(), "1 2;3 4"); + } + #[test] fn test_inference_generate_lossless_passthrough() { let body = serde_json::json!({ From a2f03f7562f98a08a8e9cf2827c06e8e56a350e0 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 3 Sep 2026 18:08:53 +0200 Subject: [PATCH 2/7] chore: remove redundant comments Signed-off-by: simondanielsson --- src/protocols/spec.rs | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/src/protocols/spec.rs b/src/protocols/spec.rs index 108418ba..d3337a05 100644 --- a/src/protocols/spec.rs +++ b/src/protocols/spec.rs @@ -2336,8 +2336,6 @@ pub enum PromptInput { String(String), } -/// Serialize a token-id sequence into a stable, prefix-preserving routing key. -/// Matches the space-separated encoding used by InferenceGenerateRequest. fn encode_token_ids(ids: &[i32]) -> String { ids.iter() .map(|id| id.to_string()) @@ -2367,9 +2365,7 @@ impl PromptInput { } /// Extract text representation for routing decisions - /// For token IDs, serialize the ids themselves so prefix-based cache-aware - /// routing can match shared prefixes (e.g. multiturn conversations); a - /// count-only key would collide every same-length prompt onto one node. + /// For token IDs, converts to a string representation pub fn extract_text_for_routing(&self) -> String { match self { PromptInput::String(s) => s.clone(), @@ -2434,8 +2430,6 @@ mod tests { #[test] fn test_prompt_input_token_ids_routing_key() { - // Token-id prompts must serialize their ids (not a count) so cache-aware - // routing can prefix-match shared conversation prefixes. let single = PromptInput::IntArray(vec![151644, 8948, 198, 2610]); assert_eq!(single.extract_text_for_routing(), "151644 8948 198 2610"); let batch = PromptInput::IntBatch(vec![vec![1, 2], vec![3, 4]]); From 171ddf45326b2cde469b760794636efd291c682e Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Mon, 7 Sep 2026 09:42:34 +0200 Subject: [PATCH 3/7] fix: old tests Signed-off-by: simondanielsson --- tests/test_prompt_input.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/test_prompt_input.rs b/tests/test_prompt_input.rs index 5f62f8ea..1379b245 100644 --- a/tests/test_prompt_input.rs +++ b/tests/test_prompt_input.rs @@ -67,7 +67,10 @@ fn test_prompt_input_single_int_array() { assert_eq!(req.prompt.len(), 1); assert!(!req.prompt.is_empty()); assert!(req.prompt.is_token_based()); - assert_eq!(req.prompt.extract_text_for_routing(), "token_ids:5"); + assert_eq!( + req.prompt.extract_text_for_routing(), + "128000 9906 11 1917 0" + ); assert_eq!(req.prompt.estimated_token_count(), 5); } @@ -98,7 +101,7 @@ fn test_prompt_input_int_batch() { assert!(req.prompt.is_token_based()); assert_eq!( req.prompt.extract_text_for_routing(), - "token_ids_batch:2:10" + "128000 9906 11 1917 0;128001 9906 11 1917 1" ); assert_eq!(req.prompt.estimated_token_count(), 10); } From 4debc7ffdb2517124d5005b98e31ddcb7882d2d5 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Mon, 7 Sep 2026 14:01:44 +0200 Subject: [PATCH 4/7] fix: encode token separators Signed-off-by: simondanielsson --- src/policies/cache_aware.rs | 55 ++++++++++++++++++++++++++++++++++++- src/protocols/spec.rs | 18 ++++++++---- tests/test_prompt_input.rs | 5 ++-- 3 files changed, 69 insertions(+), 9 deletions(-) diff --git a/src/policies/cache_aware.rs b/src/policies/cache_aware.rs index dc209a19..09979d77 100644 --- a/src/policies/cache_aware.rs +++ b/src/policies/cache_aware.rs @@ -72,6 +72,32 @@ use std::thread; use std::time::Duration; use tracing::{debug, info}; +/// Unit separator that `PromptInput` uses to terminate each encoded token id. +const TOKEN_SEPARATOR: char = '\u{1f}'; + +/// Truncate a character-level prefix match to the last complete token boundary. +/// +/// Token-id routing keys terminate each token with [`TOKEN_SEPARATOR`]. Because +/// the routing tree matches on characters, two distinct tokens that share +/// leading digits (e.g. `123456` and `123457`) would otherwise be credited a +/// partial match. Rounding back to the last separator credits only complete, +/// equal token IDs. Plain-text keys contain no separator and are returned as-is. +fn complete_token_match(text: &str, matched_char_count: usize) -> usize { + if matched_char_count == 0 || !text.contains(TOKEN_SEPARATOR) { + return matched_char_count; + } + let mut boundary = 0; + for (idx, ch) in text.chars().enumerate() { + if idx >= matched_char_count { + break; + } + if ch == TOKEN_SEPARATOR { + boundary = idx + 1; + } + } + boundary +} + /// Cache-aware routing policy /// /// Routes requests based on cache affinity when load is balanced, @@ -300,10 +326,11 @@ impl LoadBalancingPolicy for CacheAwarePolicy { // Now we work with the tree without holding the HashMap lock // Use prefix_match_with_counts to avoid redundant chars().count() calls let result = tree.prefix_match_with_counts(text); + let matched_char_count = complete_token_match(text, result.matched_char_count); let match_rate = if result.input_char_count == 0 { 0.0 } else { - result.matched_char_count as f32 / result.input_char_count as f32 + matched_char_count as f32 / result.input_char_count as f32 }; debug!( @@ -479,6 +506,32 @@ mod tests { use super::*; use crate::core::{BasicWorker, WorkerType}; + #[test] + fn test_complete_token_match_ignores_partial_token_ids() { + // "123457\u{1f}8\u{1f}" vs cached "123456\u{1f}7\u{1f}" share "12345" + // (5 chars) but zero complete tokens -> must credit 0. + let incoming = "123457\u{1f}8\u{1f}"; + assert_eq!(complete_token_match(incoming, 5), 0); + } + + #[test] + fn test_complete_token_match_credits_whole_tokens_only() { + // Matched through the second separator plus a partial third token -> + // credit only the two complete tokens (up to the last separator). + let incoming = "10\u{1f}20\u{1f}30\u{1f}"; + assert_eq!(complete_token_match(incoming, 7), 6); + // A full match ending on a separator is kept in full. + let len = incoming.chars().count(); + assert_eq!(complete_token_match(incoming, len), len); + } + + #[test] + fn test_complete_token_match_leaves_text_keys_unchanged() { + // Plain-text keys have no separator, so the character match is returned + // as-is (mid-word matches are still credited, as before). + assert_eq!(complete_token_match("hello world", 8), 8); + } + #[test] fn test_cache_aware_with_balanced_load() { // Create policy without eviction thread for testing diff --git a/src/protocols/spec.rs b/src/protocols/spec.rs index d3337a05..eadf0931 100644 --- a/src/protocols/spec.rs +++ b/src/protocols/spec.rs @@ -2337,10 +2337,10 @@ pub enum PromptInput { } fn encode_token_ids(ids: &[i32]) -> String { - ids.iter() - .map(|id| id.to_string()) - .collect::>() - .join(" ") + // Terminate every token id with a unit separator (never present in text) so + // cache-aware routing, a character radix, credits only complete token IDs + // rather than shared leading digits of two different tokens. + ids.iter().map(|id| format!("{id}\u{1f}")).collect() } impl PromptInput { @@ -2431,9 +2431,15 @@ mod tests { #[test] fn test_prompt_input_token_ids_routing_key() { let single = PromptInput::IntArray(vec![151644, 8948, 198, 2610]); - assert_eq!(single.extract_text_for_routing(), "151644 8948 198 2610"); + assert_eq!( + single.extract_text_for_routing(), + "151644\u{1f}8948\u{1f}198\u{1f}2610\u{1f}" + ); let batch = PromptInput::IntBatch(vec![vec![1, 2], vec![3, 4]]); - assert_eq!(batch.extract_text_for_routing(), "1 2;3 4"); + assert_eq!( + batch.extract_text_for_routing(), + "1\u{1f}2\u{1f};3\u{1f}4\u{1f}" + ); } #[test] diff --git a/tests/test_prompt_input.rs b/tests/test_prompt_input.rs index 1379b245..cc4077fc 100644 --- a/tests/test_prompt_input.rs +++ b/tests/test_prompt_input.rs @@ -69,7 +69,7 @@ fn test_prompt_input_single_int_array() { assert!(req.prompt.is_token_based()); assert_eq!( req.prompt.extract_text_for_routing(), - "128000 9906 11 1917 0" + "128000\u{1f}9906\u{1f}11\u{1f}1917\u{1f}0\u{1f}" ); assert_eq!(req.prompt.estimated_token_count(), 5); } @@ -101,7 +101,8 @@ fn test_prompt_input_int_batch() { assert!(req.prompt.is_token_based()); assert_eq!( req.prompt.extract_text_for_routing(), - "128000 9906 11 1917 0;128001 9906 11 1917 1" + "128000\u{1f}9906\u{1f}11\u{1f}1917\u{1f}0\u{1f}\ + ;128001\u{1f}9906\u{1f}11\u{1f}1917\u{1f}1\u{1f}" ); assert_eq!(req.prompt.estimated_token_count(), 10); } From 8c16041d3fa1578fd85a42259ba14c23f48623d2 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Mon, 7 Sep 2026 14:38:38 +0200 Subject: [PATCH 5/7] fix: make match ratio token aware Signed-off-by: simondanielsson --- src/policies/cache_aware.rs | 70 ++++++++++++++++--------------------- 1 file changed, 31 insertions(+), 39 deletions(-) diff --git a/src/policies/cache_aware.rs b/src/policies/cache_aware.rs index 09979d77..7faf1c04 100644 --- a/src/policies/cache_aware.rs +++ b/src/policies/cache_aware.rs @@ -75,27 +75,24 @@ use tracing::{debug, info}; /// Unit separator that `PromptInput` uses to terminate each encoded token id. const TOKEN_SEPARATOR: char = '\u{1f}'; -/// Truncate a character-level prefix match to the last complete token boundary. -/// -/// Token-id routing keys terminate each token with [`TOKEN_SEPARATOR`]. Because -/// the routing tree matches on characters, two distinct tokens that share -/// leading digits (e.g. `123456` and `123457`) would otherwise be credited a -/// partial match. Rounding back to the last separator credits only complete, -/// equal token IDs. Plain-text keys contain no separator and are returned as-is. -fn complete_token_match(text: &str, matched_char_count: usize) -> usize { - if matched_char_count == 0 || !text.contains(TOKEN_SEPARATOR) { - return matched_char_count; +/// Cache-match ratio, computed per token for token-id routing keys and otherwise for chars. +fn token_aware_match_rate(text: &str, matched_char_count: usize, input_char_count: usize) -> f32 { + if input_char_count == 0 { + return 0.0; } - let mut boundary = 0; - for (idx, ch) in text.chars().enumerate() { - if idx >= matched_char_count { - break; - } - if ch == TOKEN_SEPARATOR { - boundary = idx + 1; + if text.contains(TOKEN_SEPARATOR) { + let total_tokens = text.matches(TOKEN_SEPARATOR).count(); + if total_tokens == 0 { + return 0.0; } + let matched_tokens = text + .chars() + .take(matched_char_count) + .filter(|c| *c == TOKEN_SEPARATOR) + .count(); + return matched_tokens as f32 / total_tokens as f32; } - boundary + matched_char_count as f32 / input_char_count as f32 } /// Cache-aware routing policy @@ -326,12 +323,8 @@ impl LoadBalancingPolicy for CacheAwarePolicy { // Now we work with the tree without holding the HashMap lock // Use prefix_match_with_counts to avoid redundant chars().count() calls let result = tree.prefix_match_with_counts(text); - let matched_char_count = complete_token_match(text, result.matched_char_count); - let match_rate = if result.input_char_count == 0 { - 0.0 - } else { - matched_char_count as f32 / result.input_char_count as f32 - }; + let match_rate = + token_aware_match_rate(text, result.matched_char_count, result.input_char_count); debug!( "Cache match for model '{}': matched_chars={}, input_chars={}, match_rate={:.2}", @@ -507,29 +500,28 @@ mod tests { use crate::core::{BasicWorker, WorkerType}; #[test] - fn test_complete_token_match_ignores_partial_token_ids() { - // "123457\u{1f}8\u{1f}" vs cached "123456\u{1f}7\u{1f}" share "12345" - // (5 chars) but zero complete tokens -> must credit 0. + fn test_token_aware_match_rate_ignores_partial_token_ids() { + // "123457\u{1f}8\u{1f}" vs cached "123456\u{1f}7\u{1f}" share 5 chars but + // zero complete tokens -> rate 0. let incoming = "123457\u{1f}8\u{1f}"; - assert_eq!(complete_token_match(incoming, 5), 0); + let input = incoming.chars().count(); + assert_eq!(token_aware_match_rate(incoming, 5, input), 0.0); } #[test] - fn test_complete_token_match_credits_whole_tokens_only() { - // Matched through the second separator plus a partial third token -> - // credit only the two complete tokens (up to the last separator). + fn test_token_aware_match_rate_counts_whole_tokens() { + // Matched through two separators plus a partial third token -> 2 of 3. let incoming = "10\u{1f}20\u{1f}30\u{1f}"; - assert_eq!(complete_token_match(incoming, 7), 6); - // A full match ending on a separator is kept in full. - let len = incoming.chars().count(); - assert_eq!(complete_token_match(incoming, len), len); + let input = incoming.chars().count(); + assert_eq!(token_aware_match_rate(incoming, 7, input), 2.0 / 3.0); + // Full match -> 1.0. + assert_eq!(token_aware_match_rate(incoming, input, input), 1.0); } #[test] - fn test_complete_token_match_leaves_text_keys_unchanged() { - // Plain-text keys have no separator, so the character match is returned - // as-is (mid-word matches are still credited, as before). - assert_eq!(complete_token_match("hello world", 8), 8); + fn test_token_aware_match_rate_text_keys_use_chars() { + // Plain-text keys have no separator -> character ratio, unchanged. + assert_eq!(token_aware_match_rate("hello world", 6, 11), 6.0 / 11.0); } #[test] From 5744822181d829bb4a5ac80746299930f49252eb Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Mon, 7 Sep 2026 14:45:46 +0200 Subject: [PATCH 6/7] fix: export TOKEN_ID_SEPARATOR Signed-off-by: simondanielsson --- src/policies/cache_aware.rs | 10 ++++------ src/protocols/spec.rs | 12 ++++++++---- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/src/policies/cache_aware.rs b/src/policies/cache_aware.rs index 7faf1c04..8641b45a 100644 --- a/src/policies/cache_aware.rs +++ b/src/policies/cache_aware.rs @@ -63,6 +63,7 @@ use super::{get_healthy_worker_indices, CacheAwareConfig, LoadBalancingPolicy, R use crate::core::Worker; use crate::metrics::RouterMetrics; use crate::policies::normalize_model_key; +use crate::protocols::spec::TOKEN_ID_SEPARATOR; use crate::tree::Tree; use dashmap::DashMap; use rand::Rng; @@ -72,23 +73,20 @@ use std::thread; use std::time::Duration; use tracing::{debug, info}; -/// Unit separator that `PromptInput` uses to terminate each encoded token id. -const TOKEN_SEPARATOR: char = '\u{1f}'; - /// Cache-match ratio, computed per token for token-id routing keys and otherwise for chars. fn token_aware_match_rate(text: &str, matched_char_count: usize, input_char_count: usize) -> f32 { if input_char_count == 0 { return 0.0; } - if text.contains(TOKEN_SEPARATOR) { - let total_tokens = text.matches(TOKEN_SEPARATOR).count(); + if text.contains(TOKEN_ID_SEPARATOR) { + let total_tokens = text.matches(TOKEN_ID_SEPARATOR).count(); if total_tokens == 0 { return 0.0; } let matched_tokens = text .chars() .take(matched_char_count) - .filter(|c| *c == TOKEN_SEPARATOR) + .filter(|c| *c == TOKEN_ID_SEPARATOR) .count(); return matched_tokens as f32 / total_tokens as f32; } diff --git a/src/protocols/spec.rs b/src/protocols/spec.rs index eadf0931..90cf8ae4 100644 --- a/src/protocols/spec.rs +++ b/src/protocols/spec.rs @@ -2336,11 +2336,15 @@ pub enum PromptInput { String(String), } +/// Unit separator that terminates each encoded token id, so cache-aware routing +/// (a character radix) credits only complete token IDs rather than shared +/// leading digits of two different tokens. Never present in normal text. +pub(crate) const TOKEN_ID_SEPARATOR: char = '\u{1f}'; + fn encode_token_ids(ids: &[i32]) -> String { - // Terminate every token id with a unit separator (never present in text) so - // cache-aware routing, a character radix, credits only complete token IDs - // rather than shared leading digits of two different tokens. - ids.iter().map(|id| format!("{id}\u{1f}")).collect() + ids.iter() + .map(|id| format!("{id}{TOKEN_ID_SEPARATOR}")) + .collect() } impl PromptInput { From 6339867cbdb21f3166a3fe64b427fb2395f8f000 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 10 Sep 2026 15:21:11 +0200 Subject: [PATCH 7/7] fix: add explicit token id starting token Signed-off-by: simondanielsson --- src/policies/cache_aware.rs | 29 +++++++++----- src/protocols/spec.rs | 78 ++++++++++++++++++++++++++++--------- tests/test_prompt_input.rs | 6 +-- 3 files changed, 83 insertions(+), 30 deletions(-) diff --git a/src/policies/cache_aware.rs b/src/policies/cache_aware.rs index 8641b45a..a612933a 100644 --- a/src/policies/cache_aware.rs +++ b/src/policies/cache_aware.rs @@ -63,7 +63,7 @@ use super::{get_healthy_worker_indices, CacheAwareConfig, LoadBalancingPolicy, R use crate::core::Worker; use crate::metrics::RouterMetrics; use crate::policies::normalize_model_key; -use crate::protocols::spec::TOKEN_ID_SEPARATOR; +use crate::protocols::spec::{is_token_id_key, TOKEN_ID_SEPARATOR}; use crate::tree::Tree; use dashmap::DashMap; use rand::Rng; @@ -78,7 +78,7 @@ fn token_aware_match_rate(text: &str, matched_char_count: usize, input_char_coun if input_char_count == 0 { return 0.0; } - if text.contains(TOKEN_ID_SEPARATOR) { + if is_token_id_key(text) { let total_tokens = text.matches(TOKEN_ID_SEPARATOR).count(); if total_tokens == 0 { return 0.0; @@ -499,29 +499,40 @@ mod tests { #[test] fn test_token_aware_match_rate_ignores_partial_token_ids() { - // "123457\u{1f}8\u{1f}" vs cached "123456\u{1f}7\u{1f}" share 5 chars but - // zero complete tokens -> rate 0. - let incoming = "123457\u{1f}8\u{1f}"; + // "\u{1e}123457\u{1f}8\u{1f}" vs cached "\u{1e}123456\u{1f}7\u{1f}" share 6 chars + // but zero complete tokens -> rate 0. + let incoming = "\u{1e}123457\u{1f}8\u{1f}"; let input = incoming.chars().count(); - assert_eq!(token_aware_match_rate(incoming, 5, input), 0.0); + assert_eq!(token_aware_match_rate(incoming, 6, input), 0.0); } #[test] fn test_token_aware_match_rate_counts_whole_tokens() { // Matched through two separators plus a partial third token -> 2 of 3. - let incoming = "10\u{1f}20\u{1f}30\u{1f}"; + let incoming = "\u{1e}10\u{1f}20\u{1f}30\u{1f}"; let input = incoming.chars().count(); - assert_eq!(token_aware_match_rate(incoming, 7, input), 2.0 / 3.0); + assert_eq!(token_aware_match_rate(incoming, 8, input), 2.0 / 3.0); // Full match -> 1.0. assert_eq!(token_aware_match_rate(incoming, input, input), 1.0); } #[test] fn test_token_aware_match_rate_text_keys_use_chars() { - // Plain-text keys have no separator -> character ratio, unchanged. + // Plain-text keys are scored by character ratio, unchanged. assert_eq!(token_aware_match_rate("hello world", 6, 11), 6.0 / 11.0); } + #[test] + fn test_token_aware_match_rate_untagged_separator_text_uses_chars() { + // A prompt containing a raw separator is still text, not a token-id key. + let incoming = "hello\u{1f}world"; + let input = incoming.chars().count(); + assert_eq!( + token_aware_match_rate(incoming, 6, input), + 6.0 / input as f32 + ); + } + #[test] fn test_cache_aware_with_balanced_load() { // Create policy without eviction thread for testing diff --git a/src/protocols/spec.rs b/src/protocols/spec.rs index 90cf8ae4..6bf254d5 100644 --- a/src/protocols/spec.rs +++ b/src/protocols/spec.rs @@ -1,6 +1,7 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use std::collections::HashMap; +use std::fmt::Write; // # Protocol Specifications // @@ -2001,11 +2002,7 @@ impl GenerationRequest for InferenceGenerateRequest { } fn extract_text_for_routing(&self) -> String { - self.token_ids - .iter() - .map(|id| id.to_string()) - .collect::>() - .join(" ") + encode_token_ids(&self.token_ids) } } @@ -2336,15 +2333,50 @@ pub enum PromptInput { String(String), } +/// Record separator that opens every encoded token-id routing key, and separates the +/// sequences of a batch. Routing keys that do not start with it are user text. +pub(crate) const TOKEN_ID_TAG: char = '\u{1e}'; + /// Unit separator that terminates each encoded token id, so cache-aware routing /// (a character radix) credits only complete token IDs rather than shared -/// leading digits of two different tokens. Never present in normal text. +/// leading digits of two different tokens. pub(crate) const TOKEN_ID_SEPARATOR: char = '\u{1f}'; -fn encode_token_ids(ids: &[i32]) -> String { - ids.iter() - .map(|id| format!("{id}{TOKEN_ID_SEPARATOR}")) - .collect() +/// Upper bound on the characters one encoded token id occupies (`-2147483648` plus separator). +const TOKEN_ID_ENCODED_LEN: usize = 12; + +fn write_token_ids(out: &mut String, ids: &[i32]) { + for id in ids { + let _ = write!(out, "{id}{TOKEN_ID_SEPARATOR}"); + } +} + +pub(crate) fn encode_token_ids(ids: &[i32]) -> String { + let mut out = String::with_capacity(1 + ids.len() * TOKEN_ID_ENCODED_LEN); + out.push(TOKEN_ID_TAG); + write_token_ids(&mut out, ids); + out +} + +pub(crate) fn encode_token_id_batches(batches: &[Vec]) -> String { + let total_ids: usize = batches.iter().map(|b| b.len()).sum(); + let mut out = String::with_capacity(batches.len() + total_ids * TOKEN_ID_ENCODED_LEN); + for batch in batches { + out.push(TOKEN_ID_TAG); + write_token_ids(&mut out, batch); + } + out +} + +/// Whether a routing key is an encoded token-id sequence rather than user text. +/// +/// Only keys that carry the tag and consist entirely of encoded-token characters +/// qualify, so arbitrary prompt text can never be scored as if it were token ids. +pub(crate) fn is_token_id_key(key: &str) -> bool { + let mut chars = key.chars(); + chars.next() == Some(TOKEN_ID_TAG) + && chars + .all(|c| c.is_ascii_digit() || c == '-' || c == TOKEN_ID_SEPARATOR || c == TOKEN_ID_TAG) } impl PromptInput { @@ -2375,11 +2407,7 @@ impl PromptInput { PromptInput::String(s) => s.clone(), PromptInput::StringArray(arr) => arr.join(" "), PromptInput::IntArray(ids) => encode_token_ids(ids), - PromptInput::IntBatch(batches) => batches - .iter() - .map(|b| encode_token_ids(b)) - .collect::>() - .join(";"), + PromptInput::IntBatch(batches) => encode_token_id_batches(batches), } } @@ -2429,7 +2457,10 @@ mod tests { "sampling_params": {"max_tokens": 4096, "seed": 42, "logprobs": 1} }); let req: InferenceGenerateRequest = serde_json::from_value(body).unwrap(); - assert_eq!(req.extract_text_for_routing(), "151644 8948 198 2610"); + assert_eq!( + req.extract_text_for_routing(), + "\u{1e}151644\u{1f}8948\u{1f}198\u{1f}2610\u{1f}" + ); } #[test] @@ -2437,15 +2468,26 @@ mod tests { let single = PromptInput::IntArray(vec![151644, 8948, 198, 2610]); assert_eq!( single.extract_text_for_routing(), - "151644\u{1f}8948\u{1f}198\u{1f}2610\u{1f}" + "\u{1e}151644\u{1f}8948\u{1f}198\u{1f}2610\u{1f}" ); let batch = PromptInput::IntBatch(vec![vec![1, 2], vec![3, 4]]); assert_eq!( batch.extract_text_for_routing(), - "1\u{1f}2\u{1f};3\u{1f}4\u{1f}" + "\u{1e}1\u{1f}2\u{1f}\u{1e}3\u{1f}4\u{1f}" ); } + #[test] + fn test_is_token_id_key_rejects_user_text() { + assert!(is_token_id_key( + &PromptInput::IntArray(vec![1, 2]).extract_text_for_routing() + )); + assert!(!is_token_id_key("hello\u{1f}world")); + assert!(!is_token_id_key("\u{1e}hello\u{1f}world")); + assert!(!is_token_id_key("plain text")); + assert!(!is_token_id_key("")); + } + #[test] fn test_inference_generate_lossless_passthrough() { let body = serde_json::json!({ diff --git a/tests/test_prompt_input.rs b/tests/test_prompt_input.rs index cc4077fc..5be1321a 100644 --- a/tests/test_prompt_input.rs +++ b/tests/test_prompt_input.rs @@ -69,7 +69,7 @@ fn test_prompt_input_single_int_array() { assert!(req.prompt.is_token_based()); assert_eq!( req.prompt.extract_text_for_routing(), - "128000\u{1f}9906\u{1f}11\u{1f}1917\u{1f}0\u{1f}" + "\u{1e}128000\u{1f}9906\u{1f}11\u{1f}1917\u{1f}0\u{1f}" ); assert_eq!(req.prompt.estimated_token_count(), 5); } @@ -101,8 +101,8 @@ fn test_prompt_input_int_batch() { assert!(req.prompt.is_token_based()); assert_eq!( req.prompt.extract_text_for_routing(), - "128000\u{1f}9906\u{1f}11\u{1f}1917\u{1f}0\u{1f}\ - ;128001\u{1f}9906\u{1f}11\u{1f}1917\u{1f}1\u{1f}" + "\u{1e}128000\u{1f}9906\u{1f}11\u{1f}1917\u{1f}0\u{1f}\ + \u{1e}128001\u{1f}9906\u{1f}11\u{1f}1917\u{1f}1\u{1f}" ); assert_eq!(req.prompt.estimated_token_count(), 10); }