From 37311bad3a5b7a66fb53fcd55e74cbf940440d12 Mon Sep 17 00:00:00 2001 From: Jhao-Ting Chen Date: Fri, 1 May 2026 11:03:19 -0700 Subject: [PATCH] [None][fix] Fix early_stopping type and plumb through Triton ensemble + BLS After PR #13633 plumbed promptIgnoreLength correctly, early_stopping is the only sampling field that has two remaining issues: 1. Type mismatch between Triton config declaration (TYPE_BOOL, 1 byte) and the C++ extraction in getSamplingConfigFromTensors (extractOptionalSingleton, reads 4 bytes). The mismatch works accidentally for {0, 1} (adjacent memory is zero), but cannot represent the executor's documented value 2 ("stop only when all beams emit "), and is undefined behavior in principle. 2. early_stopping is missing entirely from the ensemble + BLS configs in five places, so clients hitting Triton via `ensemble`, `tensorrt_llm_bls`, `multimodal/ensemble`, or `gpt/ensemble` cannot set early_stopping at all. This pre-dates PR #8127. This change fixes both issues: Type fix (BREAKING for clients sending early_stopping as bool): - triton_backend/all_models/inflight_batcher_llm/tensorrt_llm/config.pbtxt: data_type TYPE_BOOL -> TYPE_INT32 - triton_backend/all_models/disaggregated_serving/disaggregated_serving_bls/config.pbtxt: same. Aligns the wire-protocol declaration with executor::SamplingConfig semantics (std::optional accepting 0/1/2). Clients previously sending numpy bool must now send numpy int32; behavior for values 0 and 1 is preserved. Ensemble + BLS plumbing (additive, no compat impact): - triton_backend/all_models/inflight_batcher_llm/ensemble/config.pbtxt: declare optional INT32 early_stopping input + add input_map block forwarding it to the tensorrt_llm step. - triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/config.pbtxt: declare optional INT32 early_stopping input. - triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/lib/decode.py: add early_stopping field to the Request dataclass. - triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/lib/triton_decoder.py: add early_stopping to both input_names lists and the BLS->engine name mapping. - triton_backend/all_models/multimodal/ensemble/config.pbtxt: declare + input_map. - triton_backend/all_models/gpt/ensemble/config.pbtxt: declare + input_map. - triton_backend/all_models/gpt/tensorrt_llm/config.pbtxt: declare (was missing entirely). The Python tensorrt_llm/1/model.py already forwards early_stopping to trtllm.SamplingConfig kwargs; only the wire-protocol declaration was wrong. Verified end-to-end on TinyLlama-1.1B with beam_width=4, prompt "Hello world. Goodbye." for all three Triton entry points: Path 1 (direct tensorrt_llm): early_stopping=0 -> beam_lens=[60, 60, 54, 60] early_stopping=1 -> beam_lens=[3, 0, 2, 1] Path 2 (ensemble): early_stopping=0 -> 60-token continuation of "Hello world. Goodbye..." early_stopping=1 -> "Hello world." (3 tokens) Path 3 (tensorrt_llm_bls): early_stopping=0 -> 60-token continuation of "Hello world. Goodbye..." early_stopping=1 -> "Hello world." (3 tokens) All three paths now honor early_stopping correctly. Signed-off-by: Jhao-Ting Chen --- .../disaggregated_serving_bls/config.pbtxt | 2 +- .../all_models/gpt/ensemble/config.pbtxt | 10 ++++++++ .../all_models/gpt/tensorrt_llm/config.pbtxt | 7 ++++++ .../ensemble/config.pbtxt | 10 ++++++++ .../tensorrt_llm/config.pbtxt | 2 +- .../tensorrt_llm_bls/1/lib/decode.py | 1 + .../tensorrt_llm_bls/1/lib/triton_decoder.py | 25 +++++++++++-------- .../tensorrt_llm_bls/config.pbtxt | 6 +++++ .../multimodal/ensemble/config.pbtxt | 10 ++++++++ .../inflight_batcher_llm/tests/utilsTest.cpp | 8 ++++-- 10 files changed, 66 insertions(+), 15 deletions(-) diff --git a/triton_backend/all_models/disaggregated_serving/disaggregated_serving_bls/config.pbtxt b/triton_backend/all_models/disaggregated_serving/disaggregated_serving_bls/config.pbtxt index 1d98c4be4719..267bfa2efc49 100644 --- a/triton_backend/all_models/disaggregated_serving/disaggregated_serving_bls/config.pbtxt +++ b/triton_backend/all_models/disaggregated_serving/disaggregated_serving_bls/config.pbtxt @@ -171,7 +171,7 @@ input [ }, { name: "early_stopping" - data_type: TYPE_BOOL + data_type: TYPE_INT32 dims: [ 1 ] optional: true }, diff --git a/triton_backend/all_models/gpt/ensemble/config.pbtxt b/triton_backend/all_models/gpt/ensemble/config.pbtxt index a7377d1f4eda..978579f2005b 100755 --- a/triton_backend/all_models/gpt/ensemble/config.pbtxt +++ b/triton_backend/all_models/gpt/ensemble/config.pbtxt @@ -58,6 +58,12 @@ input [ dims: [ 1 ] optional: true }, + { + name: "early_stopping" + data_type: TYPE_INT32 + dims: [ 1 ] + optional: true + }, { name: "prompt_ignore_length" data_type: TYPE_INT32 @@ -187,6 +193,10 @@ ensemble_scheduling { key: "len_penalty" value: "length_penalty" } + input_map { + key: "early_stopping" + value: "early_stopping" + } input_map { key: "prompt_ignore_length" value: "prompt_ignore_length" diff --git a/triton_backend/all_models/gpt/tensorrt_llm/config.pbtxt b/triton_backend/all_models/gpt/tensorrt_llm/config.pbtxt index 195b91e5b6b9..a7abcf997eab 100644 --- a/triton_backend/all_models/gpt/tensorrt_llm/config.pbtxt +++ b/triton_backend/all_models/gpt/tensorrt_llm/config.pbtxt @@ -71,6 +71,13 @@ input [ reshape: { shape: [ ] } optional: true }, + { + name: "early_stopping" + data_type: TYPE_INT32 + dims: [ 1 ] + reshape: { shape: [ ] } + optional: true + }, { name: "prompt_ignore_length" data_type: TYPE_INT32 diff --git a/triton_backend/all_models/inflight_batcher_llm/ensemble/config.pbtxt b/triton_backend/all_models/inflight_batcher_llm/ensemble/config.pbtxt index 8adeaea6661f..0bfa7ab90a44 100644 --- a/triton_backend/all_models/inflight_batcher_llm/ensemble/config.pbtxt +++ b/triton_backend/all_models/inflight_batcher_llm/ensemble/config.pbtxt @@ -104,6 +104,12 @@ input [ dims: [ 1 ] optional: true }, + { + name: "early_stopping" + data_type: TYPE_INT32 + dims: [ 1 ] + optional: true + }, { name: "prompt_ignore_length" data_type: TYPE_INT32 @@ -521,6 +527,10 @@ ensemble_scheduling { key: "len_penalty" value: "length_penalty" } + input_map { + key: "early_stopping" + value: "early_stopping" + } input_map { key: "prompt_ignore_length" value: "prompt_ignore_length" diff --git a/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm/config.pbtxt b/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm/config.pbtxt index 9c808d69cb4c..42a1812b8bf4 100644 --- a/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm/config.pbtxt +++ b/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm/config.pbtxt @@ -214,7 +214,7 @@ input [ }, { name: "early_stopping" - data_type: TYPE_BOOL + data_type: TYPE_INT32 dims: [ 1 ] reshape: { shape: [ ] } optional: true diff --git a/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/lib/decode.py b/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/lib/decode.py index 66bb14299509..ebd2c85c9d1a 100644 --- a/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/lib/decode.py +++ b/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/lib/decode.py @@ -76,6 +76,7 @@ class Request: top_p: Optional[np.ndarray] = None temperature: Optional[np.ndarray] = None length_penalty: Optional[np.ndarray] = None + early_stopping: Optional[np.ndarray] = None prompt_ignore_length: Optional[np.ndarray] = None repetition_penalty: Optional[np.ndarray] = None min_tokens: Optional[np.ndarray] = None diff --git a/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/lib/triton_decoder.py b/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/lib/triton_decoder.py index 55699b2559c2..67f5a5ce4abc 100644 --- a/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/lib/triton_decoder.py +++ b/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/lib/triton_decoder.py @@ -90,10 +90,11 @@ def __init__(self, "image_bytes_input", "image_url_input", "video_bytes_input", "max_tokens", "bad_words", "stop_words", "end_id", "pad_id", "top_k", "top_p", "temperature", "length_penalty", - "prompt_ignore_length", "repetition_penalty", "min_tokens", - "presence_penalty", "frequency_penalty", "seed", "return_log_probs", - "return_context_logits", "return_generation_logits", "beam_width", - "stream", "prompt_embedding_table", "prompt_vocab_size", + "prompt_ignore_length", "early_stopping", "repetition_penalty", + "min_tokens", "presence_penalty", "frequency_penalty", "seed", + "return_log_probs", "return_context_logits", + "return_generation_logits", "beam_width", "stream", + "prompt_embedding_table", "prompt_vocab_size", "prompt_table_extra_id", "embedding_bias_words", "embedding_bias_weights", "num_draft_tokens", "use_draft_logits", "lora_task_id", "lora_weights", "lora_config", @@ -104,13 +105,14 @@ def __init__(self, self.__undo_reshape_whitelist = { "max_tokens", "end_id", "pad_id", "top_k", "top_p", "temperature", - "length_penalty", "prompt_ignore_length", "repetition_penalty", - "min_tokens", "presence_penalty", "frequency_penalty", "seed", - "return_log_probs", "return_context_logits", - "return_generation_logits", "beam_width", "stream", - "prompt_vocab_size", "num_draft_tokens", "use_draft_logits", - "exclude_input_in_output", "return_perf_metrics", "lora_weights", - "lora_config", "lora_task_id", "return_num_input_tokens", + "length_penalty", "prompt_ignore_length", "early_stopping", + "repetition_penalty", "min_tokens", "presence_penalty", + "frequency_penalty", "seed", "return_log_probs", + "return_context_logits", "return_generation_logits", "beam_width", + "stream", "prompt_vocab_size", "num_draft_tokens", + "use_draft_logits", "exclude_input_in_output", + "return_perf_metrics", "lora_weights", "lora_config", + "lora_task_id", "return_num_input_tokens", "return_num_output_tokens" } @@ -456,6 +458,7 @@ def _get_llm_tensors_from_request( "temperature": "temperature", "length_penalty": "len_penalty", "prompt_ignore_length": "prompt_ignore_length", + "early_stopping": "early_stopping", "repetition_penalty": "repetition_penalty", "min_tokens": "min_tokens", "presence_penalty": "presence_penalty", diff --git a/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/config.pbtxt b/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/config.pbtxt index 8dbb66a7161c..8730c3b16209 100644 --- a/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/config.pbtxt +++ b/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/config.pbtxt @@ -134,6 +134,12 @@ input [ dims: [ 1 ] optional: true }, + { + name: "early_stopping" + data_type: TYPE_INT32 + dims: [ 1 ] + optional: true + }, { name: "prompt_ignore_length" data_type: TYPE_INT32 diff --git a/triton_backend/all_models/multimodal/ensemble/config.pbtxt b/triton_backend/all_models/multimodal/ensemble/config.pbtxt index 6a0d41e6babd..db8833ca17e1 100755 --- a/triton_backend/all_models/multimodal/ensemble/config.pbtxt +++ b/triton_backend/all_models/multimodal/ensemble/config.pbtxt @@ -123,6 +123,12 @@ input [ dims: [ 1 ] optional: true }, + { + name: "early_stopping" + data_type: TYPE_INT32 + dims: [ 1 ] + optional: true + }, { name: "prompt_ignore_length" data_type: TYPE_INT32 @@ -518,6 +524,10 @@ ensemble_scheduling { key: "len_penalty" value: "length_penalty" } + input_map { + key: "early_stopping" + value: "early_stopping" + } input_map { key: "prompt_ignore_length" value: "prompt_ignore_length" diff --git a/triton_backend/inflight_batcher_llm/tests/utilsTest.cpp b/triton_backend/inflight_batcher_llm/tests/utilsTest.cpp index fa0b84ccdcbc..58dfbb56fa8b 100644 --- a/triton_backend/inflight_batcher_llm/tests/utilsTest.cpp +++ b/triton_backend/inflight_batcher_llm/tests/utilsTest.cpp @@ -382,7 +382,11 @@ std::optional getRequest( pushTensor(inputsTensors, InputFieldsNames::topPResetIds, nvinfer1::DataType::kINT32, {1}, {25}); pushTensor(inputsTensors, InputFieldsNames::temperature, nvinfer1::DataType::kFLOAT, {1}, {0.3}); pushTensor(inputsTensors, InputFieldsNames::lengthPenalty, nvinfer1::DataType::kFLOAT, {1}, {0.4}); - pushTensor(inputsTensors, InputFieldsNames::earlyStopping, nvinfer1::DataType::kINT32, {1}, {4}); + // early_stopping is a HuggingFace-style tri-state: 0=heuristic, 1=fast, + // 2=stop only when all beams emit . Push 2 to lock in INT32 wire-format + // for the canonical-beam-search mode (was silently truncated under the old + // TYPE_BOOL config declaration; cf. PR fixing the wire-protocol type). + pushTensor(inputsTensors, InputFieldsNames::earlyStopping, nvinfer1::DataType::kINT32, {1}, {2}); pushTensor(inputsTensors, InputFieldsNames::repetitionPenalty, nvinfer1::DataType::kFLOAT, {1}, {0.8}); pushTensor(inputsTensors, InputFieldsNames::minTokens, nvinfer1::DataType::kINT32, {1}, {45}); pushTensor(inputsTensors, InputFieldsNames::beamSearchDiversityRate, nvinfer1::DataType::kFLOAT, {1}, {0.1}); @@ -580,7 +584,7 @@ void checkRequest(tensorrt_llm::executor::Request const& request, EXPECT_EQ(samplingConfig.getTopPResetIds().value(), 25); EXPECT_EQ(samplingConfig.getTemperature().value(), 0.3f); EXPECT_EQ(samplingConfig.getLengthPenalty().value(), 0.4f); - EXPECT_EQ(samplingConfig.getEarlyStopping().value(), 4); + EXPECT_EQ(samplingConfig.getEarlyStopping().value(), 2); EXPECT_EQ(samplingConfig.getRepetitionPenalty().value(), 0.8f); EXPECT_EQ(samplingConfig.getMinTokens().value(), 45); EXPECT_EQ(samplingConfig.getBeamSearchDiversityRate().value(), 0.1f);