[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_py_executor (tensorrt_llm/_torch/pyexecutor/py_executor_creator.py, 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 (ValueError)
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, gated on is_kv_cache_sender = getenv("TRTLLM_DISAGG_ROLE") != "generation" (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 isinstance(kv_cache_manager, MambaHybridCacheManager) or mamba_cache_manager is not None (ValueError)

resolve_cache_transceiver_config (kv_cache_transceiver.py) is the single place that resolves defaults and validates runtime-independent configuration: it collapses transceiver_runtime == "auto" to None, performs the pipelined-transfer auto-selection below, rejects a PYTHON runtime paired with a non-NIXL backend, and rejects kv_cache_bounce_size_mb > 0 together with enable_pipelined_transfer. It is called twice. First, at the top of create_py_executor (py_executor_creator.py), before KvCacheCreator allocates the KV cache pool or selects a manager class — two things downstream read cache_transceiver_config.transceiver_runtime and would disagree with the transceiver actually constructed later if it were still unresolved at that point: get_kv_cache_manager_cls uses it to decide between MixedMambaHybridCacheManager and the C++-compatible hybrid managers for Mamba/hybrid models, and _maybe_enable_fabric_memory_for_python_transceiver uses it to decide whether to default TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY=1 for the C++ KV cache pool (needed for MNNVL transfers out of that pool by the Python transceiver) — and that env var is read once and cached by C++ on first access, so it must be set before any pool allocation. Second, at the top of create_kv_cache_transceiver itself, which is the only call resolve_cache_transceiver_config gets on paths that don't go through create_py_executor (e.g. AutoDeploy's ad_executor.py). The second call is a no-op in the common case: once transceiver_runtime has been resolved away from "auto"/None, the auto-selection branch and the CPP-conflict branch it could otherwise re-trigger no longer match, so nothing double-logs or re-raises.

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 in create_kv_cache_transceiver rejects enable_pipelined_transfer whenever isinstance(kv_cache_manager, MambaHybridCacheManager) or a separate mamba_cache_manager was passed in. MambaHybridCacheManager is the shared base class for all three hybrid managers — CppMambaHybridCacheManager, MambaHybridCacheManagerV2, and MixedMambaHybridCacheManager — so the isinstance check alone covers all of them; the mamba_cache_manager is not None half is belt-and-suspenders for callers that pass the Mamba pool as a separate object. 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.
• token_range changes meaning. It had exactly one consumer: the sender read token_range.end to recover the request's block span ceil(end / tpb), the anchor for all the suffix arithmetic. .start was never read, since _create_kv_slice always produced [0, prompt_len). It now carries one chunk's window instead, and its presence is what makes a slice a chunk — a monolithic transfer leaves it None and 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.
• The two jobs the old .end did are now answered separately, because a non-final chunk's end is neither of them. The request's block span, which the destination projection needs, is ceil(prompt_len / tpb), rederived by the sender from the session's prompt_len. SWA's stale-block boundary, which also needs prompt_len, reads the same field.

An intermediate revision expressed all of this in block coordinates — a ChunkCoords(block_offset, block_count) field for the window plus a total_blocks field for the span. Both were dropped. total_blocks was a function of prompt_len alone, so carrying it let the two sides disagree about a value neither of them decides; and the window is a genuine range over the request's tokens, so TokenRange names it without inventing a second coordinate system for slices to speak. The window is still decided in block space (see "Deriving the chunk"), which is why both of its bounds are multiples of tokens_per_block and the sender asserts that before dividing them back out.

One invariant carries more weight now that the span is rederived rather than transmitted: _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 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:

        if self.kv_cache_transceiver:
            # A pending cancel that could not complete (a chunk was mid-write)
            # leaves the request active and still being prefilled. Its session
            # is already CANCELLED, so further chunks would only produce
            # spurious FAILED results on the receiver.
            cancel_pending_ids = set(self.canceled_req_ids)
            for req in scheduled_requests:
                if req.is_context_only_request and not req.is_finished_due_to_cancellation:
                    if self.kv_cache_transceiver.has_retired_send_session(req):
                        # The peer registration went away with the session, so
                        # no further slice can land. Checked before the branch
                        # below because start_transfer() would otherwise pin
                        # blocks that only end_transfer() can release.
                        req.state = LlmRequestState.DISAGG_TRANS_ERROR
                        continue
                    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)

                        # send KV slice for monolithic transfer or last chunk of pipelined transfer
                        self.kv_cache_transceiver.respond_and_send_async(req)

                        if self.kv_cache_transceiver.kv_transfer_timeout_ms is not None:
                            req.py_kv_transfer_start_time = time.monotonic()
                    elif (self.kv_cache_transceiver.pipeline_transfer_enabled
                          and req.state != LlmRequestState.GENERATION_COMPLETE
                          and
                          (req.py_request_id if not req.is_child else
                           req.parent_request_id) not in cancel_pending_ids):
                        # send intermediate chunk for pipelined transfer.
                        # GENERATION_COMPLETE means an error path already failed
                        # and freed this request; _update_request_states skips
                        # those, so its chunk bounds are unset.
                        self.kv_cache_transceiver.respond_and_send_async(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 (returning early if that returns None, i.e. the session was already retired), 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. It also returns early when _build_prefill_chunk returns None, which means the scheduler's chunk completed no whole block and there is nothing to send yet.
• 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.
• The retired-session check runs first and skips both branches. has_retired_send_session(req) is checked before the is_context_finished/is_finished_due_to_length split, so a request whose send session was already torn down (see "Retired send sessions and DISAGG_TRANS_ERROR" below) never reaches either respond_and_send_async call — it is failed and continues past both.
• The two branches are not gated the same way. cancel_pending_ids (a snapshot of self.canceled_req_ids taken once per call) only guards the elif — the intermediate-chunk send. The final-chunk if has no cancel_pending_ids check at all; it fires whenever the forward pass says the request's context is done, cancellation-pending or not (a fully-cancelled request is already filtered out by the outer not req.is_finished_due_to_cancellation). Once a request reaches its last chunk, start_transfer/respond_and_send_async finalize the transfer unconditionally rather than leaving it half-sent; only an earlier, still-in-flight chunk gets suppressed by a pending cancel.

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
        is_last_chunk = req.context_remaining_length == 0
        # ... both bounds round down so consecutive chunks tile block space exactly;
        # the last chunk rounds up, since nothing follows to carry its tail ...
        chunk_start_block = 0 if is_first_chunk else chunk_start_pos // tpb
        chunk_end_block = (
            (chunk_end_pos + tpb - 1) // tpb if is_last_chunk else chunk_end_pos // tpb
        )
        total_blocks = (req.prompt_len + tpb - 1) // tpb

        chunk_start = min(chunk_start_block, total_blocks)
        chunk_end = min(chunk_end_block, total_blocks)
        # ... the next chunk's start rounds down to here, so the skipped blocks
        # travel with it; the last chunk always sends ...
        if chunk_end <= chunk_start and not is_last_chunk:
            return None
        chunk_block_count = max(0, chunk_end - chunk_start)
        # ... V1 reserves the full prompt up front, V2 grows incrementally; normalize
        # both to chunk_end before projecting ...
        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 the sender rederives from the session's prompt_len. 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 window travels to the sender worker as a block-aligned TokenRange. Deciding it once 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. The sender reads that answer rather than re-deriving it; the only thing it pays for expressing the window in token space is one alignment assert on arrival before it divides the bounds back out.

Unaligned chunk boundaries round down

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   = floor(chunk_end_pos / tpb)     (ceil on the last chunk)

Both bounds floor, so consecutive chunks tile block space exactly. A block holding an unaligned boundary is not sent by the earlier chunk, whose tokens stop partway into it, and is sent whole by the next one, whose start floors into the same block. Every block is transferred exactly once, and only after the forward pass has filled it.

flowchart LR
    subgraph chunkA [Chunk A, ends at token 259]
        a1["blocks 0..7<br/>fully computed, sent"]
    end
    subgraph chunkB [Chunk B, starts at token 259]
        b1["block 8<br/>partial in A, completed here"]
        b2["blocks 9 onward"]
    end
    a1 -->|"no overlap"| b1

Loading

Three properties make this work:

• The last chunk still rounds up. Flooring it would strand the prompt's final partial block with no chunk following to carry it. Rounding up over-reads nothing there: the last chunk's tail is computed by definition, and ceil(prompt_len / tpb) is the clamp itself.
• A chunk that finishes no block sends nothing. Floor/floor collapses to an empty window whenever the scheduler cuts a chunk that starts and ends inside one block. _build_prefill_chunk returns None and respond_and_send_async returns without touching the session; the next chunk's start floors to the same place, so the deferred blocks travel with it. The final chunk cannot hit this — its end is ceil(prompt_len / tpb), always past its own start — so is_last_slice always reaches the receiver and no request can wait on a chunk that was skipped.
• The floor never over-reads. floor(chunk_end_pos / tpb) is at most the index just past the last block the forward pass filled, so resident_block_end = chunk_end stays within the source list under both allocation models.

Rounding up was the original rule, and it worked, but it put the boundary block in both chunks: the earlier chunk sent it with a stale tail past the boundary, and the next chunk rewrote it whole. That left correctness resting on write ordering — Sender._enqueue routing a peer's writes to one worker thread, which delivers them serially, so the computed copy always lands last — and briefly left a partially valid block in the destination. Flooring removes both, at the cost of deferring at most one block per boundary by one chunk. Ordering still matters for is_last_slice (see "Slice ordering"); it no longer matters for block contents.

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 a token_range 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.token_range 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.
• total_blocks = ceil(task._prompt_len / tpb) is rederived here, since it is a property of the request rather than of the slice. suffix_end_blocks is token_range.end // tpb when chunked and total_blocks otherwise; both bounds are asserted block-aligned before the division, because the producer decided them in block space and a partial block arriving on the wire would mean a bug upstream. 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: the result frame's slice id

The KV_AGENT_RESULT frame carries a single slice id, and the receiver uses it purely as an index into its own task list — so the value has to be the receiver's task index. The sender filled it with task.slice_id, its own task index, which was the same number only as long as a session had one task. With N chunks, chunk i would have addressed the receiver's task i, and the receiver only ever posts one. The field was also named slice_id with no side attached, and the parameter the receiver unpacked it into was literally called sender_slice_id, which is what let the confusion survive review.

The fix is one line at the point the frame is built:

sequenceDiagram
    participant TxSession
    participant SenderWorker
    participant RxSession
    TxSession->>SenderWorker: KVSendTask.slice_id = chunk index 0..N-1
    Note over SenderWorker: WriteMeta.slice_id = req_info.slice_id or 0<br/>the sender's chunk stays reachable via WriteMeta.task
    SenderWorker->>RxSession: KV_AGENT_RESULT carries the receiver's slice id, is_last, status
    Note over RxSession: index _kv_tasks[slice_id]

Loading

• WriteMeta keeps its single pre-existing slice_id: Optional[int] field; only its source changed. _build_kv_write_meta sets slice_id=req_info.slice_id if req_info.slice_id is not None else 0 — the peer's task index, echoed back from RecvReqInfo — instead of task.slice_id. An intermediate revision split the field into sender_slice_id/receiver_slice_id so a chunk-level RDMA failure could still name the sender's chunk in the sender's own logs. That was reverted: _deliver_kv_to_agent resolves the send task from write_meta.task, which the struct already carries, instead of indexing the session by chunk, so multi-chunk pipelining is correct with one field. The cost is that the sender's per-slice log lines no longer identify which chunk they belong to.
• _KV_RESULT_PREFIX (struct.Struct("<qqq?Bq")) is unchanged in shape and content — instance_rank, unique_rid, receiver_slice_id, is_last, status, transfer_size — it already carried transfer_size before this change. The only edit is renaming the third field from sender_slice_id to receiver_slice_id in _make_kv_result_msg, _send_failed_result_to_receiver, and _process_kv_agent_result, so the name matches what the value has always had to mean. This is a ctx/gen wire format with no version negotiation — the receiver unpacks whatever arrives against its compiled-in struct — which is why the fix stays inside the existing field rather than widening the frame.
• RxSession.process_kv_agent_result(peer_rank, receiver_slice_id, is_last_slice, status, ...) indexes _kv_tasks[receiver_slice_id] and asserts it is in range first, so a sender/receiver slice-count mismatch fails loudly instead of resolving to the wrong task. The FAILED-detail log points at the sender's own log ("reported by remote agent; see sender-side log for nixl_status") rather than trying to name the sender's chunk locally.

For a single-chunk session the old and new values coincide, which is why the monolithic path never saw this and why the change is invisible to a monolithic receiver: receiver_slice_id is 0 in every current deployment.

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) is non-abstract and defaults to False, which leaves BindKvCacheTransceiver untouched — the C++ transceiver has no pipelining, so state and transfer activity coincide there. KvCacheTransceiverV2 implements it 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 into cancel_pending_ids once per call, but the two branches are gated asymmetrically, not by one shared continue. The elif (intermediate-chunk) branch checks req.py_request_id if not req.is_child else req.parent_request_id) not in cancel_pending_ids: a session whose cancel is pending would otherwise be fed another chunk and produce a spurious FAILED result to the receiver. The if (final-chunk) branch — is_context_finished or is_finished_due_to_length — has no cancel_pending_ids check; once a request reaches its last chunk, start_transfer and respond_and_send_async run unconditionally, because a fully-cancelled request is already excluded upstream by the outer not req.is_finished_due_to_cancellation, and letting the last chunk finalize normally is simpler than adding a new half-sent state. What both branches share is the has_retired_send_session(req) check ahead of the if/elif split (see below), which is the one place a single continue really does skip both.
• Timeout detection is left exactly as the monolithic path had it. _check_kv_transfer_timeout scans requests_in_transfer(), which a context request enters on its final chunk, and respond_and_send_async overwrites py_kv_transfer_start_time on every chunk, so the clock it measures is the last chunk's. Intermediate chunks are deliberately not swept: py_kv_transfer_timed_out can only be set for a request the transfer manager already knows about, so _check_disagg_ctx_cache_transfer_status has nothing to act on while prefill is still running.

