[TRTLLM-12499][feat] Pipelined KVCache transfer for disaggregated serving in Python Cache Transceiver

[athena-nv]

Documentation (WIP) https://docs.google.com/document/d/1Z9ARCc48QNCbKfoTEhZT1W4W440x6QsbKfVh3ksKIW0/edit?tab=t.0

Status: In review

Follow up PR to remove redundant KVCM work: #17526

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.

Baseline (monolithic transfer)
  compute    [ chunk0 ][ chunk1 ][ chunk2 ][ chunk3 ]
  transfer                                            [======= whole prompt =======]
  gen unblocked at                                                                  ^

Pipelined transfer
  compute    [ chunk0 ][ chunk1 ][ chunk2 ][ chunk3 ]
  transfer             [  c0  ]  [  c1  ]  [  c2  ]  [  c3  ]
  gen unblocked at                                           ^

Only the last chunk's transfer remains on the critical path. Everything before it is paid for out of prefill compute time.

Related work by @chienchunhung for chunked KV cache transfer: Reducing KV Block Residency and Peak Memory Pressure in Disaggregated Serving

Configuration

The feature is gated behind enable_pipelined_transfer on CacheTransceiverConfig (tensorrt_llm/llmapi/llm_args.py). The flag is consumed by the Python transceiver only and has no C++ counterpart (_to_pybind does not forward it).

# Context server config
context_servers:
  cache_transceiver_config:
    backend: "NIXL"             # or "DEFAULT", which resolves to NIXL
    enable_pipelined_transfer: true
  enable_chunked_prefill: true

# Generation server config (same setting)
generation_servers:
  cache_transceiver_config:
    backend: "NIXL"
    enable_pipelined_transfer: true
  enable_chunked_prefill: true

# Disagg server config
schedule_style: generation_first

Requirements and where they are enforced

Requirement	Why	Enforced in
schedule_style: generation_first	The generation server must have registered its destination blocks before context compute finishes, otherwise there is nothing to write chunks into	_validate_disagg_config in tensorrt_llm/llmapi/disagg_utils.py (startup, from the disagg YAML) and PyExecutor._validate_request (per request)
enable_chunked_prefill: true	Chunk boundaries come from the chunked-prefill scheduler (req.py_last_context_chunk); there is no separate chunk-size knob	create_kv_cache_transceiver (ValueError)
beam_width == 1	Beam search packs a non-contiguous 1-D block layout that the chunk projection does not model	PyExecutor._validate_request (ValueError), asserted again in _build_prefill_chunk
NIXL backend + Python transceiver	Pipelining is implemented only in KvCacheTransceiverV2	resolve_cache_transceiver_config
kv_cache_bounce_size_mb == 0	The bounce path coalesces a whole request's blocks into one contiguous staging buffer, which is incompatible with per-chunk writes	resolve_cache_transceiver_config (ValueError)
pipeline_parallel_size == 1 on the sender	PP splits layers across ranks, so a chunk is not complete on any single rank when its prefill step ends	create_kv_cache_transceiver (ValueError); the generation server is unaffected and may still use PP
No Mamba/hybrid attention cache	Mamba recurrent state is one mutable slot per request, not block-addressable per chunk. Reusing that slot in concurrent slice writes cannot preserve a snapshot of each chunk's state	create_kv_cache_transceiver rejects MambaHybridCacheManager and any explicit Mamba cache manager (ValueError)

Transceiver auto-selection: when enable_pipelined_transfer is set and no explicit transceiver_runtime is given, resolve_cache_transceiver_config selects PYTHON and logs a warning. Setting transceiver_runtime='CPP' explicitly, or using a backend that resolves to something other than NIXL, raises ValueError.

The Mamba/hybrid gate applies to both MixedMambaHybridCacheManager and CppMambaHybridCacheManager. A normal transfer copies recurrent state once after prefill has produced its final value. Pipelined transfer would attach the same mamba_state_index to every chunk while prefill continues mutating that slot; because the writes are asynchronous, a slice has no stable per-chunk recurrent-state snapshot. The implementation therefore fails during transceiver creation instead of allowing a configuration that could transfer an intermediate or inconsistently observed state.

Two notes on the validation added in this PR:

• The disagg config file is validated at parse time (extract_disagg_cfg calls _validate_disagg_config), so a mismatched schedule_style fails before any worker starts rather than on the first request. parse_disagg_config_file also takes schedule_style_override so the --schedule_style CLI flag participates in the same validation instead of being applied afterward.
• _validate_request rejects only a request that carries py_disaggregated_params with a non-generation-first style. A request with no disaggregated params (an aggregate request hitting the same executor, e.g. during warmup) is allowed through.

Architecture

Sender-side chunking, monolithic receiver

Chunking is entirely a sender-side concept. The context server splits its source blocks into N chunks and issues N writes; the generation server posts a single receive covering the whole prompt and never learns how many chunks arrived.

flowchart LR
    subgraph ctx [Context server]
        exec[PyExecutor._send_kv_async] --> tcv[respond_and_send_async]
        tcv --> bld[_build_prefill_chunk]
        bld --> tx["TxSession, one per request"]
        tx --> t0["KVSendTask slice_id 0"]
        tx --> t1["KVSendTask slice_id 1"]
        tx --> tn["KVSendTask slice_id N-1, is_last_slice"]
    end
    subgraph gen [Generation server]
        rx["RxSession, one per request"] --> r0["KVRecvTask slice_id 0, whole prompt"]
    end
    t0 -->|RDMA write| r0
    t1 -->|RDMA write| r0
    tn -->|"RDMA write, is_last"| r0

Loading

This asymmetry is what keeps the change contained: request_and_receive_async on the generation side is unchanged, and the receiver completes when it has seen expected_transfers results carrying is_last_slice.

KVSlice

KVSlice (tensorrt_llm/_torch/disaggregation/base/transfer.py) previously described one whole request. It now describes one chunk of one request:

• is_last_slice was a latent field that was always True. It is now False for intermediate chunks, and it is the signal that drives finalization on both sides.
• total_blocks (new) is the full logical block span [0, total_blocks) that chunk offsets and the resident-suffix projection are expressed in. Without it, a receiver-side projection has no way to tell where a resident suffix begins.
• chunk (new, ChunkCoords) is where this slice sits in that span, as a block offset and count. Its presence is what makes a slice a chunk, so a monolithic transfer keeps its whole-request addressing untouched — including the packed beam layout, which the chunked path is not written for and which _build_prefill_chunk rejects via beam_width == 1.
• token_range (removed). It had exactly one consumer: the sender read token_range.end to recover total_blocks = ceil(end / tpb), the anchor for all the suffix arithmetic. .start was never read, since _create_kv_slice always produced [0, prompt_len). Chunking took both jobs away — a non-final chunk's end is neither the request's block span, which the destination projection needs, nor its prompt_len, which SWA needs — so they moved to the explicit total_blocks field and to a required session prompt_len respectively, leaving nothing that reads the range.

