Skip to content

[None][fix] Fix one-model MTP KV cache accounting - #17264

Open
2ez4bz wants to merge 2 commits into
NVIDIA:mainfrom
2ez4bz:dev-flashinfer-mtp-ima
Open

[None][fix] Fix one-model MTP KV cache accounting#17264
2ez4bz wants to merge 2 commits into
NVIDIA:mainfrom
2ez4bz:dev-flashinfer-mtp-ima

Conversation

@2ez4bz

@2ez4bz 2ez4bz commented Aug 4, 2026

Copy link
Copy Markdown
Collaborator

Dev Engineer Review

  • One-model MTP scheduling now synchronizes Python-side and C++ draft-token state.
  • Speculative KV accounting now uses the prompt KV boundary for the first generation step when dynamic KV-length correction is unavailable.
  • FlashInfer now 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 context-manager support to BaseResourceManager.
  • No configuration or test-list files changed.

QA Engineer Review

  • Added regression tests for prompt-boundary KV accounting and FlashInfer KV-offset restoration.
  • Added tests for one-model MTP draft-token normalization, preservation of Python draft tokens, and sequence-limit draft-length handling.
  • 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?

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.

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.

@2ez4bz

2ez4bz commented Aug 4, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@coderabbitai

coderabbitai Bot commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

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
FlashInfer tracks logical KV lengths separately from reserved generation page-table capacity. Speculative offsets update positions and decode-wrapper buffers.
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, tensorrt_llm/_torch/pyexecutor/resource_manager.py, tests/unittest/_torch/executor/test_py_executor.py, tests/unittest/_torch/attention/test_flashinfer_attention.py
One-model MTP disables drafting when the batch can exceed max_seq_len and preserves populated Python draft tokens. Tests cover draft normalization, sequence limits, page tables, offsets, and context-manager cleanup.

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
Loading

Possibly related PRs

Suggested labels: ci: full pre-merge approved

Suggested reviewers: schetlur-nv

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 13.33% 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
Title check ✅ Passed The title clearly identifies the one-model MTP KV cache accounting fix and follows the required ticket and type format.
Description check ✅ Passed The description explains the problem and solution and includes the checklist, but the Test Coverage section is not populated with specific tests.
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.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread tests/unittest/_torch/executor/test_py_executor.py Outdated
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@BowenFu

BowenFu commented Aug 5, 2026

Copy link
Copy Markdown

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.

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

coderabbitai Bot commented Aug 5, 2026

Copy link
Copy Markdown
Contributor

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.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread tests/unittest/_torch/executor/test_py_executor.py
@2ez4bz

2ez4bz commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator Author

@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.

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

coderabbitai Bot commented Aug 5, 2026

Copy link
Copy Markdown
Contributor

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.

@2ez4bz

2ez4bz commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@2ez4bz

2ez4bz commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@2ez4bz

2ez4bz commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

* 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>
@2ez4bz
2ez4bz force-pushed the dev-flashinfer-mtp-ima branch from 9ff8de8 to 275e929 Compare August 8, 2026 04:03
@2ez4bz
2ez4bz requested review from a team as code owners August 8, 2026 04:03
@coderabbitai

coderabbitai Bot commented Aug 8, 2026

Copy link
Copy Markdown
Contributor

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.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 win

Apply FlashInfer offsets to extend-context rows.

FlashInferAttentionMetadata.apply_spec_decode_kv_lens_offsets() returns when num_generations == 0. In extend-context mode, extend requests are included in num_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 win

Duplicated use_full_generation_page_table gating 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 win

Duplicated use_full_generation_page_table gating 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

📥 Commits

Reviewing files that changed from the base of the PR and between 6ba3de1 and 275e929.

📒 Files selected for processing (8)
  • tensorrt_llm/_torch/attention_backend/flashinfer.py
  • tensorrt_llm/_torch/metadata.py
  • tensorrt_llm/_torch/pyexecutor/model_engine.py
  • tensorrt_llm/_torch/pyexecutor/py_executor.py
  • tensorrt_llm/_torch/pyexecutor/resource_manager.py
  • tests/unittest/_torch/attention/test_flashinfer_attention.py
  • tests/unittest/_torch/executor/test_py_executor.py
  • tests/unittest/_torch/executor/test_pytorch_model_engine.py

Comment on lines +3313 to +3337
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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 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.py

Repository: 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.py

Repository: 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 320

Repository: 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 260

Repository: 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.py

Repository: 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 260

Repository: 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")
PY

Repository: 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"):

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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):

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this still needed?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants