diff --git a/lib/llm/src/http/service/openai.rs b/lib/llm/src/http/service/openai.rs index 479f3bd908a4..d758394c2d6d 100644 --- a/lib/llm/src/http/service/openai.rs +++ b/lib/llm/src/http/service/openai.rs @@ -43,6 +43,7 @@ use super::{ service_v2, }; use crate::engines::ValidateRequest; +use crate::preprocessor::PRESERVE_OMITTED_MAX_TOKENS_CONTEXT_KEY; use crate::protocols::openai::chat_completions::aggregator::ChatCompletionAggregator; use crate::protocols::openai::nvext::apply_header_routing_overrides; use crate::protocols::openai::{ @@ -1669,8 +1670,9 @@ async fn responses( check_ready(&state)?; // Apply template values if present. When no template and no client-supplied - // max_output_tokens, leave it as None and let the underlying engine apply its - // own default — matching the chat completions path. + // max_output_tokens, leave it as None for response echoing and let the + // backend adapter compute the dynamic generation cap from its effective + // prompt length. if let Some(template) = template { if request.inner.model.as_deref().unwrap_or("").is_empty() { request.inner.model = Some(template.model.clone()); @@ -1767,7 +1769,10 @@ async fn responses( continuous_usage_stats: false, }); - let request = context.map(|mut _req| chat_request); + let mut request = context.map(|mut _req| chat_request); + if response_params.max_output_tokens.is_none() { + request.insert(PRESERVE_OMITTED_MAX_TOKENS_CONTEXT_KEY, true); + } tracing::trace!("Getting chat completions engine for model: {}", model); diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index b30f609b02e5..361b15562505 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -206,6 +206,14 @@ static DIM_FETCH_HTTP_CLIENT: std::sync::LazyLock = .expect("dim-fetch http client construction failed") }); +pub(crate) const PRESERVE_OMITTED_MAX_TOKENS_CONTEXT_KEY: &str = + "dynamo.llm.preserve_omitted_max_tokens"; + +#[derive(Clone, Copy, Debug, Default)] +struct PreprocessRequestOptions { + preserve_omitted_max_tokens: bool, +} + pub struct OpenAIPreprocessor { mdcsum: String, formatter: Arc, @@ -240,6 +248,94 @@ pub struct OpenAIPreprocessor { } impl OpenAIPreprocessor { + fn omitted_max_tokens_default( + prompt_len: usize, + context_length: u32, + options: PreprocessRequestOptions, + ) -> Option { + if context_length == 0 || options.preserve_omitted_max_tokens { + return None; + } + Some(context_length.saturating_sub(prompt_len as u32)) + } + + fn nvext_passthrough_args( + request: &R, + ) -> Option> { + let mut nvext_passthrough = serde_json::Map::new(); + + if let Some(nvext) = request.nvext() { + if let Some(ref fields) = nvext.extra_fields { + nvext_passthrough.insert("extra_fields".to_string(), serde_json::json!(fields)); + } + if let Some(ref salt) = nvext.cache_salt { + nvext_passthrough.insert("cache_salt".to_string(), serde_json::json!(salt)); + } + if nvext.token_data.is_some() { + nvext_passthrough.insert("token_in".to_string(), serde_json::Value::Bool(true)); + } + } + + if !nvext_passthrough.contains_key("cache_salt") + && let Some(salt) = request + .unsupported_fields() + .and_then(|fields| fields.get("cache_salt")) + .and_then(|value| value.as_str()) + { + nvext_passthrough.insert("cache_salt".to_string(), serde_json::json!(salt)); + } + + if nvext_passthrough.is_empty() { + None + } else { + Some(nvext_passthrough) + } + } + + fn sampling_passthrough_args( + request: &R, + ) -> Option> { + let mut sampling_passthrough = serde_json::Map::new(); + + if let Some(fields) = request.unsupported_fields() { + for key in ["detokenize", "allowed_token_ids", "bad_words_token_ids"] { + if let Some(value) = fields.get(key) { + sampling_passthrough.insert(key.to_string(), value.clone()); + } + } + } + + if sampling_passthrough.is_empty() { + None + } else { + Some(sampling_passthrough) + } + } + + fn backend_extra_args(request: &R) -> Option { + let mut extra_args = serde_json::Map::new(); + + if let Some(nvext_passthrough) = Self::nvext_passthrough_args(request) { + extra_args.insert( + "nvext".to_string(), + serde_json::Value::Object(nvext_passthrough), + ); + } + + if let Some(sampling_passthrough) = Self::sampling_passthrough_args(request) { + extra_args.insert( + "sampling_options".to_string(), + serde_json::Value::Object(sampling_passthrough), + ); + } + + if extra_args.is_empty() { + None + } else { + Some(serde_json::Value::Object(extra_args)) + } + } + pub fn new(mdc: ModelDeploymentCard) -> Result> { let formatter = PromptFormatter::from_mdc(&mdc)?; let tokenizer = mdc.tokenizer()?; @@ -415,6 +511,23 @@ impl OpenAIPreprocessor { &self, request: &R, tracker: Option<&RequestTracker>, + ) -> Result<(PreprocessedRequest, HashMap, bool)> { + self.preprocess_request_with_options(request, tracker, PreprocessRequestOptions::default()) + .await + } + + async fn preprocess_request_with_options< + R: OAIChatLikeRequest + + AnnotationsProvider + + SamplingOptionsProvider + + StopConditionsProvider + + OutputOptionsProvider + + NvExtProvider, + >( + &self, + request: &R, + tracker: Option<&RequestTracker>, + options: PreprocessRequestOptions, ) -> Result<(PreprocessedRequest, HashMap, bool)> { let _stage_guard = StageGuard::new(STAGE_PREPROCESS, ""); let preprocess_start = Instant::now(); @@ -471,7 +584,22 @@ impl OpenAIPreprocessor { builder.router(Some(router_params.clone())); } - Ok((builder.build()?, annotations, prompt_injected_reasoning)) + let mut preprocessed = builder.build()?; + + // If omitted, allow generation up to the remaining context length. Responses requests + // preserve omission so backend adapters can compute the dynamic cap from their + // effective prompt length/tokenization. + if preprocessed.stop_conditions.max_tokens.is_none() + && let Some(max_tokens) = Self::omitted_max_tokens_default( + preprocessed.token_ids.len(), + self.context_length, + options, + ) + { + preprocessed.stop_conditions.max_tokens = Some(max_tokens); + } + + Ok((preprocessed, annotations, prompt_injected_reasoning)) } pub fn builder< @@ -563,6 +691,10 @@ impl OpenAIPreprocessor { })); } + if let Some(extra_args) = Self::backend_extra_args(request) { + builder.extra_args(Some(extra_args)); + } + // Forward mm_processor_kwargs (e.g. use_audio_in_video) to the backend. builder.mm_processor_kwargs(request.mm_processor_kwargs().cloned()); @@ -630,7 +762,7 @@ impl OpenAIPreprocessor { } } - pub async fn gather_multi_modal_data( + pub async fn gather_multi_modal_data( &self, request: &R, builder: &mut PreprocessedRequestBuilder, @@ -866,6 +998,15 @@ impl OpenAIPreprocessor { extra_args["formatted_prompt"] = serde_json::Value::String(prompt.clone()); } + if let Some(serde_json::Value::Object(backend_extra_args)) = + Self::backend_extra_args(request) + { + let extra_args_obj = extra_args + .as_object_mut() + .expect("multimodal extra_args must be an object"); + extra_args_obj.extend(backend_extra_args); + } + // Forward routing-side mm_hashes as `multi_modal_uuids` so vLLM // publishes KV events with the same key the router computes. // The kv-router parses events via parse_mm_hash_from_extra_key @@ -2417,10 +2558,16 @@ impl // create a response generator let response_generator = request.response_generator(context.id().to_string()); let tracker = Some(response_generator.tracker()); + let preprocess_options = PreprocessRequestOptions { + preserve_omitted_max_tokens: context + .get::(PRESERVE_OMITTED_MAX_TOKENS_CONTEXT_KEY) + .ok() + .is_some_and(|flag| *flag), + }; // convert the chat completion request to a common completion request let (mut common_request, annotations, prompt_injected_reasoning) = self - .preprocess_request(&request, tracker.as_deref()) + .preprocess_request_with_options(&request, tracker.as_deref(), preprocess_options) .await?; tracing::trace!(request = ?common_request, prompt_injected_reasoning, "Pre-processed request"); let trace_state = crate::agents::trace::build_agent_trace_request_end_state( @@ -2836,6 +2983,69 @@ mod tests { } } + #[test] + fn test_backend_extra_args_preserves_nvext_and_sampling_extensions() { + let request: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "hi"}], + "detokenize": false, + "allowed_token_ids": [10, 11], + "bad_words_token_ids": [[12, 13]], + "nvext": { + "cache_salt": "step_7", + "extra_fields": ["completion_token_ids"] + } + })) + .unwrap(); + + let extra_args = OpenAIPreprocessor::backend_extra_args(&request).unwrap(); + + assert_eq!(extra_args["nvext"]["cache_salt"], "step_7"); + assert_eq!( + extra_args["nvext"]["extra_fields"], + serde_json::json!(["completion_token_ids"]) + ); + assert_eq!(extra_args["sampling_options"]["detokenize"], false); + assert_eq!( + extra_args["sampling_options"]["allowed_token_ids"], + serde_json::json!([10, 11]) + ); + assert_eq!( + extra_args["sampling_options"]["bad_words_token_ids"], + serde_json::json!([[12, 13]]) + ); + } + + #[test] + fn test_internal_preserve_omitted_max_tokens_option() { + assert_eq!( + OpenAIPreprocessor::omitted_max_tokens_default( + 10, + 100, + PreprocessRequestOptions::default() + ), + Some(90) + ); + assert_eq!( + OpenAIPreprocessor::omitted_max_tokens_default( + 10, + 100, + PreprocessRequestOptions { + preserve_omitted_max_tokens: true, + }, + ), + None + ); + assert_eq!( + OpenAIPreprocessor::omitted_max_tokens_default( + 10, + 0, + PreprocessRequestOptions::default() + ), + None + ); + } + /// PRE.2 — Per-request reasoning gate. See `lib/llm/PREPROCESSOR_CASES.md`. #[test] fn test_is_reasoning_disabled_by_request() { diff --git a/lib/llm/src/protocols/common/llm_backend.rs b/lib/llm/src/protocols/common/llm_backend.rs index 8877185eedb7..25a8c6ba40ed 100644 --- a/lib/llm/src/protocols/common/llm_backend.rs +++ b/lib/llm/src/protocols/common/llm_backend.rs @@ -14,6 +14,20 @@ use dynamo_runtime::protocols::maybe_error::MaybeError; pub type TokenType = Option; pub type LogProbs = Vec; +/// Per-position prompt logprob entry reported by an engine adapter. +#[derive(Serialize, Deserialize, utoipa::ToSchema, Debug, Clone, PartialEq)] +pub struct PromptLogprobEntry { + pub logprob: f32, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub rank: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub decoded_token: Option, +} + +/// Per-token map of `token_id -> PromptLogprobEntry`. The first position +/// is `None` (no logprob exists for BOS / the very first prompt token). +pub type PromptLogprobs = Vec>>; + /// Output type discriminator for different modalities #[derive(Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq, Default)] #[serde(rename_all = "lowercase")] @@ -283,6 +297,14 @@ impl LLMEngineOutput { } } +pub(crate) fn prompt_logprobs_from_engine_data( + engine_data: Option<&serde_json::Value>, +) -> Option { + engine_data? + .get("prompt_logprobs") + .and_then(|value| serde_json::from_value(value.clone()).ok()) +} + impl MaybeError for LLMEngineOutput { fn from_err(err: impl std::error::Error + 'static) -> Self { LLMEngineOutput::error(err.to_string()) diff --git a/lib/llm/src/protocols/openai.rs b/lib/llm/src/protocols/openai.rs index e509f7832850..c15c353f9525 100644 --- a/lib/llm/src/protocols/openai.rs +++ b/lib/llm/src/protocols/openai.rs @@ -144,7 +144,6 @@ impl SamplingOptionsProvid let guided_grammar = self.get_guided_grammar(); let guided_choice = self.get_guided_choice(); let guided_whitespace_pattern = self.get_guided_whitespace_pattern(); - let guided_decoding = match common::GuidedDecodingOptions::from_optional( guided_json, guided_regex, diff --git a/lib/llm/src/protocols/openai/chat_completions.rs b/lib/llm/src/protocols/openai/chat_completions.rs index 291b837d38df..75975da6bd44 100644 --- a/lib/llm/src/protocols/openai/chat_completions.rs +++ b/lib/llm/src/protocols/openai/chat_completions.rs @@ -101,6 +101,10 @@ impl NvExtProvider for NvCreateChatCompletionRequest { fn raw_prompt(&self) -> Option { None } + + fn unsupported_fields(&self) -> Option<&std::collections::HashMap> { + Some(&self.unsupported_fields) + } } /// Implements `AnnotationsProvider` for `NvCreateChatCompletionRequest`, @@ -258,6 +262,10 @@ impl CommonExtProvider for NvCreateChatCompletionRequest { fn get_skip_special_tokens(&self) -> Option { self.common.skip_special_tokens } + + fn get_prompt_logprobs_count(&self) -> Option { + self.common.prompt_logprobs + } } /// Implements `OpenAIStopConditionsProvider` for `NvCreateChatCompletionRequest`, @@ -288,7 +296,14 @@ impl OpenAIStopConditionsProvider for NvCreateChatCompletionRequest { } fn get_stop_token_ids(&self) -> Option> { - self.inner.stop.as_ref().and_then(|stop| stop.token_ids()) + // Token IDs may be provided in the standard OpenAI `stop` array. + if let Some(ids) = self.inner.stop.as_ref().and_then(|stop| stop.token_ids()) { + return Some(ids); + } + // Also accept top-level `stop_token_ids` from passthrough clients. + self.unsupported_fields + .get("stop_token_ids") + .and_then(|v| serde_json::from_value::>(v.clone()).ok()) } /// Returns a reference to the optional `NvExt` extension, if available. @@ -320,7 +335,8 @@ impl OpenAIOutputOptionsProvider for NvCreateChatCompletionRequest { } fn get_prompt_logprobs(&self) -> Option { - None + // Top-level `prompt_logprobs` is carried through CommonExt. + self.common.prompt_logprobs } fn get_skip_special_tokens(&self) -> Option { @@ -353,6 +369,10 @@ impl ValidateRequest for NvCreateChatCompletionRequest { // validate::validate_max_tokens(self.inner.max_tokens)?; // warning depricated field validate::validate_max_completion_tokens(self.inner.max_completion_tokens)?; validate::validate_n(self.inner.n)?; + super::nvext::validate_completion_token_ids_single_choice( + self.inner.n.unwrap_or(1) as usize, + self.nvext.as_ref(), + )?; // none for modalities // none for prediction // none for audio @@ -504,14 +524,81 @@ mod tests { serde_json::from_value(scalar_token_id_stop); assert!(result.is_err()); - let unsupported_stop_token_ids = json!({ + // `stop_token_ids` is accepted and plumbed by the provider trait. + let whitelisted_stop_token_ids = json!({ "model": "test-model", "messages": [{"role": "user", "content": "Hello"}], "stop_token_ids": [576] }); let request: NvCreateChatCompletionRequest = - serde_json::from_value(unsupported_stop_token_ids) + serde_json::from_value(whitelisted_stop_token_ids) .expect("Failed to deserialize request"); + assert_eq!(request.get_stop_token_ids(), Some(vec![576])); + assert!( + ValidateRequest::validate(&request).is_ok(), + "stop_token_ids must be accepted via PASSTHROUGH_EXTRA_FIELDS" + ); + + let invalid_stop_token_ids = json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "Hello"}], + "stop_token_ids": "bad" + }); + let request: NvCreateChatCompletionRequest = + serde_json::from_value(invalid_stop_token_ids).expect("Failed to deserialize request"); + let err = ValidateRequest::validate(&request).expect_err("invalid stop_token_ids"); + assert!(err.to_string().contains("stop_token_ids")); + } + + #[test] + fn test_passthrough_token_constraints_validate() { + let request_json = json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "Hello"}], + "allowed_token_ids": [10, 11], + "bad_words_token_ids": [[12, 13]] + }); + let request: NvCreateChatCompletionRequest = + serde_json::from_value(request_json).expect("Failed to deserialize request"); + + assert_eq!( + request.unsupported_fields.get("allowed_token_ids"), + Some(&serde_json::json!([10, 11])) + ); + assert_eq!( + request.unsupported_fields.get("bad_words_token_ids"), + Some(&serde_json::json!([[12, 13]])) + ); + assert!(ValidateRequest::validate(&request).is_ok()); + } + + #[test] + fn test_completion_token_ids_rejected_for_multi_choice() { + let request_json = json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "Hello"}], + "n": 2, + "nvext": { + "extra_fields": ["completion_token_ids"] + } + }); + let request: NvCreateChatCompletionRequest = + serde_json::from_value(request_json).expect("Failed to deserialize request"); + + let err = ValidateRequest::validate(&request).expect_err("multi-choice token ids"); + assert!(err.to_string().contains("completion_token_ids")); + } + + #[test] + fn test_truncate_prompt_tokens_rejected_until_supported() { + let request_json = json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "Hello"}], + "truncate_prompt_tokens": 2 + }); + let request: NvCreateChatCompletionRequest = + serde_json::from_value(request_json).expect("Failed to deserialize request"); + assert!(ValidateRequest::validate(&request).is_err()); } } diff --git a/lib/llm/src/protocols/openai/chat_completions/delta.rs b/lib/llm/src/protocols/openai/chat_completions/delta.rs index 132f2ad9d4f5..cce5c27440be 100644 --- a/lib/llm/src/protocols/openai/chat_completions/delta.rs +++ b/lib/llm/src/protocols/openai/chat_completions/delta.rs @@ -311,12 +311,17 @@ impl crate::protocols::openai::DeltaGeneratorExt, + + /// Number of log probabilities to return per prompt token. + #[serde(default, skip_serializing_if = "Option::is_none")] + #[builder(default, setter(strip_option))] + pub prompt_logprobs: Option, } impl CommonExt { @@ -111,6 +116,11 @@ pub trait CommonExtProvider { /// Output Options fn get_skip_special_tokens(&self) -> Option; + + /// Number of prompt logprobs to request from the engine. + fn get_prompt_logprobs_count(&self) -> Option { + None + } } #[cfg(test)] @@ -206,6 +216,7 @@ mod tests { guided_decoding_backend: None, guided_whitespace_pattern: None, skip_special_tokens: None, + prompt_logprobs: None, }; assert!(common_ext.validate().is_ok()); } diff --git a/lib/llm/src/protocols/openai/completions.rs b/lib/llm/src/protocols/openai/completions.rs index e60890214bb9..0ecaf65aed95 100644 --- a/lib/llm/src/protocols/openai/completions.rs +++ b/lib/llm/src/protocols/openai/completions.rs @@ -135,6 +135,10 @@ impl NvExtProvider for NvCreateCompletionRequest { } None } + + fn unsupported_fields(&self) -> Option<&std::collections::HashMap> { + Some(&self.unsupported_fields) + } } impl AnnotationsProvider for NvCreateCompletionRequest { @@ -236,6 +240,10 @@ impl CommonExtProvider for NvCreateCompletionRequest { fn get_skip_special_tokens(&self) -> Option { self.common.skip_special_tokens } + + fn get_prompt_logprobs_count(&self) -> Option { + self.common.prompt_logprobs + } } impl OpenAIStopConditionsProvider for NvCreateCompletionRequest { @@ -252,7 +260,12 @@ impl OpenAIStopConditionsProvider for NvCreateCompletionRequest { } fn get_stop_token_ids(&self) -> Option> { - self.inner.stop.as_ref().and_then(|stop| stop.token_ids()) + if let Some(ids) = self.inner.stop.as_ref().and_then(|stop| stop.token_ids()) { + return Some(ids); + } + self.unsupported_fields + .get("stop_token_ids") + .and_then(|v| serde_json::from_value::>(v.clone()).ok()) } fn nvext(&self) -> Option<&NvExt> { @@ -409,9 +422,11 @@ impl OpenAIOutputOptionsProvider for NvCreateCompletionRequest { } fn get_prompt_logprobs(&self) -> Option { - self.inner - .echo - .and_then(|echo| if echo { Some(1) } else { None }) + self.common.prompt_logprobs.or_else(|| { + self.inner + .echo + .and_then(|echo| if echo { Some(1) } else { None }) + }) } fn get_skip_special_tokens(&self) -> Option { @@ -445,6 +460,10 @@ impl ValidateRequest for NvCreateCompletionRequest { validate::validate_temperature(self.inner.temperature)?; validate::validate_top_p(self.inner.top_p)?; validate::validate_n(self.inner.n)?; + super::nvext::validate_completion_token_ids_single_choice( + get_prompt_batch_size(&self.inner.prompt) * self.inner.n.unwrap_or(1) as usize, + self.nvext.as_ref(), + )?; // none for stream // none for stream_options validate::validate_logprobs(self.inner.logprobs)?; @@ -520,6 +539,22 @@ mod tests { } } + #[test] + fn test_prompt_logprobs_propagates() { + let request_json = json!({ + "model": "test-model", + "prompt": [1, 2, 3], + "prompt_logprobs": 3 + }); + let request: NvCreateCompletionRequest = + serde_json::from_value(request_json).expect("Failed to deserialize request"); + + let output_options = request + .extract_output_options() + .expect("Failed to extract output options"); + assert_eq!(output_options.prompt_logprobs, Some(3)); + } + #[test] fn test_prompt_embeds_only() { // Create valid embeddings: > 100 bytes (PyTorch format) @@ -733,13 +768,102 @@ mod tests { assert_eq!(request.get_stop(), Some(vec!["token_id:576".to_string()])); assert_eq!(request.get_stop_token_ids(), None); - let unsupported_stop_token_ids = json!({ + // Top-level stop_token_ids should be accepted and plumbed. + let whitelisted_stop_token_ids = json!({ "model": "test-model", "prompt": [1, 2, 3], "stop_token_ids": [576] }); - let request: NvCreateCompletionRequest = serde_json::from_value(unsupported_stop_token_ids) + let request: NvCreateCompletionRequest = serde_json::from_value(whitelisted_stop_token_ids) .expect("Failed to deserialize request"); + assert_eq!(request.get_stop_token_ids(), Some(vec![576])); + assert!( + ValidateRequest::validate(&request).is_ok(), + "stop_token_ids must be accepted via PASSTHROUGH_EXTRA_FIELDS" + ); + + let stop_conditions = request + .extract_stop_conditions() + .expect("extract stop conditions"); + assert_eq!(stop_conditions.stop_token_ids, Some(vec![576])); + + let invalid_stop_token_ids = json!({ + "model": "test-model", + "prompt": [1, 2, 3], + "stop_token_ids": "bad" + }); + let request: NvCreateCompletionRequest = + serde_json::from_value(invalid_stop_token_ids).expect("Failed to deserialize request"); + let err = ValidateRequest::validate(&request).expect_err("invalid stop_token_ids"); + assert!(err.to_string().contains("stop_token_ids")); + } + + #[test] + fn test_passthrough_token_constraints_validate() { + let request_json = json!({ + "model": "test-model", + "prompt": [1, 2, 3], + "allowed_token_ids": [10, 11], + "bad_words_token_ids": [[12, 13]] + }); + let request: NvCreateCompletionRequest = + serde_json::from_value(request_json).expect("Failed to deserialize request"); + + assert_eq!( + request.unsupported_fields.get("allowed_token_ids"), + Some(&serde_json::json!([10, 11])) + ); + assert_eq!( + request.unsupported_fields.get("bad_words_token_ids"), + Some(&serde_json::json!([[12, 13]])) + ); + assert!(ValidateRequest::validate(&request).is_ok()); + } + + #[test] + fn test_completion_token_ids_rejected_for_multi_choice() { + let request_json = json!({ + "model": "test-model", + "prompt": [1, 2, 3], + "n": 2, + "nvext": { + "extra_fields": ["completion_token_ids"] + } + }); + let request: NvCreateCompletionRequest = + serde_json::from_value(request_json).expect("Failed to deserialize request"); + + let err = ValidateRequest::validate(&request).expect_err("multi-choice token ids"); + assert!(err.to_string().contains("completion_token_ids")); + } + + #[test] + fn test_completion_token_ids_rejected_for_prompt_batch() { + let request_json = json!({ + "model": "test-model", + "prompt": ["first", "second"], + "n": 1, + "nvext": { + "extra_fields": ["completion_token_ids"] + } + }); + let request: NvCreateCompletionRequest = + serde_json::from_value(request_json).expect("Failed to deserialize request"); + + let err = ValidateRequest::validate(&request).expect_err("prompt batch token ids"); + assert!(err.to_string().contains("completion_token_ids")); + } + + #[test] + fn test_truncate_prompt_tokens_rejected_until_supported() { + let request_json = json!({ + "model": "test-model", + "prompt": [1, 2, 3], + "truncate_prompt_tokens": 2 + }); + let request: NvCreateCompletionRequest = + serde_json::from_value(request_json).expect("Failed to deserialize request"); + assert!(ValidateRequest::validate(&request).is_err()); } } diff --git a/lib/llm/src/protocols/openai/completions/delta.rs b/lib/llm/src/protocols/openai/completions/delta.rs index 471c2ba1152e..4eea61db6b10 100644 --- a/lib/llm/src/protocols/openai/completions/delta.rs +++ b/lib/llm/src/protocols/openai/completions/delta.rs @@ -232,6 +232,12 @@ impl crate::protocols::openai::DeltaGeneratorExt for } } + // Keep token IDs available for optional nvext emission only when requested. + let completion_token_ids_for_nvext = if self.options.response_fields.completion_token_ids { + Some(delta.token_ids.clone()) + } else { + None + }; let logprobs = self.create_logprobs( delta.tokens, delta.token_ids, @@ -256,12 +262,16 @@ impl crate::protocols::openai::DeltaGeneratorExt for // `NvExtResponseFieldSelection` (see `nvext.rs`). Both chat and // completions delta generators go through the same helper so the gating // rules stay in one place. + let prompt_logprobs_payload = + common::llm_backend::prompt_logprobs_from_engine_data(delta.engine_data.as_ref()); if let Some(nvext_response) = self.options.response_fields.build_response_nvext( Some(&self.tracker), delta.disaggregated_params.as_ref(), finish_reason.is_some(), delta.engine_data, stop_reason, + completion_token_ids_for_nvext.as_deref(), + prompt_logprobs_payload, ) && let Ok(nvext_json) = serde_json::to_value(&nvext_response) { response.nvext = Some(nvext_json); @@ -278,6 +288,12 @@ impl crate::protocols::openai::DeltaGeneratorExt for tokens.len() ); } + if let Some(ref tokens) = nvext_response.completion_token_ids { + tracing::debug!( + "Injected completion_token_ids into completions nvext: {} tokens", + tokens.len() + ); + } } Ok(response) diff --git a/lib/llm/src/protocols/openai/nvext.rs b/lib/llm/src/protocols/openai/nvext.rs index f22faf813304..bb5ed4dff5a5 100644 --- a/lib/llm/src/protocols/openai/nvext.rs +++ b/lib/llm/src/protocols/openai/nvext.rs @@ -10,11 +10,15 @@ use utoipa::ToSchema; use validator::{Validate, ValidationError}; pub use crate::agents::context::AgentContext; +use crate::protocols::TokenIdType; +pub use crate::protocols::common::llm_backend::PromptLogprobs; pub use crate::protocols::common::timing::TimingInfo; pub const HEADER_WORKER_INSTANCE_ID: &str = "x-worker-instance-id"; pub const HEADER_PREFILL_INSTANCE_ID: &str = "x-prefill-instance-id"; pub const HEADER_DP_RANK: &str = "x-dp-rank"; +/// Alias for data-parallel rank routing. +pub const HEADER_DP_RANK_ALIAS: &str = "x-data-parallel-rank"; pub const HEADER_PREFILL_DP_RANK: &str = "x-prefill-dp-rank"; const UNSET_DP_RANK_SENTINEL: u32 = u32::MAX; @@ -39,8 +43,10 @@ pub fn apply_header_routing_overrides(nvext: Option, headers: &HeaderMap) .and_then(|v| v.to_str().ok()) .and_then(|s| s.parse::().ok()); + // Accept the alternate data-parallel rank header used by some clients. let dp_rank = headers .get(HEADER_DP_RANK) + .or_else(|| headers.get(HEADER_DP_RANK_ALIAS)) .and_then(|v| v.to_str().ok()) .and_then(|s| s.parse::().ok()); @@ -75,6 +81,9 @@ pub fn apply_header_routing_overrides(nvext: Option, headers: &HeaderMap) pub trait NvExtProvider { fn nvext(&self) -> Option<&NvExt>; fn raw_prompt(&self) -> Option; + fn unsupported_fields(&self) -> Option<&std::collections::HashMap> { + None + } } /// Worker ID information for disaggregated serving @@ -130,6 +139,15 @@ pub struct NvExtResponse { /// If `n > 1` is supported here, this needs an indexed/per-choice shape. #[serde(skip_serializing_if = "Option::is_none")] pub stop_reason: Option, + + /// Engine-emitted output token IDs, returned when explicitly requested. + /// Streaming chunks carry delta token IDs; aggregated responses concatenate them. + #[serde(skip_serializing_if = "Option::is_none")] + pub completion_token_ids: Option>, + + /// Per-prompt-token top-k logprobs, returned on the final chunk when requested. + #[serde(skip_serializing_if = "Option::is_none")] + pub prompt_logprobs: Option, } pub(crate) fn merge_response_nvext( @@ -142,7 +160,28 @@ pub(crate) fn merge_response_nvext( match (target.as_mut(), incoming) { (Some(serde_json::Value::Object(target_obj)), serde_json::Value::Object(incoming_obj)) => { - target_obj.extend(incoming_obj); + // Token IDs are chunk deltas, so aggregation concatenates them. + for (key, value) in incoming_obj { + match key.as_str() { + "completion_token_ids" => { + let entry = target_obj + .entry(&key) + .or_insert_with(|| serde_json::Value::Array(Vec::new())); + if let (serde_json::Value::Array(acc), serde_json::Value::Array(new)) = + (entry, value) + { + acc.extend(new); + } + } + "prompt_logprobs" => { + // Prompt logprobs are final-chunk-only; keep the latest. + target_obj.insert(key, value); + } + _ => { + target_obj.insert(key, value); + } + } + } } (_, incoming) => { *target = Some(incoming); @@ -166,6 +205,10 @@ pub struct NvExtResponseFieldSelection { pub routed_experts: bool, pub engine_data: bool, pub stop_reason: bool, + /// Emit completion token IDs when requested. + pub completion_token_ids: bool, + /// Emit prompt logprobs on the final chunk when requested. + pub prompt_logprobs: bool, } impl NvExtResponseFieldSelection { @@ -183,6 +226,8 @@ impl NvExtResponseFieldSelection { "routed_experts" => selection.routed_experts = true, "engine_data" => selection.engine_data = true, "stop_reason" => selection.stop_reason = true, + "completion_token_ids" => selection.completion_token_ids = true, + "prompt_logprobs" => selection.prompt_logprobs = true, _ => {} } } @@ -216,6 +261,7 @@ impl NvExtResponseFieldSelection { /// - `timing` requires the selection flag, `finish_reason_present == true`, **and** a tracker. /// - `engine_data` requires the selection flag **and** a non-`None` `engine_data_from_backend`. /// - `stop_reason` requires the selection flag **and** a non-`None` `stop_reason_from_backend`. + #[allow(clippy::too_many_arguments)] pub fn build_response_nvext( &self, tracker: Option<&std::sync::Arc>, @@ -223,6 +269,8 @@ impl NvExtResponseFieldSelection { finish_reason_present: bool, engine_data_from_backend: Option, stop_reason_from_backend: Option, + completion_token_ids_from_backend: Option<&[TokenIdType]>, + prompt_logprobs_from_backend: Option, ) -> Option { let worker_id = if self.worker_id { tracker.and_then(|t| t.get_worker_info()) @@ -264,12 +312,28 @@ impl NvExtResponseFieldSelection { None }; + // Pass through chunk or aggregated completion token IDs as provided. + let completion_token_ids = if self.completion_token_ids { + completion_token_ids_from_backend.map(<[u32]>::to_vec) + } else { + None + }; + + // Prompt logprobs describe the full prompt, so emit them once. + let prompt_logprobs = if self.prompt_logprobs && finish_reason_present { + prompt_logprobs_from_backend + } else { + None + }; + if worker_id.is_none() && token_ids.is_none() && routed_experts.is_none() && timing.is_none() && engine_data.is_none() && stop_reason.is_none() + && completion_token_ids.is_none() + && prompt_logprobs.is_none() { return None; } @@ -281,10 +345,29 @@ impl NvExtResponseFieldSelection { routed_experts, engine_data, stop_reason, + completion_token_ids, + prompt_logprobs, }) } } +pub(crate) fn validate_completion_token_ids_single_choice( + total_choices: usize, + nvext: Option<&NvExt>, +) -> anyhow::Result<()> { + let requested = nvext + .and_then(|ext| ext.extra_fields.as_ref()) + .is_some_and(|fields| fields.iter().any(|field| field == "completion_token_ids")); + + if requested && total_choices > 1 { + anyhow::bail!( + "`nvext.extra_fields=[\"completion_token_ids\"]` requires exactly one generated choice" + ); + } + + Ok(()) +} + /// OpenAPI-facing schema for request routing constraints. /// /// Runtime serialization still uses `dynamo_kv_router::protocols::RoutingConstraints`; @@ -350,6 +433,21 @@ pub struct NvExt { #[builder(default, setter(strip_option))] pub max_thinking_tokens: Option, + /// KV prefix-cache isolation hint from RL orchestrators. + /// + /// Prime-RL's orchestrator tags every rollout request with a `cache_salt` + /// derived from the current checkpoint step (e.g. `"step_7"`). When the + /// salt changes across requests, the inference engine treats their prompt + /// prefixes as distinct cache keys even if the token sequences are + /// byte-identical — ensuring KV cache hits from the pre-weight-update + /// policy do not leak into post-update generations. + /// + /// Dynamo passes this through to the backend as a sampling-params hint. + /// Backends that do not support cache salting may ignore it. + #[serde(default, skip_serializing_if = "Option::is_none")] + #[builder(default, setter(strip_option))] + pub cache_salt: Option, + /// Extra fields to be included in the response's nvext /// This is a list of field names that should be populated in the response /// Supported fields include "worker_id", "timing", "routed_experts", "engine_data", @@ -542,6 +640,7 @@ mod tests { assert_eq!(nv_ext.backend_instance_id, None); assert_eq!(nv_ext.token_data, None); assert_eq!(nv_ext.max_thinking_tokens, None); + assert_eq!(nv_ext.cache_salt, None); assert_eq!(nv_ext.extra_fields, None); assert_eq!(nv_ext.prefill_worker_id, None); assert_eq!(nv_ext.decode_worker_id, None); @@ -827,12 +926,12 @@ mod tests { fn test_build_response_nvext_all_false_returns_none() { let sel = sel_all_false(); assert!( - sel.build_response_nvext(None, None, false, None, None) + sel.build_response_nvext(None, None, false, None, None, None, None) .is_none(), "no fields selected → None" ); assert!( - sel.build_response_nvext(None, None, true, None, None) + sel.build_response_nvext(None, None, true, None, None, None, None) .is_none(), "finish_reason alone does not force emission" ); @@ -848,7 +947,7 @@ mod tests { // finish_reason=false: worker_id still emitted (only timing is finish-gated). let out = sel - .build_response_nvext(Some(&tracker), None, false, None, None) + .build_response_nvext(Some(&tracker), None, false, None, None, None, None) .expect("worker_id should emit regardless of finish_reason"); assert!(out.worker_id.is_some()); @@ -867,7 +966,7 @@ mod tests { // timing alone + finish_reason=false → nothing to emit, returns None. assert!( - sel.build_response_nvext(Some(&tracker), None, false, None, None) + sel.build_response_nvext(Some(&tracker), None, false, None, None, None, None) .is_none(), "timing is gated on finish_reason_present" ); @@ -882,7 +981,7 @@ mod tests { let tracker = tracker_with_prefill_worker(); let out = sel - .build_response_nvext(Some(&tracker), None, true, None, None) + .build_response_nvext(Some(&tracker), None, true, None, None, None, None) .expect("timing should emit on finish"); assert!(out.timing.is_some()); @@ -899,7 +998,7 @@ mod tests { }; // finish=true but no tracker → timing not populated → None. assert!( - sel.build_response_nvext(None, None, true, None, None) + sel.build_response_nvext(None, None, true, None, None, None, None) .is_none() ); } @@ -913,7 +1012,7 @@ mod tests { let params = disagg_params_full(); let out = sel - .build_response_nvext(None, Some(¶ms), false, None, None) + .build_response_nvext(None, Some(¶ms), false, None, None, None, None) .expect("token_ids should emit when present"); assert_eq!(out.token_ids, Some(vec![11u32, 22, 33])); @@ -932,7 +1031,7 @@ mod tests { let params = serde_json::json!({ "token_ids": "not-an-array" }); assert!( - sel.build_response_nvext(None, Some(¶ms), false, None, None) + sel.build_response_nvext(None, Some(¶ms), false, None, None, None, None) .is_none(), "malformed token_ids silently suppressed; nothing else selected → None" ); @@ -947,7 +1046,7 @@ mod tests { let params = disagg_params_full(); let out = sel - .build_response_nvext(None, Some(¶ms), false, None, None) + .build_response_nvext(None, Some(¶ms), false, None, None, None, None) .expect("routed_experts should emit when present"); assert_eq!( @@ -970,6 +1069,8 @@ mod tests { true, None, Some(StopReason::String("END".to_string())), + None, + None, ) .expect("stop_reason should emit when requested and present"); @@ -988,7 +1089,7 @@ mod tests { }; assert!( - sel.build_response_nvext(None, None, true, None, None) + sel.build_response_nvext(None, None, true, None, None, None, None) .is_none() ); } @@ -1002,12 +1103,14 @@ mod tests { routed_experts: true, engine_data: false, stop_reason: false, + completion_token_ids: false, + prompt_logprobs: false, }; let tracker = tracker_with_prefill_worker(); let params = disagg_params_full(); let out = sel - .build_response_nvext(Some(&tracker), Some(¶ms), true, None, None) + .build_response_nvext(Some(&tracker), Some(¶ms), true, None, None, None, None) .expect("all fields selected and available → Some"); assert!(out.worker_id.is_some()); @@ -1039,7 +1142,97 @@ mod tests { routed_experts: true, engine_data: false, stop_reason: false, + completion_token_ids: false, + prompt_logprobs: false, } ); } + + #[test] + fn tito_parity_completion_token_ids_pass_through() { + // Per-chunk delta tokens pass through unchanged. + let sel = NvExtResponseFieldSelection { + completion_token_ids: true, + ..Default::default() + }; + let chunk_tokens: &[u32] = &[101, 102, 103]; + let out = sel + .build_response_nvext(None, None, false, None, None, Some(chunk_tokens), None) + .expect("completion_token_ids must be present when requested + provided"); + assert_eq!(out.completion_token_ids, Some(vec![101u32, 102, 103])); + // Accumulation happens in the response aggregator. + assert!(out.prompt_logprobs.is_none()); + assert!(out.engine_data.is_none()); + } + + #[test] + fn tito_parity_prompt_logprobs_final_chunk_only() { + // Prompt logprobs are emitted only with the final chunk. + let sel = NvExtResponseFieldSelection { + prompt_logprobs: true, + ..Default::default() + }; + let mut entry = std::collections::HashMap::new(); + entry.insert( + 42u32, + crate::protocols::common::llm_backend::PromptLogprobEntry { + logprob: -1.234, + rank: Some(1), + decoded_token: None, + }, + ); + let payload: crate::protocols::common::llm_backend::PromptLogprobs = + vec![None, Some(entry)]; + + // Intermediate chunk (no finish): suppressed. + assert!( + sel.build_response_nvext(None, None, false, None, None, None, Some(payload.clone())) + .is_none(), + "prompt_logprobs must be suppressed on intermediate chunks" + ); + + // Final chunk: surfaced. + let out = sel + .build_response_nvext(None, None, true, None, None, None, Some(payload.clone())) + .expect("prompt_logprobs must emit on the final chunk"); + let got = out.prompt_logprobs.expect("prompt_logprobs payload"); + assert_eq!(got.len(), 2); + assert!(got[0].is_none()); + assert_eq!( + got[1].as_ref().unwrap().get(&42u32).unwrap().logprob, + -1.234 + ); + } + + #[test] + fn tito_parity_aggregator_concatenates_completion_token_ids() { + // Aggregation concatenates per-chunk completion token IDs. + let mut target: Option = None; + // Chunk 1: tokens [10, 11, 12] + merge_response_nvext( + &mut target, + Some(serde_json::json!({ "completion_token_ids": [10, 11, 12] })), + ); + // Chunk 2 appends tokens. + merge_response_nvext( + &mut target, + Some(serde_json::json!({ "completion_token_ids": [13, 14] })), + ); + // Chunk 3 (final): one more token + non-token field that should overwrite. + merge_response_nvext( + &mut target, + Some(serde_json::json!({ + "completion_token_ids": [15], + "worker_id": { "decode_worker_id": 7 } + })), + ); + + let aggregated = target.expect("aggregator state"); + assert_eq!( + aggregated["completion_token_ids"], + serde_json::json!([10, 11, 12, 13, 14, 15]), + "completion_token_ids must concatenate across chunks" + ); + assert_eq!(aggregated["worker_id"]["decode_worker_id"], 7); + } } diff --git a/lib/llm/src/protocols/openai/responses/mod.rs b/lib/llm/src/protocols/openai/responses/mod.rs index 40352db03569..5ad484423c37 100644 --- a/lib/llm/src/protocols/openai/responses/mod.rs +++ b/lib/llm/src/protocols/openai/responses/mod.rs @@ -1192,6 +1192,16 @@ mod tests { } } + #[test] + fn test_into_chat_completion_preserves_omitted_max_output_tokens() { + let mut response_req = make_response_with_input("hi there"); + response_req.inner.max_output_tokens = None; + + let nv_req: NvCreateChatCompletionRequest = response_req.try_into().unwrap(); + + assert_eq!(nv_req.inner.max_completion_tokens, None); + } + #[test] fn test_store_mapped_to_chat_completion_request() { let mut req = make_response_with_input("audit me"); diff --git a/lib/llm/src/protocols/openai/validate.rs b/lib/llm/src/protocols/openai/validate.rs index 559dd109ac1b..ed0165fae143 100644 --- a/lib/llm/src/protocols/openai/validate.rs +++ b/lib/llm/src/protocols/openai/validate.rs @@ -97,16 +97,51 @@ pub const MAX_REPETITION_PENALTY: f32 = 2.0; // Shared Fields // -/// Validates that no unsupported fields are present in the request +/// Extra-body fields accepted for backend-specific handling. +pub const PASSTHROUGH_EXTRA_FIELDS: &[&str] = &[ + "cache_salt", + "stop_token_ids", + "detokenize", + "allowed_token_ids", + "bad_words_token_ids", +]; + +/// Validates that no unsupported fields are present in the request. +/// +/// Fields in `PASSTHROUGH_EXTRA_FIELDS` are validated by downstream handlers. pub fn validate_no_unsupported_fields( unsupported_fields: &std::collections::HashMap, ) -> Result<(), anyhow::Error> { - if !unsupported_fields.is_empty() { - let fields: Vec<_> = unsupported_fields - .keys() - .map(|s| format!("`{}`", s)) - .collect(); - anyhow::bail!("Unsupported parameter(s): {}", fields.join(", ")); + let unknown: Vec<_> = unsupported_fields + .keys() + .filter(|k| !PASSTHROUGH_EXTRA_FIELDS.contains(&k.as_str())) + .map(|s| format!("`{}`", s)) + .collect(); + if !unknown.is_empty() { + anyhow::bail!("Unsupported parameter(s): {}", unknown.join(", ")); + } + if let Some(value) = unsupported_fields.get("cache_salt") + && !value.is_string() + { + anyhow::bail!("`cache_salt` must be a string"); + } + if let Some(value) = unsupported_fields.get("stop_token_ids") { + serde_json::from_value::>(value.clone()) + .map_err(|_| anyhow::anyhow!("`stop_token_ids` must be an array of token IDs"))?; + } + if let Some(value) = unsupported_fields.get("detokenize") + && !value.is_boolean() + { + anyhow::bail!("`detokenize` must be a boolean"); + } + if let Some(value) = unsupported_fields.get("allowed_token_ids") { + serde_json::from_value::>(value.clone()) + .map_err(|_| anyhow::anyhow!("`allowed_token_ids` must be an array of token IDs"))?; + } + if let Some(value) = unsupported_fields.get("bad_words_token_ids") { + serde_json::from_value::>>(value.clone()).map_err( + |_| anyhow::anyhow!("`bad_words_token_ids` must be an array of token ID arrays"), + )?; } Ok(()) } diff --git a/lib/llm/src/protocols/unified.rs b/lib/llm/src/protocols/unified.rs index c748678126e5..e2dfbfadddfb 100644 --- a/lib/llm/src/protocols/unified.rs +++ b/lib/llm/src/protocols/unified.rs @@ -388,6 +388,10 @@ impl CommonExtProvider for UnifiedRequest { fn get_skip_special_tokens(&self) -> Option { self.inner.common.skip_special_tokens } + + fn get_prompt_logprobs_count(&self) -> Option { + self.inner.common.prompt_logprobs + } } impl OpenAIStopConditionsProvider for UnifiedRequest {