[TRTLLM-11875][feat] BREAKING: MambaCacheManager based on KVCacheManagerV2 & agentic prefix caching by VALLIS-NERIA · Pull Request #16598 · NVIDIA/TensorRT-LLM

VALLIS-NERIA · 2026-07-20T06:04:56Z

Description

This PR adds the KV-cache-manager V2 implementation for hybrid Mamba models and moves Mamba state reuse onto explicit snapshot boundaries.

Key changes:

Add MambaHybridCacheManagerV2, including recurrent-state allocation, replay, promotion, dynamic-tree capacity, and deterministic stochastic-rounding state.
Add MambaStateConfig with periodic and fixed-position snapshot policies. The legacy mamba_state_cache_interval remains accepted only as a deprecated config-file alias.
Reuse Mamba state only at planned snapshot boundaries. Hybrid models that enable block reuse without any snapshot policy now warn and disable block reuse during model-aware config validation.
Add per-conversation reuse semantics: periodic Mamba snapshots are disabled and only explicit stable conversation boundaries are committed.
Add V2 Mamba state transfer to Python/NIXL disaggregated serving using explicit slot/layer strides.
Support one-model MTP and dynamic-tree MTP with V2, including target-tree mask regeneration and accepted-state promotion.
Handle attention-free hybrid PP stages in disaggregated transfer without zero-head duplication arithmetic.

This is an API/config migration: use kv_cache_config.mamba_state_config.periodic_snapshot_interval instead of kv_cache_config.mamba_state_cache_interval in Python code.

Test Coverage

Snapshot/config resolution tests, including V2-only fixed offsets and warning-plus-disable behavior when no snapshot policy exists.
V1/V2 hybrid manager, replay, snapshot planning, state promotion, scheduler, capacity, and per-conversation unit tests.
V1/V2 Mamba disaggregated extraction and transfer tests, including affine stride layouts and attention-free PP ranks.
Dynamic-tree metadata, relocation, capacity, and accepted-leaf state-promotion tests.
Parameterized Qwen3.5 and Nemotron Super V3 MTP accuracy tests are registered in the B200 CI lists for both V1 and V2.

Local validation after merging current main:

Relevant dynamic-tree/V2 tests: 4 passed.
Snapshot model-resolution tests: 6 passed.
Disaggregated peer/rank-info tests: 25 passed.
Qwen3.5 V2 MTP GSM8K: 94.8446 (threshold 91.327), passed.
All four new V1/V2 integration node IDs collect successfully.
Nemotron Super V3 V2 MTP: MMLU 85.38 (threshold 83.483) and GSM8K 92.15 (threshold 88.697), passed.
Pre-commit checks pass for all changed files.

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 is updated if ownership changes.
Documentation is updated as needed.
Update the TAVA architecture diagram if there is a significant design change.
The reviewers assigned automatically or 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, comment /bot help.

VALLIS-NERIA · 2026-07-20T06:07:48Z

/bot run

tensorrt-cicd · 2026-07-20T06:13:22Z

PR_Github #60284 [ run ] triggered by Bot. Commit: bfc9660 Link to invocation

tensorrt-cicd · 2026-07-20T06:31:52Z

PR_Github #60284 [ run ] completed with state FAILURE. Commit: bfc9660
/LLM/main/L0_MergeRequest_PR pipeline #48641 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

VALLIS-NERIA · 2026-07-20T15:18:45Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-20T15:24:18Z

PR_Github #60398 [ run ] triggered by Bot. Commit: 73e76f7 Link to invocation

tensorrt-cicd · 2026-07-20T18:43:39Z

PR_Github #60398 [ run ] completed with state FAILURE. Commit: 73e76f7
/LLM/main/L0_MergeRequest_PR pipeline #48738 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

tensorrt-cicd · 2026-07-27T09:48:43Z

PR_Github #61884 [ run ] triggered by Bot. Commit: 78be52a Link to invocation

VALLIS-NERIA · 2026-07-27T11:39:58Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-27T11:42:19Z

PR_Github #61884 [ run ] completed with state SUCCESS. Commit: 78be52a
/LLM/main/L0_MergeRequest_PR pipeline #50078 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

