[None][fix] Fix FlashInfer shared-KV speculative decode - #17264
Conversation
|
/bot run |
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
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
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1⚔️ Resolve merge conflicts 💡
🧪 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. |
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
|
|
/bot run |
|
PR_Github #67910 [ run ] triggered by Bot. Commit: |
|
|
|
/bot run |
|
PR_Github #68018 [ run ] triggered by Bot. Commit: |
|
PR_Github #67910 [ run ] completed with state |
|
PR_Github #68018 [ run ] completed with state
|
|
/bot run |
|
PR_Github #68140 [ run ] triggered by Bot. Commit: |
|
PR_Github #68140 [ run ] completed with state
|
* Why? Overlap-scheduled speculative decoding with a shared KV cache can use draft positions beyond the logical sequence length. FlashInfer exposed only logical generation pages, and did not reflect scheduler offsets in its live KV metadata, while Gemma4 sized RoPE only to the logical limit. * What? Expose every reserved generation page for shared-KV overlap decoding while tracking logical KV lengths separately. Apply and restore overlap scheduler offsets to append positions and trtllm-gen decode lengths, and add speculative RoPE headroom to Gemma4 target and assistant models Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
* Why? Cached decode plans ignored query length and generation batch size, so speculative decoding could reuse a plan built for an incompatible launch shape. * What? Include both dimensions in the plan parameters, pass the query length to FlashInfer, and reject batches with nonuniform generation query lengths. Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
80ed499 to
9ef9c81
Compare
|
/bot run |
|
PR_Github #68210 [ run ] triggered by Bot. Commit: |
|
PR_Github #68210 [ run ] completed with state
|
|
/bot run |
|
PR_Github #68321 [ run ] triggered by Bot. Commit: |
|
PR_Github #68321 [ run ] completed with state
|
|
/bot run |
|
PR_Github #68403 [ run ] triggered by Bot. Commit: |
|
PR_Github #68403 [ run ] completed with state |
Adds recipes/gemma4-31b/trtllm/agg-gb200-agentic/deploy.yaml, derived from the B200 sibling with the GB200-specific deltas: - nodeAffinity: NVIDIA-GB200 + kubernetes.io/arch=arm64 - Rename b200 -> gb200 (DGD, ConfigMap, labels) - runtimeVersionOverride: "1.4.0" (image tag isn't pure semver) - imagePullSecrets: nvcr-docker-config - image: aarch64 TRT-LLM 1.3.0rc25 + NVIDIA/TensorRT-LLM#17264 recipes/gemma4-31b/perf/perf.yaml gains one-line comments naming the two places (podAffinity + ENDPOINT env) to retarget together when switching between b200 and gb200 variants. Validation: aiperf trace replay, 3541 requests @ c=192, MTP DL=3 output_token_throughput 8,554 tok/s (aggregate) output_token_throughput_per_user 52.9 tok/s/user (p50) time_to_first_token 2.52 s (p50) request_latency 16.9 s (p50) Signed-off-by: Aswin Visva <31215515+aswinvisva@users.noreply.github.com>
Adds recipes/gemma4-31b/trtllm/agg-gb200-agentic/deploy.yaml, derived from the B200 sibling with the GB200-specific deltas: - nodeAffinity: NVIDIA-GB200 + kubernetes.io/arch=arm64 - Rename b200 -> gb200 (DGD, ConfigMap, labels) - runtimeVersionOverride: "1.4.0" (image tag isn't pure semver) - imagePullSecrets: nvcr-docker-config - image: aarch64 TRT-LLM 1.3.0rc25 + NVIDIA/TensorRT-LLM#17264 recipes/gemma4-31b/perf/perf.yaml gains one-line comments naming the two places (podAffinity + ENDPOINT env) to retarget together when switching between b200 and gb200 variants. Validation: aiperf trace replay, 3541 requests @ c=192, MTP DL=3 output_token_throughput 8,554 tok/s (aggregate) output_token_throughput_per_user 52.9 tok/s/user (p50) time_to_first_token 2.52 s (p50) request_latency 16.9 s (p50) Signed-off-by: Aswin Visva <31215515+aswinvisva@users.noreply.github.com>
10 participants
Dev Engineer Review
KVCacheManagerandKVCacheManagerV2.BaseResourceManager.QA Engineer Review
tests/integration/test_lists/based on the available repository coverage entries.Description
Overlap-scheduled speculative decoding with a shared KV cache can use
draft positions beyond the logical sequence length. FlashInfer exposed
only logical generation pages, and did not reflect scheduler offsets in
its live KV metadata, while Gemma4 sized RoPE only to the logical limit.
Expose every reserved generation page for shared-KV overlap decoding
while tracking logical KV lengths separately. Apply and restore overlap
scheduler offsets to append positions and trtllm-gen decode lengths, and
add speculative RoPE headroom to Gemma4 target and assistant models
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.