fix(trtllm): [cherry-pick 1.5.0] stop a configured --context-length being silently discarded (#14450)

lib/sidecar/trtllm/launch/agg.sh

@@ -39,7 +39,8 @@ while [[ $# -gt 0 ]]; do
             echo "  DYN_HTTP_PORT           Dynamo frontend port (default: 8000)"
             echo "  DYN_SYSTEM_PORT         Dynamo sidecar system port (default: 8081)"
             echo "  TRTLLM_GRPC_PORT        TensorRT-LLM gRPC port (default: 50051)"
-            echo "  TRTLLM_CONTEXT_LENGTH   Registered model context length (default: 4096)"
+            echo "  TRTLLM_CONTEXT_LENGTH   Model context length, applied to both the engine and the"
+            echo "                          sidecar (default: 4096; unset when --max_seq_len is given)"
             exit 0
             ;;
         *)
@@ -62,9 +63,29 @@ trap trtllm_exit_trap EXIT
 
 TRTLLM_PYTHON="${TRTLLM_PYTHON:-python3}"
 TRTLLM_GRPC_PORT="${TRTLLM_GRPC_PORT:-50051}"
-TRTLLM_CONTEXT_LENGTH="${TRTLLM_CONTEXT_LENGTH:-4096}"
 CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}"
 
+# Keep the engine and the sidecar on one number. Started without `--max_seq_len`,
+# TensorRT-LLM reports its `max_input_len` default instead of a context length
+# and the sidecar discards it, so pass the same value to both. When the caller
+# supplies `--max_seq_len`, theirs wins and the sidecar adopts the engine's
+# report rather than overriding it with a default it was never told about.
+TRTLLM_MAX_SEQ_LEN_ARGS=()
+TRTLLM_CONTEXT_LENGTH_ARGS=()
+trtllm_max_seq_len_supplied=0
+for arg in "${EXTRA_ARGS[@]}"; do
+    case "$arg" in
+        --max_seq_len|--max_seq_len=*) trtllm_max_seq_len_supplied=1 ;;
+    esac
+done
+if [[ "$trtllm_max_seq_len_supplied" -eq 0 ]]; then
+    TRTLLM_CONTEXT_LENGTH="${TRTLLM_CONTEXT_LENGTH:-4096}"
+    TRTLLM_MAX_SEQ_LEN_ARGS=(--max_seq_len "$TRTLLM_CONTEXT_LENGTH")
+fi
+if [[ -n "$TRTLLM_CONTEXT_LENGTH" ]]; then
+    TRTLLM_CONTEXT_LENGTH_ARGS=(--context-length "$TRTLLM_CONTEXT_LENGTH")
+fi
+
 # `--grpc` needs `smg-grpc-proto`, which TRT-LLM keeps behind its optional
 # `grpc-smg` extra. Constraint copied from that extra so we resolve what
 # upstream resolves.
@@ -83,7 +104,7 @@ fi
 
 print_launch_banner "Launching TensorRT-LLM Native-gRPC Sidecar (1 GPU)" "$MODEL" "$HTTP_PORT" \
     "TensorRT-LLM gRPC: 127.0.0.1:${TRTLLM_GRPC_PORT}" \
-    "Context length:    ${TRTLLM_CONTEXT_LENGTH}"
+    "Context length:    ${TRTLLM_CONTEXT_LENGTH:-from engine report}"
 
 python3 -m dynamo.frontend &
 
@@ -93,13 +114,14 @@ CUDA_VISIBLE_DEVICES="$CUDA_VISIBLE_DEVICES" \
     --grpc \
     --host 127.0.0.1 \
     --port "$TRTLLM_GRPC_PORT" \
+    "${TRTLLM_MAX_SEQ_LEN_ARGS[@]}" \
     "${TRTLLM_GPU_MEM_ARGS[@]}" \
     "${EXTRA_ARGS[@]}" &
 
 DYN_SYSTEM_PORT="${DYN_SYSTEM_PORT:-8081}" \
     dynamo-trtllm-sidecar \
     --grpc-endpoint "127.0.0.1:${TRTLLM_GRPC_PORT}" \
     --model-path "$MODEL" \
