tensorrt_llm/_torch/pyexecutor/py_executor_creator.py (3)

1-1: Add NVIDIA Apache-2.0 header (2025).

Per repo guidelines, prepend the standard NVIDIA Apache-2.0 header.

Apply:

+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

---

30-34: Duplicate import of is_mla; remove the redundant one.

Both ._util and .config_utils export is_mla; the latter should be used. Avoid shadowing.

-from ._util import (KvCacheCreator, _adjust_torch_mem_fraction,
-                    create_py_executor_instance, instantiate_sampler, is_mla)
+from ._util import (KvCacheCreator, _adjust_torch_mem_fraction,
+                    create_py_executor_instance, instantiate_sampler)

---

8-8: Generalize update_sampler_max_seq_len and add types/docstring.

Make it work for any sampler exposing max_seq_len; no hard dependency on TRTLLMSampler.

-from typing import Optional
+from typing import Any, Optional
@@
-def update_sampler_max_seq_len(max_seq_len, sampler):
-    # Originally, TRTLLMSampler is constructed with executor_config, but
-    # _create_kv_cache_manager (via build_managers) may later overwrite executor_config.max_seq_len.
-    # Because TRTLLMSampler.sample_async still needs the updated limit and executor_config is
-    # deprecated inside TRTLLMSampler, keep TRTLLMSampler.max_seq_len updated with
-    # with executor_config.max_seq_len.
-    if isinstance(sampler, TRTLLMSampler):
-        assert hasattr(sampler, "max_seq_len")
-        sampler.max_seq_len = max_seq_len
+def update_sampler_max_seq_len(max_seq_len: int, sampler: Any) -> None:
+    """Propagate updated max_seq_len to the sampler if supported."""
+    if hasattr(sampler, "max_seq_len"):
+        sampler.max_seq_len = max_seq_len

Also applies to: 213-222

tensorrt_llm/_torch/pyexecutor/_util.py (1)

1-1: Add NVIDIA Apache-2.0 header (2025).

Please prepend the standard header.

+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tensorrt_llm/_torch/pyexecutor/sampler.py (2)

1-1: Add NVIDIA Apache-2.0 header (2025).

Please prepend the standard header.

+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

---

794-815: Guard kv_cache_config being None in TRTLLMSampler.init.

Signature allows Optional[KvCacheConfig], but code dereferences unconditionally; NPE risk if future callers pass None.

-        self.decoding_config = decoding_config if decoding_config else DecodingConfig(
-            decoding_mode)
-        max_attn_window = kv_cache_config.max_attention_window
+        self.decoding_config = decoding_config if decoding_config else DecucingConfig(
+            decoding_mode)
+        max_attn_window = getattr(kv_cache_config, "max_attention_window", None)