Reusing the vacated field to carry the chunk cursor would have given it a third, unrelated meaning, encoded as block_index * tokens_per_block and divided back out on arrival — and it cannot carry that meaning honestly anyway, because the chunk window is decided rather than observed (see "Deriving the chunk"). So a slice now describes its extent exactly one way: a chunk by its chunk, a monolithic transfer by the session's prompt_len, which is where the request's extent already lives.

Deleting the field moves one invariant: _create_kv_slice must produce ceil(prompt_len / tpb) blocks and must not include the num_extra_kv_tokens slots speculative decoding reserves, since an extra block would shift every per-layer token start. That was previously expressed as token_range.end == prompt_len; TestCreateKvSliceBlockSpan now asserts it on the block list directly, which is the thing the sender actually consumes.

SessionArgsBase.prompt_len (and the TxSession / RxSession / KVSendTask constructors) changed from Optional[int] to required. SWA needs the request's prompt_len to compute the stale-block boundary:

                stale_end = max(0, (task._prompt_len + 1 - window_size) // tpb)
                src_start = max(stale_end * tpb, src_start)
                dst_start = max(stale_end * tpb, dst_start)

Previously the slice's own extent was an acceptable stand-in for prompt_len, because a slice was the whole request. A chunk stops at its own boundary, so using it would place the sliding window in the wrong place for every chunk but the last.

TxSession and KVSendTask

TxSession already held a list of KV tasks; the list simply now has more than one entry. TxSession.send assigns slice_id = len(self.kv_tasks), so a chunk's id is its arrival order.

Session status aggregates over all tasks:

        if self._exception is not None or any(t.status == TaskStatus.ERROR for t in self.kv_tasks):
            return SessionStatus.ERROR
        if self.aux_task is not None and self.aux_task.status == TaskStatus.ERROR:
            return SessionStatus.ERROR
        kv_all_transferred = bool(self.kv_tasks) and all(
            t.status == TaskStatus.TRANSFERRED for t in self.kv_tasks
        )
        if kv_all_transferred:
            if self.aux_task is not None and self.aux_task.status == TaskStatus.TRANSFERRED:
                return SessionStatus.FULLY_TRANSFERRED
            return SessionStatus.KV_TRANSFERRED
        if self.kv_tasks and any(t.status == TaskStatus.TRANSFERRING for t in self.kv_tasks):
            return SessionStatus.TRANSFERRING
        return SessionStatus.READY if self.receiver_ready else SessionStatus.INIT

So the session is KV_TRANSFERRED only once every chunk has landed, ERROR if any chunk failed, and TRANSFERRING while any chunk is mid-write. wait_complete waits on every task, and has_transferring_tasks() (used by cancel_request) reports whether any chunk is mid-write.

Slice ordering. With N slices per session there are now two producers that can enqueue writes for the same peer: send() on the executor thread, and Sender._respond_with_kv replaying already-created slices when a peer registers late. TxSession.dispatch_lock spans both the snapshot and the enqueue loop in each path, because the receiver completes on the is_last_slice result without checking that earlier slices landed — a newer slice reaching a peer's queue ahead of an older one would let the generation server start decoding on incomplete KV.

Executor loop integration

sequenceDiagram
    participant Exec as PyExecutor
    participant Sampler as TorchSampler
    participant Tcv as KvCacheTransceiverV2
    participant Tx as TxSession
    participant Worker as Sender worker thread
    participant Rx as RxSession on gen

    loop intermediate prefill chunk
        Exec->>Exec: _forward_step(scheduled_batch)
        Exec->>Sampler: _update_requests(sample_state)
        Sampler-->>Exec: sampler_event.synchronize()
        Exec->>Tcv: respond_and_send_async(req)
        Tcv->>Tcv: _build_prefill_chunk(req)
        Tcv->>Tx: send(slice) with is_last_slice False
        Tx->>Worker: dispatch KVSendTask slice_id i
        Worker->>Rx: RDMA write then KV_AGENT_RESULT
    end

    Note over Exec: final chunk: is_context_finished or is_finished_due_to_length
    Exec->>Exec: release_index_slot(req)
    Exec->>Exec: async_transfer_manager.start_transfer(req)
    Exec->>Tcv: respond_and_send_async(req)
    Tcv->>Tx: send(slice) with is_last_slice True
    Tcv->>Tcv: _finalize_send: pack and send aux, set ContextPhaseParams
    Tcv->>Exec: req.state = DISAGG_CONTEXT_TRANS_IN_PROGRESS

Loading

Both branches live in PyExecutor._send_kv_async, which runs after each forward step:

            for req in scheduled_requests:
                if req.is_context_only_request and not req.is_finished_due_to_cancellation:
                    if req.is_context_finished or req.is_finished_due_to_length:
                        # Forward is done for this request — release the
                        # IndexMapper slot so new requests can reuse it.
                        # KV blocks stay allocated for the upcoming transfer.
                        if hasattr(self.kv_cache_manager, 'release_index_slot'):
                            self.kv_cache_manager.release_index_slot(
                                req.py_request_id)
                        # Order matters: start_transfer commits the request's blocks to the reuse
                        # tree and pins them, and must run before respond_and_send_async sends the
                        # final KV slice and (for the Python transceiver) transitions the request toward completion.
                        self.async_transfer_manager.start_transfer(req)

Points worth calling out:

• No extra synchronization is needed. TorchSampler._update_requests already calls state.sampler_event.synchronize() before _send_kv_async runs, so the chunk's KV writes are complete on the device by the time the slice is dispatched. An earlier revision carried a cuda_event on KVSlice for this; it was removed as redundant.
• respond_and_send_async handles both cases. It creates or reuses the TxSession via _get_or_create_send_session, builds a chunk with _build_prefill_chunk when pipelining is on, and only calls _finalize_send and sets DISAGG_CONTEXT_TRANS_IN_PROGRESS when slice.is_last_slice is true.
• Ordering on the final chunk. start_transfer commits the request's blocks to the reuse tree and pins them; it must run before the last slice is sent, because sending the last slice is what starts the request's transition toward completion.

Chunk projection

The hard part of sending a chunk is that "chunk" is defined in one coordinate space (global block position within the prompt) while every block list involved is a resident suffix of some other range. Sliding-window layer groups, prefix reuse, and incremental allocation during prefill all shorten a list from the front. Indexing such a list with a raw global offset yields the wrong blocks.

project_blocks_to_global_chunk (base/transfer.py) resolves this by intersecting ranges rather than indexing:

    resident_start = max(0, resident_block_end - len(block_ids))
    resident_end = resident_block_end
    chunk_start = chunk_block_offset
    chunk_end = chunk_start + chunk_block_count

    overlap_start = max(chunk_start, resident_start)
    overlap_end = min(chunk_end, resident_end)
    if overlap_start >= overlap_end:
        return block_ids[:0]

    local_start = overlap_start - resident_start
    local_end = overlap_end - resident_start
    return block_ids[local_start:local_end]

A list that does not reach the chunk at all returns empty rather than raising or silently sending the wrong blocks.

Deriving the chunk

_build_prefill_chunk (transceiver.py) turns the scheduler's token chunk into block coordinates:

        chunk_start_pos, chunk_end_pos = req.py_last_context_chunk
        tpb = self._kv_cache_manager.tokens_per_block

        # A ctx-side prefix-reuse hit starts the first chunk at
        # prepopulated_prompt_len, so no chunk covers [0, prepopulated_prompt_len).
        # ... first slice extends back to block 0 ...
        is_first_chunk = chunk_start_pos == req.prepopulated_prompt_len
        chunk_start_block = 0 if is_first_chunk else chunk_start_pos // tpb
        chunk_end_block = (chunk_end_pos + tpb - 1) // tpb
        is_last_chunk = req.context_remaining_length == 0

        prompt_blocks = (req.prompt_len + tpb - 1) // tpb
        total_blocks = prompt_blocks

        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)
        base_slice = self._create_kv_slice(
            req, resident_block_end=chunk_end
        )
        all_block_ids = base_slice.block_ids_per_layer_groups

