[None][fix] Fix FlashInfer shared-KV speculative decode by 2ez4bz · Pull Request #17264 · NVIDIA/TensorRT-LLM

2ez4bz · 2026-08-04T19:41:48Z

Dev Engineer Review

One-model MTP scheduling synchronizes Python-side and C++ draft-token state.
Drafting stops when requests can exceed the legal target-position limit.
Speculative KV accounting uses the prompt KV boundary for the first generation step when dynamic KV-length correction is unavailable.
FlashInfer separates logical KV lengths from reserved generation-page capacity.
KV-length offsets are applied and restored across speculative decoding and CUDA graph execution.
These changes prevent overstated cached lengths, invalid FlashInfer page access, and incorrect KV scheduler capacity planning.
Added reserved-block-count APIs to KVCacheManager and KVCacheManagerV2.
Added context-manager support to BaseResourceManager.
No configuration or test-list files changed.

QA Engineer Review

Added tests for one-model MTP draft-token normalization, Python draft-token preservation, and sequence-position limits.
Added tests for prompt-boundary KV accounting and FlashInfer KV-offset application and restoration.
The modified test functions are not referenced in tests/integration/test_lists/ based on the available repository coverage entries.
Verdict: needs follow-up.

Description

Why?

Overlap-scheduled speculative decoding with a shared KV cache can use
draft positions beyond the logical sequence length. FlashInfer exposed
only logical generation pages, and did not reflect scheduler offsets in
its live KV metadata, while Gemma4 sized RoPE only to the logical limit.

What?

Expose every reserved generation page for shared-KV overlap decoding
while tracking logical KV lengths separately. Apply and restore overlap
scheduler offsets to append positions and trtllm-gen decode lengths, and
add speculative RoPE headroom to Gemma4 target and assistant models

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.

ghost · 2026-08-04T19:42:00Z

/bot run

ghost · 2026-08-04T19:45:51Z

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

Walkthrough

The change updates speculative decoding across executor scheduling and FlashInfer KV-cache handling. It preserves Python draft tokens, applies and restores backend KV-length offsets, separates logical lengths from reserved page-table capacity, and adds regression coverage.

Changes

Speculative decoding and KV-cache updates

Layer / File(s)	Summary
Logical KV lengths and generation page tables `tensorrt_llm/_torch/attention_backend/flashinfer.py`, `tensorrt_llm/_torch/metadata.py`, `tensorrt_llm/_torch/pyexecutor/resource_manager.py`, `tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py`, `tests/unittest/_torch/attention/test_flashinfer_attention.py`	FlashInfer tracks logical KV lengths separately from reserved generation page-table capacity. Speculative offsets update positions and decode-wrapper buffers. Cache managers expose reserved block counts.
Executor integration and cache boundaries `tensorrt_llm/_torch/pyexecutor/model_engine.py`, `tests/unittest/_torch/executor/test_pytorch_model_engine.py`	The executor applies and restores backend offsets, enables full generation page tables when supported, and uses the prompt boundary for the first overlapping generation iteration.
One-model MTP draft state and sequence-limit safety `tensorrt_llm/_torch/pyexecutor/py_executor.py`, `tests/unittest/_torch/executor/test_py_executor.py`	One-model MTP disables drafting when target positions reach legal limits and preserves populated Python draft tokens. Tests cover draft normalization and sequence limits.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Scheduler
  participant ModelEngine
  participant FlashInfer
  participant DecodeWrapper
  Scheduler->>ModelEngine: build speculative generation metadata
  ModelEngine->>FlashInfer: apply KV-length offsets
  FlashInfer->>DecodeWrapper: publish logical KV lengths
  ModelEngine->>FlashInfer: prepare and plan decode
  ModelEngine->>FlashInfer: restore KV-length offsets

Suggested labels: api-compatible

Suggested reviewers: schetlur-nv, chienchunhung, liji-nv

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 27.78% 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 clearly identifies the primary fix for FlashInfer shared-KV speculative decoding.
Description check	✅ Passed	The description explains the problem and solution and includes the checklist, but the Test Coverage section is not filled in.

✨ Finishing Touches 💡 1

⚔️ Resolve merge conflicts 💡

Resolve merge conflict in branch dev-flashinfer-mtp-ima

🧪 Generate unit tests (beta)