Retired send sessions and DISAGG_TRANS_ERROR

A send session can be torn down before its last chunk goes out — cancellation, an RDMA failure mid-chunk, or a peer-registration timeout. TxSession.close() also calls Sender.clear_session, which drops the peer's RecvReqInfo; the generation server never re-sends it, so a session re-created after that point has no peer to write to and every task it creates would sit in INIT forever. The fix bars the re-creation instead of allowing it to silently hang.

LlmRequest gained py_kv_send_session_retired (False by default), set whenever a send session is torn down: Sender._retire_send_session (used by the cancel and completion paths in check_context_transfer_status and by cancel_request) and Sender._close_failed_sessions(..., mark_retired=True) (the failure path) both set it. _get_or_create_send_session checks it before creating a session for a rid it doesn't already have: if set, it logs a warning, sets req.state = LlmRequestState.DISAGG_TRANS_ERROR, and returns None instead of creating a session with no peer.

KvCacheTransceiverV2.has_retired_send_session(req) (transceiver.py) reports req.py_kv_send_session_retired and get_unique_rid(req) not in self._send_sessions — retired and not already re-registered, so a request whose session completed and was retired normally (every successful transfer retires its session once KV_TRANSFERRED) does not spuriously read as failed once it moves on. KvCacheTransceiver.has_retired_send_session defaults to False on the base class, so BindKvCacheTransceiver is unaffected.