total_blocks is always ceil(prompt_len / tpb) — the full prompt span, which is what the destination side is allocated for and what KVSlice.total_blocks carries to the sender worker. The chunk bounds are clamped to it. Passing resident_block_end=chunk_end also normalizes the two cache-manager allocation models: V1's full-prompt reservation is capped at the computed boundary, while V2's incrementally allocated list is already bounded there.

This is the only place chunk geometry is decided, and the resulting ChunkCoords travels to the sender worker as-is. That matters because the window is not a function of the scheduler's token bounds alone: extending the first chunk back to block 0 depends on prepopulated_prompt_len, the clamp depends on prompt_len, and the rounding rule below depends on tokens_per_block. A sender re-deriving the window from a token range would have to be handed the decided answer encoded as block_index * tpb and divide it back out — the same numbers, one lossy detour, plus an alignment precondition to validate on arrival.

Unaligned chunk boundaries round up

Nothing requires the scheduler to cut chunks on block boundaries, and the two schedulers disagree about whether it does. V1 makes it true by construction: setPrepopulatedPromptLen shrinks the first chunk so prepopulatedPromptLen + chunkSize floors to a block boundary, and TLLM_CHECKs the result to prevent cache fragmentation. The V2 Python scheduler rounds the chunk size to chunk_unit_size but not the absolute end, so a partial-block reuse hit offsets every boundary that follows — a gemma3 V2 run with tokens_per_block=32 produced a non-final chunk ending at token 259.

An earlier revision asserted alignment here and failed on exactly that. The transfer layer should not dictate scheduler chunk policy, so the rule is instead:

chunk_start_block = floor(chunk_start_pos / tpb)   (or 0 for the first chunk)
chunk_end_block   = ceil(chunk_end_pos / tpb)

An unaligned boundary therefore lands one block in both chunks. The earlier chunk sends it with a stale tail past the boundary; the next chunk, whose start floors into the same block, rewrites it whole once the tokens are computed.

flowchart LR
    subgraph chunkA [Chunk A, ends at token 259]
        a1["blocks 0..7<br/>fully computed"]
        a2["block 8<br/>rounded up, tail is stale"]
    end
    subgraph chunkB [Chunk B, starts at token 259]
        b1["block 8 again<br/>start floors to 8"]
        b2["blocks 9 onward"]
    end
    a2 -->|"same worker thread, second write waits"| b1

Loading

Four properties make this safe:

• The round up never over-reads. ceil(chunk_end_pos / tpb) is the index just past the block holding the last computed token, and that block is allocated — the forward pass just wrote into it. So resident_block_end = chunk_end stays within the source list under both allocation models.
• The overlap is exactly one block, and it is always re-sent. Chunk ends round up and the next chunk's start floors, so the two differ by one block precisely when the boundary is unaligned. Since starts floor, the following chunk always covers it. The last chunk rounds up too, so the prompt's final partial block is covered.
• Last-writer-wins is ordered, not racy. Sender._enqueue routes by (unique_rid, peer_rank) to a single worker thread, and _deliver_kv_to_agent blocks on status.wait() before taking the next item. The stale copy is fully landed before the computed copy is even submitted.
• The transient state is never observed. The generation server completes only when it has seen is_last_slice from every peer, so no decode reads the destination while a partial block is sitting there.

The cost is at most one extra block per chunk boundary, negligible at production chunk sizes. Rounding down instead was rejected: it leaves the partial tail unsent until the next chunk (which must start at floor regardless, so the block is re-sent either way), and it produces empty slices whenever a chunk falls entirely inside one block.

The first slice always starts at block 0

The scheduler's chunk sequence does not necessarily cover the whole prompt. On a context-side prefix-reuse hit, setPrepopulatedPromptLen advances context_current_position to prepopulated_prompt_len before the first chunk is cut, so py_last_context_chunk starts at P = prepopulated_prompt_len and no chunk ever spans [0, P). Those blocks are resident and valid — _create_kv_slice forces cached_per_lg = [0] * len(layer_groups) on the context side, so the base slice holds them — but the projection would drop them and the generation server would decode over whatever those pages held:

tpb = 4, prompt_len = 32 (total_blocks = 8), reuse hit P = 12 (3 blocks)

global block      0    1    2  | 3    4    5  | 6    7
base slice src    r0   r1   r2 | n3   n4   n5 | n6   n7    <- all resident, all valid
scheduler chunks              [--- chunk 0 ---][- chunk 1 -]
without the rule  --   --   -- | n3   n4   n5 | n6   n7
                  ^^^^^^^^^^^^ never sent

So the first slice extends its start back to block 0, carrying the reused prefix along with the first computed chunk. The monolithic path was never affected: it sends the whole base slice.

prepopulated_prompt_len is written exactly once per request and chunk starts increase monotonically, so only the first chunk satisfies the equality; with no reuse P == 0 and the rule is a no-op. req.is_first_context_chunk cannot substitute for it, because it compares context_current_position against prepopulated_prompt_len and the cursor has already advanced by the time _send_kv_async runs — the recorded chunk start is the pre-advance value.

