diff --git a/components/src/dynamo/trtllm/request_handlers/handler_base.py b/components/src/dynamo/trtllm/request_handlers/handler_base.py index ce726c8fe8bf..781fb80bf30e 100644 --- a/components/src/dynamo/trtllm/request_handlers/handler_base.py +++ b/components/src/dynamo/trtllm/request_handlers/handler_base.py @@ -28,7 +28,6 @@ from tensorrt_llm.executor.result import GenerationResult from tensorrt_llm.executor.utils import RequestError from tensorrt_llm.llmapi import DisaggregatedParams as LlmDisaggregatedParams -from tensorrt_llm.llmapi.disagg_utils import get_global_disagg_request_id from tensorrt_llm.llmapi.llm import SamplingParams from tensorrt_llm.sampling_params import GuidedDecodingParams from tensorrt_llm.scheduling_params import SchedulingParams @@ -60,6 +59,7 @@ from dynamo.trtllm.utils.disagg_utils import ( DisaggregatedParams, DisaggregatedParamsCodec, + get_compatible_global_disagg_request_id, ) from dynamo.trtllm.utils.request_utils import ( apply_stop_conditions_to_sampling_params, @@ -727,7 +727,7 @@ def _setup_disaggregated_params_for_mode( else: disaggregated_params = LlmDisaggregatedParams( request_type="context_only", - disagg_request_id=get_global_disagg_request_id( + disagg_request_id=get_compatible_global_disagg_request_id( self.disagg_machine_id ), ) @@ -736,8 +736,8 @@ def _setup_disaggregated_params_for_mode( # ep_disaggregated_params, so the PYTHON transceiver can track # requests across prefill/decode workers. if disaggregated_params.disagg_request_id is None: - disaggregated_params.disagg_request_id = get_global_disagg_request_id( - self.disagg_machine_id + disaggregated_params.disagg_request_id = ( + get_compatible_global_disagg_request_id(self.disagg_machine_id) ) # AGGREGATED (prefill_and_decode) mode with encoder disaggregation: diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_disagg_request_id.py b/components/src/dynamo/trtllm/tests/test_trtllm_disagg_request_id.py new file mode 100644 index 000000000000..7453d3201ebf --- /dev/null +++ b/components/src/dynamo/trtllm/tests/test_trtllm_disagg_request_id.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from unittest.mock import patch + +import pytest + +from dynamo.trtllm.utils import disagg_utils + +pytestmark = [ + pytest.mark.unit, + pytest.mark.trtllm, + pytest.mark.core, + pytest.mark.pre_merge, + pytest.mark.gpu_1, +] + + +def test_rc22_machine_id_is_split_losslessly(): + with ( + patch.object(disagg_utils, "_TRTLLM_DISAGG_ID_HAS_PROCESS_ID", True), + patch.object( + disagg_utils, + "_trtllm_get_global_disagg_request_id", + return_value=123, + ) as generate, + ): + result = disagg_utils.get_compatible_global_disagg_request_id(1020) + + assert result == 123 + generate.assert_called_once_with(15, 60) + + +def test_rc22_all_dynamo_machine_ids_are_in_range_and_unique(): + pairs = { + divmod(machine_id, disagg_utils._TRTLLM_PROCESS_ID_SPACE) + for machine_id in range(disagg_utils._DYNAMO_DISAGG_MACHINE_ID_SPACE) + } + + assert len(pairs) == disagg_utils._DYNAMO_DISAGG_MACHINE_ID_SPACE + assert all(0 <= node_id < 256 for node_id, _ in pairs) + assert all(0 <= process_id < 64 for _, process_id in pairs) + + +def test_rc21_keeps_legacy_machine_id(): + with ( + patch.object(disagg_utils, "_TRTLLM_DISAGG_ID_HAS_PROCESS_ID", False), + patch.object( + disagg_utils, + "_trtllm_get_global_disagg_request_id", + return_value=456, + ) as generate, + ): + result = disagg_utils.get_compatible_global_disagg_request_id(1020) + + assert result == 456 + generate.assert_called_once_with(1020) + + +@pytest.mark.parametrize("machine_id", [-1, 1021]) +def test_invalid_dynamo_machine_id_fails_fast(machine_id): + with pytest.raises(ValueError, match="machine_id must be in range"): + disagg_utils.get_compatible_global_disagg_request_id(machine_id) diff --git a/components/src/dynamo/trtllm/utils/disagg_utils.py b/components/src/dynamo/trtllm/utils/disagg_utils.py index 375d46f19628..461fc463e9ec 100644 --- a/components/src/dynamo/trtllm/utils/disagg_utils.py +++ b/components/src/dynamo/trtllm/utils/disagg_utils.py @@ -15,9 +15,46 @@ import base64 import dataclasses +import inspect from tensorrt_llm.executor.result import Logprob from tensorrt_llm.llmapi import DisaggregatedParams +from tensorrt_llm.llmapi.disagg_utils import ( + get_global_disagg_request_id as _trtllm_get_global_disagg_request_id, +) + +# Dynamo maps its distributed-runtime connection ID into TRT-LLM's historical +# 10-bit machine-ID space with ``connection_id % 1021``. TRT-LLM rc22 split +# that field into an 8-bit node ID and a 6-bit process ID. Preserve Dynamo's +# existing 10-bit worker slot and encode it losslessly into the new pair. +_TRTLLM_DISAGG_ID_HAS_PROCESS_ID = ( + "process_id" in inspect.signature(_trtllm_get_global_disagg_request_id).parameters +) +_TRTLLM_PROCESS_ID_SPACE = 1 << 6 +_DYNAMO_DISAGG_MACHINE_ID_SPACE = 1021 + + +def get_compatible_global_disagg_request_id(machine_id: int) -> int: + """Generate a global TRT-LLM disaggregation request ID across API versions. + + TRT-LLM <= rc21 accepts a 10-bit ``machine_id``. TRT-LLM >= rc22 accepts + an 8-bit ``node_id`` plus a 6-bit ``process_id``. Dynamo's existing + machine ID is in ``[0, 1021)``; splitting that value with ``divmod(64)`` + produces a unique, in-range pair without changing Dynamo's collision + characteristics. + """ + + if not 0 <= machine_id < _DYNAMO_DISAGG_MACHINE_ID_SPACE: + raise ValueError( + "Dynamo disagg machine_id must be in range " + f"[0, {_DYNAMO_DISAGG_MACHINE_ID_SPACE}), got {machine_id}" + ) + + if _TRTLLM_DISAGG_ID_HAS_PROCESS_ID: + node_id, process_id = divmod(machine_id, _TRTLLM_PROCESS_ID_SPACE) + return _trtllm_get_global_disagg_request_id(node_id, process_id) + + return _trtllm_get_global_disagg_request_id(machine_id) class DisaggregatedParamsCodec: diff --git a/lib/llm/src/kv_router/prefill_router/admission.rs b/lib/llm/src/kv_router/prefill_router/admission.rs index ac178274df2e..9b838701dd2a 100644 --- a/lib/llm/src/kv_router/prefill_router/admission.rs +++ b/lib/llm/src/kv_router/prefill_router/admission.rs @@ -17,7 +17,7 @@ use super::{PrefillCompletion, PrefillError, PrefillRouter}; use crate::{ kv_router::KvPushRouter, protocols::common::{ - llm_backend::{LLMEngineOutput, PreprocessedRequest}, + llm_backend::{FinishReason, LLMEngineOutput, PreprocessedRequest}, timing::RequestTracker, }, session_affinity::{AffinityTarget, SessionAffinityPushRouter}, @@ -110,6 +110,21 @@ impl PrefillRouter { tokio::spawn(async move { while prefill_response.next().await.is_some() {} }); } + // A CTX request that reaches EOS/stop during its one-token prefill step + // is already complete and does not establish a KV-cache handoff. A + // missing finish reason is equivalent to TRT-LLM's "not_finished"; + // Length also requires the normal GEN handoff. + let is_terminal = first_output + .data + .as_ref() + .and_then(|output| output.finish_reason.as_ref()) + .is_some_and(|reason| !matches!(reason, FinishReason::Length)); + if is_terminal { + return Ok(PrefillCompletion::Terminal { + output: first_output, + }); + } + let Some(output) = &first_output.data else { return Err(PrefillError::NoDisaggregatedParams( "Prefill router output has no data field".to_string(), @@ -121,7 +136,7 @@ impl PrefillRouter { )); }; - Ok(PrefillCompletion { + Ok(PrefillCompletion::Handoff { result: crate::protocols::common::preprocessor::PrefillResult { disaggregated_params, prompt_tokens_details, @@ -203,4 +218,67 @@ mod tests { assert!(result.is_err()); assert!(!tracker.record_prefill_complete()); } + + #[tokio::test] + async fn terminal_prefill_without_handoff_is_returned_to_caller() { + let output = LLMEngineOutput { + token_ids: vec![2], + finish_reason: Some(FinishReason::EoS), + ..Default::default() + }; + let result = PrefillRouter::consume_prefill_stream( + prefill_stream(vec![Annotated::from_data(output)]), + None, + ) + .await + .unwrap(); + + let PrefillCompletion::Terminal { output } = result else { + panic!("expected terminal prefill completion"); + }; + assert_eq!( + output.data.and_then(|data| data.finish_reason), + Some(FinishReason::EoS) + ); + } + + #[tokio::test] + async fn length_limited_prefill_still_requires_handoff() { + let output = LLMEngineOutput { + finish_reason: Some(FinishReason::Length), + disaggregated_params: Some(json!({"ctx_request_id": 42})), + ..Default::default() + }; + let result = PrefillRouter::consume_prefill_stream( + prefill_stream(vec![Annotated::from_data(output)]), + None, + ) + .await + .unwrap(); + + let PrefillCompletion::Handoff { result, .. } = result else { + panic!("expected prefill handoff"); + }; + assert_eq!(result.disaggregated_params, json!({"ctx_request_id": 42})); + } + + #[tokio::test] + async fn unfinished_prefill_still_requires_handoff() { + let output = LLMEngineOutput { + finish_reason: None, + disaggregated_params: Some(json!({"ctx_request_id": 42})), + ..Default::default() + }; + let result = PrefillRouter::consume_prefill_stream( + prefill_stream(vec![Annotated::from_data(output)]), + None, + ) + .await + .unwrap(); + + let PrefillCompletion::Handoff { result, .. } = result else { + panic!("expected prefill handoff"); + }; + assert_eq!(result.disaggregated_params, json!({"ctx_request_id": 42})); + } } diff --git a/lib/llm/src/kv_router/prefill_router/mod.rs b/lib/llm/src/kv_router/prefill_router/mod.rs index f30c7b101b67..cea0bc859190 100644 --- a/lib/llm/src/kv_router/prefill_router/mod.rs +++ b/lib/llm/src/kv_router/prefill_router/mod.rs @@ -5,6 +5,7 @@ use std::sync::atomic::AtomicU8; use std::sync::{Arc, OnceLock}; use anyhow::Result; +use futures::stream; use tokio_util::sync::CancellationToken; use uuid::Uuid; @@ -90,6 +91,9 @@ enum PrefillOutcome { worker_id: u64, worker_link: Option, }, + Terminal { + output: Annotated, + }, } fn extract_bootstrap_info(params: &serde_json::Value) -> Option { @@ -121,9 +125,14 @@ pub enum PrefillQueryOutcome { }, } -struct PrefillCompletion { - result: PrefillResult, - worker_link: Option, +enum PrefillCompletion { + Handoff { + result: PrefillResult, + worker_link: Option, + }, + Terminal { + output: Annotated, + }, } /// PrefillRouter is a forward-only operator that sits between Migration and the decode router. @@ -248,19 +257,27 @@ impl drop(prefill_phase_barrier); let completion = Self::consume_prefill_stream(prefill_stream, tracker).await?; - if let Some(bootstrap_info) = - extract_bootstrap_info(&completion.result.disaggregated_params) - { - PrefillOutcome::Bootstrap { - bootstrap_info, - worker_id: prepared.worker_id, - } - } else { - PrefillOutcome::Completed { - result: completion.result, - worker_id: prepared.worker_id, - worker_link: completion.worker_link, + match completion { + PrefillCompletion::Handoff { + result, + worker_link, + } => { + if let Some(bootstrap_info) = + extract_bootstrap_info(&result.disaggregated_params) + { + PrefillOutcome::Bootstrap { + bootstrap_info, + worker_id: prepared.worker_id, + } + } else { + PrefillOutcome::Completed { + result, + worker_id: prepared.worker_id, + worker_link, + } + } } + PrefillCompletion::Terminal { output } => PrefillOutcome::Terminal { output }, } }; Ok((outcome, topology_constraints)) @@ -282,6 +299,23 @@ impl } }; + // A prefill request can terminate before the backend establishes a KV + // handoff (for example, EOS on the one-token context step). Native + // disaggregated backends return that context response directly instead + // of launching a generation-only request with missing handoff IDs. + let outcome = match outcome { + PrefillOutcome::Terminal { mut output } => { + if let Some(data) = output.data.as_mut() { + data.disaggregated_params = None; + } + return Ok(dynamo_runtime::pipeline::ResponseStream::new( + Box::pin(stream::once(async move { output })), + engine_ctx, + )); + } + outcome => outcome, + }; + // NVBugs 5969206: Do NOT abort decode routing when context is killed. // In disaggregated serving, the prefill may have completed and KV transfer // is in flight. Blocking decode here orphans the transfer (no receiver) @@ -322,6 +356,9 @@ impl decode_req.migration_link = worker_link; decode_req.routing_mut().prefill_worker_id = Some(worker_id); } + PrefillOutcome::Terminal { .. } => { + unreachable!("terminal prefill outcomes return before decode routing") + } }; if let Some(topology_constraints) = topology_constraints { diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index 5aabf451b2bc..04dd9fe667e7 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -2729,6 +2729,12 @@ impl OpenAIPreprocessor { let pending = Arc::new(Mutex::new(PendingMetrics::default())); let pending_in = Arc::clone(&pending); + // Buffer raw input text so truncated tool_call blocks can be recovered as + // content if the jail drops them on finish_reason=length (Fix: GLM-5.2 + // glm47_parser allow_eof_recovery=false silently drops incomplete blocks). + let input_text_buf: Arc> = Arc::new(Mutex::new(String::new())); + let input_text_buf_in = Arc::clone(&input_text_buf); + // dynamo `Annotated` -> jail `Annotated` (buffer llm_metrics) let jail_input = stream.map(move |mut a| { if let Some(metrics) = a.data.as_mut().and_then(|nv| nv.llm_metrics.take()) { @@ -2736,6 +2742,17 @@ impl OpenAIPreprocessor { p.chunk_tokens = p.chunk_tokens.saturating_add(metrics.chunk_tokens); p.template = Some(metrics); } + // Accumulate input content for truncation recovery below. + if let Some(data) = &a.data { + for choice in &data.inner.choices { + if let Some(content) = &choice.delta.content { + input_text_buf_in + .lock() + .expect("input text buffer poisoned") + .push_str(content); + } + } + } JailAnnotated { data: a.data.map(|nv| nv.inner), id: a.id, @@ -2745,6 +2762,11 @@ impl OpenAIPreprocessor { } }); + // Track how many bytes the jail has emitted as content so we can compute + // what was silently dropped on a length-truncation finish. + let output_content_len: Arc> = Arc::new(Mutex::new(0)); + let output_content_len_track = Arc::clone(&output_content_len); + // jail `Annotated` -> dynamo `Annotated` (re-attach llm_metrics) jail_apply( tool_call_parser, @@ -2753,7 +2775,7 @@ impl OpenAIPreprocessor { uses_tool_call_structural_tag, jail_input, ) - .map(move |a| { + .flat_map(move |a| { // Stamp the accumulated metrics onto the next emitted data chunk; // data-less/synthesized chunks carry it forward (or `None`). let llm_metrics = a.data.as_ref().and_then(|_| { @@ -2765,7 +2787,7 @@ impl OpenAIPreprocessor { metrics }) }); - Annotated { + let nv_chunk = Annotated { data: a.data.map(|inner| NvCreateChatCompletionStreamResponse { inner, nvext: None, @@ -2775,7 +2797,68 @@ impl OpenAIPreprocessor { event: a.event, comment: a.comment, error: a.error.map(DynamoError::msg), + }; + + // Track output content length and detect truncated tool_call drops. + // When finish_reason=length arrives with no tool_calls in the output + // but the input contained , the jail dropped the block. + // Emit the dropped text as content before the finish chunk so the + // client sees it rather than receiving a silent empty assistant turn. + let mut extra: Option> = None; + if let Some(ref data) = nv_chunk.data { + for choice in &data.inner.choices { + // Track emitted content bytes. + if let Some(content) = &choice.delta.content { + *output_content_len_track + .lock() + .expect("output content len poisoned") += content.len(); + } + // On length finish with no tool_calls, recover dropped text. + if matches!( + choice.finish_reason, + Some(dynamo_protocols::types::FinishReason::Length) + ) && choice.delta.tool_calls.is_none() + { + let input_text = input_text_buf + .lock() + .expect("input text buffer poisoned") + .clone(); + let emitted = *output_content_len_track + .lock() + .expect("output content len poisoned"); + let dropped = if emitted < input_text.len() { + &input_text[emitted..] + } else { + "" + }; + if !dropped.is_empty() && dropped.contains("") { + tracing::debug!( + dropped_len = dropped.len(), + "glm47 streaming: preserving truncated tool_call as content" + ); + // Synthesize a content-only chunk carrying the dropped text. + let mut recovery = nv_chunk.clone(); + if let Some(ref mut rd) = recovery.data { + rd.inner.usage = None; + rd.llm_metrics = None; + for rc in &mut rd.inner.choices { + rc.delta.content = Some(dropped.to_string()); + rc.delta.tool_calls = None; + rc.finish_reason = None; + } + } + extra = Some(recovery); + } + } + } + } + + let mut out = Vec::with_capacity(2); + if let Some(e) = extra { + out.push(e); } + out.push(nv_chunk); + futures::stream::iter(out) }) } diff --git a/lib/llm/src/preprocessor/prompt.rs b/lib/llm/src/preprocessor/prompt.rs index e9efdccdfc2f..3fe9d96eee6f 100644 --- a/lib/llm/src/preprocessor/prompt.rs +++ b/lib/llm/src/preprocessor/prompt.rs @@ -35,13 +35,41 @@ pub trait MediaRequestExt { fn media_io_kwargs(&self) -> Option<&MediaDecoder>; } +/// Parse `tool_calls[*].function.arguments` from JSON string to object in a +/// serialized messages array before handing it to MiniJinja. +/// GLM-5.2's Jinja template iterates arguments with `{% for k, v in _args.items() %}` +/// which requires a dict; the OpenAI wire schema stores arguments as a JSON-object string. +pub(crate) fn normalize_tool_call_arguments(messages_json: &mut serde_json::Value) { + if let Some(msgs) = messages_json.as_array_mut() { + for msg in msgs.iter_mut() { + if let Some(tool_calls) = msg.get_mut("tool_calls").and_then(|v| v.as_array_mut()) { + for tc in tool_calls.iter_mut() { + if let Some(args_str) = tc + .pointer("/function/arguments") + .and_then(|v| v.as_str()) + { + if let Ok(parsed) = serde_json::from_str::(args_str) { + if let Some(fn_obj) = tc.get_mut("function") { + if let Some(obj) = fn_obj.as_object_mut() { + obj.insert("arguments".to_string(), parsed); + } + } + } + } + } + } + } + } +} + impl OAIChatLikeRequest for NvCreateChatCompletionRequest { fn model(&self) -> String { self.inner.model.clone() } fn messages(&self) -> Value { - let messages_json = serde_json::to_value(&self.inner.messages).unwrap(); + let mut messages_json = serde_json::to_value(&self.inner.messages).unwrap(); + normalize_tool_call_arguments(&mut messages_json); Value::from_serialize(&messages_json) } diff --git a/lib/llm/src/protocols/openai/chat_completions/aggregator.rs b/lib/llm/src/protocols/openai/chat_completions/aggregator.rs index 4db807204379..519c8a02de40 100644 --- a/lib/llm/src/protocols/openai/chat_completions/aggregator.rs +++ b/lib/llm/src/protocols/openai/chat_completions/aggregator.rs @@ -419,6 +419,20 @@ impl DeltaAggregator { choice.text = content.unwrap_or_default(); } else if is_harmony_parser(parser) && contains_harmony_protocol(&choice.text) { choice.text = content.unwrap_or_default(); + } else if matches!( + choice.finish_reason, + Some(dynamo_protocols::types::FinishReason::Length) + ) && choice.text.contains("") + { + // The parser dropped a truncated tool_call block (no end fence, hit + // max_tokens). Preserve the raw text as content so the client sees the + // partial output rather than an empty assistant turn, matching TRT-LLM + // behavior where the partial XML is returned as content. + tracing::debug!( + parser, + "preserving truncated tool_call text as content on length finish" + ); + // choice.text already holds the raw text; leave it as-is. } } } diff --git a/lib/llm/src/protocols/unified.rs b/lib/llm/src/protocols/unified.rs index 3a2f79a735f7..1962665b29fe 100644 --- a/lib/llm/src/protocols/unified.rs +++ b/lib/llm/src/protocols/unified.rs @@ -463,7 +463,8 @@ impl OAIChatLikeRequest for UnifiedRequest { } fn messages(&self) -> minijinja::value::Value { - let messages_json = serde_json::to_value(&self.inner.inner.messages).unwrap(); + let mut messages_json = serde_json::to_value(&self.inner.inner.messages).unwrap(); + crate::preprocessor::prompt::normalize_tool_call_arguments(&mut messages_json); minijinja::value::Value::from_serialize(&messages_json) }