tensorrt-cicd · 2026-07-27T11:45:37Z

PR_Github #61901 [ run ] triggered by Bot. Commit: 78be52a Link to invocation

VALLIS-NERIA · 2026-07-27T12:15:35Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-27T12:15:37Z

PR_Github #61901 [ run ] completed with state FAILURE. Commit: 78be52a
/LLM/main/L0_MergeRequest_PR pipeline #50092 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

tensorrt-cicd · 2026-07-27T12:20:51Z

PR_Github #61909 [ run ] triggered by Bot. Commit: 78be52a Link to invocation

VALLIS-NERIA · 2026-07-27T12:48:45Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-27T12:48:49Z

PR_Github #61909 [ run ] completed with state FAILURE. Commit: 78be52a
/LLM/main/L0_MergeRequest_PR pipeline #50100 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

tensorrt-cicd · 2026-07-27T12:55:10Z

PR_Github #61911 [ run ] triggered by Bot. Commit: 78be52a Link to invocation

tensorrt-cicd · 2026-07-27T16:58:16Z

PR_Github #61911 [ run ] completed with state FAILURE. Commit: 78be52a
/LLM/main/L0_MergeRequest_PR pipeline #50102 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

VALLIS-NERIA · 2026-07-28T03:35:31Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-28T03:41:07Z

PR_Github #62087 [ run ] triggered by Bot. Commit: 78be52a Link to invocation

tensorrt-cicd · 2026-07-28T03:59:30Z

PR_Github #62087 [ run ] completed with state FAILURE. Commit: 78be52a
/LLM/main/L0_MergeRequest_PR pipeline #50272 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

VALLIS-NERIA · 2026-07-28T03:59:51Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-28T04:11:41Z

PR_Github #62096 [ run ] triggered by Bot. Commit: 78be52a Link to invocation

tensorrt-cicd · 2026-07-28T06:03:05Z

PR_Github #62096 [ run ] completed with state SUCCESS. Commit: 78be52a
/LLM/main/L0_MergeRequest_PR pipeline #50279 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

VALLIS-NERIA · 2026-07-28T06:43:53Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-28T06:49:28Z

PR_Github #62139 [ run ] triggered by Bot. Commit: 78be52a Link to invocation

tensorrt-cicd · 2026-07-28T07:42:38Z

PR_Github #62139 [ run ] completed with state SUCCESS. Commit: 78be52a
/LLM/main/L0_MergeRequest_PR pipeline #50317 completed with status: 'SUCCESS'

CI Report

Link to invocation

