From 259ff891be37fa1af2c2c8c510becc8254569149 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Tue, 16 Jun 2026 13:30:25 +0800 Subject: [PATCH] [Rust Frontend] Require `ModelConfig.vocab_size` to be present (#45696) Signed-off-by: Bugen Zhao --- rust/src/chat/src/lib.rs | 4 +- rust/src/text/src/backend/hf/config.rs | 84 +++++++++++--------------- rust/src/text/src/backend/hf/mod.rs | 8 ++- rust/src/text/src/backend/mod.rs | 28 +++------ rust/src/text/src/lib.rs | 6 +- rust/src/text/src/lower.rs | 49 ++------------- rust/src/text/src/lower/logprobs.rs | 2 +- rust/src/text/src/lower/token_ids.rs | 14 +++-- 8 files changed, 69 insertions(+), 126 deletions(-) diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index e66db04c22..012307758c 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -145,8 +145,8 @@ impl ChatLlm { self.text.tokenizer_vocab_size() } - /// Model vocabulary size, else `None`. - pub fn model_vocab_size(&self) -> Option { + /// Model vocabulary size from the model config. + pub fn model_vocab_size(&self) -> usize { self.text.model_vocab_size() } diff --git a/rust/src/text/src/backend/hf/config.rs b/rust/src/text/src/backend/hf/config.rs index 65055722be..fbf796b5a7 100644 --- a/rust/src/text/src/backend/hf/config.rs +++ b/rust/src/text/src/backend/hf/config.rs @@ -91,9 +91,7 @@ impl HfSpecialTokens { #[serde(default)] pub struct ModelConfig { model_type: Option, - max_position_embeddings: Option, vocab_size: Option, - num_attention_heads: Option, num_experts: Option, moe_num_experts: Option, n_routed_experts: Option, @@ -180,29 +178,18 @@ impl ModelConfig { self.model_type.as_deref().or_else(|| self.text_config.as_deref()?.model_type()) } - /// Return the effective model vocabulary size, following the same simplified - /// text-config selection as `model_type`: the top-level config wins, - /// otherwise a single nested `text_config` may provide it. - pub fn vocab_size(&self) -> Option { - self.vocab_size.or_else(|| self.text_config.as_deref()?.vocab_size()) - } - - /// Reject partially nested `text_config` payloads that are unlikely to be - /// valid LLM configs for our current use. - /// - /// This keeps the simplified Rust-side parsing honest: if a model declares - /// `text_config`, it must at least look like a real text model config. - fn validate_text_config_selection(&self) -> Result<()> { - if let Some(text_config) = self.text_config.as_deref() - && text_config.num_attention_heads.is_none() - { - return Err(Error::Tokenizer( - "the text config extracted from the model config does not have `num_attention_heads`" - .to_string(), - )); + /// Return the effective model vocabulary size, following the same + /// simplified text-config selection as `model_type`. + pub fn vocab_size(&self) -> Result { + if let Some(vocab_size) = self.vocab_size { + Ok(vocab_size) + } else if let Some(text_config) = self.text_config.as_deref() { + text_config.vocab_size() + } else { + Err(Error::Tokenizer( + "the model config does not define `vocab_size`".to_string(), + )) } - - Ok(()) } /// Match Python's current expert-count priority on the selected text @@ -259,9 +246,7 @@ pub(super) fn load_generation_config(path: Option<&Path>) -> Result) -> Result { - let config: ModelConfig = read_json_file(path)?; - config.validate_text_config_selection()?; - Ok(config) + read_json_file(path) } fn read_json_file(path: Option<&Path>) -> Result @@ -339,12 +324,9 @@ mod tests { r#"{ "model_type": "top_level", "num_experts": 64, - "max_position_embeddings": 8192, "text_config": { "model_type": "nested", - "num_attention_heads": 32, - "num_local_experts": 8, - "max_position_embeddings": 4096 + "num_local_experts": 8 } }"#, ) @@ -352,32 +334,36 @@ mod tests { assert_eq!(config.num_experts(), 8); assert_eq!(config.model_type(), Some("top_level")); - assert_eq!( - config.effective_text_config().max_position_embeddings, - Some(4096) - ); assert!(config.is_moe()); } #[test] - fn model_config_defaults_to_non_moe_when_no_expert_metadata_exists() { - let config: ModelConfig = - serde_json::from_str(r#"{"max_position_embeddings":4096}"#).unwrap(); + fn model_config_uses_nested_vocab_size_when_top_level_is_absent() { + let config: ModelConfig = serde_json::from_str( + r#"{ + "text_config": { + "vocab_size": 151936 + } + }"#, + ) + .unwrap(); - assert_eq!(config.num_experts(), 0); - assert!(!config.is_moe()); - assert_eq!( - config.effective_text_config().max_position_embeddings, - Some(4096) - ); + assert_eq!(config.vocab_size().unwrap(), 151936); } #[test] - fn model_config_rejects_nested_text_config_without_attention_heads() { - let config: ModelConfig = - serde_json::from_str(r#"{"text_config":{"max_position_embeddings":4096}}"#).unwrap(); + fn model_config_rejects_missing_vocab_size() { + let config: ModelConfig = serde_json::from_str(r#"{}"#).unwrap(); - let error = config.validate_text_config_selection().unwrap_err(); - assert!(error.to_string().contains("does not have `num_attention_heads`"),); + let error = config.vocab_size().unwrap_err(); + assert!(error.to_string().contains("does not define `vocab_size`")); + } + + #[test] + fn model_config_defaults_to_non_moe_when_no_expert_metadata_exists() { + let config: ModelConfig = serde_json::from_str(r#"{}"#).unwrap(); + + assert_eq!(config.num_experts(), 0); + assert!(!config.is_moe()); } } diff --git a/rust/src/text/src/backend/hf/mod.rs b/rust/src/text/src/backend/hf/mod.rs index 94241ea74d..0e8a9bd3c0 100644 --- a/rust/src/text/src/backend/hf/mod.rs +++ b/rust/src/text/src/backend/hf/mod.rs @@ -36,6 +36,8 @@ pub struct HfTextBackend { /// Generation-config for sampling defaults that may be inherited when the /// user does not explicitly override them. generation_config: GenerationConfig, + /// Model vocabulary size from the selected text config. + model_vocab_size: usize, /// Model config (`config.json`). model_config: ModelConfig, } @@ -58,6 +60,7 @@ impl HfTextBackend { .and_then(|token| tokenizer.token_to_id(token.as_str())); let model_config = load_model_config(files.config_path.as_deref())?; + let model_vocab_size = model_config.vocab_size()? as usize; let generation_config = load_generation_config(files.generation_config_path.as_deref())?; let mut extra_eos_token_ids = generation_config .eos_token_id @@ -80,6 +83,7 @@ impl HfTextBackend { primary_eos_token_id, extra_eos_token_ids, generation_config, + model_vocab_size, model_config, }) } @@ -100,8 +104,8 @@ impl TextBackend for HfTextBackend { self.model_config.is_moe() } - fn model_vocab_size(&self) -> Option { - self.model_config.vocab_size().map(|v| v as usize) + fn model_vocab_size(&self) -> usize { + self.model_vocab_size } fn model_id(&self) -> &str { diff --git a/rust/src/text/src/backend/mod.rs b/rust/src/text/src/backend/mod.rs index 06be187414..8bc834aeae 100644 --- a/rust/src/text/src/backend/mod.rs +++ b/rust/src/text/src/backend/mod.rs @@ -30,9 +30,9 @@ pub struct SamplingLimits { /// `-1` means allowing requests up to the model vocabulary size. pub max_logprobs: i32, - /// Model vocabulary size from the model config, used to bound - /// `logit_bias` keys when available. - pub model_vocab_size: Option, + /// Model vocabulary size from the model config, used to bound generated + /// token IDs and logits-domain sampling controls. + pub model_vocab_size: usize, /// Tokenizer vocabulary size, used to bound `allowed_token_ids` and /// token-ID prompts. pub tokenizer_vocab_size: usize, @@ -46,19 +46,9 @@ impl SamplingLimits { /// pub const MAX_LOGPROB_TOKEN_IDS: usize = 128; - /// Return the vocabulary size used to expand `logprobs=-1`. - pub fn logprobs_vocab_size(&self) -> usize { - self.model_vocab_size.unwrap_or(self.tokenizer_vocab_size) - } - - /// Return the vocabulary size used to validate generated stop token IDs. - pub fn stop_token_vocab_size(&self) -> usize { - self.model_vocab_size.unwrap_or(self.tokenizer_vocab_size) - } - /// Return the union bound used to validate token-ID prompts. pub fn prompt_token_vocab_size(&self) -> usize { - self.tokenizer_vocab_size.max(self.model_vocab_size.unwrap_or(0)) + self.tokenizer_vocab_size.max(self.model_vocab_size) } } @@ -81,10 +71,12 @@ pub trait TextBackend: Send + Sync { Ok(SamplingHints::default()) } - /// Return the model vocabulary size from the model config, if known. Used to - /// range-check request token ids against the engine embedding table. - fn model_vocab_size(&self) -> Option { - None + /// Return the model vocabulary size from the model config. + /// + /// The permissive default exists for lightweight test backends. Production + /// backends should override it with the resolved model config value. + fn model_vocab_size(&self) -> usize { + usize::MAX } /// Return the full tokenizer vocabulary size (Python `len(tokenizer)`). diff --git a/rust/src/text/src/lib.rs b/rust/src/text/src/lib.rs index 4085f90478..a4fb86d19c 100644 --- a/rust/src/text/src/lib.rs +++ b/rust/src/text/src/lib.rs @@ -97,9 +97,9 @@ impl TextLlm { self.backend.tokenizer_vocab_size() } - /// Model vocabulary size from the model config, used to bound `logit_bias` - /// keys and token-id prompts against the engine embedding table. - pub fn model_vocab_size(&self) -> Option { + /// Model vocabulary size from the model config, used to bound generated + /// token IDs and logits-domain sampling controls. + pub fn model_vocab_size(&self) -> usize { self.backend.model_vocab_size() } diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index 082809fe5c..d75eb3b941 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -313,7 +313,7 @@ mod tests { SamplingLimits { max_model_len: 1_000_000, max_logprobs: SamplingLimits::DEFAULT_MAX_LOGPROBS, - model_vocab_size: Some(1000), + model_vocab_size: 1000, tokenizer_vocab_size: 2000, } } @@ -442,7 +442,7 @@ mod tests { vec![1500], sample_sampling_hints(), SamplingLimits { - model_vocab_size: Some(2000), + model_vocab_size: 2000, tokenizer_vocab_size: 1000, ..sample_sampling_limits() }, @@ -455,7 +455,7 @@ mod tests { vec![1500], sample_sampling_hints(), SamplingLimits { - model_vocab_size: Some(1000), + model_vocab_size: 1000, tokenizer_vocab_size: 2000, ..sample_sampling_limits() }, @@ -468,7 +468,7 @@ mod tests { vec![2000], sample_sampling_hints(), SamplingLimits { - model_vocab_size: Some(1000), + model_vocab_size: 1000, tokenizer_vocab_size: 2000, ..sample_sampling_limits() }, @@ -763,32 +763,6 @@ mod tests { assert_eq!(params.logprobs, Some(-1)); } - #[test] - fn lower_sampling_params_uses_tokenizer_vocab_when_model_vocab_is_unknown() { - let error = lower_sampling_params_with_limits( - SamplingParams { - logprobs: Some(-1), - ..Default::default() - }, - SamplingLimits { - max_logprobs: 1500, - model_vocab_size: None, - tokenizer_vocab_size: 2000, - ..sample_sampling_limits() - }, - ) - .unwrap_err(); - - assert!(matches!( - error, - Error::Logprobs(LogprobsError::TooManyCount { - parameter: "logprobs", - requested: 2000, - max_allowed: 1500, - }) - )); - } - #[test] fn lower_sampling_params_rejects_invalid_logprob_token_ids() { let error = lower_sampling_params_with_limits( @@ -874,21 +848,6 @@ mod tests { )); } - #[test] - fn lower_sampling_params_skips_logit_bias_range_when_model_vocab_is_unknown() { - lower_sampling_params_with_limits( - SamplingParams { - logit_bias: Some(HashMap::from([(1_000_000, 1.0)])), - ..Default::default() - }, - SamplingLimits { - model_vocab_size: None, - ..sample_sampling_limits() - }, - ) - .expect("logit_bias range check is skipped without model vocab size"); - } - #[test] fn lower_sampling_params_uses_generation_defaults_when_user_omits_values() { let params = lower_sampling_params( diff --git a/rust/src/text/src/lower/logprobs.rs b/rust/src/text/src/lower/logprobs.rs index 087f4dce2d..3c90f33910 100644 --- a/rust/src/text/src/lower/logprobs.rs +++ b/rust/src/text/src/lower/logprobs.rs @@ -40,7 +40,7 @@ pub(super) fn validate_logprobs( logprob_token_ids: Option<&[u32]>, sampling_limits: SamplingLimits, ) -> Result<(), LogprobsError> { - let vocab_size = sampling_limits.logprobs_vocab_size(); + let vocab_size = sampling_limits.model_vocab_size; let max_logprobs = normalize_logprobs_count(sampling_limits.max_logprobs, vocab_size, "max_logprobs")?; diff --git a/rust/src/text/src/lower/token_ids.rs b/rust/src/text/src/lower/token_ids.rs index 0e8dc8ff87..d24b46d4cc 100644 --- a/rust/src/text/src/lower/token_ids.rs +++ b/rust/src/text/src/lower/token_ids.rs @@ -58,7 +58,7 @@ pub(crate) fn validate_vocab_range( validate_param( "stop_token_ids", params.stop_token_ids.iter().copied(), - limits.stop_token_vocab_size(), + limits.model_vocab_size, )?; if let Some(token_ids) = params.allowed_token_ids.as_deref() { @@ -69,17 +69,19 @@ pub(crate) fn validate_vocab_range( )?; } - if let (Some(logit_bias), Some(vocab_size)) = - (params.logit_bias.as_ref(), limits.model_vocab_size) - { - validate_param("logit_bias", logit_bias.keys().copied(), vocab_size)?; + if let Some(logit_bias) = params.logit_bias.as_ref() { + validate_param( + "logit_bias", + logit_bias.keys().copied(), + limits.model_vocab_size, + )?; } if let Some(token_ids) = params.logprob_token_ids.as_deref() { validate_param( "logprob_token_ids", token_ids.iter().copied(), - limits.logprobs_vocab_size(), + limits.model_vocab_size, )?; }