[Rust Frontend] Require ModelConfig.vocab_size to be present

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<usize> {
+    /// Model vocabulary size from the model config.
+    pub fn model_vocab_size(&self) -> usize {
         self.text.model_vocab_size()
     }
 

rust/src/text/src/backend/hf/config.rs

@@ -91,9 +91,7 @@ impl HfSpecialTokens {
 #[serde(default)]
 pub struct ModelConfig {
     model_type: Option<String>,
-    max_position_embeddings: Option<u32>,
     vocab_size: Option<u32>,
-    num_attention_heads: Option<u32>,
     num_experts: Option<OneOrManyExpertCount>,
     moe_num_experts: Option<OneOrManyExpertCount>,
     n_routed_experts: Option<OneOrManyExpertCount>,
@@ -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<u32> {
-        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<u32> {
+        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<GenerationCo
 
 /// Load the model-side config (`config.json`) if present.
 pub fn load_model_config(path: Option<&Path>) -> Result<ModelConfig> {
-    let config: ModelConfig = read_json_file(path)?;
-    config.validate_text_config_selection()?;
-    Ok(config)
+    read_json_file(path)
 }
 
 fn read_json_file<T>(path: Option<&Path>) -> Result<T>
@@ -339,45 +324,46 @@ 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
                 }
             }"#,
         )
         .unwrap();
 
         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.vocab_size().unwrap_err();
+        assert!(error.to_string().contains("does not define `vocab_size`"));
+    }
 
-        let error = config.validate_text_config_selection().unwrap_err();
-        assert!(error.to_string().contains("does not have `num_attention_heads`"),);
+    #[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());
     }
 }

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<usize> {
-        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 {

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<usize>,
+    /// 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 {
     /// <https://github.com/vllm-project/vllm/blob/b5adb027ad03c29b46181752ba3b1cb84eff1dd4/vllm/sampling_params.py#L30-L32>
     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<usize> {
-        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)`).

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<usize> {
+    /// 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()
     }
 

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(

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")?;
 

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,
         )?;
     }