[TRTLLM-12499][feat] Pipelined KVCache transfer for disaggregated serving in Python Cache Transceiver - #15727
Conversation
📝 WalkthroughWalkthroughAdds chunked and pipelined KV cache transfer to disaggregated serving. ChangesChunked & Pipelined KV Transfer
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes 🚥 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: 6
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/disaggregation/native/transfer.py (1)
488-507: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winMark the task as in-flight before waiting on the CUDA event.
Line 492 waits while the task is still
INIT, socancel_request()can see noTRANSFERRINGtasks and free KV pages before the event completes. Move the INIT→TRANSFERRING transition before the event wait, and keep the cancelled/error abort path before synchronization.Suggested fix
- # For pipelined prefill-transfer: wait for the GPU forward - # to finish writing KV data before starting RDMA. This - # blocks only this worker thread, not the GPU or main thread. - if task._slice.cuda_event is not None: # TODO: should I sync after the task status is set to TRANSFERRING? - task._slice.cuda_event.synchronize() - - if timer: - timer.record_push_end(write_meta.peer_rank) # Hold session.lock to serialize the INIT→TRANSFERRING transition with # cancel(): prevents cancel_request() from freeing KV pages while a # worker is about to write into them. with session.lock: status = session.status if status in (SessionStatus.ERROR, SessionStatus.CANCELLED): should_abort = True else: task.status = TaskStatus.TRANSFERRING should_abort = False + + if should_abort: + ... + return + + # For pipelined prefill-transfer: wait for the GPU forward + # to finish writing KV data before starting RDMA. + if task._slice.cuda_event is not None: + task._slice.cuda_event.synchronize() + + if timer: + timer.record_push_end(write_meta.peer_rank)🤖 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/disaggregation/native/transfer.py` around lines 488 - 507, The task transition in transfer.py is happening too late in the prefill-transfer flow: `task._slice.cuda_event.synchronize()` runs while the task is still `INIT`, so `cancel_request()` can miss it and free KV pages too early. In the transfer path around `task`, `session.lock`, and `TaskStatus.TRANSFERRING`, move the INIT→TRANSFERRING state update (with the session ERROR/CANCELLED abort check) before waiting on the CUDA event, and keep the abort branch ahead of synchronization so in-flight work is visible before any blocking wait.
🧹 Nitpick comments (2)
tensorrt_llm/_torch/disaggregation/transceiver.py (2)
582-585: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRedundant session assignment.
_get_or_create_send_sessionalready inserts the session intoself._send_sessions, so re-assigning the return value is redundant (and could mask a future divergence between the two code paths). Mirror the simpler form used inrespond_and_send_async.🤖 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/disaggregation/transceiver.py` around lines 582 - 585, The send-session initialization in transceiver logic has a redundant assignment because _get_or_create_send_session already stores the session in self._send_sessions. Update the rid-not-in-self._send_sessions branch in transceiver.py to follow the same pattern as respond_and_send_async by simply invoking _get_or_create_send_session(req) for its side effects, then keep setting _ever_had_send_session and _pipelined_chunk_offsets[rid] as before.
602-604: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNumerous open TODOs in the pipelined-send path before merge.
send_prefill_chunkandrespond_and_send_asynccarry several unresolvedTODO(athenac)questions on correctness-critical fields (token_range,mamba_state_index,layer_range, thereq.statetransition, the offset accumulation "might be a faulty calculation", and the redundancy between the two methods). Since the PR is marked WIP, these need resolution before this is production-ready. I can help draft the offset/metadata handling and consolidate the shared logic into a single helper.Also applies to: 608-611, 628-631, 657-666, 675-675
🤖 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/disaggregation/transceiver.py` around lines 602 - 604, The pipelined-send path in transceiver.py still contains unresolved correctness TODOs in send_prefill_chunk and respond_and_send_async, especially around token_range, mamba_state_index, layer_range, req.state transitions, and the offset accumulation logic. Resolve these TODO(athenac) questions by verifying the metadata semantics, fixing the offset calculation, and making the state update explicit and correct before merge. Also remove the duplicated logic between send_prefill_chunk and respond_and_send_async by consolidating the shared send/metadata assembly into a single helper so the two paths stay consistent.
🤖 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/disaggregation/native/transfer.py`:
- Around line 728-745: The chunked destination slicing in transfer logic is now
using chunk offsets, but the token alignment still assumes each chunk maps to a
suffix ending at token_range.end. Update the code around the chunked path in
transfer.py and the downstream token-start calculation to derive starts from
chunk_block_offset, or require callers to provide per-chunk KVSlice.token_range
for each chunk. Make sure the block selection and token-range alignment stay
consistent for prefix-cache and SWA cases so the written blocks match the
intended chunk.
- Around line 523-524: The abort/result notification path in transfer.py still
uses write_meta.slice_id, which can conflict with the receiver’s single-task
slice handling. Update the abort send logic in the relevant transfer routine to
mirror the success path by reporting receiver_slice_id as 0 for aborts too, so
the receiver does not see a later-chunk slice ID and hit its slice assertion.
Keep the existing task/event unblocking behavior intact while ensuring the
aborted/failure result is always sent to receiver slice 0.
- Around line 736-743: The chunk-to-destination mapping in transfer.py is too
strict for exhausted layer groups: when len(src_block_ids) is 0, the current
bounds check in the chunk slicing logic still raises on advanced chunk_offset
values. Update the chunk handling around the dst_block_ids slice so empty source
chunks become a no-op and do not trigger the out-of-bounds error; keep the
existing bounds validation for non-empty chunks in the same chunk
offset/full_dst_block_ids path.
In `@tensorrt_llm/_torch/disaggregation/transceiver.py`:
- Line 648: The `respond_and_send_async` skip guard in `transceiver.py` needs
both a lint fix and a logic check: move the `return` onto its own line to
satisfy E701, and verify the condition around `rid in self._send_sessions and
rid not in self._pipelined_chunk_offsets` correctly prevents duplicate sends
while pipelined chunks are still outstanding. If needed, adjust the guard so the
full `_create_kv_slices` resend path only runs when it is safe, using the
existing `_send_sessions` and `_pipelined_chunk_offsets` state to avoid
duplicate transfer.
- Around line 591-611: The chunking logic in send_prefill_chunk() and the
_pipelined_chunk_offsets update can split KV slices on token boundaries that are
not aligned to tokens_per_block, which causes the boundary block to be resent
and offsets to drift. Adjust the prefill chunk selection so every chunk boundary
lands on a KV block boundary (or clamp the sliding-window fallback so it only
overlaps when it evenly divides tokens_per_block), and then recompute
_pipelined_chunk_offsets from the actual block count in the chunk.
In `@tests/unittest/disaggregated/test_kv_transfer.py`:
- Around line 1789-1825: The send/receive flow is using the wrong API shape:
TxSession.send() and RxSession.receive() should be called with a fully populated
KVSlice rather than extra kwargs, and they do not return futures. Update the
test setup around KVSlice, sender_session.send(), and
receiver_sessions/RxSession.receive() to set chunk_block_offset and cuda_event
on the slice object before calling send/receive, then replace the .result()
waits with wait_complete()/wait_complete(blocking=True) on the session or slice
as appropriate.
---
Outside diff comments:
In `@tensorrt_llm/_torch/disaggregation/native/transfer.py`:
- Around line 488-507: The task transition in transfer.py is happening too late
in the prefill-transfer flow: `task._slice.cuda_event.synchronize()` runs while
the task is still `INIT`, so `cancel_request()` can miss it and free KV pages
too early. In the transfer path around `task`, `session.lock`, and
`TaskStatus.TRANSFERRING`, move the INIT→TRANSFERRING state update (with the
session ERROR/CANCELLED abort check) before waiting on the CUDA event, and keep
the abort branch ahead of synchronization so in-flight work is visible before
any blocking wait.
---
Nitpick comments:
In `@tensorrt_llm/_torch/disaggregation/transceiver.py`:
- Around line 582-585: The send-session initialization in transceiver logic has
a redundant assignment because _get_or_create_send_session already stores the
session in self._send_sessions. Update the rid-not-in-self._send_sessions branch
in transceiver.py to follow the same pattern as respond_and_send_async by simply
invoking _get_or_create_send_session(req) for its side effects, then keep
setting _ever_had_send_session and _pipelined_chunk_offsets[rid] as before.
- Around line 602-604: The pipelined-send path in transceiver.py still contains
unresolved correctness TODOs in send_prefill_chunk and respond_and_send_async,
especially around token_range, mamba_state_index, layer_range, req.state
transitions, and the offset accumulation logic. Resolve these TODO(athenac)
questions by verifying the metadata semantics, fixing the offset calculation,
and making the state update explicit and correct before merge. Also remove the
duplicated logic between send_prefill_chunk and respond_and_send_async by
consolidating the shared send/metadata assembly into a single helper so the two
paths stay consistent.
🪄 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: 784a62c4-609a-4545-bc83-7d70311b9078
📒 Files selected for processing (11)
tensorrt_llm/_torch/disaggregation/base/transfer.pytensorrt_llm/_torch/disaggregation/native/transfer.pytensorrt_llm/_torch/disaggregation/transceiver.pytensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/llmapi/llm_args.pytests/integration/defs/accuracy/test_disaggregated_serving.pytests/integration/test_lists/test-db/l0_dgx_b200.ymltests/unittest/disaggregated/test_chunked_transfer.pytests/unittest/disaggregated/test_kv_transfer.pytests/unittest/llmapi/test_llm_args.py
|
@coderabbitai check the test coverage, including function test, perf test, stress test and multi nodes test. |
|
Tip For best results, initiate chat on the files or code changes. 🧩 Analysis chain🏁 Script executed: #!/bin/bash
set -euo pipefail
echo "== Changed/mentioned test files =="
for f in \
tests/unittest/disaggregated/test_chunked_transfer.py \
tests/unittest/disaggregated/test_kv_transfer.py \
tests/unittest/llmapi/test_llm_args.py \
tests/integration/defs/accuracy/test_disaggregated_serving.py \
tests/integration/test_lists/test-db/l0_dgx_b200.yml \
tensorrt_llm/_torch/disaggregation/transceiver.py \
tensorrt_llm/_torch/pyexecutor/py_executor.py \
tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py
do
if [ -f "$f" ]; then
echo "-- outline: $f"
ast-grep outline "$f" --view expanded || true
echo
else
echo "Missing: $f"
fi
done
echo "== Search for perf/stress/multi-node/disaggregated related tests =="
rg -n -i -C2 \
'(perf|performance|benchmark|stress|soak|multi[-_ ]node|multinode|disaggregated serving|disaggregated|chunked|pipelined)' \
tests tensorrt_llm/_torch | sed -n '1,260p'
echo
echo "== Inspect relevant slices from changed tests =="
sed -n '1,260p' tests/unittest/disaggregated/test_chunked_transfer.py
echo "-----"
sed -n '1,260p' tests/unittest/disaggregated/test_kv_transfer.py
echo "-----"
sed -n '1,220p' tests/integration/defs/accuracy/test_disaggregated_serving.py
echo "-----"
sed -n '1,220p' tests/unittest/llmapi/test_llm_args.py
echo "-----"
sed -n '1,220p' tests/integration/test_lists/test-db/l0_dgx_b200.yml
echo
echo "== Inspect relevant implementation slices for testability points =="
sed -n '1,260p' tensorrt_llm/_torch/disaggregation/transceiver.py
echo "-----"
sed -n '1,220p' tensorrt_llm/_torch/pyexecutor/py_executor.py
echo "-----"
sed -n '1,220p' tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.pyLength of output: 144072 🏁 Script executed: #!/bin/bash
set -euo pipefail
echo "== Direct test references for new pipelined/chunked entry points =="
rg -n -C2 \
'_maybe_send_prefill_chunk|send_prefill_chunk|enable_pipelined_transfer|respond_and_send_async|request_and_receive_async|_create_kv_slices|create_kv_cache_transceiver' \
tests tensorrt_llm/_torch | sed -n '1,260p'
echo
echo "== Disaggregated multi-node / stress / soak style tests in test tree =="
rg -n -i -C2 \
'(multi[-_ ]node|multinode|soak|stress|long[-_ ]run|longevity|burn[-_ ]in)' \
tests | sed -n '1,220p'
echo
echo "== Show pipelined accuracy test and nearby YAML entries =="
sed -n '720,820p' tests/integration/defs/accuracy/test_disaggregated_serving.py
echo "-----"
rg -n -C2 'chunked_kv_transfer|pipelined_kv_transfer' tests/integration/test_lists/test-db/l0_dgx_b200.ymlLength of output: 49120
What is covered wellFunction / unit tests
Integration / accuracy
Gaps I found1. Some new entry points are not directly tested
So the core mechanics are exercised indirectly, but several newly added integration points are still missing targeted tests. 2. Pipelined accuracy test does not appear to be scheduled in the shown pre-merge list
So pipelined accuracy exists in code, but it does not look wired into this pre-merge list yet. 3. No PR-specific perf validation
Given the PR objective is explicitly performance-oriented, this is a notable gap. 4. No PR-specific stress coverage
But I did not find stress coverage specifically targeting:
5. No PR-specific multi-node coverage
But I did not find multi-node tests specifically exercising this PR’s new chunked/pipelined KV transfer path. Overall assessment
Recommended follow-ups
If you want, I can turn this into a concrete review comment with prioritized test additions. |
7705b98 to
37f1bdd
Compare
37f1bdd to
7173ea6
Compare
d3de936 to
aa98b71
Compare
51febc2 to
1c0ef43
Compare
Signed-off-by: Athena Cai <athenac@nvidia.com>
1c0ef43 to
9ac8af3
Compare
|
|
||
| if self.kv_cache_transceiver.kv_transfer_timeout_ms is not None: | ||
| req.py_kv_transfer_start_time = time.time() | ||
| if req.is_context_only_request and not req.is_finished_due_to_cancellation: |
There was a problem hiding this comment.
Need to handle cancellation + pipeline transfer in a follow up PR and test it.
There was a problem hiding this comment.
+1 to the good point about cancellation handling.
The first non-final send creates a live session, but DISAGG_CONTEXT_TRANS_IN_PROGRESS is set only for the final slice, so _try_cancel_request() returns success between chunks and termination can free source KV during NIXL. Also, after cancel_request() returns false for an active task, status polling removes the CANCELLED session without proving physical drain, and AUX is not included in has_transferring_tasks().
We may need to make ownership begin with the first slice and retain exact per-peer KV/AUX operations until quiescent.
| chunk_start_pos, chunk_end_pos = req.py_last_context_chunk | ||
| tpb = self._kv_cache_manager.tokens_per_block | ||
|
|
||
| chunk_start_block = chunk_start_pos // tpb |
There was a problem hiding this comment.
on a ctx-side prefix-reuse hit, context_current_position starts at prepopulated_prompt_len, so the first chunk starts there and no chunk ever covers [0, prepopulated_prompt_len). We should only not send the reused prefill if the gen side has it.
There was a problem hiding this comment.
This is a good catch, thanks. I ran into this while I was implementing early reused cache transfer and I'm thinking of just pulling that change onto this PR since the code diff is quite small.
The executor needs to schedule the transfer of the reused cache at some point, so it might as well just schedule it as soon as prepoulated prompt length is set.
There was a problem hiding this comment.
I'm fine with pulling it in this PR if the change is small.
Signed-off-by: Athena Cai <athenac@nvidia.com>
Signed-off-by: Athena Cai <athenac@nvidia.com>
chienchunhung
left a comment
There was a problem hiding this comment.
Thanks for the PR; nice work!
I took the first pass and left some comments; happy to take another pass once this PR resolves the conflicts with the base and the comments are addressed.
| block_ids, | ||
| chunk_block_offset=chunk_start, | ||
| chunk_block_count=chunk_block_count, | ||
| resident_block_end=chunk_end, |
There was a problem hiding this comment.
IIUC this may not work with V2’s incremental allocation.
For example, consider an eight-block prompt where only the first two blocks have been computed. V2 returns the two pages belonging to logical blocks 0 and 1, but this projection assumes any two-page list is the suffix of the full prompt and treats them as blocks 6 and 7. It therefore finds no overlap with the first chunk and sends an empty slice.
The current test misses this because it allocates all eight blocks before sending any chunk. Please track the pages’ actual logical range, or reject pipelining with V2 until that is supported, and test allocation and transfer one chunk at a time.
|
|
||
| if self.kv_cache_transceiver.kv_transfer_timeout_ms is not None: | ||
| req.py_kv_transfer_start_time = time.time() | ||
| if req.is_context_only_request and not req.is_finished_due_to_cancellation: |
There was a problem hiding this comment.
+1 to the good point about cancellation handling.
The first non-final send creates a live session, but DISAGG_CONTEXT_TRANS_IN_PROGRESS is set only for the final slice, so _try_cancel_request() returns success between chunks and termination can free source KV during NIXL. Also, after cancel_request() returns false for an active task, status polling removes the CANCELLED session without proving physical drain, and AUX is not included in has_transferring_tasks().
We may need to make ownership begin with the first slice and retain exact per-peer KV/AUX operations until quiescent.
| @@ -3713,7 +3723,15 @@ class CacheTransceiverConfig(StrictBaseModel, PybindMirror): | |||
| "Per-region size in MiB of the native-disagg KV-cache bounce buffer (one for send, one for recv). Bounce coalesces a request's scattered per-block KV into one contiguous fabric-VMM buffer and issues a single multi-rail NIXL write. The size doubles as the on/off switch: 0 (default) keeps the per-block path, >0 enables bounce at that capacity. Only used by the Python (v2) transceiver." | |||
| ) | |||
|
|
|||
| enable_pipelined_transfer: bool = Field( | |||
There was a problem hiding this comment.
Adding this nested public configuration field requires regenerating tensorrt_llm/usage/llm_args_golden_manifest.json; the field is currently absent.
Please run the generator, commit the manifest, and request telemetry/privacy CODEOWNER review for the new field.
| - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_pipelined_kv_transfer_nixl_python_accuracy[enable_block_reuse=False-disable_overlap_scheduler=False] | ||
| - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_pipelined_kv_transfer_nixl_python_accuracy[enable_block_reuse=True-disable_overlap_scheduler=False] | ||
| - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_pipelined_kv_transfer_nixl_python_accuracy[enable_block_reuse=False-disable_overlap_scheduler=True] | ||
| - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_pipelined_kv_transfer_nixl_python_accuracy[enable_block_reuse=True-disable_overlap_scheduler=True] |
There was a problem hiding this comment.
The four B200 test-list entries do not match the collected pytest node IDs, so these selectors cannot run the intended parameterized cases.
The stacked decorators produce disable_overlap_scheduler=...-enable_block_reuse=..., but the YAML lists the parameters in reverse order, so the claimed B200 coverage is not scheduled.
| req: The completed context request whose KV cache to transfer. | ||
| """ | ||
| if self.kv_transfer_timeout_ms is not None and req.py_kv_transfer_start_time is None: | ||
| req.py_kv_transfer_start_time = time.time() |
There was a problem hiding this comment.
Since #15238, main's timeout checks compare against time.monotonic() (assignments at py_executor.py:6051/6089), while this new line uses time.time() and will merge cleanly on rebase. Unless I'm misreading, elapsed then stays negative and the ctx-side kv_transfer_timeout_ms never fires — this check runs unconditionally, so non-pipelined sends too. Could all three new time.time() (here + two in py_executor.py) switch to monotonic() during the rebase?
Also, in pipelined mode the clock starts at the first chunk, while the request is not added to requests_in_transfer() until the final chunk. The remaining prefill compute therefore counts toward the timeout and may make the request time out immediately when it is registered. Is that the intended end-to-end timeout semantics?
| # does not support UCX, MPI, or MOONCAKE). | ||
| runtime = cache_transceiver_config.transceiver_runtime | ||
| use_python = runtime == "PYTHON" | ||
| if (runtime is None and cache_transceiver_config.enable_pipelined_transfer): |
There was a problem hiding this comment.
A post-rebase ordering question: main's KvCacheCreator._maybe_enable_fabric_memory_for_python_transceiver() (#16832) runs before KV-pool allocation and enables the fabric-memory default only when the config says transceiver_runtime == "PYTHON". The default "auto" (e.g. Llama) resolves to None, and the auto-selection here only flips the local use_python without persisting it — after the pool is already allocated. So enable_pipelined_transfer + default runtime would get a V1 C++ pool without fabric memory, affecting MNNVL transfers. The integration test sets "PYTHON" explicitly, so this path isn't covered.
Could "pipelined ⇒ PYTHON" be resolved and persisted before pool routing — via_resolve_transceiver_runtime_auto on the normal path, with a fallback or fail-fast for paths like AutoDeploy that skip it?
| result_msg = _make_kv_result_msg( | ||
| self._instance_rank, | ||
| write_meta.unique_rid, | ||
| 0, # receiver slice_id: always the single monolithic task |
There was a problem hiding this comment.
Could we explicitly distinguish sender_slice_id from receiver_slice_id in this protocol?
Before this PR , both sides had a single slice/task, so write_meta.slice_id == receiver_slice_id == 0 and the distinction was invisible. Now write_meta.slice_id identifies a sender chunk, while the ID carried by KV_AGENT_RESULT is used by RxSession.process_kv_agent_result() for some log.
Hardcoding 0 is correct for the current monolithic receiver, but the receiver still names the argument sender_slice_id and uses it at line 1888 as a local task index, which makes the contract misleading.
Could we rename the receiver-side argument to receiver_slice_id and ideally carry both IDs explicitly, e.g.:
sender_slice_id = write_meta.slice_id
receiver_slice_id = req_info.slice_id
This would also make per-chunk error reporting, bounce-buffer ownership less error-prone.
| req.state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS | ||
|
|
There was a problem hiding this comment.
Before this PR, context forward and KV transfer were mutually exclusive request phases: a request entered DISAGG_CONTEXT_TRANS_IN_PROGRESS only after context computation had finished.
With pipelined transfer, this is no longer true. After the first non-final session.send(slice), context forward and KV transfer can be active at the same time.
Setting DISAGG_CONTEXT_TRANS_IN_PROGRESS only for the last slice therefore means the request state no longer fully represents whether transfer resources are in flight.
This is especially problematic for cancellation and resource cleanup: code that checks only LlmRequestState may conclude that the request is not transferring and release its KV pages while a chunk is still being read by NIXL.
Could we model transfer activity independently from the compute/request phase, for example with a separate kv_transfer_state or ownership flag? A combined state
such as TRANS_AND_CONTEXT_FORWARD_OVERLAP could work, but a separate transfer dimension may avoid a cross-product of context, transfer, cancellation, and error
states.
| chunk_start = min(chunk_start_block, total_blocks) | ||
| chunk_end = min(chunk_end_block, total_blocks) | ||
| chunk_block_count = max(0, chunk_end - chunk_start) | ||
| chunk_block_ids = [ |
There was a problem hiding this comment.
Sliding-window source lists are already trimmed to a suffix of the whole prompt but are projected here as a suffix of the current chunk, while the receiver anchors on the whole prompt, every middle chunk then sends uncomputed blocks, silently, since only counts are compared.
|
|
||
| if self.kv_cache_transceiver.kv_transfer_timeout_ms is not None and req.py_kv_transfer_start_time is None: | ||
| req.py_kv_transfer_start_time = time.time() | ||
| elif self.kv_cache_transceiver.pipeline_transfer_enabled: |
There was a problem hiding this comment.
This fires for any context-only request still in the batch, including ones an error path already failed and freed, whose chunk bounds are then still unset, please gate it on a request that is genuinely still prefilling.
| @@ -6043,28 +5738,10 @@ def kv_connector_request_finished(req: LlmRequest): | |||
| if self.kv_cache_transceiver: | |||
| self._check_disagg_ctx_cache_transfer_status(0) | |||
There was a problem hiding this comment.
This sweep now runs against requests that are still mid-prefill, so one failed chunk frees a live request's KV while sibling writes may still be on the wire, an intermediate failure needs to stop prefill and drain in flight work first.
| @@ -459,6 +459,7 @@ def _apply_aux(self, session, req: LlmRequest): | |||
| req.context_phase_params.draft_tokens = draft_tokens | |||
|
|
|||
| def _get_or_create_send_session(self, req: LlmRequest) -> TxSessionBase: | |||
There was a problem hiding this comment.
Tearing down a send session also drops the peer's receive registration, yet the next chunk silently re-creates a session with nobody to send to, so the request hangs and its blocks leak.
| "transceiver instead of the C++ transceiver to enable " | ||
| "pipelined KV cache transfer. " | ||
| "Set transceiver_runtime='CPP' to disable this auto-selection.") | ||
| use_python = True |
There was a problem hiding this comment.
The Python runtime is chosen locally but never persisted, so the two earlier decisions keyed off that same field still take the C++ path, hybrid Mamba models end up with the C++ cache manager and block reuse left on.
There was a problem hiding this comment.
The replay path snapshots work under the lock but enqueues outside it, so a concurrently added chunk can jump ahead of older ones and break the ordering the last-chunk marker relies on, the loop also has no per-item error isolation.
| "beam_width > 1 is not supported when enable_pipelined_transfer is set." | ||
| ) | ||
|
|
||
| disagg_params = request.py_disaggregated_params |
There was a problem hiding this comment.
The scheduling style this depends on is an orchestrator setting while the flag is per worker with no startup cross-check, so a mismatched deployment comes up healthy and then rejects every request, including the server's own health probe.
| chunk_end_block = (chunk_end_pos + tpb - 1) // tpb | ||
| is_last_chunk = req.context_remaining_length == 0 | ||
|
|
||
| base_slice = self._create_kv_slice(req) |
There was a problem hiding this comment.
Every chunk rebuilds the request's entire block map and discards all but the current chunk, turning a one-time cost into a per-iteration one on the executor's critical path.
fredricz-20070104
left a comment
There was a problem hiding this comment.
Review summary - CONCERNS
Verdict: Cannot merge as-is — the branch is in conflict (mergeable_state=dirty) and has two outstanding CHANGES_REQUESTED on the head commit while the PR is still WIP (benchmarking). The feature logic and tests are largely solid, but there is one correctness concern worth confirming before merge.
Concerns
-
[MAJOR]
native/transfer.py(~_deliver_kv_to_agent, lines 488-528) - cuda_event sync happens while task is still INIT- What is wrong: The worker waits on
task._slice.cuda_event.synchronize()before thesession.lock-heldINIT→TRANSFERRINGtransition. (The sync line is just outside this diff, so please confirm the current ordering — coderabbitai flagged the same at 488-507.) - How it fails: While the worker is blocked in
synchronize(), the task is stillINIT. A concurrentcancel_request()iterateskv_tasks, sees noTRANSFERRINGtask, and frees/reuses the request's KV pages. When the event completes the worker issues the RDMA into freed pages → use-after-free / corrupted KV on the wire. - Suggested fix: Perform the
INIT→TRANSFERRINGtransition (undersession.lock, with the CANCELLED/ERROR abort check) beforecuda_event.synchronize(), and return early on the abort path.
- What is wrong: The worker waits on
-
[MAJOR] Merge blocker — dirty branch + change requests.
mergeable=false,mergeable_state=dirty. Rebase onto main, resolve conflicts, and address the CHANGES_REQUESTED fromShixiaowei02andnv-xtfon commit4c6c035. PR description still reads 'Status: Benchmarking performance impact'.
Minor notes (non-blocking)
tests/unittest/disaggregated/test_kv_transfer.py-test_transfer_worker_chunkedis missing theif not torch.cuda.is_available(): pytest.skip(...)guard thattest_transfer_worker_pipelinedhas; it will error rather than skip on a CUDA-less runner._build_prefill_chunkassumes chunk boundaries are KV-block-aligned (floor start / ceil end). This is guaranteed by chunked-prefill's block-multiple invariant and validated later byderive_chunk_block_coords, but consider an explicit assert in_build_prefill_chunkto fail fast if a non-aligned boundary ever reaches it.
QA view
- Test coverage: adequate — new unit tests (
test_chunked_transfer.py, additions totest_kv_transfer.py) cover projection math, TxSession/RxSession state machines, chunked & pipelined send paths, and the requirement validators; a new integration accuracy test is added. Uncovered: the concurrent cancel-during-cuda_event-sync race in finding #1. - SM coverage: architecture-independent (pure Python/RDMA transfer logic, no
get_sm_version/arch guards). Integration test runs Hopper+ (skip_pre_hopper) and is listed on B200. - Test code: missing CUDA skip guard on
test_transfer_worker_chunked; otherwise clean (context-managed server launch, explicit cleanup). - Test time: significant — four new disaggregated GSM8K accuracy parametrizations on
l0_dgx_b200, each launching a ctx+gen server pair. - Needs
/qa-verify: yes — WIP multi-GPU feature, dirty branch, and the cancel/cuda_event race has no reproducing test; re-run the disagg accuracy suite after conflicts are resolved.
Possible new issues
- Non-block-aligned chunk boundaries could overlap/transfer partial blocks (mitigated by chunked-prefill's block-multiple invariant).
prompt_lenchanged fromOptional=Noneto required acrossSessionArgsBase/KVSendTask/TxSession/RxSession; external callers not shown may break._build_kv_write_metanow usestask._prompt_lenforslice_endinstead oftoken_range.end; recheck the speculative-decoding extra-draft-block path for a one-block start-offset shift.
What I could not verify
- The exact placement of
cuda_event.synchronize()relative to the lock transition (line lies just outside the diff). - All external callers of the constructors whose
prompt_lenbecame required. - Runtime accuracy/perf of the pipelined path (integration test not executed here).
Automated review by NVCortex Lite, run by @fredricz-20070104.
Documentation (WIP) https://docs.google.com/document/d/1Z9ARCc48QNCbKfoTEhZT1W4W440x6QsbKfVh3ksKIW0/edit?tab=t.0
Status: Benchmarking performance impact
Description
Summary
Implements pipelined prefill-transfer for disaggregated serving. Instead of waiting for all prefill chunks to complete before starting KV cache transfer, each chunk's KV data is transferred to the generation server immediately after its prefill completes. This overlaps GPU compute with RDMA transfer, hiding transfer latency behind prefill computation. When transfer time per chunk is less than prefill time per chunk (typical for 100+ Gbps NIC with long-context workloads), transfer latency is nearly fully hidden.
Configuration
Feature is gated behind enable_pipelined_transfer flag
Generation first scheduling is required so that generation server can receive chunks before context compute completes
Enable_chunked_prefill is required, number of blocks transferred in a chunk is determined by the chunked prefill scheduling
Added assertions for both of these requirements when enable_pipelined_transfer is set
Pipelined transfer of KV Cache
Related work by @chienchunhung for chunked KV cache transfer Reducing KV Block Residency and Peak Memory Pressure in Disaggregated Serving
Implementation:
Feature is gated behind the
enable_pipelined_transferflag onCacheTransceiverConfig(tensorrt_llm/llmapi/llm_args.py). The flag is consumed by the Python transceiver only and has no C++ counterpart (_to_pybinddoes not forward it).Requirements (all validated when
enable_pipelined_transferis set):schedule_style=generation_first) is required so the generation server can receive chunks before context compute completes. Enforced inPyExecutor._validate_request(raisesValueError).enable_chunked_prefillis required; the number of blocks transferred in a chunk is determined by the chunked-prefill scheduling (req.py_last_context_chunk). Enforced increate_kv_cache_transceiver(raisesValueError).enable_chunked_prefillis plumbed throughcreate_py_executor_instance→create_kv_cache_transceiver(and the AutoDeploy shimad_executor.py).beam_width == 1; beam search is unsupported for chunked transfer. Enforced inPyExecutor._validate_request(raisesValueError) and asserted in_build_prefill_chunk.Transceiver auto-selection: pipelined transfer is only implemented in the Python transceiver (the C++ transceiver does not support it). When
enable_pipelined_transferis set and no explicittransceiver_runtimeis given, the Python transceiver is auto-selected for the NIXL/DEFAULT backends. Settingtransceiver_runtime='CPP'explicitly, or using a backend other than NIXL/DEFAULT, raisesValueError.KVSlicegains atotal_blocksfield describing the full logical block span[0, total_blocks)that chunk offsets and the resident-suffix projection are expressed in.SessionArgsBase.prompt_len(and theTxSession/RxSession/KVSendTaskconstructors) are now required (previouslyOptional).KVSlice.is_last_slicewas a latent feature that was alwaysTruebefore this feature, now it is set toFalsefor intermediate slices.PyExecutor):is_context_finished/is_finished_due_to_length): release the IndexMapper slot, thenstart_transfer(commits blocks to the reuse tree and pins them) beforerespond_and_send_asyncsends the last KV slice.respond_and_send_asyncsends the chunk during ongoing prefill.respond_and_send_asynccreates (or reuses) aTxSessionand, for the pipelined path, calls_build_prefill_chunkwhich builds aKVSlice. Performs a projection onto each layer group’s resident blocks so that the correct chunks are sent when SWA is used. Only whenslice.is_last_sliceis true does it call_finalize_sendand set the request state toDISAGG_CONTEXT_TRANS_IN_PROGRESS.PyExecutor.update_requests_send_kv_result_to_receiver: Update receiver with status with slice_id always set to 0 (the receiver has a single monolithic task, so sender-side chunking is transparent to it). Intermediate chunk results are sent (not suppressed) so RDMA failures propagate immediately.derive_chunk_block_coordsderives the global chunk block offset/count from the block-alignedtoken_range.Failed transfer
Upon failed transfer, the per-slice
KVSendTask, theTxSession, and (indirectly) the request state are updated.Test coverage
Unit tests
tests/unittest/disaggregated/test_chunked_transfer.py(new) — exercises the session state machine with realTxSession/RxSessionclasses and stub sender/receiver objects:project_blocks_to_global_chunkno-ops when a chunk is outside a short layer group's resident range, and maps prefix-reuse suffixes by overlap rather than raw offset indexing._build_kv_write_metaprojects an asymmetric (SWA) layer group's suffix blocks onto the overlapping global chunk, checking resultingsrc_ptrs/dst_ptrs/sizes.TxSession/RxSessionmulti-slice status: notKV_TRANSFERREDuntil all tasks complete,ERRORon any task failure,wait_completeblocks on all tasks and returnsFAILEDon partial failure, and mid-chunk failure reportsERROR.pipeline_transfer_enabledreflects the flag;ValueErrorwhen chunked prefill is missing (create_kv_cache_transceiver) or when the request is not generation-first (_validate_request);respond_and_send_asyncsends the built chunk and only finalizes on the last chunk.tests/unittest/disaggregated/test_kv_transfer.py— added pipelined prefill chunk creation coverage:_send_prefill_chunksdrives the real_build_prefill_chunkpath per chunk. Parametrizedtest_send_prefill_chunks_basiccovers no-chunking, even/uneven splits, empty blocks, and chunk-larger-than-total; integrity check reassembles per-layer-group block IDs; multi-layer-group test checks projection of a shorter layer group;mamba_state_indexpropagation is preserved.test_transfer_with_gen_prefix_offsetis parametrizedsingle_slice/sender_chunked.test_transfer_worker_chunkedandtest_transfer_worker_pipelinedrun end-to-end sender-chunked and incremental-pipelined transfers for V1 and V2 KV cache managers and verify block-data equality.Integration / accuracy tests
tests/integration/defs/accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_pipelined_kv_transfer_nixl_python_accuracy— GSM8K accuracy over the NIXL Python transceiver withenable_pipelined_transfer=True,enable_chunked_prefill=True,schedule_style=generation_first, and capped prefill chunk size. Parametrized overenable_block_reuseanddisable_overlap_scheduler(4 combinations). Registered intests/integration/test_lists/test-db/l0_dgx_b200.yml.