@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 259ff891be

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

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

Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Bugen Zhao committed Jun 16, 2026 at 13:30 UTC 259ff891be37fa1af2c2c8c510becc8254569149

8 files changed +69 -126

rust/src/chat/src/lib.rs

+2 -2

@@ -145,8 +145,8 @@ impl ChatLlm {

145 self.text.tokenizer_vocab_size()

146 }

147

148 - /// Model vocabulary size, else `None`.

149 - pub fn model_vocab_size(&self) -> Option<usize> {

148 + /// Model vocabulary size from the model config.

149 + pub fn model_vocab_size(&self) -> usize {

150 self.text.model_vocab_size()

151 }

152

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

+35 -49

@@ -91,9 +91,7 @@ impl HfSpecialTokens {

91 #[serde(default)]

92 pub struct ModelConfig {

93 model_type: Option<String>,

94 - max_position_embeddings: Option<u32>,

94 vocab_size: Option<u32>,

96 - num_attention_heads: Option<u32>,

95 num_experts: Option<OneOrManyExpertCount>,

96 moe_num_experts: Option<OneOrManyExpertCount>,

97 n_routed_experts: Option<OneOrManyExpertCount>,

@@ -180,29 +178,18 @@ impl ModelConfig {

178 self.model_type.as_deref().or_else(|| self.text_config.as_deref()?.model_type())

179 }

180

183 - /// Return the effective model vocabulary size, following the same simplified

184 - /// text-config selection as `model_type`: the top-level config wins,

185 - /// otherwise a single nested `text_config` may provide it.

186 - pub fn vocab_size(&self) -> Option<u32> {

187 - self.vocab_size.or_else(|| self.text_config.as_deref()?.vocab_size())

188 - }

189 -

190 - /// Reject partially nested `text_config` payloads that are unlikely to be

191 - /// valid LLM configs for our current use.

192 - ///

193 - /// This keeps the simplified Rust-side parsing honest: if a model declares

194 - /// `text_config`, it must at least look like a real text model config.

195 - fn validate_text_config_selection(&self) -> Result<()> {

196 - if let Some(text_config) = self.text_config.as_deref()

197 - && text_config.num_attention_heads.is_none()

198 - {

199 - return Err(Error::Tokenizer(

200 - "the text config extracted from the model config does not have `num_attention_heads`"

201 - .to_string(),

202 - ));

181 + /// Return the effective model vocabulary size, following the same

182 + /// simplified text-config selection as `model_type`.

183 + pub fn vocab_size(&self) -> Result<u32> {

184 + if let Some(vocab_size) = self.vocab_size {

185 + Ok(vocab_size)

186 + } else if let Some(text_config) = self.text_config.as_deref() {

187 + text_config.vocab_size()

188 + } else {

189 + Err(Error::Tokenizer(

190 + "the model config does not define `vocab_size`".to_string(),

191 + ))

192 }

204 -

205 - Ok(())

193 }

194

195 /// 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

246

247 /// Load the model-side config (`config.json`) if present.

248 pub fn load_model_config(path: Option<&Path>) -> Result<ModelConfig> {

262 - let config: ModelConfig = read_json_file(path)?;

263 - config.validate_text_config_selection()?;

264 - Ok(config)

249 + read_json_file(path)

250 }

251

252 fn read_json_file<T>(path: Option<&Path>) -> Result<T>

@@ -339,12 +324,9 @@ mod tests {

324 r#"{

325 "model_type": "top_level",

326 "num_experts": 64,

342 - "max_position_embeddings": 8192,

327 "text_config": {

328 "model_type": "nested",

345 - "num_attention_heads": 32,

346 - "num_local_experts": 8,

347 - "max_position_embeddings": 4096

329 + "num_local_experts": 8

330 }

331 }"#,

332 )

@@ -352,32 +334,36 @@ mod tests {

334

335 assert_eq!(config.num_experts(), 8);

336 assert_eq!(config.model_type(), Some("top_level"));

355 - assert_eq!(

356 - config.effective_text_config().max_position_embeddings,

357 - Some(4096)

358 - );

337 assert!(config.is_moe());

338 }

339

340 #[test]

363 - fn model_config_defaults_to_non_moe_when_no_expert_metadata_exists() {

364 - let config: ModelConfig =

365 - serde_json::from_str(r#"{"max_position_embeddings":4096}"#).unwrap();

341 + fn model_config_uses_nested_vocab_size_when_top_level_is_absent() {

342 + let config: ModelConfig = serde_json::from_str(

343 + r#"{

344 + "text_config": {

345 + "vocab_size": 151936

346 + }

347 + }"#,

348 + )

349 + .unwrap();

350

367 - assert_eq!(config.num_experts(), 0);

368 - assert!(!config.is_moe());

369 - assert_eq!(

370 - config.effective_text_config().max_position_embeddings,

371 - Some(4096)

372 - );

351 + assert_eq!(config.vocab_size().unwrap(), 151936);

352 }

353

354 #[test]

376 - fn model_config_rejects_nested_text_config_without_attention_heads() {

377 - let config: ModelConfig =

378 - serde_json::from_str(r#"{"text_config":{"max_position_embeddings":4096}}"#).unwrap();

355 + fn model_config_rejects_missing_vocab_size() {

356 + let config: ModelConfig = serde_json::from_str(r#"{}"#).unwrap();

357 +

358 + let error = config.vocab_size().unwrap_err();

359 + assert!(error.to_string().contains("does not define `vocab_size`"));

360 + }

361

380 - let error = config.validate_text_config_selection().unwrap_err();

381 - assert!(error.to_string().contains("does not have `num_attention_heads`"),);

362 + #[test]

363 + fn model_config_defaults_to_non_moe_when_no_expert_metadata_exists() {

364 + let config: ModelConfig = serde_json::from_str(r#"{}"#).unwrap();

365 +

366 + assert_eq!(config.num_experts(), 0);

367 + assert!(!config.is_moe());

368 }

369 }

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

+6 -2

@@ -36,6 +36,8 @@ pub struct HfTextBackend {

36 /// Generation-config for sampling defaults that may be inherited when the

37 /// user does not explicitly override them.

38 generation_config: GenerationConfig,

39 + /// Model vocabulary size from the selected text config.

40 + model_vocab_size: usize,

41 /// Model config (`config.json`).

42 model_config: ModelConfig,

43 }

@@ -58,6 +60,7 @@ impl HfTextBackend {

60 .and_then(|token| tokenizer.token_to_id(token.as_str()));

62 let model_config = load_model_config(files.config_path.as_deref())?;

63 + let model_vocab_size = model_config.vocab_size()? as usize;

64 let generation_config = load_generation_config(files.generation_config_path.as_deref())?;

65 let mut extra_eos_token_ids = generation_config

66 .eos_token_id

@@ -80,6 +83,7 @@ impl HfTextBackend {

83 primary_eos_token_id,

84 extra_eos_token_ids,

85 generation_config,

86 + model_vocab_size,

87 model_config,

88 })

89 }

@@ -100,8 +104,8 @@ impl TextBackend for HfTextBackend {

104 self.model_config.is_moe()

105 }

106

103 - fn model_vocab_size(&self) -> Option<usize> {

104 - self.model_config.vocab_size().map(|v| v as usize)

107 + fn model_vocab_size(&self) -> usize {

108 + self.model_vocab_size

109 }

110

111 fn model_id(&self) -> &str {

rust/src/text/src/backend/mod.rs

+10 -18

@@ -30,9 +30,9 @@ pub struct SamplingLimits {

30 /// `-1` means allowing requests up to the model vocabulary size.

31 pub max_logprobs: i32,

33 - /// Model vocabulary size from the model config, used to bound

34 - /// `logit_bias` keys when available.

35 - pub model_vocab_size: Option<usize>,

33 + /// Model vocabulary size from the model config, used to bound generated

34 + /// token IDs and logits-domain sampling controls.

35 + pub model_vocab_size: usize,

36 /// Tokenizer vocabulary size, used to bound `allowed_token_ids` and

37 /// token-ID prompts.

38 pub tokenizer_vocab_size: usize,

@@ -46,19 +46,9 @@ impl SamplingLimits {

46 /// <https://github.com/vllm-project/vllm/blob/b5adb027ad03c29b46181752ba3b1cb84eff1dd4/vllm/sampling_params.py#L30-L32>

47 pub const MAX_LOGPROB_TOKEN_IDS: usize = 128;

49 - /// Return the vocabulary size used to expand `logprobs=-1`.

50 - pub fn logprobs_vocab_size(&self) -> usize {

51 - self.model_vocab_size.unwrap_or(self.tokenizer_vocab_size)

52 - }

53 -

54 - /// Return the vocabulary size used to validate generated stop token IDs.

55 - pub fn stop_token_vocab_size(&self) -> usize {

56 - self.model_vocab_size.unwrap_or(self.tokenizer_vocab_size)

57 - }

58 -

49 /// Return the union bound used to validate token-ID prompts.

50 pub fn prompt_token_vocab_size(&self) -> usize {

61 - self.tokenizer_vocab_size.max(self.model_vocab_size.unwrap_or(0))

51 + self.tokenizer_vocab_size.max(self.model_vocab_size)

52 }

53 }

@@ -81,10 +71,12 @@ pub trait TextBackend: Send + Sync {

71 Ok(SamplingHints::default())

72 }

84 - /// Return the model vocabulary size from the model config, if known. Used to

85 - /// range-check request token ids against the engine embedding table.

86 - fn model_vocab_size(&self) -> Option<usize> {

87 - None

74 + /// Return the model vocabulary size from the model config.

75 + ///

76 + /// The permissive default exists for lightweight test backends. Production

77 + /// backends should override it with the resolved model config value.

78 + fn model_vocab_size(&self) -> usize {

79 + usize::MAX

80 }

82 /// Return the full tokenizer vocabulary size (Python `len(tokenizer)`).

rust/src/text/src/lib.rs

+3 -3

@@ -97,9 +97,9 @@ impl TextLlm {

97 self.backend.tokenizer_vocab_size()

98 }

100 - /// Model vocabulary size from the model config, used to bound `logit_bias`

101 - /// keys and token-id prompts against the engine embedding table.

102 - pub fn model_vocab_size(&self) -> Option<usize> {

100 + /// Model vocabulary size from the model config, used to bound generated

101 + /// token IDs and logits-domain sampling controls.

102 + pub fn model_vocab_size(&self) -> usize {

103 self.backend.model_vocab_size()

104 }

105

rust/src/text/src/lower.rs

+4 -45

@@ -313,7 +313,7 @@ mod tests {

313 SamplingLimits {

314 max_model_len: 1_000_000,

315 max_logprobs: SamplingLimits::DEFAULT_MAX_LOGPROBS,

316 - model_vocab_size: Some(1000),

316 + model_vocab_size: 1000,

317 tokenizer_vocab_size: 2000,

318 }

319 }

@@ -442,7 +442,7 @@ mod tests {

442 vec![1500],

443 sample_sampling_hints(),

444 SamplingLimits {

445 - model_vocab_size: Some(2000),

445 + model_vocab_size: 2000,

446 tokenizer_vocab_size: 1000,

447 ..sample_sampling_limits()

448 },

@@ -455,7 +455,7 @@ mod tests {

455 vec![1500],

456 sample_sampling_hints(),

457 SamplingLimits {

458 - model_vocab_size: Some(1000),

458 + model_vocab_size: 1000,

459 tokenizer_vocab_size: 2000,

460 ..sample_sampling_limits()

461 },

@@ -468,7 +468,7 @@ mod tests {

468 vec![2000],

469 sample_sampling_hints(),

470 SamplingLimits {

471 - model_vocab_size: Some(1000),

471 + model_vocab_size: 1000,

472 tokenizer_vocab_size: 2000,

473 ..sample_sampling_limits()

474 },

@@ -763,32 +763,6 @@ mod tests {

763 assert_eq!(params.logprobs, Some(-1));

764 }

765

766 - #[test]

767 - fn lower_sampling_params_uses_tokenizer_vocab_when_model_vocab_is_unknown() {

768 - let error = lower_sampling_params_with_limits(

769 - SamplingParams {

770 - logprobs: Some(-1),

771 - ..Default::default()

772 - },

773 - SamplingLimits {

774 - max_logprobs: 1500,

775 - model_vocab_size: None,

776 - tokenizer_vocab_size: 2000,

777 - ..sample_sampling_limits()

778 - },

779 - )

780 - .unwrap_err();

781 -

782 - assert!(matches!(

783 - error,

784 - Error::Logprobs(LogprobsError::TooManyCount {

785 - parameter: "logprobs",

786 - requested: 2000,

787 - max_allowed: 1500,

788 - })

789 - ));

790 - }

791 -

766 #[test]

767 fn lower_sampling_params_rejects_invalid_logprob_token_ids() {

768 let error = lower_sampling_params_with_limits(

@@ -874,21 +848,6 @@ mod tests {

848 ));

849 }

850

877 - #[test]

878 - fn lower_sampling_params_skips_logit_bias_range_when_model_vocab_is_unknown() {

879 - lower_sampling_params_with_limits(

880 - SamplingParams {

881 - logit_bias: Some(HashMap::from([(1_000_000, 1.0)])),

882 - ..Default::default()

883 - },

884 - SamplingLimits {

885 - model_vocab_size: None,

886 - ..sample_sampling_limits()

887 - },

888 - )

889 - .expect("logit_bias range check is skipped without model vocab size");

890 - }

891 -

851 #[test]

852 fn lower_sampling_params_uses_generation_defaults_when_user_omits_values() {

853 let params = lower_sampling_params(

rust/src/text/src/lower/logprobs.rs

+1 -1

@@ -40,7 +40,7 @@ pub(super) fn validate_logprobs(

40 logprob_token_ids: Option<&[u32]>,

41 sampling_limits: SamplingLimits,

42 ) -> Result<(), LogprobsError> {

43 - let vocab_size = sampling_limits.logprobs_vocab_size();

43 + let vocab_size = sampling_limits.model_vocab_size;

44 let max_logprobs =

45 normalize_logprobs_count(sampling_limits.max_logprobs, vocab_size, "max_logprobs")?;

rust/src/text/src/lower/token_ids.rs

+8 -6

@@ -58,7 +58,7 @@ pub(crate) fn validate_vocab_range(

58 validate_param(

59 "stop_token_ids",

60 params.stop_token_ids.iter().copied(),

61 - limits.stop_token_vocab_size(),

61 + limits.model_vocab_size,

62 )?;

64 if let Some(token_ids) = params.allowed_token_ids.as_deref() {

@@ -69,17 +69,19 @@ pub(crate) fn validate_vocab_range(

69 )?;

70 }

72 - if let (Some(logit_bias), Some(vocab_size)) =

73 - (params.logit_bias.as_ref(), limits.model_vocab_size)

74 - {

75 - validate_param("logit_bias", logit_bias.keys().copied(), vocab_size)?;

72 + if let Some(logit_bias) = params.logit_bias.as_ref() {

73 + validate_param(

74 + "logit_bias",

75 + logit_bias.keys().copied(),

76 + limits.model_vocab_size,

77 + )?;

78 }

80 if let Some(token_ids) = params.logprob_token_ids.as_deref() {

81 validate_param(

82 "logprob_token_ids",

83 token_ids.iter().copied(),

82 - limits.logprobs_vocab_size(),

84 + limits.model_vocab_size,

85 )?;

86 }