[None][fix] Fix one-model MTP KV cache accounting - #17264
Conversation
|
/bot run |
|
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 (4)
🚧 Files skipped from review as they are similar to previous changes (4)
WalkthroughOne-model speculative decoding preserves Python draft tokens across iterations. First-iteration cache accounting uses the prompt boundary when dynamic KV lengths are unavailable. Tests cover draft-token state, cached-token accounting, and position IDs. ChangesSpeculative decoding updates
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
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
📒 Files selected for processing (4)
tensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/py_executor.pytests/unittest/_torch/executor/test_py_executor.pytests/unittest/_torch/executor/test_pytorch_model_engine.py
|
PR_Github #63826 [ run ] triggered by Bot. Commit: |
|
PR_Github #63826 [ run ] completed with state
|
|
Could you clarify how these two hunks relate?
The |
0baecbe to
990299f
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. |
There was a problem hiding this comment.
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
📒 Files selected for processing (4)
tensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/py_executor.pytests/unittest/_torch/executor/test_py_executor.pytests/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
|
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 The first-generation I’m keeping them together because both are needed for this MTP + FlashInfer + block-reuse fix. |
* Why? The first overlapped MTP generation step could overstate its cached length and access an invalid FlashInfer page, causing a CUDA illegal memory access. The KV scheduler also lacked the Python-side draft-token count needed for correct capacity planning. * What? Keep the Python and C++ draft-token state synchronized, and use the prompt KV boundary when generation immediately follows context for backends without dynamic KV-length correction. Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
990299f to
9ff8de8
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. |
|
/bot run |
|
PR_Github #64088 [ run ] triggered by Bot. Commit: |
|
PR_Github #64088 [ run ] completed with state
|
|
/bot run |
|
PR_Github #64146 [ run ] triggered by Bot. Commit: |
|
PR_Github #64146 [ run ] completed with state
|
|
/bot run |
|
PR_Github #64208 [ run ] triggered by Bot. Commit: |
|
PR_Github #64208 [ run ] completed with state
|
3 participants
Dev Engineer Review
kv_lens_cuda.QA Engineer Review
test_first_speculative_generation_uses_prompt_kv_boundary.test_py_executor.py.tests/integration/test_lists/based on the available repository coverage entries.Description
The first overlapped MTP generation step could overstate its cached length and access an invalid FlashInfer page, causing a CUDA illegal memory access. The KV scheduler also lacked the Python-side draft-token count needed for correct capacity planning.
Keep the Python and C++ draft-token state synchronized, and use the prompt KV boundary when generation immediately follows context for backends without dynamic KV-length correction.
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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.