[None][fix] use prompt lookahead for MTP Eagle chunked prefill by yizhang-nv · Pull Request #18295 · NVIDIA/TensorRT-LLM

yizhang-nv · 2026-08-27T05:24:53Z

Description

MTP Eagle left-shifts each context chunk before the first draft-layer forward. For a non-final chunk, the shifted tail must be the immediate next prompt token. It previously used the target model's sampled token, so linear and dynamic-tree MTP Eagle could persist draft-layer KV that did not match canonical full-prompt prefill.

This change:

stages exactly one immediate prompt lookahead token per active context request in a fixed-address SpecMetadata tensor with shape [max_num_requests]; -1 marks final chunks with no valid lookahead;
derives the token from the request's context cursor, independently of KV cache manager V1 or V2;
uses the prompt token for non-final context rows in linear, dynamic-tree, and external shared-target-KV MTP Eagle, while final chunks fall back to the sampled token and generation rows remain unchanged; and
leaves vanilla MTP, cache keys, and V2 draft reuse unchanged for follow-up work.

The external shared-target-KV path does not persist independent draft KV, but it uses the same token/hidden-state/position contract so all MTP-Eagle rollout paths now start from the correctly aligned token.

Acceptance Measurements

Measured on one B200 with DeepSeek-V3-Lite/bf16, TP=1, MTP Eagle one-model, max draft length 1, chunked prefill enabled, prompt length 512, chunk size 128, and output length 128. The workload used three prompts sourced from the TensorRT-LLM documentation and two runs per configuration.

Configuration	Accepted / drafted tokens	Acceptance rate	Acceptance length
Before	338 / 430	78.60%	1.7860
After	339 / 429	79.02%	1.7902

The aggregate delta is +0.42 percentage points in acceptance rate and +0.0042 in acceptance length. Per-run acceptance-rate ranges overlap (before: 76.96%-80.28%; after: 78.60%-79.44%), so this should be treated as neutral/small with no observed regression, rather than as a statistically meaningful improvement.

Test Coverage

pytest -s -q tests/unittest/_torch/executor/test_pytorch_model_engine.py tests/unittest/_torch/speculative/test_eagle3.py tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py
- 128 passed, 22 skipped, 26 subtests passed
pre-commit run --files on all seven changed files
- all hooks passed

The focused tests cover prompt-boundary extraction, token ID 0, sentinel fallback and stale-value clearing, CUDA graph metadata allocation, linear and dynamic-tree context input, and external shared-target-KV context/generation alignment.

PR Checklist

The fix is independent of KV cache manager V1/V2.
Linear, dynamic-tree, and external shared-target-KV MTP-Eagle paths are covered.
Vanilla MTP and V2 draft reuse are intentionally unchanged.
No public API, dependency, ownership, documentation, or architecture-diagram change is introduced.

Dev Engineer Review

The implementation stores one immediate prompt lookahead token per active context request.
-1 marks final chunks and preserves sampled-token fallback behavior.
The implementation is independent of KV cache manager V1 and V2.
Linear, dynamic-tree, and shared-target-KV MTP Eagle paths use the lookahead token.
Vanilla MTP, cache keys, and V2 draft reuse remain unchanged.
The SpecMetadata API now uses a flat one-token-per-request buffer.
No configuration or test-list files changed.
Focused tests passed: 128 passed, 22 skipped, and 26 subtests passed.
Pre-commit checks passed for all seven changed files.
Reported measurements improved acceptance rate from 78.60% to 79.02%, with overlapping per-run ranges and no observed regression.

QA Engineer Review

Modified test files:
- test_pytorch_model_engine.py: updated prompt-boundary and overlap-mock coverage.
- test_eagle3.py: updated one-token lookahead, sentinel, CUDA-graph, draft-input, and dynamic-tree coverage.
- test_mtp.py: updated shared-KV draft-input coverage for multiple context requests.
The tests cover non-final lookahead tokens, final-chunk sentinel fallback, accepted tokens, draft IDs, recurrent states, and positions.
No corresponding test-db/ or qa/ test-list changes were provided.
The modified test functions are not mapped in the provided test-list data.
Verdict: needs follow-up. CBTS coverage mapping is unavailable.

coderabbitai · 2026-08-27T05:35:44Z

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

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 5a723d80-53b5-4e90-bb01-0ff484e71dfe

📥 Commits

Reviewing files that changed from the base of the PR and between 07c5f21 and c834e7a.

📒 Files selected for processing (7)

