[None][fix] Fix one-model MTP KV cache accounting - #17264
Conversation
|
/bot run |
WalkthroughThe 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. ChangesSpeculative decoding and KV-cache updates
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
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
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. |
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
|
* 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>
* Why? With shared-KV one-model MTP and overlap scheduling, FlashInfer built generation page tables from the host logical length. Device-resident acceptance state could advance an append into the next reserved page, which was absent from the table and caused an illegal memory access. Speculative positions near the sequence limit could also exceed the model's RoPE table. * What? Stage every reserved generation page while tracking logical KV lengths separately for positions and decode. Propagate overlap corrections through FlashInfer's CUDA graph state, and disable drafting batch-wide when speculative positions could exceed the sequence limit. Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
9ff8de8 to
275e929
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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
2968-2992: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winApply FlashInfer offsets to extend-context rows.
FlashInferAttentionMetadata.apply_spec_decode_kv_lens_offsets()returns whennum_generations == 0. In extend-context mode, extend requests are included innum_contexts, so their rows are[num_contexts - num_chunked_ctx_requests:num_contexts]and are not corrected. Add equivalent chunked-context handling in both preprocessing and restoration.🤖 Prompt for 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. In `@tensorrt_llm/_torch/pyexecutor/model_engine.py` around lines 2968 - 2992, Update both preprocessing and restoration flows around apply_spec_decode_kv_lens_offsets so FlashInfer metadata also adjusts chunked extend-context rows. When num_chunked_ctx_requests > 0, apply previous_kv_lens_offsets_cuda to rows [num_ctx_requests - num_chunked_ctx_requests:num_ctx_requests], including cases with zero generations; preserve the existing generation-row handling for other requests.
🧹 Nitpick comments (2)
tensorrt_llm/_torch/pyexecutor/model_engine.py (2)
3485-3492: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicated
use_full_generation_page_tablegating condition.This exact four-clause expression is repeated verbatim at Line 5098-5101 in
_prepare_tp_inputs. Both sites gate the same correctness-critical feature (whether generation rows get the full reservation-width page table). If one site's condition changes in the future (e.g. a new gating factor) without the other, the two code paths silently diverge on this KV-cache-accounting decision. Extract a small helper, e.g.self._should_use_full_generation_page_table(spec_config, attn_metadata).♻️ Proposed refactor
+ def _should_use_full_generation_page_table( + self, spec_config: Optional[DecodingBaseConfig], + attn_metadata: AttentionMetadata) -> bool: + return (self.enable_spec_decode and not self._disable_overlap_scheduler + and getattr(spec_config, '_use_shared_kv_cache', False) + and hasattr(attn_metadata, 'apply_spec_decode_kv_lens_offsets')) + def _prepare_incremental_update_metadata( self, ... attn_metadata.kv_cache_params = KVCacheParams( use_cache=True, num_cached_tokens_per_seq=num_cached_tokens_per_seq, num_extra_kv_tokens=get_num_extra_kv_tokens(spec_config), - use_full_generation_page_table=( - self.enable_spec_decode and not self._disable_overlap_scheduler - and getattr(spec_config, '_use_shared_kv_cache', False) and - hasattr(attn_metadata, 'apply_spec_decode_kv_lens_offsets'))) + use_full_generation_page_table=self._should_use_full_generation_page_table( + spec_config, attn_metadata))🤖 Prompt for 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. In `@tensorrt_llm/_torch/pyexecutor/model_engine.py` around lines 3485 - 3492, Extract the repeated four-clause `use_full_generation_page_table` condition from the KVCacheParams construction and `_prepare_tp_inputs` into a shared helper such as `_should_use_full_generation_page_table(spec_config, attn_metadata)`. Replace both inline expressions with calls to this helper while preserving the existing gating behavior.
5094-5101: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicated
use_full_generation_page_tablegating condition (second site).Same expression as Line 3489-3492 in
_prepare_incremental_update_metadata. See the comment on that range; both should call one shared helper.🤖 Prompt for 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. In `@tensorrt_llm/_torch/pyexecutor/model_engine.py` around lines 5094 - 5101, Extract the repeated use_full_generation_page_table eligibility expression into a shared helper, then update both _prepare_incremental_update_metadata and this KVCacheParams construction to call it. Preserve the existing checks for spec decoding, overlap scheduling, shared KV cache, and apply_spec_decode_kv_lens_offsets.
🤖 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 `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 3313-3337: Pass the resolved maximum sequence length into the
AutoDeploy PyExecutor constructor so self.max_seq_len is initialized, using the
engine cache sequence interface’s max_seq_len value. Ensure
_one_model_mtp_batch_needs_zero_draft can compare max_target_position against a
valid limit during one-model MTP drafting.
---
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 2968-2992: Update both preprocessing and restoration flows around
apply_spec_decode_kv_lens_offsets so FlashInfer metadata also adjusts chunked
extend-context rows. When num_chunked_ctx_requests > 0, apply
previous_kv_lens_offsets_cuda to rows [num_ctx_requests -
num_chunked_ctx_requests:num_ctx_requests], including cases with zero
generations; preserve the existing generation-row handling for other requests.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 3485-3492: Extract the repeated four-clause
`use_full_generation_page_table` condition from the KVCacheParams construction
and `_prepare_tp_inputs` into a shared helper such as
`_should_use_full_generation_page_table(spec_config, attn_metadata)`. Replace
both inline expressions with calls to this helper while preserving the existing
gating behavior.
- Around line 5094-5101: Extract the repeated use_full_generation_page_table
eligibility expression into a shared helper, then update both
_prepare_incremental_update_metadata and this KVCacheParams construction to call
it. Preserve the existing checks for spec decoding, overlap scheduling, shared
KV cache, and apply_spec_decode_kv_lens_offsets.
🪄 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: c80bbdf5-195b-4a3d-93f2-146f12daecfb
📒 Files selected for processing (8)
tensorrt_llm/_torch/attention_backend/flashinfer.pytensorrt_llm/_torch/metadata.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/resource_manager.pytests/unittest/_torch/attention/test_flashinfer_attention.pytests/unittest/_torch/executor/test_py_executor.pytests/unittest/_torch/executor/test_pytorch_model_engine.py
| def _one_model_mtp_batch_needs_zero_draft( | ||
| self, scheduled_batch: ScheduledRequests) -> bool: | ||
| """Return whether drafting could produce an out-of-range position.""" | ||
| spec_config = self.model_engine.spec_config | ||
| runtime_draft_len = self.model_engine.runtime_draft_len | ||
| if (runtime_draft_len == 0 or spec_config is None | ||
| or not spec_config.spec_dec_mode.is_mtp_eagle_one_model()): | ||
| return False | ||
|
|
||
| # With overlap, the host request has not incorporated the previous | ||
| # iteration's accepted tokens yet. _preprocess_inputs adds that count | ||
| # to every target position on device. Use the maximum possible count | ||
| # because reading the exact value here would synchronize the GPU. | ||
| max_pending_tokens = (0 if self.disable_overlap_scheduler else | ||
| self.model_engine.max_draft_len + 1) | ||
| target_position_width = spec_config.get_runtime_tokens_per_gen_step( | ||
| runtime_draft_len) | ||
|
|
||
| for request in scheduled_batch.generation_requests: | ||
| max_target_position = (request.max_beam_num_tokens - 1 + | ||
| max_pending_tokens + target_position_width - | ||
| 1) | ||
| if max_target_position >= self.max_seq_len: | ||
| return True | ||
| return False |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Find every PyExecutor construction and check whether max_seq_len is passed.
rg -nP -C 6 'PyExecutor\(' --type=py -g '!**/tests/**'Repository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- candidate files ---'
git ls-files | rg '(^|/)(py_executor|.*executor.*)\.py$|CODING_GUIDELINES\.md$'
printf '%s\n' '--- PyExecutor declarations and references ---'
rg -n -C 4 'class PyExecutor|PyExecutor\(|max_seq_len\s*=' --type py -g '!**/tests/**' .
printf '%s\n' '--- target implementation ---'
sed -n '3260,3360p' tensorrt_llm/_torch/pyexecutor/py_executor.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- PyExecutor declaration ---'
rg -n -C 8 '^class PyExecutor|def __init__\(' tensorrt_llm/_torch/pyexecutor/py_executor.py | head -n 180
printf '%s\n' '--- PyExecutor creator ---'
rg -n -C 10 'PyExecutor|create_executor|create_py_executor|max_seq_len' tensorrt_llm/_torch/pyexecutor/py_executor_creator.py tensorrt_llm/executor/base_worker.py
printf '%s\n' '--- exact PyExecutor references outside tests ---'
rg -n --glob '*.py' --glob '!**/tests/**' '\bPyExecutor\b' tensorrt_llm examples scripts
printf '%s\n' '--- target method and nearby callers ---'
rg -n -C 12 '_one_model_mtp_batch_needs_zero_draft|needs_zero_draft' tensorrt_llm/_torch/pyexecutor/py_executor.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- PyExecutor constructor ---'
sed -n '504,650p' tensorrt_llm/_torch/pyexecutor/py_executor.py
printf '%s\n' '--- executor instance factory ---'
rg -n -C 18 '^def create_py_executor_instance|create_py_executor_instance\(' tensorrt_llm/_torch/pyexecutor/_util.py tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
printf '%s\n' '--- runtime-size derivation ---'
rg -n -C 10 'def get_runtime_sizes|max_seq_len\s*=' tensorrt_llm/llmapi/llm_args.py tensorrt_llm/_torch/pyexecutor/py_executor_creator.py | head -n 260
printf '%s\n' '--- one-model MTP configuration and model-engine limit ---'
rg -n -C 10 'is_mtp_eagle_one_model|model_engine\.max_seq_len|max_seq_len.*None|None.*max_seq_len' tensorrt_llm/_torch/pyexecutor tensorrt_llm/llmapi/llm_args.py | head -n 320Repository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- all executor-instance factory calls ---'
rg -n -U -C 3 'create_py_executor_instance\(' --glob '*.py' . | rg -v '^[^:]+:[0-9]+:.*tests/' | head -n 220
printf '%s\n' '--- max_seq_len assignment in PyExecutor ---'
rg -n -C 6 'self\.max_seq_len|max_seq_len\s*=' tensorrt_llm/_torch/pyexecutor/py_executor.py | head -n 180
printf '%s\n' '--- target method callers ---'
rg -n -C 8 '_one_model_mtp_batch_needs_zero_draft' tensorrt_llm/_torch/pyexecutor/py_executor.py
printf '%s\n' '--- model-engine max_seq_len definition ---'
rg -n -C 8 'class .*ModelEngine|def max_seq_len|max_seq_len\s*:' tensorrt_llm/_torch/pyexecutor/model_engine.py tensorrt_llm/_torch/pyexecutor -g '*.py' | head -n 260Repository: NVIDIA/TensorRT-LLM
Length of output: 30197
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- factory implementation and direct constructor call ---'
sed -n '2550,2665p' tensorrt_llm/_torch/pyexecutor/_util.py
printf '%s\n' '--- all direct PyExecutor constructor expressions ---'
rg -n -U -C 5 '\bPyExecutor\s*\(' --glob '*.py' . | head -n 220
printf '%s\n' '--- model-engine initialization and inference ---'
sed -n '303,390p' tensorrt_llm/_torch/pyexecutor/model_engine.py
sed -n '2800,2865p' tensorrt_llm/_torch/pyexecutor/model_engine.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 18068
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- auto-deploy direct constructor arguments ---'
sed -n '1290,1365p' tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
printf '%s\n' '--- shared factory constructor arguments ---'
sed -n '2880,2965p' tensorrt_llm/_torch/pyexecutor/_util.py
printf '%s\n' '--- auto-deploy max_seq_len sources ---'
rg -n -C 8 'max_seq_len|spec_config|is_mtp_eagle_one_model' tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py | head -n 260Repository: NVIDIA/TensorRT-LLM
Length of output: 20218
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- AutoDeploy max_seq_len declaration and defaults ---'
rg -n -C 8 'class .*Args|class .*Config|max_seq_len\s*:' tensorrt_llm/_torch/auto_deploy tensorrt_llm/llmapi/llm_args.py | head -n 260
printf '%s\n' '--- AutoDeploy sequence-info max_seq_len handling ---'
rg -n -C 10 'class CachedSequenceInterface|max_seq_len' tensorrt_llm/_torch/auto_deploy tensorrt_llm/_torch | head -n 260
printf '%s\n' '--- deterministic comparison probe ---'
python3 - <<'PY'
max_target_position = 10
max_seq_len = None
try:
max_target_position >= max_seq_len
except Exception as exc:
print(type(exc).__name__, str(exc))
else:
print("comparison did not raise")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 45361
Pass max_seq_len to the AutoDeploy PyExecutor constructor. The direct constructor call omits it, so self.max_seq_len remains None. When one-model MTP drafting is active, _one_model_mtp_batch_needs_zero_draft() raises TypeError during max_target_position >= self.max_seq_len. Pass the resolved engine limit, such as engine.cache_seq_interface.info.max_seq_len, or add a None guard.
🤖 Prompt for 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.
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py` around lines 3313 - 3337, Pass
the resolved maximum sequence length into the AutoDeploy PyExecutor constructor
so self.max_seq_len is initialized, using the engine cache sequence interface’s
max_seq_len value. Ensure _one_model_mtp_batch_needs_zero_draft can compare
max_target_position against a valid limit during one-model MTP drafting.
| logical_num_blocks: List[int], | ||
| num_contexts: int) -> List[int]: | ||
| """Keep context rows logical and expose every reserved generation page.""" | ||
| if hasattr(kv_cache_manager, "kv_cache_map"): |
There was a problem hiding this comment.
Check type of kv_cache_manager + see if kv_cache_map is guaranteed to exist.
|
|
||
| def _publish_decode_wrapper_kv_lens(self, decode_wrapper) -> None: | ||
| """Publish device-logical lengths after a trtllm-gen decode plan.""" | ||
| kv_lens_buffer = getattr(decode_wrapper, "_kv_lens_buffer", None) |
There was a problem hiding this comment.
Ditto here re tighetning type of decode_wrapper + potentially getting rid of getattr call.
| start = self.num_contexts | ||
| end = start + self.num_generations | ||
| if self._is_shared_kv_draft_view: | ||
| # The external assistant is Q-only and does not append its query to |
There was a problem hiding this comment.
What does "assistant" refer to here?
| @@ -3303,6 +3303,39 @@ def _handle_dynamic_draft_len(self, | |||
| if spec_config is not None and spec_config.is_linear_tree else | |||
| self.model_engine.max_total_draft_tokens) | |||
|
|
|||
| if self._one_model_mtp_batch_needs_zero_draft(scheduled_batch): | |||
There was a problem hiding this comment.
Is this still needed?
Dev Engineer Review
BaseResourceManager.QA Engineer Review
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.