-    --context-length "$TRTLLM_CONTEXT_LENGTH" &
+    "${TRTLLM_CONTEXT_LENGTH_ARGS[@]}" &
 
 wait_any_exit

lib/sidecar/trtllm/src/args.rs

@@ -18,9 +18,14 @@ pub(crate) struct Args {
 
     /// Model maximum sequence length (input + output). Used to register the
     /// context length and to derive a default `max_tokens` when a request omits
-    /// one. TensorRT-LLM's `GetModelInfo` gRPC does not report this on current
-    /// releases (it returns zero), so supply it here; otherwise requests that
-    /// omit `max_tokens` are rejected. See the note in `convert.rs`.
+    /// one. A value supplied here takes precedence over the context length
+    /// TensorRT-LLM's `GetModelInfo` gRPC reports; that report is used only when
+    /// this argument is omitted, and a disagreement is logged at WARN. Supply
+    /// this whenever the engine was started without `--max_seq_len`, because
+    /// TensorRT-LLM then reports its `max_input_len` default instead of a real
+    /// context length and `client::model_info` discards it. With neither
+    /// source, requests that omit `max_tokens` are rejected. See the note in
+    /// `convert.rs`.
     #[arg(long, env = "TRTLLM_CONTEXT_LENGTH")]
     pub context_length: Option<u32>,
 }

lib/sidecar/trtllm/src/client.rs

@@ -56,7 +56,14 @@ impl TrtllmClient {
     }
 
     /// Queries `GetModelInfo` and returns the reported maximum sequence length
