[None][feat] BREAKING: add per-model KV cache manager v2 auto selection

[yizhang-nv]

Summary by CodeRabbit

• New Features

  • Added an “auto” option for KV cache manager selection, allowing the app to choose a model-specific default when supported.
  • Improved argument handling so model defaults can influence cache manager behavior during startup.

• Bug Fixes

  • Tightened cache manager selection logic for more reliable behavior.
  • Ensured models with built-in cache preferences use those defaults consistently.

• Tests

  • Expanded coverage for default resolution, explicit overrides, and validation of accepted values.

Description

Add "auto" as the default value of KvCacheConfig.use_kv_cache_manager_v2 so models can opt in to KV cache manager V2 through their model defaults. Explicit user-provided true or false values continue to take precedence, while models without a V2 default resolve to false.

DeepSeek V4 opts in to V2. Gemma and DeepSeek R1 defaults are unchanged.

Test Coverage

• pytest targeted KV cache auto-resolution and DeepSeek V4 model-default tests: 7 passed.
• Clean C++ build for SM80, SM90, SM100, and SM103 (A100, H100/H200, B200, and B300): passed.
• pre-commit run --files <changed files>: passed.
• DeepSeek R1 compatibility validation on 4x B200 with a temporary V2 model default: confirmed KVCacheManagerV2 and KVCacheV2Scheduler; MMLU and GSM8K completed and passed their accuracy thresholds. The remaining CNN/DailyMail evaluation could not start because the local cached test split contained zero samples. The temporary R1 default was reverted before this PR.

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.

[yizhang-nv]

/bot run --disable-fail-fast

[coderabbitai]

📝 Walkthrough

Walkthrough

Changes KvCacheConfig.use_kv_cache_manager_v2 to accept bool | Literal["auto"] with a default of "auto", add resolution logic that consults model-specific defaults during model loading, update DeepseekV4 model defaults and executor routing checks, update the CLI example default, and add corresponding golden manifest and test updates.

Changes

KV Cache Manager V2 Auto Resolution

Layer / File(s)	Summary
Config field and schema updates tensorrt_llm/llmapi/llm_args.py, tensorrt_llm/usage/llm_args_golden_manifest.json, tests/unittest/llmapi/test_llm_args.py	use_kv_cache_manager_v2 field type changed to bool | Literal["auto"] with default "auto"; golden manifest schema updated to allow the "auto" literal; declaration test validates default, boolean handling, and invalid-value errors.
Auto-resolution helper and tests tensorrt_llm/llmapi/llm_utils.py, tests/unittest/llmapi/test_llm_args.py	New _resolve_kv_cache_manager_v2_auto helper resolves "auto" into a boolean based on model defaults, raising ValueError for invalid values; tests cover model-default fallback and user-override behavior.
Model loader wiring tensorrt_llm/_torch/pyexecutor/model_loader.py	Imports and calls the resolver when use_kv_cache_manager_v2 is "auto", initializes model_defaults safely, and logs the resolved value with the model class name.
DeepseekV4 model default and executor routing tensorrt_llm/_torch/models/modeling_deepseekv4.py, tensorrt_llm/_torch/pyexecutor/_util.py, tests/unittest/_torch/modeling/test_modeling_deepseekv4.py	DeepseekV4 defaults now set use_kv_cache_manager_v2: True; executor routing tightens the needs_v2 check to is True and updates fallback comments; test updated accordingly.
CLI example default update examples/llm-api/quickstart_advanced.py	--use_kv_cache_manager_v2 argparse default changed from False to 'auto' with updated help text.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Suggested reviewers: juney-nvidia, aswinvisva, moraxu, JunyiXu-nv, ziyixiong-nv

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 27.27% 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
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.
Title check	✅ Passed	The title matches the main change and follows the required ticket/type pattern, with BREAKING marking the API change.
Description check	✅ Passed	The description follows the template with clear Description, Test Coverage, and PR Checklist sections.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

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

338-357: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Auto-resolution log never fires for the exact scenario this PR targets (e.g. DeepseekV4).

apply_model_defaults_to_llm_args (called at lines 342-343) mutates llm_args.kv_cache_config.use_kv_cache_manager_v2 in place via setattr for every field in new_args, overriding the "auto" sentinel with the model's explicit default (e.g. True for DeepseekV4ForCausalLM) whenever the user hasn't explicitly set it. By line 349, use_kv_cache_manager_v2 is therefore already True/False, not "auto" — so if use_kv_cache_manager_v2 == "auto": at line 351 is False and the "Resolved use_kv_cache_manager_v2='auto' to %s for %s" log never prints for models that supply an explicit use_kv_cache_manager_v2 default. The final resolved boolean is still correct, but the dedicated auto-resolution log — the primary observability feature this PR adds — is dead code for this case.

