[https://nvbugs/6441022][fix] Enable CUDA graph for final context token computation - #16746
Conversation
WalkthroughThe change adds promoted final-context handling for CUDA graph execution, updates graph-key length selection, preserves KV-cache reuse state across boundary cases, and adds unit and integration coverage for guided decoding, multimodal requests, context logits, offload, and stream synchronization. ChangesPromoted final-context execution
Estimated code review effort: 4 (Complex) | ~60 minutes Suggested labels: Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant ScheduledRequests
participant PyTorchModelEngine
participant CUDAGraphRunner
ScheduledRequests->>PyTorchModelEngine: identify final-token context rows
PyTorchModelEngine->>PyTorchModelEngine: build promoted context graph batch
PyTorchModelEngine->>CUDAGraphRunner: request graph with promoted context ids
CUDAGraphRunner-->>PyTorchModelEngine: return graph and execution key
PyTorchModelEngine->>CUDAGraphRunner: replay promoted-context graph
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 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 `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 3721-3731: Move the “the request has no previous tensor”
explanatory comment from above the conditional so it appears immediately before
the `elif new_tokens_device is None or request.is_dummy or request.py_batch_idx
is None` branch. Keep the `is_promoted_context` branch and its logic unchanged.
🪄 Autofix (Beta)
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: 01bca682-2dee-4678-bad3-7332ccdbae4d
📒 Files selected for processing (7)
cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpptensorrt_llm/_torch/pyexecutor/cuda_graph_runner.pytensorrt_llm/_torch/pyexecutor/model_engine.pytests/integration/defs/kv_cache/test_final_single_token_context_cuda_graph.pytests/integration/test_lists/test-db/l0_h100.ymltests/unittest/_torch/executor/test_pytorch_model_engine.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py
|
[by Codex] @thorjohnsen Could you please review this PR? Thank you! |
039a183 to
c40fdab
Compare
|
/bot run --disable-fail-fast |
dpitman-nvda
left a comment
dpitman-nvda
left a comment
There was a problem hiding this comment.
Approval for infra
There was a problem hiding this comment.
Don't really see the purpose of adding test cases under cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp matters here. My opinion is we don't need them at all.
Promoting the context batch here to enable CUDA graph looks good to me.
It seems like the test cases won't get captured by how you specified them under the .yml files. I think you will need to update them.
Approving to unblock.
Hi @eopXD , thanks for the feedback. The purpose of the C++ tests is to pin the upstream KV-cache-manager contract that graph promotion depends on: final-token reuse must leave exactly one context token, preserve semantic request state/type and feature metadata, and correctly onboard shared offloaded blocks without breaking ownership or cancellation. Removing all of them would leave that cross-language boundary implicit. |
|
/bot run --disable-fail-fast |
|
PR_Github #62253 [ run ] triggered by Bot. Commit: |
There was a problem hiding this comment.
Overall, this looks good to me, and I’m approving to unblock.
One design consideration for the future: the PR places context requests in generation_requests and passes promoted_context_request_ids through several layers to preserve their original semantics. This works for the current change, though it may be somewhat fragile to maintain—future code may check only which list contains the request and accidentally apply generation behavior.
|
/bot run --disable-fail-fast |
|
PR_Github #62327 [ run ] triggered by Bot. Commit: |
|
PR_Github #62253 [ run ] completed with state |
|
PR_Github #62327 [ run ] completed with state
|
…ken contexts Promote eligible final single-token context rows through a temporary decode-shaped execution view after KV preparation. Commit the view only when the existing graph runner finds a matching graph, and retain semantic eager fallback otherwise. Preserve request state, request type, KV ownership, sampling order, and lifecycle updates. Cover context logits, guided decoding, zero-runtime-draft target execution, overlap token sourcing, sparse graph keys, local offload, changed-tail reuse, and TP2 rank-local replay. Tests added: cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp: - KVCacheManagerTest.AddSequenceBatchLeavesOneFinalContextTokenAfterReuse - KVCacheManagerTest.AddSequenceBatchPreservesDraftTokensOnFinalContextAfterReuse - KVCacheManagerTest.AddSequenceBatchPreservesGuidanceAndContextLogitsAfterReuse - KVCacheManagerTest.AddSequenceBatchOnboardsOffloadedPrefixForFinalContextToken - KVCacheManagerTest.AddSequenceBatchLeavesOneFinalMultimodalContextTokenAfterReuse tests/integration/defs/kv_cache/test_final_single_token_context_cuda_graph.py: - test_final_token_reuse_cuda_graph - test_changed_final_token_reuse_cuda_graph - test_final_token_reuse_cuda_graph_tp2 - test_context_logits_after_final_token_reuse - test_guided_decoding_after_final_token_reuse - test_zero_runtime_draft_speculation_after_final_token_reuse tests/unittest/_torch/executor/test_pytorch_model_engine.py: - SingleTokenContextGraphBatchTestCase.test_generation_only_is_identity - SingleTokenContextGraphBatchTestCase.test_eligible_batch_has_independent_lists_and_stable_order - SingleTokenContextGraphBatchTestCase.test_structural_fallbacks_return_semantic_batch - SingleTokenContextGraphBatchTestCase.test_context_shape_and_mode_fallback_matrix - SingleTokenContextGraphBatchTestCase.test_context_logits_use_final_token_graph_candidate - SingleTokenContextGraphBatchTestCase.test_generation_only_request_in_context_list_falls_back - SingleTokenContextGraphBatchTestCase.test_generation_shape_fallback_matrix - SingleTokenContextGraphBatchTestCase.test_mixed_one_and_two_token_contexts_fall_back_together - SingleTokenContextGraphBatchTestCase.test_mrope_delta_is_supported_by_decode_provider - SingleTokenContextGraphBatchTestCase.test_multimodal_context_requires_compatible_decode_token - SingleTokenContextGraphBatchTestCase.test_multimodal_pending_event_is_rechecked - SingleTokenContextGraphBatchTestCase.test_multimodal_decode_compatibility_uses_final_prompt_token - SingleTokenContextGraphBatchTestCase.test_sparse_sequence_mode_uses_promoted_context_cursor - SingleTokenContextGraphBatchTestCase.test_graph_key_forwards_promoted_context_ids - SingleTokenContextGraphBatchTestCase.test_graph_lookup_forwards_promoted_context_ids - SingleTokenContextGraphBatchTestCase.test_forward_commits_candidate_only_on_graph_hit - SingleTokenContextGraphBatchTestCase.test_forward_graph_miss_uses_semantic_eager_batch - SingleTokenContextGraphBatchTestCase.test_zero_runtime_draft_speculation_commits_graph_candidate - SingleTokenContextGraphBatchTestCase.test_zero_runtime_draft_speculation_graph_miss_is_semantic_eager - SingleTokenContextGraphBatchTestCase.test_forward_allows_guided_context_logits_on_graph_hit - SingleTokenContextGraphBatchTestCase.test_multimodal_graph_miss_preserves_semantic_payload - SingleTokenContextGraphBatchTestCase.test_generation_only_forward_does_not_call_new_selector - SingleTokenContextGraphBatchTestCase.test_global_incompatibilities_bypass_candidate_selection - PyTorchModelEngineTestCase.test_promoted_context_uses_prompt_token_during_overlap - PyTorchModelEngineTestCase.test_promoted_context_precedes_speculative_overlap_generation - PyTorchModelEngineTestCase.test_promoted_mrope_context_uses_decode_state_contract - PyTorchModelEngineTestCase.test_kv_cache_manager_with_execution_stream - PyTorchModelEngineTestCase.test_cuda_graph_replay_observes_execution_stream_dependency tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py: - test_v2_resume_restores_offsets_only_after_execution_stream_ready Signed-off-by: Simeng Liu <simengl@nvidia.com>
Signed-off-by: Simeng Liu <simengl@nvidia.com>
Signed-off-by: Simeng Liu <simengl@nvidia.com>
Signed-off-by: Simeng Liu <simengl@nvidia.com>
c25ef42 to
7a69bd7
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #62577 [ run ] triggered by Bot. Commit: |
|
PR_Github #62577 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62813 [ run ] triggered by Bot. Commit: |
|
PR_Github #62813 [ run ] completed with state
|
|
/bot run |
|
PR_Github #63063 [ run ] triggered by Bot. Commit: |
|
PR_Github #63063 [ run ] completed with state |
…en computation (#16746) Signed-off-by: Simeng Liu <simengl@nvidia.com>
10 participants
Summary by CodeRabbit
New Features
Bug Fixes
Tests
Description
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.