Both projections still hold with chunk_start = 0. On the source, _create_kv_slice(..., resident_block_end=chunk_end) first removes uncomputed V1 pages and trims VSWA groups to the overlap between the computed prefix and the final prompt window. On the destination (resident_block_end = total_blocks) the overlap is [max(0, G_b), chunk_end), which no-ops whatever the generation server already has. And when the whole prompt fits in one chunk alongside a reuse hit, the slice degenerates to a chunk spanning [0, total_blocks): suffix_end_blocks equals total_blocks and the destination projection is the identity, so the write is addressed byte-for-byte like a non-pipelined transfer. test_whole_prompt_chunk_addresses_like_a_monolithic_slice pins that by building the same slice with and without chunk and comparing the resulting WriteMeta.

This gap was invisible in CI because it is masked whenever the generation server has the same prefix cached: its RecvReqInfo block list is trimmed by cache_skip, dst_start rises above the gap, and _align_kv_blocks would have trimmed those blocks anyway. The bug bites only when the two cache states diverge — generation-side reuse off, a cold generation cache, different eviction pressure, or different DP routing.

Two projections with different resident_block_end

The same helper is called on both sides of the transfer with a deliberately different end bound:

Side	Projected in	resident_block_end	Why
Source blocks	_build_prefill_chunk	chunk_end	The source is normalized to the computed prefix: V1's full-prompt allocation is capped at chunk_end, while V2 is already incremental.
Destination blocks	_build_kv_write_meta	total_blocks	The generation server allocated the whole prompt up front when it posted its single receive, so its list is a suffix of the full span.

Using total_blocks for the source was the bug fixed in Fix resident chunk end calculation (0628b97): it placed resident_start too far left, so every intermediate chunk selected blocks from the wrong position in a partially allocated source list.

VSWA chunk projection: V1 cache manager

Consider a 16-block prompt, a final 4-block VSWA suffix [12,16), and an intermediate chunk [11,13). Intervals are end-exclusive, so this chunk has computed logical blocks 11 and 12 and overlaps the final window only at block 12. Page names pN below are illustrative physical pages holding logical block N.

global block        0 ... 10 | 11 | 12 | 13 | 14 | 15
current chunk                 [---- 11..13 ----)
final VSWA suffix                   [------ 12..16 ------)
required overlap                    [12,13)

V1 reserves physical pages for the entire prompt before all chunks have been computed. The source must therefore be capped before taking the VSWA suffix:

flowchart LR
    raw["V1 raw pages<br/>p0 ... p15<br/>full prompt reserved"]
    cap["Cap at chunk_end = 13<br/>p0 ... p12"]
    trim["Trim below final stale_end = 12<br/>p12"]
    project["Project chunk [11,13)<br/>p12"]
    dst["Write destination<br/>logical block 12"]
    raw --> cap --> trim --> project --> dst

Loading

Without the cap, trimming first would produce p12 p13 p14 p15. Reinterpreting that list as a suffix ending at block 13 could select future pages such as p14 and silently pair one with destination block 12.

VSWA chunk projection: V2 cache manager

V2 grows and evicts incrementally. At chunk_end = 13, its live 4-block VSWA range is [9,13), so it exposes p9 p10 p11 p12 and has no future pages. It still needs the same final-window trim; otherwise the sender can pair an earlier computed page with destination block 12.

flowchart LR
    raw["V2 raw live pages<br/>p9 p10 p11 p12<br/>range [9,13)"]
    cap["Cap at chunk_end = 13<br/>no change"]
    trim["Trim below final stale_end = 12<br/>p12"]
    project["Project chunk [11,13)<br/>p12"]
    dst["Write destination<br/>logical block 12"]
    raw --> cap --> trim --> project --> dst

Loading

After normalization, both managers send exactly one computed source page for the same logical position the receiver requested:

V1 source  p12  ──>  destination block 12
V2 source  p12  ──>  destination block 12

Sender-worker side

_build_kv_write_meta (native/transfer.py) applies the destination projection and then reduces everything to a token-space alignment:

• slice.chunk selects the path. It is set only by _build_prefill_chunk, so the sender reads the chunk window rather than inferring one, and a monolithic transfer — including a packed beam layout, which the chunked path does not model — cannot accidentally be routed through it.
• suffix_end_blocks is chunk.block_offset + chunk.block_count when chunked and total_blocks otherwise; per-layer token starts follow from (suffix_end_blocks - n_blocks) * tpb.
• Those starts are then floored by req_info.dst_start_token (generation-side prefix reuse) and by the SWA stale_end, and _align_kv_blocks trims both arrays to the shared token overlap. That single overlap computation covers all four cases: no prefix cache, context-side prefix cache, generation-side prefix cache, and a chunk that falls entirely inside the generation server's already-cached prefix (which produces an empty transfer).

Wire protocol: sender_slice_id and receiver_slice_id

The KV_AGENT_RESULT frame used to carry a single slice id. The receiver used it purely as an index into its own task list, so the field was always semantically the receiver's id — but it was named sender_slice_id on the receive side, and once chunking existed the sender hardcoded 0 into it while its real chunk id sat unused in write_meta.slice_id. The consequences were a per-chunk RDMA failure that always logged slice=0, and two send paths that disagreed about what the field meant (_send_failed_result_to_receiver already sent info.slice_id, while _send_kv_result_to_receiver sent a literal 0).

The frame now carries both ids explicitly:

sequenceDiagram
    participant TxSession
    participant SenderWorker
    participant RxSession
    TxSession->>SenderWorker: KVSendTask.slice_id = chunk index 0..N-1
    Note over SenderWorker: WriteMeta.sender_slice_id = task.slice_id<br/>WriteMeta.receiver_slice_id = req_info.slice_id
    SenderWorker->>RxSession: KV_AGENT_RESULT with both ids, is_last, status
    Note over RxSession: index _kv_tasks[receiver_slice_id]<br/>log sender_slice_id only

Loading

• _KV_RESULT_PREFIX widened from struct.Struct("<qqq?B") to struct.Struct("<qqqq?B"); the field order is instance_rank, unique_rid, sender_slice_id, receiver_slice_id, is_last, status. This is a ctx/gen wire format with no version negotiation — the receiver unpacks whatever arrives against its compiled-in struct, so both servers must run matching builds.
• NO_SLICE_ID = -1 means "no sender slice for this result", used by _send_failed_result_to_receiver when a session fails before any KVSendTask exists.
• WriteMeta.sender_slice_id = task.slice_id identifies the sender's chunk and is carried for logging and cross-side correlation only. WriteMeta.receiver_slice_id = req_info.slice_id if req_info.slice_id is not None else 0 is the peer's own task index, echoed back from RecvReqInfo, and is what resolves the task.
• RxSession.process_kv_agent_result(peer_rank, receiver_slice_id, sender_slice_id, is_last_slice, status, ...) indexes _kv_tasks[receiver_slice_id] and includes sender_slice_id in the assertion message, the bounce-scatter-failure path, the perf warning, the completion debug line, and the FAILED detail — so a chunk-level RDMA failure is attributable to a specific chunk.

