[https://nvbugs/6185234][fix] register deepseek_v32 / kimi_k2 with transformers AutoConfig

tensorrt_llm/_torch/configs/__init__.py

@@ -1,3 +1,27 @@
 from tensorrt_llm._torch.configs.deepseek_v3 import DeepseekV3Config
 
+
+def _register_custom_configs_with_transformers() -> None:
+    # Make AutoConfig.from_pretrained / AutoTokenizer.from_pretrained accept
+    # model_types that TRT-LLM understands but upstream transformers does not
+    # (DeepSeek-V3.2 and Kimi K2 both ship config.json with these model_types
+    # and rely on TRT-LLM's local DeepseekV3Config workaround).
+    #
+    # Without this, transformers 5.5.x falls back to a bare PreTrainedConfig
+    # that lacks attributes like `max_position_embeddings`, and
+    # AutoTokenizer.from_pretrained then raises AttributeError before any
+    # tokenizer can be constructed. Bypass AutoConfig.register's model_type
+    # consistency check (DeepseekV3Config.model_type is "deepseek_v3") by
+    # writing into the underlying mapping directly.
+    from transformers.models.auto.configuration_auto import CONFIG_MAPPING
+
+    for model_type in ("deepseek_v32", "kimi_k2"):
+        if model_type in CONFIG_MAPPING:
+            continue
+        CONFIG_MAPPING.register(model_type, DeepseekV3Config, exist_ok=True)
+
+
+_register_custom_configs_with_transformers()
+del _register_custom_configs_with_transformers
+
 __all__ = ["DeepseekV3Config"]

tensorrt_llm/_torch/models/__init__.py

@@ -1,5 +1,10 @@
 import transformers
 
+# Importing _torch.configs triggers AutoConfig registration for TRT-LLM-only
+# model_types (deepseek_v32, kimi_k2) so AutoTokenizer.from_pretrained works
+# under transformers >= 5.5; see _torch/configs/__init__.py.
+import tensorrt_llm._torch.configs  # noqa: F401
+
 from .modeling_auto import AutoModelForCausalLM
 from .modeling_bert import BertForSequenceClassification
 from .modeling_clip import CLIPVisionModel

tests/unittest/_torch/test_custom_config_registration.py

@@ -0,0 +1,52 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Regression tests for the transformers AutoConfig / AutoTokenizer dispatch
+for TRT-LLM-only model_types (deepseek_v32, kimi_k2).
+
+Before this registration, transformers >= 5.5 falls back to a bare
+PreTrainedConfig that lacks `max_position_embeddings`, and
+AutoTokenizer.from_pretrained then raises AttributeError on it. The
+broken test that motivated this — perf/test_perf_sanity.py disagg gen_only
+on GB200 — only runs in L0_PostMerge because it needs 12 GB200 GPUs across
+3 nodes, so a cheap pre-merge unit test is the right place to catch
+regressions.
+"""
+
+import json
+
+import pytest
+
+import tensorrt_llm  # noqa: F401  triggers AutoConfig registration
+from tensorrt_llm._torch.configs import DeepseekV3Config
+
+
+@pytest.mark.parametrize("model_type", ["deepseek_v32", "kimi_k2"])
+def test_custom_model_type_registered_with_autoconfig(model_type):
+    from transformers.models.auto.configuration_auto import CONFIG_MAPPING
+
+    assert model_type in CONFIG_MAPPING
+    assert CONFIG_MAPPING[model_type] is DeepseekV3Config
+
+
+@pytest.mark.parametrize("model_type", ["deepseek_v32", "kimi_k2"])
+def test_autoconfig_from_pretrained_resolves_to_local_config(tmp_path, model_type):
+    # Mirrors what the benchmark_serving subprocess does under the hood:
+    # AutoTokenizer.from_pretrained -> AutoConfig.from_pretrained. Without
+    # the registration this fails through to a bare PreTrainedConfig that
+    # lacks `max_position_embeddings`.
+    from transformers import AutoConfig
+
+    model_dir = tmp_path / model_type
+    model_dir.mkdir()
+    (model_dir / "config.json").write_text(
+        json.dumps(
+            {
+                "model_type": model_type,
+                "max_position_embeddings": 16384,
+            }
+        )
+    )
+
+    cfg = AutoConfig.from_pretrained(str(model_dir))
+    assert isinstance(cfg, DeepseekV3Config)
+    assert cfg.max_position_embeddings == 16384