…dings Migrates tensorrt_llm/runtime/kv_cache_manager_v2 from pure Python to a C++ implementation under cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/ with nanobind bindings compiled into bindings.so, preserving the same public API. The dispatcher __init__.py selects the backend via TLLM_KV_CACHE_MANAGER_V2_BACKEND (default "cpp"); both backends pass the shared test suite. CODEOWNERS assigns the new C++ tree to trt-llm-kv-cache-manager-devs. Includes ports of subsequent main features: commit-min-snapshot + SWA-slot reservation, SHA-256 block-key hashing, CUDA-graph request IDs, event manager and stats API to C++, uint64 ReuseScope salt/lora_id, resume-utilization KV constraints, per-conversation KV cache block reuse (PlannedDropHandle), the reuse-probe _KVCache refactor (compute the ReuseMatch once in the manager and pass it into the cache, NVIDIA#14333), SSM-snapshot iteration stats (NVIDIA#16598), the simplified plan_committed_block_drop (db117e6276), and an MPI teardown fix for the unittest/bindings CI shard. Migration planning docs (TODO.md, MIGRATION_PLAN_CPP.md, CPP_MIGRATION_PLAN_MAIN_15633.md) are kept on a separate branch. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

…n the precheck Main's d6a2d25 (NVIDIA#16598) grew _resolve_kv_cache_manager_v2_auto a third optional param and taught it to read llm_args.cache_transceiver_config to gate a V2 model default on the NIXL Python transceiver. The precheck's contract test pinned the old 2-param signature, so it failed in CI, and the driver's shim lacked cache_transceiver_config, so a V2-default model would have silently fallen back to V1 via the except path. Resolve the transceiver runtime before the V2 setting (the resolver treats an unresolved 'auto' runtime as non-PYTHON), hand the resolver the cache config through the shim, and re-anchor the contract test on the call shape the driver actually uses: first two positional params fixed, any later addition must carry a default. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>

…dings Migrates tensorrt_llm/runtime/kv_cache_manager_v2 from pure Python to a C++ implementation under cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/ with nanobind bindings compiled into bindings.so, preserving the same public API. The dispatcher __init__.py selects the backend via TLLM_KV_CACHE_MANAGER_V2_BACKEND (default "cpp"); both backends pass the shared test suite. CODEOWNERS assigns the new C++ tree to trt-llm-kv-cache-manager-devs. Includes ports of subsequent main features: commit-min-snapshot + SWA-slot reservation, SHA-256 block-key hashing, CUDA-graph request IDs, event manager and stats API to C++, uint64 ReuseScope salt/lora_id, resume-utilization KV constraints, per-conversation KV cache block reuse (PlannedDropHandle), the reuse-probe _KVCache refactor (compute the ReuseMatch once in the manager and pass it into the cache, NVIDIA#14333), SSM-snapshot iteration stats (NVIDIA#16598), the simplified plan_committed_block_drop (db117e6276), and an MPI teardown fix for the unittest/bindings CI shard. Migration planning docs (TODO.md, MIGRATION_PLAN_CPP.md, CPP_MIGRATION_PLAN_MAIN_15633.md) are kept on a separate branch. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

github-actions Bot assigned VALLIS-NERIA Jul 20, 2026

VALLIS-NERIA force-pushed the agent/v2-mamba-snapshot-reuse-core branch 2 times, most recently from 94db6dc to 6ffdf55 Compare July 20, 2026 10:19

VALLIS-NERIA added the api-breaking Accepted LLM API contract change that is backwards-incompatible label Jul 20, 2026

VALLIS-NERIA changed the title ~~[TRTLLM-11875][feat] MambaCacheManager based on KVCacheManagerV2 & agentic prefix caching~~ [TRTLLM-11875][feat] BREAKING: MambaCacheManager based on KVCacheManagerV2 & agentic prefix caching Jul 20, 2026

VALLIS-NERIA force-pushed the agent/v2-mamba-snapshot-reuse-core branch from 73e76f7 to f7770b5 Compare July 21, 2026 05:12

VALLIS-NERIA marked this pull request as ready for review July 21, 2026 05:14

VALLIS-NERIA requested review from a team as code owners July 21, 2026 05:14

VALLIS-NERIA requested review from QiJune, SimengLiu-nv, arysef, eopXD and pcastonguay July 21, 2026 05:14

VALLIS-NERIA merged commit d6a2d25 into NVIDIA:main Jul 28, 2026
7 checks passed

Uh oh!

Conversation

VALLIS-NERIA commented Jul 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

VALLIS-NERIA commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

VALLIS-NERIA commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 27, 2026

Uh oh!

VALLIS-NERIA commented Jul 27, 2026

Uh oh!

tensorrt-cicd commented Jul 27, 2026

Uh oh!

tensorrt-cicd commented Jul 27, 2026

Uh oh!

VALLIS-NERIA commented Jul 27, 2026

Uh oh!

tensorrt-cicd commented Jul 27, 2026

Uh oh!

tensorrt-cicd commented Jul 27, 2026

Uh oh!

VALLIS-NERIA commented Jul 27, 2026

Uh oh!

tensorrt-cicd commented Jul 27, 2026

Uh oh!

tensorrt-cicd commented Jul 27, 2026

Uh oh!

tensorrt-cicd commented Jul 27, 2026

Uh oh!

VALLIS-NERIA commented Jul 28, 2026

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

VALLIS-NERIA commented Jul 28, 2026

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

VALLIS-NERIA commented Jul 28, 2026

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

11 participants

VALLIS-NERIA commented Jul 20, 2026 •

edited

Loading