This is behavior-neutral for the monolithic receiver: receiver_slice_id is 0 in every current deployment. It is naming plus an explicit protocol field, laying the groundwork for a future multi-task receiver.

KV transfer state

Before pipelined transfer, DISAGG_CONTEXT_TRANS_IN_PROGRESS was a faithful proxy for "the fabric may be reading this request's KV pages". It no longer is: after the first non-final session.send(slice), chunks are in flight while the request is still in CONTEXT_INIT.

Two failures followed from the gate keying entirely off request state. Cancelling a request mid-prefill took the "nothing to cancel" path in _try_cancel_request, so _handle_responses terminated it and free_resources released pages a KVSendTask might still be reading, while the TxSession leaked in _send_sessions and the receiver was never notified. Separately, respond_and_send_async starts the timeout clock on the first chunk, but _check_kv_transfer_timeout only walked async_transfer_manager.requests_in_transfer() — which the request does not enter until start_transfer on the last chunk — so the entire pipelined phase was unmonitored while its clock ran.

The fix is two orthogonal dimensions instead of one overloaded state. LlmRequestState keeps meaning "compute and response phase". Transfer activity is answered by the component that actually owns the resources — the transceiver's session maps — so it cannot drift out of sync the way a mirrored request field would.

flowchart LR
    subgraph phase [Request phase - LlmRequestState]
        ctxInit[CONTEXT_INIT] --> transProg[DISAGG_CONTEXT_TRANS_IN_PROGRESS]
        transProg --> complete[DISAGG_CONTEXT_COMPLETE]
    end
    subgraph transfer [Transfer ownership - transceiver session maps]
        nosession[no session] --> active[session in _send_sessions]
        active --> torn[session closed and deleted]
    end
    ctxInit -.->|first non-final send| active
    torn -.->|safe to free KV| complete

Loading

Session membership is the right record: _get_or_create_send_session inserts before the first send, and every teardown path (cancel_request, the completed and cancelled loops in check_context_transfer_status, _close_failed_sessions) deletes it.

Concretely:

• KvCacheTransceiver.has_inflight_transfer(req) and has_any_inflight_transfer() are non-abstract and default to False, which leaves BindKvCacheTransceiver untouched — the C++ transceiver has no pipelining, so state and transfer activity coincide there. KvCacheTransceiverV2 implements both from session membership; get_unique_rid returning None for a non-disagg request naturally yields False.
• _is_request_in_transmission returns True when either the state says so or has_inflight_transfer(request) does. Its only caller is _try_cancel_request, so the blast radius is contained: a mid-prefill cancel now routes through KvCacheTransceiverV2.cancel_request, which cancels the TxSession, notifies the receiver, and returns False while any task is TRANSFERRING. The existing retry in _handle_canceled_requests then holds the KV pages until the write drains — the same behavior the monolithic path already relies on.
• _send_kv_async snapshots canceled_req_ids before the loop and gates only the intermediate-chunk branch on it. A session whose _terminal_status is already CANCELLED would otherwise be fed another chunk and produce a spurious FAILED result to the receiver. The last-chunk branch is untouched so nothing can be stranded in _requests_in_transfer.
• PyExecutor._has_any_inflight_kv_transfer() ORs kv_cache_transceiver.has_any_inflight_transfer() into async_transfer_manager.has_any_inflight_requests(), and replaces the latter at the four loop-level gates that would otherwise suppress timeout checking during the pipelined phase (the three _check_kv_transfer_timeout call sites and the KV-pressure RuntimeError).
• _check_kv_transfer_timeout gains a sweep over active context-only requests that are absent from requests_in_transfer, have a non-None py_kv_transfer_start_time, and have an in-flight transfer.
• _check_disagg_ctx_cache_transfer_status gains the matching recovery sweep: for a request flagged py_kv_transfer_timed_out but not yet known to the transfer manager, call cancel_request and set DISAGG_TRANS_ERROR only when it returns True (session closed, nothing mid-write). That reuses the established lever — _check_cache_transfer_errors picks it up for non-ADP and _handle_disagg_cache_errors_synced votes on it under ADP — so no new divergence path is introduced.

[coderabbitai]

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

Walkthrough

Adds 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.

Changes

Pipelined KV transfer

