[TRTLLM-12499][feat] Pipelined KVCache transfer for disaggregated serving in Python Cache Transceiver - #15727
Conversation
|
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:
WalkthroughAdds sender-side pipelined KV transfer for disaggregated serving. The change introduces block projection, separate sender and receiver identifiers, session retirement, configuration validation, executor integration, and unit and integration coverage. ChangesPipelined KV transfer
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant PyExecutor
participant KvCacheTransceiverV2
participant TxSession
participant RxSession
PyExecutor->>KvCacheTransceiverV2: send prefill chunk
KvCacheTransceiverV2->>TxSession: send projected KVSlice
TxSession->>RxSession: deliver KV result with sender and receiver IDs
RxSession-->>PyExecutor: resolve receiver task and report completion
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Description checkExplanation The description is detailed and relevant. It explains the motivation, configuration requirements, architecture, lifecycle behavior, validation, tests, and known coverage gaps. The template checklist confirmation is not reproduced, but the substantive description and test coverage sections are complete. ✨ 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: 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
athena-nv
changed the title
37f1bdd to
7173ea6
Compare
d3de936 to
aa98b71
Compare
…ving in Python Cache Transceiver 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. Only the last chunk's transfer remains on the critical path. The feature is gated behind `enable_pipelined_transfer` on `CacheTransceiverConfig` and is implemented in `KvCacheTransceiverV2` only. It requires `schedule_style: generation_first`, `enable_chunked_prefill: true`, `beam_width == 1`, the NIXL backend, `kv_cache_bounce_size_mb == 0`, `pipeline_parallel_size == 1` on the sender, and a non-Mamba/hybrid cache manager. Each requirement is enforced at startup or per request. Squashed from 15 commits: - Chunking is sender-side only; the generation server posts a single receive covering the whole prompt and completes on `is_last_slice`. - `KVSlice` now describes one chunk rather than one whole request, gaining `total_blocks` and a meaningful `is_last_slice`. `prompt_len` became required on the session args so SWA can compute the stale-block boundary. - `project_blocks_to_global_chunk` intersects ranges instead of indexing, so resident-suffix block lists (sliding window groups, prefix reuse, incremental allocation) project correctly onto a global chunk. - The first slice always extends back to block 0, so a context-side prefix-reuse hit does not leave `[0, prepopulated_prompt_len)` unsent. - Source blocks are capped at the computed chunk boundary before SWA trimming, normalizing V1's full-prompt reservation against V2's incremental allocation. - `KV_AGENT_RESULT` carries `sender_slice_id` and `receiver_slice_id` separately, making per-chunk RDMA failures attributable. Behavior-neutral for the monolithic receiver. - KV transfer activity is modeled by transceiver session membership rather than `LlmRequestState`, so mid-prefill cancellation and transfer-timeout monitoring work during the pipelined phase. - A retired send session cannot be silently re-created, since closing it drops the peer's `RecvReqInfo` and the receiver never re-registers. - `TxSession.dispatch_lock` serializes chunk dispatch across the executor thread and the late-peer replay path, so a newer slice cannot reach a peer's queue ahead of an older one. - Transceiver configuration resolution happens early and idempotently, and backend/runtime compatibility validation is centralized. Signed-off-by: Athena Cai <athenac@nvidia.com> Simpilify _build_kv_write_meta logic Signed-off-by: Athena Cai <athenac@nvidia.com> Carry the pipelined chunk window as a TokenRange The chunk cursor on KVSlice goes back to the existing TokenRange rather than a new ChunkCoords dataclass, so the slice keeps one field for "how far does this slice reach" instead of gaining a second vocabulary for it. The window is still decided in block space by _build_prefill_chunk, so the range it emits is block-aligned and the sender asserts that before dividing it back out. TokenRange now admits an empty range, which a chunk clamped past the end of the prompt produces. Signed-off-by: Athena Cai <athenac@nvidia.com> Drop cancellation, timeout, and failure-path chunked transfer tests Removes the error-path coverage from test_chunked_transfer.py: the cancelled-request send gates, the transfer-timeout sweeps, the session ERROR/FAILED status cases, and the retired-send-session block. The file now covers chunk projection, slice-id addressing, dispatch ordering, and the pipelined config gates only. Signed-off-by: Athena Cai <athenac@nvidia.com> Restore WriteMeta's single slice_id field WriteMeta carried a sender_slice_id/receiver_slice_id pair so the sender could log its own chunk index while addressing the peer's task on the wire. Only the peer's index is needed, so the field goes back to the pre-existing slice_id and _deliver_kv_to_agent resolves the send task from write_meta.task instead of indexing the session by chunk. Sender logs lose per-chunk attribution. Signed-off-by: Athena Cai <athenac@nvidia.com> Remove KVSlice.total_blocks Signed-off-by: Athena Cai <athenac@nvidia.com> Round down at block boundaries Signed-off-by: Athena Cai <athenac@nvidia.com>
Signed-off-by: Athena Cai <athenac@nvidia.com>
… request state. 2. Executor records the failed request. 3. After the current sampling/state-update phase—the next safe boundary—it sets DISAGG_TRANS_ERROR and runs normal cleanup. 4. The request is excluded before the next scheduling pass. Signed-off-by: Athena Cai <athenac@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #70523 [ run ] triggered by Bot. Commit: |
|
PR_Github #70523 [ run ] completed with state
|
Signed-off-by: Athena Cai <athenac@nvidia.com>
|
/bot run |
|
PR_Github #70592 [ run ] triggered by Bot. Commit: |
|
PR_Github #70592 [ run ] completed with state
|
Signed-off-by: Athena Cai <athenac@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #70813 [ run ] triggered by Bot. Commit: |
chienchunhung
left a comment
chienchunhung
left a comment
There was a problem hiding this comment.
Regarding the ownership lifecycle, we don't need duplicate the implementation in this PR now that others (#17720 and #18041) started enforcing them. Let's make sure TRTLLM-15184 documents the remaining timeout/cancellation issue; specifically retaining Tx KV, Rx KV, and auxiliary resources until all physical operations drain.
|
PR_Github #70813 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #70891 [ run ] triggered by Bot. Commit: |
|
PR_Github #70891 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #70960 [ run ] triggered by Bot. Commit: |
|
PR_Github #70960 [ run ] completed with state |
| """Return whether this rank supports pipelined KV transfer.""" | ||
| if not cfg.enable_pipelined_transfer: | ||
| return False | ||
| blockers = [] |
There was a problem hiding this comment.
Could this list also cover the capacity scheduler policy? Under max utilization the C++ scheduler can pause a context request past its first chunk and the executor frees its blocks in the same iteration, while an earlier chunk is still being read, and the KV connector already blocks that policy in the same way.
There was a problem hiding this comment.
Will be addressed in subsequent PRs.
| == LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS): | ||
| return True | ||
| return (self.kv_cache_transceiver is not None | ||
| and self.kv_cache_transceiver.has_inflight_transfer(request)) |
There was a problem hiding this comment.
With a session existing from the first chunk this answer now depends on whether a worker thread has picked the newest task up yet, so tensor parallel ranks can disagree and split the batch; returning false unconditionally for a context only request and letting the rank uniform retirement finish it would avoid that.
There was a problem hiding this comment.
Will be addressed in subsequent PRs.
SimengLiu-nv
left a comment
SimengLiu-nv
left a comment
There was a problem hiding this comment.
LGTM for KVCM changes.
15 participants
Summary
Implements pipelined prefill transfer for disaggregated serving. Instead of waiting for every
prefill chunk to finish before transferring the request's KV cache, the context server sends the
completed KV blocks after each chunk. This overlaps GPU prefill compute with RDMA transfer and
leaves only the final transfer on the critical path when transfer is faster than prefill.
Chunking remains sender-side only. The context server creates multiple
KVSendTasks in oneTxSession; the generation server posts one receive for the whole request and completes thatreceive when the final-slice result arrives.
Related work by @chienchunhung:
KV cache transfer: Reducing KV Block Residency and Peak Memory Pressure in Disaggregated Serving
Configuration
The feature is gated by
CacheTransceiverConfig.enable_pipelined_transfer.resolve_cache_transceiver_config()intensorrt_llm/_torch/pyexecutor/config_utils.pyhandles runtime-independent selection:transceiver_runtime: CPPis rejected.create_py_executor()and again increate_kv_cache_transceiver()for callers that bypass the executor-creation path.Runtime-dependent restrictions are checked by
KvCacheTransceiverV2._resolve_pipelined_transfer()after the mapping and cache manager exist:KvCacheTransceiverV2resolve_cache_transceiver_config()pipeline_parallel_size == 1_resolve_pipelined_transfer()kv_cache_bounce_size_mb == 0_resolve_pipelined_transfer()_resolve_pipelined_transfer()enable_chunked_prefill: truepy_last_context_chunkPyExecutor._validate_request()beam_width == 1PyExecutor._validate_request()and_build_prefill_chunk()schedule_style == generation_firstPyExecutor._validate_request()The last three checks are request checks, limited to disaggregated context-only requests using a
pipelined transceiver. Aggregate/warmup requests and generation-only requests are not rejected.
The disaggregated YAML parser and
serve.pydo not perform feature-specific startup validation;the CLI schedule-style override continues to be applied after parsing.
Architecture
Sender-side chunks, monolithic receiver
flowchart LR subgraph ctx [Context server] exec[PyExecutor._send_kv_async] --> transceiver[respond_and_send_async] transceiver --> build[_build_prefill_chunk] build --> tx["TxSession: one per request"] tx --> c0["KVSendTask 0"] tx --> c1["KVSendTask 1"] tx --> cn["KVSendTask N-1, is_last_slice"] end subgraph gen [Generation server] rx["RxSession: one per request"] --> task["one whole-request KVRecvTask"] end c0 -->|RDMA write| task c1 -->|RDMA write| task cn -->|final RDMA write| taskrequest_and_receive_async()is unchanged: the receiver allocates and advertises the fulldestination once. Each sender result echoes the receiver's task id, and
is_last_slicetells thereceiver when no more chunks remain.
Slice model
KVSlice.token_rangedistinguishes the two paths:token_range=Noneand coversprompt_len.is_last_slice=Trueonly for the final emitted slice.SessionArgsBase.prompt_lenis the request-wide extent. It is independent of a chunk'stoken_rangeand is required for prompt-wide calculations such as the final sliding-windowboundary.
Deriving a chunk
KvCacheTransceiverV2._build_prefill_chunk()converts the scheduler'sreq.py_last_context_chunktoken bounds to block coordinates:The first slice starts at block zero so a context-side prefix-reuse hit is included even though
the scheduler starts computing at
prepopulated_prompt_len. Non-final partial blocks are deferredto the next chunk; the final chunk rounds up so the prompt tail is not stranded. If a non-final
chunk completes no transferable block, no slice is emitted.
The final slice is still emitted when its token range is empty. This preserves the final-slice
signal for cases where all remaining attention groups were already sent or contain no blocks.
Sliding-window attention
SWA layers are deliberately not pipelined. Pages can leave the active window between prefill
chunks, so sending each intermediate resident suffix can either transfer stale pages or lose pages
that are still needed in the generation server's final window.
For a layer group whose
sliding_window_size < prompt_len:that group; it first trims the receiver list to the source window and then aligns the two lists.
Non-windowed groups continue to transfer chunk by chunk. A model with mixed full-attention and SWA
groups therefore pipelines the full-attention KV while deferring only the SWA groups.
Native sender addressing
project_blocks_to_global_chunk()now lives next to its only production consumer intensorrt_llm/_torch/disaggregation/native/transfer.py. It intersects a global chunk intervalwith a block list represented as a resident suffix; a non-overlapping list produces an empty
projection.
Sender._build_kv_write_meta()uses the slice end as the coordinate-space extent:For partial non-SWA chunks, the receiver's whole-prompt block list is projected to the same global
chunk. For the final SWA slice, that projection is skipped so the complete active window remains
addressable. Generation-side prefix reuse and sender/receiver window differences are then handled
by
_trim_receiver_window_head()and_align_kv_blocks().This keeps the addressing correct for:
Result protocol and slice ids
The sender and receiver have different slice namespaces: the sender has one task per prefill
chunk, while the receiver currently has one whole-request task.
WriteMetatherefore carries both:slice_id: the local sender task/chunk id, used for sender-side task lookup and diagnostics;receiver_slice_id: the peer's task id copied fromRecvReqInfo, used inKV_AGENT_RESULT.The binary wire prefix remains
<qqq?Bq>:RxSession.process_kv_agent_result()validates and indexesreceiver_slice_id. The protocol shapedoes not change; current receivers advertise task zero, so every sender chunk reports
receiver_slice_id == 0.TxSession completion and registration races
A
TxSessionowns all KV tasks and the optional auxiliary task for a request._has_last_sliceis set when a final slice is appended._has_last_slice, at least one task, and every KV task transferred.This prevents an intermediate chunk from making the session appear complete before the final
task exists.
FULLY_TRANSFERREDadditionally requires the generation-first auxiliary transfer.has_transferring_tasks()includes both KV and auxiliary tasks.TxSession.lockatomically appends a task and snapshots currently registered peers. The listener'slate-registration path uses the same lock to save peer info and snapshot existing tasks. Actual
metadata construction and queue dispatch happen after releasing the lock, keeping transfer work
outside the critical section while ensuring a task is seen by either the send path or the replay
path.
Executor and lifecycle integration
Sending from the executor loop
PyExecutor._send_kv_async()runs after each forward step:GENERATION_COMPLETE.async_transfer_manager.start_transfer()before sending, then send thefinal slice and start the final-transfer timeout clock.
respond_and_send_async()without moving the request into thetransfer manager.
respond_and_send_async()creates or reuses theTxSession, builds the next chunk, and finalizesauxiliary state only when the emitted slice is final. Request state changes to
DISAGG_CONTEXT_TRANS_IN_PROGRESSonly at that point; session membership tracks ownership duringearlier chunks.
No additional CUDA synchronization is introduced. The existing sampler event synchronization
completes before
_send_kv_async()dispatches the slice.Cancellation, failure, and safe retirement
Pipelining allows fabric writes while the request is still in
CONTEXT_INIT, so request statealone cannot answer whether its KV pages are safe to release.
KvCacheTransceiverV2.has_inflight_transfer()derives ownership from the send/receive sessionmaps.
_is_request_in_transmission()combines request phase with session ownership, routingmid-prefill cancellation through the transceiver.
TRANSFERRINGuntil the physical write finishes. INIT KV and auxiliary tasks are failedimmediately.
py_kv_send_session_retiredprevents recreating a sender session after teardown has removed thepeer registration. A later send attempt transitions the request to
DISAGG_TRANS_ERROR.In TP/PP deployments, logical outcome and physical quiescence are reconciled across ranks.
_ctx_consensus_outcome()packs cancelled, failed, completed, and locally quiesced request ids intothe existing outcome allgather, avoiding a second collective per sweep:
transferring KV or auxiliary task.
A bounded wait timeout is nonterminal: the session remains live and retains its pages because a
peer write may still be active. The polling path logs the timeout and retries on a later sweep.
Tests
The updated unit coverage exercises:
TxSession._has_last_slicecompletion gating;GENERATION_COMPLETEsend suppression;The old parser-level disaggregated-config tests were removed because feature-specific validation
now occurs in the executor/transceiver, not in
extract_disagg_cfg().Change inventory since
1a80e56eebc420684d7fc90dc14ad8c3a633e42ctensorrt_llm/_torch/disaggregation/base/transfer.pytensorrt_llm/_torch/disaggregation/native/transfer.pytensorrt_llm/_torch/disaggregation/transceiver.pytensorrt_llm/_torch/pyexecutor/config_utils.pytensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.pytensorrt_llm/_torch/pyexecutor/llm_request.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/py_executor_creator.pytensorrt_llm/commands/serve.pytensorrt_llm/llmapi/disagg_utils.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_disagg_utils.pytests/unittest/disaggregated/test_kv_transfer.pytests/unittest/disaggregated/test_transceiver_bounded_polling.pyThe cumulative diff is 781 insertions and 976 deletions across 17 files.