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

[longlee0622]

transformers 5.5.3 (PR #13994) no longer carries default max_position_embeddings on bare PreTrainedConfig. For model types that TRT-LLM understands locally but upstream transformers doesn't recognise (deepseek_v32, kimi_k2 — DeepSeek-V3.2 / Kimi-K2 checkpoints), AutoConfig.from_pretrained falls through to the bare config, and AutoTokenizer.from_pretrained then raises
'PreTrainedConfig' object has no attribute 'max_position_embeddings' before any tokenizer is constructed.

Caught by GB200 disagg perf-sanity
disagg_upload-gen_only-gb200_deepseek-v32-fp4_1k1k_con1_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL in L0_PostMerge #2722: the benchmark_serving subprocess prints argparse Namespace(...) and exits 1 on the very next get_tokenizer -> AutoTokenizer.from_pretrained(... trust_remote_code=True) call. CTX/GEN servers log the same warning via load_hf_tokenizer. Pre-merge missed it because pre-merge DeepSeek-V3.2 coverage (TestDeepSeekV32 on DGX-B200) goes through LLM(...) which has tokenizer-load fallbacks, and only the post-merge GB200 perf-sanity stage runs trtllm-serve + benchmark_serving.

Fix: register deepseek_v32 and kimi_k2 against the local DeepseekV3Config in transformers.models.auto.configuration_auto.CONFIG_MAPPING at TRT-LLM import time, mirroring the existing _CONFIG_REGISTRY in tensorrt_llm/_torch/pyexecutor/config_utils.py. Registration goes through the underlying mapping (not AutoConfig.register) to bypass its config.model_type == key consistency check — DeepseekV3Config keeps model_type = "deepseek_v3" so existing in-engine paths are unaffected.

Trigger is placed in tensorrt_llm/_torch/models/__init__.py (already eagerly imported from tensorrt_llm/__init__.py) to avoid adding to tensorrt_llm/__init__.py's E402 legacy-ruff baseline.

Pre-merge regression coverage added in
tests/unittest/_torch/test_custom_config_registration.py: both registration presence and an end-to-end
AutoConfig.from_pretrained(<tmpdir config.json with model_type>) call that hits the exact dispatch the broken post-merge test triggered.

Summary by CodeRabbit

• New Features

  • Added support for deepseek_v32 and kimi_k2 model types with Transformers' AutoConfig and AutoTokenizer.
  • Improved compatibility with Transformers 5.5+.

• Tests

  • Added regression tests for model type registration.

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[longlee0622]

/bot run

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR registers TensorRT-LLM's DeepseekV3Config with Hugging Face Transformers' AutoConfig for two additional model types (deepseek_v32, kimi_k2) by mutating CONFIG_MAPPING at import time. A side-effect import in models/__init__.py triggers registration when the package loads, and regression tests verify both CONFIG_MAPPING registration and AutoConfig.from_pretrained dispatch.

Changes

AutoConfig Registration for Custom Model Types

Layer / File(s)	Summary
Config registration and wiring tensorrt_llm/_torch/configs/__init__.py, tensorrt_llm/_torch/models/__init__.py	A private helper registers DeepseekV3Config under deepseek_v32 and kimi_k2 in Transformers' CONFIG_MAPPING, executes immediately, and is deleted. Side-effect import in models/__init__.py triggers this registration on tensorrt_llm import to enable AutoTokenizer.from_pretrained dispatch.
Registration verification tests tests/unittest/_torch/test_custom_config_registration.py	Module-level docstring explains the failure mode. Two parametrized tests verify both model types are registered in CONFIG_MAPPING and that AutoConfig.from_pretrained resolves to DeepseekV3Config with preserved config attributes.

🎯 2 (Simple) | ⏱️ ~12 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly describes the primary change: registering deepseek_v32 and kimi_k2 model types with transformers AutoConfig, which is the core fix.
Description check	✅ Passed	The PR description comprehensively explains the issue, root cause, solution, and implementation details; test coverage is documented with regression tests.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (1)

tests/unittest/_torch/test_custom_config_registration.py (1)

1-52: QA test list updates are not required for this unittest-only change.

Per coding guidelines, since this PR only adds unit tests under tests/unittest/, QA test list updates are unnecessary. The existing integration tests already provide regression coverage:

• llm_function_core.txt includes accuracy/test_disaggregated_serving.py::TestKimiK2 and TestDeepSeekV32Exp tests that exercise serving/tokenizer paths
• llm_perf_multinode.txt includes perf/test_perf_sanity.py::test_e2e[gb200_deepseek-v32-fp4_...] and test_e2e[gb200_kimi-k2-thinking-fp4_...] entries

These existing tests will serve as the regression signal for AutoConfig/AutoTokenizer dispatch in production serving scenarios.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/test_custom_config_registration.py` around lines 1 -
52, The PR adds only unit tests under tests/unittest and therefore should not
include QA test-list updates; update the PR description/metadata to remove any
QA list modifications and explicitly state that no QA test list changes are
required because the new tests are unit-only (referencing the new tests
test_custom_model_type_registered_with_autoconfig and
test_autoconfig_from_pretrained_resolves_to_local_config which live under
tests/unittest/_torch/test_custom_config_registration.py) so that reviewers know
existing integration tests (llm_function_core.txt and llm_perf_multinode.txt)
provide regression coverage.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@tests/unittest/_torch/test_custom_config_registration.py`:
- Around line 1-52: The PR adds only unit tests under tests/unittest and
therefore should not include QA test-list updates; update the PR
description/metadata to remove any QA list modifications and explicitly state
that no QA test list changes are required because the new tests are unit-only
(referencing the new tests test_custom_model_type_registered_with_autoconfig and
test_autoconfig_from_pretrained_resolves_to_local_config which live under
tests/unittest/_torch/test_custom_config_registration.py) so that reviewers know
existing integration tests (llm_function_core.txt and llm_perf_multinode.txt)
provide regression coverage.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: ba5fab7b-6d05-4d65-acbe-cb992b5a5bf3

📥 Commits

Reviewing files that changed from the base of the PR and between 89e1410 and 1c47176.

📒 Files selected for processing (3)

• tensorrt_llm/_torch/configs/__init__.py
• tensorrt_llm/_torch/models/__init__.py
• tests/unittest/_torch/test_custom_config_registration.py

[tensorrt-cicd]

PR_Github #49116 [ run ] triggered by Bot. Commit: 1c47176 Link to invocation

[tensorrt-cicd]

PR_Github #49116 [ run ] completed with state SUCCESS. Commit: 1c47176
/LLM/main/L0_MergeRequest_PR pipeline #38820 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[Hudayday]

/bot run --disable-fail-fast --extra-stage "DGX_B200-8_GPUs-PyTorch-1, DGX_B200-8_GPUs-PyTorch-2, DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1, DGX_B200-8_GPUs-PyTorch-PerfSanity-Post-Merge-4"

[tensorrt-cicd]

PR_Github #49188 [ run ] triggered by Bot. Commit: 1c47176 Link to invocation

[tensorrt-cicd]

PR_Github #49188 [ run ] completed with state SUCCESS. Commit: 1c47176
/LLM/main/L0_MergeRequest_PR pipeline #38865 completed with status: 'SUCCESS'

CI Report

Link to invocation