feat(RL): add nvext Tokens-in-Tokens-Out and RL related response protocol and frontend support

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

lib/llm/src/preprocessor.rs

@@ -206,6 +206,14 @@ static DIM_FETCH_HTTP_CLIENT: std::sync::LazyLock<reqwest::Client> =
             .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<dyn OAIPromptFormatter>,
@@ -240,6 +248,94 @@ pub struct OpenAIPreprocessor {
 }
 
 impl OpenAIPreprocessor {
+    fn omitted_max_tokens_default(
+        prompt_len: usize,
+        context_length: u32,
+        options: PreprocessRequestOptions,
+    ) -> Option<u32> {
+        if context_length == 0 || options.preserve_omitted_max_tokens {
+            return None;
+        }
+        Some(context_length.saturating_sub(prompt_len as u32))
+    }
+
+    fn nvext_passthrough_args<R: NvExtProvider>(
+        request: &R,
+    ) -> Option<serde_json::Map<String, serde_json::Value>> {
+        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<R: NvExtProvider>(
+        request: &R,
+    ) -> Option<serde_json::Map<String, serde_json::Value>> {
+        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<R: NvExtProvider>(request: &R) -> Option<serde_json::Value> {
+        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<Arc<Self>> {
         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<String, String>, 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<String, String>, 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<R: OAIChatLikeRequest>(
+    pub async fn gather_multi_modal_data<R: OAIChatLikeRequest + NvExtProvider>(
         &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::<bool>(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() {

lib/llm/src/protocols/common/llm_backend.rs

@@ -14,6 +14,20 @@ use dynamo_runtime::protocols::maybe_error::MaybeError;
 pub type TokenType = Option<String>;
 pub type LogProbs = Vec<f64>;
 
+/// 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<u32>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub decoded_token: Option<String>,
+}
+
+/// 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<Option<std::collections::HashMap<TokenIdType, PromptLogprobEntry>>>;
+
 /// 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<PromptLogprobs> {
+    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())

lib/llm/src/protocols/openai.rs

@@ -144,7 +144,6 @@ impl<T: OpenAISamplingOptionsProvider + CommonExtProvider> 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,