Create PR with unit tests

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

ghost

Actionable comments posted: 1

🤖 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 `@tests/unittest/_torch/executor/test_py_executor.py`:
- Around line 2010-2013: Update the draft-token assertions in the affected
executor test to verify values, not only lengths: compare both
gen.py_draft_tokens and disagg_gen.py_draft_tokens against [0] *
self.MAX_TOTAL_DRAFT_TOKENS. Keep the existing num_draft_tokens assertions
unchanged.

🪄 Autofix

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: cfdc40e2-408d-41fd-8732-7561ca155193

📥 Commits

Reviewing files that changed from the base of the PR and between be93500 and 0baecbe.

📒 Files selected for processing (4)

tensorrt_llm/_torch/pyexecutor/model_engine.py
tensorrt_llm/_torch/pyexecutor/py_executor.py
tests/unittest/_torch/executor/test_py_executor.py
tests/unittest/_torch/executor/test_pytorch_model_engine.py

ghost · 2026-08-04T19:48:08Z

PR_Github #63826 [ run ] triggered by Bot. Commit: 0baecbe Link to invocation

ghost · 2026-08-04T21:06:37Z

PR_Github #63826 [ run ] completed with state SUCCESS. Commit: 0baecbe
/LLM/main/L0_MergeRequest_PR pipeline #51767 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

ghost · 2026-08-05T01:30:25Z

Could you clarify how these two hunks relate?

On the one-model path, py_draft_tokens now contains dummy values when _prepare_tp_inputs first reads it. Please explain why setting only the C++ draft_tokens count would not satisfy scheduler accounting.
Is the py_decoding_iter == 0 cached_token_num correction required because that list is now non-empty? If the changes are coupled, please document the invariant and add a focused test. If they fix independent bugs, separate commits would make them easier to validate and revert.

The kv_lens_cuda backend guard looks appropriately scoped.

ghost · 2026-08-05T16:54:07Z

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

ghost

Actionable comments posted: 1

🤖 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 `@tests/unittest/_torch/executor/test_py_executor.py`:
- Around line 2036-2037: Update the sampler draft setup in the relevant executor
test to use a nonempty list shorter than self.MAX_TOTAL_DRAFT_TOKENS, while
retaining the assertion that request.draft_tokens uses the full scheduler budget
and request.py_draft_tokens preserves the supplied shorter list.

🪄 Autofix

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: 7989f080-973b-451a-9155-656f54711e18

📥 Commits

Reviewing files that changed from the base of the PR and between 1dfb7b1 and 990299f.

📒 Files selected for processing (4)

tensorrt_llm/_torch/pyexecutor/model_engine.py
tensorrt_llm/_torch/pyexecutor/py_executor.py
tests/unittest/_torch/executor/test_py_executor.py
tests/unittest/_torch/executor/test_pytorch_model_engine.py

🚧 Files skipped from review as they are similar to previous changes (3)

tests/unittest/_torch/executor/test_pytorch_model_engine.py
tensorrt_llm/_torch/pyexecutor/py_executor.py
tensorrt_llm/_torch/pyexecutor/model_engine.py

ghost · 2026-08-05T17:24:33Z

@BowenFu

These changes address two parts of the same MTP + FlashInfer + block-reuse failure path.

The python and C++ draft-token representations must both be populated because they have different consumers: the C++ micro-batch scheduler uses draft_tokens, while python scheduling and model-input preparation use py_draft_tokens.

The first-generation cached_token_num correction handles the same context-to-generation transition: a previous overlap tensor exists, but no speculative target forward has populated draft-token KV entries yet. Without using the prompt boundary, FlashInfer can access an invalid reused page. The focused tests cover the two state assumptions independently.

I’m keeping them together because both are needed for this MTP + FlashInfer + block-reuse fix.

ghost · 2026-08-05T17:26:14Z

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

ghost · 2026-08-05T17:27:35Z

/bot run

ghost · 2026-08-05T17:33:09Z

PR_Github #64088 [ run ] triggered by Bot. Commit: 9ff8de8 Link to invocation

ghost · 2026-08-05T19:45:06Z

PR_Github #64088 [ run ] completed with state SUCCESS. Commit: 9ff8de8
/LLM/main/L0_MergeRequest_PR pipeline #52015 completed with status: 'FAILURE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, then re-trigger CI with the same bot command (no rebase needed).

