iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Rust Frontend] Require ModelConfig.vocab_size to be present (#45696)

Browse Source 
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
This commit is contained in:
Bugen Zhao
2026-06-16 05:30:25 +00:00
committed by GitHub
parent 6607a80dab
commit 259ff891be
8 changed files with 69 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/src/chat/src/lib.rs
+2 -2
View File
@@ -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
+35 -49
View File
@@ -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,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());
}
}
rust/src/text/src/backend/hf/mod.rs
+6 -2
View File
@@ -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
+10 -18
View File
@@ -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
+3 -3
View File
@@ -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
+4 -45
View File
@@ -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
+1 -1
View File
@@ -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
+8 -6
View File
@@ -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,
)?;
}