tensorrt_llm/_torch/pyexecutor/model_engine.py
tensorrt_llm/_torch/speculative/eagle3.py
tensorrt_llm/_torch/speculative/interface.py
tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py
tests/unittest/_torch/executor/test_pytorch_model_engine.py
tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py
tests/unittest/_torch/speculative/test_eagle3.py

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

tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py
tensorrt_llm/_torch/pyexecutor/model_engine.py
tests/unittest/_torch/speculative/test_eagle3.py
tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py
tests/unittest/_torch/executor/test_pytorch_model_engine.py
tensorrt_llm/_torch/speculative/eagle3.py
tensorrt_llm/_torch/speculative/interface.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

Walkthrough

The change replaces padded prompt lookahead rows with one immediate token per context request. Boundary cases use INVALID_PROMPT_LOOKAHEAD_TOKEN. MTP Eagle and dynamic-tree preparation consume the flattened representation.

Changes

Context prompt lookahead

Layer / File(s)	Summary
Prompt lookahead metadata and input selection `tensorrt_llm/_torch/speculative/interface.py`	`SpecMetadata` stores one sentinel-filled CUDA token per request. Context input preparation uses the lookahead token when valid and the accepted token when the sentinel marks the final chunk.
Context lookahead collection `tensorrt_llm/_torch/pyexecutor/model_engine.py`, `tests/unittest/_torch/executor/test_pytorch_model_engine.py`	The model engine extracts one token after each context chunk. Prompt-boundary chunks receive one invalid-token sentinel. Tests cover valid tokens, boundaries, and metadata setup.
Eagle draft input integration and validation `tensorrt_llm/_torch/speculative/eagle3.py`, `tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py`, `tests/unittest/_torch/speculative/test_eagle3.py`, `tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py`	MTP Eagle and dynamic-tree step-0 preparation forward flattened prompt lookahead tokens. Shared-KV preparation preserves accepted-token fallback for sentinel entries. Tests cover initialization, fallback behavior, dynamic-tree inputs, and updated draft outputs.

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

Merge Risk: ⚪ Minimal · up to c834e

This localized change aligns MTP Eagle chunked-prefill inputs with the next prompt token while preserving generation and other MTP paths; no actionable merge-blocking risk remains after normal checks and review.

Sequence Diagram(s)

sequenceDiagram
  participant ModelEngine
  participant SpecMetadata
  participant MTPEagleWorker
  participant DynamicTree
  participant ContextInputPreparation

  ModelEngine->>ModelEngine: Extract one token after each context chunk
  ModelEngine->>SpecMetadata: Store token or INVALID_PROMPT_LOOKAHEAD_TOKEN
  MTPEagleWorker->>SpecMetadata: Read flat lookahead buffer
  MTPEagleWorker->>ContextInputPreparation: Pass MTP lookahead tokens
  DynamicTree->>SpecMetadata: Read flat lookahead buffer
  DynamicTree->>ContextInputPreparation: Pass step-0 lookahead tokens
  ContextInputPreparation->>ContextInputPreparation: Select lookahead or accepted-token fallback

Suggested reviewers: lori-ren, mikeiovine, nv-xtf

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 56.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 25 functions across 7 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description check	✅ Passed	The description explains the issue, solution, acceptance measurements, focused test coverage, and checklist decisions. It is complete and directly aligned with the requested template.
Title check	✅ Passed	The title uses the required [None][fix] format and clearly identifies the primary change: using prompt lookahead for MTP Eagle chunked prefill.
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.

Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

yizhang-nv · 2026-08-27T06:32:12Z

/bot run

tensorrt-cicd · 2026-08-27T06:38:47Z

PR_Github #69655 [ run ] triggered by Bot. Commit: 345f71d Link to invocation

yizhang-nv · 2026-08-27T07:11:35Z

/bot run

tensorrt-cicd · 2026-08-27T07:16:50Z

PR_Github #69663 [ run ] triggered by Bot. Commit: 32a287f Link to invocation

tensorrt-cicd · 2026-08-27T07:18:46Z

PR_Github #69655 [ run ] completed with state ABORTED. Commit: 345f71d

Link to invocation

zhaoyangwang-nvidia

Could you add measured acceptance-rate / acceptance-length numbers, before and after this change, to the description — along with the prompt length and chunk size used?

tensorrt-cicd · 2026-08-27T09:40:12Z

PR_Github #69663 [ run ] completed with state SUCCESS. Commit: 32a287f
/LLM/main/L0_MergeRequest_PR pipeline #56967 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

coderabbitai · 2026-08-28T10:03:18Z

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.

