fix(trtllm): cherry-pick rc22 disagg request ID compat into feat/aa_glm5.2

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:

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)

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:

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

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<TraceLink>,
     },
+    Terminal {
+        output: Annotated<LLMEngineOutput>,
+    },
 }
 
 fn extract_bootstrap_info(params: &serde_json::Value) -> Option<BootstrapInfo> {
@@ -121,9 +125,14 @@ pub enum PrefillQueryOutcome {
     },
 }
 
-struct PrefillCompletion {
-    result: PrefillResult,
-    worker_link: Option<TraceLink>,
+enum PrefillCompletion {
+    Handoff {
+        result: PrefillResult,
+        worker_link: Option<TraceLink>,
+    },
+    Terminal {
+        output: Annotated<LLMEngineOutput>,
+    },
 }
 
 /// 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 {