[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: Benchmarking performance impact

Description

Summary

Implements pipelined prefill-transfer for disaggregated serving. Instead of waiting for all prefill chunks to complete before starting KV cache transfer, each chunk's KV data is transferred to the generation server immediately after its prefill completes. This overlaps GPU compute with RDMA transfer, hiding transfer latency behind prefill computation. When transfer time per chunk is less than prefill time per chunk (typical for 100+ Gbps NIC with long-context workloads), transfer latency is nearly fully hidden.

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.
• token_range for a pipelined chunk is block-aligned and encodes the chunk's global position. For a monolithic transfer it stays [0, prompt_len + num_extra_kv_tokens) and is not required to be block-aligned.

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 token_range.end was an acceptable stand-in for prompt_len. For a non-final chunk it is not — token_range.end is the chunk boundary, and 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, and the resulting token_range is block-aligned by construction. 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.

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 token_range.start == 0 with is_last_slice, so is_chunked evaluates false and the write takes the monolithic path — byte-for-byte a non-pipelined transfer.

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:

• derive_chunk_block_coords(token_range, tpb) recovers (chunk_offset, chunk_block_count) from the block-aligned token_range, raising ValueError if it is not block-aligned.
• The is_chunked predicate — not is_last_slice, or token_range.start > 0 — decides whether to derive chunk coordinates at all. A monolithic transfer's token_range ends at prompt_len + num_extra_kv_tokens and need not be block-aligned, so it must not go through derive_chunk_block_coords.
• suffix_end_blocks is chunk_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.

Dev Engineer Review

• Adds pipelined KV cache transfer for the Python NIXL transceiver.
• Sends each completed prefill chunk while GPU computation continues.
• Adds chunk projection, multi-slice sessions, slice-ID correlation, transfer-state tracking, cancellation, timeout handling, and retired-session handling.
• Validates chunked prefill, generation_first, beam width 1, Python NIXL, zero bounce buffer, sender PP size 1, and unsupported Mamba or hybrid cache configurations.
• Updates configuration parsing, AutoDeploy integration, and the KV-transfer wire protocol.
• Follow-up remains required for performance benchmarks, stress testing, multi-node coverage, and targeted tests for send_prefill_chunk, _maybe_send_prefill_chunk, create_kv_cache_transceiver, and request_and_receive_async.

QA Engineer Review

• Adds unit coverage for chunk projection, slice-ID routing, replay ordering, session failures, configuration validation, prefix reuse, cancellation, scheduling, timeout handling, and retired sessions.
• Adds end-to-end coverage through test_transfer_worker_chunked, test_transfer_worker_pipelined, and test_transfer_worker_pipelined_ctx_prefix_reuse.
• Adds pipelined accuracy coverage for Llama 3.1 8B and Gemma 3 1B.
• The provided repository diff is empty, so current test-list coverage cannot be independently verified.
• Verdict: needs follow-up because CBTS coverage data and current test-list entries are unavailable. Performance, stress, long-running stability, and multi-node validation also remain uncovered.

[coderabbitai]

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#17361: Modifies sender-side KV block mapping for sliding-window receivers.
• NVIDIA/TensorRT-LLM#17107: Changes executor transfer-progress polling logic.
• NVIDIA/TensorRT-LLM#17137: Changes KV-transfer timeout and session handling.

Suggested reviewers: bo-nv, bowenfu, brnguyen2

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 62.36% 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
Title check	✅ Passed	The title clearly identifies the feature: pipelined KV cache transfer for disaggregated serving through the Python cache transceiver.
Description check	✅ Passed	The description clearly explains the feature, configuration requirements, architecture, implementation, limitations, and current benchmarking status.
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.

✨ Finishing Touches 💡 1

⚔️ Resolve merge conflicts 💡

• Resolve merge conflict in branch trtllm-12499-pipelined-kvcache-transfer

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

[nv-xtf]

A post-rebase ordering question: main's KvCacheCreator._maybe_enable_fabric_memory_for_python_transceiver() (#16832) runs before KV-pool allocation and enables the fabric-memory default only when the config says transceiver_runtime == "PYTHON". The default "auto" (e.g. Llama) resolves to None, and the auto-selection here only flips the local use_python without persisting it — after the pool is already allocated. So enable_pipelined_transfer + default runtime would get a V1 C++ pool without fabric memory, affecting MNNVL transfers. The integration test sets "PYTHON" explicitly, so this path isn't covered.
Could "pipelined ⇒ PYTHON" be resolved and persisted before pool routing — via_resolve_transceiver_runtime_auto on the normal path, with a fallback or fail-fast for paths like AutoDeploy that skip it?

[fredricz-20070104]

Review summary - CONCERNS

Verdict: Cannot merge as-is — the branch is in conflict (mergeable_state=dirty) and has two outstanding CHANGES_REQUESTED on the head commit while the PR is still WIP (benchmarking). The feature logic and tests are largely solid, but there is one correctness concern worth confirming before merge.

Concerns

1. [MAJOR] native/transfer.py (~_deliver_kv_to_agent, lines 488-528) - cuda_event sync happens while task is still INIT

   • What is wrong: The worker waits on task._slice.cuda_event.synchronize() before the session.lock-held INIT→TRANSFERRING transition. (The sync line is just outside this diff, so please confirm the current ordering — coderabbitai flagged the same at 488-507.)
   • How it fails: While the worker is blocked in synchronize(), the task is still INIT. A concurrent cancel_request() iterates kv_tasks, sees no TRANSFERRING task, and frees/reuses the request's KV pages. When the event completes the worker issues the RDMA into freed pages → use-after-free / corrupted KV on the wire.
   • Suggested fix: Perform the INIT→TRANSFERRING transition (under session.lock, with the CANCELLED/ERROR abort check) before cuda_event.synchronize(), and return early on the abort path.

2. [MAJOR] Merge blocker — dirty branch + change requests. mergeable=false, mergeable_state=dirty. Rebase onto main, resolve conflicts, and address the CHANGES_REQUESTED from Shixiaowei02 and nv-xtf on commit 4c6c035. PR description still reads 'Status: Benchmarking performance impact'.

Minor notes (non-blocking)

• tests/unittest/disaggregated/test_kv_transfer.py - test_transfer_worker_chunked is missing the if not torch.cuda.is_available(): pytest.skip(...) guard that test_transfer_worker_pipelined has; it will error rather than skip on a CUDA-less runner.
• _build_prefill_chunk assumes chunk boundaries are KV-block-aligned (floor start / ceil end). This is guaranteed by chunked-prefill's block-multiple invariant and validated later by derive_chunk_block_coords, but consider an explicit assert in _build_prefill_chunk to fail fast if a non-aligned boundary ever reaches it.

QA view

• Test coverage: adequate — new unit tests (test_chunked_transfer.py, additions to test_kv_transfer.py) cover projection math, TxSession/RxSession state machines, chunked & pipelined send paths, and the requirement validators; a new integration accuracy test is added. Uncovered: the concurrent cancel-during-cuda_event-sync race in finding #1.
• SM coverage: architecture-independent (pure Python/RDMA transfer logic, no get_sm_version/arch guards). Integration test runs Hopper+ (skip_pre_hopper) and is listed on B200.
• Test code: missing CUDA skip guard on test_transfer_worker_chunked; otherwise clean (context-managed server launch, explicit cleanup).
• Test time: significant — four new disaggregated GSM8K accuracy parametrizations on l0_dgx_b200, each launching a ctx+gen server pair.
• Needs /qa-verify: yes — WIP multi-GPU feature, dirty branch, and the cancel/cuda_event race has no reproducing test; re-run the disagg accuracy suite after conflicts are resolved.

Possible new issues

• Non-block-aligned chunk boundaries could overlap/transfer partial blocks (mitigated by chunked-prefill's block-multiple invariant).
• prompt_len changed from Optional=None to required across SessionArgsBase/KVSendTask/TxSession/RxSession; external callers not shown may break.
• _build_kv_write_meta now uses task._prompt_len for slice_end instead of token_range.end; recheck the speculative-decoding extra-draft-block path for a one-block start-offset shift.

What I could not verify

• The exact placement of cuda_event.synchronize() relative to the lock transition (line lies just outside the diff).
• All external callers of the constructors whose prompt_len became required.
• Runtime accuracy/perf of the pipelined path (integration test not executed here).

Automated review by NVCortex Lite, run by @fredricz-20070104.

[athena-nv]

@coderabbitai resume

[coderabbitai]

✅ Action performed

Reviews resumed.

[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: 5

🧹 Nitpick comments (3)

tensorrt_llm/llmapi/disagg_utils.py (1)

78-79: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Annotate the public parser return type.

Add -> DisaggServerConfig to parse_disagg_config_file. _validate_disagg_config is private, so a docstring is not required by the repository guidelines.

🤖 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/llmapi/disagg_utils.py` around lines 78 - 79, Update the public
parse_disagg_config_file function signature to explicitly return
DisaggServerConfig. Leave the private _validate_disagg_config function
unchanged, including its lack of a docstring.

Source: Coding guidelines

tensorrt_llm/_torch/disaggregation/native/transfer.py (1)

206-215: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Annotate the rewritten _make_kv_result_msg signature.

The signature was expanded in this change but stays unannotated. Add parameter and return types so the new sender/receiver ID contract is explicit at the call sites.

♻️ Proposed annotations

 def _make_kv_result_msg(
-    instance_rank,
-    unique_rid,
-    sender_slice_id,
-    receiver_slice_id,
-    is_last_slice,
-    agent_result,
-    transfer_size=0,
-    tail=None,
-):
+    instance_rank: int,
+    unique_rid: int,
+    sender_slice_id: int,
+    receiver_slice_id: int,
+    is_last_slice: bool,
+    agent_result: AgentResult,
+    transfer_size: int = 0,
+    tail: Optional[list[bytes]] = None,
+) -> list:

Based on coding guidelines: "Annotate every function, use None for procedures".

🤖 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 206 -
215, Annotate the `_make_kv_result_msg` parameters and return value, including
the sender and receiver slice ID types and the optional `tail` parameter. Use
`None` as the return annotation if the function is a procedure, and follow the
surrounding module’s established type aliases or conventions.

Source: Coding guidelines

tensorrt_llm/_torch/disaggregation/transceiver.py (1)

571-597: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Route send-side failure teardown through _retire_send_session.

_close_failed_sessions now duplicates the teardown that _retire_send_session performs, and it only does so when mark_retired=True. The two paths also differ in robustness: _close_failed_sessions uses del reqs[rid] and reqs[rid].state, which raise KeyError if _send_reqs has no entry, while _retire_send_session uses .pop(). Today _send_reqs[rid] is always populated by _build_prefill_chunk before any failure can be collected, so this is a maintainability concern rather than a live bug.

Consider keeping _close_failed_sessions for the receive path and calling _retire_send_session for send-side failures, so one function owns sender retirement.

🤖 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 571 - 597,
The send-side failure path should use _retire_send_session as the single owner
of sender teardown. Update callers handling send failures to invoke
_retire_send_session for each failed request, while retaining
_close_failed_sessions for receive-side cleanup; remove any duplicated
send-session close, removal, or retirement handling from that path.

🤖 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 2014-2021: Update the assertion in the lock block before indexing
self._kv_tasks to require receiver_slice_id >= 0 as well as being less than
len(self._kv_tasks). Preserve the existing mismatch error context, then use the
validated receiver_slice_id to select the task.

In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 3000-3007: Update the call to create_kv_cache_transceiver in
create_py_executor to pass the effective chunked-prefill state (ctx_chunk_config
is not None or the resolved enable_chunked_context) instead of
llm_args.enable_chunked_prefill, ensuring the factory reflects the scheduler’s
actual configuration.

In `@tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py`:
- Around line 137-142: Update the warning in the enable_pipelined_transfer
handling to state that users must explicitly set transceiver_runtime="PYTHON" or
disable enable_pipelined_transfer; remove the incorrect guidance that
transceiver_runtime="CPP" disables pipelined transfer. Keep the automatic
cache_transceiver_config.transceiver_runtime assignment unchanged.

In `@tests/integration/defs/accuracy/test_disaggregated_serving.py`:
- Around line 1593-1595: Add the skip_pre_hopper decorator to the new test’s
decorator list alongside skip_less_device and the existing parameterization
decorators, matching the guards on TestGemma3_1BInstruct::test_auto_dtype and
test_kv_cache_v2_nixl_python.

In `@tests/integration/test_lists/test-db/l0_dgx_b200.yml`:
- Around line 21-26: Reduce the pre_merge entries under the “Disaggregated
Serving: Pipelined KV Transfer (multi-GPU)” section in l0_dgx_b200.yml by
retaining only representative Llama and Gemma combinations; move the
lower-priority node IDs to the post_merge tier instead of adding all five to
pre_merge.

---

Nitpick comments:
In `@tensorrt_llm/_torch/disaggregation/native/transfer.py`:
- Around line 206-215: Annotate the `_make_kv_result_msg` parameters and return
value, including the sender and receiver slice ID types and the optional `tail`
parameter. Use `None` as the return annotation if the function is a procedure,
and follow the surrounding module’s established type aliases or conventions.

In `@tensorrt_llm/_torch/disaggregation/transceiver.py`:
- Around line 571-597: The send-side failure path should use
_retire_send_session as the single owner of sender teardown. Update callers
handling send failures to invoke _retire_send_session for each failed request,
while retaining _close_failed_sessions for receive-side cleanup; remove any
duplicated send-session close, removal, or retirement handling from that path.

In `@tensorrt_llm/llmapi/disagg_utils.py`:
- Around line 78-79: Update the public parse_disagg_config_file function
signature to explicitly return DisaggServerConfig. Leave the private
_validate_disagg_config function unchanged, including its lack of a docstring.

🪄 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: 9b760343-8b54-47b9-80fe-44654fe3da71

📥 Commits

Reviewing files that changed from the base of the PR and between 4b69c59 and 085a8d6.

📒 Files selected for processing (21)

• 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/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 (1)

• tensorrt_llm/llmapi/llm_args.py

[coderabbitai]

🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win

Check the lower bound of receiver_slice_id.

The assertion checks only the upper bound. A negative value indexes self._kv_tasks from the end and resolves the wrong task instead of failing. NO_SLICE_ID is -1 and is carried in the same frame, so a future sender-side mix-up of the two fields would corrupt completion silently rather than raise.

🛡️ Proposed guard

-            assert receiver_slice_id < len(self._kv_tasks), (
+            assert 0 <= receiver_slice_id < len(self._kv_tasks), (
                 f"Receiver got receiver_slice_id={receiver_slice_id} (sender_slice_id="
                 f"{sender_slice_id}) but only has {len(self._kv_tasks)} receive task(s) "
                 f"for request {self.request_id}. Sender/receiver slice count mismatch."
             )

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	with self.lock:
	self.kv_cache_size_bytes += transfer_size
	assert sender_slice_id < len(self._kv_tasks), (
	f"Receiver got slice_id={sender_slice_id} from sender but only has "
	f"{len(self._kv_tasks)} receive task(s) for request {self.request_id}. "
	f"Sender/receiver slice count mismatch."
	assert receiver_slice_id < len(self._kv_tasks), (
	f"Receiver got receiver_slice_id={receiver_slice_id} (sender_slice_id="
	f"{sender_slice_id}) but only has {len(self._kv_tasks)} receive task(s) "
	f"for request {self.request_id}. Sender/receiver slice count mismatch."
	)
	task = self._kv_tasks[sender_slice_id]
	task = self._kv_tasks[receiver_slice_id]
	with self.lock:
	self.kv_cache_size_bytes += transfer_size
	assert 0 <= receiver_slice_id < len(self._kv_tasks), (
	f"Receiver got receiver_slice_id={receiver_slice_id} (sender_slice_id="
	f"{sender_slice_id}) but only has {len(self._kv_tasks)} receive task(s) "
	f"for request {self.request_id}. Sender/receiver slice count mismatch."
	)
	task = self._kv_tasks[receiver_slice_id]

🤖 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 2014 -
2021, Update the assertion in the lock block before indexing self._kv_tasks to
require receiver_slice_id >= 0 as well as being less than len(self._kv_tasks).
Preserve the existing mismatch error context, then use the validated
receiver_slice_id to select the task.

[coderabbitai]

🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift

Pass the effective chunked-prefill state.

create_py_executor() can set enable_chunked_context = False after validating the original LLM arguments. This call still passes llm_args.enable_chunked_prefill, so the factory accepts pipelined transfer although the scheduler has no chunked prefill. Pass ctx_chunk_config is not None, or pass the resolved enable_chunked_context through create_py_executor_instance.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/pyexecutor/_util.py` around lines 3000 - 3007, Update the
call to create_kv_cache_transceiver in create_py_executor to pass the effective
chunked-prefill state (ctx_chunk_config is not None or the resolved
enable_chunked_context) instead of llm_args.enable_chunked_prefill, ensuring the
factory reflects the scheduler’s actual configuration.

[coderabbitai]

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Correct the runtime guidance.

When enable_pipelined_transfer is enabled, transceiver_runtime="CPP" raises ValueError at Lines 143-146. It does not disable pipelined transfer. Tell users to set transceiver_runtime="PYTHON" explicitly, or to disable enable_pipelined_transfer.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py` around lines 137 -
142, Update the warning in the enable_pipelined_transfer handling to state that
users must explicitly set transceiver_runtime="PYTHON" or disable
enable_pipelined_transfer; remove the incorrect guidance that
transceiver_runtime="CPP" disables pipelined transfer. Keep the automatic
cache_transceiver_config.transceiver_runtime assignment unchanged.

[coderabbitai]

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Add @skip_pre_hopper to match the sibling NIXL Python tests.

TestGemma3_1BInstruct::test_auto_dtype and TestGemma3_1BInstruct::test_kv_cache_v2_nixl_python both carry @skip_pre_hopper. This new test uses the same NIXL Python transceiver path but has no architecture guard. On pre-Hopper hardware the test runs instead of skipping.

🩹 Proposed fix

     `@pytest.mark.skip_less_device`(2)
+    `@skip_pre_hopper`
     `@parametrize_with_ids`("enable_block_reuse", [True])
     `@parametrize_with_ids`("disable_overlap_scheduler", [False])
     def test_pipelined_kv_transfer_nixl_python_accuracy(

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	@pytest.mark.skip_less_device(2)
	@parametrize_with_ids("enable_block_reuse", [True])
	@parametrize_with_ids("disable_overlap_scheduler", [False])
	`@pytest.mark.skip_less_device`(2)
	`@skip_pre_hopper`
	`@parametrize_with_ids`("enable_block_reuse", [True])
	`@parametrize_with_ids`("disable_overlap_scheduler", [False])

🤖 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/integration/defs/accuracy/test_disaggregated_serving.py` around lines
1593 - 1595, Add the skip_pre_hopper decorator to the new test’s decorator list
alongside skip_less_device and the existing parameterization decorators,
matching the guards on TestGemma3_1BInstruct::test_auto_dtype and
test_kv_cache_v2_nixl_python.

[coderabbitai]

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Confirm the surrounding YAML structure and tier for the new entries.
fd -t f 'l0_dgx_b200.yml' --exec sed -n '1,40p' {}

# Show the parametrize decorators so the generated ids can be derived.
rg -n -B8 'def test_pipelined_kv_transfer_nixl_python_accuracy' \
  tests/integration/defs/accuracy/test_disaggregated_serving.py

# Confirm the ids are not also present in other test-db lists.
fd -t f -e yml . tests/integration/test_lists/test-db \
  --exec rg -n 'test_pipelined_kv_transfer_nixl_python_accuracy' {}

Repository: NVIDIA/TensorRT-LLM

Length of output: 4301

---

🏁 Script executed:

#!/bin/bash
set -eu

python3 - <<'PY'
from pathlib import Path
import re

yaml = Path("tests/integration/test_lists/test-db/l0_dgx_b200.yml").read_text()
source = Path("tests/integration/defs/accuracy/test_disaggregated_serving.py").read_text()

# Extract the first relevant YAML block and its five entries.
block = yaml.split("- condition:", 1)[1]
block = block.split("- condition:", 1)[0]
entries = [
    line.strip()[2:]
    for line in block.splitlines()
    if "test_pipelined_kv_transfer_nixl_python_accuracy[" in line
]
print("test-db entries:", len(entries))
for entry in entries:
    print(entry)

# Extract the focused method bodies and their immediate decorator stacks.
for cls in ("TestLlama3_1_8BInstruct", "TestGemma3_1BInstruct"):
    cls_match = re.search(
        rf"class {cls}\b(?P<body>.*?)(?=\nclass |\Z)", source, re.S
    )
    if not cls_match:
        print(f"{cls}: class not found")
        continue
    body = cls_match.group("body")
    method_match = re.search(
        r"(?P<decorators>(?:    @[^\n]+\n)+)"
        r"    def test_pipelined_kv_transfer_nixl_python_accuracy\((?P<signature>.*?)\n    \):"
        r"(?P<method_body>.*?)(?=\n    def |\nclass |\Z)",
        body,
        re.S,
    )
    if not method_match:
        print(f"{cls}: method not found")
        continue
    decorators = [
        line.strip()
        for line in method_match.group("decorators").splitlines()
        if "parametrize_with_ids" in line
    ]
    print(f"{cls} parameter decorators:")
    for decorator in decorators:
        print(" ", decorator)

    # Pytest applies stacked parametrizations in the observed ID order:
    # the lower decorator's parameter is the first ID fragment.
    params = []
    for decorator in reversed(decorators):
        m = re.search(r'"([^"]+)",\s*\[([^\]]+)\]', decorator)
        if not m:
            raise SystemExit(f"Cannot parse decorator: {decorator}")
        values = [v.strip() for v in m.group(2).split(",")]
        params.append((m.group(1), [v for v in values]))
    generated = []
    def visit(i, parts):
        if i == len(params):
            generated.append(
                f"{cls}::test_pipelined_kv_transfer_nixl_python_accuracy["
                + "-".join(f"{name}={value}" for name, value in parts)
                + "]"
            )
            return
        name, values = params[i]
        for value in values:
            visit(i + 1, parts + [(name, value)])
    visit(0, [])
    print("derived IDs:")
    for item in generated:
        print(" ", item)

# Confirm the five exact entries are in only one test-db file.
matches = []
needle = "test_pipelined_kv_transfer_nixl_python_accuracy"
for path in Path("tests/integration/test_lists/test-db").rglob("*.yml"):
    text = path.read_text()
    count = text.count(needle)
    if count:
        matches.append((str(path), count))
print("test-db files containing target:", matches)
PY

sed -n '720,775p' tests/integration/defs/accuracy/test_disaggregated_serving.py
sed -n '1565,1615p' tests/integration/defs/accuracy/test_disaggregated_serving.py

Repository: NVIDIA/TensorRT-LLM

Length of output: 5651

---

🏁 Script executed:

#!/bin/bash
set -eu

printf '%s\n' '--- class declarations and helper ---'
rg -n '^(class |def parametrize_with_ids|parametrize_with_ids\s*=|from .*parametrize_with_ids)' \
  tests/integration/defs/accuracy/test_disaggregated_serving.py \
  tests/integration/defs -g '*.py' | head -80

printf '%s\n' '--- focused Llama method ---'
sed -n '735,790p' tests/integration/defs/accuracy/test_disaggregated_serving.py

printf '%s\n' '--- focused Gemma method ---'
sed -n '1588,1665p' tests/integration/defs/accuracy/test_disaggregated_serving.py

printf '%s\n' '--- possible CBTS artifacts ---'
find . -maxdepth 5 \( -name 'cbts_touchmap.sqlite' -o -iname '*cbts*coverage*' -o -iname '*coverage*report*' \) -print

Repository: NVIDIA/TensorRT-LLM

Length of output: 14281

---

🏁 Script executed:

#!/bin/bash
set -eu

helper_file="$(rg -l 'def parametrize_with_ids' tests/integration -g '*.py' | head -1)"
printf 'helper_file=%s\n' "$helper_file"
rg -n -A35 -B5 'def parametrize_with_ids' "$helper_file"

Repository: NVIDIA/TensorRT-LLM

Length of output: 1862

---

Reduce the pre-merge matrix for disaggregated accuracy tests.

l0_dgx_b200.yml adds four Llama and one Gemma node IDs to the stage: pre_merge tier. Each test launches one context server, one generation server, and a full GSM8K evaluation. Move lower-priority combinations to post-merge or retain only representative combinations.

Test coverage summary: no test functions changed; five test-list entries were added and none were removed. All five entries match collected node IDs. Coverage verdict: needs follow-up because no cbts_touchmap.sqlite or CBTS coverage report is available.

🤖 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/integration/test_lists/test-db/l0_dgx_b200.yml` around lines 21 - 26,
Reduce the pre_merge entries under the “Disaggregated Serving: Pipelined KV
Transfer (multi-GPU)” section in l0_dgx_b200.yml by retaining only
representative Llama and Gemma combinations; move the lower-priority node IDs to
the post_merge tier instead of adding all five to pre_merge.

Source: Path instructions

[coderabbitai]

🧹 Nitpick comments (3)

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

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

Remove the stale _chunk_size_blocks setup.

pipeline_transfer_enabled reads only _enable_pipelined_transfer. KvCacheTransceiverV2 no longer defines _chunk_size_blocks, so line 663 configures an attribute the property never reads. Removing it keeps the test aligned with the current property.

♻️ 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 661 -
666, Remove the unnecessary transceiver._chunk_size_blocks assignment from the
test for KvCacheTransceiverV2.pipeline_transfer_enabled, leaving the
_enable_pipelined_transfer setup and assertion unchanged.

---

833-853: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Set pipeline_transfer_enabled explicitly and add a monolithic-branch test.

transceiver is a MagicMock, so transceiver.pipeline_transfer_enabled is truthy no matter what _enable_pipelined_transfer is set to. These tests therefore take the pipelined branch through mock truthiness, not through the flag. No test covers the else branch at line 760 of tensorrt_llm/_torch/disaggregation/transceiver.py, which is the default non-pipelined path that calls _create_kv_slice.

Set the property on the mock directly, and add a test for the monolithic path.

💚 Proposed test additions

     transceiver = MagicMock()
-    transceiver._enable_pipelined_transfer = True
+    transceiver.pipeline_transfer_enabled = True
     transceiver.kv_transfer_timeout_ms = None
     transceiver._get_or_create_send_session.return_value = session
     transceiver._build_prefill_chunk.return_value = last_slice

Add a companion test for the non-pipelined path:

def test_monolithic_transfer_sends_full_slice():
    """With pipelining off, respond_and_send_async sends one full-prompt slice."""
    from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2

    session = MagicMock()
    full_slice = KVSlice(
        is_last_slice=True,
        block_ids_per_layer_groups=[np.array([0, 1], dtype=np.int64)],
    )

    transceiver = MagicMock()
    transceiver.pipeline_transfer_enabled = False
    transceiver._get_or_create_send_session.return_value = session
    transceiver._create_kv_slice.return_value = full_slice

    request = SimpleNamespace(
        py_disaggregated_params=DisaggregatedParams(disagg_request_id=42),
        request_id=42,
        prompt_len=8,
        py_beam_width=1,
        py_kv_transfer_start_time=None,
        set_kv_cache_transfer_start=lambda _ts: None,
    )

    KvCacheTransceiverV2.respond_and_send_async(transceiver, request)

    transceiver._create_kv_slice.assert_called_once_with(request)
    transceiver._build_prefill_chunk.assert_not_called()
    session.send.assert_called_once_with(full_slice)
    transceiver._finalize_send.assert_called_once_with(request, session)

🤖 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 833 -
853, Set pipeline_transfer_enabled directly on the MagicMock in the existing
pipelined test so branch selection uses the intended flag rather than mock
truthiness. Add a companion test for KvCacheTransceiverV2.respond_and_send_async
with pipeline_transfer_enabled=False, asserting _create_kv_slice(request) is
called, _build_prefill_chunk is not called, and the resulting full slice is sent
and finalized.

---

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

Test coverage: add two missing cases.

• The new file adds all listed test_* functions; no tests were modified or removed.
• Unit tests require no per-test list entry. The B200 test list contains five pipelined accuracy entries.
• Coverage is insufficient. Add a non-block-aligned chunk_end_pos case for _build_prefill_chunk, and test the respond_and_send_async non-pipelined branch (_create_kv_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_chunked_transfer.py` around lines 15 - 19,
Extend the tests in the chunked transfer test module with two cases: cover a
non-block-aligned chunk_end_pos for _build_prefill_chunk, and cover the
non-pipelined respond_and_send_async path, including its _create_kv_slice
behavior. Keep the existing pipelined accuracy entries and test structure
unchanged.

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.

Nitpick comments:
In `@tests/unittest/disaggregated/test_chunked_transfer.py`:
- Around line 661-666: Remove the unnecessary transceiver._chunk_size_blocks
assignment from the test for KvCacheTransceiverV2.pipeline_transfer_enabled,
leaving the _enable_pipelined_transfer setup and assertion unchanged.
- Around line 833-853: Set pipeline_transfer_enabled directly on the MagicMock
in the existing pipelined test so branch selection uses the intended flag rather
than mock truthiness. Add a companion test for
KvCacheTransceiverV2.respond_and_send_async with
pipeline_transfer_enabled=False, asserting _create_kv_slice(request) is called,
_build_prefill_chunk is not called, and the resulting full slice is sent and
finalized.
- Around line 15-19: Extend the tests in the chunked transfer test module with
two cases: cover a non-block-aligned chunk_end_pos for _build_prefill_chunk, and
cover the non-pipelined respond_and_send_async path, including its
_create_kv_slice behavior. Keep the existing pipelined accuracy entries and test
structure unchanged.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: e64bf3a5-860b-4669-bf05-3aeefe7622e2

📥 Commits

Reviewing files that changed from the base of the PR and between 085a8d6 and 661dee6.

📒 Files selected for processing (4)

• tensorrt_llm/_torch/disaggregation/transceiver.py
• tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py
• tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py
• tests/unittest/disaggregated/test_chunked_transfer.py

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

• tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py