[TRTLLM-12026][feat] Support MTP with block reuse enabled for hybrid models by VALLIS-NERIA · Pull Request #12896 · NVIDIA/TensorRT-LLM

VALLIS-NERIA · 2026-04-09T17:36:39Z

Summary

Extend KVCacheManager C++ layer with mutable cache block ID accessors and logging for linear cache memory budget calculation
Refactor MambaCacheManager to support block-level cache operations (add/remove/reuse) instead of flat tensor management
Update scheduler to handle linear attention cache reuse alongside KV cache reuse
Wire through enable_cache_reuse flag for mamba cache in executor and resource manager
Add integration test for mamba2 hybrid model cache reuse
Add unit tests for KV cache manager with linear attention metadata

Test plan

Integration test for mamba2 hybrid model with cache reuse
Existing KV cache reuse tests still pass

🤖 Generated with Claude Code

Summary by CodeRabbit

Release Notes

New Features
- Added speculative decoding support for Mamba models with improved cache management.
Improvements
- Enhanced KV cache estimation for linear-attention layers.
- Better error reporting and validation messaging for cache block operations.
- Optimized cache block handling for recurrent state scenarios.
Tests
- Added MTP speculative decoding test configurations and block reuse validation.

VALLIS-NERIA · 2026-04-11T03:49:03Z

/bot run

tensorrt-cicd · 2026-04-11T03:55:06Z

PR_Github #42776 [ run ] triggered by Bot. Commit: 95923ce Link to invocation

tensorrt-cicd · 2026-04-11T04:18:17Z

PR_Github #42776 [ run ] completed with state FAILURE. Commit: 95923ce
/LLM/main/L0_MergeRequest_PR pipeline #33453 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

Link to invocation

VALLIS-NERIA · 2026-04-11T04:23:20Z

/bot run

tensorrt-cicd · 2026-04-11T04:28:55Z

PR_Github #42778 [ run ] triggered by Bot. Commit: 95923ce Link to invocation

tensorrt-cicd · 2026-04-11T04:48:36Z

PR_Github #42778 [ run ] completed with state FAILURE. Commit: 95923ce
/LLM/main/L0_MergeRequest_PR pipeline #33455 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

Link to invocation

VALLIS-NERIA · 2026-04-11T04:52:32Z

/bot run

tensorrt-cicd · 2026-04-11T04:58:32Z

PR_Github #42779 [ run ] triggered by Bot. Commit: 95923ce Link to invocation

tensorrt-cicd · 2026-04-11T05:18:24Z

PR_Github #42779 [ run ] completed with state FAILURE. Commit: 95923ce
/LLM/main/L0_MergeRequest_PR pipeline #33456 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

Link to invocation

VALLIS-NERIA · 2026-04-11T05:21:31Z

/bot run

tensorrt-cicd · 2026-04-11T05:28:33Z

PR_Github #42782 [ run ] triggered by Bot. Commit: 95923ce Link to invocation

tensorrt-cicd · 2026-04-11T05:50:30Z

PR_Github #42782 [ run ] completed with state FAILURE. Commit: 95923ce
/LLM/main/L0_MergeRequest_PR pipeline #33459 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

Link to invocation

VALLIS-NERIA · 2026-04-20T05:48:40Z

/bot run

tensorrt-cicd · 2026-04-20T05:55:04Z

PR_Github #44336 [ run ] triggered by Bot. Commit: 51ffff3 Link to invocation

tensorrt-cicd · 2026-04-20T06:08:46Z

PR_Github #44336 [ run ] completed with state FAILURE. Commit: 51ffff3
/LLM/main/L0_MergeRequest_PR pipeline #34754 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

Link to invocation

VALLIS-NERIA · 2026-04-20T06:12:03Z

/bot run

tensorrt-cicd · 2026-04-20T06:18:22Z

PR_Github #44352 [ run ] triggered by Bot. Commit: 51ffff3 Link to invocation

tensorrt-cicd · 2026-04-20T06:30:36Z

PR_Github #44352 [ run ] completed with state FAILURE. Commit: 51ffff3
/LLM/main/L0_MergeRequest_PR pipeline #34770 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

Link to invocation