_send_kv_async checks self.kv_cache_transceiver.has_retired_send_session(req) first, ahead of both the final-chunk and intermediate-chunk branches: if true, it sets req.state = LlmRequestState.DISAGG_TRANS_ERROR and continues past both, explicitly before start_transfer could otherwise pin blocks that only end_transfer can release. LlmRequestState.DISAGG_TRANS_ERROR itself is not new — the C++ core already defines it, py_executor.py already had several pre-existing consumers of it (e.g. _end_transfer_and_maybe_terminate, _get_disagg_reqs_in_error_state, the generation-side error-state sweeps), and transceiver.py's _close_failed_sessions already set it on the pre-pipelining RDMA-failure path. This PR adds two new producers that both flow from session retirement: _get_or_create_send_session (fails the request instead of re-creating a session with no peer) and the has_retired_send_session check in _send_kv_async. Both reuse the existing state and existing downstream handling rather than inventing a new one, so a request whose peer registration is gone is failed and drained through the same machinery as any other disaggregated transfer error, instead of stalling until a timeout.

Change inventory by file

Non-test files touched by this branch, for evaluating which changes are load-bearing vs. removable noise. Sizes are cumulative insertions/deletions against the merge base (git diff <base> --numstat). "Core" = the feature does not work without it. "Supporting infra" = validation/config plumbing/error handling that makes the feature robust but isn't the transfer mechanism itself. "Adjacent" = could plausibly be split into its own PR.