Capture whether the setting was "auto" before model defaults are applied, and gate the log on that captured flag instead of re-reading the already-mutated value.

🐛 Proposed fix

+        use_kv_cache_manager_v2_was_auto = (
+            llm_args.kv_cache_config.use_kv_cache_manager_v2 == "auto")
         model_defaults = {}
         if model_cls and hasattr(model_cls, 'get_model_defaults'):
             model_defaults = model_cls.get_model_defaults(llm_args) or {}
             if model_defaults:
                 applied_defaults = apply_model_defaults_to_llm_args(
                     llm_args, model_defaults)
                 if applied_defaults:
                     logger.info(
                         f"Applied model defaults for {model_cls.__name__}: {applied_defaults}"
                     )

-        use_kv_cache_manager_v2 = llm_args.kv_cache_config.use_kv_cache_manager_v2
         _resolve_kv_cache_manager_v2_auto(llm_args, model_defaults)
-        if use_kv_cache_manager_v2 == "auto":
+        if use_kv_cache_manager_v2_was_auto:
             logger.info(
                 "Resolved use_kv_cache_manager_v2='auto' to %s for %s",
                 llm_args.kv_cache_config.use_kv_cache_manager_v2,
                 model_cls.__name__
                 if model_cls is not None else "unknown model")

🤖 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 `@tensorrt_llm/_torch/pyexecutor/model_loader.py` around lines 338 - 357,
Capture the original `llm_args.kv_cache_config.use_kv_cache_manager_v2` value
before calling `apply_model_defaults_to_llm_args` in `model_loader.py`, since
that helper mutates the field in place. Use that pre-মutation flag to detect
whether the user had `"auto"` set, and gate the `logger.info` message on it
instead of re-reading the updated value. Keep the existing resolution logic in
`_resolve_kv_cache_manager_v2_auto` unchanged, but ensure the auto-resolution
log still fires for models like `DeepseekV4ForCausalLM` that provide an explicit
default.

🧹 Nitpick comments (2)

tests/unittest/_torch/modeling/test_modeling_deepseekv4.py (1)

151-162: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Test correctly covers the static defaults dict; add end-to-end coverage separately.

This unit test is sufficient for validating DeepseekV4ForCausalLM.get_model_defaults()'s static return value. However, it doesn't cover the end-to-end auto-resolution behavior in model_loader.py (e.g., that "auto" + DeepseekV4 model default correctly resolves to True and is logged). Given the log-gating bug flagged in model_loader.py (lines 338-357), consider adding a regression test — likely in tests/unittest/llmapi/test_llm_args.py or a dedicated model_loader test — asserting the resolved value and log behavior for a model with an explicit use_kv_cache_manager_v2 default.

As per path instructions, "Keep feedback actionable: suggest concrete list file names and whether coverage is sufficient, insufficient, or needs follow-up outside the PR."

🤖 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/modeling/test_modeling_deepseekv4.py` around lines 151
- 162, The current test only verifies the static defaults returned by
DeepseekV4ForCausalLM.get_model_defaults(), but it does not cover the end-to-end
auto-resolution path in model_loader.py. Add a separate regression test in
tests/unittest/llmapi/test_llm_args.py or a dedicated model_loader test that
exercises a DeepseekV4 model with an explicit use_kv_cache_manager_v2 default,
asserts that "auto" resolves to True, and verifies the expected logging behavior
from the auto-resolution flow. Keep the existing test as-is since it is
sufficient for the static defaults dict, and use the new test to cover the
log-gating behavior in model_loader and the LlmArgs/DeepseekV4ForCausalLM
interaction.

Source: Path instructions

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

18-19: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Private helper imported across module boundary.

_resolve_kv_cache_manager_v2_auto is underscore-prefixed in llm_utils.py, signaling it's module-private, yet it's now imported and called from model_loader.py. Consider renaming it to a public name (drop the leading underscore) in llm_utils.py since it's part of the cross-module contract now.

As per coding guidelines, "Variables and functions not part of a class's or module's public interface should be prefixed with an underscore."

🤖 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 `@tensorrt_llm/_torch/pyexecutor/model_loader.py` around lines 18 - 19, The
imported helper _resolve_kv_cache_manager_v2_auto is being used across a module
boundary even though its leading underscore marks it as private. Update the
symbol in llm_utils.py to a public name without the underscore, and then change
model_loader.py to import and call the renamed public helper consistently. Keep
apply_model_defaults_to_llm_args unchanged and make sure the new name reflects
the intended cross-module API.

Source: Coding guidelines

🤖 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.

Inline comments:
In `@examples/llm-api/quickstart_advanced.py`:
- Around line 132-138: The --use_kv_cache_manager_v2 argument in
quickstart_advanced.py is not truly tri-state because action='store_true' only
produces auto or True, so users cannot explicitly disable it when model defaults
enable it. Update the parser setup for this option to accept auto/true/false as
a string-valued CLI flag, and adjust the downstream handling in the code that
consumes this argument so it maps those values to the corresponding
enabled/disabled/auto behavior.

---

Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/model_loader.py`:
- Around line 338-357: Capture the original
`llm_args.kv_cache_config.use_kv_cache_manager_v2` value before calling
`apply_model_defaults_to_llm_args` in `model_loader.py`, since that helper
mutates the field in place. Use that pre-মutation flag to detect whether the
user had `"auto"` set, and gate the `logger.info` message on it instead of
re-reading the updated value. Keep the existing resolution logic in
`_resolve_kv_cache_manager_v2_auto` unchanged, but ensure the auto-resolution
log still fires for models like `DeepseekV4ForCausalLM` that provide an explicit
default.

