diff --git a/src/policies/cache_aware.rs b/src/policies/cache_aware.rs index dc209a19..a612933a 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::{is_token_id_key, TOKEN_ID_SEPARATOR}; use crate::tree::Tree; use dashmap::DashMap; use rand::Rng; @@ -72,6 +73,26 @@ use std::thread; use std::time::Duration; use tracing::{debug, info}; +/// 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 is_token_id_key(text) { + 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_ID_SEPARATOR) + .count(); + return matched_tokens as f32 / total_tokens as f32; + } + matched_char_count as f32 / input_char_count as f32 +} + /// Cache-aware routing policy /// /// Routes requests based on cache affinity when load is balanced, @@ -300,11 +321,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 match_rate = if result.input_char_count == 0 { - 0.0 - } else { - result.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}", @@ -479,6 +497,42 @@ mod tests { use super::*; use crate::core::{BasicWorker, WorkerType}; + #[test] + fn test_token_aware_match_rate_ignores_partial_token_ids() { + // "\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, 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 = "\u{1e}10\u{1f}20\u{1f}30\u{1f}"; + let input = incoming.chars().count(); + 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 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 3ed54dc8..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,6 +2333,52 @@ 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. +pub(crate) const TOKEN_ID_SEPARATOR: char = '\u{1f}'; + +/// 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 { /// Get the number of items in the PromptInput pub fn len(&self) -> usize { @@ -2363,16 +2406,8 @@ impl PromptInput { 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) => encode_token_id_batches(batches), } } @@ -2422,7 +2457,35 @@ 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] + fn test_prompt_input_token_ids_routing_key() { + let single = PromptInput::IntArray(vec![151644, 8948, 198, 2610]); + assert_eq!( + single.extract_text_for_routing(), + "\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(), + "\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] diff --git a/tests/test_prompt_input.rs b/tests/test_prompt_input.rs index 5f62f8ea..5be1321a 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(), + "\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); } @@ -98,7 +101,8 @@ 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" + "\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); }