File	What changed	Why / purpose	Size (+/-)	Category
tensorrt_llm/_torch/disaggregation/base/transfer.py	Adds project_blocks_to_global_chunk() (range-intersection projector), repurposes KVSlice.token_range from "the whole request" to "this chunk's block-aligned window" (and relaxes TokenRange to allow an empty range), and makes SessionArgsBase.prompt_len: int required (was Optional[int]).	Defines the shared data model a chunk is expressed in — token_range marks a slice as a chunk vs. whole-request and carries its window, prompt_len on the session is where the request's own extent now lives, and project_blocks_to_global_chunk is the primitive both sender and receiver use to map a global chunk window onto a resident-suffix block list (needed for SWA/prefix-reuse/incremental allocation).	+72/-16	core
tensorrt_llm/_torch/disaggregation/native/transfer.py	Sources WriteMeta.slice_id from req_info.slice_id (the peer's task index) instead of the sender's chunk index and renames the wire field to receiver_slice_id in _KV_RESULT_PREFIX/_make_kv_result_msg/process_kv_agent_result, adds _send_kv_result_to_receiver() helper, adds chunk-aware projection in _build_kv_write_meta (rederives total_blocks from prompt_len, takes suffix_end_blocks from task._slice.token_range), makes TxSession/RxSession/KVSendTask take required prompt_len, adds TxSession.dispatch_lock around send()/send_aux()/_respond_with_kv replay, and extends TxSession.status to aggregate ERROR/TRANSFERRING over all kv_tasks.	Implements the actual per-chunk transfer mechanics: correct source/destination block addressing for a chunk, a result frame that addresses the receiver's single task rather than the sender's chunk, and ordering guarantees (dispatch_lock) so a late chunk can't race an in-flight one and violate is_last_slice semantics.	+213/-120	core
tensorrt_llm/_torch/disaggregation/transceiver.py	Adds _build_prefill_chunk() (derives the block-aligned TokenRange from req.py_last_context_chunk/prepopulated_prompt_len/tokens_per_block, returning None for a chunk that completes no block), extends _create_kv_slice(req, resident_block_end=...), adds pipeline_transfer_enabled property, has_inflight_transfer(), has_retired_send_session(), _retire_send_session(), rewrites _close_failed_sessions to take mark_retired, and updates respond_and_send_async to skip an empty chunk, send intermediate chunks, or finalize only on is_last_slice.	This is where scheduler token-chunk bounds are turned into a block-aligned window, where session/request lifecycle is reconciled with pipelining (retired-session guard, DISAGG_TRANS_ERROR on stale re-creation), and where respond_and_send_async gains the branch that ships an intermediate vs. final chunk.	+204/-44	core
tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py	Adds resolve_cache_transceiver_config() (collapses "auto", auto-selects PYTHON runtime when enable_pipelined_transfer is set, rejects CPP+pipelined, rejects non-NIXL backend, rejects kv_cache_bounce_size_mb>0+pipelined), calls it at the top of create_kv_cache_transceiver, adds Mamba/hybrid-cache and pipeline_parallel_size==1 rejection checks, and adds non-abstract pipeline_transfer_enabled/has_inflight_transfer/has_retired_send_session defaults to the KvCacheTransceiver base class (so BindKvCacheTransceiver is unaffected).	Centralizes all runtime-independent config validation/auto-selection for the feature (NIXL-only, Python-transceiver-only, no bounce buffer, no PP, no Mamba) so invalid combinations fail fast at startup rather than mid-run.	+101/-9	supporting infra
tensorrt_llm/_torch/pyexecutor/llm_request.py	Adds LlmRequest.py_kv_send_session_retired = False field.	Backing flag consumed by _get_or_create_send_session/has_retired_send_session to detect and fail requests whose send session was torn down before its last chunk instead of silently re-creating a session with no peer.	+4/-0	core
tensorrt_llm/_torch/pyexecutor/py_executor.py	Adds _validate_request checks (beam_width==1, schedule_style==generation_first when pipelining), rewrites the _send_kv_async context-request loop to check has_retired_send_session first, then branch on final-chunk (is_context_finished/is_finished_due_to_length) vs. intermediate-chunk (pipeline_transfer_enabled + not cancel-pending) sends, and extends _is_request_in_transmission to also consult kv_cache_transceiver.has_inflight_transfer(request).	Wires pipelined sends into the executor's per-forward-step loop and fixes the two correctness gaps pipelining introduces: cancellation of a mid-prefill request (transfer ownership no longer tracked purely by LlmRequestState) and per-request validation of the two hard preconditions.	+67/-20	core
tensorrt_llm/_torch/pyexecutor/py_executor_creator.py	Calls resolve_cache_transceiver_config(cache_transceiver_config) early in create_py_executor (before KV cache pool/manager selection), and raises ValueError if enable_pipelined_transfer is set without enable_chunked_prefill.	Ensures config resolution happens before downstream consumers (get_kv_cache_manager_cls, fabric-memory env-var default) read transceiver_runtime, and enforces the chunked-prefill precondition (chunk boundaries come from the chunked-prefill scheduler).	+11/-0	supporting infra
tensorrt_llm/commands/serve.py	Changes disaggregated() to call parse_disagg_config_file(config_file, schedule_style_override=schedule_style) instead of parsing then overwriting disagg_cfg.schedule_style after the fact.	Routes the --schedule_style CLI flag through the same validation path (_validate_disagg_config) as the YAML-declared value, so a mismatched schedule_style with enable_pipelined_transfer=True is caught at startup regardless of whether it came from the file or the CLI.	+2/-3	supporting infra
tensorrt_llm/llmapi/disagg_utils.py	Adds _validate_disagg_config() (checks schedule_style is valid and rejects enable_pipelined_transfer=True with a non-generation_first schedule_style), calls it from extract_disagg_cfg, adds schedule_style_override param to parse_disagg_config_file, and adds an empty-YAML guard.	Fail-fast startup validation of the schedule_style: generation_first requirement — the pipelining precondition that the generation server must have registered destination blocks before prefill compute finishes.	+44/-3	supporting infra
tensorrt_llm/llmapi/llm_args.py	Adds CacheTransceiverConfig.enable_pipelined_transfer: bool = False field, with a comment in _to_pybind() noting it has no C++ counterpart and isn't forwarded.	The user-facing config flag that gates the entire feature.	+8/-0	core
tensorrt_llm/usage/llm_args_golden_manifest.json	Adds the cache_transceiver_config.enable_pipelined_transfer entry (bool, no allowed values/converter).	Mechanical: required companion regeneration (generate_llm_args_golden_manifest.py) whenever an LlmArgs-nested field is added — keeps the golden manifest in sync with llm_args.py.	+7/-0	supporting infra

None of the above looks like unrelated cleanup that could plausibly split into its own PR — even the serve.py and disagg_utils.py changes, which read like general validation/robustness plumbing, exist specifically to enforce the schedule_style=generation_first precondition this feature requires, and commands/serve.py's one-line refactor exists only so the CLI flag goes through that new validation path. The closest thing to a separable line item is the result-frame rename (sender_slice_id → receiver_slice_id) inside native/transfer.py, which is a pure rename with no behavior change; it is kept here because it documents the value the accompanying one-line fix now puts in that field, and splitting the rename from the fix would leave one of the two halves looking arbitrary.

[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 and concisely describes the main change: pipelined KV cache transfer for disaggregated serving in the Python cache transceiver.
Description check	✅ Passed	The description thoroughly explains the feature, requirements, design, configuration, and implementation, but it omits explicit Test Coverage and PR Checklist sections.

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

[Shixiaowei02]

Besides, they are some comments based on agent suggestions.

[Shixiaowei02]

Thanks for turning all of this around so quickly. Most of the last round is properly addressed. Here are some comments based on the agent review workflow.

[Shixiaowei02]

This function is unchanged from main, and on main there is no reason to split it: it is called once, when prefill ends, so taking ownership of the request for the fabric and doing the end of context work happen at the same moment.

Pipelining separates those two moments, and the way the PR covers the gap is by adding a second ownership record on the transceiver and a flag on the request. Those two are new here, so I think they need resolving in this PR rather than later, one way or the other.

The clean version is to split this into an acquire that runs on the first chunk and a completion that runs on the last. Then the places that ask whether a request is still owned by the fabric all become correct on their own, the extra record and the flag can go, and the timeout clock gets a natural place to start. I realise that touches the connector path too, so if you would rather keep the extra records, could you say in a comment why they are the right call and what keeps them in step with the manager? Right now there are three answers to the same question and they do not all agree.