-    /// (input + output) as the registration context length, if positive.
+    /// (input + output) as the registration context length, if it is usable.
+    ///
+    /// TensorRT-LLM answers `max_seq_len` with `args.max_seq_len or
+    /// args.max_input_len`, and on the PyTorch path `max_seq_len` is unset
+    /// while `max_input_len` holds its 1024 default. An equal pair is that
+    /// fallback's signature: the number describes the argument defaults rather
+    /// than the model, so it is discarded in favour of `--context-length` or
+    /// the context length the frontend reads from the model itself.
     pub(crate) async fn model_info(&self) -> Result<Option<u32>, DynamoError> {
         let info = tokio::time::timeout(
             RPC_TIMEOUT,
@@ -70,7 +77,20 @@ impl TrtllmClient {
         })?
         .map(tonic::Response::into_inner)
         .map_err(|status| status_to_dynamo("GetModelInfo", status))?;
-        Ok(u32::try_from(info.max_seq_len).ok().filter(|len| *len > 0))
+        let max_seq_len = u32::try_from(info.max_seq_len).ok().filter(|len| *len > 0);
+        let max_input_len = u32::try_from(info.max_input_len)
+            .ok()
+            .filter(|len| *len > 0);
+        if max_seq_len.is_some() && max_seq_len == max_input_len {
+            tracing::warn!(
+                reported_context_length = max_seq_len,
+                "TensorRT-LLM reported max_seq_len == max_input_len, which is how it \
+                 answers when the engine was started without --max_seq_len; ignoring \
+                 the report"
+            );
+            return Ok(None);
+        }
+        Ok(max_seq_len)
     }
 
     pub(crate) async fn abort(&self, request_id: String) -> Result<(), DynamoError> {

lib/sidecar/trtllm/src/convert.rs

@@ -82,13 +82,19 @@ pub(crate) fn build_generate_request(
 // natural value. We mirror the in-process backend's text-only default,
 // `max(1, context_length - prompt_len)` (components/src/dynamo/trtllm
 // `_default_max_tokens`); the sidecar rejects multimodal before dispatch, so
-// `token_ids.len()` is the true prompt length. `context_length` comes from
-// `--context-length` (GetModelInfo returns zero on current releases).
+// `token_ids.len()` is the true prompt length. `context_length` is resolved in
+// `engine::start`, where `--context-length` wins over the `GetModelInfo` report.
+// Only `/v1/chat/completions` and `/v1/responses` reach this fallback: they set
+// `PRESERVE_OMITTED_MAX_TOKENS_CONTEXT_KEY`, which stops the frontend supplying
+// its own default (`preprocessor::omitted_max_tokens_default`). `/v1/completions`
+// is already defaulted by the frontend and never lands here.
 //
 // Remove when https://github.com/NVIDIA/TensorRT-LLM/issues/16549 lands (gRPC
-// `max_tokens` made optional): drop this fallback, the `--context-length` arg,
-// and the context-length plumbing in `engine.rs`, and forward an omitted
-// `max_tokens` as unset.
+// `max_tokens` made optional): drop this fallback and forward an omitted
+// `max_tokens` as unset. Keep `--context-length` and the plumbing in
+// `engine.rs` — they also feed `LlmRegistration.context_length`, which the
+// frontend registers as the served context window and which this fallback does
+// not govern.
 fn max_tokens(
     request: &PreprocessedRequest,
     context_length: Option<u32>,

lib/sidecar/trtllm/src/engine.rs

@@ -139,15 +139,38 @@ impl LLMEngine for TrtllmSidecarEngine {
         let client = TrtllmClient::connect(&self.endpoint, self.transport).await?;
         let connection_count = client.connection_count();
 
-        // Prefer a server-reported context length; fall back to the configured
-        // `--context-length`. GetModelInfo returns zero on current TRT-LLM
-        // releases, so the argument is currently the only source. The resolved
-        // value backs the default-`max_tokens` path in `convert::max_tokens`.
+        // `GetModelInfo` reports the engine's `--max_seq_len`, which is unset by
+        // default; `client::model_info` discards the value TensorRT-LLM
+        // substitutes for it. A configured `--context-length` wins over what
+        // survives that check.
         let mut model = self.model.clone();
-        match client.model_info().await {
-            Ok(Some(context_length)) => model.context_length = Some(context_length),
-            Ok(None) => {}
-            Err(error) => tracing::warn!(%error, "GetModelInfo failed; using --context-length"),
+        let reported = match client.model_info().await {
+            Ok(reported) => reported,
+            Err(error) => {
+                match model.context_length {
+                    Some(configured) => tracing::warn!(
+                        %error,
+                        configured_context_length = configured,
+                        "GetModelInfo failed; using the configured --context-length"
+                    ),
+                    None => tracing::warn!(
+                        %error,
+                        "GetModelInfo failed and no --context-length was configured; \
+                         no context length is available"
+                    ),
+                }
+                None
+            }
+        };
+        match (model.context_length, reported) {
+            (Some(configured), Some(reported)) if configured != reported => tracing::warn!(
+                configured_context_length = configured,
+                engine_context_length = reported,
+                "--context-length disagrees with the context length TensorRT-LLM reported; \
+                 using the configured --context-length"
+            ),
+            (None, Some(reported)) => model.context_length = Some(reported),
+            _ => {}
         }
         if let Some(context_length) = model.context_length {
             let _ = self.context_length.set(context_length);
@@ -160,6 +183,7 @@ impl LLMEngine for TrtllmSidecarEngine {
             endpoint = %self.endpoint,
             connections = connection_count,
             model = %model.source,
+            context_length = ?model.context_length,
             "TensorRT-LLM gRPC is ready"
         );
         Ok(model.engine_config())

lib/sidecar/trtllm/src/model.rs

@@ -11,7 +11,8 @@ pub(crate) struct ConfiguredModel {
     /// HF repo name or local path used for tokenization and templates.
     pub source: String,
     /// Maximum sequence length (input + output), from the `--context-length`
-    /// argument or a server `GetModelInfo` report, if known.
+    /// argument if it was supplied, else from a server `GetModelInfo` report,
+    /// if known.
     pub context_length: Option<u32>,
 }
 

lib/sidecar/trtllm/src/tests.rs

@@ -37,6 +37,16 @@ struct FakeTrtllm {
     peers: Arc<Mutex<Vec<SocketAddr>>>,
     reject: Arc<AtomicBool>,
     hang: Arc<AtomicBool>,
+    /// `(max_seq_len, max_input_len)` for `GetModelInfo`; unset reports 4096
+    /// with `max_input_len` absent.
+    model_info: Arc<Mutex<Option<(i32, i32)>>>,
+}
+
+impl FakeTrtllm {
+    async fn reporting(self, max_seq_len: i32, max_input_len: i32) -> Self {
+        *self.model_info.lock().await = Some((max_seq_len, max_input_len));
+        self
+    }
 }
 
 #[tonic::async_trait]
@@ -147,9 +157,11 @@ impl pb::trtllm_service_server::TrtllmService for FakeTrtllm {
         &self,
         _request: Request<pb::GetModelInfoRequest>,
     ) -> Result<Response<pb::GetModelInfoResponse>, Status> {
+        let (max_seq_len, max_input_len) = self.model_info.lock().await.unwrap_or((4096, 0));
         Ok(Response::new(pb::GetModelInfoResponse {
             model_id: "fake-model".to_string(),
-            max_seq_len: 4096,
+            max_seq_len,
+            max_input_len,
             vocab_size: 32000,
             ..Default::default()
         }))
@@ -256,12 +268,20 @@ fn transport(connections: usize) -> GrpcTransportConfig {
 }
 
 fn engine(endpoint: &str, connections: usize) -> TrtllmSidecarEngine {
+    engine_with_context_length(endpoint, connections, None)
+}
+
+fn engine_with_context_length(
+    endpoint: &str,
+    connections: usize,
+    context_length: Option<u32>,
+) -> TrtllmSidecarEngine {
     TrtllmSidecarEngine::new(
         GrpcEndpoint::parse(endpoint, "--grpc-endpoint").expect("valid test endpoint"),
         transport(connections),
         ConfiguredModel {
             source: "model-source".to_string(),
-            context_length: None,
+            context_length,
         },
     )
 }
@@ -632,6 +652,55 @@ async fn aggregated_generation_streams_delta_then_terminal() {
     );
 }
 
+#[tokio::test]
+async fn configured_context_length_overrides_the_engine_report() {
+    let server = FakeServer::start(FakeTrtllm::default()).await;
+    // The fake's GetModelInfo reports 4096, standing in for a release that
+    // under-reports `max_seq_len`; the operator configured 8192.
+    let engine = engine_with_context_length(&server.endpoint, 1, Some(8192));
+    let config = engine.start(0).await.expect("start");
+    assert_eq!(config.llm.unwrap().context_length, Some(8192));
+
+    let mut omits_max_tokens = request();
+    omits_max_tokens.stop_conditions.max_tokens = None;
+    let outputs = collect(&engine, omits_max_tokens).await;
+    assert_eq!(
+        outputs.last().unwrap().finish_reason,
+        Some(FinishReason::Stop)
+    );
+
+    let requests = server.service.requests.lock().await;
+    let sent = requests.first().expect("recorded request");
+    // 8192 configured minus request()'s three prompt tokens; the reported
+    // 4096 would give 4093.
+    assert_eq!(sent.max_tokens, 8189);
+}
+
+#[tokio::test]
+async fn a_max_seq_len_equal_to_max_input_len_is_not_registered() {
+    // TensorRT-LLM answers `max_seq_len` with `max_input_len` when the engine
+    // was started without `--max_seq_len`, so an equal pair carries no model
+    // information and must not reach registration.
+    let server = FakeServer::start(FakeTrtllm::default().reporting(1024, 1024).await).await;
+    let engine = engine(&server.endpoint, 1);
+    let config = engine.start(0).await.expect("start");
+    // Nothing is registered, so the frontend keeps the context length it read
+    // from the model itself instead of being pinned to 1024. An omitted
+    // `max_tokens` then has no source to derive from and is rejected, which
+    // `omitted_max_tokens_without_context_length_is_rejected` covers.
+    assert_eq!(config.llm.unwrap().context_length, None);
+}
+
+#[tokio::test]
+async fn a_distinct_max_seq_len_is_registered() {
+    // `max_seq_len != max_input_len` means the engine was given an explicit
+    // `--max_seq_len`, so the report is real and is adopted.
+    let server = FakeServer::start(FakeTrtllm::default().reporting(2048, 1024).await).await;
+    let engine = engine(&server.endpoint, 1);
+    let config = engine.start(0).await.expect("start");
+    assert_eq!(config.llm.unwrap().context_length, Some(2048));
+}
+
 #[tokio::test]
 async fn grpc_request_errors_are_propagated() {
     let service = FakeTrtllm::default();