⚠️ 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

ghost · 2026-08-06T00:47:15Z

/bot run

ghost · 2026-08-06T00:53:24Z

PR_Github #64146 [ run ] triggered by Bot. Commit: 9ff8de8 Link to invocation

ghost · 2026-08-06T02:58:19Z

PR_Github #64146 [ run ] completed with state FAILURE. Commit: 9ff8de8
/LLM/main/L0_MergeRequest_PR pipeline #52064 completed with status: 'FAILURE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, then re-trigger CI with the same bot command (no rebase needed).

⚠️ 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

ghost · 2026-08-06T03:58:32Z

/bot run

ghost · 2026-08-06T04:20:08Z

PR_Github #64208 [ run ] triggered by Bot. Commit: 9ff8de8 Link to invocation

ghost · 2026-08-06T07:06:44Z

PR_Github #64208 [ run ] completed with state FAILURE. Commit: 9ff8de8
/LLM/main/L0_MergeRequest_PR pipeline #52120 completed with status: 'FAILURE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, then re-trigger CI with the same bot command (no rebase needed).

⚠️ 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

ghost · 2026-08-20T18:22:49Z

PR_Github #67910 [ run ] triggered by Bot. Commit: 80ed499 Link to invocation

ghost

lgtm

ghost · 2026-08-20T20:02:20Z

max_draft_len=5 + forced acceptance length also surfaces another IMA, which should be fixed once we update flashinfer to include flashinfer-ai/flashinfer#4382 .

ghost · 2026-08-20T23:46:44Z

/bot run

ghost · 2026-08-20T23:53:05Z

PR_Github #68018 [ run ] triggered by Bot. Commit: 80ed499 Link to invocation

ghost · 2026-08-20T23:53:12Z

PR_Github #67910 [ run ] completed with state ABORTED. Commit: 80ed499

Link to invocation

ghost · 2026-08-21T01:37:43Z

PR_Github #68018 [ run ] completed with state SUCCESS. Commit: 80ed499
/LLM/main/L0_MergeRequest_PR pipeline #55471 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

ghost · 2026-08-21T04:36:56Z

/bot run

ghost · 2026-08-21T04:42:42Z

PR_Github #68140 [ run ] triggered by Bot. Commit: 80ed499 Link to invocation

ghost · 2026-08-21T04:58:44Z

PR_Github #68140 [ run ] completed with state SUCCESS. Commit: 80ed499
/LLM/main/L0_MergeRequest_PR pipeline #55588 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

* Why? Overlap-scheduled speculative decoding with a shared KV cache can use draft positions beyond the logical sequence length. FlashInfer exposed only logical generation pages, and did not reflect scheduler offsets in its live KV metadata, while Gemma4 sized RoPE only to the logical limit. * What? Expose every reserved generation page for shared-KV overlap decoding while tracking logical KV lengths separately. Apply and restore overlap scheduler offsets to append positions and trtllm-gen decode lengths, and add speculative RoPE headroom to Gemma4 target and assistant models Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>

* Why? Cached decode plans ignored query length and generation batch size, so speculative decoding could reuse a plan built for an incompatible launch shape. * What? Include both dimensions in the plan parameters, pass the query length to FlashInfer, and reject batches with nonuniform generation query lengths. Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>

ghost · 2026-08-21T07:32:35Z

/bot run

ghost · 2026-08-21T07:38:14Z

PR_Github #68210 [ run ] triggered by Bot. Commit: 9ef9c81 Link to invocation

ghost · 2026-08-21T10:19:48Z

PR_Github #68210 [ run ] completed with state FAILURE. Commit: 9ef9c81
/LLM/main/L0_MergeRequest_PR pipeline #55650 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

ghost · 2026-08-21T15:58:09Z

/bot run

ghost · 2026-08-21T16:04:19Z

PR_Github #68321 [ run ] triggered by Bot. Commit: 9ef9c81 Link to invocation

ghost · 2026-08-21T19:11:07Z

PR_Github #68321 [ run ] completed with state SUCCESS. Commit: 9ef9c81
/LLM/main/L0_MergeRequest_PR pipeline #55748 completed with status: 'UNSTABLE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, then re-trigger CI with the same bot command (no rebase needed).