yizhang-nv · 2026-08-28T10:04:07Z

@zhaoyangwang-nvidia Thanks for the review. I added the measured acceptance-rate and acceptance-length results, including prompt length and chunk size, to the PR description.

On one B200 with DeepSeek-V3-Lite BF16, TP=1, MTP Eagle one-model, max draft length 1, prompt length 512, chunk size 128, and output length 128, using 3 documentation prompts and 2 runs per configuration:

before: 338/430 accepted draft tokens, AR 78.60%, acceptance length 1.7860
after: 339/429 accepted draft tokens, AR 79.02%, acceptance length 1.7902

That is +0.42 pp AR and +0.0042 acceptance length. The per-run AR ranges overlap (76.96%-80.28% before and 78.60%-79.44% after), so I described the impact as neutral/small with no observed regression rather than claiming a statistically meaningful improvement.

I also addressed the inline cleanup suggestions, rebased onto the latest origin/main, and reran the related suites: 128 passed, 22 skipped, and 26 subtests passed.

yizhang-nv · 2026-08-28T10:34:29Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-28T10:41:06Z

PR_Github #69963 [ run ] triggered by Bot. Commit: 52a4c92 Link to invocation

tensorrt-cicd · 2026-08-28T18:18:17Z

PR_Github #69963 [ run ] completed with state SUCCESS. Commit: 52a4c92
/LLM/main/L0_MergeRequest_PR pipeline #57244 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 · 2026-08-29T06:52:58Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-29T06:58:10Z

PR_Github #70131 [ run ] triggered by Bot. Commit: 52a4c92 Link to invocation

tensorrt-cicd · 2026-08-29T08:45:19Z

PR_Github #70131 [ run ] completed with state SUCCESS. Commit: 52a4c92
/LLM/main/L0_MergeRequest_PR pipeline #57395 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

yizhang-nv · 2026-08-31T04:22:52Z

/bot run --disable-fail-fast

yizhang-nv · 2026-08-31T09:55:28Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-31T10:02:44Z

PR_Github #70324 [ run ] triggered by Bot. Commit: 52a4c92 Link to invocation

tensorrt-cicd · 2026-08-31T12:38:36Z

PR_Github #70324 [ run ] completed with state SUCCESS. Commit: 52a4c92
/LLM/main/L0_MergeRequest_PR pipeline #57562 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 · 2026-08-31T15:26:39Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-31T15:32:53Z

PR_Github #70378 [ run ] triggered by Bot. Commit: 52a4c92 Link to invocation

tensorrt-cicd · 2026-08-31T17:55:40Z

PR_Github #70378 [ run ] completed with state SUCCESS. Commit: 52a4c92
/LLM/main/L0_MergeRequest_PR pipeline #57611 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 · 2026-08-31T19:29:16Z

/bot run --disable-fail-fast

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>

coderabbitai · 2026-09-01T03:30:40Z

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.

yizhang-nv · 2026-09-01T03:30:54Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-09-01T03:38:12Z

PR_Github #70561 [ run ] triggered by Bot. Commit: c834e7a Link to invocation

tensorrt-cicd · 2026-09-01T07:45:28Z

PR_Github #70561 [ run ] completed with state FAILURE. Commit: c834e7a
/LLM/main/L0_MergeRequest_PR pipeline #57773 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 · 2026-09-01T07:45:43Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-09-01T07:53:11Z

PR_Github #70629 [ run ] triggered by Bot. Commit: c834e7a Link to invocation

tensorrt-cicd · 2026-09-01T08:38:06Z

PR_Github #70629 [ run ] completed with state SUCCESS. Commit: c834e7a
/LLM/main/L0_MergeRequest_PR pipeline #57832 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 · 2026-09-01T08:46:41Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-09-01T08:57:51Z

PR_Github #70655 [ run ] triggered by Bot. Commit: c834e7a Link to invocation

tensorrt-cicd · 2026-09-01T11:53:07Z

PR_Github #70655 [ run ] completed with state FAILURE. Commit: c834e7a
/LLM/main/L0_MergeRequest_PR pipeline #57858 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 · 2026-09-01T11:59:17Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-09-01T12:05:32Z

PR_Github #70686 [ run ] triggered by Bot. Commit: c834e7a Link to invocation

tensorrt-cicd · 2026-09-01T12:52:55Z

PR_Github #70686 [ run ] completed with state SUCCESS. Commit: c834e7a
/LLM/main/L0_MergeRequest_PR pipeline #57884 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 · 2026-09-01T13:03:25Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-09-01T13:09:23Z