---

Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/model_loader.py`:
- Around line 18-19: The imported helper _resolve_kv_cache_manager_v2_auto is
being used across a module boundary even though its leading underscore marks it
as private. Update the symbol in llm_utils.py to a public name without the
underscore, and then change model_loader.py to import and call the renamed
public helper consistently. Keep apply_model_defaults_to_llm_args unchanged and
make sure the new name reflects the intended cross-module API.

In `@tests/unittest/_torch/modeling/test_modeling_deepseekv4.py`:
- Around line 151-162: The current test only verifies the static defaults
returned by DeepseekV4ForCausalLM.get_model_defaults(), but it does not cover
the end-to-end auto-resolution path in model_loader.py. Add a separate
regression test in tests/unittest/llmapi/test_llm_args.py or a dedicated
model_loader test that exercises a DeepseekV4 model with an explicit
use_kv_cache_manager_v2 default, asserts that "auto" resolves to True, and
verifies the expected logging behavior from the auto-resolution flow. Keep the
existing test as-is since it is sufficient for the static defaults dict, and use
the new test to cover the log-gating behavior in model_loader and the
LlmArgs/DeepseekV4ForCausalLM interaction.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 8a31798e-a203-4fac-8728-a663b4c935e2

📥 Commits

Reviewing files that changed from the base of the PR and between 790c866 and 67a9fe0.

📒 Files selected for processing (9)

• examples/llm-api/quickstart_advanced.py
• tensorrt_llm/_torch/models/modeling_deepseekv4.py
• tensorrt_llm/_torch/pyexecutor/_util.py
• tensorrt_llm/_torch/pyexecutor/model_loader.py
• tensorrt_llm/llmapi/llm_args.py
• tensorrt_llm/llmapi/llm_utils.py
• tensorrt_llm/usage/llm_args_golden_manifest.json
• tests/unittest/_torch/modeling/test_modeling_deepseekv4.py
• tests/unittest/llmapi/test_llm_args.py

[tensorrt-cicd]

PR_Github #56894 [ run ] triggered by Bot. Commit: 5ea1d1a Link to invocation

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #56985 [ run ] triggered by Bot. Commit: 5ea1d1a Link to invocation

[tensorrt-cicd]

PR_Github #56894 [ run ] completed with state ABORTED. Commit: 5ea1d1a

Link to invocation

[tensorrt-cicd]

PR_Github #56985 [ run ] completed with state FAILURE. Commit: 5ea1d1a
/LLM/main/L0_MergeRequest_PR pipeline #45786 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

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #57363 [ run ] triggered by Bot. Commit: 5ea1d1a Link to invocation

[tensorrt-cicd]

PR_Github #57363 [ run ] completed with state FAILURE. Commit: 5ea1d1a
/LLM/main/L0_MergeRequest_PR pipeline #46114 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

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #57678 [ run ] triggered by Bot. Commit: e9bcb46 Link to invocation

[tensorrt-cicd]

PR_Github #57678 [ run ] completed with state SUCCESS. Commit: e9bcb46
/LLM/main/L0_MergeRequest_PR pipeline #46398 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

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #57917 [ run ] triggered by Bot. Commit: e9bcb46 Link to invocation

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58398 [ run ] triggered by Bot. Commit: 384a077 Link to invocation

[QiJune]

LGTM

[tensorrt-cicd]

PR_Github #58398 [ run ] completed with state SUCCESS. Commit: 384a077
/LLM/main/L0_MergeRequest_PR pipeline #47018 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

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58859 [ run ] triggered by Bot. Commit: 384a077 Link to invocation

[tensorrt-cicd]

PR_Github #58859 [ run ] completed with state SUCCESS. Commit: 384a077
/LLM/main/L0_MergeRequest_PR pipeline #47409 completed with status: 'SUCCESS'

CI Report

Link to invocation