VALLIS-NERIA · 2026-04-20T06:36:40Z

/bot run

VALLIS-NERIA · 2026-04-21T03:54:25Z

/bot run

tensorrt-cicd · 2026-04-21T04:00:31Z

PR_Github #44611 [ run ] triggered by Bot. Commit: 3b49991 Link to invocation

tensorrt-cicd · 2026-04-21T07:12:39Z

PR_Github #44611 [ run ] completed with state SUCCESS. Commit: 3b49991
/LLM/main/L0_MergeRequest_PR pipeline #34995 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

Link to invocation

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

VALLIS-NERIA · 2026-05-08T01:37:16Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-08T01:56:17Z

PR_Github #47273 [ run ] triggered by Bot. Commit: 274fd8a Link to invocation

tensorrt-cicd · 2026-05-08T10:31:01Z

PR_Github #47273 [ run ] completed with state SUCCESS. Commit: 274fd8a
/LLM/main/L0_MergeRequest_PR pipeline #37220 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-05-08T10:44:01Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-08T10:49:22Z

PR_Github #47387 [ run ] triggered by Bot. Commit: 274fd8a Link to invocation

tensorrt-cicd · 2026-05-08T17:15:20Z

PR_Github #47387 [ run ] completed with state SUCCESS. Commit: 274fd8a
/LLM/main/L0_MergeRequest_PR pipeline #37318 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-05-09T00:15:40Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-09T00:21:34Z

PR_Github #47454 [ run ] triggered by Bot. Commit: 274fd8a Link to invocation

tensorrt-cicd · 2026-05-09T04:56:23Z

PR_Github #47454 [ run ] completed with state SUCCESS. Commit: 274fd8a
/LLM/main/L0_MergeRequest_PR pipeline #37376 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-05-09T08:10:07Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-09T08:16:17Z

PR_Github #47496 [ run ] triggered by Bot. Commit: 274fd8a Link to invocation

tensorrt-cicd · 2026-05-09T13:05:54Z

PR_Github #47496 [ run ] completed with state FAILURE. Commit: 274fd8a
/LLM/main/L0_MergeRequest_PR pipeline #37416 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

Link to invocation

VALLIS-NERIA · 2026-05-09T14:10:14Z

/bot run --disable-fail-fast

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

VALLIS-NERIA · 2026-05-09T14:39:50Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-09T14:45:10Z

PR_Github #47527 [ run ] triggered by Bot. Commit: 7d49440 Link to invocation

Signed-off-by: xiweny <13230610+VALLIS-NERIA@users.noreply.github.com>

VALLIS-NERIA · 2026-05-10T06:46:36Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-10T06:51:54Z

PR_Github #47568 [ run ] triggered by Bot. Commit: 752f6a8 Link to invocation

tensorrt-cicd · 2026-05-10T21:24:36Z

PR_Github #47568 [ run ] completed with state SUCCESS. Commit: 752f6a8
/LLM/main/L0_MergeRequest_PR pipeline #37480 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

Link to invocation

VALLIS-NERIA · 2026-05-10T21:26:37Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-10T21:32:01Z

PR_Github #47620 [ run ] triggered by Bot. Commit: 752f6a8 Link to invocation

tensorrt-cicd · 2026-05-10T22:45:07Z

PR_Github #47620 [ run ] completed with state SUCCESS. Commit: 752f6a8
/LLM/main/L0_MergeRequest_PR pipeline #37525 completed with status: 'SUCCESS'

CI Report

Link to invocation

