[TRTLLM-11875][feat] Support fine-grained context chunk management by VALLIS-NERIA · Pull Request #16427 · NVIDIA/TensorRT-LLM

VALLIS-NERIA · 2026-07-15T13:15:49Z

Summary by CodeRabbit

New Features
- Added snapshot-aware context chunking to align processing with recurrent-state snapshot points.
- Exposed expected snapshot points through the Python request interface.
- Improved cache reuse scheduling across hybrid model execution.
Bug Fixes
- Improved chunk sizing under context and compute-capacity limits.
- Corrected KV-head selection when earlier layers report zero heads.
- Added validation for incompatible hybrid cache transceiver configurations.
- Improved cache buffer sizing and alignment for FlashMLA setups.

Description

This PR has two goals:

Fix several independent KV-cache correctness issues.
Add expect_snapshot_points as the manager-to-scheduler contract and move the existing C++ Mamba interval-reuse chunking onto that contract.

Snapshot-aware context chunking

Store expect_snapshot_points on the C++ request object and expose it through nanobind, so the C++ and Python schedulers share the same per-request field.
Have CppMambaHybridCacheManager publish regular states_snapshot_interval boundaries before scheduling when block reuse is enabled. Prompt end remains the implicit final boundary.
Make the C++ MicroBatchScheduler, Python PyMicroBatchScheduler, and KVCacheV2Scheduler honor these absolute boundaries under FORCE_CHUNK.
Preserve the existing fallback for requests without snapshot points and block-align capacity and max-context truncation.
Add focused C++ and Python coverage for interval publication, successive boundaries, prompt-end handling, and capacity truncation.

KV-cache fixes

Calculate the default max_tokens_in_buffer only after the effective tokens_per_block is known, including the FlashMLA override, and round it up to a block boundary.
Publish the first non-zero per-layer KV-head count in hybrid-model RankInfo, because layer 0 may be a linear-attention layer with no KV heads.
Reject MixedMambaHybridCacheManager with a non-Python transceiver runtime, since its separately managed attention-KV and Mamba-state pools require the Python routing path.

This PR does not add the V2 Mamba cache manager, save-last reuse, or V2 Mamba disaggregated transfer.

Test Coverage

Full native rebuild for SM100 and SM103
C++ ForceChunkTest.*: 15 passed
Python TestForceChunkPolicy: 14 passed
Focused interval, KV-cache buffer, hybrid RankInfo, and runtime-guard tests: 13 passed
Mixed Mamba Python transceiver coverage across all 8 dtype combinations, plus runtime guards: 10 passed
Native request snapshot-point default, assignment, and copy probe
Complete pre-commit hook set

PR Checklist

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

VALLIS-NERIA · 2026-07-15T13:16:09Z

/bot run

tensorrt-cicd · 2026-07-15T13:22:30Z

PR_Github #59453 [ run ] triggered by Bot. Commit: 7014727 Link to invocation

tensorrt-cicd · 2026-07-15T15:07:43Z

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

Author notes explaining the purpose and scope of each production change block and its focused regression coverage.

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

VALLIS-NERIA · 2026-07-16T07:56:44Z

/bot run

tensorrt-cicd · 2026-07-16T08:02:50Z

PR_Github #59667 [ run ] triggered by Bot. Commit: 5a5b53d Link to invocation

tensorrt-cicd · 2026-07-16T08:19:04Z

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

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

tensorrt-cicd · 2026-07-16T14:44:04Z

PR_Github #59725 [ run ] triggered by Bot. Commit: 85bb37f Link to invocation

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

VALLIS-NERIA · 2026-07-16T15:21:40Z

/bot run

tensorrt-cicd · 2026-07-16T15:27:51Z

PR_Github #59738 [ run ] triggered by Bot. Commit: 605d535 Link to invocation

tensorrt-cicd · 2026-07-16T15:29:47Z

PR_Github #59725 [ run ] completed with state ABORTED. Commit: 85bb37f

Link to invocation

tensorrt-cicd · 2026-07-16T16:03:42Z

PR_Github #59738 [ run ] completed with state FAILURE. Commit: 605d535
/LLM/main/L0_MergeRequest_PR pipeline #48164 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-17T00:48:47Z

/bot run

tensorrt-cicd · 2026-07-17T00:55:18Z

PR_Github #59811 [ run ] triggered by Bot. Commit: f0f16f2 Link to invocation

tensorrt-cicd · 2026-07-17T01:37:14Z

PR_Github #59811 [ run ] completed with state FAILURE. Commit: f0f16f2
/LLM/main/L0_MergeRequest_PR pipeline #48221 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-17T01:39:50Z

/bot run

tensorrt-cicd · 2026-07-17T01:45:20Z

PR_Github #59823 [ run ] triggered by Bot. Commit: f0f16f2 Link to invocation

tensorrt-cicd · 2026-07-17T02:25:19Z

PR_Github #59823 [ run ] completed with state SUCCESS. Commit: f0f16f2
/LLM/main/L0_MergeRequest_PR pipeline #48234 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-17T02:59:40Z

/bot run

tensorrt-cicd · 2026-07-17T03:08:18Z

PR_Github #59850 [ run ] triggered by Bot. Commit: f0f16f2 Link to invocation

