[https://nvbugs/6293536][fix] order KV cache transfers with overlap scheduler - #15349
Closed
VALLIS-NERIA wants to merge 2 commits into
Closed
[https://nvbugs/6293536][fix] order KV cache transfers with overlap scheduler#15349VALLIS-NERIA wants to merge 2 commits into
VALLIS-NERIA wants to merge 2 commits into
Conversation
VALLIS-NERIA
changed the title
VALLIS-NERIA
force-pushed
the
nvbug-6293536-kv-transfer-ordering
branch
from
June 14, 2026 11:32
0c409cc to
ad5b49f
Compare
Collaborator
Author
|
Update after force-push to
|
VALLIS-NERIA
force-pushed
the
nvbug-6293536-kv-transfer-ordering
branch
from
June 14, 2026 11:51
ad5b49f to
1a3eea5
Compare
Collaborator
Author
|
Update after force-push to
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #54402 [ run ] triggered by Bot. Commit: |
Contributor
📝 WalkthroughWalkthroughAdds a ChangesKV cache pending host transfer query and stream synchronization
Sequence Diagram(s)sequenceDiagram
participant PyExecutor
participant execution_stream as CUDA execution_stream
participant ResourceManager
participant KVCacheManager as KVCacheManager (Python)
participant BaseKVCacheManager as BaseKVCacheManager (C++)
PyExecutor->>execution_stream: record Event → forward_done_event
PyExecutor->>ResourceManager: prepare_resources(scheduled_batch, forward_done_event)
ResourceManager->>BaseKVCacheManager: has_pending_host_transfers()
BaseKVCacheManager-->>ResourceManager: bool
alt has pending host transfers
ResourceManager->>execution_stream: forward_done_event.wait()
end
ResourceManager->>KVCacheManager: refresh_blocks(scheduled_batch)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/resource_manager.py (1)
990-993: ⚡ Quick winAdd explicit
-> Nonereturn annotations to touchedprepare_resourcesmethods.These signatures were modified and should include explicit return annotations for typing consistency.
As per coding guidelines: “Always annotate functions. Make the return type
Noneif the function does not return anything.”✍️ Suggested patch
- def prepare_resources( - self, - scheduled_batch: ScheduledRequests, - forward_done_event: Optional[torch.cuda.Event] = None): + def prepare_resources( + self, + scheduled_batch: ScheduledRequests, + forward_done_event: Optional[torch.cuda.Event] = None) -> None: @@ - def prepare_resources( - self, - scheduled_batch: ScheduledRequests, - forward_done_event: Optional[torch.cuda.Event] = None): + def prepare_resources( + self, + scheduled_batch: ScheduledRequests, + forward_done_event: Optional[torch.cuda.Event] = None) -> None:Also applies to: 3852-3855
🤖 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/resource_manager.py` around lines 990 - 993, Add explicit `-> None` return type annotations to the `prepare_resources` method signatures in the file tensorrt_llm/_torch/pyexecutor/resource_manager.py. At the first location (lines 990-993), add `-> None` between the closing parenthesis of the method parameters and the colon to complete the type annotation. Apply the same change at the second location (lines 3852-3855) where another `prepare_resources` method signature appears. Both method signatures should follow the pattern of ending with `-> None:` to match the coding guideline that all functions must have explicit return type annotations.Source: Coding guidelines
🤖 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 3280-3283: The forward_done_event is recorded at line 3280 after
_prepare_and_schedule_batch() has already executed, but that method can invoke
_prepare_disagg_gen_init() which directly calls the KV cache manager's
prepare_resources() without the event, causing potential race conditions in
overlap mode. Move the _record_execution_stream_tail_event() call to execute
before _prepare_and_schedule_batch() is called, or thread the recorded event
through to _prepare_disagg_gen_init() so that the event is passed to all
prepare_resources() calls in both the main path and the disagg-gen-init path,
ensuring proper synchronization with the previous forward operation.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py`:
- Around line 990-993: Add explicit `-> None` return type annotations to the
`prepare_resources` method signatures in the file
tensorrt_llm/_torch/pyexecutor/resource_manager.py. At the first location (lines
990-993), add `-> None` between the closing parenthesis of the method parameters
and the colon to complete the type annotation. Apply the same change at the
second location (lines 3852-3855) where another `prepare_resources` method
signature appears. Both method signatures should follow the pattern of ending
with `-> None:` to match the coding guideline that all functions must have
explicit return type annotations.
🪄 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: 8caf2160-83af-40ef-9198-f7ab43cfd1cb
📒 Files selected for processing (7)
cpp/include/tensorrt_llm/batch_manager/kvCacheManager.hcpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.hcpp/tensorrt_llm/batch_manager/kvCacheManager.cppcpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cppcpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpptensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/resource_manager.py
Comment on lines
+3280
to
+3283
| forward_done_event = self._record_execution_stream_tail_event( | ||
| ) | ||
| self.resource_manager.prepare_resources( | ||
| scheduled_batch, forward_done_event) |
Contributor
There was a problem hiding this comment.
Thread the tail event through the earlier disagg-gen-init prepare path.
Line 3280 records the event only after _prepare_and_schedule_batch() has already run, but that path can call _prepare_disagg_gen_init() and directly invoke KV cache manager prepare_resources(...) without the event. In overlap mode, that still lets KVCM prepare/refresh work race ahead of the previous forward queued on execution_stream.
Potential direction
- scheduled_batch, iter_stats = self._prepare_and_schedule_batch()
+ forward_done_event = self._record_execution_stream_tail_event()
+ scheduled_batch, iter_stats = self._prepare_and_schedule_batch(
+ forward_done_event
+ )
...
- forward_done_event = self._record_execution_stream_tail_event(
- )
self.resource_manager.prepare_resources(
scheduled_batch, forward_done_event)- def _prepare_and_schedule_batch(self):
+ def _prepare_and_schedule_batch(
+ self,
+ forward_done_event: Optional[torch.cuda.Event] = None):
...
- self._prepare_disagg_gen_init(fitting_disagg_gen_init_requests)
+ self._prepare_disagg_gen_init(
+ fitting_disagg_gen_init_requests, forward_done_event)- def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests):
+ def _prepare_disagg_gen_init(
+ self,
+ fitting_disagg_gen_init_requests,
+ forward_done_event: Optional[torch.cuda.Event] = None):
...
- self.resource_manager.resource_managers[
- resource_mgr_type].prepare_resources(
- disagg_gen_init_to_prepare)
+ manager = self.resource_manager.resource_managers[
+ resource_mgr_type]
+ if resource_mgr_type == ResourceManagerType.KV_CACHE_MANAGER:
+ manager.prepare_resources(
+ disagg_gen_init_to_prepare, forward_done_event)
+ else:
+ manager.prepare_resources(disagg_gen_init_to_prepare)🤖 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 3280 - 3283, The
forward_done_event is recorded at line 3280 after _prepare_and_schedule_batch()
has already executed, but that method can invoke _prepare_disagg_gen_init()
which directly calls the KV cache manager's prepare_resources() without the
event, causing potential race conditions in overlap mode. Move the
_record_execution_stream_tail_event() call to execute before
_prepare_and_schedule_batch() is called, or thread the recorded event through to
_prepare_disagg_gen_init() so that the event is passed to all
prepare_resources() calls in both the main path and the disagg-gen-init path,
ensuring proper synchronization with the previous forward operation.
Collaborator
|
PR_Github #54402 [ run ] completed with state
|
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
VALLIS-NERIA
force-pushed
the
nvbug-6293536-kv-transfer-ordering
branch
from
June 16, 2026 04:02
1a3eea5 to
219e2a7
Compare
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #54482 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #54482 [ run ] completed with state
|
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #54805 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #54805 [ run ] completed with state
|
Collaborator
Author
|
We have a real fix now #15546 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Refs NVBUG 6293536
Summary
refresh_blocks()onhas_pending_transfers()so iterations without KV transfers keep the normal overlap path.Details
The overlap scheduler can enter KVCM prepare while previous forward work is still queued on the execution stream. The previous implementation synchronized transfer streams through the buffer-manager stream only. This PR adds explicit event ordering for the transfer-manager path while keeping BufferManager off the execution stream.
Validation
python ./scripts/build_wheel.py --trt_root /usr/local/tensorrt --use_ccache -a "100-real;103-real" -s -j 64 --benchmarks --use_ccache --extra-cmake-vars NIXL_ROOT=/opt/nvidia/nvda_nixlpython3 -m py_compile tensorrt_llm/_torch/pyexecutor/resource_manager.py tensorrt_llm/_torch/pyexecutor/py_executor.pypython3 overlap-kvcm-repro.py --worker-url http://localhost:8000 --n-prefixes 12 --isl-reps 850 --concurrency 4 --rounds 100(ok=400 err=0).nvcudmpgenerated after the repro run.Summary by CodeRabbit
New Features
Improvements