diff --git a/lib/sidecar/trtllm/launch/agg.sh b/lib/sidecar/trtllm/launch/agg.sh index e0a9252e2f25..90cad14f2d86 100755 --- a/lib/sidecar/trtllm/launch/agg.sh +++ b/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,6 +114,7 @@ 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[@]}" & @@ -100,6 +122,6 @@ 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 diff --git a/lib/sidecar/trtllm/src/args.rs b/lib/sidecar/trtllm/src/args.rs index b85f04cdc120..5d9ce90b36a6 100644 --- a/lib/sidecar/trtllm/src/args.rs +++ b/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, } diff --git a/lib/sidecar/trtllm/src/client.rs b/lib/sidecar/trtllm/src/client.rs index 5fca4a4b1aca..155b374b0e32 100644 --- a/lib/sidecar/trtllm/src/client.rs +++ b/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, 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> { diff --git a/lib/sidecar/trtllm/src/convert.rs b/lib/sidecar/trtllm/src/convert.rs index 22d673e29140..fc8db3faae1c 100644 --- a/lib/sidecar/trtllm/src/convert.rs +++ b/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, diff --git a/lib/sidecar/trtllm/src/engine.rs b/lib/sidecar/trtllm/src/engine.rs index 5e2fc8aec12a..0daf9983cb54 100644 --- a/lib/sidecar/trtllm/src/engine.rs +++ b/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()) diff --git a/lib/sidecar/trtllm/src/model.rs b/lib/sidecar/trtllm/src/model.rs index b2b3795b62a5..d3e50129a9a0 100644 --- a/lib/sidecar/trtllm/src/model.rs +++ b/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, } diff --git a/lib/sidecar/trtllm/src/tests.rs b/lib/sidecar/trtllm/src/tests.rs index 6da223a631ed..50dc426487e8 100644 --- a/lib/sidecar/trtllm/src/tests.rs +++ b/lib/sidecar/trtllm/src/tests.rs @@ -37,6 +37,16 @@ struct FakeTrtllm { peers: Arc>>, reject: Arc, hang: Arc, + /// `(max_seq_len, max_input_len)` for `GetModelInfo`; unset reports 4096 + /// with `max_input_len` absent. + model_info: Arc>>, +} + +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, ) -> Result, 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, +) -> 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();