…models (NVIDIA#12896) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Signed-off-by: xiweny <13230610+VALLIS-NERIA@users.noreply.github.com>

github-actions Bot assigned VALLIS-NERIA Apr 9, 2026

VALLIS-NERIA changed the title ~~[None][feat] Enable mamba/linear attention cache reuse in scheduler~~ [None][feat] MTP + cache reuse for hybrid linear models Apr 13, 2026

VALLIS-NERIA force-pushed the user/xiweny/mamba_cache_reuse branch from 95923ce to 29a1d27 Compare April 14, 2026 12:46

VALLIS-NERIA added the Release Blocker PRs that blocking the final release build or branching out the release branch label Apr 20, 2026

VALLIS-NERIA force-pushed the user/xiweny/mamba_cache_reuse branch from f3f5eba to 51ffff3 Compare April 20, 2026 05:47

VALLIS-NERIA changed the title ~~[None][feat] MTP + cache reuse for hybrid linear models~~ [TRTLLM-12026][feat] Support MTP with block reuse enabled for hybrid models Apr 20, 2026

VALLIS-NERIA force-pushed the user/xiweny/mamba_cache_reuse branch from 51ffff3 to ff89391 Compare April 21, 2026 03:53

VALLIS-NERIA removed the Release Blocker PRs that blocking the final release build or branching out the release branch label Apr 24, 2026

pcastonguay approved these changes May 7, 2026

View reviewed changes

eopXD approved these changes May 7, 2026

View reviewed changes

mikeiovine approved these changes May 7, 2026

View reviewed changes

Merge branch 'main' into user/xiweny/mamba_cache_reuse

274fd8a

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

VALLIS-NERIA enabled auto-merge (squash) May 9, 2026 00:36

Merge branch 'main' into user/xiweny/mamba_cache_reuse

7d49440

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

Merge branch 'main' into user/xiweny/mamba_cache_reuse

752f6a8

Signed-off-by: xiweny <13230610+VALLIS-NERIA@users.noreply.github.com>

VALLIS-NERIA merged commit 091ad7b into NVIDIA:main May 10, 2026
6 checks passed

Conversation

VALLIS-NERIA commented Apr 9, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Summary by CodeRabbit

Release Notes

Uh oh!

VALLIS-NERIA commented Apr 11, 2026

Uh oh!

tensorrt-cicd commented Apr 11, 2026

Uh oh!

tensorrt-cicd commented Apr 11, 2026

Uh oh!

VALLIS-NERIA commented Apr 11, 2026

Uh oh!

tensorrt-cicd commented Apr 11, 2026

Uh oh!

tensorrt-cicd commented Apr 11, 2026

Uh oh!

VALLIS-NERIA commented Apr 11, 2026

Uh oh!

tensorrt-cicd commented Apr 11, 2026

Uh oh!

tensorrt-cicd commented Apr 11, 2026

Uh oh!

VALLIS-NERIA commented Apr 11, 2026

Uh oh!

tensorrt-cicd commented Apr 11, 2026

Uh oh!

tensorrt-cicd commented Apr 11, 2026

Uh oh!

VALLIS-NERIA commented Apr 20, 2026

Uh oh!

tensorrt-cicd commented Apr 20, 2026

Uh oh!

tensorrt-cicd commented Apr 20, 2026

Uh oh!

VALLIS-NERIA commented Apr 20, 2026

Uh oh!

tensorrt-cicd commented Apr 20, 2026

Uh oh!

tensorrt-cicd commented Apr 20, 2026

Uh oh!

VALLIS-NERIA commented Apr 20, 2026

Uh oh!

VALLIS-NERIA commented Apr 21, 2026

Uh oh!

tensorrt-cicd commented Apr 21, 2026

Uh oh!

tensorrt-cicd commented Apr 21, 2026

Uh oh!

VALLIS-NERIA commented May 8, 2026

Uh oh!

tensorrt-cicd commented May 8, 2026

Uh oh!

tensorrt-cicd commented May 8, 2026

Uh oh!

VALLIS-NERIA commented May 8, 2026

Uh oh!

tensorrt-cicd commented May 8, 2026

Uh oh!

tensorrt-cicd commented May 8, 2026

Uh oh!

VALLIS-NERIA commented May 9, 2026

Uh oh!

tensorrt-cicd commented May 9, 2026

Uh oh!

tensorrt-cicd commented May 9, 2026

Uh oh!

VALLIS-NERIA commented May 9, 2026

Uh oh!

tensorrt-cicd commented May 9, 2026

Uh oh!

tensorrt-cicd commented May 9, 2026

Uh oh!

VALLIS-NERIA commented May 9, 2026

Uh oh!

VALLIS-NERIA commented May 9, 2026

Uh oh!

VALLIS-NERIA commented Apr 9, 2026 •

edited by coderabbitai Bot

Loading