From 3c551499acbe7219ce9e1d81541539df86b19c9f Mon Sep 17 00:00:00 2001 From: Andrew Chen <48723787+chuenchen309@users.noreply.github.com> Date: Fri, 17 Jul 2026 08:35:03 +0800 Subject: [PATCH] fix(trt_llm): BEI max_num_tokens upgrade is a no-op _bei_specfic_migration logs that it is raising max_num_tokens to the BEI minimum, but rebinds the local name instead of assigning the field: self = self.model_copy(update={"max_num_tokens": BEI_REQUIRED_MAX_NUM_TOKENS}) The copy is discarded when the method returns, so the value stays as the user set it: TrussTRTLLMBuildConfiguration(base_model="encoder", max_num_tokens=1024, ...) WARNING build.max_num_tokens=1024, upgrading to 16384 .max_num_tokens # 1024 Two things show this is a slip rather than intent. The next two lines of the same method assign fields directly and do stick (plugin_configuration.paged_kv_cache = False). And the only other model_copy in the file assigns the result back to a field: self.runtime = self.runtime.model_copy(update={...}) Assign the field. Non-encoder builds are unaffected. Co-Authored-By: Claude Opus 4.8 (1M context) --- truss/base/trt_llm_config.py | 4 +--- truss/tests/trt_llm/test_trt_llm_config.py | 27 ++++++++++++++++++++++ 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/truss/base/trt_llm_config.py b/truss/base/trt_llm_config.py index 2c10adb75..211e04f82 100644 --- a/truss/base/trt_llm_config.py +++ b/truss/base/trt_llm_config.py @@ -382,9 +382,7 @@ def _bei_specfic_migration(self): logger.warning( f"build.max_num_tokens={self.max_num_tokens}, upgrading to {BEI_REQUIRED_MAX_NUM_TOKENS}" ) - self = self.model_copy( - update={"max_num_tokens": BEI_REQUIRED_MAX_NUM_TOKENS} - ) + self.max_num_tokens = BEI_REQUIRED_MAX_NUM_TOKENS # set page_kv_cache and use_paged_context_fmha to false for encoder self.plugin_configuration.paged_kv_cache = False self.plugin_configuration.use_paged_context_fmha = False diff --git a/truss/tests/trt_llm/test_trt_llm_config.py b/truss/tests/trt_llm/test_trt_llm_config.py index e005c22cc..90b7eed5a 100644 --- a/truss/tests/trt_llm/test_trt_llm_config.py +++ b/truss/tests/trt_llm/test_trt_llm_config.py @@ -3,6 +3,7 @@ import pydantic import pytest +from truss.base.constants import BEI_REQUIRED_MAX_NUM_TOKENS from truss.base.trt_llm_config import ( TRTLLMConfiguration, TRTLLMConfigurationV1, @@ -246,3 +247,29 @@ def test_trt_llm_config_additional_fields(trtllm_config_v2): assert config.inference_stack == "v2" assert isinstance(config.build, TrussTRTLLMBuildConfiguration) + + +def test_trt_llm_build_config_encoder_upgrades_max_num_tokens(trtllm_config): + """BEI requires a minimum max_num_tokens, and the encoder path enforces it. + + The upgrade rebound the local `self` instead of assigning the field, so it + logged the upgrade while leaving max_num_tokens at the user's value. + """ + build = copy.deepcopy(trtllm_config["trt_llm"]["build"]) + build["base_model"] = "encoder" + build["max_num_tokens"] = 1024 + + build_config = TrussTRTLLMBuildConfiguration(**build) + + assert build_config.max_num_tokens == BEI_REQUIRED_MAX_NUM_TOKENS + + +def test_trt_llm_build_config_decoder_keeps_max_num_tokens(trtllm_config): + """Only the encoder path upgrades max_num_tokens.""" + build = copy.deepcopy(trtllm_config["trt_llm"]["build"]) + build["base_model"] = "decoder" + build["max_num_tokens"] = 1024 + + build_config = TrussTRTLLMBuildConfiguration(**build) + + assert build_config.max_num_tokens == 1024