[None][fix] use prompt lookahead for MTP Eagle chunked prefill - #18295
[None][fix] use prompt lookahead for MTP Eagle chunked prefill#18295yizhang-nv wants to merge 5 commits into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (7)
🚧 Files skipped from review as they are similar to previous changes (7)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. WalkthroughThe change replaces padded prompt lookahead rows with one immediate token per context request. Boundary cases use ChangesContext prompt lookahead
Estimated code review effort: 3 (Moderate) | ~20 minutes Merge Risk: ⚪ Minimal · up to This localized change aligns MTP Eagle chunked-prefill inputs with the immediate prompt token while preserving final-chunk and generation behavior; 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: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Description checkExplanation The description includes a clear problem statement, solution details, acceptance measurements, focused test coverage, and a completed checklist. It explains the affected MTP Eagle paths and the intentionally unchanged behavior.
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
/bot run |
|
PR_Github #69655 [ run ] triggered by Bot. Commit: |
|
/bot run |
|
PR_Github #69663 [ run ] triggered by Bot. Commit: |
|
PR_Github #69655 [ run ] completed with state |
zhaoyangwang-nvidia
left a comment
zhaoyangwang-nvidia
left a comment
There was a problem hiding this comment.
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?
| if spec_metadata.context_prompt_lookahead_tokens is not None: | ||
| # No-cache context inputs contain the complete prompt, so | ||
| # there is never a valid token beyond the current chunk. | ||
| spec_metadata.populate_context_prompt_lookahead( |
There was a problem hiding this comment.
_set_up_spec_metadata's no_cache branch constructs a fresh SpecMetadata on every call rather than reusing self.spec_metadata, so the buffer here was just filled with INVALID_PROMPT_LOOKAHEAD_TOKEN by allocate_context_prompt_lookahead and this call always writes -1 over -1. Is there a path where a stale value could actually reach here? If not, I'd drop the block rather than leave a no-op that implies one exists.
|
PR_Github #69663 [ run ] completed with state
|
Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
32a287f to
52a4c92
Compare
|
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. |
|
@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:
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. |
|
/bot run --disable-fail-fast |
|
PR_Github #69963 [ run ] triggered by Bot. Commit: |
|
PR_Github #69963 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #70131 [ run ] triggered by Bot. Commit: |
|
PR_Github #70131 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #70249 [ run ] triggered by Bot. Commit: |
|
PR_Github #70249 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #70324 [ run ] triggered by Bot. Commit: |
|
PR_Github #70324 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #70378 [ run ] triggered by Bot. Commit: |
|
PR_Github #70378 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
6 participants
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:
SpecMetadatatensor with shape[max_num_requests];-1marks final chunks with no valid lookahead;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.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.pypre-commit run --fileson all seven changed filesThe 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
Dev Engineer Review
-1marks final chunks and preserves sampled-token fallback behavior.SpecMetadataAPI now uses a flat one-token-per-request buffer.QA Engineer Review
test_pytorch_model_engine.py: updated prompt-boundary and overlap-mock coverage.test_eagle3.py: updated flattened lookahead, sentinel, CUDA-graph, draft-input, and dynamic-tree coverage.test_mtp.py: updated shared-KV draft-input coverage for multiple context requests.test-db/orqa/test-list changes were provided.