Layer / File(s)	Summary
Transfer contracts and block projection tensorrt_llm/_torch/disaggregation/base/transfer.py	Adds block-coordinate helpers, KVSlice.total_blocks, and required prompt_len values.
Native protocol and session handling tensorrt_llm/_torch/disaggregation/native/transfer.py	Separates sender and receiver slice IDs, updates the binary result format, serializes dispatch, projects chunk metadata, and centralizes failure results.
Prefill chunk construction and retirement tensorrt_llm/_torch/disaggregation/transceiver.py	Builds projected prefill chunks, tracks in-flight transfers, and retires sender sessions after completion, failure, or cancellation.
Configuration and executor integration tensorrt_llm/llmapi/*, tensorrt_llm/_torch/pyexecutor/*, tensorrt_llm/commands/serve.py, tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py	Adds pipelined-transfer settings and validation. Executor scheduling and timeout paths include transceiver-owned transfers.
Validation coverage tests/unittest/disaggregated/*, tests/integration/defs/accuracy/test_disaggregated_serving.py, tests/integration/test_lists/test-db/l0_dgx_b200.yml, tests/unittest/_torch/executor/*	Covers projection, protocol IDs, replay ordering, lifecycle failures, configuration restrictions, cancellation, timeout behavior, end-to-end transfer, and accuracy cases.

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

Loading

Possibly related PRs

• NVIDIA/TensorRT-LLM#17480: Shares Python/NIXL transceiver configuration and validation changes.

Suggested reviewers: brnguyen2, qijune, bowenfu

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 62.01% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly identifies the feature and matches the primary changes for pipelined KV cache transfer in Python disaggregated serving.
Description check	✅ Passed	The description clearly explains the motivation, design, configuration requirements, implementation details, and validation approach, but it does not explicitly complete the template checklist.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

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 win

Mark the task as in-flight before waiting on the CUDA event.

Line 492 waits while the task is still INIT, so cancel_request() can see no TRANSFERRING tasks 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 value

Redundant session assignment.

_get_or_create_send_session already inserts the session into self._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 in respond_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 win

Numerous open TODOs in the pipelined-send path before merge.

send_prefill_chunk and respond_and_send_async carry several unresolved TODO(athenac) questions on correctness-critical fields (token_range, mamba_state_index, layer_range, the req.state transition, 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

📥 Commits

Reviewing files that changed from the base of the PR and between a3026c9 and 9aed65c.

📒 Files selected for processing (11)

• tensorrt_llm/_torch/disaggregation/base/transfer.py
• tensorrt_llm/_torch/disaggregation/native/transfer.py
• tensorrt_llm/_torch/disaggregation/transceiver.py
• tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py
• tensorrt_llm/_torch/pyexecutor/py_executor.py
• tensorrt_llm/llmapi/llm_args.py
• tests/integration/defs/accuracy/test_disaggregated_serving.py
• tests/integration/test_lists/test-db/l0_dgx_b200.yml
• tests/unittest/disaggregated/test_chunked_transfer.py
• tests/unittest/disaggregated/test_kv_transfer.py
• tests/unittest/llmapi/test_llm_args.py

[xinhe-nv]

@coderabbitai check the test coverage, including function test, perf test, stress test and multi nodes test.

[coderabbitai]

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.py

Length 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.yml

Length of output: 49120

---

@xinhe-nv I checked the coverage. Short version: function coverage is decent, but perf/stress/multi-node coverage for this new feature is still incomplete.

What is covered well

Function / unit tests

• tests/unittest/disaggregated/test_chunked_transfer.py
  • Covers KVSendTask, TxSession, RxSession
  • Covers multi-slice status/error handling
  • Covers enable_pipelined_transfer
  • Covers the “skip re-send” branch in respond_and_send_async
• tests/unittest/disaggregated/test_kv_transfer.py
  • Covers _create_kv_slices
  • Covers chunk integrity across layer groups
  • Covers mamba_state_index propagation
  • Adds end-to-end chunked and pipelined transfer tests

Integration / accuracy

• tests/integration/defs/accuracy/test_disaggregated_serving.py
  • test_chunked_kv_transfer_nixl_python_accuracy at Line 744
  • test_pipelined_kv_transfer_nixl_python_accuracy at Line 787

Gaps I found

1. Some new entry points are not directly tested
I did not find direct tests for:

• tensorrt_llm._torch.disaggregation.transceiver.KvCacheTransceiverV2.send_prefill_chunk
• tensorrt_llm._torch.pyexecutor.py_executor.PyExecutor._maybe_send_prefill_chunk
• tensorrt_llm._torch.pyexecutor.kv_cache_transceiver.create_kv_cache_transceiver
• tensorrt_llm._torch.disaggregation.transceiver.KvCacheTransceiverV2.request_and_receive_async

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

• tests/integration/test_lists/test-db/l0_dgx_b200.yml includes the chunked accuracy test at Line 20 and Line 21
• I did not find the new test_pipelined_kv_transfer_nixl_python_accuracy in that YAML

So pipelined accuracy exists in code, but it does not look wired into this pre-merge list yet.

3. No PR-specific perf validation
I found general perf infrastructure in the repo, but nothing added here that validates the PR goal:

• no latency assertion
• no overlap/hidden-transfer metric
• no benchmark comparing pipelined vs non-pipelined transfer

Given the PR objective is explicitly performance-oriented, this is a notable gap.

4. No PR-specific stress coverage
The repo has general stress suites, including disaggregated stress entries in:

• tests/integration/test_lists/qa/llm_function_stress.txt

But I did not find stress coverage specifically targeting:

• chunked KV transfer
• pipelined KV transfer
• repeated chunk-send / cancel / failure / long-run stability for this path

5. No PR-specific multi-node coverage
The repo does have multi-node disaggregated tests, e.g.:

• tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py
• tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes_service_discovery.py
• multi-node perf YAMLs under tests/integration/test_lists/test-db/

But I did not find multi-node tests specifically exercising this PR’s new chunked/pipelined KV transfer path.

Overall assessment

• Function test coverage: good
• Accuracy coverage: present for both chunked and pipelined
• Pre-merge coverage: incomplete, because pipelined accuracy does not appear scheduled
• Perf coverage: missing for this feature
• Stress coverage: missing for this feature
• Multi-node coverage: missing for this feature

Recommended follow-ups

1. Add direct unit tests for:
   • send_prefill_chunk
   • _maybe_send_prefill_chunk
   • create_kv_cache_transceiver auto-selection / warnings
   • request_and_receive_async
2. Add the pipelined accuracy test to the appropriate test-db YAML.
3. Add at least one perf sanity test for:
   • pipelined vs non-pipelined transfer
   • visible transfer latency / overlap benefit
4. Add a stress case for repeated chunked/pipelined disaggregated serving.
5. Add a multi-node disaggregated case specifically for pipelined KV transfer.

If you want, I can turn this into a concrete review comment with prioritized test additions.

[athena-nv]

/bot run

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (4)

tests/unittest/disaggregated/test_chunked_transfer.py (2)

657-666: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Remove the unused _chunk_size_blocks stub attribute.

KvCacheTransceiverV2.pipeline_transfer_enabled reads only _enable_pipelined_transfer. The production class defines no _chunk_size_blocks attribute. Setting it here suggests a dependency that does not exist.

♻️ Proposed cleanup

     transceiver = MagicMock()
     transceiver._enable_pipelined_transfer = False
-    transceiver._chunk_size_blocks = 64
 
     result = KvCacheTransceiverV2.pipeline_transfer_enabled.fget(transceiver)

🤖 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 `@tests/unittest/disaggregated/test_chunked_transfer.py` around lines 657 -
666, Remove the unused transceiver._chunk_size_blocks = 64 stub assignment from
test_pipelined_transfer_disabled_by_default; keep the test focused on the
_enable_pipelined_transfer value read by
KvCacheTransceiverV2.pipeline_transfer_enabled.

---

15-51: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add the remaining pipelined-transfer coverage.

Coverage: 55 test functions cover projection, slice IDs, session lifecycle, cancellation, retirement, and configuration. The accuracy tests are registered in tests/integration/test_lists/test-db/l0_dgx_b200.yml, and unit tests are collected by directory.

Add tests for:

• _close_failed_sessions(..., mark_retired=True) when rid is absent from _send_reqs.
• Non-block-aligned chunk ends, asserting successive destination writes do not overlap.
• Multiple chunks through a real respond_and_send_async and TxSession; current tests mock the session and chunk builder.

Run pytest tests/unittest/.

🤖 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 `@tests/unittest/disaggregated/test_chunked_transfer.py` around lines 15 - 51,
Add unit tests in the existing chunked-transfer test module covering three gaps:
verify _close_failed_sessions with mark_retired=True retires sessions whose rid
is absent from _send_reqs; exercise non-block-aligned chunk ends and assert
successive destination writes do not overlap; and drive multiple chunks through
the real respond_and_send_async and TxSession paths without mocking the session
or chunk builder. Follow the existing fixtures and assertions, then run pytest
tests/unittest/.

Sources: Coding guidelines, Path instructions

tests/unittest/disaggregated/test_kv_transfer.py (2)

1996-2001: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Set token_range on the receiver slice for parity with the chunked helper.

add_and_verify_chunked_request posts the receive slice with token_range=TokenRange(start=0, end=request_len) (Line 1774). The pipelined helper omits token_range. The docstring states the receiver posts one slice covering the whole prompt, so make the range explicit. This keeps the two helpers comparable and removes reliance on the default None.

♻️ Proposed change

     for recv_session, block_ids_per_groups in zip(receiver_sessions, gen_block_ids):
         full_slice = KVSlice(
             is_last_slice=True,
             block_ids_per_layer_groups=block_ids_per_groups,
+            token_range=TokenRange(start=0, end=request_len),
         )
         recv_session.receive(full_slice)

🤖 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 `@tests/unittest/disaggregated/test_kv_transfer.py` around lines 1996 - 2001,
Update the receiver-side KVSlice construction in the pipelined helper to set
token_range explicitly to TokenRange(start=0, end=request_len), matching
add_and_verify_chunked_request and covering the full prompt while leaving the
existing final-slice and block-group fields unchanged.

---

1962-1995: 📐 Maintainability & Code Quality | 🔵 Trivial

Test coverage summary

Added test functions in tests/unittest/disaggregated/test_kv_transfer.py:

• test_build_prefill_chunk_projects_incremental_source_against_full_prompt
• test_build_prefill_chunk_normalizes_swa_source_to_computed_prefix (2 cases)
• test_send_prefill_chunks_basic (5 cases)
• test_send_prefill_chunks_integrity_check (3 cases)
• test_send_prefill_chunks_multiple_layer_groups
• test_send_prefill_chunks_preserves_mamba_state_index
• test_send_prefill_chunks_none_mamba_state_index
• test_transfer_worker_chunked (V1 and V2)
• test_transfer_worker_pipelined (V1 and V2)
• test_transfer_worker_pipelined_ctx_prefix_reuse (V1 and V2)

Modified test functions:

• test_transfer_with_gen_prefix_offset: new chunk_size_blocks parameter with single_slice and sender_chunked cases.
• test_session_cancel_after_send: now asserts wait_complete() == WaitResult.FAILED instead of an exception.

New helpers: _send_prefill_chunks, _setup_chunked_request, _verify_and_cleanup_chunked, add_and_verify_chunked_request, add_and_verify_pipelined_request.

Test list status: the new tests live in an existing module in tests/unittest/disaggregated/, so they run with the existing unit-test job for that directory. No entry was added under tests/integration/test_lists/test-db/ or tests/integration/test_lists/qa/ for these unit tests, which matches the existing pattern for this module.

Covered paths: chunk projection against a growing source, SWA normalization, chunk counts including empty and oversized chunk sizes, block integrity with a context-side reuse prefix, layer-group suffix handling, mamba_state_index propagation, cancellation state, and end-to-end chunked and pipelined transfer for V1 and V2.

Gaps: no test drives send_prefill_chunk or _maybe_send_prefill_chunk on the executor path; no test covers a mid-transfer failure or timeout during pipelined sending; no multi-rank (TP or PP > 1) pipelined configuration; PIPELINED_TEST_CONFIGS and CHUNKED_TEST_CONFIGS both cover only tp1/pp1.

Verdict: needs follow-up. The unit-level chunk projection coverage is strong, but the executor entry points and multi-rank pipelined paths are untested. Do you want me to open an issue to track the executor-path and multi-rank pipelined tests?

As per path instructions: "Always produce a test coverage summary, even if no issues are found."

🤖 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 `@tests/unittest/disaggregated/test_kv_transfer.py` around lines 1962 - 1995,
Extend test coverage for the missing executor paths and multi-rank pipelined
transfers. Add tests that invoke send_prefill_chunk and
_maybe_send_prefill_chunk through the executor, including failure or timeout
behavior during pipelined sending, and add TP/PP > 1 cases to
PIPELINED_TEST_CONFIGS or the corresponding pipelined test setup while
preserving existing V1/V2 coverage.

Source: Path instructions

🤖 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/transceiver.py`:
- Around line 732-741: The chunk construction in _build_prefill_chunk must
ensure every non-final chunk ends on a tokens-per-block boundary; assert that
non-final chunk_end_pos is divisible by tpb, or clamp chunk_end_block down for
non-final chunks to prevent overlapping destination blocks. In
tensorrt_llm/_torch/disaggregation/base/transfer.py lines 58-74, verify
derive_chunk_block_coords is only called with block-aligned pipelined chunks and
document the producer alignment precondition; no other behavior should change.
- Around line 601-610: Update _close_failed_sessions to tolerate failed session
IDs that have no corresponding entry in reqs, matching the asymmetry handling in
_retire_send_session. Only update request state, retirement flags, and delete
the request when reqs contains rid; always close and remove the associated
session so reconciliation continues without KeyError or leaked sessions.

---

Nitpick comments:
In `@tests/unittest/disaggregated/test_chunked_transfer.py`:
- Around line 657-666: Remove the unused transceiver._chunk_size_blocks = 64
stub assignment from test_pipelined_transfer_disabled_by_default; keep the test
focused on the _enable_pipelined_transfer value read by
KvCacheTransceiverV2.pipeline_transfer_enabled.
- Around line 15-51: Add unit tests in the existing chunked-transfer test module
covering three gaps: verify _close_failed_sessions with mark_retired=True
retires sessions whose rid is absent from _send_reqs; exercise non-block-aligned
chunk ends and assert successive destination writes do not overlap; and drive
multiple chunks through the real respond_and_send_async and TxSession paths
without mocking the session or chunk builder. Follow the existing fixtures and
assertions, then run pytest tests/unittest/.

In `@tests/unittest/disaggregated/test_kv_transfer.py`:
- Around line 1996-2001: Update the receiver-side KVSlice construction in the
pipelined helper to set token_range explicitly to TokenRange(start=0,
end=request_len), matching add_and_verify_chunked_request and covering the full
prompt while leaving the existing final-slice and block-group fields unchanged.
- Around line 1962-1995: Extend test coverage for the missing executor paths and
multi-rank pipelined transfers. Add tests that invoke send_prefill_chunk and
_maybe_send_prefill_chunk through the executor, including failure or timeout
behavior during pipelined sending, and add TP/PP > 1 cases to
PIPELINED_TEST_CONFIGS or the corresponding pipelined test setup while
preserving existing V1/V2 coverage.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: dc203880-13e7-4d84-bb84-3e48776a9ec6

📥 Commits

Reviewing files that changed from the base of the PR and between 5c6d65b and 7c87a01.

📒 Files selected for processing (22)

• tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
• tensorrt_llm/_torch/disaggregation/base/transfer.py
• tensorrt_llm/_torch/disaggregation/native/transfer.py
• tensorrt_llm/_torch/disaggregation/transceiver.py
• tensorrt_llm/_torch/pyexecutor/_util.py
• tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py
• tensorrt_llm/_torch/pyexecutor/llm_request.py
• tensorrt_llm/_torch/pyexecutor/py_executor.py
• tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
• tensorrt_llm/commands/serve.py
• tensorrt_llm/llmapi/disagg_utils.py
• tensorrt_llm/llmapi/llm_args.py
• tensorrt_llm/usage/llm_args_golden_manifest.json
• tests/integration/defs/accuracy/test_disaggregated_serving.py
• tests/integration/test_lists/test-db/l0_dgx_b200.yml
• tests/unittest/_torch/executor/test_disagg_index_mapper_early_release.py
• tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py
• tests/unittest/disaggregated/test_bounce.py
• tests/unittest/disaggregated/test_chunked_transfer.py
• tests/unittest/disaggregated/test_disagg_utils.py
• tests/unittest/disaggregated/test_kv_transfer.py
• tests/unittest/disaggregated/test_transceiver_bounded_polling.py

🚧 Files skipped from review as they are similar to previous changes (17)

• tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py
• tests/unittest/disaggregated/test_transceiver_bounded_polling.py
• tensorrt_llm/usage/llm_args_golden_manifest.json
• tensorrt_llm/_torch/pyexecutor/_util.py
• tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
• tensorrt_llm/_torch/pyexecutor/llm_request.py
• tests/unittest/_torch/executor/test_disagg_index_mapper_early_release.py
• tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
• tests/unittest/disaggregated/test_disagg_utils.py
• tests/integration/defs/accuracy/test_disaggregated_serving.py
• tests/unittest/disaggregated/test_bounce.py
• tensorrt_llm/commands/serve.py
• tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py
• tensorrt_llm/_torch/pyexecutor/py_executor.py
• tensorrt_llm/llmapi/llm_args.py
• tensorrt_llm/_torch/disaggregation/native/transfer.py
• tests/integration/test_lists/test-db/l0_dgx_b200.yml

[athena-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #65352 [ run ] triggered by Bot. Commit: f6ee378 Link to invocation

[tensorrt-cicd]

PR_Github #65352 [ run ] completed with state SUCCESS. Commit: f6ee378
/LLM/main/L0_MergeRequest_PR pipeline #53119 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[athena-nv]

/bot run --disable-fail-fast

[Shixiaowei02]

Thanks for the effort.

A scope note, not a correctness one. About 350 of the ~1000 changed source lines are the feature. Most of the rest is refactoring that would be easier to review, and to revert, on its own. Three things worth doing before the next round.

The chunked-prefill check reads the wrong value. It reads the declared setting. The executor factory later turns chunked context off for star attention and for MLA on unsupported SMs, and that later value is what the scheduler uses. So on those models the check passes, chunking is off, and the pipelined path runs over one degenerate chunk. Moving the check to where the effective value is decided fixes it, and drops the new parameter, the call-site reflow, and the AutoDeploy hunk.

The result frame does not need a new field. The real problem is just that the sender puts its own task index on the wire, where the receiver expects its own. The receiver posts exactly one task, so send the peer's index in the existing field and keep the sender's index local for logging. That avoids widening a binary contract with no version negotiation, and stops an unrelated KV-bounce test from being rewritten.

The new any-inflight helper looks unnecessary. The timeout sweep already runs once per iteration from the scheduling path, so two of the four sites are redundant. The other two are on the pipeline-parallel loop, which this feature rejects on the sending side. It does change flag-off behaviour, and that looks like a latent fix for pipeline parallelism — worth landing with its own test. Note the per-request variant is a different method, and the only thing keeping pages from being freed under an active write. That one should stay.

Three smaller ones: the YAML validation is not the enforcement point — the per-request check covers every deployment, and the YAML one never sees the config that takes effect; keeping the slice's existing extent field would also stop an untouched pre-merge test from failing to build a session; and the lock change is ~90 lines of pure re-indentation, easier to read as its own commit.

[Shixiaowei02]

This file carries more than the feature needs. Could we only make these changes?

@dataclass
class ChunkCoords:
    """Position of one pipelined chunk in a request's global block space."""

    block_offset: int
    block_count: int

    def __post_init__(self):
        if self.block_offset < 0 or self.block_count < 0:
            raise ValueError(f"Invalid chunk: offset={self.block_offset}, count={self.block_count}")

...

class KVSlice:
    ...
    chunk: Optional[ChunkCoords] = None

[athena-nv]

What do you mean by this? Are you saying I should move the project_blocks_to_global_chunk helper to another file? Or do you mean I should reduce the comments?

[athena-nv]

The _has_any_inflight_kv_transfer helper is meant to track transfer state of the request independently of it's prefill compute state, since DISAGG_CONTEXT_TRANS_IN_PROGRESS is not set until the last chunk transfer is scheduled.
I can reduce the usage of this helper though.

[chienchunhung]

The configured timeout is not monitored until the final chunk is submitted, leaving the earlier chunks transfer uncovered, e.g., taking longer or hanging indefinitely.

We should monitor the request for every chunk transfer and define how a stalled NIXL operation is drained or terminated before its KV memory can be reused. For example, add a test that blocks the first chunk and verifies that the request times out safely before the final chunk.

[athena-nv]

UIUC, the last chunk is blocked behind the earlier chunks so the transfer timeout of the last chunk would catch any previous chunks which are hanging.
I will file a ticket for this for now.

[chienchunhung]

I think it is fine to only monitor timeout for the final chunk transfer, which transitively detects earlier chunk transfer hang, under FIFO sequence.

The deeper concern, however, is once a timeout is detected for the last chunk transfer, how do we correctly mark the corresponding transfer (not just the final one but potentially some of the earlier ones) terminal. Today, cancel_request() marks the session CANCELLED but returns false while the earlier task is still TRANSFERRING. On the next status sweep, the cancelled session is retired without rechecking whether that task has physically drained; a subsequent retry can then release the request’s KV ownership.

PS: I am flagging this as a potential transfer lifecycle/ownership risk; addressing this completely may go beyond the scope of this PR.

[chienchunhung]

The --schedule_style override works in the launcher, but fleet and MPI workers reload the original YAML and validate it before applying the override. For example, a YAML file containing context_first and pipelining can be launched with --schedule_style generation_first; the launcher accepts it, but its workers read context_first again and fail during startup. Please pass the resolved schedule style to every worker before parsing and validation, and add a test for this configuration.