⚠️ 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

ghost · 2026-08-21T20:13:31Z

/bot run

ghost · 2026-08-21T20:20:25Z

PR_Github #68403 [ run ] triggered by Bot. Commit: 9ef9c81 Link to invocation

ghost · 2026-08-21T22:57:50Z

PR_Github #68403 [ run ] completed with state SUCCESS. Commit: 9ef9c81
/LLM/main/L0_MergeRequest_PR pipeline #55822 completed with status: 'SUCCESS'

CI Report

Link to invocation

2ez4bz requested review from a team as code owners August 4, 2026 19:41

2ez4bz requested review from JunyiXu-nv, Tabrizian, allisonlim-nv and chienchunhung August 4, 2026 19:41

github-actions Bot assigned 2ez4bz Aug 4, 2026

coderabbitai Bot reviewed Aug 4, 2026

View reviewed changes

Comment thread tests/unittest/_torch/executor/test_py_executor.py Outdated

2ez4bz force-pushed the dev-flashinfer-mtp-ima branch from 0baecbe to 990299f Compare August 5, 2026 16:53

coderabbitai Bot reviewed Aug 5, 2026

View reviewed changes

Comment thread tests/unittest/_torch/executor/test_py_executor.py Outdated

2ez4bz force-pushed the dev-flashinfer-mtp-ima branch from 990299f to 9ff8de8 Compare August 5, 2026 17:26

coderabbitai Bot mentioned this pull request Aug 6, 2026

[None][fix] Add KV cache V2 recompute pause path #15252

Merged

4 tasks

coderabbitai Bot mentioned this pull request Aug 6, 2026

[None][fix] fix disagg overlap slot headroom without MTP #17282

Merged

1 task

aswinvisva approved these changes Aug 20, 2026

View reviewed changes

allisonlim-nv approved these changes Aug 20, 2026

View reviewed changes

William Zhang added 2 commits August 21, 2026 00:32

2ez4bz force-pushed the dev-flashinfer-mtp-ima branch from 80ed499 to 9ef9c81 Compare August 21, 2026 07:32

chang-l added the ci: full pre-merge approved label Aug 21, 2026

2ez4bz merged commit 5264ed5 into NVIDIA:main Aug 21, 2026

yuxianq mentioned this pull request Sep 3, 2026

[https://nvbugs/6689016][fix] Limit eager FlashInfer plan cache growth #18557

Open

1 task

Uh oh!

Conversation

2ez4bz commented Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Dev Engineer Review

QA Engineer Review

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

ghost commented Aug 4, 2026

Uh oh!

ghost commented Aug 4, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Reviews paused

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (1 warning)

Uh oh!

ghost left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ghost commented Aug 4, 2026

Uh oh!

ghost commented Aug 4, 2026

Uh oh!

ghost commented Aug 5, 2026 • edited by BowenFu Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ghost commented Aug 5, 2026

Uh oh!

ghost left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ghost commented Aug 5, 2026

Uh oh!

ghost commented Aug 5, 2026

Uh oh!

ghost commented Aug 5, 2026

Uh oh!

ghost commented Aug 5, 2026

Uh oh!

ghost commented Aug 5, 2026

Uh oh!

ghost commented Aug 6, 2026

Uh oh!

ghost commented Aug 6, 2026

Uh oh!

ghost commented Aug 6, 2026

Uh oh!

ghost commented Aug 6, 2026

Uh oh!

ghost commented Aug 6, 2026

Uh oh!

ghost commented Aug 6, 2026

Uh oh!

ghost commented Aug 20, 2026

Uh oh!

ghost left a comment

Choose a reason for hiding this comment

Uh oh!

ghost commented Aug 20, 2026

Uh oh!

ghost commented Aug 20, 2026

Uh oh!

ghost commented Aug 20, 2026

Uh oh!

ghost commented Aug 20, 2026

Uh oh!

ghost commented Aug 21, 2026

Uh oh!

ghost commented Aug 21, 2026

Uh oh!

ghost commented Aug 21, 2026

Uh oh!

ghost commented Aug 21, 2026

2ez4bz commented Aug 4, 2026 •

edited

Loading

ghost commented Aug 4, 2026 •

edited by coderabbitai Bot

Loading

ghost commented Aug 5, 2026 •

edited by BowenFu

Loading