From 127a14a7703cc70a21089d2656a8477f4e92e62c Mon Sep 17 00:00:00 2001 From: William Arnold Date: Mon, 13 Apr 2026 09:57:39 -0700 Subject: [PATCH 01/20] feat: cherry-pick tokenize/detokenize endpoints and token ID support from #7699 Cherry-picks the following from jthomson04/tokenize-endpoint: - POST /v1/tokenize and /v1/detokenize HTTP endpoints - Tokenizer trait: encode_with_special_tokens(), convert_ids_to_tokens() - return_tokens_as_token_ids parameter for chat completions - Multi-instance tokenize fix (discovery watcher) - Jail logprobs preservation through tool-call jailing --- components/src/dynamo/vllm/handlers.py | 30 +- lib/llm/src/discovery/watcher.rs | 22 + lib/llm/src/entrypoint/input/text.rs | 1 + lib/llm/src/http/service/openai.rs | 766 +++++++++++++++++- lib/llm/src/http/service/service_v2.rs | 1 + lib/llm/src/model_card.rs | 25 +- lib/llm/src/preprocessor/prompt/template.rs | 102 ++- lib/llm/src/protocols/anthropic/types.rs | 1 + lib/llm/src/protocols/common.rs | 4 + lib/llm/src/protocols/openai.rs | 7 + .../src/protocols/openai/chat_completions.rs | 10 + .../openai/chat_completions/delta.rs | 14 +- .../protocols/openai/chat_completions/jail.rs | 215 ++++- lib/llm/src/protocols/openai/responses/mod.rs | 1 + lib/llm/src/protocols/openai/tokenization.rs | 124 +++ lib/llm/src/protocols/unified.rs | 1 + lib/llm/src/tokenizers.rs | 28 +- lib/llm/src/tokenizers/fastokens.rs | 26 +- lib/llm/src/tokenizers/hf.rs | 31 +- lib/llm/src/tokenizers/tiktoken.rs | 53 +- .../tests/parallel_tool_call_integration.rs | 1 + lib/llm/tests/preprocessor.rs | 2 + lib/llm/tests/test_common_ext.rs | 3 + lib/llm/tests/test_streaming_usage.rs | 2 + lib/llm/tests/tool_choice.rs | 1 + lib/llm/tests/tool_choice_finish_reasons.rs | 1 + 26 files changed, 1387 insertions(+), 85 deletions(-) create mode 100644 lib/llm/src/protocols/openai/tokenization.rs diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index 96aa983f9c58..d148e21bba2e 100644 --- a/components/src/dynamo/vllm/handlers.py +++ b/components/src/dynamo/vllm/handlers.py @@ -1359,7 +1359,8 @@ def _build_completion_usage( @staticmethod def _extract_logprobs( - output, num_output_tokens_so_far: int, tokenizer=None + output, num_output_tokens_so_far: int, tokenizer=None, + return_tokens_as_token_ids: bool = False, ) -> tuple[list[float] | None, list[list[dict]] | None]: """ Extract logprobs from vLLM CompletionOutput for new tokens. @@ -1401,12 +1402,15 @@ def _extract_logprobs( # Build top_logprobs list for this token position token_top_logprobs = [] for tok_id, logprob_info in token_logprobs_dict.items(): - token_str = getattr(logprob_info, "decoded_token", None) - if not token_str and tokenizer: - try: - token_str = tokenizer.decode([tok_id]) - except Exception: - token_str = None + if return_tokens_as_token_ids: + token_str = f"token_id:{tok_id}" + else: + token_str = getattr(logprob_info, "decoded_token", None) + if not token_str and tokenizer: + try: + token_str = tokenizer.decode([tok_id]) + except Exception: + token_str = None token_top_logprobs.append( { "rank": ( @@ -1468,6 +1472,7 @@ async def generate_tokens( embedding_sequence_length=None, trace_headers=None, priority=0, + return_tokens_as_token_ids=False, ): try: # Log LoRA usage for this generation (debug level to avoid log spam) @@ -1511,7 +1516,8 @@ async def generate_tokens( # Extract logprobs for new tokens if available tokenizer = getattr(self.engine_client, "tokenizer", None) log_probs, top_logprobs = self._extract_logprobs( - output, num_output_tokens_so_far, tokenizer=tokenizer + output, num_output_tokens_so_far, tokenizer=tokenizer, + return_tokens_as_token_ids=return_tokens_as_token_ids, ) if log_probs is not None: out["log_probs"] = log_probs @@ -1698,6 +1704,13 @@ async def _generate_token_mode(self, request, context, request_id): trace_headers = build_trace_headers(context) + output_options = request.get("output_options", {}) + return_tokens_as_token_ids = bool( + output_options.get("return_tokens_as_token_ids") + ) + + print(f"[DEBUG] output_options={output_options}, return_tokens_as_token_ids={return_tokens_as_token_ids}", flush=True) + async with self._abort_monitor(context, request_id): try: async for tok in self.generate_tokens( @@ -1709,6 +1722,7 @@ async def _generate_token_mode(self, request, context, request_id): embedding_sequence_length=embedding_sequence_length, trace_headers=trace_headers, priority=priority, + return_tokens_as_token_ids=return_tokens_as_token_ids, ): if prefill_result is not None and "completion_usage" in tok: tok["completion_usage"][ diff --git a/lib/llm/src/discovery/watcher.rs b/lib/llm/src/discovery/watcher.rs index a9019f5b1096..fee02ebb0b58 100644 --- a/lib/llm/src/discovery/watcher.rs +++ b/lib/llm/src/discovery/watcher.rs @@ -380,6 +380,7 @@ impl ModelWatcher { if let Some(model) = self.manager.get_model(&model_name) && model.has_worker_set(&ws_key) { + self.resolve_card_local_files(&model_name, card); self.manager .save_model_card(&mcid.to_path(), card.clone())?; tracing::debug!( @@ -396,6 +397,7 @@ impl ModelWatcher { .registering_worker_sets .insert(registration_key.clone()) { + self.resolve_card_local_files(&model_name, card); self.manager .save_model_card(&mcid.to_path(), card.clone())?; tracing::debug!( @@ -414,6 +416,26 @@ impl ModelWatcher { result } + /// If an existing card for the same model has already-downloaded local files, + /// point this card's URL-backed files at the same local directory. This avoids + /// re-downloading config files for every worker that joins an existing WorkerSet. + fn resolve_card_local_files( + &self, + model_name: &str, + card: &mut ModelDeploymentCard, + ) { + let local_dir = self + .manager + .get_model_cards() + .iter() + .find(|c| c.name() == model_name) + .and_then(|c| c.local_file_dir().map(|p| p.to_path_buf())); + + if let Some(dir) = local_dir { + card.update_dir(&dir); + } + } + /// Build a complete WorkerSet with all engines for this (model, namespace) /// and add it to the Model. async fn do_worker_set_registration( diff --git a/lib/llm/src/entrypoint/input/text.rs b/lib/llm/src/entrypoint/input/text.rs index 1c0138fd34b3..5821176496a9 100644 --- a/lib/llm/src/entrypoint/input/text.rs +++ b/lib/llm/src/entrypoint/input/text.rs @@ -115,6 +115,7 @@ async fn main_loop( nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; diff --git a/lib/llm/src/http/service/openai.rs b/lib/llm/src/http/service/openai.rs index 41fa5cd69a6f..9417f15a9fe6 100644 --- a/lib/llm/src/http/service/openai.rs +++ b/lib/llm/src/http/service/openai.rs @@ -43,6 +43,8 @@ use super::{ service_v2, }; use crate::engines::ValidateRequest; +use crate::model_card::ModelDeploymentCard; +use crate::preprocessor::prompt::PromptFormatter; use crate::protocols::openai::chat_completions::aggregator::ChatCompletionAggregator; use crate::protocols::openai::nvext::apply_header_routing_overrides; use crate::protocols::openai::{ @@ -55,6 +57,10 @@ use crate::protocols::openai::{ embeddings::{NvCreateEmbeddingRequest, NvCreateEmbeddingResponse}, images::{NvCreateImageRequest, NvImagesResponse}, responses::{NvCreateResponse, NvResponse, ResponseParams, chat_completion_to_response}, + tokenization::{ + DetokenizeRequest, DetokenizeResponse, TokenizeChatRequest, TokenizeCompletionRequest, + TokenizeRequest, TokenizeResponse, + }, videos::{NvCreateVideoRequest, NvVideosResponse}, }; use crate::protocols::unified::UnifiedRequest; @@ -303,6 +309,176 @@ pub async fn smart_json_error_middleware(request: Request, next: Next) -> } } +fn bad_request>(message: T) -> ErrorResponse { + let code = StatusCode::BAD_REQUEST; + ( + code, + Json(ErrorMessage { + message: message.into(), + error_type: map_error_code_to_error_type(code), + code: code.as_u16(), + }), + ) +} + +fn resolve_tokenizer_model_name( + state: &Arc, + requested_model: Option<&str>, +) -> Result { + if let Some(model) = requested_model { + if state.manager().has_model_any(model) { + return Ok(model.to_string()); + } + return Err(ErrorMessage::model_not_found()); + } + + // Preserve this order: without an explicit model, prefer `model_display_names()` first because + // those names are known to the serving layer; only if that yields one choice do we fall back + // to `get_model_cards()`, whose card-only metadata may exist before a model is fully + // registered. This keeps tokenizer endpoints usable from cards alone, but card names may not + // map to serving-capable models, so explicit requests still gate on `has_model_any()` / + // `ErrorMessage::model_not_found()` and ambiguous fallback still returns `bad_request()`. + let served_models = state.manager().model_display_names(); + if served_models.len() == 1 { + return Ok(served_models.into_iter().next().unwrap()); + } + + let card_models: HashSet = state + .manager() + .get_model_cards() + .into_iter() + .map(|card| card.display_name) + .collect(); + if card_models.len() == 1 { + return Ok(card_models.into_iter().next().unwrap()); + } + + Err(bad_request( + "Model must be specified when more than one model is served.", + )) +} + +fn resolve_model_card( + state: &Arc, + requested_model: Option<&str>, +) -> Result<(String, ModelDeploymentCard), ErrorResponse> { + let model = resolve_tokenizer_model_name(state, requested_model)?; + let card = state + .manager() + .get_model_cards() + .into_iter() + .find(|card| card.display_name == model) + .ok_or_else(|| { + ErrorMessage::internal_server_error(&format!( + "Tokenizer metadata is not available for model '{}'", + model + )) + })?; + Ok((model, card)) +} + +fn extract_assistant_content_text( + content: &dynamo_protocols::types::ChatCompletionRequestAssistantMessageContent, +) -> String { + use dynamo_protocols::types::ChatCompletionRequestAssistantMessageContent as Content; + use dynamo_protocols::types::ChatCompletionRequestAssistantMessageContentPart as Part; + + match content { + Content::Text(text) => text.clone(), + Content::Array(parts) => parts + .iter() + .filter_map(|part| match part { + Part::Text(text) => Some(text.text.clone()), + Part::Refusal(_) => None, + }) + .collect::>() + .join(""), + } +} + +fn apply_continue_final_message( + rendered_prompt: String, + messages: &[dynamo_protocols::types::ChatCompletionRequestMessage], +) -> Result { + use dynamo_protocols::types::ChatCompletionRequestMessage as Message; + + let Some(Message::Assistant(message)) = messages.last() else { + return Err(bad_request( + "Cannot set `continue_final_message` to True when the final message is not from the assistant.", + )); + }; + + let Some(content) = message.content.as_ref() else { + return Err(bad_request( + "Cannot set `continue_final_message` to True when the final assistant message has no content.", + )); + }; + + let final_message = extract_assistant_content_text(content); + let trimmed_final_message = final_message.trim(); + if trimmed_final_message.is_empty() { + return Err(bad_request( + "Cannot set `continue_final_message` to True when the final assistant message content is empty.", + )); + } + + // Use rfind to locate the last occurrence of the assistant content in the rendered prompt, + // then truncate everything after it (e.g. EOS tokens, generation prompts). This assumes the + // final assistant message text appears only once at the end of the rendered output. + let Some(final_msg_loc) = rendered_prompt.rfind(trimmed_final_message) else { + return Err(ErrorMessage::internal_server_error( + "Failed to trim rendered prompt for `continue_final_message`.", + )); + }; + + Ok(rendered_prompt[..final_msg_loc + trimmed_final_message.len()].to_string()) +} + +fn make_tokenize_chat_completion_request( + model: String, + request: &TokenizeChatRequest, +) -> NvCreateChatCompletionRequest { + let inner = dynamo_protocols::types::CreateChatCompletionRequest { + model, + messages: request.messages.clone(), + tools: request.tools.clone(), + ..Default::default() + }; + + NvCreateChatCompletionRequest { + inner, + common: Default::default(), + nvext: None, + chat_template_args: Some(request.merged_chat_template_kwargs()), + media_io_kwargs: request.media_io_kwargs.clone(), + return_tokens_as_token_ids: None, + unsupported_fields: Default::default(), + } +} + +fn render_tokenize_chat_prompt( + card: &ModelDeploymentCard, + model: String, + request: &TokenizeChatRequest, +) -> Result { + request.validate().map_err(bad_request)?; + + let formatter = + PromptFormatter::from_mdc_with_chat_template(card, request.chat_template.as_deref()) + .map_err(|err| ErrorMessage::from_anyhow(err, "Failed to build chat formatter"))?; + let wrapped_request = make_tokenize_chat_completion_request(model, request); + let mut prompt = match formatter { + PromptFormatter::OAI(formatter) => formatter.render(&wrapped_request), + } + .map_err(|err| ErrorMessage::from_anyhow(err, "Failed to render chat prompt"))?; + + if request.continue_final_message { + prompt = apply_continue_final_message(prompt, &request.messages)?; + } + + Ok(prompt) +} + /// Return the request ID for the current request. /// /// The canonical request ID is set by `make_inference_request_span()` and stored @@ -1874,6 +2050,102 @@ struct ModelListing { owned_by: String, } +async fn tokenize( + State(state): State>, + Json(request): Json, +) -> Result { + check_ready(&state)?; + + let (_, card) = resolve_model_card(&state, request.model())?; + let tokenizer = card + .tokenizer() + .map_err(|err| ErrorMessage::from_anyhow(err, "Failed to load tokenizer"))?; + + let (tokens, token_strs) = match request { + TokenizeRequest::Completion(TokenizeCompletionRequest { + prompt, + add_special_tokens, + return_token_strs, + .. + }) => { + let encoding = tokenizer + .encode_with_special_tokens(&prompt, add_special_tokens) + .map_err(|err| ErrorMessage::from_anyhow(err, "Failed to tokenize prompt"))?; + let token_ids = encoding.token_ids().to_vec(); + let token_strs = if return_token_strs { + Some(tokenizer.convert_ids_to_tokens(&token_ids).map_err(|err| { + ErrorMessage::from_anyhow(err, "Failed to resolve token strings") + })?) + } else { + None + }; + (token_ids, token_strs) + } + TokenizeRequest::Chat(request) => { + let model = request + .model + .clone() + .unwrap_or_else(|| card.display_name.clone()); + let prompt = render_tokenize_chat_prompt(&card, model, &request)?; + let encoding = tokenizer + .encode_with_special_tokens(&prompt, request.add_special_tokens) + .map_err(|err| { + ErrorMessage::from_anyhow(err, "Failed to tokenize rendered chat prompt") + })?; + let token_ids = encoding.token_ids().to_vec(); + let token_strs = if request.return_token_strs { + Some(tokenizer.convert_ids_to_tokens(&token_ids).map_err(|err| { + ErrorMessage::from_anyhow(err, "Failed to resolve token strings") + })?) + } else { + None + }; + (token_ids, token_strs) + } + }; + + Ok(Json(TokenizeResponse { + count: tokens.len(), + max_model_len: card.context_length, + tokens, + token_strs, + }) + .into_response()) +} + +async fn detokenize( + State(state): State>, + Json(request): Json, +) -> Result { + check_ready(&state)?; + + let (_, card) = resolve_model_card(&state, request.model.as_deref())?; + let tokenizer = card + .tokenizer() + .map_err(|err| ErrorMessage::from_anyhow(err, "Failed to load tokenizer"))?; + let prompt = tokenizer + .decode(&request.tokens, false) + .map_err(|err| ErrorMessage::from_anyhow(err, "Failed to detokenize prompt"))?; + + Ok(Json(DetokenizeResponse { prompt }).into_response()) +} + +pub fn tokenization_router(state: Arc) -> (Vec, Router) { + let tokenize_path = "/tokenize"; + let detokenize_path = "/detokenize"; + let docs = vec![ + RouteDoc::new(axum::http::Method::POST, tokenize_path), + RouteDoc::new(axum::http::Method::POST, detokenize_path), + ]; + let router = Router::new() + .route(tokenize_path, post(tokenize)) + .route(detokenize_path, post(detokenize)) + .layer(middleware::from_fn(smart_json_error_middleware)) + .layer(axum::extract::DefaultBodyLimit::max(get_body_limit())) + .with_state(state); + (docs, router) +} + /// Create an Axum [`Router`] for the OpenAI API Completions endpoint /// If not path is provided, the default path is `/v1/completions` pub fn completions_router( @@ -2443,19 +2715,121 @@ pub fn audios_router( mod tests { use super::*; - use crate::discovery::ModelManagerError; + use crate::discovery::{ModelManager, ModelManagerError}; use crate::protocols::openai::chat_completions::NvCreateChatCompletionRequest; use crate::protocols::openai::common_ext::CommonExt; use crate::protocols::openai::completions::NvCreateCompletionRequest; use crate::protocols::openai::responses::NvCreateResponse; + use crate::protocols::openai::tokenization::DetokenizeRequest; + use axum::extract::State; use dynamo_protocols::types::responses::{CreateResponse, Input, PromptConfig}; use dynamo_protocols::types::{ + ChatCompletionRequestAssistantMessage, ChatCompletionRequestAssistantMessageContent, ChatCompletionRequestMessage, ChatCompletionRequestUserMessage, - ChatCompletionRequestUserMessageContent, CreateChatCompletionRequest, - CreateCompletionRequest, + ChatCompletionRequestUserMessageContent, ChatCompletionTool, CreateChatCompletionRequest, + CreateCompletionRequest, FunctionObject, }; + use dynamo_runtime::discovery::{MockDiscovery, SharedMockRegistry}; + use std::collections::HashMap as StdHashMap; + use tokio_util::sync::CancellationToken; const BACKUP_ERROR_MESSAGE: &str = "Failed to generate completions"; + const TOKENIZE_MODEL_PATH: &str = concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/data/sample-models/mock-llama-3.1-8b-instruct" + ); + const DETOKENIZE_MODEL_PATH: &str = concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/data/sample-models/TinyLlama_v1.1" + ); + + fn make_tokenize_state_with_path( + model_name: &str, + model_path: &str, + ) -> (Arc, ModelDeploymentCard) { + let mut card = ModelDeploymentCard::load_from_disk(model_path, None).unwrap(); + card.set_name(model_name); + + let manager = Arc::new(ModelManager::new()); + manager + .save_model_card(&format!("__test_model_card_{model_name}"), card.clone()) + .unwrap(); + manager + .add_prefill_model(model_name, card.mdcsum()) + .unwrap(); + + let discovery = Arc::new(MockDiscovery::new(None, SharedMockRegistry::new())); + let state = Arc::new(service_v2::State::new( + manager, + discovery, + CancellationToken::new(), + )); + (state, card) + } + + fn make_tokenize_state(model_name: &str) -> (Arc, ModelDeploymentCard) { + make_tokenize_state_with_path(model_name, TOKENIZE_MODEL_PATH) + } + + fn make_tokenize_state_without_card(model_name: &str) -> Arc { + let manager = Arc::new(ModelManager::new()); + manager + .add_prefill_model(model_name, "missing-card") + .unwrap(); + + let discovery = Arc::new(MockDiscovery::new(None, SharedMockRegistry::new())); + Arc::new(service_v2::State::new( + manager, + discovery, + CancellationToken::new(), + )) + } + + async fn response_json(response: Response) -> T { + let body = response.into_body(); + let bytes = axum::body::to_bytes(body, usize::MAX).await.unwrap(); + serde_json::from_slice(&bytes).unwrap() + } + + fn sample_chat_messages() -> Vec { + vec![ + ChatCompletionRequestMessage::User(ChatCompletionRequestUserMessage { + content: ChatCompletionRequestUserMessageContent::Text("Hi there!".to_string()), + name: None, + }), + ChatCompletionRequestMessage::Assistant(ChatCompletionRequestAssistantMessage { + content: Some(ChatCompletionRequestAssistantMessageContent::Text( + "Nice to meet you!".to_string(), + )), + ..Default::default() + }), + ChatCompletionRequestMessage::User(ChatCompletionRequestUserMessage { + content: ChatCompletionRequestUserMessageContent::Text( + "Can I ask a question?".to_string(), + ), + name: None, + }), + ] + } + + fn sample_tools() -> Vec { + vec![ChatCompletionTool { + r#type: dynamo_protocols::types::ChatCompletionToolType::Function, + function: FunctionObject { + name: "get_weather".to_string(), + description: None, + parameters: Some(serde_json::json!({ + "type": "object", + "properties": { + "location": { + "type": "string" + } + } + })), + strict: None, + }, + }] + } fn http_error_from_engine(code: u16) -> Result<(), anyhow::Error> { Err(HttpError { @@ -2641,6 +3015,7 @@ mod tests { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_chat_completion_required_fields(&request); @@ -2673,6 +3048,7 @@ mod tests { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_chat_completion_required_fields(&request); @@ -2889,6 +3265,7 @@ mod tests { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; @@ -2919,6 +3296,7 @@ mod tests { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_chat_completion_fields_generic(&request); @@ -2948,6 +3326,7 @@ mod tests { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_chat_completion_fields_generic(&request); @@ -2977,6 +3356,7 @@ mod tests { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_chat_completion_fields_generic(&request); @@ -3008,6 +3388,7 @@ mod tests { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_chat_completion_fields_generic(&request); @@ -3037,6 +3418,7 @@ mod tests { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_chat_completion_fields_generic(&request); @@ -3343,6 +3725,384 @@ mod tests { ); } + #[test] + fn test_apply_continue_final_message_trims_rendered_prompt() { + let rendered_prompt = "USER: Hi\nASSISTANT: Sure.<|assistant|>".to_string(); + let messages = vec![ + ChatCompletionRequestMessage::User(ChatCompletionRequestUserMessage { + content: ChatCompletionRequestUserMessageContent::Text("Hi".to_string()), + name: None, + }), + ChatCompletionRequestMessage::Assistant(ChatCompletionRequestAssistantMessage { + content: Some(ChatCompletionRequestAssistantMessageContent::Text( + "Sure.".to_string(), + )), + ..Default::default() + }), + ]; + + let trimmed = apply_continue_final_message(rendered_prompt, &messages).unwrap(); + assert_eq!(trimmed, "USER: Hi\nASSISTANT: Sure."); + } + + #[tokio::test] + async fn test_tokenize_completion_route_matches_tokenizer() { + let (state, card) = make_tokenize_state("test-model"); + let tokenizer = card.tokenizer().unwrap(); + let prompt = "This is a completion tokenize test."; + + for add_special_tokens in [false, true] { + let response = tokenize( + State(state.clone()), + Json(TokenizeRequest::Completion(TokenizeCompletionRequest { + model: Some("test-model".to_string()), + prompt: prompt.to_string(), + add_special_tokens, + return_token_strs: false, + })), + ) + .await + .unwrap(); + let body: TokenizeResponse = response_json(response).await; + let expected = tokenizer + .encode_with_special_tokens(prompt, add_special_tokens) + .unwrap(); + + assert_eq!(body.tokens, expected.token_ids()); + assert_eq!(body.count, expected.token_ids().len()); + assert_eq!(body.max_model_len, card.context_length); + assert!(body.token_strs.is_none()); + } + } + + #[tokio::test] + async fn test_tokenize_chat_route_matches_rendered_prompt() { + let (state, card) = make_tokenize_state("test-model"); + let tokenizer = card.tokenizer().unwrap(); + let messages = sample_chat_messages(); + + for add_generation_prompt in [false, true] { + for add_special_tokens in [false, true] { + let request = TokenizeChatRequest { + model: Some("test-model".to_string()), + messages: messages.clone(), + add_generation_prompt, + return_token_strs: false, + continue_final_message: false, + add_special_tokens, + chat_template: None, + chat_template_kwargs: None, + media_io_kwargs: None, + mm_processor_kwargs: None, + tools: None, + }; + let prompt = + render_tokenize_chat_prompt(&card, "test-model".to_string(), &request).unwrap(); + let expected = tokenizer + .encode_with_special_tokens(&prompt, add_special_tokens) + .unwrap(); + + let response = tokenize( + State(state.clone()), + Json(TokenizeRequest::Chat(TokenizeChatRequest { + model: Some("test-model".to_string()), + messages: messages.clone(), + add_generation_prompt, + return_token_strs: false, + continue_final_message: false, + add_special_tokens, + chat_template: None, + chat_template_kwargs: None, + media_io_kwargs: None, + mm_processor_kwargs: None, + tools: None, + })), + ) + .await + .unwrap(); + let body: TokenizeResponse = response_json(response).await; + assert_eq!(body.tokens, expected.token_ids()); + assert_eq!(body.count, expected.token_ids().len()); + } + } + } + + #[tokio::test] + async fn test_tokenize_chat_route_with_tools_and_continue_final_message() { + let (state, card) = make_tokenize_state("test-model"); + let tokenizer = card.tokenizer().unwrap(); + let mut messages = sample_chat_messages(); + messages.push(ChatCompletionRequestMessage::Assistant( + ChatCompletionRequestAssistantMessage { + content: Some(ChatCompletionRequestAssistantMessageContent::Text( + "Sure,".to_string(), + )), + ..Default::default() + }, + )); + let tools = sample_tools(); + + let request = TokenizeChatRequest { + model: Some("test-model".to_string()), + messages: messages.clone(), + add_generation_prompt: false, + return_token_strs: false, + continue_final_message: true, + add_special_tokens: true, + chat_template: None, + chat_template_kwargs: None, + media_io_kwargs: None, + mm_processor_kwargs: None, + tools: Some(tools.clone()), + }; + let prompt = + render_tokenize_chat_prompt(&card, "test-model".to_string(), &request).unwrap(); + let expected = tokenizer.encode_with_special_tokens(&prompt, true).unwrap(); + + let response = tokenize( + State(state), + Json(TokenizeRequest::Chat(TokenizeChatRequest { + model: Some("test-model".to_string()), + messages: messages.clone(), + add_generation_prompt: false, + return_token_strs: false, + continue_final_message: true, + add_special_tokens: true, + chat_template: None, + chat_template_kwargs: None, + media_io_kwargs: None, + mm_processor_kwargs: None, + tools: Some(tools), + })), + ) + .await + .unwrap(); + let body: TokenizeResponse = response_json(response).await; + assert_eq!(body.tokens, expected.token_ids()); + } + + #[tokio::test] + async fn test_tokenize_route_returns_token_strings() { + let (state, card) = make_tokenize_state("test-model"); + let tokenizer = card.tokenizer().unwrap(); + let prompt = "Return token strings please."; + + let response = tokenize( + State(state), + Json(TokenizeRequest::Completion(TokenizeCompletionRequest { + model: Some("test-model".to_string()), + prompt: prompt.to_string(), + add_special_tokens: true, + return_token_strs: true, + })), + ) + .await + .unwrap(); + let body: TokenizeResponse = response_json(response).await; + let expected = tokenizer.encode_with_special_tokens(prompt, true).unwrap(); + let expected_token_strs = tokenizer + .convert_ids_to_tokens(expected.token_ids()) + .unwrap(); + + assert_eq!(body.tokens, expected.token_ids()); + assert_eq!(body.token_strs, Some(expected_token_strs)); + } + + #[tokio::test] + async fn test_tokenize_chat_rejects_incompatible_flags() { + let (state, _) = make_tokenize_state("test-model"); + let messages = sample_chat_messages(); + + let error = tokenize( + State(state), + Json(TokenizeRequest::Chat(TokenizeChatRequest { + model: Some("test-model".to_string()), + messages, + add_generation_prompt: true, + return_token_strs: false, + continue_final_message: true, + add_special_tokens: false, + chat_template: None, + chat_template_kwargs: None, + media_io_kwargs: None, + mm_processor_kwargs: None, + tools: None, + })), + ) + .await + .unwrap_err(); + let response = error.into_response(); + let body: ErrorMessage = response_json(response).await; + assert!(body.message.contains( + "Cannot set both `continue_final_message` and `add_generation_prompt` to True." + )); + } + + #[tokio::test] + async fn test_detokenize_route_round_trips_prompt() { + let (state, card) = make_tokenize_state_with_path("test-model", DETOKENIZE_MODEL_PATH); + let tokenizer = card.tokenizer().unwrap(); + let prompt = "This is a detokenize test prompt."; + let tokens = tokenizer + .encode_with_special_tokens(prompt, false) + .unwrap() + .token_ids() + .to_vec(); + + let response = detokenize( + State(state), + Json(DetokenizeRequest { + model: Some("test-model".to_string()), + tokens, + }), + ) + .await + .unwrap(); + let body: DetokenizeResponse = response_json(response).await; + assert_eq!(body.prompt, prompt); + } + + #[tokio::test] + async fn test_tokenize_route_rejects_models_without_card_metadata() { + let state = make_tokenize_state_without_card("test-model"); + let error = tokenize( + State(state), + Json(TokenizeRequest::Completion(TokenizeCompletionRequest { + model: Some("test-model".to_string()), + prompt: "hello".to_string(), + add_special_tokens: true, + return_token_strs: false, + })), + ) + .await + .unwrap_err(); + + let response = error.into_response(); + let body: ErrorMessage = response_json(response).await; + assert!( + body.message + .contains("Tokenizer metadata is not available for model 'test-model'") + ); + } + + #[tokio::test] + async fn test_detokenize_route_rejects_models_without_card_metadata() { + let state = make_tokenize_state_without_card("test-model"); + let error = detokenize( + State(state), + Json(DetokenizeRequest { + model: Some("test-model".to_string()), + tokens: vec![1, 2, 3], + }), + ) + .await + .unwrap_err(); + + let response = error.into_response(); + let body: ErrorMessage = response_json(response).await; + assert!( + body.message + .contains("Tokenizer metadata is not available for model 'test-model'") + ); + } + + #[tokio::test] + async fn test_tokenize_route_defaults_model_when_only_one_is_served() { + let (state, card) = make_tokenize_state("test-model"); + let tokenizer = card.tokenizer().unwrap(); + let prompt = "Single served model default."; + let expected = tokenizer.encode_with_special_tokens(prompt, true).unwrap(); + + let response = tokenize( + State(state), + Json(TokenizeRequest::Completion(TokenizeCompletionRequest { + model: None, + prompt: prompt.to_string(), + add_special_tokens: true, + return_token_strs: false, + })), + ) + .await + .unwrap(); + let body: TokenizeResponse = response_json(response).await; + assert_eq!(body.tokens, expected.token_ids()); + } + + #[tokio::test] + async fn test_tokenize_chat_route_supports_request_chat_template_override() { + let (state, card) = make_tokenize_state("test-model"); + let tokenizer = card.tokenizer().unwrap(); + let messages = sample_chat_messages(); + let custom_template = concat!( + "{% for message in messages %}", + "[[{{ message['role'] }}]] {{ message['content'] }}\n", + "{% endfor %}", + "{% if add_generation_prompt %}[[assistant]] {% endif %}" + ); + + let request = TokenizeChatRequest { + model: Some("test-model".to_string()), + messages: messages.clone(), + add_generation_prompt: true, + return_token_strs: false, + continue_final_message: false, + add_special_tokens: false, + chat_template: Some(custom_template.to_string()), + chat_template_kwargs: Some(StdHashMap::from([( + "unused_value".to_string(), + serde_json::Value::String("ignored".to_string()), + )])), + media_io_kwargs: None, + mm_processor_kwargs: Some(StdHashMap::from([( + "image".to_string(), + serde_json::json!({"size": "ignored"}), + )])), + tools: None, + }; + let prompt = + render_tokenize_chat_prompt(&card, "test-model".to_string(), &request).unwrap(); + let expected = tokenizer + .encode_with_special_tokens(&prompt, false) + .unwrap(); + + let response = tokenize( + State(state), + Json(TokenizeRequest::Chat(TokenizeChatRequest { + model: Some("test-model".to_string()), + messages, + add_generation_prompt: true, + return_token_strs: false, + continue_final_message: false, + add_special_tokens: false, + chat_template: Some(custom_template.to_string()), + chat_template_kwargs: Some(StdHashMap::from([( + "unused_value".to_string(), + serde_json::Value::String("ignored".to_string()), + )])), + media_io_kwargs: None, + mm_processor_kwargs: Some(StdHashMap::from([( + "image".to_string(), + serde_json::json!({"size": "ignored"}), + )])), + tools: None, + })), + ) + .await + .unwrap(); + let body: TokenizeResponse = response_json(response).await; + assert_eq!(body.tokens, expected.token_ids()); + } + + #[test] + fn test_tokenization_router_registers_root_paths() { + let (state, _) = make_tokenize_state("test-model"); + let (docs, _) = tokenization_router(state); + let docs = docs.iter().map(|doc| doc.to_string()).collect::>(); + + assert!(docs.contains(&"POST /tokenize".to_string())); + assert!(docs.contains(&"POST /detokenize".to_string())); + } + // ── streaming dispatch tests ────────────────────────────────────── use std::collections::{HashMap, HashSet}; diff --git a/lib/llm/src/http/service/service_v2.rs b/lib/llm/src/http/service/service_v2.rs index 3d29fe0de6f1..8405a0ae09f1 100644 --- a/lib/llm/src/http/service/service_v2.rs +++ b/lib/llm/src/http/service/service_v2.rs @@ -525,6 +525,7 @@ impl HttpServiceConfigBuilder { config.drt_metrics, ), super::openai::list_models_router(state.clone(), var(HTTP_SVC_MODELS_PATH_ENV).ok()), + super::openai::tokenization_router(state.clone()), super::health::health_check_router(state.clone(), var(HTTP_SVC_HEALTH_PATH_ENV).ok()), super::health::live_check_router(state.clone(), var(HTTP_SVC_LIVE_PATH_ENV).ok()), super::busy_threshold::busy_threshold_router(state.clone(), None), diff --git a/lib/llm/src/model_card.rs b/lib/llm/src/model_card.rs index 771d850208ae..22d015c345c3 100644 --- a/lib/llm/src/model_card.rs +++ b/lib/llm/src/model_card.rs @@ -597,6 +597,29 @@ impl ModelDeploymentCard { Ok(()) } + /// Return the local directory that contains the model config files, if any + /// file has already been resolved to a local path. + pub(crate) fn local_file_dir(&self) -> Option<&Path> { + if let Some(TokenizerKind::HfTokenizerJson(cf) | TokenizerKind::TikTokenModel(cf)) = + &self.tokenizer + { + if let Some(dir) = cf.path().and_then(|p| p.parent()) { + return Some(dir); + } + } + if let Some(ModelInfoType::HfConfigJson(cf)) = &self.model_info { + if let Some(dir) = cf.path().and_then(|p| p.parent()) { + return Some(dir); + } + } + if let Some(PromptFormatterArtifact::HfTokenizerConfigJson(cf)) = &self.prompt_formatter { + if let Some(dir) = cf.path().and_then(|p| p.parent()) { + return Some(dir); + } + } + None + } + /// Are all the files we need (tokenizer.json, etc) available locally? fn has_local_files(&self) -> bool { let has_model_info = self @@ -633,7 +656,7 @@ impl ModelDeploymentCard { } /// Update the directory for files like tokenizer.json be in here. - fn update_dir(&mut self, dir: &Path) { + pub(crate) fn update_dir(&mut self, dir: &Path) { if let Some(model_info) = self.model_info.as_mut() { model_info.update_dir(dir); } diff --git a/lib/llm/src/preprocessor/prompt/template.rs b/lib/llm/src/preprocessor/prompt/template.rs index fbdf2da1af4a..07610a7dd10b 100644 --- a/lib/llm/src/preprocessor/prompt/template.rs +++ b/lib/llm/src/preprocessor/prompt/template.rs @@ -4,6 +4,7 @@ use std::{collections::HashSet, sync::Arc}; use anyhow::{Context, Ok, Result}; +use either::Either; use minijinja::Environment; use crate::model_card::{ModelDeploymentCard, PromptContextMixin, PromptFormatterArtifact}; @@ -19,11 +20,19 @@ use tokcfg::ChatTemplateValue; impl PromptFormatter { pub fn from_mdc(mdc: &ModelDeploymentCard) -> Result { + Self::from_mdc_with_chat_template(mdc, None) + } + + pub fn from_mdc_with_chat_template( + mdc: &ModelDeploymentCard, + chat_template_override: Option<&str>, + ) -> Result { // Special handling for DeepSeek-V3.2(-Speciale) which doesn't provide Jinja chat_template let name_lower = mdc.display_name.to_lowercase(); if name_lower.contains("deepseek") && name_lower.contains("v3.2") && !name_lower.contains("exp") + && chat_template_override.is_none() { tracing::info!("Detected DeepSeek V3.2 model (non-Exp), using native Rust formatter"); return Ok(Self::OAI(Arc::new( @@ -54,57 +63,64 @@ impl PromptFormatter { crate::log_json_err(&file.display().to_string(), &contents, err) })?; + if let Some(chat_template) = chat_template_override { + config.chat_template = + Some(ChatTemplateValue(Either::Left(chat_template.to_string()))); + } + // Some HF model (i.e. meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8) // stores the chat template in a separate file, we check if the file exists and // put the chat template into config as normalization. // This may also be a custom template provided via CLI flag. - match mdc.chat_template_file.as_ref() { - Some(PromptFormatterArtifact::HfChatTemplateJinja { - file: checked_file, - .. - }) => { - let Some(path) = checked_file.path() else { - anyhow::bail!( - "HfChatTemplateJinja for {} is a URL, cannot load", - mdc.display_name - ); - }; - let chat_template = std::fs::read_to_string(path) - .with_context(|| format!("fs:read_to_string '{}'", path.display()))?; - config.chat_template = Some(ChatTemplateValue(either::Left(chat_template))); - } - Some(PromptFormatterArtifact::HfChatTemplateJson { - file: checked_file, - .. - }) => { - let Some(path) = checked_file.path() else { - anyhow::bail!( - "HfChatTemplateJson for {} is a URL, cannot load", - mdc.display_name - ); - }; - let raw = std::fs::read_to_string(path) - .with_context(|| format!("fs:read_to_string '{}'", path.display()))?; - let wrapper: serde_json::Value = - serde_json::from_str(&raw).with_context(|| { - format!("Failed to parse '{}' as JSON", path.display()) - })?; - let field = wrapper.get("chat_template").ok_or_else(|| { - anyhow::anyhow!( - "'{}' does not contain a 'chat_template' field", - path.display() - ) - })?; - let value = serde_json::from_value::(field.clone()) - .with_context(|| { - format!( - "Failed to deserialize 'chat_template' in '{}'", + if chat_template_override.is_none() { + match mdc.chat_template_file.as_ref() { + Some(PromptFormatterArtifact::HfChatTemplateJinja { + file: checked_file, + .. + }) => { + let Some(path) = checked_file.path() else { + anyhow::bail!( + "HfChatTemplateJinja for {} is a URL, cannot load", + mdc.display_name + ); + }; + let chat_template = std::fs::read_to_string(path) + .with_context(|| format!("fs:read_to_string '{}'", path.display()))?; + config.chat_template = Some(ChatTemplateValue(either::Left(chat_template))); + } + Some(PromptFormatterArtifact::HfChatTemplateJson { + file: checked_file, + .. + }) => { + let Some(path) = checked_file.path() else { + anyhow::bail!( + "HfChatTemplateJson for {} is a URL, cannot load", + mdc.display_name + ); + }; + let raw = std::fs::read_to_string(path) + .with_context(|| format!("fs:read_to_string '{}'", path.display()))?; + let wrapper: serde_json::Value = + serde_json::from_str(&raw).with_context(|| { + format!("Failed to parse '{}' as JSON", path.display()) + })?; + let field = wrapper.get("chat_template").ok_or_else(|| { + anyhow::anyhow!( + "'{}' does not contain a 'chat_template' field", path.display() ) })?; - config.chat_template = Some(value); + let value = serde_json::from_value::(field.clone()) + .with_context(|| { + format!( + "Failed to deserialize 'chat_template' in '{}'", + path.display() + ) + })?; + config.chat_template = Some(value); + } + _ => {} } - _ => {} } Self::from_parts( config, diff --git a/lib/llm/src/protocols/anthropic/types.rs b/lib/llm/src/protocols/anthropic/types.rs index 179525d9be95..f4383a612853 100644 --- a/lib/llm/src/protocols/anthropic/types.rs +++ b/lib/llm/src/protocols/anthropic/types.rs @@ -137,6 +137,7 @@ impl TryFrom for NvCreateChatCompletionRequest { None }, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }) } diff --git a/lib/llm/src/protocols/common.rs b/lib/llm/src/protocols/common.rs index 730d56fe54f9..1f7d7035a269 100644 --- a/lib/llm/src/protocols/common.rs +++ b/lib/llm/src/protocols/common.rs @@ -514,6 +514,10 @@ pub struct OutputOptions { /// the tokenizer. This is useful for inspecting the behavior of prompt /// templates that are applied during the backend preprocessing. pub formatted_prompt: Option, + + /// When true, logprob token fields are returned as "token_id:" + /// instead of the decoded text. vLLM-specific extension for NeMo-RL. + pub return_tokens_as_token_ids: Option, } // Struct for log probability information diff --git a/lib/llm/src/protocols/openai.rs b/lib/llm/src/protocols/openai.rs index 42ef621f8797..0ddf2049996e 100644 --- a/lib/llm/src/protocols/openai.rs +++ b/lib/llm/src/protocols/openai.rs @@ -20,6 +20,7 @@ pub mod images; pub mod models; pub mod nvext; pub mod responses; +pub mod tokenization; pub mod tools; pub mod validate; pub mod videos; @@ -90,6 +91,10 @@ pub(crate) trait OpenAIOutputOptionsProvider { fn get_skip_special_tokens(&self) -> Option; fn get_formatted_prompt(&self) -> Option; + + fn get_return_tokens_as_token_ids(&self) -> Option { + None + } } impl SamplingOptionsProvider for T { @@ -203,12 +208,14 @@ impl OutputOptionsProvider for T { let prompt_logprobs = self.get_prompt_logprobs(); let skip_special_tokens = self.get_skip_special_tokens(); let formatted_prompt = self.get_formatted_prompt(); + let return_tokens_as_token_ids = self.get_return_tokens_as_token_ids(); Ok(common::OutputOptions { logprobs, prompt_logprobs, skip_special_tokens, formatted_prompt, + return_tokens_as_token_ids, }) } } diff --git a/lib/llm/src/protocols/openai/chat_completions.rs b/lib/llm/src/protocols/openai/chat_completions.rs index 8a77038d5834..f41324d30806 100644 --- a/lib/llm/src/protocols/openai/chat_completions.rs +++ b/lib/llm/src/protocols/openai/chat_completions.rs @@ -59,6 +59,12 @@ pub struct NvCreateChatCompletionRequest { #[serde(default, skip_serializing_if = "Option::is_none")] pub media_io_kwargs: Option, + /// When true, logprob token fields are returned as "token_id:" instead + /// of the decoded text. This is a vLLM-specific extension used by NeMo-RL + /// to extract per-token IDs for RL training. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub return_tokens_as_token_ids: Option, + /// Catch-all for unsupported fields - checked during validation #[serde(flatten, default, skip_serializing)] pub unsupported_fields: std::collections::HashMap, @@ -330,6 +336,10 @@ impl OpenAIOutputOptionsProvider for NvCreateChatCompletionRequest { fn get_formatted_prompt(&self) -> Option { None } + + fn get_return_tokens_as_token_ids(&self) -> Option { + self.return_tokens_as_token_ids + } } /// Implements `ValidateRequest` for `NvCreateChatCompletionRequest`, diff --git a/lib/llm/src/protocols/openai/chat_completions/delta.rs b/lib/llm/src/protocols/openai/chat_completions/delta.rs index 8bf31756bfa4..01f8c8f2ae28 100644 --- a/lib/llm/src/protocols/openai/chat_completions/delta.rs +++ b/lib/llm/src/protocols/openai/chat_completions/delta.rs @@ -82,6 +82,7 @@ impl NvCreateChatCompletionRequest { enable_logprobs: self.inner.logprobs.unwrap_or(false) || self.inner.top_logprobs.unwrap_or(0) > 0, enable_tracking, + return_tokens_as_token_ids: self.return_tokens_as_token_ids.unwrap_or(false), runtime_config: ModelRuntimeConfig::default(), }; @@ -100,6 +101,8 @@ pub struct DeltaGeneratorOptions { pub enable_logprobs: bool, /// Determines whether request tracking (timing, KV hit rate) should be enabled. pub enable_tracking: bool, + /// When true, logprob token fields use "token_id:" format instead of decoded text. + pub return_tokens_as_token_ids: bool, pub runtime_config: ModelRuntimeConfig, } @@ -210,16 +213,22 @@ impl DeltaGenerator { .map(|(_, lp)| lp as f32) .collect::>(); + let return_as_ids = self.options.return_tokens_as_token_ids; let content = top_logprobs.map(|top_logprobs| { toks.iter() .zip(tok_lps) .zip(top_logprobs) .map(|(((t, tid), lp), top_lps)| { + let token_str = if return_as_ids { + format!("token_id:{}", tid) + } else { + t.clone() + }; let converted = convert_backend_top_logprobs(&top_lps, t, *tid, lp); dynamo_protocols::types::ChatCompletionTokenLogprob { - token: t.clone(), + token: token_str.clone(), logprob: lp, - bytes: token_to_utf8_bytes(t), + bytes: token_to_utf8_bytes(&token_str), top_logprobs: converted, } }) @@ -525,6 +534,7 @@ mod tests { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), } } diff --git a/lib/llm/src/protocols/openai/chat_completions/jail.rs b/lib/llm/src/protocols/openai/chat_completions/jail.rs index 9ea6ca5a9f42..88c2f333ae23 100644 --- a/lib/llm/src/protocols/openai/chat_completions/jail.rs +++ b/lib/llm/src/protocols/openai/chat_completions/jail.rs @@ -100,6 +100,10 @@ struct ChoiceJailState { is_jailed: bool, /// Accumulated content for this choice while jailed accumulated_content: String, + /// Accumulated logprobs for this choice while jailed. + /// Logprobs from each jailed chunk are appended so the full token-level + /// log-probability information is preserved when the jail emits. + accumulated_logprobs: Option, /// Buffer for partial marker matches across chunks partial_match_buffer: String, /// Stream finish reason @@ -145,6 +149,7 @@ impl ChoiceJailState { index, is_jailed: starts_jailed, accumulated_content: String::new(), + accumulated_logprobs: None, partial_match_buffer: String::new(), stream_finish_reason: None, emitted_tool_calls_count: 0, @@ -152,16 +157,41 @@ impl ChoiceJailState { } } - /// Add content to this choice's accumulation - fn accumulate(&mut self, content: &str) { + /// Add content and logprobs to this choice's accumulation + fn accumulate(&mut self, content: &str, logprobs: Option<&ChatChoiceLogprobs>) { if self.is_jailed { self.accumulated_content.push_str(content); + // Accumulate logprobs so they are preserved across jailed chunks. + if let Some(lp) = logprobs { + let state_lps = self.accumulated_logprobs.get_or_insert(ChatChoiceLogprobs { + content: None, + refusal: None, + }); + if let Some(content_lps) = &lp.content { + state_lps + .content + .get_or_insert_with(Vec::new) + .extend(content_lps.clone()); + } + if let Some(refusal_lps) = &lp.refusal { + state_lps + .refusal + .get_or_insert_with(Vec::new) + .extend(refusal_lps.clone()); + } + } } } + /// Consume the accumulated logprobs, replacing them with `None`. + fn take_accumulated_logprobs(&mut self) -> Option { + self.accumulated_logprobs.take() + } + /// End jailing and return the accumulated content fn end_jail(&mut self) -> String { self.is_jailed = false; + self.accumulated_logprobs = None; std::mem::take(&mut self.accumulated_content) } @@ -235,6 +265,8 @@ impl ChoiceJailState { if jail_stream.should_start_jail(trailing_part) { self.is_jailed = true; self.accumulated_content = trailing_part.to_string(); + // No logprobs to seed here — they were already emitted with the tool call + self.accumulated_logprobs = None; } else { #[allow(deprecated)] let trailing_choice = create_choice_stream( @@ -253,6 +285,8 @@ impl ChoiceJailState { // Start jailing with the marker and suffix self.is_jailed = true; self.accumulated_content = full_content; + // Seed accumulated logprobs with this chunk's logprobs + self.accumulated_logprobs = choice.logprobs.clone(); } self.partial_match_buffer.clear(); @@ -301,6 +335,8 @@ impl ChoiceJailState { // Start jailing with the combined content self.is_jailed = true; self.accumulated_content = combined_content; + // Seed accumulated logprobs with this chunk's logprobs + self.accumulated_logprobs = choice.logprobs.clone(); self.partial_match_buffer.clear(); } else { // No markers - emit everything @@ -322,25 +358,31 @@ impl ChoiceJailState { } } } else { - // Already jailed - accumulate and check for unjail - self.accumulate(content); + // Already jailed - accumulate content AND logprobs, then check for unjail + self.accumulate(content, choice.logprobs.as_ref()); let (should_end, split_pos) = jail_stream.should_end_jail(&self.accumulated_content).await; if should_end { + // Take accumulated logprobs before borrowing accumulated_content + let jail_logprobs = self.take_accumulated_logprobs(); + // Split the content let (jailed_part, trailing_part) = self.accumulated_content.split_at(split_pos); + let trailing_owned = trailing_part.to_string(); + let jailed_owned = jailed_part.to_string(); - // Create the unjailed choice - let unjailed_choice = jail_stream + // Create the unjailed choice, using accumulated logprobs + let mut unjailed_choice = jail_stream .create_tool_call_choice( choice.index, - jailed_part, + &jailed_owned, choice, self.emitted_tool_calls_count, ) .await; + unjailed_choice.logprobs = jail_logprobs; // Determine emission type based on whether tool calls were parsed if unjailed_choice.delta.tool_calls.is_some() { @@ -353,7 +395,6 @@ impl ChoiceJailState { } // End jailing before processing trailing content - let trailing_owned = trailing_part.to_string(); self.end_jail(); // Handle trailing content if any @@ -393,7 +434,7 @@ impl ChoiceJailState { None, self.stream_finish_reason, // For the accumulated content, assign the original stream finish reason, otherwise it will get lost None, - None, + self.accumulated_logprobs.clone(), ); let mut final_choice = jail_stream @@ -404,6 +445,8 @@ impl ChoiceJailState { self.emitted_tool_calls_count, ) .await; + // Attach the full accumulated logprobs to the final choice + final_choice.logprobs = self.take_accumulated_logprobs(); // Preserve any pending reasoning content collected while jailed. if let Some(pending_reasoning) = self.pending_reasoning_content.take() { @@ -928,7 +971,7 @@ impl JailedStream { Some(tool_call_chunks), None, None, - None, + base_choice.logprobs.clone(), ); return choice; } @@ -1413,6 +1456,158 @@ mod tests { .collect() } + /// Helper: build a single-choice stream chunk with text content and logprobs + #[allow(deprecated)] + fn text_chunk_with_logprobs(text: &str) -> Annotated { + let logprobs = ChatChoiceLogprobs { + content: Some( + text.chars() + .enumerate() + .map(|(i, c)| dynamo_protocols::types::ChatCompletionTokenLogprob { + token: c.to_string(), + logprob: -(i as f32 + 1.0) * 0.1, + bytes: Some(c.to_string().into_bytes()), + top_logprobs: vec![], + }) + .collect(), + ), + refusal: None, + }; + + let choice = ChatChoiceStream { + index: 0, + delta: ChatCompletionStreamResponseDelta { + role: Some(Role::Assistant), + content: Some(dynamo_protocols::types::ChatCompletionMessageContent::Text( + text.to_string(), + )), + tool_calls: None, + function_call: None, + refusal: None, + reasoning_content: None, + }, + finish_reason: None, + stop_reason: None, + logprobs: Some(logprobs), + }; + + Annotated { + data: Some(NvCreateChatCompletionStreamResponse { + inner: CreateChatCompletionStreamResponse { + id: "id-42".to_string(), + object: "chat.completion.chunk".to_string(), + created: 0, + model: "test-model".to_string(), + choices: vec![choice], + usage: None, + service_tier: None, + system_fingerprint: None, + }, + nvext: None, + }), + id: None, + event: None, + comment: None, + error: None, + } + } + + /// Collect all logprobs from jailed stream output choices + fn collect_logprobs( + responses: &[Annotated], + ) -> Vec> { + responses + .iter() + .flat_map(|r| r.data.iter()) + .flat_map(|d| d.inner.choices.iter()) + .map(|c| c.logprobs.clone()) + .collect() + } + + #[tokio::test] + async fn test_tool_call_preserves_logprobs_single_chunk() { + let jail = JailedStream::builder().tool_call_parser("hermes").build(); + + let chunks = vec![text_chunk_with_logprobs( + "\n{\"name\": \"get_weather\", \"arguments\": {\"location\": \"SF\"}}\n", + )]; + + let input_stream = Box::pin(stream::iter(chunks)); + let output_stream = jail.apply_with_finish_reason(input_stream); + + let responses: Vec<_> = output_stream.collect().await; + let tool_calls = collect_tool_calls(&responses); + assert_eq!(tool_calls.len(), 1, "Expected 1 tool call, got {:?}", tool_calls); + assert_eq!(tool_calls[0].0, "get_weather"); + + // Logprobs must be preserved even though the entire output is a tool call + let all_logprobs = collect_logprobs(&responses); + let has_some_logprobs = all_logprobs.iter().any(|lp| lp.is_some()); + assert!( + has_some_logprobs, + "Logprobs should be preserved for tool call responses, got all None: {:?}", + all_logprobs + ); + } + + #[tokio::test] + async fn test_tool_call_preserves_logprobs_multiple_chunks() { + let jail = JailedStream::builder().tool_call_parser("hermes").build(); + + let chunks = vec![ + text_chunk_with_logprobs("\n{\"name\": \"get_weather\", \"arguments\""), + text_chunk_with_logprobs(": {\"location\": \"SF\"}}\n"), + ]; + + let input_stream = Box::pin(stream::iter(chunks)); + let output_stream = jail.apply_with_finish_reason(input_stream); + + let responses: Vec<_> = output_stream.collect().await; + let tool_calls = collect_tool_calls(&responses); + assert!(!tool_calls.is_empty(), "Expected tool calls, got none"); + + let all_logprobs = collect_logprobs(&responses); + let has_some_logprobs = all_logprobs.iter().any(|lp| lp.is_some()); + assert!( + has_some_logprobs, + "Logprobs should be preserved for tool call responses across chunks, got all None", + ); + } + + #[tokio::test] + async fn test_tool_call_with_text_preserves_logprobs() { + let jail = JailedStream::builder().tool_call_parser("hermes").build(); + + let chunks = vec![text_chunk_with_logprobs( + "Let me check.\n\n{\"name\": \"get_weather\", \"arguments\": {\"location\": \"SF\"}}\n", + )]; + + let input_stream = Box::pin(stream::iter(chunks)); + let output_stream = jail.apply_with_finish_reason(input_stream); + + let responses: Vec<_> = output_stream.collect().await; + let tool_calls = collect_tool_calls(&responses); + assert_eq!(tool_calls.len(), 1); + + let all_logprobs = collect_logprobs(&responses); + let has_some_logprobs = all_logprobs.iter().any(|lp| lp.is_some()); + assert!( + has_some_logprobs, + "Logprobs should be preserved for mixed text+tool_call responses", + ); + + // Verify the logprobs content is non-empty + let logprob_entries: Vec<_> = all_logprobs + .iter() + .filter_map(|lp| lp.as_ref()) + .filter_map(|lp| lp.content.as_ref()) + .collect(); + assert!( + logprob_entries.iter().any(|entries| !entries.is_empty()), + "Logprobs content should have entries", + ); + } + #[tokio::test] async fn test_multi_tool_call_single_chunk() { let jail = JailedStream::builder().tool_call_parser("hermes").build(); diff --git a/lib/llm/src/protocols/openai/responses/mod.rs b/lib/llm/src/protocols/openai/responses/mod.rs index 5750c66ee985..3b03d1bf9a05 100644 --- a/lib/llm/src/protocols/openai/responses/mod.rs +++ b/lib/llm/src/protocols/openai/responses/mod.rs @@ -532,6 +532,7 @@ impl TryFrom for NvCreateChatCompletionRequest { nvext: resp.nvext, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }) } diff --git a/lib/llm/src/protocols/openai/tokenization.rs b/lib/llm/src/protocols/openai/tokenization.rs new file mode 100644 index 000000000000..95559684ad89 --- /dev/null +++ b/lib/llm/src/protocols/openai/tokenization.rs @@ -0,0 +1,124 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::HashMap; + +use serde::{Deserialize, Serialize}; + +use crate::preprocessor::media::MediaDecoder; +use crate::types::TokenIdType; + +fn default_true() -> bool { + true +} + +fn default_false() -> bool { + false +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct TokenizeCompletionRequest { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub model: Option, + pub prompt: String, + #[serde(default = "default_true")] + pub add_special_tokens: bool, + #[serde(default = "default_false")] + pub return_token_strs: bool, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct TokenizeChatRequest { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub model: Option, + pub messages: Vec, + #[serde(default = "default_true")] + pub add_generation_prompt: bool, + #[serde(default = "default_false")] + pub return_token_strs: bool, + #[serde(default = "default_false")] + pub continue_final_message: bool, + #[serde(default = "default_false")] + pub add_special_tokens: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub chat_template: Option, + #[serde( + default, + skip_serializing_if = "Option::is_none", + alias = "chat_template_args" + )] + pub chat_template_kwargs: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub media_io_kwargs: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub mm_processor_kwargs: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tools: Option>, +} + +impl TokenizeChatRequest { + pub fn validate(&self) -> Result<(), String> { + if self.continue_final_message && self.add_generation_prompt { + return Err( + "Cannot set both `continue_final_message` and `add_generation_prompt` to True." + .to_string(), + ); + } + + Ok(()) + } + + pub fn merged_chat_template_kwargs(&self) -> HashMap { + let mut kwargs = self.chat_template_kwargs.clone().unwrap_or_default(); + kwargs.insert( + "add_generation_prompt".to_string(), + serde_json::Value::Bool(self.add_generation_prompt), + ); + kwargs.insert( + "continue_final_message".to_string(), + serde_json::Value::Bool(self.continue_final_message), + ); + kwargs + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(untagged)] +#[allow(clippy::large_enum_variant)] +pub enum TokenizeRequest { + Completion(TokenizeCompletionRequest), + Chat(TokenizeChatRequest), +} + +impl TokenizeRequest { + pub fn model(&self) -> Option<&str> { + match self { + Self::Completion(request) => request.model.as_deref(), + Self::Chat(request) => request.model.as_deref(), + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct TokenizeResponse { + pub count: usize, + pub max_model_len: u32, + pub tokens: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub token_strs: Option>, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct DetokenizeRequest { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub model: Option, + pub tokens: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct DetokenizeResponse { + pub prompt: String, +} diff --git a/lib/llm/src/protocols/unified.rs b/lib/llm/src/protocols/unified.rs index e2f8b97355e6..e36f5c4c8023 100644 --- a/lib/llm/src/protocols/unified.rs +++ b/lib/llm/src/protocols/unified.rs @@ -530,6 +530,7 @@ mod tests { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; diff --git a/lib/llm/src/tokenizers.rs b/lib/llm/src/tokenizers.rs index 21b411cc9de1..6329c68e5f79 100644 --- a/lib/llm/src/tokenizers.rs +++ b/lib/llm/src/tokenizers.rs @@ -60,6 +60,14 @@ pub mod traits { pub trait Encoder: Send + Sync { fn encode(&self, input: &str) -> Result; fn encode_batch(&self, inputs: &[&str]) -> Result>; + + fn encode_with_special_tokens( + &self, + input: &str, + _add_special_tokens: bool, + ) -> Result { + self.encode(input) + } } /// Implementations **must** use lossy UTF-8 conversion (e.g. `String::from_utf8_lossy`) @@ -71,8 +79,12 @@ pub mod traits { } pub trait Tokenizer: Encoder + Decoder { - // fn get_vocab_size(&self) -> usize; - // fn make_unique_clone(&self) -> Box; + fn convert_ids_to_tokens(&self, token_ids: &[TokenIdType]) -> Result> { + token_ids + .iter() + .map(|id| self.decode(std::slice::from_ref(id), false)) + .collect() + } } } @@ -93,6 +105,18 @@ impl Tokenizer { Ok(Tokenizer(create_tokenizer_from_file(file_path)?)) } + pub fn encode_with_special_tokens( + &self, + input: &str, + add_special_tokens: bool, + ) -> Result { + self.0.encode_with_special_tokens(input, add_special_tokens) + } + + pub fn convert_ids_to_tokens(&self, token_ids: &[TokenIdType]) -> Result> { + self.0.convert_ids_to_tokens(token_ids) + } + /// Create a stateful sequence object for decoding token_ids into text pub fn decode_stream( &self, diff --git a/lib/llm/src/tokenizers/fastokens.rs b/lib/llm/src/tokenizers/fastokens.rs index 83c5c7e3bade..3384dc435e88 100644 --- a/lib/llm/src/tokenizers/fastokens.rs +++ b/lib/llm/src/tokenizers/fastokens.rs @@ -39,16 +39,28 @@ impl FastTokenizer { impl Encoder for FastTokenizer { fn encode(&self, input: &str) -> Result { + self.encode_with_special_tokens(input, false) + } + + fn encode_batch(&self, inputs: &[&str]) -> Result> { + inputs.par_iter().map(|input| self.encode(input)).collect() + } + + fn encode_with_special_tokens( + &self, + input: &str, + add_special_tokens: bool, + ) -> Result { + if add_special_tokens { + return self.hf_decoder.encode_with_special_tokens(input, true); + } + let ids = self .fast_encoder .encode(input) .map_err(|e| Error::msg(format!("Fastokens encode error: {e}")))?; Ok(Encoding::Sp(ids)) } - - fn encode_batch(&self, inputs: &[&str]) -> Result> { - inputs.par_iter().map(|input| self.encode(input)).collect() - } } impl Decoder for FastTokenizer { @@ -57,7 +69,11 @@ impl Decoder for FastTokenizer { } } -impl Tokenizer for FastTokenizer {} +impl Tokenizer for FastTokenizer { + fn convert_ids_to_tokens(&self, token_ids: &[TokenIdType]) -> Result> { + self.hf_decoder.convert_ids_to_tokens(token_ids) + } +} #[cfg(test)] mod tests { diff --git a/lib/llm/src/tokenizers/hf.rs b/lib/llm/src/tokenizers/hf.rs index 0ca8f3fec838..9c45871c5bf4 100644 --- a/lib/llm/src/tokenizers/hf.rs +++ b/lib/llm/src/tokenizers/hf.rs @@ -27,13 +27,7 @@ impl HuggingFaceTokenizer { impl Encoder for HuggingFaceTokenizer { fn encode(&self, input: &str) -> Result { - // This self.tokenizer is the library - let encoding = self - .tokenizer - .encode(input, false) - .map_err(|err| Error::msg(format!("Error tokenizing input: {err}")))?; - - Ok(Encoding::Hf(Box::new(encoding))) + self.encode_with_special_tokens(input, false) } fn encode_batch(&self, inputs: &[&str]) -> Result> { @@ -49,6 +43,20 @@ impl Encoder for HuggingFaceTokenizer { Ok(encodings) } + + fn encode_with_special_tokens( + &self, + input: &str, + add_special_tokens: bool, + ) -> Result { + // This self.tokenizer is the library + let encoding = self + .tokenizer + .encode(input, add_special_tokens) + .map_err(|err| Error::msg(format!("Error tokenizing input: {err}")))?; + + Ok(Encoding::Hf(Box::new(encoding))) + } } impl Decoder for HuggingFaceTokenizer { @@ -63,7 +71,14 @@ impl Decoder for HuggingFaceTokenizer { } } -impl Tokenizer for HuggingFaceTokenizer {} +impl Tokenizer for HuggingFaceTokenizer { + fn convert_ids_to_tokens(&self, token_ids: &[TokenIdType]) -> Result> { + Ok(token_ids + .iter() + .map(|&id| self.tokenizer.id_to_token(id).unwrap_or_default()) + .collect()) + } +} impl From for HuggingFaceTokenizer { fn from(tokenizer: HfTokenizer) -> Self { diff --git a/lib/llm/src/tokenizers/tiktoken.rs b/lib/llm/src/tokenizers/tiktoken.rs index 798af8f1e76e..67c4c72f6221 100644 --- a/lib/llm/src/tokenizers/tiktoken.rs +++ b/lib/llm/src/tokenizers/tiktoken.rs @@ -24,6 +24,8 @@ const KIMI_PATTERN: &str = r#"[\p{Han}]+|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p pub struct TikTokenTokenizer { bpe: CoreBPE, special_token_ids: HashSet, + decoder_tokens: FxHashMap>, + special_tokens_decoder: FxHashMap>, } impl TikTokenTokenizer { @@ -39,6 +41,14 @@ impl TikTokenTokenizer { special_tokens: FxHashMap, ) -> Result { let encoder = parse_tiktoken_file(path)?; + let decoder_tokens: FxHashMap> = encoder + .iter() + .map(|(bytes, &id)| (id, bytes.clone())) + .collect(); + let special_tokens_decoder: FxHashMap> = special_tokens + .iter() + .map(|(token, &id)| (id, token.as_bytes().to_vec())) + .collect(); let special_token_ids: HashSet = special_tokens.values().copied().collect(); let bpe = CoreBPE::new(encoder, special_tokens, pattern) @@ -47,6 +57,8 @@ impl TikTokenTokenizer { Ok(Self { bpe, special_token_ids, + decoder_tokens, + special_tokens_decoder, }) } @@ -62,9 +74,17 @@ impl TikTokenTokenizer { let pattern = detect_bpe_pattern(directory)?; let encoder = parse_tiktoken_file(path)?; + let decoder_tokens: FxHashMap> = encoder + .iter() + .map(|(bytes, &id)| (id, bytes.clone())) + .collect(); // Use max rank + 1 (not len) to avoid ID collisions with sparse/non-contiguous ranks let num_base_tokens = encoder.values().max().map_or(0, |&m| m + 1) as usize; let special_tokens = load_special_tokens(directory, num_base_tokens)?; + let special_tokens_decoder: FxHashMap> = special_tokens + .iter() + .map(|(token, &id)| (id, token.as_bytes().to_vec())) + .collect(); let special_token_ids: HashSet = special_tokens.values().copied().collect(); let bpe = CoreBPE::new(encoder, special_tokens, pattern) @@ -73,19 +93,33 @@ impl TikTokenTokenizer { Ok(Self { bpe, special_token_ids, + decoder_tokens, + special_tokens_decoder, }) } } impl Encoder for TikTokenTokenizer { fn encode(&self, input: &str) -> Result { - let token_ids: Vec = self.bpe.encode_with_special_tokens(input); - Ok(Encoding::Sp(token_ids)) + self.encode_with_special_tokens(input, true) } fn encode_batch(&self, inputs: &[&str]) -> Result> { inputs.par_iter().map(|input| self.encode(input)).collect() } + + fn encode_with_special_tokens( + &self, + input: &str, + add_special_tokens: bool, + ) -> Result { + let token_ids: Vec = if add_special_tokens { + self.bpe.encode_with_special_tokens(input) + } else { + self.bpe.encode_ordinary(input) + }; + Ok(Encoding::Sp(token_ids)) + } } impl Decoder for TikTokenTokenizer { @@ -109,7 +143,20 @@ impl Decoder for TikTokenTokenizer { } } -impl Tokenizer for TikTokenTokenizer {} +impl Tokenizer for TikTokenTokenizer { + fn convert_ids_to_tokens(&self, token_ids: &[TokenIdType]) -> Result> { + Ok(token_ids + .iter() + .map(|id| { + self.decoder_tokens + .get(id) + .or_else(|| self.special_tokens_decoder.get(id)) + .map(|bytes| String::from_utf8_lossy(bytes).into_owned()) + .unwrap_or_default() + }) + .collect()) + } +} /// Parse a tiktoken model file (base64-encoded token + rank per line). fn parse_tiktoken_file(path: &str) -> Result, u32>> { diff --git a/lib/llm/tests/parallel_tool_call_integration.rs b/lib/llm/tests/parallel_tool_call_integration.rs index 2827239d4754..4eeb89026f74 100644 --- a/lib/llm/tests/parallel_tool_call_integration.rs +++ b/lib/llm/tests/parallel_tool_call_integration.rs @@ -92,6 +92,7 @@ fn create_mock_chat_completion_request() -> NvCreateChatCompletionRequest { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), } } diff --git a/lib/llm/tests/preprocessor.rs b/lib/llm/tests/preprocessor.rs index d5766e814525..a66ea9921dc5 100644 --- a/lib/llm/tests/preprocessor.rs +++ b/lib/llm/tests/preprocessor.rs @@ -260,6 +260,7 @@ impl Request { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), } } @@ -650,6 +651,7 @@ mod context_length_validation { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), } } diff --git a/lib/llm/tests/test_common_ext.rs b/lib/llm/tests/test_common_ext.rs index 8e49c7377b09..b2366001eab1 100644 --- a/lib/llm/tests/test_common_ext.rs +++ b/lib/llm/tests/test_common_ext.rs @@ -69,6 +69,7 @@ fn test_sampling_parameters_include_stop_str_in_output_extraction() { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; @@ -299,6 +300,7 @@ fn test_serialization_preserves_structure() { }), chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; @@ -351,6 +353,7 @@ fn test_sampling_parameters_extraction() { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; diff --git a/lib/llm/tests/test_streaming_usage.rs b/lib/llm/tests/test_streaming_usage.rs index 5ccd6476a666..679e3691b2c9 100644 --- a/lib/llm/tests/test_streaming_usage.rs +++ b/lib/llm/tests/test_streaming_usage.rs @@ -191,6 +191,7 @@ fn create_chat_request( nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), } } @@ -521,6 +522,7 @@ fn create_nonstreaming_chat_request() -> NvCreateChatCompletionRequest { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), } } diff --git a/lib/llm/tests/tool_choice.rs b/lib/llm/tests/tool_choice.rs index 2c182c55c0af..796747579305 100644 --- a/lib/llm/tests/tool_choice.rs +++ b/lib/llm/tests/tool_choice.rs @@ -40,6 +40,7 @@ fn create_test_request() -> NvCreateChatCompletionRequest { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), } } diff --git a/lib/llm/tests/tool_choice_finish_reasons.rs b/lib/llm/tests/tool_choice_finish_reasons.rs index a6f3b7998901..1947b455ea22 100644 --- a/lib/llm/tests/tool_choice_finish_reasons.rs +++ b/lib/llm/tests/tool_choice_finish_reasons.rs @@ -33,6 +33,7 @@ fn create_test_request() -> NvCreateChatCompletionRequest { nvext: None, chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), } } From 49a2140fc23f00cf70ec864654570ceb78e6ea18 Mon Sep 17 00:00:00 2001 From: William Arnold Date: Mon, 13 Apr 2026 09:59:30 -0700 Subject: [PATCH 02/20] feat(sglang): support return_tokens_as_token_ids for token-based logprobs Add return_tokens_as_token_ids support to the SGLang decode handler, mirroring what PR #7699 added for vLLM. When enabled, logprob token fields are returned as "token_id:" instead of decoded text. Changes: - decode_handler.py: Read return_tokens_as_token_ids from output_options, pass through _process_token_stream to _extract_logprobs, format token strings accordingly - sglang_processor.py: Forward return_tokens_as_token_ids through _build_dynamo_preproc output_options - vllm/handlers.py: Remove debug print left in cherry-picked code --- .../src/dynamo/frontend/sglang_processor.py | 1 + .../request_handlers/llm/decode_handler.py | 34 +++++++++++++++---- components/src/dynamo/vllm/handlers.py | 2 -- 3 files changed, 29 insertions(+), 8 deletions(-) diff --git a/components/src/dynamo/frontend/sglang_processor.py b/components/src/dynamo/frontend/sglang_processor.py index bad1db7a460a..dcb7a4e5dcf4 100644 --- a/components/src/dynamo/frontend/sglang_processor.py +++ b/components/src/dynamo/frontend/sglang_processor.py @@ -199,6 +199,7 @@ def _build_dynamo_preproc( "logprobs": logprobs_val, "prompt_logprobs": None, "skip_special_tokens": True, + "return_tokens_as_token_ids": request.get("return_tokens_as_token_ids"), }, "eos_token_ids": [eos_token_id] if eos_token_id is not None else [], "annotations": [], diff --git a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py index 1c7c447e828b..8c8f7754159e 100644 --- a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py +++ b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py @@ -190,7 +190,9 @@ def _build_logprob_kwargs(request: Dict[str, Any]) -> Dict[str, Any]: @staticmethod def _extract_logprobs( - meta_info: Dict[str, Any], num_output_logprobs_so_far: int + meta_info: Dict[str, Any], + num_output_logprobs_so_far: int, + return_tokens_as_token_ids: bool = False, ) -> tuple: """Extract logprobs from SGLang meta_info for new tokens. @@ -234,11 +236,17 @@ def _extract_logprobs( continue position_list = [] for rank_idx, entry in enumerate(position_entries): + tok_id = entry[1] + token_str = ( + f"token_id:{tok_id}" + if return_tokens_as_token_ids + else entry[2] + ) position_list.append( { "rank": rank_idx + 1, - "token_id": entry[1], - "token": entry[2], + "token_id": tok_id, + "token": token_str, "logprob": float(entry[0]), } ) @@ -272,6 +280,11 @@ async def generate( priority = (request.get("routing") or {}).get("priority") logprob_kwargs = self._build_logprob_kwargs(request) + output_options = request.get("output_options", {}) + return_tokens_as_token_ids = bool( + output_options.get("return_tokens_as_token_ids") + ) + if self.serving_mode == DisaggregationMode.DECODE: # Check if bootstrap_info is pre-computed in the request (from frontend) bootstrap_info = request.get("bootstrap_info") @@ -310,7 +323,9 @@ async def generate( ) if not self.use_sglang_tokenizer: - async for out in self._process_token_stream(decode, context): + async for out in self._process_token_stream( + decode, context, return_tokens_as_token_ids + ): yield out else: async for out in self._process_text_stream(decode, context): @@ -342,7 +357,9 @@ async def generate( **self._priority_kwargs(priority), ) if not self.use_sglang_tokenizer: - async for out in self._process_token_stream(agg, context): + async for out in self._process_token_stream( + agg, context, return_tokens_as_token_ids + ): yield out else: async for out in self._process_text_stream(agg, context): @@ -352,6 +369,7 @@ async def _process_token_stream( self, stream_source: AsyncGenerator[Dict[str, Any], None], context: Context, + return_tokens_as_token_ids: bool = False, ) -> AsyncGenerator[Dict[str, Any], None]: """Process token-based stream output. @@ -408,7 +426,11 @@ async def _process_token_stream( log_probs, top_logprobs, num_output_logprobs_so_far, - ) = self._extract_logprobs(res["meta_info"], num_output_logprobs_so_far) + ) = self._extract_logprobs( + res["meta_info"], + num_output_logprobs_so_far, + return_tokens_as_token_ids=return_tokens_as_token_ids, + ) if log_probs is not None: out["log_probs"] = log_probs if top_logprobs is not None: diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index d148e21bba2e..0b21364f92d0 100644 --- a/components/src/dynamo/vllm/handlers.py +++ b/components/src/dynamo/vllm/handlers.py @@ -1709,8 +1709,6 @@ async def _generate_token_mode(self, request, context, request_id): output_options.get("return_tokens_as_token_ids") ) - print(f"[DEBUG] output_options={output_options}, return_tokens_as_token_ids={return_tokens_as_token_ids}", flush=True) - async with self._abort_monitor(context, request_id): try: async for tok in self.generate_tokens( From 5664bf55fe44b476bd232b1d6932bac65f39bb6b Mon Sep 17 00:00:00 2001 From: William Arnold Date: Tue, 14 Apr 2026 09:57:44 -0700 Subject: [PATCH 03/20] feat: wire up return_tokens_as_token_ids for /v1/completions endpoint Add the return_tokens_as_token_ids field to NvCreateCompletionRequest and implement get_return_tokens_as_token_ids() so the completions endpoint has parity with chat completions for token-based logprobs. --- lib/llm/src/grpc/service/openai.rs | 1 + lib/llm/src/http/service/openai.rs | 7 +++++++ lib/llm/src/protocols/openai/completions.rs | 9 +++++++++ lib/llm/tests/openai_completions.rs | 1 + lib/llm/tests/test_streaming_usage.rs | 1 + 5 files changed, 19 insertions(+) diff --git a/lib/llm/src/grpc/service/openai.rs b/lib/llm/src/grpc/service/openai.rs index 561ff47dc142..6b0430fcd5c0 100644 --- a/lib/llm/src/grpc/service/openai.rs +++ b/lib/llm/src/grpc/service/openai.rs @@ -339,6 +339,7 @@ impl TryFrom for NvCreateCompletionRequest { common: Default::default(), nvext: None, metadata: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }) } diff --git a/lib/llm/src/http/service/openai.rs b/lib/llm/src/http/service/openai.rs index 9417f15a9fe6..3d9760fb89eb 100644 --- a/lib/llm/src/http/service/openai.rs +++ b/lib/llm/src/http/service/openai.rs @@ -3086,6 +3086,7 @@ mod tests { common: Default::default(), nvext: None, metadata: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; @@ -3110,6 +3111,7 @@ mod tests { common: Default::default(), nvext: None, metadata: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_completion_fields_generic(&request); @@ -3133,6 +3135,7 @@ mod tests { common: Default::default(), nvext: None, metadata: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_completion_fields_generic(&request); @@ -3156,6 +3159,7 @@ mod tests { common: Default::default(), nvext: None, metadata: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_completion_fields_generic(&request); @@ -3181,6 +3185,7 @@ mod tests { .unwrap(), nvext: None, metadata: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_completion_fields_generic(&request); @@ -3204,6 +3209,7 @@ mod tests { common: Default::default(), nvext: None, metadata: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; let result = validate_completion_fields_generic(&request); @@ -3235,6 +3241,7 @@ mod tests { "session": {"id": "session-1", "timestamp": 1640995200} }) .into(), + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; diff --git a/lib/llm/src/protocols/openai/completions.rs b/lib/llm/src/protocols/openai/completions.rs index 0537277522e9..98b5084d05d3 100644 --- a/lib/llm/src/protocols/openai/completions.rs +++ b/lib/llm/src/protocols/openai/completions.rs @@ -40,6 +40,11 @@ pub struct NvCreateCompletionRequest { #[serde(skip_serializing_if = "Option::is_none")] pub metadata: Option, + /// When true, logprob token fields are returned as "token_id:" + /// instead of the decoded text. + #[serde(skip_serializing_if = "Option::is_none")] + pub return_tokens_as_token_ids: Option, + /// Catch-all for unsupported fields - checked during validation #[serde(flatten, default, skip_serializing)] pub unsupported_fields: std::collections::HashMap, @@ -416,6 +421,10 @@ impl OpenAIOutputOptionsProvider for NvCreateCompletionRequest { fn get_formatted_prompt(&self) -> Option { None } + + fn get_return_tokens_as_token_ids(&self) -> Option { + self.return_tokens_as_token_ids + } } /// Implements `ValidateRequest` for `NvCreateCompletionRequest`, diff --git a/lib/llm/tests/openai_completions.rs b/lib/llm/tests/openai_completions.rs index 2d772b6af6e6..fd23d916ec44 100644 --- a/lib/llm/tests/openai_completions.rs +++ b/lib/llm/tests/openai_completions.rs @@ -29,6 +29,7 @@ impl CompletionSample { common: Default::default(), nvext: None, metadata: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), }; diff --git a/lib/llm/tests/test_streaming_usage.rs b/lib/llm/tests/test_streaming_usage.rs index 679e3691b2c9..273459c4837c 100644 --- a/lib/llm/tests/test_streaming_usage.rs +++ b/lib/llm/tests/test_streaming_usage.rs @@ -495,6 +495,7 @@ fn create_cmpl_request(include_usage: Option, stream: bool) -> NvCreateCom common: Default::default(), nvext: None, metadata: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), } } From 1fe31519db71fb6410298055323960fe2a6cc576 Mon Sep 17 00:00:00 2001 From: William Arnold Date: Thu, 16 Apr 2026 13:04:40 -0700 Subject: [PATCH 04/20] style: cargo fmt --- lib/llm/src/discovery/watcher.rs | 6 +----- lib/llm/src/preprocessor/prompt/template.rs | 18 +++++++++------- .../protocols/openai/chat_completions/jail.rs | 21 ++++++++++++------- lib/llm/src/tokenizers.rs | 5 ++++- 4 files changed, 30 insertions(+), 20 deletions(-) diff --git a/lib/llm/src/discovery/watcher.rs b/lib/llm/src/discovery/watcher.rs index e63866308386..3e34f63dc859 100644 --- a/lib/llm/src/discovery/watcher.rs +++ b/lib/llm/src/discovery/watcher.rs @@ -562,11 +562,7 @@ impl ModelWatcher { /// If an existing card for the same model has already-downloaded local files, /// point this card's URL-backed files at the same local directory. This avoids /// re-downloading config files for every worker that joins an existing WorkerSet. - fn resolve_card_local_files( - &self, - model_name: &str, - card: &mut ModelDeploymentCard, - ) { + fn resolve_card_local_files(&self, model_name: &str, card: &mut ModelDeploymentCard) { let local_dir = self .manager .get_model_cards() diff --git a/lib/llm/src/preprocessor/prompt/template.rs b/lib/llm/src/preprocessor/prompt/template.rs index 07610a7dd10b..66449b8559e7 100644 --- a/lib/llm/src/preprocessor/prompt/template.rs +++ b/lib/llm/src/preprocessor/prompt/template.rs @@ -84,9 +84,12 @@ impl PromptFormatter { mdc.display_name ); }; - let chat_template = std::fs::read_to_string(path) - .with_context(|| format!("fs:read_to_string '{}'", path.display()))?; - config.chat_template = Some(ChatTemplateValue(either::Left(chat_template))); + let chat_template = + std::fs::read_to_string(path).with_context(|| { + format!("fs:read_to_string '{}'", path.display()) + })?; + config.chat_template = + Some(ChatTemplateValue(either::Left(chat_template))); } Some(PromptFormatterArtifact::HfChatTemplateJson { file: checked_file, @@ -98,10 +101,11 @@ impl PromptFormatter { mdc.display_name ); }; - let raw = std::fs::read_to_string(path) - .with_context(|| format!("fs:read_to_string '{}'", path.display()))?; - let wrapper: serde_json::Value = - serde_json::from_str(&raw).with_context(|| { + let raw = std::fs::read_to_string(path).with_context(|| { + format!("fs:read_to_string '{}'", path.display()) + })?; + let wrapper: serde_json::Value = serde_json::from_str(&raw) + .with_context(|| { format!("Failed to parse '{}' as JSON", path.display()) })?; let field = wrapper.get("chat_template").ok_or_else(|| { diff --git a/lib/llm/src/protocols/openai/chat_completions/jail.rs b/lib/llm/src/protocols/openai/chat_completions/jail.rs index 88c2f333ae23..d2b1390cd65f 100644 --- a/lib/llm/src/protocols/openai/chat_completions/jail.rs +++ b/lib/llm/src/protocols/openai/chat_completions/jail.rs @@ -1463,12 +1463,14 @@ mod tests { content: Some( text.chars() .enumerate() - .map(|(i, c)| dynamo_protocols::types::ChatCompletionTokenLogprob { - token: c.to_string(), - logprob: -(i as f32 + 1.0) * 0.1, - bytes: Some(c.to_string().into_bytes()), - top_logprobs: vec![], - }) + .map( + |(i, c)| dynamo_protocols::types::ChatCompletionTokenLogprob { + token: c.to_string(), + logprob: -(i as f32 + 1.0) * 0.1, + bytes: Some(c.to_string().into_bytes()), + top_logprobs: vec![], + }, + ) .collect(), ), refusal: None, @@ -1537,7 +1539,12 @@ mod tests { let responses: Vec<_> = output_stream.collect().await; let tool_calls = collect_tool_calls(&responses); - assert_eq!(tool_calls.len(), 1, "Expected 1 tool call, got {:?}", tool_calls); + assert_eq!( + tool_calls.len(), + 1, + "Expected 1 tool call, got {:?}", + tool_calls + ); assert_eq!(tool_calls[0].0, "get_weather"); // Logprobs must be preserved even though the entire output is a tool call diff --git a/lib/llm/src/tokenizers.rs b/lib/llm/src/tokenizers.rs index f414e955e992..07d88d988024 100644 --- a/lib/llm/src/tokenizers.rs +++ b/lib/llm/src/tokenizers.rs @@ -137,7 +137,10 @@ pub mod traits { fn convert_ids_to_tokens(&self, token_ids: &[TokenIdType]) -> Result> { token_ids .iter() - .map(|id| self.decode(std::slice::from_ref(id), false).map(String::from)) + .map(|id| { + self.decode(std::slice::from_ref(id), false) + .map(String::from) + }) .collect() } } From e8286f9762ae5e87f4ef9909739bf6dfd904a494 Mon Sep 17 00:00:00 2001 From: William Arnold Date: Thu, 16 Apr 2026 13:14:11 -0700 Subject: [PATCH 05/20] fix(clippy): collapse nested if-let in local_file_dir --- lib/llm/src/model_card.rs | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/lib/llm/src/model_card.rs b/lib/llm/src/model_card.rs index 22d015c345c3..c7f4d1a50a69 100644 --- a/lib/llm/src/model_card.rs +++ b/lib/llm/src/model_card.rs @@ -602,20 +602,19 @@ impl ModelDeploymentCard { pub(crate) fn local_file_dir(&self) -> Option<&Path> { if let Some(TokenizerKind::HfTokenizerJson(cf) | TokenizerKind::TikTokenModel(cf)) = &self.tokenizer + && let Some(dir) = cf.path().and_then(|p| p.parent()) { - if let Some(dir) = cf.path().and_then(|p| p.parent()) { - return Some(dir); - } + return Some(dir); } - if let Some(ModelInfoType::HfConfigJson(cf)) = &self.model_info { - if let Some(dir) = cf.path().and_then(|p| p.parent()) { - return Some(dir); - } + if let Some(ModelInfoType::HfConfigJson(cf)) = &self.model_info + && let Some(dir) = cf.path().and_then(|p| p.parent()) + { + return Some(dir); } - if let Some(PromptFormatterArtifact::HfTokenizerConfigJson(cf)) = &self.prompt_formatter { - if let Some(dir) = cf.path().and_then(|p| p.parent()) { - return Some(dir); - } + if let Some(PromptFormatterArtifact::HfTokenizerConfigJson(cf)) = &self.prompt_formatter + && let Some(dir) = cf.path().and_then(|p| p.parent()) + { + return Some(dir); } None } From d4f71213c1d2c3621c42aa24a058ba460262dc66 Mon Sep 17 00:00:00 2001 From: William Arnold Date: Thu, 16 Apr 2026 16:40:42 -0700 Subject: [PATCH 06/20] fix(completions): honor return_tokens_as_token_ids in logprob tokens The /v1/completions path was building the logprob tokens list from the decoded token strings, ignoring return_tokens_as_token_ids. Plumb the flag through DeltaGeneratorOptions and emit "token_id:" strings in the tokens field when set, mirroring what chat_completions/delta.rs already does for its selected-token field. Fixes 100% fallback_tokenize overhead on clients that rely on the "token_id:N" format to skip client-side retokenization. --- .../src/protocols/openai/completions/delta.rs | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/lib/llm/src/protocols/openai/completions/delta.rs b/lib/llm/src/protocols/openai/completions/delta.rs index f40f5105aedb..7930872b82aa 100644 --- a/lib/llm/src/protocols/openai/completions/delta.rs +++ b/lib/llm/src/protocols/openai/completions/delta.rs @@ -75,6 +75,7 @@ impl NvCreateCompletionRequest { .unwrap_or(false), enable_logprobs: self.inner.logprobs.unwrap_or(0) > 0, enable_tracking, + return_tokens_as_token_ids: self.return_tokens_as_token_ids.unwrap_or(false), }; DeltaGenerator::new(self.inner.model.clone(), options, request_id) @@ -87,6 +88,8 @@ pub struct DeltaGeneratorOptions { pub continuous_usage_stats: bool, pub enable_logprobs: bool, pub enable_tracking: bool, + /// When true, logprob token fields use "token_id:" format instead of decoded text. + pub return_tokens_as_token_ids: bool, } pub struct DeltaGenerator { @@ -170,6 +173,7 @@ impl DeltaGenerator { .map(|(_, lp)| lp as f32) .collect::>(); + let return_as_ids = self.options.return_tokens_as_token_ids; let top_lps = top_logprobs.map_or(vec![], |top_logprobs| { toks.iter() .zip(tok_lps.iter()) @@ -181,8 +185,19 @@ impl DeltaGenerator { .collect() }); + let tokens_out: Vec = toks + .iter() + .map(|(t, tid)| { + if return_as_ids { + format!("token_id:{}", tid) + } else { + t.clone() + } + }) + .collect(); + Some(dynamo_protocols::types::Logprobs { - tokens: toks.iter().map(|(t, _)| t.clone()).collect(), + tokens: tokens_out, token_logprobs: tok_lps.into_iter().map(Some).collect(), text_offset: vec![], top_logprobs: top_lps, From c62b8037d6c811e7f1624c74016c7586e69a16a5 Mon Sep 17 00:00:00 2001 From: William Arnold Date: Thu, 16 Apr 2026 18:38:47 -0700 Subject: [PATCH 07/20] fix(mypy): annotate token_str as str | None in _extract_logprobs --- components/src/dynamo/vllm/handlers.py | 1 + 1 file changed, 1 insertion(+) diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index 6fbd315a987a..b1d8dc590bc2 100644 --- a/components/src/dynamo/vllm/handlers.py +++ b/components/src/dynamo/vllm/handlers.py @@ -1478,6 +1478,7 @@ def _extract_logprobs( # Build top_logprobs list for this token position token_top_logprobs = [] for tok_id, logprob_info in token_logprobs_dict.items(): + token_str: str | None if return_tokens_as_token_ids: token_str = f"token_id:{tok_id}" else: From baabe55d789ffc5c936e6a8307dc249ea02a7bba Mon Sep 17 00:00:00 2001 From: William Arnold Date: Fri, 17 Apr 2026 15:25:08 -0700 Subject: [PATCH 08/20] feat(sglang): reject logprobs >= 1 unless DYN_SGL_ALLOW_TOP_LOGPROBS set SGLang's tokenizer manager detokenizes top-k tokens per-position serially, causing O(N) latency per generated token. Silently dropping the top_logprobs feature is worse than surfacing the limitation, so raise a clear ValueError when callers request logprobs>=1 (or prompt_logprobs>=1) and pin top_logprobs_num=0 as a belt-and-suspenders guard. Escape hatch: DYN_SGL_ALLOW_TOP_LOGPROBS=1 restores the previous passthrough for use once upstream batches detokenize_top_logprobs_tokens. Update CLAUDE.md to document the gate. --- components/src/dynamo/sglang/CLAUDE.md | 6 ++ .../request_handlers/llm/decode_handler.py | 90 ++++++++++++------- 2 files changed, 62 insertions(+), 34 deletions(-) diff --git a/components/src/dynamo/sglang/CLAUDE.md b/components/src/dynamo/sglang/CLAUDE.md index b074cb218c58..19f34ae0002a 100644 --- a/components/src/dynamo/sglang/CLAUDE.md +++ b/components/src/dynamo/sglang/CLAUDE.md @@ -225,6 +225,12 @@ absolute sequence position where logprob computation starts: `-1` (default) = ou only (`len(prompt) - 1`), `0` = from prompt start. We set it to 0 when `prompt_logprobs` is requested. +**Top-logprobs gate**: `logprobs >= 1` (or `prompt_logprobs >= 1`) raises `ValueError` +by default. SGLang's tokenizer manager detokenizes top-k tokens per-position serially, +causing severe latency degradation (O(N) per generated token). Callers must use +`logprobs=0` for chosen-token-only logprobs. Set `DYN_SGL_ALLOW_TOP_LOGPROBS=1` to +override once upstream batches `detokenize_top_logprobs_tokens`. + **Streaming behavior** (`_extract_logprobs`): Dynamo forces `stream_output=True` (args.py:374), making `output_ids` disjoint per chunk. diff --git a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py index 3177f5b85c83..cbd2d0ad74b3 100644 --- a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py +++ b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py @@ -3,6 +3,7 @@ import asyncio import logging +import os import time from typing import Any, AsyncGenerator, Dict, Optional @@ -17,6 +18,25 @@ from dynamo.sglang.publisher import DynamoSglangPublisher from dynamo.sglang.request_handlers.handler_base import BaseWorkerHandler +# Escape hatch: set to "1" (or any truthy value) to allow top_logprobs_num >= 1. +# Default-off because SGLang's tokenizer manager detokenizes top-k tokens +# per-position serially (O(N) per generated token), causing severe latency +# degradation. Flip once upstream batches detokenize_top_logprobs_tokens. +_ALLOW_TOP_LOGPROBS_ENV = "DYN_SGL_ALLOW_TOP_LOGPROBS" + +_TOP_LOGPROBS_UNSUPPORTED_MSG = ( + "Dynamo's SGLang backend does not currently support logprobs >= 1 due to " + "an O(N) per-position detokenization in the upstream sglang tokenizer " + "manager. Use logprobs=0 for chosen-token logprobs, or set " + "DYN_SGL_ALLOW_TOP_LOGPROBS=1 to override at your own risk. " + "Track the upstream fix at https://github.com/sgl-project/sglang/issues/." +) + + +def _top_logprobs_allowed() -> bool: + """Return True if the DYN_SGL_ALLOW_TOP_LOGPROBS escape hatch is enabled.""" + return os.environ.get(_ALLOW_TOP_LOGPROBS_ENV, "").lower() not in ("", "0", "false") + def _extract_media_urls(mm_data: Dict[str, Any], media_key: str) -> list[str] | None: """Normalize multimodal URL items from the frontend wire format.""" @@ -149,48 +169,50 @@ def _build_logprob_kwargs(request: Dict[str, Any]) -> Dict[str, Any]: if not output_options: return kwargs - logprobs_value = output_options.get("logprobs") - if logprobs_value is not None: + allow_top = _top_logprobs_allowed() + + def _parse(name: str, value: Any) -> Optional[int]: try: - parsed = int(logprobs_value) - if parsed < 0: - logging.warning( - f"Invalid logprobs value: {logprobs_value} " - "(must be non-negative), ignoring" - ) - else: - kwargs["return_logprob"] = True - kwargs["top_logprobs_num"] = parsed + parsed = int(value) except (ValueError, TypeError): logging.warning( - f"Invalid logprobs value: {logprobs_value} " - "(must be integer), ignoring" + f"Invalid {name} value: {value} (must be integer), ignoring" + ) + return None + if parsed < 0: + logging.warning( + f"Invalid {name} value: {value} (must be non-negative), ignoring" ) + return None + if parsed >= 1 and not allow_top: + raise ValueError(_TOP_LOGPROBS_UNSUPPORTED_MSG) + return parsed + + logprobs_value = output_options.get("logprobs") + if logprobs_value is not None: + parsed = _parse("logprobs", logprobs_value) + if parsed is not None: + kwargs["return_logprob"] = True + kwargs["top_logprobs_num"] = parsed prompt_logprobs_value = output_options.get("prompt_logprobs") if prompt_logprobs_value is not None: - try: - parsed = int(prompt_logprobs_value) - if parsed < 0: - logging.warning( - f"Invalid prompt_logprobs value: {prompt_logprobs_value} " - "(must be non-negative), ignoring" - ) - else: - kwargs["return_logprob"] = True - # SGLang has a single top_logprobs_num for both prompt - # and output tokens, so take the max of the two. - kwargs["top_logprobs_num"] = max( - kwargs.get("top_logprobs_num", 0), parsed - ) - # logprob_start_len=0 computes from prompt start; - # omitting it (or -1) computes output tokens only. - kwargs["logprob_start_len"] = 0 - except (ValueError, TypeError): - logging.warning( - f"Invalid prompt_logprobs value: {prompt_logprobs_value} " - "(must be integer), ignoring" + parsed = _parse("prompt_logprobs", prompt_logprobs_value) + if parsed is not None: + kwargs["return_logprob"] = True + # SGLang has a single top_logprobs_num for both prompt + # and output tokens, so take the max of the two. + kwargs["top_logprobs_num"] = max( + kwargs.get("top_logprobs_num", 0), parsed ) + # logprob_start_len=0 computes from prompt start; + # omitting it (or -1) computes output tokens only. + kwargs["logprob_start_len"] = 0 + + # Belt-and-suspenders: if return_logprob was requested and the gate is + # not open, pin top_logprobs_num=0 so no future code path can flip it on. + if kwargs.get("return_logprob") and not allow_top: + kwargs["top_logprobs_num"] = 0 return kwargs From d013b5ceb372035ea3145bf25150ba69f3137156 Mon Sep 17 00:00:00 2001 From: William Arnold Date: Tue, 21 Apr 2026 14:28:35 -0700 Subject: [PATCH 09/20] fix(completions): enable logprobs when logprobs=0 (chosen-token only) response_generator gated logprobs on `logprobs.unwrap_or(0) > 0`, which evaluated to false for logprobs=0. create_logprobs then short-circuited to None, so the chosen-token logprobs the backend already computed were dropped from the response. Switch to `logprobs.is_some()` so Some(0) (chosen-token-only) enables the field while None preserves the off-by-default behavior. Required by the SGLang backend's logprobs>=1 rejection policy, which forces logprobs=0 as the only supported positive value. --- lib/llm/src/protocols/openai/completions/delta.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/llm/src/protocols/openai/completions/delta.rs b/lib/llm/src/protocols/openai/completions/delta.rs index 8d15051b7f76..3533c7d21972 100644 --- a/lib/llm/src/protocols/openai/completions/delta.rs +++ b/lib/llm/src/protocols/openai/completions/delta.rs @@ -60,7 +60,7 @@ impl NvCreateCompletionRequest { .as_ref() .map(|opts| opts.continuous_usage_stats) .unwrap_or(false), - enable_logprobs: self.inner.logprobs.unwrap_or(0) > 0, + enable_logprobs: self.inner.logprobs.is_some(), response_fields, return_tokens_as_token_ids: self.return_tokens_as_token_ids.unwrap_or(false), }; From 5fc895298bd4d5c6392d9bf562db9cd3489e41ec Mon Sep 17 00:00:00 2001 From: William Arnold Date: Fri, 1 May 2026 10:51:44 -0700 Subject: [PATCH 10/20] feat(completions): propagate SGLang stop_reason to /v1/completions responses Adds stop_reason (matched stop sequence, token ID, or token ID list) to the completion choice schema, plumbed end-to-end: - protocols: replace upstream async_openai Choice/CreateCompletionResponse with custom dynamo types that include stop_reason; add StopReason::IntArray for SGLang's matched-token-list shape - Rust delta/aggregator/backend: thread stop_reason through choice creation and stream aggregation - SGLang decode handler: extract finish_reason.matched (str | int | list[int]) into stop_reason; reject bool and list-of-str - frontend post-processor: forward stop_reason on chat choices Tests: - Python: parametrized _extract_sglang_stop_reason coverage and end-to-end post-processor passthrough - Rust: choice_from_postprocessor preserves stop_reason; serde JSON shape on Choice --- .../src/dynamo/frontend/sglang_prepost.py | 16 ++++- .../tests/test_sglang_processor_unit.py | 13 ++++ .../request_handlers/llm/decode_handler.py | 25 ++++++++ .../tests/test_sglang_decode_handler.py | 21 ++++++- lib/llm/src/backend.rs | 2 +- lib/llm/src/engines.rs | 4 +- lib/llm/src/protocols/common.rs | 5 ++ lib/llm/src/protocols/openai/completions.rs | 1 + .../openai/completions/aggregator.rs | 23 +++++++- .../src/protocols/openai/completions/delta.rs | 32 +++++++++- lib/llm/tests/kserve_service.rs | 2 +- lib/protocols/src/types/chat.rs | 3 +- lib/protocols/src/types/completion.rs | 59 ++++++++++++++++++- 13 files changed, 192 insertions(+), 14 deletions(-) diff --git a/components/src/dynamo/frontend/sglang_prepost.py b/components/src/dynamo/frontend/sglang_prepost.py index c34018385e29..76b7c7213805 100644 --- a/components/src/dynamo/frontend/sglang_prepost.py +++ b/components/src/dynamo/frontend/sglang_prepost.py @@ -487,24 +487,31 @@ def process_output(self, engine_response: dict[str, Any]) -> dict[str, Any] | No raw_ids = engine_response.get("token_ids") token_ids = raw_ids if isinstance(raw_ids, list) else list(raw_ids or []) finish_reason = engine_response.get("finish_reason") + stop_reason = engine_response.get("stop_reason") delta_text = self._incremental_decode(token_ids) if token_ids else "" if self._fast_plain_text: if delta_text: - return { + choice = { "index": 0, "delta": {"role": "assistant", "content": delta_text}, "finish_reason": finish_reason, "logprobs": None, } + if stop_reason is not None: + choice["stop_reason"] = stop_reason + return choice elif finish_reason: - return { + choice = { "index": 0, "delta": {}, "finish_reason": finish_reason, "logprobs": None, } + if stop_reason is not None: + choice["stop_reason"] = stop_reason + return choice return None # -- Reasoning parsing -- @@ -699,11 +706,14 @@ def process_output(self, engine_response: dict[str, Any]) -> dict[str, Any] | No effective_finish = "tool_calls" if has_content or effective_finish: - return { + choice = { "index": 0, "delta": delta if has_content else {}, "finish_reason": effective_finish, "logprobs": None, } + if stop_reason is not None: + choice["stop_reason"] = stop_reason + return choice return None diff --git a/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py b/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py index cb1a73c588d4..bddee938b82f 100644 --- a/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py +++ b/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py @@ -1050,6 +1050,19 @@ def test_finish_reason_only(self, tokenizer): assert choice is not None assert choice["finish_reason"] == "stop" + def test_stop_reason_passthrough(self, tokenizer): + """Backend stop_reason is included on the emitted choice.""" + post = SglangStreamingPostProcessor( + tokenizer=tokenizer, tool_call_parser=None, reasoning_parser=None + ) + + choice = post.process_output( + {"token_ids": [], "finish_reason": "stop", "stop_reason": "END"} + ) + + assert choice is not None + assert choice["stop_reason"] == "END" + def test_lookback_trimming(self, tokenizer): """Verify _all_token_ids doesn't grow unbounded.""" post = SglangStreamingPostProcessor( diff --git a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py index cbd2d0ad74b3..b2179c8ed05a 100644 --- a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py +++ b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py @@ -59,6 +59,25 @@ def _extract_media_urls(mm_data: Dict[str, Any], media_key: str) -> list[str] | return urls or None +def _extract_sglang_stop_reason(finish_reason: Dict[str, Any] | None) -> Any | None: + """Extract SGLang's matched stop value for Dynamo's stop_reason field.""" + + if not finish_reason: + return None + + matched = finish_reason.get("matched") + if isinstance(matched, bool): + return None + if isinstance(matched, (str, int)): + return matched + if isinstance(matched, list) and all( + isinstance(item, int) and not isinstance(item, bool) for item in matched + ): + return matched + + return None + + class DecodeWorkerHandler(BaseWorkerHandler): """Handler for decode workers in both aggregated and disaggregated serving modes.""" @@ -444,6 +463,9 @@ async def _process_token_stream( out["finish_reason"] = normalize_finish_reason( finish_reason["type"] ) + stop_reason = _extract_sglang_stop_reason(finish_reason) + if stop_reason is not None: + out["stop_reason"] = stop_reason # With stream_output=True, output_ids contains only new tokens (disjoint) output_ids = res.get("output_ids", []) @@ -546,6 +568,9 @@ async def _process_text_stream( "delta": {"role": "assistant", "content": delta}, "finish_reason": finish_reason_type, } + stop_reason = _extract_sglang_stop_reason(finish_reason) + if stop_reason is not None: + choice_data["stop_reason"] = stop_reason response = { "id": res["meta_info"]["id"], diff --git a/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py b/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py index 7e3d546549b8..8e9bd88f5b6f 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py +++ b/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py @@ -3,7 +3,10 @@ import pytest -from dynamo.sglang.request_handlers.llm.decode_handler import _extract_media_urls +from dynamo.sglang.request_handlers.llm.decode_handler import ( + _extract_media_urls, + _extract_sglang_stop_reason, +) pytestmark = [ pytest.mark.unit, @@ -34,3 +37,19 @@ def test_extract_media_urls_returns_none_for_missing_or_invalid_items(): assert ( _extract_media_urls({"image_url": [{"ignored": "value"}]}, "image_url") is None ) + + +@pytest.mark.parametrize( + ("finish_reason", "expected"), + [ + ({"type": "stop", "matched": "END"}, "END"), + ({"type": "stop", "matched": 128001}, 128001), + ({"type": "stop", "matched": [128001, 128009]}, [128001, 128009]), + ({"type": "stop", "matched": True}, None), + ({"type": "stop", "matched": ["END"]}, None), + ({"type": "length"}, None), + (None, None), + ], +) +def test_extract_sglang_stop_reason(finish_reason, expected): + assert _extract_sglang_stop_reason(finish_reason) == expected diff --git a/lib/llm/src/backend.rs b/lib/llm/src/backend.rs index 5f682976537e..160a598fa326 100644 --- a/lib/llm/src/backend.rs +++ b/lib/llm/src/backend.rs @@ -268,7 +268,7 @@ impl // which we don't want to propagate to `data.finish_reason`. if finish_reason.is_some() { data.finish_reason = finish_reason; - data.stop_reason = stop_reason; + data.stop_reason = stop_reason.or(data.stop_reason); } data.text = text; data.tokens = Some(tokens); diff --git a/lib/llm/src/engines.rs b/lib/llm/src/engines.rs index 74dd84e2d3b8..7ead7108034b 100644 --- a/lib/llm/src/engines.rs +++ b/lib/llm/src/engines.rs @@ -192,11 +192,11 @@ impl let mut id = 1; for c in chars_string.chars() { tokio::time::sleep(*TOKEN_ECHO_DELAY).await; - let response = deltas.create_choice(0, Some(c.to_string()), None, None); + let response = deltas.create_choice(0, Some(c.to_string()), None, None, None); yield Annotated{ id: Some(id.to_string()), data: Some(response), event: None, comment: None, error: None }; id += 1; } - let response = deltas.create_choice(0, None, Some(dynamo_protocols::types::CompletionFinishReason::Stop), None); + let response = deltas.create_choice(0, None, Some(dynamo_protocols::types::CompletionFinishReason::Stop), None, None); yield Annotated { id: Some(id.to_string()), data: Some(response), event: None, comment: None, error: None }; }; diff --git a/lib/llm/src/protocols/common.rs b/lib/llm/src/protocols/common.rs index da595e9f321a..9ba9a8915e80 100644 --- a/lib/llm/src/protocols/common.rs +++ b/lib/llm/src/protocols/common.rs @@ -18,6 +18,7 @@ use derive_builder::Builder; use serde::{Deserialize, Serialize}; use super::TokenIdType; +use dynamo_protocols::types::StopReason; /// Maximum nesting depth allowed in guided_grammar EBNF strings. const MAX_GRAMMAR_NESTING_DEPTH: usize = 500; @@ -608,6 +609,10 @@ pub struct Delta { pub finish_reason: Option, + /// The stop string or token that triggered the stop condition. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub stop_reason: Option, + // new token_ids pub token_ids: Option>, diff --git a/lib/llm/src/protocols/openai/completions.rs b/lib/llm/src/protocols/openai/completions.rs index 98b5084d05d3..8bed1ab0602f 100644 --- a/lib/llm/src/protocols/openai/completions.rs +++ b/lib/llm/src/protocols/openai/completions.rs @@ -397,6 +397,7 @@ impl TryFrom for dynamo_protocols::types::C index, logprobs, finish_reason, + stop_reason: response.delta.stop_reason, }; Ok(choice) diff --git a/lib/llm/src/protocols/openai/completions/aggregator.rs b/lib/llm/src/protocols/openai/completions/aggregator.rs index 6bea09f59adb..ab5d72ccc784 100644 --- a/lib/llm/src/protocols/openai/completions/aggregator.rs +++ b/lib/llm/src/protocols/openai/completions/aggregator.rs @@ -31,6 +31,7 @@ struct DeltaChoice { index: u32, text: String, finish_reason: Option, + stop_reason: Option, logprobs: Option, } @@ -100,6 +101,7 @@ impl DeltaAggregator { index: choice.index, text: "".to_string(), finish_reason: None, + stop_reason: None, logprobs: None, }); @@ -121,6 +123,10 @@ impl DeltaAggregator { None => None, }; + if let Some(stop_reason) = choice.stop_reason { + state_choice.stop_reason = Some(stop_reason); + } + // Update logprobs if let Some(logprobs) = &choice.logprobs { let state_lps = state_choice.logprobs.get_or_insert( @@ -187,6 +193,7 @@ impl From for dynamo_protocols::types::Choice { index: delta.index, text: delta.text, finish_reason, + stop_reason: delta.stop_reason, logprobs: delta.logprobs, } } @@ -255,6 +262,7 @@ mod tests { index, text: text.to_string(), finish_reason, + stop_reason: None, logprobs, }], object: "text_completion".to_string(), @@ -334,8 +342,15 @@ mod tests { // One will have a MessageRole and no FinishReason, // the other will have a FinishReason and no MessageRole let annotated_delta1 = create_test_delta(0, "Hello,", None, Some(-0.1)); - let annotated_delta2 = + let mut annotated_delta2 = create_test_delta(0, " world!", Some("stop".to_string()), Some(-0.2)); + annotated_delta2 + .data + .as_mut() + .expect("delta data") + .inner + .choices[0] + .stop_reason = Some(dynamo_protocols::types::StopReason::Int(128001)); // Create a stream let annotated_deltas = vec![annotated_delta1, annotated_delta2]; @@ -357,6 +372,10 @@ mod tests { choice.finish_reason, Some(dynamo_protocols::types::CompletionFinishReason::Stop) ); + assert_eq!( + choice.stop_reason, + Some(dynamo_protocols::types::StopReason::Int(128001)) + ); assert_eq!(choice.logprobs.as_ref().unwrap().tokens.len(), 2); assert_eq!( choice.logprobs.as_ref().unwrap().token_logprobs, @@ -378,12 +397,14 @@ mod tests { index: 0, text: "Choice 0".to_string(), finish_reason: Some(dynamo_protocols::types::CompletionFinishReason::Stop), + stop_reason: None, logprobs: None, }, dynamo_protocols::types::Choice { index: 1, text: "Choice 1".to_string(), finish_reason: Some(dynamo_protocols::types::CompletionFinishReason::Stop), + stop_reason: None, logprobs: None, }, ], diff --git a/lib/llm/src/protocols/openai/completions/delta.rs b/lib/llm/src/protocols/openai/completions/delta.rs index 3533c7d21972..ec5a07b75eee 100644 --- a/lib/llm/src/protocols/openai/completions/delta.rs +++ b/lib/llm/src/protocols/openai/completions/delta.rs @@ -197,6 +197,7 @@ impl DeltaGenerator { index: u32, text: Option, finish_reason: Option, + stop_reason: Option, logprobs: Option, ) -> NvCreateCompletionResponse { // todo - update for tool calling @@ -214,6 +215,7 @@ impl DeltaGenerator { text: text.unwrap_or_default(), index, finish_reason, + stop_reason, logprobs, }], usage: if self.options.enable_usage && self.options.continuous_usage_stats { @@ -309,7 +311,13 @@ impl crate::protocols::openai::DeltaGeneratorExt for // create choice let index = delta.index.unwrap_or(0); - let mut response = self.create_choice(index, delta.text.clone(), finish_reason, logprobs); + let mut response = self.create_choice( + index, + delta.text.clone(), + finish_reason, + delta.stop_reason, + logprobs, + ); // Record finish for timing/ITL accounting even when timing is not returned to the client. // Kept at call site because it's a side effect on the tracker — not a gating decision. @@ -392,6 +400,7 @@ mod tests { common: Default::default(), nvext: None, metadata: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), } } @@ -437,6 +446,27 @@ mod tests { assert!(response.nvext.is_none()); } + #[test] + fn test_choice_from_postprocessor_preserves_stop_reason() { + let request = create_test_request(); + let mut generator = request.response_generator("req-stop-reason".to_string()); + let mut output = final_backend_output(); + output.stop_reason = Some(dynamo_protocols::types::StopReason::String( + "END".to_string(), + )); + + let response = generator + .choice_from_postprocessor(output) + .expect("choice generation"); + + assert_eq!( + response.inner.choices[0].stop_reason, + Some(dynamo_protocols::types::StopReason::String( + "END".to_string() + )) + ); + } + #[test] fn test_timing_extra_field_emits_timing_on_final_chunk() { use crate::protocols::openai::nvext::NvExt; diff --git a/lib/llm/tests/kserve_service.rs b/lib/llm/tests/kserve_service.rs index 0f24f450b88a..dfde0cd9097b 100644 --- a/lib/llm/tests/kserve_service.rs +++ b/lib/llm/tests/kserve_service.rs @@ -119,7 +119,7 @@ pub mod kserve_test { let stream = stream! { tokio::time::sleep(std::time::Duration::from_millis(10)).await; for word in word_list { - yield Annotated::from_data(generator.create_choice(0, Some(word.to_string()), None, None)); + yield Annotated::from_data(generator.create_choice(0, Some(word.to_string()), None, None, None)); } }; diff --git a/lib/protocols/src/types/chat.rs b/lib/protocols/src/types/chat.rs index c9cef996b9df..38fdfad11a75 100644 --- a/lib/protocols/src/types/chat.rs +++ b/lib/protocols/src/types/chat.rs @@ -51,7 +51,6 @@ pub use async_openai::types::chat::{ ChatCompletionRequestToolMessageContentPart, ChatCompletionResponseMessageAudio, ChatCompletionTokenLogprob, - Choice, CompletionFinishReason, CompletionTokensDetails, CompletionUsage, @@ -281,11 +280,13 @@ pub struct ChatCompletionTool { /// Inference backends (vLLM, SGLang) report which stop condition triggered: /// - `String`: a matched user-provided stop sequence /// - `Int`: a matched stop token ID +/// - `IntArray`: matched stop token IDs reported as a sequence #[derive(Debug, Deserialize, Serialize, Clone, PartialEq)] #[serde(untagged)] pub enum StopReason { String(String), Int(i64), + IntArray(Vec), } /// Reasoning content from a previous assistant turn. diff --git a/lib/protocols/src/types/completion.rs b/lib/protocols/src/types/completion.rs index d86a99f0a40c..cf55a51c9383 100644 --- a/lib/protocols/src/types/completion.rs +++ b/lib/protocols/src/types/completion.rs @@ -13,10 +13,47 @@ use serde::{Deserialize, Serialize}; use crate::error::OpenAIError; -use super::{ChatCompletionStreamOptions, Choice, CompletionUsage, Prompt, Stop}; +use super::{ + ChatCompletionStreamOptions, CompletionFinishReason, CompletionUsage, Logprobs, Prompt, Stop, + StopReason, +}; -// Re-export response type from upstream (identical) -pub use async_openai::types::completions::CreateCompletionResponse; +/// Completion choice with inference-serving extensions. +/// +/// Extends upstream `Choice` with: +/// - `stop_reason`: the matched stop sequence, token ID, or token ID sequence +/// reported by inference backends +#[derive(Debug, Deserialize, Serialize, Clone, PartialEq)] +pub struct Choice { + pub text: String, + pub index: u32, + #[serde(skip_serializing_if = "Option::is_none")] + pub logprobs: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub finish_reason: Option, + /// Matched stop condition from the backend. + #[serde(skip_serializing_if = "Option::is_none")] + pub stop_reason: Option, +} + +/// Non-streaming or streaming text completion response. +#[derive(Debug, Deserialize, Clone, PartialEq, Serialize)] +pub struct CreateCompletionResponse { + /// A unique identifier for the completion. + pub id: String, + pub choices: Vec, + /// The Unix timestamp (in seconds) of when the completion was created. + pub created: u32, + + /// The model used for completion. + pub model: String, + /// This fingerprint represents the backend configuration that the model runs with. + pub system_fingerprint: Option, + + /// The object type, which is always "text_completion". + pub object: String, + pub usage: Option, +} /// Custom deserializer for the echo parameter that only accepts booleans. /// Rejects integers and strings with clear error messages. @@ -170,4 +207,20 @@ mod tests { assert!(err_msg.contains("string")); assert!(err_msg.contains("echo parameter")); } + + #[test] + fn completion_choice_serializes_stop_reason() { + let choice = Choice { + text: "hello".to_string(), + index: 0, + logprobs: None, + finish_reason: Some(CompletionFinishReason::Stop), + stop_reason: Some(StopReason::String("END".to_string())), + }; + + let value = serde_json::to_value(choice).expect("serialize choice"); + + assert_eq!(value["finish_reason"], "stop"); + assert_eq!(value["stop_reason"], "END"); + } } From 316b08b20d92a7c6f4761b90fe70c94a8fff633e Mon Sep 17 00:00:00 2001 From: William Arnold Date: Fri, 1 May 2026 11:57:12 -0700 Subject: [PATCH 11/20] fix(tests): add return_tokens_as_token_ids: None to test fixtures --- lib/llm/src/protocols/openai/chat_completions/delta.rs | 1 + lib/llm/src/protocols/openai/completions/delta.rs | 1 + 2 files changed, 2 insertions(+) diff --git a/lib/llm/src/protocols/openai/chat_completions/delta.rs b/lib/llm/src/protocols/openai/chat_completions/delta.rs index 1c9f4658acd4..3f4700c4f768 100644 --- a/lib/llm/src/protocols/openai/chat_completions/delta.rs +++ b/lib/llm/src/protocols/openai/chat_completions/delta.rs @@ -598,6 +598,7 @@ mod tests { ), chat_template_args: None, media_io_kwargs: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), } } diff --git a/lib/llm/src/protocols/openai/completions/delta.rs b/lib/llm/src/protocols/openai/completions/delta.rs index 3960199956de..e45470c54f40 100644 --- a/lib/llm/src/protocols/openai/completions/delta.rs +++ b/lib/llm/src/protocols/openai/completions/delta.rs @@ -451,6 +451,7 @@ mod tests { .unwrap(), ), metadata: None, + return_tokens_as_token_ids: None, unsupported_fields: Default::default(), } } From 6231dabb4d8edc699067724434e4fee801a07fdd Mon Sep 17 00:00:00 2001 From: William Arnold Date: Fri, 1 May 2026 17:48:51 -0700 Subject: [PATCH 12/20] fix(sglang): trim token logprob branch scope --- .../tests/test_sglang_decode_handler.py | 47 ++ lib/llm/src/discovery/watcher.rs | 25 - lib/llm/src/http/service/openai.rs | 761 +----------------- lib/llm/src/http/service/service_v2.rs | 1 - lib/llm/src/model_card.rs | 24 +- lib/llm/src/preprocessor/prompt/template.rs | 141 ++-- lib/llm/src/protocols/common.rs | 2 +- lib/llm/src/protocols/openai.rs | 11 +- .../src/protocols/openai/chat_completions.rs | 3 +- .../openai/chat_completions/delta.rs | 3 +- .../src/protocols/openai/completions/delta.rs | 58 +- lib/llm/src/protocols/openai/tokenization.rs | 124 --- lib/tokenizers/src/fastokens.rs | 26 +- lib/tokenizers/src/hf.rs | 31 +- lib/tokenizers/src/lib.rs | 31 +- lib/tokenizers/src/tiktoken.rs | 53 +- 16 files changed, 195 insertions(+), 1146 deletions(-) delete mode 100644 lib/llm/src/protocols/openai/tokenization.rs diff --git a/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py b/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py index 87fef3097cac..86d6629d5032 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py +++ b/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py @@ -113,6 +113,53 @@ def test_build_sampling_params_passes_n_for_sglang_tokenizer_requests(): assert sampling_params["max_new_tokens"] == 8 +def test_build_logprob_kwargs_allows_chosen_token_logprobs(monkeypatch): + monkeypatch.delenv("DYN_SGL_ALLOW_TOP_LOGPROBS", raising=False) + + kwargs = DecodeWorkerHandler._build_logprob_kwargs( + {"output_options": {"logprobs": 0}} + ) + + assert kwargs == {"return_logprob": True, "top_logprobs_num": 0} + + +def test_build_logprob_kwargs_rejects_top_logprobs_by_default(monkeypatch): + monkeypatch.delenv("DYN_SGL_ALLOW_TOP_LOGPROBS", raising=False) + + with pytest.raises(ValueError, match="does not currently support logprobs >= 1"): + DecodeWorkerHandler._build_logprob_kwargs({"output_options": {"logprobs": 1}}) + + +def test_build_logprob_kwargs_allows_top_logprobs_with_escape_hatch(monkeypatch): + monkeypatch.setenv("DYN_SGL_ALLOW_TOP_LOGPROBS", "1") + + kwargs = DecodeWorkerHandler._build_logprob_kwargs( + {"output_options": {"logprobs": 2}} + ) + + assert kwargs == {"return_logprob": True, "top_logprobs_num": 2} + + +def test_extract_logprobs_formats_top_tokens_as_token_ids(): + log_probs, top_logprobs, total = DecodeWorkerHandler._extract_logprobs( + { + "output_token_logprobs": [(-0.1, 101, "a")], + "output_top_logprobs": [[(-0.1, 101, "a"), (-0.2, 102, "b")]], + }, + 0, + return_tokens_as_token_ids=True, + ) + + assert log_probs == [-0.1] + assert top_logprobs == [ + [ + {"rank": 1, "token_id": 101, "token": "token_id:101", "logprob": -0.1}, + {"rank": 2, "token_id": 102, "token": "token_id:102", "logprob": -0.2}, + ] + ] + assert total == 1 + + @pytest.mark.asyncio async def test_process_token_stream_tracks_logprobs_per_choice_index(): handler = _new_decode_handler() diff --git a/lib/llm/src/discovery/watcher.rs b/lib/llm/src/discovery/watcher.rs index 85baa326514c..ee8430b9b5eb 100644 --- a/lib/llm/src/discovery/watcher.rs +++ b/lib/llm/src/discovery/watcher.rs @@ -485,7 +485,6 @@ impl ModelWatcher { "Checksum mismatch for worker in namespace {namespace}" )); } - self.resolve_card_local_files(&model_name, card); self.manager .save_model_card(&mcid.to_path(), card.clone())?; tracing::debug!( @@ -512,14 +511,6 @@ impl ModelWatcher { ) .await? { - self.resolve_card_local_files(&model_name, card); - self.manager - .save_model_card(&mcid.to_path(), card.clone())?; - tracing::debug!( - model_name = card.name(), - namespace = namespace, - "WorkerSet registration in progress, skipping" - ); return Ok(()); } @@ -650,22 +641,6 @@ impl ModelWatcher { Ok(false) } - /// If an existing card for the same model has already-downloaded local files, - /// point this card's URL-backed files at the same local directory. This avoids - /// re-downloading config files for every worker that joins an existing WorkerSet. - fn resolve_card_local_files(&self, model_name: &str, card: &mut ModelDeploymentCard) { - let local_dir = self - .manager - .get_model_cards() - .iter() - .find(|c| c.name() == model_name) - .and_then(|c| c.local_file_dir().map(|p| p.to_path_buf())); - - if let Some(dir) = local_dir { - card.update_dir(&dir); - } - } - /// Build a complete WorkerSet with all engines for this (model, namespace) /// and add it to the Model. async fn do_worker_set_registration( diff --git a/lib/llm/src/http/service/openai.rs b/lib/llm/src/http/service/openai.rs index 1e31d4a71119..a65441c33677 100644 --- a/lib/llm/src/http/service/openai.rs +++ b/lib/llm/src/http/service/openai.rs @@ -43,8 +43,6 @@ use super::{ service_v2, }; use crate::engines::ValidateRequest; -use crate::model_card::ModelDeploymentCard; -use crate::preprocessor::prompt::PromptFormatter; use crate::protocols::openai::chat_completions::aggregator::ChatCompletionAggregator; use crate::protocols::openai::nvext::apply_header_routing_overrides; use crate::protocols::openai::{ @@ -57,10 +55,6 @@ use crate::protocols::openai::{ embeddings::{NvCreateEmbeddingRequest, NvCreateEmbeddingResponse}, images::{NvCreateImageRequest, NvImagesResponse}, responses::{NvCreateResponse, NvResponse, ResponseParams, chat_completion_to_response}, - tokenization::{ - DetokenizeRequest, DetokenizeResponse, TokenizeChatRequest, TokenizeCompletionRequest, - TokenizeRequest, TokenizeResponse, - }, videos::{NvCreateVideoRequest, NvVideosResponse}, }; use crate::protocols::unified::UnifiedRequest; @@ -333,176 +327,6 @@ pub async fn smart_json_error_middleware(request: Request, next: Next) -> } } -fn bad_request>(message: T) -> ErrorResponse { - let code = StatusCode::BAD_REQUEST; - ( - code, - Json(ErrorMessage { - message: message.into(), - error_type: map_error_code_to_error_type(code), - code: code.as_u16(), - }), - ) -} - -fn resolve_tokenizer_model_name( - state: &Arc, - requested_model: Option<&str>, -) -> Result { - if let Some(model) = requested_model { - if state.manager().has_model_any(model) { - return Ok(model.to_string()); - } - return Err(ErrorMessage::model_not_found()); - } - - // Preserve this order: without an explicit model, prefer `model_display_names()` first because - // those names are known to the serving layer; only if that yields one choice do we fall back - // to `get_model_cards()`, whose card-only metadata may exist before a model is fully - // registered. This keeps tokenizer endpoints usable from cards alone, but card names may not - // map to serving-capable models, so explicit requests still gate on `has_model_any()` / - // `ErrorMessage::model_not_found()` and ambiguous fallback still returns `bad_request()`. - let served_models = state.manager().model_display_names(); - if served_models.len() == 1 { - return Ok(served_models.into_iter().next().unwrap()); - } - - let card_models: HashSet = state - .manager() - .get_model_cards() - .into_iter() - .map(|card| card.display_name) - .collect(); - if card_models.len() == 1 { - return Ok(card_models.into_iter().next().unwrap()); - } - - Err(bad_request( - "Model must be specified when more than one model is served.", - )) -} - -fn resolve_model_card( - state: &Arc, - requested_model: Option<&str>, -) -> Result<(String, ModelDeploymentCard), ErrorResponse> { - let model = resolve_tokenizer_model_name(state, requested_model)?; - let card = state - .manager() - .get_model_cards() - .into_iter() - .find(|card| card.display_name == model) - .ok_or_else(|| { - ErrorMessage::internal_server_error(&format!( - "Tokenizer metadata is not available for model '{}'", - model - )) - })?; - Ok((model, card)) -} - -fn extract_assistant_content_text( - content: &dynamo_protocols::types::ChatCompletionRequestAssistantMessageContent, -) -> String { - use dynamo_protocols::types::ChatCompletionRequestAssistantMessageContent as Content; - use dynamo_protocols::types::ChatCompletionRequestAssistantMessageContentPart as Part; - - match content { - Content::Text(text) => text.clone(), - Content::Array(parts) => parts - .iter() - .filter_map(|part| match part { - Part::Text(text) => Some(text.text.clone()), - Part::Refusal(_) => None, - }) - .collect::>() - .join(""), - } -} - -fn apply_continue_final_message( - rendered_prompt: String, - messages: &[dynamo_protocols::types::ChatCompletionRequestMessage], -) -> Result { - use dynamo_protocols::types::ChatCompletionRequestMessage as Message; - - let Some(Message::Assistant(message)) = messages.last() else { - return Err(bad_request( - "Cannot set `continue_final_message` to True when the final message is not from the assistant.", - )); - }; - - let Some(content) = message.content.as_ref() else { - return Err(bad_request( - "Cannot set `continue_final_message` to True when the final assistant message has no content.", - )); - }; - - let final_message = extract_assistant_content_text(content); - let trimmed_final_message = final_message.trim(); - if trimmed_final_message.is_empty() { - return Err(bad_request( - "Cannot set `continue_final_message` to True when the final assistant message content is empty.", - )); - } - - // Use rfind to locate the last occurrence of the assistant content in the rendered prompt, - // then truncate everything after it (e.g. EOS tokens, generation prompts). This assumes the - // final assistant message text appears only once at the end of the rendered output. - let Some(final_msg_loc) = rendered_prompt.rfind(trimmed_final_message) else { - return Err(ErrorMessage::internal_server_error( - "Failed to trim rendered prompt for `continue_final_message`.", - )); - }; - - Ok(rendered_prompt[..final_msg_loc + trimmed_final_message.len()].to_string()) -} - -fn make_tokenize_chat_completion_request( - model: String, - request: &TokenizeChatRequest, -) -> NvCreateChatCompletionRequest { - let inner = dynamo_protocols::types::CreateChatCompletionRequest { - model, - messages: request.messages.clone(), - tools: request.tools.clone(), - ..Default::default() - }; - - NvCreateChatCompletionRequest { - inner, - common: Default::default(), - nvext: None, - chat_template_args: Some(request.merged_chat_template_kwargs()), - media_io_kwargs: request.media_io_kwargs.clone(), - return_tokens_as_token_ids: None, - unsupported_fields: Default::default(), - } -} - -fn render_tokenize_chat_prompt( - card: &ModelDeploymentCard, - model: String, - request: &TokenizeChatRequest, -) -> Result { - request.validate().map_err(bad_request)?; - - let formatter = - PromptFormatter::from_mdc_with_chat_template(card, request.chat_template.as_deref()) - .map_err(|err| ErrorMessage::from_anyhow(err, "Failed to build chat formatter"))?; - let wrapped_request = make_tokenize_chat_completion_request(model, request); - let mut prompt = match formatter { - PromptFormatter::OAI(formatter) => formatter.render(&wrapped_request), - } - .map_err(|err| ErrorMessage::from_anyhow(err, "Failed to render chat prompt"))?; - - if request.continue_final_message { - prompt = apply_continue_final_message(prompt, &request.messages)?; - } - - Ok(prompt) -} - /// Return the request ID for the current request. /// /// The canonical request ID is set by `make_inference_request_span()` and stored @@ -2163,105 +1987,6 @@ struct ModelListing { max_output_tokens: Option, } -async fn tokenize( - State(state): State>, - Json(request): Json, -) -> Result { - check_ready(&state)?; - - let (_, card) = resolve_model_card(&state, request.model())?; - let tokenizer = card - .tokenizer() - .map_err(|err| ErrorMessage::from_anyhow(err, "Failed to load tokenizer"))?; - - let (tokens, token_strs) = match request { - TokenizeRequest::Completion(TokenizeCompletionRequest { - prompt, - add_special_tokens, - return_token_strs, - .. - }) => { - let encoding = tokenizer - .encode_with_special_tokens(&prompt, add_special_tokens) - .map_err(|err| ErrorMessage::from_anyhow(err, "Failed to tokenize prompt"))?; - let token_ids = encoding.token_ids().to_vec(); - let token_strs = if return_token_strs { - Some(tokenizer.convert_ids_to_tokens(&token_ids).map_err(|err| { - ErrorMessage::from_anyhow(err, "Failed to resolve token strings") - })?) - } else { - None - }; - (token_ids, token_strs) - } - TokenizeRequest::Chat(request) => { - let model = request - .model - .clone() - .unwrap_or_else(|| card.display_name.clone()); - let prompt = render_tokenize_chat_prompt(&card, model, &request)?; - let encoding = tokenizer - .encode_with_special_tokens(&prompt, request.add_special_tokens) - .map_err(|err| { - ErrorMessage::from_anyhow(err, "Failed to tokenize rendered chat prompt") - })?; - let token_ids = encoding.token_ids().to_vec(); - let token_strs = if request.return_token_strs { - Some(tokenizer.convert_ids_to_tokens(&token_ids).map_err(|err| { - ErrorMessage::from_anyhow(err, "Failed to resolve token strings") - })?) - } else { - None - }; - (token_ids, token_strs) - } - }; - - Ok(Json(TokenizeResponse { - count: tokens.len(), - max_model_len: card.context_length, - tokens, - token_strs, - }) - .into_response()) -} - -async fn detokenize( - State(state): State>, - Json(request): Json, -) -> Result { - check_ready(&state)?; - - let (_, card) = resolve_model_card(&state, request.model.as_deref())?; - let tokenizer = card - .tokenizer() - .map_err(|err| ErrorMessage::from_anyhow(err, "Failed to load tokenizer"))?; - let prompt = tokenizer - .decode(&request.tokens, false) - .map_err(|err| ErrorMessage::from_anyhow(err, "Failed to detokenize prompt"))?; - - Ok(Json(DetokenizeResponse { - prompt: prompt.into(), - }) - .into_response()) -} - -pub fn tokenization_router(state: Arc) -> (Vec, Router) { - let tokenize_path = "/tokenize"; - let detokenize_path = "/detokenize"; - let docs = vec![ - RouteDoc::new(axum::http::Method::POST, tokenize_path), - RouteDoc::new(axum::http::Method::POST, detokenize_path), - ]; - let router = Router::new() - .route(tokenize_path, post(tokenize)) - .route(detokenize_path, post(detokenize)) - .layer(middleware::from_fn(smart_json_error_middleware)) - .layer(axum::extract::DefaultBodyLimit::max(get_body_limit())) - .with_state(state); - (docs, router) -} - /// Create an Axum [`Router`] for the OpenAI API Completions endpoint /// If not path is provided, the default path is `/v1/completions` pub fn completions_router( @@ -2918,121 +2643,19 @@ pub fn audios_router( mod tests { use super::*; - use crate::discovery::{ModelManager, ModelManagerError}; + use crate::discovery::ModelManagerError; use crate::protocols::openai::chat_completions::NvCreateChatCompletionRequest; use crate::protocols::openai::common_ext::CommonExt; use crate::protocols::openai::completions::NvCreateCompletionRequest; use crate::protocols::openai::responses::NvCreateResponse; - use crate::protocols::openai::tokenization::DetokenizeRequest; - use axum::extract::State; use dynamo_protocols::types::responses::{CreateResponse, Input, PromptConfig}; use dynamo_protocols::types::{ - ChatCompletionRequestAssistantMessage, ChatCompletionRequestAssistantMessageContent, ChatCompletionRequestMessage, ChatCompletionRequestUserMessage, - ChatCompletionRequestUserMessageContent, ChatCompletionTool, CreateChatCompletionRequest, - CreateCompletionRequest, FunctionObject, + ChatCompletionRequestUserMessageContent, CreateChatCompletionRequest, + CreateCompletionRequest, }; - use dynamo_runtime::discovery::{MockDiscovery, SharedMockRegistry}; - use std::collections::HashMap as StdHashMap; - use tokio_util::sync::CancellationToken; const BACKUP_ERROR_MESSAGE: &str = "Failed to generate completions"; - const TOKENIZE_MODEL_PATH: &str = concat!( - env!("CARGO_MANIFEST_DIR"), - "/tests/data/sample-models/mock-llama-3.1-8b-instruct" - ); - const DETOKENIZE_MODEL_PATH: &str = concat!( - env!("CARGO_MANIFEST_DIR"), - "/tests/data/sample-models/TinyLlama_v1.1" - ); - - fn make_tokenize_state_with_path( - model_name: &str, - model_path: &str, - ) -> (Arc, ModelDeploymentCard) { - let mut card = ModelDeploymentCard::load_from_disk(model_path, None).unwrap(); - card.set_name(model_name); - - let manager = Arc::new(ModelManager::new()); - manager - .save_model_card(&format!("__test_model_card_{model_name}"), card.clone()) - .unwrap(); - manager - .add_prefill_model(model_name, card.mdcsum()) - .unwrap(); - - let discovery = Arc::new(MockDiscovery::new(None, SharedMockRegistry::new())); - let state = Arc::new(service_v2::State::new( - manager, - discovery, - CancellationToken::new(), - )); - (state, card) - } - - fn make_tokenize_state(model_name: &str) -> (Arc, ModelDeploymentCard) { - make_tokenize_state_with_path(model_name, TOKENIZE_MODEL_PATH) - } - - fn make_tokenize_state_without_card(model_name: &str) -> Arc { - let manager = Arc::new(ModelManager::new()); - manager - .add_prefill_model(model_name, "missing-card") - .unwrap(); - - let discovery = Arc::new(MockDiscovery::new(None, SharedMockRegistry::new())); - Arc::new(service_v2::State::new( - manager, - discovery, - CancellationToken::new(), - )) - } - - async fn response_json(response: Response) -> T { - let body = response.into_body(); - let bytes = axum::body::to_bytes(body, usize::MAX).await.unwrap(); - serde_json::from_slice(&bytes).unwrap() - } - - fn sample_chat_messages() -> Vec { - vec![ - ChatCompletionRequestMessage::User(ChatCompletionRequestUserMessage { - content: ChatCompletionRequestUserMessageContent::Text("Hi there!".to_string()), - name: None, - }), - ChatCompletionRequestMessage::Assistant(ChatCompletionRequestAssistantMessage { - content: Some(ChatCompletionRequestAssistantMessageContent::Text( - "Nice to meet you!".to_string(), - )), - ..Default::default() - }), - ChatCompletionRequestMessage::User(ChatCompletionRequestUserMessage { - content: ChatCompletionRequestUserMessageContent::Text( - "Can I ask a question?".to_string(), - ), - name: None, - }), - ] - } - - fn sample_tools() -> Vec { - vec![ChatCompletionTool { - r#type: dynamo_protocols::types::ChatCompletionToolType::Function, - function: FunctionObject { - name: "get_weather".to_string(), - description: None, - parameters: Some(serde_json::json!({ - "type": "object", - "properties": { - "location": { - "type": "string" - } - } - })), - strict: None, - }, - }] - } fn http_error_from_engine(code: u16) -> Result<(), anyhow::Error> { Err(HttpError { @@ -3997,384 +3620,6 @@ mod tests { ); } - #[test] - fn test_apply_continue_final_message_trims_rendered_prompt() { - let rendered_prompt = "USER: Hi\nASSISTANT: Sure.<|assistant|>".to_string(); - let messages = vec![ - ChatCompletionRequestMessage::User(ChatCompletionRequestUserMessage { - content: ChatCompletionRequestUserMessageContent::Text("Hi".to_string()), - name: None, - }), - ChatCompletionRequestMessage::Assistant(ChatCompletionRequestAssistantMessage { - content: Some(ChatCompletionRequestAssistantMessageContent::Text( - "Sure.".to_string(), - )), - ..Default::default() - }), - ]; - - let trimmed = apply_continue_final_message(rendered_prompt, &messages).unwrap(); - assert_eq!(trimmed, "USER: Hi\nASSISTANT: Sure."); - } - - #[tokio::test] - async fn test_tokenize_completion_route_matches_tokenizer() { - let (state, card) = make_tokenize_state("test-model"); - let tokenizer = card.tokenizer().unwrap(); - let prompt = "This is a completion tokenize test."; - - for add_special_tokens in [false, true] { - let response = tokenize( - State(state.clone()), - Json(TokenizeRequest::Completion(TokenizeCompletionRequest { - model: Some("test-model".to_string()), - prompt: prompt.to_string(), - add_special_tokens, - return_token_strs: false, - })), - ) - .await - .unwrap(); - let body: TokenizeResponse = response_json(response).await; - let expected = tokenizer - .encode_with_special_tokens(prompt, add_special_tokens) - .unwrap(); - - assert_eq!(body.tokens, expected.token_ids()); - assert_eq!(body.count, expected.token_ids().len()); - assert_eq!(body.max_model_len, card.context_length); - assert!(body.token_strs.is_none()); - } - } - - #[tokio::test] - async fn test_tokenize_chat_route_matches_rendered_prompt() { - let (state, card) = make_tokenize_state("test-model"); - let tokenizer = card.tokenizer().unwrap(); - let messages = sample_chat_messages(); - - for add_generation_prompt in [false, true] { - for add_special_tokens in [false, true] { - let request = TokenizeChatRequest { - model: Some("test-model".to_string()), - messages: messages.clone(), - add_generation_prompt, - return_token_strs: false, - continue_final_message: false, - add_special_tokens, - chat_template: None, - chat_template_kwargs: None, - media_io_kwargs: None, - mm_processor_kwargs: None, - tools: None, - }; - let prompt = - render_tokenize_chat_prompt(&card, "test-model".to_string(), &request).unwrap(); - let expected = tokenizer - .encode_with_special_tokens(&prompt, add_special_tokens) - .unwrap(); - - let response = tokenize( - State(state.clone()), - Json(TokenizeRequest::Chat(TokenizeChatRequest { - model: Some("test-model".to_string()), - messages: messages.clone(), - add_generation_prompt, - return_token_strs: false, - continue_final_message: false, - add_special_tokens, - chat_template: None, - chat_template_kwargs: None, - media_io_kwargs: None, - mm_processor_kwargs: None, - tools: None, - })), - ) - .await - .unwrap(); - let body: TokenizeResponse = response_json(response).await; - assert_eq!(body.tokens, expected.token_ids()); - assert_eq!(body.count, expected.token_ids().len()); - } - } - } - - #[tokio::test] - async fn test_tokenize_chat_route_with_tools_and_continue_final_message() { - let (state, card) = make_tokenize_state("test-model"); - let tokenizer = card.tokenizer().unwrap(); - let mut messages = sample_chat_messages(); - messages.push(ChatCompletionRequestMessage::Assistant( - ChatCompletionRequestAssistantMessage { - content: Some(ChatCompletionRequestAssistantMessageContent::Text( - "Sure,".to_string(), - )), - ..Default::default() - }, - )); - let tools = sample_tools(); - - let request = TokenizeChatRequest { - model: Some("test-model".to_string()), - messages: messages.clone(), - add_generation_prompt: false, - return_token_strs: false, - continue_final_message: true, - add_special_tokens: true, - chat_template: None, - chat_template_kwargs: None, - media_io_kwargs: None, - mm_processor_kwargs: None, - tools: Some(tools.clone()), - }; - let prompt = - render_tokenize_chat_prompt(&card, "test-model".to_string(), &request).unwrap(); - let expected = tokenizer.encode_with_special_tokens(&prompt, true).unwrap(); - - let response = tokenize( - State(state), - Json(TokenizeRequest::Chat(TokenizeChatRequest { - model: Some("test-model".to_string()), - messages: messages.clone(), - add_generation_prompt: false, - return_token_strs: false, - continue_final_message: true, - add_special_tokens: true, - chat_template: None, - chat_template_kwargs: None, - media_io_kwargs: None, - mm_processor_kwargs: None, - tools: Some(tools), - })), - ) - .await - .unwrap(); - let body: TokenizeResponse = response_json(response).await; - assert_eq!(body.tokens, expected.token_ids()); - } - - #[tokio::test] - async fn test_tokenize_route_returns_token_strings() { - let (state, card) = make_tokenize_state("test-model"); - let tokenizer = card.tokenizer().unwrap(); - let prompt = "Return token strings please."; - - let response = tokenize( - State(state), - Json(TokenizeRequest::Completion(TokenizeCompletionRequest { - model: Some("test-model".to_string()), - prompt: prompt.to_string(), - add_special_tokens: true, - return_token_strs: true, - })), - ) - .await - .unwrap(); - let body: TokenizeResponse = response_json(response).await; - let expected = tokenizer.encode_with_special_tokens(prompt, true).unwrap(); - let expected_token_strs = tokenizer - .convert_ids_to_tokens(expected.token_ids()) - .unwrap(); - - assert_eq!(body.tokens, expected.token_ids()); - assert_eq!(body.token_strs, Some(expected_token_strs)); - } - - #[tokio::test] - async fn test_tokenize_chat_rejects_incompatible_flags() { - let (state, _) = make_tokenize_state("test-model"); - let messages = sample_chat_messages(); - - let error = tokenize( - State(state), - Json(TokenizeRequest::Chat(TokenizeChatRequest { - model: Some("test-model".to_string()), - messages, - add_generation_prompt: true, - return_token_strs: false, - continue_final_message: true, - add_special_tokens: false, - chat_template: None, - chat_template_kwargs: None, - media_io_kwargs: None, - mm_processor_kwargs: None, - tools: None, - })), - ) - .await - .unwrap_err(); - let response = error.into_response(); - let body: ErrorMessage = response_json(response).await; - assert!(body.message.contains( - "Cannot set both `continue_final_message` and `add_generation_prompt` to True." - )); - } - - #[tokio::test] - async fn test_detokenize_route_round_trips_prompt() { - let (state, card) = make_tokenize_state_with_path("test-model", DETOKENIZE_MODEL_PATH); - let tokenizer = card.tokenizer().unwrap(); - let prompt = "This is a detokenize test prompt."; - let tokens = tokenizer - .encode_with_special_tokens(prompt, false) - .unwrap() - .token_ids() - .to_vec(); - - let response = detokenize( - State(state), - Json(DetokenizeRequest { - model: Some("test-model".to_string()), - tokens, - }), - ) - .await - .unwrap(); - let body: DetokenizeResponse = response_json(response).await; - assert_eq!(body.prompt, prompt); - } - - #[tokio::test] - async fn test_tokenize_route_rejects_models_without_card_metadata() { - let state = make_tokenize_state_without_card("test-model"); - let error = tokenize( - State(state), - Json(TokenizeRequest::Completion(TokenizeCompletionRequest { - model: Some("test-model".to_string()), - prompt: "hello".to_string(), - add_special_tokens: true, - return_token_strs: false, - })), - ) - .await - .unwrap_err(); - - let response = error.into_response(); - let body: ErrorMessage = response_json(response).await; - assert!( - body.message - .contains("Tokenizer metadata is not available for model 'test-model'") - ); - } - - #[tokio::test] - async fn test_detokenize_route_rejects_models_without_card_metadata() { - let state = make_tokenize_state_without_card("test-model"); - let error = detokenize( - State(state), - Json(DetokenizeRequest { - model: Some("test-model".to_string()), - tokens: vec![1, 2, 3], - }), - ) - .await - .unwrap_err(); - - let response = error.into_response(); - let body: ErrorMessage = response_json(response).await; - assert!( - body.message - .contains("Tokenizer metadata is not available for model 'test-model'") - ); - } - - #[tokio::test] - async fn test_tokenize_route_defaults_model_when_only_one_is_served() { - let (state, card) = make_tokenize_state("test-model"); - let tokenizer = card.tokenizer().unwrap(); - let prompt = "Single served model default."; - let expected = tokenizer.encode_with_special_tokens(prompt, true).unwrap(); - - let response = tokenize( - State(state), - Json(TokenizeRequest::Completion(TokenizeCompletionRequest { - model: None, - prompt: prompt.to_string(), - add_special_tokens: true, - return_token_strs: false, - })), - ) - .await - .unwrap(); - let body: TokenizeResponse = response_json(response).await; - assert_eq!(body.tokens, expected.token_ids()); - } - - #[tokio::test] - async fn test_tokenize_chat_route_supports_request_chat_template_override() { - let (state, card) = make_tokenize_state("test-model"); - let tokenizer = card.tokenizer().unwrap(); - let messages = sample_chat_messages(); - let custom_template = concat!( - "{% for message in messages %}", - "[[{{ message['role'] }}]] {{ message['content'] }}\n", - "{% endfor %}", - "{% if add_generation_prompt %}[[assistant]] {% endif %}" - ); - - let request = TokenizeChatRequest { - model: Some("test-model".to_string()), - messages: messages.clone(), - add_generation_prompt: true, - return_token_strs: false, - continue_final_message: false, - add_special_tokens: false, - chat_template: Some(custom_template.to_string()), - chat_template_kwargs: Some(StdHashMap::from([( - "unused_value".to_string(), - serde_json::Value::String("ignored".to_string()), - )])), - media_io_kwargs: None, - mm_processor_kwargs: Some(StdHashMap::from([( - "image".to_string(), - serde_json::json!({"size": "ignored"}), - )])), - tools: None, - }; - let prompt = - render_tokenize_chat_prompt(&card, "test-model".to_string(), &request).unwrap(); - let expected = tokenizer - .encode_with_special_tokens(&prompt, false) - .unwrap(); - - let response = tokenize( - State(state), - Json(TokenizeRequest::Chat(TokenizeChatRequest { - model: Some("test-model".to_string()), - messages, - add_generation_prompt: true, - return_token_strs: false, - continue_final_message: false, - add_special_tokens: false, - chat_template: Some(custom_template.to_string()), - chat_template_kwargs: Some(StdHashMap::from([( - "unused_value".to_string(), - serde_json::Value::String("ignored".to_string()), - )])), - media_io_kwargs: None, - mm_processor_kwargs: Some(StdHashMap::from([( - "image".to_string(), - serde_json::json!({"size": "ignored"}), - )])), - tools: None, - })), - ) - .await - .unwrap(); - let body: TokenizeResponse = response_json(response).await; - assert_eq!(body.tokens, expected.token_ids()); - } - - #[test] - fn test_tokenization_router_registers_root_paths() { - let (state, _) = make_tokenize_state("test-model"); - let (docs, _) = tokenization_router(state); - let docs = docs.iter().map(|doc| doc.to_string()).collect::>(); - - assert!(docs.contains(&"POST /tokenize".to_string())); - assert!(docs.contains(&"POST /detokenize".to_string())); - } - // ── streaming dispatch tests ────────────────────────────────────── use std::collections::{HashMap, HashSet}; diff --git a/lib/llm/src/http/service/service_v2.rs b/lib/llm/src/http/service/service_v2.rs index cb78f30686ec..1eeaefdff912 100644 --- a/lib/llm/src/http/service/service_v2.rs +++ b/lib/llm/src/http/service/service_v2.rs @@ -532,7 +532,6 @@ impl HttpServiceConfigBuilder { } else { super::openai::list_models_router(state.clone(), var(HTTP_SVC_MODELS_PATH_ENV).ok()) }, - super::openai::tokenization_router(state.clone()), super::health::health_check_router(state.clone(), var(HTTP_SVC_HEALTH_PATH_ENV).ok()), super::health::live_check_router(state.clone(), var(HTTP_SVC_LIVE_PATH_ENV).ok()), super::busy_threshold::busy_threshold_router(state.clone(), None), diff --git a/lib/llm/src/model_card.rs b/lib/llm/src/model_card.rs index 2d17e55f665f..9a231569c685 100644 --- a/lib/llm/src/model_card.rs +++ b/lib/llm/src/model_card.rs @@ -689,28 +689,6 @@ impl ModelDeploymentCard { Ok(()) } - /// Return the local directory that contains the model config files, if any - /// file has already been resolved to a local path. - pub(crate) fn local_file_dir(&self) -> Option<&Path> { - if let Some(TokenizerKind::HfTokenizerJson(cf) | TokenizerKind::TikTokenModel(cf)) = - &self.tokenizer - && let Some(dir) = cf.path().and_then(|p| p.parent()) - { - return Some(dir); - } - if let Some(ModelInfoType::HfConfigJson(cf)) = &self.model_info - && let Some(dir) = cf.path().and_then(|p| p.parent()) - { - return Some(dir); - } - if let Some(PromptFormatterArtifact::HfTokenizerConfigJson(cf)) = &self.prompt_formatter - && let Some(dir) = cf.path().and_then(|p| p.parent()) - { - return Some(dir); - } - None - } - /// Are all the files we need (tokenizer.json, etc) available locally? fn has_local_files(&self) -> bool { let has_model_info = self @@ -747,7 +725,7 @@ impl ModelDeploymentCard { } /// Update the directory for files like tokenizer.json be in here. - pub(crate) fn update_dir(&mut self, dir: &Path) { + fn update_dir(&mut self, dir: &Path) { if let Some(model_info) = self.model_info.as_mut() { model_info.update_dir(dir); } diff --git a/lib/llm/src/preprocessor/prompt/template.rs b/lib/llm/src/preprocessor/prompt/template.rs index 7d148bb8cceb..b9b86c599f2c 100644 --- a/lib/llm/src/preprocessor/prompt/template.rs +++ b/lib/llm/src/preprocessor/prompt/template.rs @@ -4,7 +4,6 @@ use std::{collections::HashSet, sync::Arc}; use anyhow::{Context, Ok, Result}; -use either::Either; use minijinja::Environment; use crate::model_card::{ModelDeploymentCard, PromptContextMixin, PromptFormatterArtifact}; @@ -20,13 +19,6 @@ use tokcfg::ChatTemplateValue; impl PromptFormatter { pub fn from_mdc(mdc: &ModelDeploymentCard) -> Result { - Self::from_mdc_with_chat_template(mdc, None) - } - - pub fn from_mdc_with_chat_template( - mdc: &ModelDeploymentCard, - chat_template_override: Option<&str>, - ) -> Result { // Special handling for DeepSeek models whose HF repos don't ship a Jinja chat_template. // // Prefer the authoritative `model_type` from config.json — it's set by @@ -37,10 +29,6 @@ impl PromptFormatter { // An empty `model_type` string (rare but legal in the JSON) carries // no signal — normalize it to `None` so the display-name fallback // still runs instead of being silently suppressed. - // - // Skip native-formatter routing when an explicit chat_template override - // is provided — the caller is asking for a specific template, not the - // built-in one. let model_type_lower = mdc .model_info .as_ref() @@ -49,25 +37,21 @@ impl PromptFormatter { .filter(|s| !s.is_empty()); let display_name_lower = mdc.display_name.to_lowercase(); - if chat_template_override.is_none() { - if is_deepseek_v4(&model_type_lower, &display_name_lower) { - tracing::info!( - model_type = ?model_type_lower, - display_name = %mdc.display_name, - "Detected DeepSeek V4 model, using native Rust formatter", - ); - return Ok(Self::OAI(Arc::new( - super::deepseek_v4::DeepSeekV4Formatter::new_thinking(), - ))); - } - if is_deepseek_v3_2_non_exp(&model_type_lower, &display_name_lower) { - tracing::info!( - "Detected DeepSeek V3.2 model (non-Exp), using native Rust formatter" - ); - return Ok(Self::OAI(Arc::new( - super::deepseek_v32::DeepSeekV32Formatter::new_thinking(), - ))); - } + if is_deepseek_v4(&model_type_lower, &display_name_lower) { + tracing::info!( + model_type = ?model_type_lower, + display_name = %mdc.display_name, + "Detected DeepSeek V4 model, using native Rust formatter", + ); + return Ok(Self::OAI(Arc::new( + super::deepseek_v4::DeepSeekV4Formatter::new_thinking(), + ))); + } + if is_deepseek_v3_2_non_exp(&model_type_lower, &display_name_lower) { + tracing::info!("Detected DeepSeek V3.2 model (non-Exp), using native Rust formatter"); + return Ok(Self::OAI(Arc::new( + super::deepseek_v32::DeepSeekV32Formatter::new_thinking(), + ))); } match mdc @@ -93,68 +77,57 @@ impl PromptFormatter { crate::log_json_err(&file.display().to_string(), &contents, err) })?; - if let Some(chat_template) = chat_template_override { - config.chat_template = - Some(ChatTemplateValue(Either::Left(chat_template.to_string()))); - } - // Some HF model (i.e. meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8) // stores the chat template in a separate file, we check if the file exists and // put the chat template into config as normalization. // This may also be a custom template provided via CLI flag. - if chat_template_override.is_none() { - match mdc.chat_template_file.as_ref() { - Some(PromptFormatterArtifact::HfChatTemplateJinja { - file: checked_file, - .. - }) => { - let Some(path) = checked_file.path() else { - anyhow::bail!( - "HfChatTemplateJinja for {} is a URL, cannot load", - mdc.display_name - ); - }; - let chat_template = - std::fs::read_to_string(path).with_context(|| { - format!("fs:read_to_string '{}'", path.display()) - })?; - config.chat_template = - Some(ChatTemplateValue(either::Left(chat_template))); - } - Some(PromptFormatterArtifact::HfChatTemplateJson { - file: checked_file, - .. - }) => { - let Some(path) = checked_file.path() else { - anyhow::bail!( - "HfChatTemplateJson for {} is a URL, cannot load", - mdc.display_name - ); - }; - let raw = std::fs::read_to_string(path).with_context(|| { - format!("fs:read_to_string '{}'", path.display()) + match mdc.chat_template_file.as_ref() { + Some(PromptFormatterArtifact::HfChatTemplateJinja { + file: checked_file, + .. + }) => { + let Some(path) = checked_file.path() else { + anyhow::bail!( + "HfChatTemplateJinja for {} is a URL, cannot load", + mdc.display_name + ); + }; + let chat_template = std::fs::read_to_string(path) + .with_context(|| format!("fs:read_to_string '{}'", path.display()))?; + config.chat_template = Some(ChatTemplateValue(either::Left(chat_template))); + } + Some(PromptFormatterArtifact::HfChatTemplateJson { + file: checked_file, + .. + }) => { + let Some(path) = checked_file.path() else { + anyhow::bail!( + "HfChatTemplateJson for {} is a URL, cannot load", + mdc.display_name + ); + }; + let raw = std::fs::read_to_string(path) + .with_context(|| format!("fs:read_to_string '{}'", path.display()))?; + let wrapper: serde_json::Value = + serde_json::from_str(&raw).with_context(|| { + format!("Failed to parse '{}' as JSON", path.display()) })?; - let wrapper: serde_json::Value = serde_json::from_str(&raw) - .with_context(|| { - format!("Failed to parse '{}' as JSON", path.display()) - })?; - let field = wrapper.get("chat_template").ok_or_else(|| { - anyhow::anyhow!( - "'{}' does not contain a 'chat_template' field", + let field = wrapper.get("chat_template").ok_or_else(|| { + anyhow::anyhow!( + "'{}' does not contain a 'chat_template' field", + path.display() + ) + })?; + let value = serde_json::from_value::(field.clone()) + .with_context(|| { + format!( + "Failed to deserialize 'chat_template' in '{}'", path.display() ) })?; - let value = serde_json::from_value::(field.clone()) - .with_context(|| { - format!( - "Failed to deserialize 'chat_template' in '{}'", - path.display() - ) - })?; - config.chat_template = Some(value); - } - _ => {} + config.chat_template = Some(value); } + _ => {} } Self::from_parts( config, diff --git a/lib/llm/src/protocols/common.rs b/lib/llm/src/protocols/common.rs index 9ba9a8915e80..031d34a7f3bc 100644 --- a/lib/llm/src/protocols/common.rs +++ b/lib/llm/src/protocols/common.rs @@ -517,7 +517,7 @@ pub struct OutputOptions { pub formatted_prompt: Option, /// When true, logprob token fields are returned as "token_id:" - /// instead of the decoded text. vLLM-specific extension for NeMo-RL. + /// instead of decoded text. pub return_tokens_as_token_ids: Option, } diff --git a/lib/llm/src/protocols/openai.rs b/lib/llm/src/protocols/openai.rs index 0ddf2049996e..225ef28c61a8 100644 --- a/lib/llm/src/protocols/openai.rs +++ b/lib/llm/src/protocols/openai.rs @@ -20,7 +20,6 @@ pub mod images; pub mod models; pub mod nvext; pub mod responses; -pub mod tokenization; pub mod tools; pub mod validate; pub mod videos; @@ -237,6 +236,7 @@ pub(crate) fn convert_backend_top_logprobs( selected_token: &str, selected_token_id: TokenIdType, selected_logprob: f32, + return_tokens_as_token_ids: bool, ) -> Vec { let mut found_selected = false; let mut result: Vec = top_lps @@ -254,10 +254,15 @@ pub(crate) fn convert_backend_top_logprobs( .collect(); if !found_selected { + let token = if return_tokens_as_token_ids { + format!("token_id:{}", selected_token_id) + } else { + selected_token.to_string() + }; result.push(dynamo_protocols::types::TopLogprobs { - token: selected_token.to_string(), + bytes: token_to_utf8_bytes(&token), + token, logprob: selected_logprob, - bytes: token_to_utf8_bytes(selected_token), }); } result diff --git a/lib/llm/src/protocols/openai/chat_completions.rs b/lib/llm/src/protocols/openai/chat_completions.rs index f41324d30806..4b0a87160958 100644 --- a/lib/llm/src/protocols/openai/chat_completions.rs +++ b/lib/llm/src/protocols/openai/chat_completions.rs @@ -60,8 +60,7 @@ pub struct NvCreateChatCompletionRequest { pub media_io_kwargs: Option, /// When true, logprob token fields are returned as "token_id:" instead - /// of the decoded text. This is a vLLM-specific extension used by NeMo-RL - /// to extract per-token IDs for RL training. + /// of decoded text. #[serde(default, skip_serializing_if = "Option::is_none")] pub return_tokens_as_token_ids: Option, diff --git a/lib/llm/src/protocols/openai/chat_completions/delta.rs b/lib/llm/src/protocols/openai/chat_completions/delta.rs index 3f4700c4f768..e6e76d8ad7a5 100644 --- a/lib/llm/src/protocols/openai/chat_completions/delta.rs +++ b/lib/llm/src/protocols/openai/chat_completions/delta.rs @@ -212,7 +212,8 @@ impl DeltaGenerator { } else { t.clone() }; - let converted = convert_backend_top_logprobs(&top_lps, t, *tid, lp); + let converted = + convert_backend_top_logprobs(&top_lps, t, *tid, lp, return_as_ids); dynamo_protocols::types::ChatCompletionTokenLogprob { token: token_str.clone(), logprob: lp, diff --git a/lib/llm/src/protocols/openai/completions/delta.rs b/lib/llm/src/protocols/openai/completions/delta.rs index e45470c54f40..7cec35531fb2 100644 --- a/lib/llm/src/protocols/openai/completions/delta.rs +++ b/lib/llm/src/protocols/openai/completions/delta.rs @@ -167,7 +167,8 @@ impl DeltaGenerator { .zip(tok_lps.iter()) .zip(top_logprobs.iter()) .map(|(((t, tid), lp), top_lps)| { - let converted = convert_backend_top_logprobs(top_lps, t, *tid, *lp); + let converted = + convert_backend_top_logprobs(top_lps, t, *tid, *lp, return_as_ids); serde_json::to_value(converted).unwrap() }) .collect() @@ -512,6 +513,61 @@ mod tests { ); } + #[test] + fn test_logprobs_zero_emits_chosen_token_logprob() { + let mut request = create_test_request(); + request.inner.logprobs = Some(0); + let mut generator = request.response_generator("req-logprobs-zero".to_string()); + let mut output = final_backend_output(); + output.log_probs = Some(vec![-0.5]); + + let response = generator + .choice_from_postprocessor(output) + .expect("choice generation"); + let logprobs = response.inner.choices[0] + .logprobs + .as_ref() + .expect("logprobs"); + + assert_eq!(logprobs.tokens, vec!["hello"]); + assert_eq!(logprobs.token_logprobs, vec![Some(-0.5)]); + assert!(logprobs.top_logprobs.is_empty()); + } + + #[test] + fn test_return_token_ids_formats_selected_top_logprob_fallback() { + let mut request = create_test_request(); + request.inner.logprobs = Some(1); + request.return_tokens_as_token_ids = Some(true); + let generator = request.response_generator("req-token-id-logprobs".to_string()); + + let logprobs = generator + .create_logprobs( + vec![Some("hello".to_string())], + vec![123], + Some(vec![-0.5]), + Some(vec![vec![common::llm_backend::TopLogprob { + rank: 1, + token_id: 999, + token: Some("other".to_string()), + logprob: -1.0, + bytes: None, + }]]), + ) + .expect("logprobs"); + + assert_eq!(logprobs.tokens, vec!["token_id:123"]); + let top_logprobs = logprobs.top_logprobs[0] + .as_array() + .expect("top_logprobs array"); + let selected = top_logprobs + .iter() + .find(|item| item["token_id"] == 123) + .expect("selected token fallback"); + assert_eq!(selected["token"], "token_id:123"); + assert_eq!(selected["bytes"], serde_json::json!(b"token_id:123")); + } + #[test] fn test_timing_extra_field_emits_timing_on_final_chunk() { use crate::protocols::openai::nvext::NvExt; diff --git a/lib/llm/src/protocols/openai/tokenization.rs b/lib/llm/src/protocols/openai/tokenization.rs deleted file mode 100644 index 95559684ad89..000000000000 --- a/lib/llm/src/protocols/openai/tokenization.rs +++ /dev/null @@ -1,124 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -use std::collections::HashMap; - -use serde::{Deserialize, Serialize}; - -use crate::preprocessor::media::MediaDecoder; -use crate::types::TokenIdType; - -fn default_true() -> bool { - true -} - -fn default_false() -> bool { - false -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct TokenizeCompletionRequest { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub model: Option, - pub prompt: String, - #[serde(default = "default_true")] - pub add_special_tokens: bool, - #[serde(default = "default_false")] - pub return_token_strs: bool, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct TokenizeChatRequest { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub model: Option, - pub messages: Vec, - #[serde(default = "default_true")] - pub add_generation_prompt: bool, - #[serde(default = "default_false")] - pub return_token_strs: bool, - #[serde(default = "default_false")] - pub continue_final_message: bool, - #[serde(default = "default_false")] - pub add_special_tokens: bool, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub chat_template: Option, - #[serde( - default, - skip_serializing_if = "Option::is_none", - alias = "chat_template_args" - )] - pub chat_template_kwargs: Option>, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub media_io_kwargs: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub mm_processor_kwargs: Option>, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub tools: Option>, -} - -impl TokenizeChatRequest { - pub fn validate(&self) -> Result<(), String> { - if self.continue_final_message && self.add_generation_prompt { - return Err( - "Cannot set both `continue_final_message` and `add_generation_prompt` to True." - .to_string(), - ); - } - - Ok(()) - } - - pub fn merged_chat_template_kwargs(&self) -> HashMap { - let mut kwargs = self.chat_template_kwargs.clone().unwrap_or_default(); - kwargs.insert( - "add_generation_prompt".to_string(), - serde_json::Value::Bool(self.add_generation_prompt), - ); - kwargs.insert( - "continue_final_message".to_string(), - serde_json::Value::Bool(self.continue_final_message), - ); - kwargs - } -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(untagged)] -#[allow(clippy::large_enum_variant)] -pub enum TokenizeRequest { - Completion(TokenizeCompletionRequest), - Chat(TokenizeChatRequest), -} - -impl TokenizeRequest { - pub fn model(&self) -> Option<&str> { - match self { - Self::Completion(request) => request.model.as_deref(), - Self::Chat(request) => request.model.as_deref(), - } - } -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -pub struct TokenizeResponse { - pub count: usize, - pub max_model_len: u32, - pub tokens: Vec, - #[serde(skip_serializing_if = "Option::is_none")] - pub token_strs: Option>, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct DetokenizeRequest { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub model: Option, - pub tokens: Vec, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -pub struct DetokenizeResponse { - pub prompt: String, -} diff --git a/lib/tokenizers/src/fastokens.rs b/lib/tokenizers/src/fastokens.rs index 4295e69fcb91..93e855cc5c58 100644 --- a/lib/tokenizers/src/fastokens.rs +++ b/lib/tokenizers/src/fastokens.rs @@ -39,28 +39,16 @@ impl FastTokenizer { impl Encoder for FastTokenizer { fn encode(&self, input: &str) -> Result { - self.encode_with_special_tokens(input, false) - } - - fn encode_batch(&self, inputs: &[&str]) -> Result> { - inputs.par_iter().map(|input| self.encode(input)).collect() - } - - fn encode_with_special_tokens( - &self, - input: &str, - add_special_tokens: bool, - ) -> Result { - if add_special_tokens { - return self.hf_decoder.encode_with_special_tokens(input, true); - } - let ids = self .fast_encoder .encode(input) .map_err(|e| Error::msg(format!("Fastokens encode error: {e}")))?; Ok(Encoding::Sp(ids)) } + + fn encode_batch(&self, inputs: &[&str]) -> Result> { + inputs.par_iter().map(|input| self.encode(input)).collect() + } } impl Decoder for FastTokenizer { @@ -69,11 +57,7 @@ impl Decoder for FastTokenizer { } } -impl Tokenizer for FastTokenizer { - fn convert_ids_to_tokens(&self, token_ids: &[TokenIdType]) -> Result> { - self.hf_decoder.convert_ids_to_tokens(token_ids) - } -} +impl Tokenizer for FastTokenizer {} #[cfg(test)] mod tests { diff --git a/lib/tokenizers/src/hf.rs b/lib/tokenizers/src/hf.rs index f0f3e3764a46..080a775719fe 100644 --- a/lib/tokenizers/src/hf.rs +++ b/lib/tokenizers/src/hf.rs @@ -27,7 +27,13 @@ impl HuggingFaceTokenizer { impl Encoder for HuggingFaceTokenizer { fn encode(&self, input: &str) -> Result { - self.encode_with_special_tokens(input, false) + // This self.tokenizer is the library + let encoding = self + .tokenizer + .encode(input, false) + .map_err(|err| Error::msg(format!("Error tokenizing input: {err}")))?; + + Ok(Encoding::Hf(Box::new(encoding))) } fn encode_batch(&self, inputs: &[&str]) -> Result> { @@ -43,20 +49,6 @@ impl Encoder for HuggingFaceTokenizer { Ok(encodings) } - - fn encode_with_special_tokens( - &self, - input: &str, - add_special_tokens: bool, - ) -> Result { - // This self.tokenizer is the library - let encoding = self - .tokenizer - .encode(input, add_special_tokens) - .map_err(|err| Error::msg(format!("Error tokenizing input: {err}")))?; - - Ok(Encoding::Hf(Box::new(encoding))) - } } impl Decoder for HuggingFaceTokenizer { @@ -71,14 +63,7 @@ impl Decoder for HuggingFaceTokenizer { } } -impl Tokenizer for HuggingFaceTokenizer { - fn convert_ids_to_tokens(&self, token_ids: &[TokenIdType]) -> Result> { - Ok(token_ids - .iter() - .map(|&id| self.tokenizer.id_to_token(id).unwrap_or_default()) - .collect()) - } -} +impl Tokenizer for HuggingFaceTokenizer {} impl From for HuggingFaceTokenizer { fn from(tokenizer: HfTokenizer) -> Self { diff --git a/lib/tokenizers/src/lib.rs b/lib/tokenizers/src/lib.rs index f97a235cff6c..95494b4f73f0 100644 --- a/lib/tokenizers/src/lib.rs +++ b/lib/tokenizers/src/lib.rs @@ -63,14 +63,6 @@ pub mod traits { pub trait Encoder: Send + Sync { fn encode(&self, input: &str) -> Result; fn encode_batch(&self, inputs: &[&str]) -> Result>; - - fn encode_with_special_tokens( - &self, - input: &str, - _add_special_tokens: bool, - ) -> Result { - self.encode(input) - } } /// Result of decoding token IDs to text. @@ -136,15 +128,8 @@ pub mod traits { } pub trait Tokenizer: Encoder + Decoder { - fn convert_ids_to_tokens(&self, token_ids: &[TokenIdType]) -> Result> { - token_ids - .iter() - .map(|id| { - self.decode(std::slice::from_ref(id), false) - .map(String::from) - }) - .collect() - } + // fn get_vocab_size(&self) -> usize; + // fn make_unique_clone(&self) -> Box; } } @@ -239,18 +224,6 @@ impl Tokenizer { Ok(Tokenizer(create_tokenizer_from_file(file_path)?)) } - pub fn encode_with_special_tokens( - &self, - input: &str, - add_special_tokens: bool, - ) -> Result { - self.0.encode_with_special_tokens(input, add_special_tokens) - } - - pub fn convert_ids_to_tokens(&self, token_ids: &[TokenIdType]) -> Result> { - self.0.convert_ids_to_tokens(token_ids) - } - /// Create a stateful sequence object for decoding token_ids into text pub fn decode_stream( &self, diff --git a/lib/tokenizers/src/tiktoken.rs b/lib/tokenizers/src/tiktoken.rs index cc049c13001d..7082acb0f6a6 100644 --- a/lib/tokenizers/src/tiktoken.rs +++ b/lib/tokenizers/src/tiktoken.rs @@ -24,8 +24,6 @@ const KIMI_PATTERN: &str = r#"[\p{Han}]+|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p pub struct TikTokenTokenizer { bpe: CoreBPE, special_token_ids: HashSet, - decoder_tokens: FxHashMap>, - special_tokens_decoder: FxHashMap>, } impl TikTokenTokenizer { @@ -41,14 +39,6 @@ impl TikTokenTokenizer { special_tokens: FxHashMap, ) -> Result { let encoder = parse_tiktoken_file(path)?; - let decoder_tokens: FxHashMap> = encoder - .iter() - .map(|(bytes, &id)| (id, bytes.clone())) - .collect(); - let special_tokens_decoder: FxHashMap> = special_tokens - .iter() - .map(|(token, &id)| (id, token.as_bytes().to_vec())) - .collect(); let special_token_ids: HashSet = special_tokens.values().copied().collect(); let bpe = CoreBPE::new(encoder, special_tokens, pattern) @@ -57,8 +47,6 @@ impl TikTokenTokenizer { Ok(Self { bpe, special_token_ids, - decoder_tokens, - special_tokens_decoder, }) } @@ -74,17 +62,9 @@ impl TikTokenTokenizer { let pattern = detect_bpe_pattern(directory)?; let encoder = parse_tiktoken_file(path)?; - let decoder_tokens: FxHashMap> = encoder - .iter() - .map(|(bytes, &id)| (id, bytes.clone())) - .collect(); // Use max rank + 1 (not len) to avoid ID collisions with sparse/non-contiguous ranks let num_base_tokens = encoder.values().max().map_or(0, |&m| m + 1) as usize; let special_tokens = load_special_tokens(directory, num_base_tokens)?; - let special_tokens_decoder: FxHashMap> = special_tokens - .iter() - .map(|(token, &id)| (id, token.as_bytes().to_vec())) - .collect(); let special_token_ids: HashSet = special_tokens.values().copied().collect(); let bpe = CoreBPE::new(encoder, special_tokens, pattern) @@ -93,33 +73,19 @@ impl TikTokenTokenizer { Ok(Self { bpe, special_token_ids, - decoder_tokens, - special_tokens_decoder, }) } } impl Encoder for TikTokenTokenizer { fn encode(&self, input: &str) -> Result { - self.encode_with_special_tokens(input, true) + let token_ids: Vec = self.bpe.encode_with_special_tokens(input); + Ok(Encoding::Sp(token_ids)) } fn encode_batch(&self, inputs: &[&str]) -> Result> { inputs.par_iter().map(|input| self.encode(input)).collect() } - - fn encode_with_special_tokens( - &self, - input: &str, - add_special_tokens: bool, - ) -> Result { - let token_ids: Vec = if add_special_tokens { - self.bpe.encode_with_special_tokens(input) - } else { - self.bpe.encode_ordinary(input) - }; - Ok(Encoding::Sp(token_ids)) - } } impl Decoder for TikTokenTokenizer { @@ -153,20 +119,7 @@ impl Decoder for TikTokenTokenizer { } } -impl Tokenizer for TikTokenTokenizer { - fn convert_ids_to_tokens(&self, token_ids: &[TokenIdType]) -> Result> { - Ok(token_ids - .iter() - .map(|id| { - self.decoder_tokens - .get(id) - .or_else(|| self.special_tokens_decoder.get(id)) - .map(|bytes| String::from_utf8_lossy(bytes).into_owned()) - .unwrap_or_default() - }) - .collect()) - } -} +impl Tokenizer for TikTokenTokenizer {} /// Parse a tiktoken model file (base64-encoded token + rank per line). fn parse_tiktoken_file(path: &str) -> Result, u32>> { From f32b200023d44c5c238e293882909348c2fcf733 Mon Sep 17 00:00:00 2001 From: William Arnold Date: Tue, 5 May 2026 11:22:10 -0700 Subject: [PATCH 13/20] fix(openai): format top logprobs as token ids --- lib/llm/src/protocols/openai.rs | 12 ++++++++++-- lib/llm/src/protocols/openai/completions/delta.rs | 7 ++++++- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/lib/llm/src/protocols/openai.rs b/lib/llm/src/protocols/openai.rs index 225ef28c61a8..455ef04da87e 100644 --- a/lib/llm/src/protocols/openai.rs +++ b/lib/llm/src/protocols/openai.rs @@ -242,9 +242,17 @@ pub(crate) fn convert_backend_top_logprobs( let mut result: Vec = top_lps .iter() .map(|top_lp| { - let tok = top_lp.token.clone().unwrap_or_default(); + let tok = if return_tokens_as_token_ids { + format!("token_id:{}", top_lp.token_id) + } else { + top_lp.token.clone().unwrap_or_default() + }; found_selected = found_selected || top_lp.token_id == selected_token_id; - let bytes = top_lp.bytes.clone().or_else(|| token_to_utf8_bytes(&tok)); + let bytes = if return_tokens_as_token_ids { + token_to_utf8_bytes(&tok) + } else { + top_lp.bytes.clone().or_else(|| token_to_utf8_bytes(&tok)) + }; dynamo_protocols::types::TopLogprobs { token: tok, logprob: top_lp.logprob as f32, diff --git a/lib/llm/src/protocols/openai/completions/delta.rs b/lib/llm/src/protocols/openai/completions/delta.rs index 7cec35531fb2..ac0e35682977 100644 --- a/lib/llm/src/protocols/openai/completions/delta.rs +++ b/lib/llm/src/protocols/openai/completions/delta.rs @@ -560,9 +560,14 @@ mod tests { let top_logprobs = logprobs.top_logprobs[0] .as_array() .expect("top_logprobs array"); + let other = top_logprobs + .iter() + .find(|item| item["token"] == "token_id:999") + .expect("top token_id formatting"); + assert_eq!(other["bytes"], serde_json::json!(b"token_id:999")); let selected = top_logprobs .iter() - .find(|item| item["token_id"] == 123) + .find(|item| item["token"] == "token_id:123") .expect("selected token fallback"); assert_eq!(selected["token"], "token_id:123"); assert_eq!(selected["bytes"], serde_json::json!(b"token_id:123")); From fd3c9e976dcec5dc128afaf0e971df11952c1d5e Mon Sep 17 00:00:00 2001 From: William Arnold Date: Tue, 5 May 2026 21:10:22 -0700 Subject: [PATCH 14/20] fix(openai): return stop_reason via nvext --- .../src/dynamo/frontend/sglang_prepost.py | 16 +-- .../src/dynamo/frontend/sglang_processor.py | 13 +- .../tests/test_sglang_processor_unit.py | 61 ++++++++- components/src/dynamo/frontend/utils.py | 9 ++ .../src/dynamo/frontend/vllm_processor.py | 5 + docs/components/frontend/nvext.md | 10 +- lib/llm/src/audit/stream.rs | 4 - lib/llm/src/engines.rs | 14 ++- lib/llm/src/http/service/openai.rs | 5 - lib/llm/src/perf/logprobs.rs | 3 - .../protocols/anthropic/stream_converter.rs | 3 - lib/llm/src/protocols/anthropic/types.rs | 1 - .../openai/chat_completions/aggregator.rs | 60 ++++++--- .../openai/chat_completions/delta.rs | 32 +++-- .../protocols/openai/chat_completions/jail.rs | 18 --- lib/llm/src/protocols/openai/completions.rs | 1 - .../openai/completions/aggregator.rs | 53 ++++---- .../src/protocols/openai/completions/delta.rs | 41 +++--- lib/llm/src/protocols/openai/nvext.rs | 118 ++++++++++++++++-- lib/llm/src/protocols/openai/responses/mod.rs | 3 - .../openai/responses/stream_converter.rs | 2 - lib/llm/tests/aggregators.rs | 1 - lib/llm/tests/http-service.rs | 2 +- lib/llm/tests/http_metrics.rs | 2 +- lib/llm/tests/kserve_service.rs | 2 +- lib/llm/tests/logprob_analysis_integration.rs | 2 - lib/llm/tests/postprocessor_parsing_stream.rs | 2 - lib/llm/tests/test_jail.rs | 6 - lib/llm/tests/test_reasoning_parser.rs | 1 - lib/llm/tests/test_streaming_tool_parsers.rs | 1 - lib/llm/tests/tool_choice.rs | 1 - lib/protocols/src/types/chat.rs | 13 +- lib/protocols/src/types/completion.rs | 50 ++------ 33 files changed, 339 insertions(+), 216 deletions(-) diff --git a/components/src/dynamo/frontend/sglang_prepost.py b/components/src/dynamo/frontend/sglang_prepost.py index 5d7e5de1a579..25539023d312 100644 --- a/components/src/dynamo/frontend/sglang_prepost.py +++ b/components/src/dynamo/frontend/sglang_prepost.py @@ -660,31 +660,24 @@ def process_output(self, engine_response: dict[str, Any]) -> dict[str, Any] | No raw_ids = engine_response.get("token_ids") token_ids = raw_ids if isinstance(raw_ids, list) else list(raw_ids or []) finish_reason = engine_response.get("finish_reason") - stop_reason = engine_response.get("stop_reason") delta_text = self._incremental_decode(token_ids) if token_ids else "" if self._fast_plain_text: if delta_text: - choice = { + return { "index": 0, "delta": {"role": "assistant", "content": delta_text}, "finish_reason": finish_reason, "logprobs": None, } - if stop_reason is not None: - choice["stop_reason"] = stop_reason - return choice elif finish_reason: - choice = { + return { "index": 0, "delta": {}, "finish_reason": finish_reason, "logprobs": None, } - if stop_reason is not None: - choice["stop_reason"] = stop_reason - return choice return None # -- Reasoning parsing -- @@ -896,14 +889,11 @@ def process_output(self, engine_response: dict[str, Any]) -> dict[str, Any] | No effective_finish = "tool_calls" if has_content or effective_finish: - choice = { + return { "index": 0, "delta": delta if has_content else {}, "finish_reason": effective_finish, "logprobs": None, } - if stop_reason is not None: - choice["stop_reason"] = stop_reason - return choice return None diff --git a/components/src/dynamo/frontend/sglang_processor.py b/components/src/dynamo/frontend/sglang_processor.py index 3f3bb09a9441..ff987410bdc7 100644 --- a/components/src/dynamo/frontend/sglang_processor.py +++ b/components/src/dynamo/frontend/sglang_processor.py @@ -39,7 +39,13 @@ create_parsers, preprocess_chat_request, ) -from .utils import PreprocessError, extract_mm_urls, random_uuid, worker_warmup +from .utils import ( + PreprocessError, + extract_mm_urls, + nvext_extra_field_requested, + random_uuid, + worker_warmup, +) logger = logging.getLogger(__name__) @@ -482,6 +488,7 @@ async def _generate_and_stream( new_ids = engine_response["token_ids"] raw_finish = engine_response.get("finish_reason") finish_reason = _map_finish_reason(raw_finish) + stop_reason = engine_response.get("stop_reason") if usage := engine_response.get("completion_usage"): pending_usage = usage @@ -517,6 +524,10 @@ async def _generate_and_stream( } if pending_usage: dynamo_out["usage"] = pending_usage + if stop_reason is not None and nvext_extra_field_requested( + request, "stop_reason" + ): + dynamo_out["nvext"] = {"stop_reason": stop_reason} yield dynamo_out diff --git a/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py b/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py index 4d85c858f8a6..1a68b0256513 100644 --- a/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py +++ b/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py @@ -10,6 +10,7 @@ """ +import asyncio import json import sys import types @@ -34,12 +35,18 @@ ) from dynamo.frontend.sglang_processor import ( SglangPreprocessWorkerResult, + SglangProcessor, _build_dynamo_preproc, _init_worker, _map_finish_reason, _runtime_config_parser_name, ) -from dynamo.frontend.utils import PreprocessError, random_call_id, random_uuid +from dynamo.frontend.utils import ( + PreprocessError, + nvext_extra_field_requested, + random_call_id, + random_uuid, +) # Needs sglang packages (gpu_1 container). No need for parallel marker. pytestmark = [ @@ -1465,8 +1472,8 @@ def test_finish_reason_only(self, tokenizer): assert choice is not None assert choice["finish_reason"] == "stop" - def test_stop_reason_passthrough(self, tokenizer): - """Backend stop_reason is included on the emitted choice.""" + def test_stop_reason_not_emitted_on_choice(self, tokenizer): + """Backend stop_reason is not part of the OpenAI choice shape.""" post = SglangStreamingPostProcessor( tokenizer=tokenizer, tool_call_parser=None, reasoning_parser=None ) @@ -1476,7 +1483,46 @@ def test_stop_reason_passthrough(self, tokenizer): ) assert choice is not None - assert choice["stop_reason"] == "END" + assert "stop_reason" not in choice + + def test_stop_reason_emits_in_nvext_when_requested(self, tokenizer): + """Frontend emits backend stop_reason under nvext when requested.""" + + class FakeRouter: + async def generate(self, *args, **kwargs): + yield { + "token_ids": [], + "finish_reason": "stop", + "stop_reason": "END", + } + + async def collect(): + processor = SglangProcessor( + tokenizer=tokenizer, + router=FakeRouter(), + tool_call_parser_name=None, + reasoning_parser_name=None, + eos_token_id=None, + ) + post = SglangStreamingPostProcessor( + tokenizer=tokenizer, tool_call_parser=None, reasoning_parser=None + ) + request = { + "model": "test-model", + "nvext": {"extra_fields": ["stop_reason"]}, + } + return [ + item + async for item in processor._generate_and_stream( + "req-stop", request, {}, [], post + ) + ] + + items = asyncio.run(collect()) + + assert len(items) == 1 + assert items[0]["nvext"]["stop_reason"] == "END" + assert "stop_reason" not in items[0]["choices"][0] def test_lookback_trimming(self, tokenizer): """Verify _all_token_ids doesn't grow unbounded.""" @@ -1595,6 +1641,13 @@ def test_preprocess_error(self): # FRONTEND.8 err = PreprocessError("n=2 unsupported") assert "n=2" in str(err) + def test_nvext_extra_field_requested(self): + assert nvext_extra_field_requested( + {"nvext": {"extra_fields": ["stop_reason"]}}, "stop_reason" + ) + assert not nvext_extra_field_requested({"nvext": {}}, "stop_reason") + assert not nvext_extra_field_requested({}, "stop_reason") + # --------------------------------------------------------------------------- # SglangPreprocessWorkerResult picklability diff --git a/components/src/dynamo/frontend/utils.py b/components/src/dynamo/frontend/utils.py index 2ef3f3f0ed71..310cf5c2ae9e 100644 --- a/components/src/dynamo/frontend/utils.py +++ b/components/src/dynamo/frontend/utils.py @@ -20,6 +20,15 @@ def random_call_id() -> str: return f"call_{uuid.uuid4().int & _MASK_64_BITS:016x}" +def nvext_extra_field_requested(request: dict[str, Any], field: str) -> bool: + """Return whether a request opted into a response nvext field.""" + nvext = request.get("nvext") + if not isinstance(nvext, dict): + return False + extra_fields = nvext.get("extra_fields") + return isinstance(extra_fields, list) and field in extra_fields + + def worker_warmup() -> bool: """Dummy task to ensure a ProcessPoolExecutor worker is fully initialized.""" return True diff --git a/components/src/dynamo/frontend/vllm_processor.py b/components/src/dynamo/frontend/vllm_processor.py index 8962bfd02ec5..abb61b2826c3 100644 --- a/components/src/dynamo/frontend/vllm_processor.py +++ b/components/src/dynamo/frontend/vllm_processor.py @@ -49,6 +49,7 @@ extract_mm_urls, handle_engine_error, make_internal_error, + nvext_extra_field_requested, random_uuid, ) @@ -645,6 +646,10 @@ async def _generate_and_stream( } if usage := engine_response.get("completion_usage"): dynamo_out["usage"] = usage + if stop_reason is not None and nvext_extra_field_requested( + request, "stop_reason" + ): + dynamo_out["nvext"] = {"stop_reason": stop_reason} yield dynamo_out _nvtx.end_range(rng_stream) diff --git a/docs/components/frontend/nvext.md b/docs/components/frontend/nvext.md index 2accb4fff98d..e697c66f2b0a 100644 --- a/docs/components/frontend/nvext.md +++ b/docs/components/frontend/nvext.md @@ -35,13 +35,19 @@ Include `nvext` as a top-level field alongside standard OpenAI-compatible fields | `backend_instance_id` | `u64` | `None` | Router | Routes the request to a specific backend instance. | | `token_data` | `u32[]` | `None` | Preprocessor | Pre-tokenized prompt tokens. When provided with `backend_instance_id`, tokenization is skipped. | | `max_thinking_tokens` | `u32` | `None` | Backend | Maximum thinking tokens allowed (passed through to backends). | -| `extra_fields` | `string[]` | `None` | Response builder | Fields to include in the response `nvext`. Supported: `"worker_id"`, `"timing"`, `"routed_experts"`. | +| `extra_fields` | `string[]` | `None` | Response builder | Fields to include in the response `nvext`. Supported: `"worker_id"`, `"timing"`, `"routed_experts"`, `"engine_data"`, `"stop_reason"`. | | `prefill_worker_id` | `u64` | `None` | Router | Routes the request to a specific prefill worker (disaggregated serving). | | `decode_worker_id` | `u64` | `None` | Router | Routes the request to a specific decode worker (disaggregated serving). | | `agent_context` | object | `None` | Preprocessor | Passive workflow and program identity for agent traces. See [Agent Context](#agent-context). | | `agent_hints` | object | `None` | Router | Per-request hints for scheduling and load balancing. See [Agent Hints](#agent-hints). | | `session_control` | object | `None` | Router | Session lifecycle and sticky routing for subagent KV isolation. See [Session Control](#session-control). | +Related root-level Dynamo output option: + +| Field | Type | Default | Consumed By | Description | +|-------|------|---------|-------------|-------------| +| `return_tokens_as_token_ids` | `bool` | `false` | Response builder | Formats logprob token strings as `token_id:` instead of decoded text. | + ### Header Overrides Routing fields can also be set via HTTP headers, which take priority over `nvext` values: @@ -195,6 +201,8 @@ When the client requests response metadata via `extra_fields`, the response incl | `worker_id` | `extra_fields: ["worker_id"]` | Prefill/decode worker IDs and data parallel ranks that processed the request. | | `timing` | `extra_fields: ["timing"]` | Per-request timing information (TTFT, ITL, queue time, etc.). | | `routed_experts` | `extra_fields: ["routed_experts"]` | Routed expert capture payload returned by SGLang-backed requests. | +| `engine_data` | `extra_fields: ["engine_data"]` | Opaque backend-provided engine metadata. | +| `stop_reason` | `extra_fields: ["stop_reason"]` | Backend-specific matched stop condition, returned under `nvext` because it is not part of the OpenAI completions schema. | | `token_ids` | Automatic (GAIE Stage 1) | Tokenized prompt for reuse in Stage 2 query-only mode. | ### Example response `nvext` diff --git a/lib/llm/src/audit/stream.rs b/lib/llm/src/audit/stream.rs index 2663cdd7eeb0..5e0e2d69f2fe 100644 --- a/lib/llm/src/audit/stream.rs +++ b/lib/llm/src/audit/stream.rs @@ -221,7 +221,6 @@ pub fn final_response_to_one_chunk_stream( index: idx as u32, delta, finish_reason: ch.finish_reason, - stop_reason: ch.stop_reason.clone(), logprobs: ch.logprobs.clone(), }; choices.push(choice); @@ -278,7 +277,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }; @@ -319,7 +317,6 @@ mod tests { reasoning_content: None, }, finish_reason: Some(FinishReason::Stop), - stop_reason: None, logprobs: None, }; @@ -452,7 +449,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, } }], diff --git a/lib/llm/src/engines.rs b/lib/llm/src/engines.rs index 7ead7108034b..be5656488c35 100644 --- a/lib/llm/src/engines.rs +++ b/lib/llm/src/engines.rs @@ -159,12 +159,13 @@ impl for c in prompt.chars() { // we are returning characters not tokens, so there will be some postprocessing overhead tokio::time::sleep(*TOKEN_ECHO_DELAY).await; - let response = deltas.create_choice(0, Some(c.to_string()), None, None, None); + let response = deltas.create_choice(0, Some(c.to_string()), None, None); yield Annotated{ id: Some(id.to_string()), data: Some(response), event: None, comment: None, error: None }; id += 1; } - let response = deltas.create_choice(0, None, Some(dynamo_protocols::types::FinishReason::Stop), None, None); + let response = + deltas.create_choice(0, None, Some(dynamo_protocols::types::FinishReason::Stop), None); yield Annotated { id: Some(id.to_string()), data: Some(response), event: None, comment: None, error: None }; }; @@ -192,11 +193,16 @@ impl let mut id = 1; for c in chars_string.chars() { tokio::time::sleep(*TOKEN_ECHO_DELAY).await; - let response = deltas.create_choice(0, Some(c.to_string()), None, None, None); + let response = deltas.create_choice(0, Some(c.to_string()), None, None); yield Annotated{ id: Some(id.to_string()), data: Some(response), event: None, comment: None, error: None }; id += 1; } - let response = deltas.create_choice(0, None, Some(dynamo_protocols::types::CompletionFinishReason::Stop), None, None); + let response = deltas.create_choice( + 0, + None, + Some(dynamo_protocols::types::CompletionFinishReason::Stop), + None, + ); yield Annotated { id: Some(id.to_string()), data: Some(response), event: None, comment: None, error: None }; }; diff --git a/lib/llm/src/http/service/openai.rs b/lib/llm/src/http/service/openai.rs index a65441c33677..43dae9dc766e 100644 --- a/lib/llm/src/http/service/openai.rs +++ b/lib/llm/src/http/service/openai.rs @@ -3765,7 +3765,6 @@ mod tests { reasoning_content: reasoning.map(|s| s.to_string()), }, finish_reason: finish, - stop_reason: None, logprobs: None, } } @@ -3797,7 +3796,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, } } @@ -3899,7 +3897,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }; @@ -3971,7 +3968,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }; @@ -4008,7 +4004,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }; diff --git a/lib/llm/src/perf/logprobs.rs b/lib/llm/src/perf/logprobs.rs index 02c93cdf1ecf..6873ad49a27c 100644 --- a/lib/llm/src/perf/logprobs.rs +++ b/lib/llm/src/perf/logprobs.rs @@ -963,7 +963,6 @@ mod tests { reasoning_content: None, }, finish_reason: Some(FinishReason::Stop), - stop_reason: None, logprobs: Some(ChatChoiceLogprobs { content: Some(token_logprobs), refusal: None, @@ -999,7 +998,6 @@ mod tests { reasoning_content: None, }, finish_reason: Some(FinishReason::Stop), - stop_reason: None, logprobs: Some(ChatChoiceLogprobs { content: Some(token_logprobs), refusal: None, @@ -1353,7 +1351,6 @@ mod tests { reasoning_content: None, }, finish_reason: Some(FinishReason::Stop), - stop_reason: None, logprobs: None, // No logprobs }], created: 1234567890, diff --git a/lib/llm/src/protocols/anthropic/stream_converter.rs b/lib/llm/src/protocols/anthropic/stream_converter.rs index 426ff3e7164b..bb8220445a79 100644 --- a/lib/llm/src/protocols/anthropic/stream_converter.rs +++ b/lib/llm/src/protocols/anthropic/stream_converter.rs @@ -756,7 +756,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }], created: 0, @@ -799,7 +798,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }], created: 0, @@ -946,7 +944,6 @@ mod tests { reasoning_content: Some(text.into()), }, finish_reason: None, - stop_reason: None, logprobs: None, }], created: 0, diff --git a/lib/llm/src/protocols/anthropic/types.rs b/lib/llm/src/protocols/anthropic/types.rs index 0d94644748c8..a199713ec90f 100644 --- a/lib/llm/src/protocols/anthropic/types.rs +++ b/lib/llm/src/protocols/anthropic/types.rs @@ -805,7 +805,6 @@ mod tests { reasoning_content: None, }, finish_reason: Some(dynamo_protocols::types::FinishReason::Stop), - stop_reason: None, logprobs: None, }], created: 1726000000, diff --git a/lib/llm/src/protocols/openai/chat_completions/aggregator.rs b/lib/llm/src/protocols/openai/chat_completions/aggregator.rs index 5def709448d8..0849b557fe21 100644 --- a/lib/llm/src/protocols/openai/chat_completions/aggregator.rs +++ b/lib/llm/src/protocols/openai/chat_completions/aggregator.rs @@ -11,10 +11,10 @@ use crate::protocols::{ Annotated, codec::{Message, SseCodecError}, convert_sse_stream, - openai::ParsingOptions, + openai::{ParsingOptions, nvext::merge_response_nvext}, }; -use dynamo_protocols::types::{ChatCompletionMessageContent, StopReason}; +use dynamo_protocols::types::ChatCompletionMessageContent; use dynamo_runtime::engine::DataStream; /// Aggregates a stream of [`NvCreateChatCompletionStreamResponse`]s into a single @@ -52,8 +52,6 @@ struct DeltaChoice { role: Option, /// The reason the completion was finished (if applicable). finish_reason: Option, - /// The stop string or token that triggered the stop condition. - stop_reason: Option, /// Optional log probabilities for the chat choice. logprobs: Option, // Tool-call chunks accumulated in the order they arrived from the stream, @@ -238,10 +236,7 @@ impl DeltaAggregator { aggregator.system_fingerprint = Some(system_fingerprint); } - // Aggregate nvext field (take the last non-None value) - if delta.nvext.is_some() { - aggregator.nvext = delta.nvext; - } + merge_response_nvext(&mut aggregator.nvext, delta.nvext); // Aggregate choices incrementally. for choice in delta.inner.choices { @@ -254,7 +249,6 @@ impl DeltaAggregator { text: "".to_string(), role: choice.delta.role, finish_reason: None, - stop_reason: None, logprobs: None, tool_call_chunks: BTreeMap::new(), tool_calls: None, @@ -307,11 +301,6 @@ impl DeltaAggregator { state_choice.finish_reason = Some(finish_reason); } - // Update stop reason if provided. - if let Some(stop_reason) = choice.stop_reason { - state_choice.stop_reason = Some(stop_reason); - } - // Update logprobs if let Some(logprobs) = &choice.logprobs { let state_lps = state_choice.logprobs.get_or_insert( @@ -466,7 +455,6 @@ impl From for dynamo_protocols::types::ChatChoice { }, index: delta.index, finish_reason, - stop_reason: delta.stop_reason, logprobs: delta.logprobs, } } @@ -581,7 +569,6 @@ mod tests { index, delta, finish_reason, - stop_reason: None, logprobs, }; @@ -631,7 +618,6 @@ mod tests { index, delta, finish_reason, - stop_reason: None, logprobs: None, }; let data = NvCreateChatCompletionStreamResponse { @@ -1041,6 +1027,43 @@ mod tests { assert_eq!(choice.message.role, dynamo_protocols::types::Role::User); } + #[tokio::test] + async fn test_multiple_deltas_merge_nvext_fields() { + let mut annotated_delta1 = create_test_delta( + 0, + "Hello", + Some(dynamo_protocols::types::Role::Assistant), + None, + None, + None, + ); + annotated_delta1.data.as_mut().expect("delta data").nvext = + Some(serde_json::json!({ "engine_data": { "trace_id": "abc" } })); + let mut annotated_delta2 = create_test_delta( + 0, + " world", + None, + Some(dynamo_protocols::types::FinishReason::Stop), + None, + None, + ); + annotated_delta2.data.as_mut().expect("delta data").nvext = + Some(serde_json::json!({ "stop_reason": 128001 })); + + let stream = Box::pin(stream::iter(vec![annotated_delta1, annotated_delta2])); + let response = DeltaAggregator::apply(stream, ParsingOptions::default()) + .await + .expect("aggregate stream"); + + assert_eq!( + response.nvext, + Some(serde_json::json!({ + "engine_data": { "trace_id": "abc" }, + "stop_reason": 128001, + })) + ); + } + #[allow(deprecated)] #[tokio::test] async fn test_multiple_choices() { @@ -1068,7 +1091,6 @@ mod tests { reasoning_content: None, }, finish_reason: Some(dynamo_protocols::types::FinishReason::Stop), - stop_reason: None, logprobs: None, }, dynamo_protocols::types::ChatChoiceStream { @@ -1084,7 +1106,6 @@ mod tests { reasoning_content: None, }, finish_reason: Some(dynamo_protocols::types::FinishReason::Stop), - stop_reason: None, logprobs: None, }, ], @@ -1557,7 +1578,6 @@ mod tests { text: String::new(), role: Some(dynamo_protocols::types::Role::Assistant), finish_reason: Some(dynamo_protocols::types::FinishReason::Stop), - stop_reason: None, logprobs: None, tool_call_chunks: BTreeMap::new(), tool_calls: None, diff --git a/lib/llm/src/protocols/openai/chat_completions/delta.rs b/lib/llm/src/protocols/openai/chat_completions/delta.rs index e6e76d8ad7a5..6f42b2f54556 100644 --- a/lib/llm/src/protocols/openai/chat_completions/delta.rs +++ b/lib/llm/src/protocols/openai/chat_completions/delta.rs @@ -237,18 +237,15 @@ impl DeltaGenerator { /// * `text` - The text content for the response. /// * `finish_reason` - The reason why the response finished (e.g., stop, length, etc.). /// * `logprobs` - Optional log probabilities of the generated tokens. - /// * `stop_reason` - Optional stop string or token that triggered the stop. /// /// # Returns /// * An [`dynamo_protocols::types::CreateChatCompletionStreamResponse`] instance representing the choice. - #[allow(deprecated)] pub fn create_choice( &mut self, index: u32, text: Option, finish_reason: Option, logprobs: Option, - stop_reason: Option, ) -> NvCreateChatCompletionStreamResponse { let delta = dynamo_protocols::types::ChatCompletionStreamResponseDelta { content: text.map(dynamo_protocols::types::ChatCompletionMessageContent::Text), @@ -267,7 +264,6 @@ impl DeltaGenerator { index, delta, finish_reason, - stop_reason, logprobs, }; @@ -401,16 +397,11 @@ impl crate::protocols::openai::DeltaGeneratorExt None, }; + let stop_reason = delta.stop_reason.clone(); // Create the streaming response. let index = delta.index.unwrap_or(0); - let mut stream_response = self.create_choice( - index, - delta.text, - finish_reason, - logprobs, - delta.stop_reason, - ); + let mut stream_response = self.create_choice(index, delta.text, finish_reason, logprobs); // Record finish for timing/ITL accounting even when timing is not returned to the client. // Kept at call site because it's a side effect on the tracker — not a gating decision. @@ -429,6 +420,7 @@ impl crate::protocols::openai::DeltaGeneratorExt>, finish_reason: Option, - stop_reason: Option, logprobs: Option, ) -> ChatChoiceStream { #[allow(deprecated)] @@ -140,7 +139,6 @@ fn create_choice_stream( reasoning_content: None, }, finish_reason, - stop_reason, logprobs, } } @@ -228,7 +226,6 @@ impl ChoiceJailState { &prefix, None, choice.finish_reason, - None, choice.logprobs.clone(), ); emissions.push(ChoiceEmission::PassThrough(prefix_choice)); @@ -279,7 +276,6 @@ impl ChoiceJailState { trailing_part, None, choice.finish_reason, - None, choice.logprobs.clone(), ); emissions.push(ChoiceEmission::Trailing(trailing_choice)); @@ -310,7 +306,6 @@ impl ChoiceJailState { &prefix, None, choice.finish_reason, - None, choice.logprobs.clone(), ); emissions.push(ChoiceEmission::PassThrough(prefix_choice)); @@ -352,7 +347,6 @@ impl ChoiceJailState { &content, None, choice.finish_reason, - None, choice.logprobs.clone(), ); emissions.push(ChoiceEmission::PassThrough(pass_through_choice)); @@ -415,7 +409,6 @@ impl ChoiceJailState { &trailing_owned, None, choice.finish_reason, - None, choice.logprobs.clone(), ); emissions.push(ChoiceEmission::Trailing(trailing_choice)); @@ -438,7 +431,6 @@ impl ChoiceJailState { &self.accumulated_content, None, self.stream_finish_reason, // For the accumulated content, assign the original stream finish reason, otherwise it will get lost - None, self.accumulated_logprobs.clone(), ); @@ -666,7 +658,6 @@ impl JailedStream { index: choice.index, delta: choice.delta.clone(), finish_reason: choice.finish_reason, - stop_reason: choice.stop_reason.clone(), logprobs: choice.logprobs.clone(), }; all_emissions.push(ChoiceEmission::PassThrough(pass_through_choice)); @@ -980,7 +971,6 @@ impl JailedStream { normal_text.as_deref().unwrap_or(""), None, base_choice.finish_reason, - base_choice.stop_reason.clone(), base_choice.logprobs.clone(), ); } @@ -1005,7 +995,6 @@ impl JailedStream { normal_text.as_deref().unwrap_or(""), Some(tool_call_chunks), None, - None, base_choice.logprobs.clone(), ) } @@ -1033,7 +1022,6 @@ impl JailedStream { content, None, base_choice.finish_reason, - base_choice.stop_reason.clone(), base_choice.logprobs.clone(), ) } @@ -1052,7 +1040,6 @@ impl JailedStream { "", None, base_choice.finish_reason, - base_choice.stop_reason.clone(), base_choice.logprobs.clone(), ) } @@ -1171,7 +1158,6 @@ impl JailedStream { "", Some(tool_call_chunks), base_choice.finish_reason, - None, base_choice.logprobs.clone(), ) } else if filter_dropped_all { @@ -1183,7 +1169,6 @@ impl JailedStream { "", None, base_choice.finish_reason, - base_choice.stop_reason.clone(), base_choice.logprobs.clone(), ) } else { @@ -1194,7 +1179,6 @@ impl JailedStream { accumulated_content, None, base_choice.finish_reason, - base_choice.stop_reason.clone(), base_choice.logprobs.clone(), ) } @@ -1568,7 +1552,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }; @@ -1669,7 +1652,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: Some(logprobs), }; diff --git a/lib/llm/src/protocols/openai/completions.rs b/lib/llm/src/protocols/openai/completions.rs index 8bed1ab0602f..98b5084d05d3 100644 --- a/lib/llm/src/protocols/openai/completions.rs +++ b/lib/llm/src/protocols/openai/completions.rs @@ -397,7 +397,6 @@ impl TryFrom for dynamo_protocols::types::C index, logprobs, finish_reason, - stop_reason: response.delta.stop_reason, }; Ok(choice) diff --git a/lib/llm/src/protocols/openai/completions/aggregator.rs b/lib/llm/src/protocols/openai/completions/aggregator.rs index ab5d72ccc784..90dfe93f9733 100644 --- a/lib/llm/src/protocols/openai/completions/aggregator.rs +++ b/lib/llm/src/protocols/openai/completions/aggregator.rs @@ -12,7 +12,7 @@ use crate::protocols::{ codec::{Message, SseCodecError}, common::FinishReason, convert_sse_stream, - openai::ParsingOptions, + openai::{ParsingOptions, nvext::merge_response_nvext}, }; /// Aggregates a stream of [`CompletionResponse`]s into a single [`CompletionResponse`]. @@ -31,7 +31,6 @@ struct DeltaChoice { index: u32, text: String, finish_reason: Option, - stop_reason: Option, logprobs: Option, } @@ -86,10 +85,7 @@ impl DeltaAggregator { if let Some(system_fingerprint) = delta.inner.system_fingerprint { aggregator.system_fingerprint = Some(system_fingerprint); } - // Aggregate nvext field (take the last non-None value) - if delta.nvext.is_some() { - aggregator.nvext = delta.nvext; - } + merge_response_nvext(&mut aggregator.nvext, delta.nvext); // handle the choices for choice in delta.inner.choices { @@ -101,7 +97,6 @@ impl DeltaAggregator { index: choice.index, text: "".to_string(), finish_reason: None, - stop_reason: None, logprobs: None, }); @@ -123,10 +118,6 @@ impl DeltaAggregator { None => None, }; - if let Some(stop_reason) = choice.stop_reason { - state_choice.stop_reason = Some(stop_reason); - } - // Update logprobs if let Some(logprobs) = &choice.logprobs { let state_lps = state_choice.logprobs.get_or_insert( @@ -193,7 +184,6 @@ impl From for dynamo_protocols::types::Choice { index: delta.index, text: delta.text, finish_reason, - stop_reason: delta.stop_reason, logprobs: delta.logprobs, } } @@ -262,7 +252,6 @@ mod tests { index, text: text.to_string(), finish_reason, - stop_reason: None, logprobs, }], object: "text_completion".to_string(), @@ -344,13 +333,8 @@ mod tests { let annotated_delta1 = create_test_delta(0, "Hello,", None, Some(-0.1)); let mut annotated_delta2 = create_test_delta(0, " world!", Some("stop".to_string()), Some(-0.2)); - annotated_delta2 - .data - .as_mut() - .expect("delta data") - .inner - .choices[0] - .stop_reason = Some(dynamo_protocols::types::StopReason::Int(128001)); + annotated_delta2.data.as_mut().expect("delta data").nvext = + Some(serde_json::json!({ "stop_reason": 128001 })); // Create a stream let annotated_deltas = vec![annotated_delta1, annotated_delta2]; @@ -373,8 +357,8 @@ mod tests { Some(dynamo_protocols::types::CompletionFinishReason::Stop) ); assert_eq!( - choice.stop_reason, - Some(dynamo_protocols::types::StopReason::Int(128001)) + response.nvext, + Some(serde_json::json!({ "stop_reason": 128001 })) ); assert_eq!(choice.logprobs.as_ref().unwrap().tokens.len(), 2); assert_eq!( @@ -383,6 +367,29 @@ mod tests { ); } + #[tokio::test] + async fn test_multiple_deltas_merge_nvext_fields() { + let mut annotated_delta1 = create_test_delta(0, "Hello,", None, None); + annotated_delta1.data.as_mut().expect("delta data").nvext = + Some(serde_json::json!({ "engine_data": { "trace_id": "abc" } })); + let mut annotated_delta2 = create_test_delta(0, " world!", Some("stop".to_string()), None); + annotated_delta2.data.as_mut().expect("delta data").nvext = + Some(serde_json::json!({ "stop_reason": 128001 })); + + let stream = Box::pin(stream::iter(vec![annotated_delta1, annotated_delta2])); + let response = DeltaAggregator::apply(stream, ParsingOptions::default()) + .await + .expect("aggregate stream"); + + assert_eq!( + response.nvext, + Some(serde_json::json!({ + "engine_data": { "trace_id": "abc" }, + "stop_reason": 128001, + })) + ); + } + #[tokio::test] async fn test_multiple_choices() { // Create a delta with multiple choices @@ -397,14 +404,12 @@ mod tests { index: 0, text: "Choice 0".to_string(), finish_reason: Some(dynamo_protocols::types::CompletionFinishReason::Stop), - stop_reason: None, logprobs: None, }, dynamo_protocols::types::Choice { index: 1, text: "Choice 1".to_string(), finish_reason: Some(dynamo_protocols::types::CompletionFinishReason::Stop), - stop_reason: None, logprobs: None, }, ], diff --git a/lib/llm/src/protocols/openai/completions/delta.rs b/lib/llm/src/protocols/openai/completions/delta.rs index ac0e35682977..38cc8d66b642 100644 --- a/lib/llm/src/protocols/openai/completions/delta.rs +++ b/lib/llm/src/protocols/openai/completions/delta.rs @@ -198,7 +198,6 @@ impl DeltaGenerator { index: u32, text: Option, finish_reason: Option, - stop_reason: Option, logprobs: Option, ) -> NvCreateCompletionResponse { // todo - update for tool calling @@ -216,7 +215,6 @@ impl DeltaGenerator { text: text.unwrap_or_default(), index, finish_reason, - stop_reason, logprobs, }], usage: if self.options.enable_usage && self.options.continuous_usage_stats { @@ -309,16 +307,11 @@ impl crate::protocols::openai::DeltaGeneratorExt for ); let finish_reason = delta.finish_reason.map(Into::into); + let stop_reason = delta.stop_reason.clone(); // create choice let index = delta.index.unwrap_or(0); - let mut response = self.create_choice( - index, - delta.text.clone(), - finish_reason, - delta.stop_reason, - logprobs, - ); + let mut response = self.create_choice(index, delta.text.clone(), finish_reason, logprobs); // Record finish for timing/ITL accounting even when timing is not returned to the client. // Kept at call site because it's a side effect on the tracker — not a gating decision. @@ -337,6 +330,7 @@ impl crate::protocols::openai::DeltaGeneratorExt for delta.disaggregated_params.as_ref(), finish_reason.is_some(), delta.engine_data, + stop_reason, ) && let Ok(nvext_json) = serde_json::to_value(&nvext_response) { response.nvext = Some(nvext_json); @@ -493,7 +487,7 @@ mod tests { } #[test] - fn test_choice_from_postprocessor_preserves_stop_reason() { + fn test_stop_reason_is_suppressed_without_nvext_extra_field() { let request = create_test_request(); let mut generator = request.response_generator("req-stop-reason".to_string()); let mut output = final_backend_output(); @@ -505,12 +499,27 @@ mod tests { .choice_from_postprocessor(output) .expect("choice generation"); - assert_eq!( - response.inner.choices[0].stop_reason, - Some(dynamo_protocols::types::StopReason::String( - "END".to_string() - )) - ); + let response_json = serde_json::to_value(&response).expect("serialize response"); + assert!(response_json["choices"][0].get("stop_reason").is_none()); + assert!(response_json.get("nvext").is_none()); + } + + #[test] + fn test_stop_reason_emits_in_nvext_when_requested() { + let request = create_test_request_with_extra_fields(vec!["stop_reason".to_string()]); + let mut generator = request.response_generator("req-stop-reason-nvext".to_string()); + let mut output = final_backend_output(); + output.stop_reason = Some(dynamo_protocols::types::StopReason::String( + "END".to_string(), + )); + + let response = generator + .choice_from_postprocessor(output) + .expect("choice generation"); + + let response_json = serde_json::to_value(&response).expect("serialize response"); + assert!(response_json["choices"][0].get("stop_reason").is_none()); + assert_eq!(response_json["nvext"]["stop_reason"], "END"); } #[test] diff --git a/lib/llm/src/protocols/openai/nvext.rs b/lib/llm/src/protocols/openai/nvext.rs index 90919836fc29..8459b3085699 100644 --- a/lib/llm/src/protocols/openai/nvext.rs +++ b/lib/llm/src/protocols/openai/nvext.rs @@ -3,6 +3,7 @@ use axum::http::HeaderMap; use derive_builder::Builder; +use dynamo_protocols::types::StopReason; use serde::{Deserialize, Serialize}; use utoipa::ToSchema; use validator::{Validate, ValidationError}; @@ -120,6 +121,29 @@ pub struct NvExtResponse { /// Dynamo does not inspect this; it is forwarded as-is to the client. #[serde(skip_serializing_if = "Option::is_none")] pub engine_data: Option, + + /// Backend-specific matched stop condition. This is not part of the + /// OpenAI response schema, so it is only returned under nvext when requested. + #[serde(skip_serializing_if = "Option::is_none")] + pub stop_reason: Option, +} + +pub(crate) fn merge_response_nvext( + target: &mut Option, + incoming: Option, +) { + let Some(incoming) = incoming else { + return; + }; + + match (target.as_mut(), incoming) { + (Some(serde_json::Value::Object(target_obj)), serde_json::Value::Object(incoming_obj)) => { + target_obj.extend(incoming_obj); + } + (_, incoming) => { + *target = Some(incoming); + } + } } /// Response nvext fields requested for a given request. @@ -137,6 +161,7 @@ pub struct NvExtResponseFieldSelection { pub token_ids: bool, pub routed_experts: bool, pub engine_data: bool, + pub stop_reason: bool, } impl NvExtResponseFieldSelection { @@ -153,6 +178,7 @@ impl NvExtResponseFieldSelection { "timing" => selection.timing = true, "routed_experts" => selection.routed_experts = true, "engine_data" => selection.engine_data = true, + "stop_reason" => selection.stop_reason = true, _ => {} } } @@ -185,12 +211,14 @@ impl NvExtResponseFieldSelection { /// `disaggregated_params` (cloned as-is, no validation). /// - `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`. pub fn build_response_nvext( &self, tracker: Option<&std::sync::Arc>, disaggregated_params: Option<&serde_json::Value>, finish_reason_present: bool, engine_data_from_backend: Option, + stop_reason_from_backend: Option, ) -> Option { let worker_id = if self.worker_id { tracker.and_then(|t| t.get_worker_info()) @@ -226,11 +254,18 @@ impl NvExtResponseFieldSelection { None }; + let stop_reason = if self.stop_reason { + stop_reason_from_backend.and_then(|reason| serde_json::to_value(reason).ok()) + } 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() { return None; } @@ -241,6 +276,7 @@ impl NvExtResponseFieldSelection { token_ids, routed_experts, engine_data, + stop_reason, }) } } @@ -292,7 +328,7 @@ pub struct NvExt { /// 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", - /// which map to fields in NvExtResponse. + /// "stop_reason", which map to fields in NvExtResponse. #[serde(default, skip_serializing_if = "Option::is_none")] #[builder(default, setter(strip_option))] pub extra_fields: Option>, @@ -709,6 +745,22 @@ mod tests { ); } + #[test] + fn test_nvext_response_field_selection_stop_reason_only() { + let nvext = NvExt::builder() + .extra_fields(vec!["stop_reason".to_string()]) + .build() + .unwrap(); + + assert_eq!( + NvExtResponseFieldSelection::from_nvext(Some(&nvext)), + NvExtResponseFieldSelection { + stop_reason: true, + ..Default::default() + } + ); + } + // Helpers for build_response_nvext tests ----------------------------- fn sel_all_false() -> NvExtResponseFieldSelection { @@ -736,11 +788,13 @@ 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).is_none(), + sel.build_response_nvext(None, None, false, None, None) + .is_none(), "no fields selected → None" ); assert!( - sel.build_response_nvext(None, None, true, None).is_none(), + sel.build_response_nvext(None, None, true, None, None) + .is_none(), "finish_reason alone does not force emission" ); } @@ -755,7 +809,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) + .build_response_nvext(Some(&tracker), None, false, None, None) .expect("worker_id should emit regardless of finish_reason"); assert!(out.worker_id.is_some()); @@ -774,7 +828,7 @@ mod tests { // timing alone + finish_reason=false → nothing to emit, returns None. assert!( - sel.build_response_nvext(Some(&tracker), None, false, None) + sel.build_response_nvext(Some(&tracker), None, false, None, None) .is_none(), "timing is gated on finish_reason_present" ); @@ -789,7 +843,7 @@ mod tests { let tracker = tracker_with_prefill_worker(); let out = sel - .build_response_nvext(Some(&tracker), None, true, None) + .build_response_nvext(Some(&tracker), None, true, None, None) .expect("timing should emit on finish"); assert!(out.timing.is_some()); @@ -805,7 +859,10 @@ mod tests { ..Default::default() }; // finish=true but no tracker → timing not populated → None. - assert!(sel.build_response_nvext(None, None, true, None).is_none()); + assert!( + sel.build_response_nvext(None, None, true, None, None) + .is_none() + ); } #[test] @@ -817,7 +874,7 @@ mod tests { let params = disagg_params_full(); let out = sel - .build_response_nvext(None, Some(¶ms), false, None) + .build_response_nvext(None, Some(¶ms), false, None, None) .expect("token_ids should emit when present"); assert_eq!(out.token_ids, Some(vec![11u32, 22, 33])); @@ -836,7 +893,7 @@ mod tests { let params = serde_json::json!({ "token_ids": "not-an-array" }); assert!( - sel.build_response_nvext(None, Some(¶ms), false, None) + sel.build_response_nvext(None, Some(¶ms), false, None, None) .is_none(), "malformed token_ids silently suppressed; nothing else selected → None" ); @@ -851,7 +908,7 @@ mod tests { let params = disagg_params_full(); let out = sel - .build_response_nvext(None, Some(¶ms), false, None) + .build_response_nvext(None, Some(¶ms), false, None, None) .expect("routed_experts should emit when present"); assert_eq!( @@ -860,6 +917,43 @@ mod tests { ); } + #[test] + fn test_build_response_nvext_stop_reason_when_requested() { + let sel = NvExtResponseFieldSelection { + stop_reason: true, + ..Default::default() + }; + + let out = sel + .build_response_nvext( + None, + None, + true, + None, + Some(StopReason::String("END".to_string())), + ) + .expect("stop_reason should emit when requested and present"); + + assert_eq!(out.stop_reason, Some(serde_json::json!("END"))); + assert!(out.worker_id.is_none()); + assert!(out.timing.is_none()); + assert!(out.token_ids.is_none()); + assert!(out.routed_experts.is_none()); + } + + #[test] + fn test_build_response_nvext_stop_reason_suppressed_when_absent() { + let sel = NvExtResponseFieldSelection { + stop_reason: true, + ..Default::default() + }; + + assert!( + sel.build_response_nvext(None, None, true, None, None) + .is_none() + ); + } + #[test] fn test_build_response_nvext_combined_emission() { let sel = NvExtResponseFieldSelection { @@ -868,12 +962,13 @@ mod tests { token_ids: true, routed_experts: true, engine_data: false, + stop_reason: false, }; let tracker = tracker_with_prefill_worker(); let params = disagg_params_full(); let out = sel - .build_response_nvext(Some(&tracker), Some(¶ms), true, None) + .build_response_nvext(Some(&tracker), Some(¶ms), true, None, None) .expect("all fields selected and available → Some"); assert!(out.worker_id.is_some()); @@ -904,6 +999,7 @@ mod tests { token_ids: false, // only enabled via query_instance_id routed_experts: true, engine_data: false, + stop_reason: false, } ); } diff --git a/lib/llm/src/protocols/openai/responses/mod.rs b/lib/llm/src/protocols/openai/responses/mod.rs index cbc97aa80dd0..ae1c868af84a 100644 --- a/lib/llm/src/protocols/openai/responses/mod.rs +++ b/lib/llm/src/protocols/openai/responses/mod.rs @@ -2103,7 +2103,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }], created: now, @@ -2164,7 +2163,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }], created: now, @@ -2572,7 +2570,6 @@ thinking audio: None, }, finish_reason: Some(FinishReason::Stop), - stop_reason: None, logprobs: None, }], created: 0, diff --git a/lib/llm/src/protocols/openai/responses/stream_converter.rs b/lib/llm/src/protocols/openai/responses/stream_converter.rs index ff631bff01f7..d4a3a4948939 100644 --- a/lib/llm/src/protocols/openai/responses/stream_converter.rs +++ b/lib/llm/src/protocols/openai/responses/stream_converter.rs @@ -726,7 +726,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }], created: 0, @@ -756,7 +755,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }], created: 0, diff --git a/lib/llm/tests/aggregators.rs b/lib/llm/tests/aggregators.rs index 5ce5e86eb95f..367522b89118 100644 --- a/lib/llm/tests/aggregators.rs +++ b/lib/llm/tests/aggregators.rs @@ -172,7 +172,6 @@ fn make_stream_delta( reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }] } else { diff --git a/lib/llm/tests/http-service.rs b/lib/llm/tests/http-service.rs index 71d0c15b47bf..9bdc40d7ea6a 100644 --- a/lib/llm/tests/http-service.rs +++ b/lib/llm/tests/http-service.rs @@ -85,7 +85,7 @@ impl let stream = stream! { tokio::time::sleep(std::time::Duration::from_millis(max_tokens)).await; for i in 0..10 { - let output = generator.create_choice(i, Some(format!("choice {i}")), None, None, None); + let output = generator.create_choice(i, Some(format!("choice {i}")), None, None); yield Annotated::from_data(output); } diff --git a/lib/llm/tests/http_metrics.rs b/lib/llm/tests/http_metrics.rs index dd54a9011ffc..2d295e89359b 100644 --- a/lib/llm/tests/http_metrics.rs +++ b/lib/llm/tests/http_metrics.rs @@ -55,7 +55,7 @@ impl // output_sequence_tokens is properly recorded (the histogram only // records when osl > 0, which requires the annotation to be present). for i in 0..5 { - let output = generator.create_choice(i, Some(format!("Mock response {i}")), None, None, None); + let output = generator.create_choice(i, Some(format!("Mock response {i}")), None, None); let mut annotated = Annotated::from_data(output); let metrics = LLMMetricAnnotation { input_tokens: 5, diff --git a/lib/llm/tests/kserve_service.rs b/lib/llm/tests/kserve_service.rs index dfde0cd9097b..0f24f450b88a 100644 --- a/lib/llm/tests/kserve_service.rs +++ b/lib/llm/tests/kserve_service.rs @@ -119,7 +119,7 @@ pub mod kserve_test { let stream = stream! { tokio::time::sleep(std::time::Duration::from_millis(10)).await; for word in word_list { - yield Annotated::from_data(generator.create_choice(0, Some(word.to_string()), None, None, None)); + yield Annotated::from_data(generator.create_choice(0, Some(word.to_string()), None, None)); } }; diff --git a/lib/llm/tests/logprob_analysis_integration.rs b/lib/llm/tests/logprob_analysis_integration.rs index 1adcf655ae06..f684991fbf5d 100644 --- a/lib/llm/tests/logprob_analysis_integration.rs +++ b/lib/llm/tests/logprob_analysis_integration.rs @@ -388,7 +388,6 @@ fn create_response_with_linear_probs( reasoning_content: None, }, finish_reason: Some(FinishReason::Stop), - stop_reason: None, logprobs: Some(ChatChoiceLogprobs { content: Some(token_logprobs), refusal: None, @@ -469,7 +468,6 @@ fn create_multi_choice_response( reasoning_content: None, }, finish_reason: Some(FinishReason::Stop), - stop_reason: None, logprobs: Some(ChatChoiceLogprobs { content: Some(token_logprobs), refusal: None, diff --git a/lib/llm/tests/postprocessor_parsing_stream.rs b/lib/llm/tests/postprocessor_parsing_stream.rs index 2bb48cd4cbfc..9a5ef8fb3305 100644 --- a/lib/llm/tests/postprocessor_parsing_stream.rs +++ b/lib/llm/tests/postprocessor_parsing_stream.rs @@ -289,7 +289,6 @@ fn mock_content_chunk(content: &str) -> NvCreateChatCompletionStreamResponse { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }; NvCreateChatCompletionStreamResponse { @@ -324,7 +323,6 @@ fn mock_final_chunk() -> NvCreateChatCompletionStreamResponse { reasoning_content: None, }, finish_reason: Some(FinishReason::Stop), - stop_reason: None, logprobs: None, }; NvCreateChatCompletionStreamResponse { diff --git a/lib/llm/tests/test_jail.rs b/lib/llm/tests/test_jail.rs index 7a8da4aa4dcd..9c79eead37ab 100644 --- a/lib/llm/tests/test_jail.rs +++ b/lib/llm/tests/test_jail.rs @@ -45,7 +45,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }; @@ -88,7 +87,6 @@ mod tests { reasoning_content: None, }, finish_reason: Some(FinishReason::Stop), - stop_reason: None, logprobs: None, }; @@ -135,7 +133,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, }; @@ -181,7 +178,6 @@ mod tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, } }) @@ -229,7 +225,6 @@ mod tests { reasoning_content: None, }, finish_reason: Some(FinishReason::Stop), - stop_reason: None, logprobs: None, } }) @@ -2405,7 +2400,6 @@ mod parallel_jail_tests { reasoning_content: None, }, finish_reason: None, - stop_reason: None, logprobs: None, } }) diff --git a/lib/llm/tests/test_reasoning_parser.rs b/lib/llm/tests/test_reasoning_parser.rs index 49113b8a660a..c5d0202816b8 100644 --- a/lib/llm/tests/test_reasoning_parser.rs +++ b/lib/llm/tests/test_reasoning_parser.rs @@ -34,7 +34,6 @@ fn create_mock_response_chunk( reasoning_content, }, finish_reason: None, - stop_reason: None, logprobs: None, }; diff --git a/lib/llm/tests/test_streaming_tool_parsers.rs b/lib/llm/tests/test_streaming_tool_parsers.rs index d8b95b9d96bd..2fb30575ce8e 100644 --- a/lib/llm/tests/test_streaming_tool_parsers.rs +++ b/lib/llm/tests/test_streaming_tool_parsers.rs @@ -1290,7 +1290,6 @@ mod tests { reasoning_content: None, }, finish_reason, - stop_reason: None, logprobs: None, }; Annotated { diff --git a/lib/llm/tests/tool_choice.rs b/lib/llm/tests/tool_choice.rs index 3a1dd71151bb..0a720f985a36 100644 --- a/lib/llm/tests/tool_choice.rs +++ b/lib/llm/tests/tool_choice.rs @@ -491,7 +491,6 @@ fn make_text_chunk( } else { None }, - stop_reason: None, logprobs: None, }], created: 1234567890, diff --git a/lib/protocols/src/types/chat.rs b/lib/protocols/src/types/chat.rs index 38fdfad11a75..4fcf8629a058 100644 --- a/lib/protocols/src/types/chat.rs +++ b/lib/protocols/src/types/chat.rs @@ -51,6 +51,7 @@ pub use async_openai::types::chat::{ ChatCompletionRequestToolMessageContentPart, ChatCompletionResponseMessageAudio, ChatCompletionTokenLogprob, + Choice, CompletionFinishReason, CompletionTokensDetails, CompletionUsage, @@ -695,9 +696,6 @@ pub struct ChatChoice { pub index: u32, pub message: ChatCompletionResponseMessage, pub finish_reason: Option, - /// Matched stop condition from the backend. - #[serde(skip_serializing_if = "Option::is_none")] - pub stop_reason: Option, pub logprobs: Option, } @@ -746,19 +744,12 @@ pub struct ChatCompletionStreamResponseDeltaFunctionCall { pub arguments: Option, } -/// Streaming chat choice with stop reason support. -/// -/// Extends upstream `ChatChoiceStream` with: -/// - `stop_reason`: the matched stop sequence (string) or stop token ID (integer) -/// reported by inference backends +/// Streaming chat choice. #[derive(Debug, Deserialize, Clone, PartialEq, Serialize)] pub struct ChatChoiceStream { pub index: u32, pub delta: ChatCompletionStreamResponseDelta, pub finish_reason: Option, - /// Matched stop condition from the backend. - #[serde(skip_serializing_if = "Option::is_none")] - pub stop_reason: Option, pub logprobs: Option, } diff --git a/lib/protocols/src/types/completion.rs b/lib/protocols/src/types/completion.rs index cf55a51c9383..a14683dffedc 100644 --- a/lib/protocols/src/types/completion.rs +++ b/lib/protocols/src/types/completion.rs @@ -13,47 +13,10 @@ use serde::{Deserialize, Serialize}; use crate::error::OpenAIError; -use super::{ - ChatCompletionStreamOptions, CompletionFinishReason, CompletionUsage, Logprobs, Prompt, Stop, - StopReason, -}; +use super::{ChatCompletionStreamOptions, Prompt, Stop}; -/// Completion choice with inference-serving extensions. -/// -/// Extends upstream `Choice` with: -/// - `stop_reason`: the matched stop sequence, token ID, or token ID sequence -/// reported by inference backends -#[derive(Debug, Deserialize, Serialize, Clone, PartialEq)] -pub struct Choice { - pub text: String, - pub index: u32, - #[serde(skip_serializing_if = "Option::is_none")] - pub logprobs: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub finish_reason: Option, - /// Matched stop condition from the backend. - #[serde(skip_serializing_if = "Option::is_none")] - pub stop_reason: Option, -} - -/// Non-streaming or streaming text completion response. -#[derive(Debug, Deserialize, Clone, PartialEq, Serialize)] -pub struct CreateCompletionResponse { - /// A unique identifier for the completion. - pub id: String, - pub choices: Vec, - /// The Unix timestamp (in seconds) of when the completion was created. - pub created: u32, - - /// The model used for completion. - pub model: String, - /// This fingerprint represents the backend configuration that the model runs with. - pub system_fingerprint: Option, - - /// The object type, which is always "text_completion". - pub object: String, - pub usage: Option, -} +// Re-export response type from upstream (identical) +pub use async_openai::types::completions::CreateCompletionResponse; /// Custom deserializer for the echo parameter that only accepts booleans. /// Rejects integers and strings with clear error messages. @@ -209,18 +172,19 @@ mod tests { } #[test] - fn completion_choice_serializes_stop_reason() { + fn completion_choice_serializes_openai_shape() { + use crate::types::{Choice, CompletionFinishReason}; + let choice = Choice { text: "hello".to_string(), index: 0, logprobs: None, finish_reason: Some(CompletionFinishReason::Stop), - stop_reason: Some(StopReason::String("END".to_string())), }; let value = serde_json::to_value(choice).expect("serialize choice"); assert_eq!(value["finish_reason"], "stop"); - assert_eq!(value["stop_reason"], "END"); + assert_eq!(value["text"], "hello"); } } From 2dc32157f9a65f6f583499d7445d8cab08a8ce12 Mon Sep 17 00:00:00 2001 From: William Arnold Date: Thu, 7 May 2026 07:50:04 -0700 Subject: [PATCH 15/20] chore(openai): drop vllm logprob changes --- .../src/dynamo/frontend/vllm_processor.py | 5 --- components/src/dynamo/vllm/handlers.py | 33 +++++-------------- docs/components/frontend/nvext.md | 2 +- lib/llm/src/protocols/openai/nvext.rs | 3 ++ 4 files changed, 12 insertions(+), 31 deletions(-) diff --git a/components/src/dynamo/frontend/vllm_processor.py b/components/src/dynamo/frontend/vllm_processor.py index abb61b2826c3..8962bfd02ec5 100644 --- a/components/src/dynamo/frontend/vllm_processor.py +++ b/components/src/dynamo/frontend/vllm_processor.py @@ -49,7 +49,6 @@ extract_mm_urls, handle_engine_error, make_internal_error, - nvext_extra_field_requested, random_uuid, ) @@ -646,10 +645,6 @@ async def _generate_and_stream( } if usage := engine_response.get("completion_usage"): dynamo_out["usage"] = usage - if stop_reason is not None and nvext_extra_field_requested( - request, "stop_reason" - ): - dynamo_out["nvext"] = {"stop_reason": stop_reason} yield dynamo_out _nvtx.end_range(rng_stream) diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index 6e4544bf6e0c..2d64f860b79e 100644 --- a/components/src/dynamo/vllm/handlers.py +++ b/components/src/dynamo/vllm/handlers.py @@ -1851,10 +1851,7 @@ def _build_completion_usage( @staticmethod def _extract_logprobs( - output, - num_output_tokens_so_far: int, - tokenizer=None, - return_tokens_as_token_ids: bool = False, + output, num_output_tokens_so_far: int, tokenizer=None ) -> tuple[list[float] | None, list[list[dict]] | None]: """ Extract logprobs from vLLM CompletionOutput for new tokens. @@ -1896,16 +1893,12 @@ def _extract_logprobs( # Build top_logprobs list for this token position token_top_logprobs = [] for tok_id, logprob_info in token_logprobs_dict.items(): - token_str: str | None - if return_tokens_as_token_ids: - token_str = f"token_id:{tok_id}" - else: - token_str = getattr(logprob_info, "decoded_token", None) - if not token_str and tokenizer: - try: - token_str = tokenizer.decode([tok_id]) - except Exception: - token_str = None + token_str = getattr(logprob_info, "decoded_token", None) + if not token_str and tokenizer: + try: + token_str = tokenizer.decode([tok_id]) + except Exception: + token_str = None token_top_logprobs.append( { "rank": ( @@ -1967,7 +1960,6 @@ async def generate_tokens( embedding_sequence_length=None, trace_headers=None, priority=0, - return_tokens_as_token_ids=False, ): try: # Log LoRA usage for this generation (debug level to avoid log spam) @@ -2017,10 +2009,7 @@ async def generate_tokens( # Extract logprobs for new tokens if available tokenizer = getattr(self.engine_client, "tokenizer", None) log_probs, top_logprobs = self._extract_logprobs( - output, - previous_total_toks, - tokenizer=tokenizer, - return_tokens_as_token_ids=return_tokens_as_token_ids, + output, previous_total_toks, tokenizer=tokenizer ) if log_probs is not None: out["log_probs"] = log_probs @@ -2266,11 +2255,6 @@ async def _generate_token_mode(self, request, context, request_id): trace_headers = build_trace_headers(context) - output_options = request.get("output_options", {}) - return_tokens_as_token_ids = bool( - output_options.get("return_tokens_as_token_ids") - ) - # In disagg decode mode, defer engine_client.abort() until the first # token so we don't abort while a NIXL KV transfer is still in flight # on the decode worker (which can crash EngineCore). The guard's @@ -2293,7 +2277,6 @@ async def _generate_token_mode(self, request, context, request_id): embedding_sequence_length=embedding_sequence_length, trace_headers=trace_headers, priority=priority, - return_tokens_as_token_ids=return_tokens_as_token_ids, ): if abort_guard is not None: abort_guard.signal_first_token() diff --git a/docs/components/frontend/nvext.md b/docs/components/frontend/nvext.md index e697c66f2b0a..7acd64f99d3f 100644 --- a/docs/components/frontend/nvext.md +++ b/docs/components/frontend/nvext.md @@ -202,7 +202,7 @@ When the client requests response metadata via `extra_fields`, the response incl | `timing` | `extra_fields: ["timing"]` | Per-request timing information (TTFT, ITL, queue time, etc.). | | `routed_experts` | `extra_fields: ["routed_experts"]` | Routed expert capture payload returned by SGLang-backed requests. | | `engine_data` | `extra_fields: ["engine_data"]` | Opaque backend-provided engine metadata. | -| `stop_reason` | `extra_fields: ["stop_reason"]` | Backend-specific matched stop condition, returned under `nvext` because it is not part of the OpenAI completions schema. | +| `stop_reason` | `extra_fields: ["stop_reason"]` | Backend-specific matched stop condition, returned under `nvext` because it is not part of the OpenAI completions schema. Dynamo currently serves this as a response-level field for single-choice requests; supporting `n > 1` will require an indexed per-choice shape. | | `token_ids` | Automatic (GAIE Stage 1) | Tokenized prompt for reuse in Stage 2 query-only mode. | ### Example response `nvext` diff --git a/lib/llm/src/protocols/openai/nvext.rs b/lib/llm/src/protocols/openai/nvext.rs index 8459b3085699..a088fd522b11 100644 --- a/lib/llm/src/protocols/openai/nvext.rs +++ b/lib/llm/src/protocols/openai/nvext.rs @@ -124,6 +124,9 @@ pub struct NvExtResponse { /// Backend-specific matched stop condition. This is not part of the /// OpenAI response schema, so it is only returned under nvext when requested. + /// + /// This is response-level for Dynamo's current single-choice serving paths. + /// If `n > 1` is supported here, this needs an indexed/per-choice shape. #[serde(skip_serializing_if = "Option::is_none")] pub stop_reason: Option, } From 0397befd801c8a6895724554c659ca3a39b20ae0 Mon Sep 17 00:00:00 2001 From: William Arnold Date: Thu, 7 May 2026 08:34:47 -0700 Subject: [PATCH 16/20] docs(sglang): link top logprobs upstream fix --- .../src/dynamo/sglang/request_handlers/llm/decode_handler.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py index 5d112245e349..ebf62b0be7ab 100644 --- a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py +++ b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py @@ -22,7 +22,8 @@ # Escape hatch: set to "1" (or any truthy value) to allow top_logprobs_num >= 1. # Default-off because SGLang's tokenizer manager detokenizes top-k tokens # per-position serially (O(N) per generated token), causing severe latency -# degradation. Flip once upstream batches detokenize_top_logprobs_tokens. +# degradation. Flip once upstream lands batched top-logprob detokenization: +# https://github.com/sgl-project/sglang/pull/24447 _ALLOW_TOP_LOGPROBS_ENV = "DYN_SGL_ALLOW_TOP_LOGPROBS" _TOP_LOGPROBS_UNSUPPORTED_MSG = ( @@ -30,7 +31,7 @@ "an O(N) per-position detokenization in the upstream sglang tokenizer " "manager. Use logprobs=0 for chosen-token logprobs, or set " "DYN_SGL_ALLOW_TOP_LOGPROBS=1 to override at your own risk. " - "Track the upstream fix at https://github.com/sgl-project/sglang/issues/." + "Track the upstream fix at https://github.com/sgl-project/sglang/pull/24447." ) From ce704c1eed41a9b3f132511819c8839d5a9a7bb8 Mon Sep 17 00:00:00 2001 From: William Arnold Date: Thu, 7 May 2026 11:28:26 -0700 Subject: [PATCH 17/20] feat(openai): accept token id stop arrays --- .../src/dynamo/frontend/sglang_processor.py | 8 ++- .../tests/test_sglang_processor_unit.py | 24 ++++++++ components/src/dynamo/sglang/protocol.py | 1 + docs/components/frontend/nvext.md | 5 ++ lib/llm/src/backend.rs | 35 ++++++++++-- lib/llm/src/protocols/common.rs | 5 ++ lib/llm/src/protocols/openai.rs | 11 ++++ .../src/protocols/openai/chat_completions.rs | 9 +-- lib/llm/src/protocols/openai/completions.rs | 41 +++++++++++-- lib/llm/src/protocols/openai/validate.rs | 12 ++++ lib/llm/src/protocols/unified.rs | 17 ++++-- lib/llm/tests/test_stop_behavior.rs | 19 +++++++ lib/protocols/src/types/chat.rs | 57 ++++++++++++++++++- lib/protocols/src/types/completion.rs | 27 +++++++++ 14 files changed, 248 insertions(+), 23 deletions(-) diff --git a/components/src/dynamo/frontend/sglang_processor.py b/components/src/dynamo/frontend/sglang_processor.py index e5678c992794..26c179e58294 100644 --- a/components/src/dynamo/frontend/sglang_processor.py +++ b/components/src/dynamo/frontend/sglang_processor.py @@ -203,13 +203,17 @@ def _build_dynamo_preproc( max_tokens = request.get("max_completion_tokens") or request.get("max_tokens") stop = request.get("stop") + stop_token_ids = request.get("stop_token_ids", []) if isinstance(stop, str): stop = [stop] + elif isinstance(stop, list) and all( + isinstance(item, int) and not isinstance(item, bool) for item in stop + ): + stop_token_ids = [*stop_token_ids, *stop] + stop = [] elif stop is None: stop = [] - stop_token_ids = request.get("stop_token_ids", []) - # Handle logprobs logprobs_val = None logprobs = request.get("logprobs") diff --git a/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py b/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py index 26e041a2a924..6c5b4a3b0f63 100644 --- a/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py +++ b/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py @@ -259,6 +259,30 @@ def test_model_name_and_token_ids(self): assert result["model"] == "my-model" assert result["token_ids"] == [10, 20, 30] + def test_stop_token_id_array_maps_to_stop_token_ids(self): + """Integer stop arrays are token-id stops, not string stops.""" + result = _build_dynamo_preproc( + {"model": "test", "stop": [576]}, + [1], + "test", + None, + ) + + assert result["stop_conditions"]["stop"] == [] + assert result["stop_conditions"]["stop_token_ids"] == [576] + + def test_token_id_display_string_remains_string_stop(self): + """token_id:N strings are output display strings, not token-id stops.""" + result = _build_dynamo_preproc( + {"model": "test", "stop": ["token_id:576"]}, + [1], + "test", + None, + ) + + assert result["stop_conditions"]["stop"] == ["token_id:576"] + assert result["stop_conditions"]["stop_token_ids"] == [] + # --------------------------------------------------------------------------- # _map_finish_reason diff --git a/components/src/dynamo/sglang/protocol.py b/components/src/dynamo/sglang/protocol.py index 2c4e713791a3..32cbf11205f1 100644 --- a/components/src/dynamo/sglang/protocol.py +++ b/components/src/dynamo/sglang/protocol.py @@ -19,6 +19,7 @@ class StopConditions(BaseModel): max_tokens: Optional[int] = None stop: Optional[List[str]] = None + stop_token_ids: Optional[List[TokenIdType]] = None stop_token_ids_hidden: Optional[List[TokenIdType]] = None min_tokens: Optional[int] = None ignore_eos: Optional[bool] = None diff --git a/docs/components/frontend/nvext.md b/docs/components/frontend/nvext.md index a0777f3449a6..3f1c193c1ef8 100644 --- a/docs/components/frontend/nvext.md +++ b/docs/components/frontend/nvext.md @@ -48,6 +48,11 @@ Related root-level Dynamo output option: |-------|------|---------|-------------|-------------| | `return_tokens_as_token_ids` | `bool` | `false` | Response builder | Formats logprob token strings as `token_id:` instead of decoded text. | +`return_tokens_as_token_ids` only changes returned logprob token display. To stop on +token IDs, pass integer IDs in the normal `stop` array, for example +`"stop": [576]`. Strings such as `"token_id:576"` remain literal string stop +sequences and are not parsed as token IDs. + ### Header Overrides Routing fields can also be set via HTTP headers, which take priority over `nvext` values: diff --git a/lib/llm/src/backend.rs b/lib/llm/src/backend.rs index 94a609a34275..2df3cdf374c9 100644 --- a/lib/llm/src/backend.rs +++ b/lib/llm/src/backend.rs @@ -262,6 +262,13 @@ impl // System EOS token - no stop_reason (user didn't request this stop) (Some(FinishReason::Stop), None) } + Some(StopTrigger::UserStopTokenDetected(token_id)) => { + // User-provided token stop (hidden from output) + ( + Some(FinishReason::Stop), + Some(StopReason::Int((*token_id).into())), + ) + } Some(StopTrigger::HiddenStopSequenceDetected(seq)) => { // User-provided stop sequence (hidden from output) ( @@ -435,6 +442,10 @@ pub struct Decoder { // minimum number of tokens have been generated hidden_stop_ids: HashSet, + // user-provided token stop IDs, kept separate from system/EOS stop IDs so + // stop_reason can report user-triggered token stops without reporting EOS. + user_stop_ids: HashSet, + // text sequences that if found in the response will trigger a stop condition after the // minimum number of tokens have been generated (excluded from output) hidden_stop_sequences: Vec, @@ -461,6 +472,7 @@ pub struct Decoder { pub enum StopTrigger { MaxTokensLimit, HiddenStopTokenDetected(TokenIdType), + UserStopTokenDetected(TokenIdType), HiddenStopSequenceDetected(String), VisibleStopSequenceDetected(String), } @@ -501,12 +513,19 @@ impl Decoder { include_stop_str_in_output: bool, tracker: Option>, ) -> Self { - let hidden_stop_ids: HashSet = stop_condition + let user_stop_ids: HashSet = stop_condition + .stop_token_ids + .unwrap_or_default() + .iter() + .copied() + .collect(); + let system_stop_ids: HashSet = stop_condition .stop_token_ids_hidden .unwrap_or_default() .iter() .copied() .collect(); + let hidden_stop_ids = user_stop_ids.union(&system_stop_ids).copied().collect(); // Categorize stop sequences based on include_stop_str_in_output: // - When true: user-provided stop sequences go to visible (included in output) @@ -529,6 +548,7 @@ impl Decoder { decode_stream, tracker, hidden_stop_ids, + user_stop_ids, hidden_stop_sequences, visible_stop_sequences, min_tokens: stop_condition.min_tokens.unwrap_or(0), @@ -565,12 +585,15 @@ impl Decoder { return Ok(StepResult::ok(token)); } - // check for hidden stop tokens - eos takes precedence + // Check token stops. User-provided token IDs take precedence over + // system/EOS IDs so stop_reason only reports stops the caller requested. if self.hidden_stop_ids.contains(&token_id) { - return Ok(StepResult::with_stop_trigger( - None, - StopTrigger::HiddenStopTokenDetected(token_id), - )); + let trigger = if self.user_stop_ids.contains(&token_id) { + StopTrigger::UserStopTokenDetected(token_id) + } else { + StopTrigger::HiddenStopTokenDetected(token_id) + }; + return Ok(StepResult::with_stop_trigger(None, trigger)); } // check stop sequences - the jail will always hold at least the largest stop sequence diff --git a/lib/llm/src/protocols/common.rs b/lib/llm/src/protocols/common.rs index 031d34a7f3bc..257bcc170bd2 100644 --- a/lib/llm/src/protocols/common.rs +++ b/lib/llm/src/protocols/common.rs @@ -240,6 +240,10 @@ pub struct StopConditions { /// List of tokens that stop the generation when they are /// generated. The returned output will NOT contain the stop tokens. + pub stop_token_ids: Option>, + + /// List of hidden/system tokens that stop generation when they are + /// generated. The returned output will NOT contain the stop tokens. pub stop_token_ids_hidden: Option>, /// The minimum number of tokens to generate @@ -260,6 +264,7 @@ impl StopConditions { pub fn apply_ignore_eos(&mut self) { if self.ignore_eos.unwrap_or(false) { self.stop = None; + self.stop_token_ids = None; self.stop_token_ids_hidden = None; } } diff --git a/lib/llm/src/protocols/openai.rs b/lib/llm/src/protocols/openai.rs index 466a20a8f83f..0234ccc6f59e 100644 --- a/lib/llm/src/protocols/openai.rs +++ b/lib/llm/src/protocols/openai.rs @@ -63,6 +63,10 @@ pub(crate) trait OpenAIStopConditionsProvider { fn get_stop(&self) -> Option>; + fn get_stop_token_ids(&self) -> Option> { + None + } + fn nvext(&self) -> Option<&nvext::NvExt>; /// Get ignore_eos from CommonExt if the type supports it. @@ -180,6 +184,7 @@ impl StopConditionsProvider for T { let max_tokens = self.get_max_tokens(); let min_tokens = self.get_min_tokens(); let stop = self.get_stop(); + let stop_token_ids = self.get_stop_token_ids(); let max_thinking_tokens = self.get_max_thinking_tokens(); if let Some(stop) = &stop @@ -187,6 +192,11 @@ impl StopConditionsProvider for T { { anyhow::bail!("stop conditions must be less than 4") } + if let Some(stop_token_ids) = &stop_token_ids + && stop_token_ids.len() > 4 + { + anyhow::bail!("stop token IDs must be less than 4") + } // Use the trait method to get ignore_eos, which handles precedence let ignore_eos = self.get_ignore_eos(); @@ -195,6 +205,7 @@ impl StopConditionsProvider for T { max_tokens, min_tokens, stop, + stop_token_ids, stop_token_ids_hidden: None, ignore_eos, max_thinking_tokens, diff --git a/lib/llm/src/protocols/openai/chat_completions.rs b/lib/llm/src/protocols/openai/chat_completions.rs index 4b0a87160958..b91a2eac93af 100644 --- a/lib/llm/src/protocols/openai/chat_completions.rs +++ b/lib/llm/src/protocols/openai/chat_completions.rs @@ -290,10 +290,11 @@ impl OpenAIStopConditionsProvider for NvCreateChatCompletionRequest { /// * `Some(Vec)` if stop conditions are set. /// * `None` if no stop conditions are defined. fn get_stop(&self) -> Option> { - self.inner.stop.as_ref().map(|stop| match stop { - dynamo_protocols::types::Stop::String(s) => vec![s.clone()], - dynamo_protocols::types::Stop::StringArray(arr) => arr.clone(), - }) + self.inner.stop.as_ref().and_then(|stop| stop.strings()) + } + + fn get_stop_token_ids(&self) -> Option> { + self.inner.stop.as_ref().and_then(|stop| stop.token_ids()) } /// Returns a reference to the optional `NvExt` extension, if available. diff --git a/lib/llm/src/protocols/openai/completions.rs b/lib/llm/src/protocols/openai/completions.rs index 98b5084d05d3..baf144bd83ac 100644 --- a/lib/llm/src/protocols/openai/completions.rs +++ b/lib/llm/src/protocols/openai/completions.rs @@ -248,11 +248,11 @@ impl OpenAIStopConditionsProvider for NvCreateCompletionRequest { } fn get_stop(&self) -> Option> { - use dynamo_protocols::types::Stop; - self.inner.stop.as_ref().map(|s| match s { - Stop::String(s) => vec![s.clone()], - Stop::StringArray(arr) => arr.clone(), - }) + self.inner.stop.as_ref().and_then(|stop| stop.strings()) + } + + fn get_stop_token_ids(&self) -> Option> { + self.inner.stop.as_ref().and_then(|stop| stop.token_ids()) } fn nvext(&self) -> Option<&NvExt> { @@ -669,6 +669,7 @@ mod tests { let request: NvCreateCompletionRequest = serde_json::from_value(null_stop).expect("Failed to deserialize request"); assert_eq!(request.get_stop(), None); + assert_eq!(request.get_stop_token_ids(), None); let one_stop = json!({ "model": "test-model", @@ -678,6 +679,7 @@ mod tests { let request: NvCreateCompletionRequest = serde_json::from_value(one_stop).expect("Failed to deserialize request"); assert_eq!(request.get_stop(), Some(vec!["foo".to_string()])); + assert_eq!(request.get_stop_token_ids(), None); let many_stops = json!({ "model": "test-model", @@ -690,5 +692,34 @@ mod tests { request.get_stop(), Some(vec!["foo".to_string(), "bar".to_string()]) ); + assert_eq!(request.get_stop_token_ids(), None); + + let token_id_stop = json!({ + "model": "test-model", + "prompt": [1, 2, 3], + "stop": [576] + }); + let request: NvCreateCompletionRequest = + serde_json::from_value(token_id_stop).expect("Failed to deserialize request"); + assert_eq!(request.get_stop(), None); + assert_eq!(request.get_stop_token_ids(), Some(vec![576])); + + let stop_conditions = request + .extract_stop_conditions() + .expect("extract stop conditions"); + assert_eq!(stop_conditions.stop, None); + assert_eq!(stop_conditions.stop_token_ids, Some(vec![576])); + assert_eq!(stop_conditions.stop_token_ids_hidden, None); + + let token_id_display_string_stop = json!({ + "model": "test-model", + "prompt": [1, 2, 3], + "stop": ["token_id:576"] + }); + let request: NvCreateCompletionRequest = + serde_json::from_value(token_id_display_string_stop) + .expect("Failed to deserialize request"); + assert_eq!(request.get_stop(), Some(vec!["token_id:576".to_string()])); + assert_eq!(request.get_stop_token_ids(), None); } } diff --git a/lib/llm/src/protocols/openai/validate.rs b/lib/llm/src/protocols/openai/validate.rs index 237e84bc75be..559dd109ac1b 100644 --- a/lib/llm/src/protocols/openai/validate.rs +++ b/lib/llm/src/protocols/openai/validate.rs @@ -380,6 +380,18 @@ pub fn validate_stop(stop: &Option) -> Result<(), } } } + dynamo_protocols::types::Stop::TokenIdArray(token_ids) => { + if token_ids.is_empty() { + anyhow::bail!("Stop token IDs array cannot be empty"); + } + if token_ids.len() > MAX_STOP_SEQUENCES { + anyhow::bail!( + "Maximum of {} stop token IDs allowed, got {}", + MAX_STOP_SEQUENCES, + token_ids.len() + ); + } + } } } Ok(()) diff --git a/lib/llm/src/protocols/unified.rs b/lib/llm/src/protocols/unified.rs index 3380bc040c4e..c748678126e5 100644 --- a/lib/llm/src/protocols/unified.rs +++ b/lib/llm/src/protocols/unified.rs @@ -404,10 +404,19 @@ impl OpenAIStopConditionsProvider for UnifiedRequest { } fn get_stop(&self) -> Option> { - self.inner.inner.stop.as_ref().map(|stop| match stop { - dynamo_protocols::types::Stop::String(s) => vec![s.clone()], - dynamo_protocols::types::Stop::StringArray(arr) => arr.clone(), - }) + self.inner + .inner + .stop + .as_ref() + .and_then(|stop| stop.strings()) + } + + fn get_stop_token_ids(&self) -> Option> { + self.inner + .inner + .stop + .as_ref() + .and_then(|stop| stop.token_ids()) } fn nvext(&self) -> Option<&NvExt> { diff --git a/lib/llm/tests/test_stop_behavior.rs b/lib/llm/tests/test_stop_behavior.rs index 1a45638efd65..e61e2612eec3 100644 --- a/lib/llm/tests/test_stop_behavior.rs +++ b/lib/llm/tests/test_stop_behavior.rs @@ -136,3 +136,22 @@ fn stop_token_priority_over_sequence() { Some(StopTrigger::HiddenStopTokenDetected(id)) if id == STOP )); } + +#[test] +fn user_stop_token_reports_distinct_trigger() { + let tokenizer: Arc = Arc::new(TestTokenizer); + let decode_stream = tokenizers::DecodeStream::new(tokenizer, &[], false); + let stop_conditions = StopConditions { + stop_token_ids: Some(vec![STOP]), + stop_token_ids_hidden: Some(vec![EOS]), + ..Default::default() + }; + let mut decoder = Decoder::new(decode_stream, stop_conditions, false, None); + let result = decoder.process_token_ids(&[HI, STOP]).unwrap(); + + assert_eq!(result.text.as_deref(), Some("hi")); + assert!(matches!( + result.stop_trigger, + Some(StopTrigger::UserStopTokenDetected(id)) if id == STOP + )); +} diff --git a/lib/protocols/src/types/chat.rs b/lib/protocols/src/types/chat.rs index 4fcf8629a058..5fdc4ff15dfc 100644 --- a/lib/protocols/src/types/chat.rs +++ b/lib/protocols/src/types/chat.rs @@ -77,8 +77,61 @@ pub use async_openai::types::chat::{ WebSearchUserLocationType, }; -// Upstream renamed Stop -> StopConfiguration; re-export under old name for compat -pub use async_openai::types::chat::StopConfiguration as Stop; +/// OpenAI stop configuration, with Dynamo's token-id stop extension. +/// +/// The standard OpenAI shape accepts a string or string array. Dynamo also +/// accepts an integer array, e.g. `"stop": [576]`, to express token-id stop +/// conditions for tokenized in/out workflows. Strings like `"token_id:576"` +/// remain ordinary string stops; the `token_id:` format is only an output +/// display format for logprobs. +#[derive(Debug, Deserialize, Serialize, Clone, PartialEq)] +#[serde(untagged)] +pub enum Stop { + String(String), + StringArray(Vec), + TokenIdArray(Vec), +} + +impl Stop { + pub fn strings(&self) -> Option> { + match self { + Stop::String(s) => Some(vec![s.clone()]), + Stop::StringArray(arr) => Some(arr.clone()), + Stop::TokenIdArray(_) => None, + } + } + + pub fn token_ids(&self) -> Option> { + match self { + Stop::TokenIdArray(arr) => Some(arr.clone()), + Stop::String(_) | Stop::StringArray(_) => None, + } + } +} + +impl From for Stop { + fn from(value: String) -> Self { + Stop::String(value) + } +} + +impl From<&str> for Stop { + fn from(value: &str) -> Self { + Stop::String(value.to_string()) + } +} + +impl From> for Stop { + fn from(value: Vec) -> Self { + Stop::StringArray(value) + } +} + +impl From> for Stop { + fn from(value: Vec) -> Self { + Stop::TokenIdArray(value) + } +} // Upstream renamed FinishReason (streaming) -- re-export pub use async_openai::types::chat::FinishReason; diff --git a/lib/protocols/src/types/completion.rs b/lib/protocols/src/types/completion.rs index a14683dffedc..c06d607147bd 100644 --- a/lib/protocols/src/types/completion.rs +++ b/lib/protocols/src/types/completion.rs @@ -187,4 +187,31 @@ mod tests { assert_eq!(value["finish_reason"], "stop"); assert_eq!(value["text"], "hello"); } + + #[test] + fn stop_accepts_token_id_array() { + let json = r#"{"model": "test_model", "prompt": [1, 2, 3], "stop": [576]}"#; + let request: CreateCompletionRequest = serde_json::from_str(json).unwrap(); + + assert_eq!(request.stop, Some(Stop::TokenIdArray(vec![576]))); + } + + #[test] + fn stop_token_id_display_string_remains_string_stop() { + let json = r#"{"model": "test_model", "prompt": [1, 2, 3], "stop": ["token_id:576"]}"#; + let request: CreateCompletionRequest = serde_json::from_str(json).unwrap(); + + assert_eq!( + request.stop, + Some(Stop::StringArray(vec!["token_id:576".to_string()])) + ); + } + + #[test] + fn stop_rejects_single_token_id() { + let json = r#"{"model": "test_model", "prompt": [1, 2, 3], "stop": 576}"#; + let result: Result = serde_json::from_str(json); + + assert!(result.is_err()); + } } From 4c97a5cf61dc996ac1bcffcdfe01192d37a2d21f Mon Sep 17 00:00:00 2001 From: William Arnold Date: Thu, 7 May 2026 11:46:10 -0700 Subject: [PATCH 18/20] fix(sglang): keep stop reasons in nvext --- .../request_handlers/llm/decode_handler.py | 118 ++++++++++++-- .../tests/test_sglang_decode_handler.py | 149 +++++++++++++++++- lib/protocols/src/types/chat.rs | 36 +++++ lib/protocols/src/types/completion.rs | 14 ++ 4 files changed, 305 insertions(+), 12 deletions(-) diff --git a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py index ebf62b0be7ab..55d9748e8ec9 100644 --- a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py +++ b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py @@ -61,7 +61,64 @@ def _extract_media_urls(mm_data: Dict[str, Any], media_key: str) -> list[str] | return urls or None -def _extract_sglang_stop_reason(finish_reason: Dict[str, Any] | None) -> Any | None: +def _nvext_extra_field_requested(request: Dict[str, Any], field: str) -> bool: + nvext = request.get("nvext") + if not isinstance(nvext, dict): + return False + extra_fields = nvext.get("extra_fields") + if not isinstance(extra_fields, list): + return False + return field in extra_fields + + +def _user_stop_token_ids(request: Dict[str, Any]) -> set[int]: + stop_conditions = request.get("stop_conditions") + if isinstance(stop_conditions, dict): + return { + token_id + for token_id in (stop_conditions.get("stop_token_ids") or []) + if isinstance(token_id, int) and not isinstance(token_id, bool) + } + + stop = request.get("stop") + if isinstance(stop, list) and all( + isinstance(item, int) and not isinstance(item, bool) for item in stop + ): + return set(stop) + + return { + token_id + for token_id in (request.get("stop_token_ids") or []) + if isinstance(token_id, int) and not isinstance(token_id, bool) + } + + +def _openai_stop_sampling_params(request: Dict[str, Any]) -> Dict[str, Any]: + stop = request.get("stop") + if isinstance(stop, str): + return {"stop": stop} + if isinstance(stop, list): + if stop and all( + isinstance(item, int) and not isinstance(item, bool) for item in stop + ): + return {"stop_token_ids": stop} + if stop and all(isinstance(item, str) for item in stop): + return {"stop": stop} + + stop_token_ids = [ + token_id + for token_id in (request.get("stop_token_ids") or []) + if isinstance(token_id, int) and not isinstance(token_id, bool) + ] + if stop_token_ids: + return {"stop_token_ids": stop_token_ids} + return {} + + +def _extract_sglang_stop_reason( + finish_reason: Dict[str, Any] | None, + user_stop_token_ids: set[int] | None = None, +) -> Any | None: """Extract SGLang's matched stop value for Dynamo's stop_reason field.""" if not finish_reason: @@ -70,11 +127,19 @@ def _extract_sglang_stop_reason(finish_reason: Dict[str, Any] | None) -> Any | N matched = finish_reason.get("matched") if isinstance(matched, bool): return None - if isinstance(matched, (str, int)): + if isinstance(matched, str): + return matched + if isinstance(matched, int): + if user_stop_token_ids is not None and matched not in user_stop_token_ids: + return None return matched if isinstance(matched, list) and all( isinstance(item, int) and not isinstance(item, bool) for item in matched ): + if user_stop_token_ids is not None and any( + item not in user_stop_token_ids for item in matched + ): + return None return matched return None @@ -183,6 +248,7 @@ def _build_sampling_params(self, request: Dict[str, Any]) -> Dict[str, Any]: "top_k": request.get("top_k"), "n": request.get("n"), "max_new_tokens": request.get("max_tokens"), + **_openai_stop_sampling_params(request), **self._get_guided_decoding_params(request.get("guided_decoding")), } @@ -361,6 +427,7 @@ async def generate( return_tokens_as_token_ids = bool( output_options.get("return_tokens_as_token_ids") ) + user_stop_token_ids = _user_stop_token_ids(request) lora_path = self._resolve_lora(request) if lora_path: @@ -407,11 +474,19 @@ async def generate( if not self.use_sglang_tokenizer: async for out in self._process_token_stream( - decode, context, return_tokens_as_token_ids + decode, + context, + return_tokens_as_token_ids, + user_stop_token_ids=user_stop_token_ids, ): yield out else: - async for out in self._process_text_stream(decode, context): + async for out in self._process_text_stream( + decode, + context, + request=request, + user_stop_token_ids=user_stop_token_ids, + ): yield out else: # Extract image/video URLs for multimodal requests. SGLang's mm_data_processor @@ -443,11 +518,19 @@ async def generate( ) if not self.use_sglang_tokenizer: async for out in self._process_token_stream( - agg, context, return_tokens_as_token_ids + agg, + context, + return_tokens_as_token_ids, + user_stop_token_ids=user_stop_token_ids, ): yield out else: - async for out in self._process_text_stream(agg, context): + async for out in self._process_text_stream( + agg, + context, + request=request, + user_stop_token_ids=user_stop_token_ids, + ): yield out async def _process_token_stream( @@ -455,6 +538,7 @@ async def _process_token_stream( stream_source: AsyncGenerator[Dict[str, Any], None], context: Context, return_tokens_as_token_ids: bool = False, + user_stop_token_ids: set[int] | None = None, ) -> AsyncGenerator[Dict[str, Any], None]: """Process token-based stream output. @@ -498,7 +582,9 @@ async def _process_token_stream( out["finish_reason"] = normalize_finish_reason( finish_reason["type"] ) - stop_reason = _extract_sglang_stop_reason(finish_reason) + stop_reason = _extract_sglang_stop_reason( + finish_reason, user_stop_token_ids + ) if stop_reason is not None: out["stop_reason"] = stop_reason @@ -558,6 +644,8 @@ async def _process_text_stream( self, stream_source: AsyncGenerator[Dict[str, Any], None], context: Context, + request: Dict[str, Any] | None = None, + user_stop_token_ids: set[int] | None = None, ) -> AsyncGenerator[Dict[str, Any], None]: """Process text-based stream output in OpenAI format. @@ -568,6 +656,7 @@ async def _process_text_stream( Yields: OpenAI-formatted chat completion chunk dicts. """ + request = request or {} # SGLang text chunks are cumulative per choice. Keep independent text # offsets so interleaved n>1 choices do not compute deltas from each # other's previous text. @@ -609,9 +698,9 @@ async def _process_text_stream( "delta": {"role": "assistant", "content": delta}, "finish_reason": finish_reason_type, } - stop_reason = _extract_sglang_stop_reason(finish_reason) - if stop_reason is not None: - choice_data["stop_reason"] = stop_reason + stop_reason = _extract_sglang_stop_reason( + finish_reason, user_stop_token_ids + ) response = { "id": res["meta_info"]["id"], @@ -620,13 +709,20 @@ async def _process_text_stream( "model": self.config.server_args.served_model_name, "object": "chat.completion.chunk", } + response_nvext: dict[str, Any] = {} + if stop_reason is not None and _nvext_extra_field_requested( + request, "stop_reason" + ): + response_nvext["stop_reason"] = stop_reason routed_experts = res["meta_info"].get("routed_experts") if routed_experts is not None: # Base64-encode tensor bytes to match sglang's output format. routed_experts = pybase64.b64encode( routed_experts.numpy().tobytes() ).decode("utf-8") - response["nvext"] = {"routed_experts": routed_experts} + response_nvext["routed_experts"] = routed_experts + if response_nvext: + response["nvext"] = response_nvext if not context.is_stopped(): yield response text_counts_per_choice[index] = next_count diff --git a/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py b/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py index 86d6629d5032..09e83e7287e1 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py +++ b/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py @@ -11,6 +11,8 @@ DecodeWorkerHandler, _extract_media_urls, _extract_sglang_stop_reason, + _openai_stop_sampling_params, + _user_stop_token_ids, ) from dynamo.sglang.request_handlers.multimodal.worker_handler import StreamProcessor @@ -61,6 +63,62 @@ def test_extract_sglang_stop_reason(finish_reason, expected): assert _extract_sglang_stop_reason(finish_reason) == expected +def test_extract_sglang_stop_reason_filters_hidden_token_ids(): + finish_reason = {"type": "stop", "matched": 128001} + + assert _extract_sglang_stop_reason(finish_reason, {576}) is None + assert _extract_sglang_stop_reason(finish_reason, {128001}) == 128001 + + +def test_extract_sglang_stop_reason_filters_hidden_token_id_arrays(): + finish_reason = {"type": "stop", "matched": [128001, 128009]} + + assert _extract_sglang_stop_reason(finish_reason, {128001}) is None + assert _extract_sglang_stop_reason(finish_reason, {128001, 128009}) == [ + 128001, + 128009, + ] + + +def test_user_stop_token_ids_ignores_hidden_ids(): + assert _user_stop_token_ids( + { + "stop_conditions": { + "stop_token_ids": [576], + "stop_token_ids_hidden": [128001], + } + } + ) == {576} + + +def test_user_stop_token_ids_handles_null_fields(): + assert _user_stop_token_ids({"stop_conditions": {"stop_token_ids": None}}) == set() + assert _user_stop_token_ids({"stop_token_ids": None}) == set() + + +def test_user_stop_token_ids_accepts_stop_token_id_array(): + assert _user_stop_token_ids({"stop": [576]}) == {576} + + +def test_user_stop_token_ids_treats_token_id_display_as_string_stop(): + assert _user_stop_token_ids({"stop": ["token_id:576"]}) == set() + + +def test_openai_stop_sampling_params_preserves_string_stops(): + assert _openai_stop_sampling_params({"stop": "END"}) == {"stop": "END"} + assert _openai_stop_sampling_params({"stop": ["END"]}) == {"stop": ["END"]} + assert _openai_stop_sampling_params({"stop": ["token_id:576"]}) == { + "stop": ["token_id:576"] + } + + +def test_openai_stop_sampling_params_maps_token_id_stop_array(): + assert _openai_stop_sampling_params({"stop": [576]}) == {"stop_token_ids": [576]} + assert _openai_stop_sampling_params({"stop_token_ids": [576]}) == { + "stop_token_ids": [576] + } + + def _new_decode_handler(*, use_sglang_tokenizer: bool = False): handler = DecodeWorkerHandler.__new__(DecodeWorkerHandler) handler.use_sglang_tokenizer = use_sglang_tokenizer @@ -105,12 +163,19 @@ def test_build_sampling_params_passes_n_for_sglang_tokenizer_requests(): handler = _new_decode_handler(use_sglang_tokenizer=True) sampling_params = handler._build_sampling_params( - {"temperature": 0.2, "top_p": 0.9, "n": 2, "max_tokens": 8} + { + "temperature": 0.2, + "top_p": 0.9, + "n": 2, + "max_tokens": 8, + "stop": [576], + } ) assert sampling_params["n"] == 2 assert sampling_params["temperature"] == 0.2 assert sampling_params["max_new_tokens"] == 8 + assert sampling_params["stop_token_ids"] == [576] def test_build_logprob_kwargs_allows_chosen_token_logprobs(monkeypatch): @@ -253,6 +318,88 @@ async def test_process_text_stream_tracks_delta_per_choice_index(): ] +@pytest.mark.asyncio +async def test_process_text_stream_stop_reason_uses_response_nvext(): + handler = _new_decode_handler() + + chunks = await _collect( + handler._process_text_stream( + _stream( + [ + { + "index": 0, + "text": "Hello", + "meta_info": { + "id": "request-1", + "finish_reason": {"type": "stop", "matched": "END"}, + }, + } + ] + ), + _Context(), + request={"nvext": {"extra_fields": ["stop_reason"]}}, + ) + ) + + assert "stop_reason" not in chunks[0]["choices"][0] + assert chunks[0]["nvext"]["stop_reason"] == "END" + + +@pytest.mark.asyncio +async def test_process_text_stream_stop_reason_requires_nvext_extra_field(): + handler = _new_decode_handler() + + chunks = await _collect( + handler._process_text_stream( + _stream( + [ + { + "index": 0, + "text": "Hello", + "meta_info": { + "id": "request-1", + "finish_reason": {"type": "stop", "matched": "END"}, + }, + } + ] + ), + _Context(), + ) + ) + + assert "stop_reason" not in chunks[0]["choices"][0] + assert "nvext" not in chunks[0] + + +@pytest.mark.asyncio +async def test_process_token_stream_suppresses_hidden_stop_token_reason(): + handler = _new_decode_handler() + + chunks = await _collect( + handler._process_token_stream( + _stream( + [ + { + "index": 0, + "output_ids": [128001], + "meta_info": { + "id": "request-1", + "finish_reason": {"type": "stop", "matched": 128001}, + "prompt_tokens": 1, + "completion_tokens": 1, + "cached_tokens": None, + }, + } + ] + ), + _Context(), + user_stop_token_ids={576}, + ) + ) + + assert "stop_reason" not in chunks[0] + + @pytest.mark.asyncio async def test_multimodal_stream_keeps_reading_after_one_choice_finishes(): chunks = await _collect( diff --git a/lib/protocols/src/types/chat.rs b/lib/protocols/src/types/chat.rs index 5fdc4ff15dfc..d83f87efc1d0 100644 --- a/lib/protocols/src/types/chat.rs +++ b/lib/protocols/src/types/chat.rs @@ -133,6 +133,17 @@ impl From> for Stop { } } +impl From for Stop { + fn from(value: async_openai::types::chat::StopConfiguration) -> Self { + match value { + async_openai::types::chat::StopConfiguration::String(value) => Stop::String(value), + async_openai::types::chat::StopConfiguration::StringArray(value) => { + Stop::StringArray(value) + } + } + } +} + // Upstream renamed FinishReason (streaming) -- re-export pub use async_openai::types::chat::FinishReason; @@ -823,6 +834,31 @@ pub struct CreateChatCompletionStreamResponse { mod tests { use super::*; + #[test] + fn stop_accepts_token_id_array() { + let stop: Stop = serde_json::from_value(serde_json::json!([576])).unwrap(); + + assert_eq!(stop, Stop::TokenIdArray(vec![576])); + } + + #[test] + fn stop_token_id_display_string_remains_string_stop() { + let stop: Stop = serde_json::from_value(serde_json::json!(["token_id:576"])).unwrap(); + + assert_eq!(stop, Stop::StringArray(vec!["token_id:576".to_string()])); + } + + #[test] + fn stop_converts_from_upstream_stop_configuration() { + let upstream = + async_openai::types::chat::StopConfiguration::StringArray(vec!["END".to_string()]); + + assert_eq!( + Stop::from(upstream), + Stop::StringArray(vec!["END".to_string()]) + ); + } + #[test] fn tool_call_defaults_type_on_deserialize() { let tool_call: ChatCompletionMessageToolCall = serde_json::from_value(serde_json::json!({ diff --git a/lib/protocols/src/types/completion.rs b/lib/protocols/src/types/completion.rs index c06d607147bd..4fb664f979eb 100644 --- a/lib/protocols/src/types/completion.rs +++ b/lib/protocols/src/types/completion.rs @@ -207,6 +207,20 @@ mod tests { ); } + #[test] + fn builder_accepts_upstream_stop_configuration() { + let upstream_stop = async_openai::types::chat::StopConfiguration::String("END".to_string()); + + let request = CreateCompletionRequestArgs::default() + .model("test_model") + .prompt(Prompt::String("hello".to_string())) + .stop(upstream_stop) + .build() + .unwrap(); + + assert_eq!(request.stop, Some(Stop::String("END".to_string()))); + } + #[test] fn stop_rejects_single_token_id() { let json = r#"{"model": "test_model", "prompt": [1, 2, 3], "stop": 576}"#; From eda50cb0e01d332bdb8aa8a23a258f467459c946 Mon Sep 17 00:00:00 2001 From: William Arnold Date: Thu, 7 May 2026 14:34:46 -0700 Subject: [PATCH 19/20] test(openai): encode stop input contract --- .../tests/test_sglang_processor_unit.py | 36 +++++++- .../tests/test_sglang_decode_handler.py | 14 +-- .../src/protocols/openai/chat_completions.rs | 86 ++++++++++++++++++- lib/llm/src/protocols/openai/completions.rs | 26 +++++- lib/protocols/src/types/chat.rs | 29 ++++++- lib/protocols/src/types/completion.rs | 25 +++++- 6 files changed, 200 insertions(+), 16 deletions(-) diff --git a/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py b/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py index 6c5b4a3b0f63..d6c9b289cf53 100644 --- a/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py +++ b/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py @@ -262,17 +262,49 @@ def test_model_name_and_token_ids(self): def test_stop_token_id_array_maps_to_stop_token_ids(self): """Integer stop arrays are token-id stops, not string stops.""" result = _build_dynamo_preproc( - {"model": "test", "stop": [576]}, + {"model": "test", "stop": [32, 34]}, [1], "test", None, ) assert result["stop_conditions"]["stop"] == [] - assert result["stop_conditions"]["stop_token_ids"] == [576] + assert result["stop_conditions"]["stop_token_ids"] == [32, 34] + + def test_string_stops_remain_string_stops(self): + """String stops are forwarded as string stops.""" + result = _build_dynamo_preproc( + {"model": "test", "stop": " The"}, + [1], + "test", + None, + ) + + assert result["stop_conditions"]["stop"] == [" The"] + assert result["stop_conditions"]["stop_token_ids"] == [] + + result = _build_dynamo_preproc( + {"model": "test", "stop": ["A", "B"]}, + [1], + "test", + None, + ) + + assert result["stop_conditions"]["stop"] == ["A", "B"] + assert result["stop_conditions"]["stop_token_ids"] == [] def test_token_id_display_string_remains_string_stop(self): """token_id:N strings are output display strings, not token-id stops.""" + result = _build_dynamo_preproc( + {"model": "test", "stop": "token_id:576"}, + [1], + "test", + None, + ) + + assert result["stop_conditions"]["stop"] == ["token_id:576"] + assert result["stop_conditions"]["stop_token_ids"] == [] + result = _build_dynamo_preproc( {"model": "test", "stop": ["token_id:576"]}, [1], diff --git a/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py b/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py index 09e83e7287e1..aae6a08c8554 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py +++ b/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py @@ -97,7 +97,7 @@ def test_user_stop_token_ids_handles_null_fields(): def test_user_stop_token_ids_accepts_stop_token_id_array(): - assert _user_stop_token_ids({"stop": [576]}) == {576} + assert _user_stop_token_ids({"stop": [32, 34]}) == {32, 34} def test_user_stop_token_ids_treats_token_id_display_as_string_stop(): @@ -113,9 +113,11 @@ def test_openai_stop_sampling_params_preserves_string_stops(): def test_openai_stop_sampling_params_maps_token_id_stop_array(): - assert _openai_stop_sampling_params({"stop": [576]}) == {"stop_token_ids": [576]} - assert _openai_stop_sampling_params({"stop_token_ids": [576]}) == { - "stop_token_ids": [576] + assert _openai_stop_sampling_params({"stop": [32, 34]}) == { + "stop_token_ids": [32, 34] + } + assert _openai_stop_sampling_params({"stop_token_ids": [32, 34]}) == { + "stop_token_ids": [32, 34] } @@ -168,14 +170,14 @@ def test_build_sampling_params_passes_n_for_sglang_tokenizer_requests(): "top_p": 0.9, "n": 2, "max_tokens": 8, - "stop": [576], + "stop": [32, 34], } ) assert sampling_params["n"] == 2 assert sampling_params["temperature"] == 0.2 assert sampling_params["max_new_tokens"] == 8 - assert sampling_params["stop_token_ids"] == [576] + assert sampling_params["stop_token_ids"] == [32, 34] def test_build_logprob_kwargs_allows_chosen_token_logprobs(monkeypatch): diff --git a/lib/llm/src/protocols/openai/chat_completions.rs b/lib/llm/src/protocols/openai/chat_completions.rs index b91a2eac93af..c88cfed8d9b4 100644 --- a/lib/llm/src/protocols/openai/chat_completions.rs +++ b/lib/llm/src/protocols/openai/chat_completions.rs @@ -391,7 +391,8 @@ impl ValidateRequest for NvCreateChatCompletionRequest { #[cfg(test)] mod tests { use super::*; - use crate::protocols::common::OutputOptionsProvider; + use crate::engines::ValidateRequest; + use crate::protocols::common::{OutputOptionsProvider, StopConditionsProvider}; use serde_json::json; #[test] @@ -436,4 +437,87 @@ mod tests { assert_eq!(output_options.skip_special_tokens, Some(skip_value)); } } + + #[test] + fn test_stop_contract() { + let one_stop = json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "Hello"}], + "stop": " The" + }); + let request: NvCreateChatCompletionRequest = + serde_json::from_value(one_stop).expect("Failed to deserialize request"); + assert_eq!(request.get_stop(), Some(vec![" The".to_string()])); + assert_eq!(request.get_stop_token_ids(), None); + + let many_stops = json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "Hello"}], + "stop": ["A", "B"] + }); + let request: NvCreateChatCompletionRequest = + serde_json::from_value(many_stops).expect("Failed to deserialize request"); + assert_eq!( + request.get_stop(), + Some(vec!["A".to_string(), "B".to_string()]) + ); + assert_eq!(request.get_stop_token_ids(), None); + + let token_id_stops = json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "Hello"}], + "stop": [32, 34] + }); + let request: NvCreateChatCompletionRequest = + serde_json::from_value(token_id_stops).expect("Failed to deserialize request"); + assert_eq!(request.get_stop(), None); + assert_eq!(request.get_stop_token_ids(), Some(vec![32, 34])); + + let stop_conditions = request + .extract_stop_conditions() + .expect("extract stop conditions"); + assert_eq!(stop_conditions.stop, None); + assert_eq!(stop_conditions.stop_token_ids, Some(vec![32, 34])); + + let token_id_display_string_stop = json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "Hello"}], + "stop": "token_id:576" + }); + let request: NvCreateChatCompletionRequest = + serde_json::from_value(token_id_display_string_stop) + .expect("Failed to deserialize request"); + assert_eq!(request.get_stop(), Some(vec!["token_id:576".to_string()])); + assert_eq!(request.get_stop_token_ids(), None); + + let token_id_display_string_array_stop = json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "Hello"}], + "stop": ["token_id:576"] + }); + let request: NvCreateChatCompletionRequest = + serde_json::from_value(token_id_display_string_array_stop) + .expect("Failed to deserialize request"); + assert_eq!(request.get_stop(), Some(vec!["token_id:576".to_string()])); + assert_eq!(request.get_stop_token_ids(), None); + + let scalar_token_id_stop = json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "Hello"}], + "stop": 576 + }); + let result: Result = + serde_json::from_value(scalar_token_id_stop); + assert!(result.is_err()); + + let unsupported_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) + .expect("Failed to deserialize request"); + assert!(ValidateRequest::validate(&request).is_err()); + } } diff --git a/lib/llm/src/protocols/openai/completions.rs b/lib/llm/src/protocols/openai/completions.rs index baf144bd83ac..e60890214bb9 100644 --- a/lib/llm/src/protocols/openai/completions.rs +++ b/lib/llm/src/protocols/openai/completions.rs @@ -697,20 +697,31 @@ mod tests { let token_id_stop = json!({ "model": "test-model", "prompt": [1, 2, 3], - "stop": [576] + "stop": [32, 34] }); let request: NvCreateCompletionRequest = serde_json::from_value(token_id_stop).expect("Failed to deserialize request"); assert_eq!(request.get_stop(), None); - assert_eq!(request.get_stop_token_ids(), Some(vec![576])); + assert_eq!(request.get_stop_token_ids(), Some(vec![32, 34])); let stop_conditions = request .extract_stop_conditions() .expect("extract stop conditions"); assert_eq!(stop_conditions.stop, None); - assert_eq!(stop_conditions.stop_token_ids, Some(vec![576])); + assert_eq!(stop_conditions.stop_token_ids, Some(vec![32, 34])); assert_eq!(stop_conditions.stop_token_ids_hidden, None); + let token_id_display_string_scalar_stop = json!({ + "model": "test-model", + "prompt": [1, 2, 3], + "stop": "token_id:576" + }); + let request: NvCreateCompletionRequest = + serde_json::from_value(token_id_display_string_scalar_stop) + .expect("Failed to deserialize request"); + assert_eq!(request.get_stop(), Some(vec!["token_id:576".to_string()])); + assert_eq!(request.get_stop_token_ids(), None); + let token_id_display_string_stop = json!({ "model": "test-model", "prompt": [1, 2, 3], @@ -721,5 +732,14 @@ mod tests { .expect("Failed to deserialize request"); 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!({ + "model": "test-model", + "prompt": [1, 2, 3], + "stop_token_ids": [576] + }); + let request: NvCreateCompletionRequest = serde_json::from_value(unsupported_stop_token_ids) + .expect("Failed to deserialize request"); + assert!(ValidateRequest::validate(&request).is_err()); } } diff --git a/lib/protocols/src/types/chat.rs b/lib/protocols/src/types/chat.rs index d83f87efc1d0..75026c9ff22d 100644 --- a/lib/protocols/src/types/chat.rs +++ b/lib/protocols/src/types/chat.rs @@ -836,18 +836,43 @@ mod tests { #[test] fn stop_accepts_token_id_array() { - let stop: Stop = serde_json::from_value(serde_json::json!([576])).unwrap(); + let stop: Stop = serde_json::from_value(serde_json::json!([32, 34])).unwrap(); - assert_eq!(stop, Stop::TokenIdArray(vec![576])); + assert_eq!(stop, Stop::TokenIdArray(vec![32, 34])); + } + + #[test] + fn stop_accepts_string_and_string_array() { + let stop: Stop = serde_json::from_value(serde_json::json!(" The")).unwrap(); + + assert_eq!(stop, Stop::String(" The".to_string())); + + let stop: Stop = serde_json::from_value(serde_json::json!(["A", "B"])).unwrap(); + + assert_eq!( + stop, + Stop::StringArray(vec!["A".to_string(), "B".to_string()]) + ); } #[test] fn stop_token_id_display_string_remains_string_stop() { + let stop: Stop = serde_json::from_value(serde_json::json!("token_id:576")).unwrap(); + + assert_eq!(stop, Stop::String("token_id:576".to_string())); + let stop: Stop = serde_json::from_value(serde_json::json!(["token_id:576"])).unwrap(); assert_eq!(stop, Stop::StringArray(vec!["token_id:576".to_string()])); } + #[test] + fn stop_rejects_single_token_id() { + let result = serde_json::from_value::(serde_json::json!(576)); + + assert!(result.is_err()); + } + #[test] fn stop_converts_from_upstream_stop_configuration() { let upstream = diff --git a/lib/protocols/src/types/completion.rs b/lib/protocols/src/types/completion.rs index 4fb664f979eb..4c1ecd1ff584 100644 --- a/lib/protocols/src/types/completion.rs +++ b/lib/protocols/src/types/completion.rs @@ -190,14 +190,35 @@ mod tests { #[test] fn stop_accepts_token_id_array() { - let json = r#"{"model": "test_model", "prompt": [1, 2, 3], "stop": [576]}"#; + let json = r#"{"model": "test_model", "prompt": [1, 2, 3], "stop": [32, 34]}"#; let request: CreateCompletionRequest = serde_json::from_str(json).unwrap(); - assert_eq!(request.stop, Some(Stop::TokenIdArray(vec![576]))); + assert_eq!(request.stop, Some(Stop::TokenIdArray(vec![32, 34]))); + } + + #[test] + fn stop_accepts_string_and_string_array() { + let one_stop = r#"{"model": "test_model", "prompt": "hello", "stop": " The"}"#; + let request: CreateCompletionRequest = serde_json::from_str(one_stop).unwrap(); + + assert_eq!(request.stop, Some(Stop::String(" The".to_string()))); + + let many_stops = r#"{"model": "test_model", "prompt": "hello", "stop": ["A", "B"]}"#; + let request: CreateCompletionRequest = serde_json::from_str(many_stops).unwrap(); + + assert_eq!( + request.stop, + Some(Stop::StringArray(vec!["A".to_string(), "B".to_string()])) + ); } #[test] fn stop_token_id_display_string_remains_string_stop() { + let json = r#"{"model": "test_model", "prompt": [1, 2, 3], "stop": "token_id:576"}"#; + let request: CreateCompletionRequest = serde_json::from_str(json).unwrap(); + + assert_eq!(request.stop, Some(Stop::String("token_id:576".to_string()))); + let json = r#"{"model": "test_model", "prompt": [1, 2, 3], "stop": ["token_id:576"]}"#; let request: CreateCompletionRequest = serde_json::from_str(json).unwrap(); From 3d5b5f867893fed9bcad4f241e8121adb9a224ad Mon Sep 17 00:00:00 2001 From: William Arnold Date: Thu, 7 May 2026 14:53:01 -0700 Subject: [PATCH 20/20] fix(openai): update echo choice calls --- lib/llm/src/engines.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/lib/llm/src/engines.rs b/lib/llm/src/engines.rs index 4d59f0bcad63..87530d3d2b2d 100644 --- a/lib/llm/src/engines.rs +++ b/lib/llm/src/engines.rs @@ -183,7 +183,7 @@ impl break; } tokio::time::sleep(*TOKEN_ECHO_DELAY).await; - let response = deltas.create_choice(0, Some(c.to_string()), None, None, None); + let response = deltas.create_choice(0, Some(c.to_string()), None, None); yield Annotated { id: Some(id.to_string()), data: Some(response), @@ -200,7 +200,6 @@ impl None, Some(dynamo_protocols::types::FinishReason::Stop), None, - None, ); yield Annotated { id: Some(id.to_string()),