tensorrt-cicd · 2026-07-17T03:28:18Z

PR_Github #59850 [ run ] completed with state FAILURE. Commit: f0f16f2
/LLM/main/L0_MergeRequest_PR pipeline #48260 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-17T05:15:34Z

/bot run

tensorrt-cicd · 2026-07-17T05:21:20Z

PR_Github #59883 [ run ] triggered by Bot. Commit: f0f16f2 Link to invocation

BowenFu

LGTM — the fine-grained context-chunk changes preserve the existing default scheduling path: the FORCE_CHUNK capacity-overflow now rounds down to a unit boundary (equivalent to the prior zeroing since chunkSize <= unit), and the mamba snapshot-range addition is a no-op boundary for existing configs. Core scheduler change, but existing behavior is unchanged.

tensorrt-cicd · 2026-07-17T05:42:03Z

PR_Github #59883 [ run ] completed with state FAILURE. Commit: f0f16f2
/LLM/main/L0_MergeRequest_PR pipeline #48286 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-17T05:51:45Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-17T05:58:27Z

PR_Github #59897 [ run ] triggered by Bot. Commit: f0f16f2 Link to invocation

tensorrt-cicd · 2026-07-17T09:47:41Z

PR_Github #59897 [ run ] completed with state SUCCESS. Commit: f0f16f2
/LLM/main/L0_MergeRequest_PR pipeline #48300 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-17T10:51:33Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-17T10:57:16Z

PR_Github #59966 [ run ] triggered by Bot. Commit: f0f16f2 Link to invocation

tensorrt-cicd · 2026-07-17T14:49:30Z

PR_Github #59966 [ run ] completed with state SUCCESS. Commit: f0f16f2
/LLM/main/L0_MergeRequest_PR pipeline #48360 completed with status: 'SUCCESS'

CI Report

Link to invocation

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

VALLIS-NERIA commented Jul 16, 2026

View reviewed changes

VALLIS-NERIA mentioned this pull request Jul 16, 2026

[None][feat] Add V2 Mamba hybrid cache with snapshot reuse #16471

Closed

1 task

VALLIS-NERIA added 7 commits July 16, 2026 15:07

[None][fix] Harden scheduling and disagg lifecycle (By Agent)

f2115d5

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

[None][chore] Apply resource manager test formatting (By Agent)

16e608e

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

[None][fix] Address scheduling and transfer review (By Agent)

3456ec3

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

[None][fix] Align hybrid Mamba C++ disagg routing (By Agent)

75651dd

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

[None][test] Wait for hybrid transfer cancellation (By Agent)

fb13a38

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

[None][test] Cover mixed Mamba Python transfer (By Agent)

47e8a7b

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

[None][test] Cover hybrid runtime auto fallback (By Agent)

5a5b53d

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

VALLIS-NERIA force-pushed the agent/mamba-scheduling-disagg-lifecycle branch from 7014727 to 5a5b53d Compare July 16, 2026 07:42

VALLIS-NERIA changed the title ~~[None][fix] Harden scheduling and disaggregated transfer lifecycle~~ [None][feat] Support fine-grained context chunk management Jul 16, 2026

VALLIS-NERIA mentioned this pull request Jul 16, 2026

[None][fix] Require Python transceiver for mixed Mamba #16428

Closed

1 task

VALLIS-NERIA mentioned this pull request Jul 16, 2026

[None][feat] Add V2 Mamba hybrid cache with snapshot reuse VALLIS-NERIA/TensorRT-LLM#8

Draft

1 task

[None][fix] Address PR1 review feedback (By Agent)

d935140

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

VALLIS-NERIA commented Jul 16, 2026

View reviewed changes

VALLIS-NERIA marked this pull request as ready for review July 16, 2026 12:25

VALLIS-NERIA requested review from a team as code owners July 16, 2026 12:25

VALLIS-NERIA requested review from JunyiXu-nv, Tabrizian and zhaoyangwang-nvidia July 16, 2026 12:25

VALLIS-NERIA changed the title ~~[None][feat] Support fine-grained context chunk management~~ [TRTLLM-11875][feat] Support fine-grained context chunk management Jul 16, 2026

[None][refactor] Consolidate Mamba snapshot point preparation (By Agent)

605d535

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

Merge branch 'main' into agent/mamba-scheduling-disagg-lifecycle

f0f16f2

BowenFu approved these changes Jul 17, 2026

View reviewed changes

Shixiaowei02 approved these changes Jul 17, 2026

View reviewed changes

Shixiaowei02 requested a review from bo-nv July 17, 2026 06:05

bo-nv approved these changes Jul 17, 2026

View reviewed changes

Uh oh!

Conversation

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

Uh oh!

Summary by CodeRabbit

Description

Snapshot-aware context chunking

KV-cache fixes

Test Coverage

PR Checklist

Uh oh!

VALLIS-NERIA commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

VALLIS-NERIA left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

VALLIS-NERIA commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

VALLIS-NERIA commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

VALLIS-NERIA commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

VALLIS-NERIA commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

VALLIS-NERIA commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

VALLIS-NERIA commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

BowenFu left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

VALLIS-NERIA commented Jul 17, 2026

Uh oh!

VALLIS-NERIA commented Jul 15, 2026 •

edited by coderabbitai Bot

Loading