PR_Github #70693 [ run ] triggered by Bot. Commit: c834e7a Link to invocation

tensorrt-cicd · 2026-09-01T14:00:27Z

PR_Github #70693 [ run ] completed with state SUCCESS. Commit: c834e7a
/LLM/main/L0_MergeRequest_PR pipeline #57891 completed with status: 'SUCCESS'

CI Report

Link to invocation

yizhang-nv requested review from a team as code owners August 27, 2026 05:24

yizhang-nv requested review from lori-ren, mikeiovine, nv-xtf and reasonsolo August 27, 2026 05:24

github-actions Bot assigned yizhang-nv Aug 27, 2026

yizhang-nv requested a review from zhaoyangwang-nvidia August 27, 2026 07:13

reasonsolo approved these changes Aug 27, 2026

View reviewed changes

zhaoyangwang-nvidia reviewed Aug 27, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/pyexecutor/model_engine.py Outdated

Comment thread tensorrt_llm/_torch/speculative/interface.py Outdated

Comment thread tensorrt_llm/_torch/speculative/eagle3.py Outdated

Comment thread tensorrt_llm/_torch/speculative/eagle3.py Outdated

zhaoyangwang-nvidia approved these changes Aug 27, 2026

View reviewed changes

Shixiaowei02 approved these changes Aug 27, 2026

View reviewed changes

yizhang-nv force-pushed the codex/fix-mtp-eagle-chunk-lookahead branch from 32a287f to 52a4c92 Compare August 28, 2026 10:03

yuxianq added the ci: full pre-merge approved label Aug 31, 2026

yizhang-nv added 5 commits August 31, 2026 20:07

[None][fix] use prompt lookahead for MTP Eagle prefill

2297293

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>

[None][fix] use sentinel for MTP prompt lookahead

aa3767b

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>

[None][fix] simplify MTP Eagle prompt lookahead

95a12b2

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>

[None][fix] clarify MTP Eagle draft input

a0cab97

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>

fix MTP Eagle chunked prefill lookahead

c834e7a

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>

yizhang-nv force-pushed the codex/fix-mtp-eagle-chunk-lookahead branch from 52a4c92 to c834e7a Compare September 1, 2026 03:30

yizhang-nv merged commit e00d95f into NVIDIA:main Sep 1, 2026
10 checks passed

Conversation

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

Uh oh!

Description

Acceptance Measurements

Test Coverage

PR Checklist

Dev Engineer Review

QA Engineer Review

Uh oh!

coderabbitai Bot commented Aug 27, 2026 • edited 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!

yizhang-nv commented Aug 27, 2026

Uh oh!

tensorrt-cicd commented Aug 27, 2026

Uh oh!

yizhang-nv commented Aug 27, 2026

Uh oh!

tensorrt-cicd commented Aug 27, 2026

Uh oh!

tensorrt-cicd commented Aug 27, 2026

Uh oh!

zhaoyangwang-nvidia left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tensorrt-cicd commented Aug 27, 2026

Uh oh!

coderabbitai Bot commented Aug 28, 2026

Uh oh!

yizhang-nv commented Aug 28, 2026

Uh oh!

yizhang-nv commented Aug 28, 2026

Uh oh!

tensorrt-cicd commented Aug 28, 2026

Uh oh!

tensorrt-cicd commented Aug 28, 2026

Uh oh!

yizhang-nv commented Aug 29, 2026

Uh oh!

tensorrt-cicd commented Aug 29, 2026

Uh oh!

tensorrt-cicd commented Aug 29, 2026

Uh oh!

yizhang-nv commented Aug 31, 2026

Uh oh!

yizhang-nv commented Aug 31, 2026

Uh oh!

tensorrt-cicd commented Aug 31, 2026

Uh oh!

tensorrt-cicd commented Aug 31, 2026

Uh oh!

yizhang-nv commented Aug 31, 2026

Uh oh!

tensorrt-cicd commented Aug 31, 2026

Uh oh!

tensorrt-cicd commented Aug 31, 2026

Uh oh!

yizhang-nv commented Aug 31, 2026

Uh oh!

coderabbitai Bot commented Sep 1, 2026

Uh oh!

yizhang-nv commented Sep 1, 2026

Uh oh!

tensorrt-cicd commented Sep 1, 2026

Uh oh!

tensorrt-cicd commented Sep 1, 2026

Uh oh!

yizhang-nv commented Sep 1, 2026

Uh oh!

tensorrt-cicd commented Sep 1, 2026

yizhang-nv commented Aug 27, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Aug 27, 2026 •

edited

Loading