[TRTLLM-12499][perf] Bounded block-ID retrieval for pipelined KV transfer chunks - #17526
Conversation
…ving in Python Cache Transceiver Instead of waiting for all prefill chunks to complete before starting KV cache transfer, each chunk's KV data is transferred to the generation server immediately after its prefill completes. This overlaps GPU compute with RDMA transfer, hiding transfer latency behind prefill computation. Only the last chunk's transfer remains on the critical path. The feature is gated behind `enable_pipelined_transfer` on `CacheTransceiverConfig` and is implemented in `KvCacheTransceiverV2` only. It requires `schedule_style: generation_first`, `enable_chunked_prefill: true`, `beam_width == 1`, the NIXL backend, `kv_cache_bounce_size_mb == 0`, `pipeline_parallel_size == 1` on the sender, and a non-Mamba/hybrid cache manager. Each requirement is enforced at startup or per request. Squashed from 15 commits: - Chunking is sender-side only; the generation server posts a single receive covering the whole prompt and completes on `is_last_slice`. - `KVSlice` now describes one chunk rather than one whole request, gaining `total_blocks` and a meaningful `is_last_slice`. `prompt_len` became required on the session args so SWA can compute the stale-block boundary. - `project_blocks_to_global_chunk` intersects ranges instead of indexing, so resident-suffix block lists (sliding window groups, prefix reuse, incremental allocation) project correctly onto a global chunk. - The first slice always extends back to block 0, so a context-side prefix-reuse hit does not leave `[0, prepopulated_prompt_len)` unsent. - Source blocks are capped at the computed chunk boundary before SWA trimming, normalizing V1's full-prompt reservation against V2's incremental allocation. - `KV_AGENT_RESULT` carries `sender_slice_id` and `receiver_slice_id` separately, making per-chunk RDMA failures attributable. Behavior-neutral for the monolithic receiver. - KV transfer activity is modeled by transceiver session membership rather than `LlmRequestState`, so mid-prefill cancellation and transfer-timeout monitoring work during the pipelined phase. - A retired send session cannot be silently re-created, since closing it drops the peer's `RecvReqInfo` and the receiver never re-registers. - `TxSession.dispatch_lock` serializes chunk dispatch across the executor thread and the late-peer replay path, so a newer slice cannot reach a peer's queue ahead of an older one. - Transceiver configuration resolution happens early and idempotently, and backend/runtime compatibility validation is centralized. Signed-off-by: Athena Cai <athenac@nvidia.com>
Signed-off-by: Athena Cai <athenac@nvidia.com>
_build_prefill_chunk now asks each layer group only for the block range the current chunk needs instead of rebuilding a whole-request slice and projecting it. A windowed group's range starts at the request's final-window stale boundary: the sender worker floors src_start there without dropping the blocks under it, so a run reaching further back would be paired with the window tail's destination blocks. A chunk entirely below the window skips the cache query altogether. V1 gains a bounded C++/nanobind query and keeps the block-ID to pool-slot translation get_block_ids does. V2 cuts the range out of the aggregated page list, since neither of its backends has a bounded query. Signed-off-by: Athena Cai <athenac@nvidia.com>
WalkthroughThis change adds bounded cache block range APIs and exposes them to Python. It implements pipelined disaggregated KV transfer with chunk projection, independent sender and receiver slice IDs, session retirement, transfer ownership tracking, configuration validation, and expanded unit and integration tests. ChangesCache range APIs
Pipelined transfer setup
Transfer lifecycle
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant PyExecutor
participant KvCacheTransceiverV2
participant KVCacheManager
participant TxSession
participant RxSession
PyExecutor->>KvCacheTransceiverV2: submit active prefill request
KvCacheTransceiverV2->>KVCacheManager: retrieve resident block range
KVCacheManager-->>KvCacheTransceiverV2: return block IDs
KvCacheTransceiverV2->>TxSession: create and dispatch chunk
TxSession->>RxSession: transfer chunk result with sender and receiver slice IDs
RxSession-->>PyExecutor: report completion or failure
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1⚔️ Resolve merge conflicts 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (5)
tensorrt_llm/_torch/disaggregation/base/transfer.py (1)
101-124: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUpdate the
total_blocksderivation described in the docstring.The docstring still states
total_blocks = ceil(token_range.end / tpb). The sender no longer computes it that way.Sender._build_kv_write_metaintensorrt_llm/_torch/disaggregation/native/transfer.pyreadstask._slice.total_blocksfirst, and falls back toceil(prompt_len / tpb)when the new field isNone. For a non-final pipelined chunk the two formulas differ, so the documented rule is wrong for exactly the case the new field was added to cover.Describe the new field in the docstring:
total_blocksis the full prompt block count that destination projection uses, and the per-layer token start is derived from the chunk suffix boundary rather than fromtoken_range.end.🤖 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/base/transfer.py` around lines 101 - 124, Update the docstring near block-coordinate derivation to describe total_blocks as the full prompt block count used by destination projection, rather than ceil(token_range.end / tpb). State that per-layer token starts are derived from the chunk suffix boundary using the block count, while preserving the existing cached-prefix and beam-search descriptions.tensorrt_llm/_torch/disaggregation/native/transfer.py (2)
189-204: 🗄️ Data Integrity & Integration | 🔵 TrivialPlan the rollout for the changed
KV_AGENT_RESULTframe.
_KV_RESULT_PREFIXgained a field, so its packed size changed. A context server on the new build and a generation server on the old build will not interoperate. The receiver's_process_kv_agent_resultcalls_KV_RESULT_PREFIX.unpack, a size mismatch raisesstruct.error, and_start_listeneronly logs that exception. The receive task then staysTRANSFERRINGuntil the KV transfer timeout fires, so the failure looks like a hang rather than a version mismatch.The inline comment already states that both servers must run matching builds. Add the same requirement to the release or upgrade notes for disaggregated deployments, and upgrade context and generation servers together rather than in a rolling fashion.
🤖 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 189 - 204, Add a release or upgrade note for disaggregated deployments stating that context and generation servers must run matching builds because the changed KV_AGENT_RESULT frame is not wire-compatible. Explicitly instruct operators to upgrade both server roles together rather than using a rolling upgrade.
1172-1189: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winMove the failure-result send out of
session.lock.Line 1180 calls
_send_failed_result_to_receiverwhile the listener thread holds bothsession.dispatch_lockandsession.lock. That helper performs a ZMQ DEALERsend, which can block when the socket reaches its high-water mark. While it blocks,TxSession.cancel,TxSession.set_exception, andTxSession.sendall stall onsession.lock.Capture the decision inside the lock and perform the send after releasing it. The rest of the method already follows that pattern: it snapshots
tasksundersession.lockand builds metadata outside it.♻️ Proposed change
with session.dispatch_lock: + abort_receiver = False with session.lock: self._save_peer_req_info(info) tasks = list(session.kv_tasks) # No tasks: no worker will send KV_AGENT_RESULT FAILED to the receiver. # Send it directly to unblock the receiver's TRANSFERRING task event; # CANCEL_SESSION alone would leave it stuck indefinitely. if not tasks and session.status in (SessionStatus.ERROR, SessionStatus.CANCELLED): - self._send_failed_result_to_receiver(info) - return + abort_receiver = True + if abort_receiver: + self._send_failed_result_to_receiver(info) + return for task in tasks:🤖 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 1172 - 1189, In the listener method containing the task snapshot, record a local flag inside session.lock when tasks are empty and the session is ERROR or CANCELLED, then release the lock before calling _send_failed_result_to_receiver(info). Preserve the early-return behavior after performing the send outside session.lock, while leaving task processing unchanged.tensorrt_llm/_torch/disaggregation/transceiver.py (1)
699-712: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winDocument the invariant that keeps a mid-prefill session from being retired early.
_send_reqs[rid]is now populated on the first chunk, socheck_context_transfer_statuscan observe this request while later chunks are still pending._collect_donecallssession.is_completed(), andTxSession.statusreportsKV_TRANSFERREDas soon as every already createdkv_taskis transferred. For a request whose first chunk landed and whose second chunk is not yet sent, that condition is true.The only thing preventing an early
_retire_send_sessionisTxSession._need_aux: withschedule_style=generation_first,is_completed()requiresFULLY_TRANSFERRED, and the aux task is created only in_finalize_sendon the last chunk. Pipelined transfer is separately validated to requiregeneration_firstindisagg_utils._validate_disagg_configandPyExecutor._validate_request.That invariant spans three files and is load-bearing: if it ever breaks, a request's KV is silently truncated mid-prefill. Record it here and assert it once per session.
♻️ Suggested guard
assert req.py_beam_width == 1, "beam_width > 1 is not supported for chunked KV transfer" + # A partially sent session must not look complete to + # check_context_transfer_status. TxSession only withholds completion + # until _finalize_send creates the aux task, and it only does that for + # generation_first. Pipelined transfer is validated to require that + # schedule style; re-assert it here so the coupling cannot silently break. + assert self._need_aux_transfer(req), ( + "pipelined transfer requires schedule_style=generation_first" + ) rid = get_unique_rid(req)🤖 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 699 - 712, Document and enforce in _build_prefill_chunk the invariant that pipelined prefill sessions use generation_first, so TxSession._need_aux remains set until _finalize_send creates the auxiliary task for the final chunk. Add the assertion only when the session is first registered in _send_reqs, preserving subsequent chunk handling and preventing check_context_transfer_status/_collect_done from retiring the session early.tests/unittest/disaggregated/test_kv_transfer.py (1)
2166-2166: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd
strict=Trueto these twozip()calls for consistency with the sibling helper.
add_and_verify_chunked_requestuseszip(..., strict=True)at lines 1944 and 1955.add_and_verify_pipelined_requestomits it here. Ifsender_sessionsandctx_block_idsever differ in length,ziptruncates silently and a rank's transfer is never driven, so the test passes without exercising it. Thev2_tp2_pp1_pipelinedandv1_tp2_pp1_pipelinedconfigurations run with two ranks, so the truncation would be silent there.This is a consistency and test-reliability point, not a Ruff B905 finding; B905 is not enabled in this repository.
♻️ Proposed change
- for sender_session, block_ids_per_groups in zip(sender_sessions, ctx_block_ids): + for sender_session, block_ids_per_groups in zip(sender_sessions, ctx_block_ids, strict=True): @@ - for recv_session, block_ids_per_groups in zip(receiver_sessions, gen_block_ids): + for recv_session, block_ids_per_groups in zip( + receiver_sessions, gen_block_ids, strict=True + ):Also applies to: 2178-2178
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/disaggregated/test_kv_transfer.py` at line 2166, Update both zip() calls in add_and_verify_pipelined_request to pass strict=True, matching the existing calls in add_and_verify_chunked_request and ensuring mismatched sender_sessions and ctx_block_ids lengths fail instead of truncating.Source: Learnings
🤖 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/resource/cache_reuse.py`:
- Around line 215-242: Update get_block_ids_range to use a bounded page-index
retrieval API on KVCacheManagerV2, requesting only the [block_begin, block_end)
range instead of materializing the full aggregated list and slicing it locally.
Preserve placeholder entries and the existing gap-trimming behavior, and add the
corresponding manager API implementation so the Python/NIXL pipelined-transfer
path performs bounded retrieval.
In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 3008-3015: Update the enable_chunked_prefill argument in the
create_kv_cache_transceiver call within create_py_executor to pass whether
ctx_chunk_config is not None, rather than llm_args.enable_chunked_prefill, so
the transceiver receives the scheduler’s effective chunking state.
In `@tests/integration/defs/accuracy/test_disaggregated_serving.py`:
- Around line 1593-1597: Add the `@skip_pre_hopper` decorator to
test_pipelined_kv_transfer_nixl_python, matching the other tests in
TestGemma3_1BInstruct and the corresponding Llama test while preserving the
existing decorators.
In `@tests/integration/test_lists/test-db/l0_dgx_b200.yml`:
- Around line 21-26: Reverse the parameter order in all five parametrized test
IDs under the Disaggregated Serving section, placing enable_block_reuse before
disable_overlap_scheduler while preserving each value combination and test
target.
In `@tests/unittest/disaggregated/test_chunked_transfer.py`:
- Around line 678-682: Escape the literal periods in the pytest.raises match
patterns at the three affected assertions around create_kv_cache_transceiver,
including the cases near lines 680, 706, and 795. Use re.escape or regex-escaped
raw strings so each assertion matches the exact intended error message rather
than treating dots as wildcards.
---
Nitpick comments:
In `@tensorrt_llm/_torch/disaggregation/base/transfer.py`:
- Around line 101-124: Update the docstring near block-coordinate derivation to
describe total_blocks as the full prompt block count used by destination
projection, rather than ceil(token_range.end / tpb). State that per-layer token
starts are derived from the chunk suffix boundary using the block count, while
preserving the existing cached-prefix and beam-search descriptions.
In `@tensorrt_llm/_torch/disaggregation/native/transfer.py`:
- Around line 189-204: Add a release or upgrade note for disaggregated
deployments stating that context and generation servers must run matching builds
because the changed KV_AGENT_RESULT frame is not wire-compatible. Explicitly
instruct operators to upgrade both server roles together rather than using a
rolling upgrade.
- Around line 1172-1189: In the listener method containing the task snapshot,
record a local flag inside session.lock when tasks are empty and the session is
ERROR or CANCELLED, then release the lock before calling
_send_failed_result_to_receiver(info). Preserve the early-return behavior after
performing the send outside session.lock, while leaving task processing
unchanged.
In `@tensorrt_llm/_torch/disaggregation/transceiver.py`:
- Around line 699-712: Document and enforce in _build_prefill_chunk the
invariant that pipelined prefill sessions use generation_first, so
TxSession._need_aux remains set until _finalize_send creates the auxiliary task
for the final chunk. Add the assertion only when the session is first registered
in _send_reqs, preserving subsequent chunk handling and preventing
check_context_transfer_status/_collect_done from retiring the session early.
In `@tests/unittest/disaggregated/test_kv_transfer.py`:
- Line 2166: Update both zip() calls in add_and_verify_pipelined_request to pass
strict=True, matching the existing calls in add_and_verify_chunked_request and
ensuring mismatched sender_sessions and ctx_block_ids lengths fail instead of
truncating.
🪄 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: 60eee558-02c3-478e-8bc1-7b3695a2565c
📒 Files selected for processing (29)
cpp/include/tensorrt_llm/batch_manager/kvCacheManager.hcpp/tensorrt_llm/batch_manager/kvCacheManager.cppcpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cppcpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpptensorrt_llm/_torch/auto_deploy/shim/ad_executor.pytensorrt_llm/_torch/disaggregation/base/transfer.pytensorrt_llm/_torch/disaggregation/native/transfer.pytensorrt_llm/_torch/disaggregation/resource/cache_reuse.pytensorrt_llm/_torch/disaggregation/transceiver.pytensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.pytensorrt_llm/_torch/pyexecutor/llm_request.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/py_executor_creator.pytensorrt_llm/_torch/pyexecutor/resource_manager.pytensorrt_llm/commands/serve.pytensorrt_llm/llmapi/disagg_utils.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/usage/llm_args_golden_manifest.jsontests/integration/defs/accuracy/test_disaggregated_serving.pytests/integration/test_lists/test-db/l0_dgx_b200.ymltests/unittest/_torch/executor/test_disagg_index_mapper_early_release.pytests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.pytests/unittest/disaggregated/test_bounce.pytests/unittest/disaggregated/test_cache_reuse_adapter.pytests/unittest/disaggregated/test_chunked_transfer.pytests/unittest/disaggregated/test_disagg_utils.pytests/unittest/disaggregated/test_kv_transfer.pytests/unittest/disaggregated/test_transceiver_bounded_polling.py
| def get_block_ids_range(self, req, group_idx, lg, block_begin, block_end): # noqa: ARG002 | ||
| if block_begin < 0 or block_end < 0: | ||
| raise ValueError("block range bounds must be non-negative") | ||
| if block_begin > block_end: | ||
| raise ValueError("block_begin must not exceed block_end") | ||
| # Neither V2 backend has a bounded query, so read the whole aggregated | ||
| # list -- keeping the placeholders, which is what makes an entry's index | ||
| # its block ordinal -- and cut the range out of it here. | ||
| all_block_ids = np.fromiter( | ||
| self._mgr.kv_cache_map[req.py_request_id].get_aggregated_page_indices( | ||
| group_idx, valid_only=False | ||
| ), | ||
| dtype=np.int64, | ||
| ) | ||
| if block_end > all_block_ids.size: | ||
| raise ValueError( | ||
| f"block_end={block_end} exceeds the {all_block_ids.size} allocated blocks; " | ||
| "the result would not end at block_end, and callers recover block " | ||
| "ordinals from its length" | ||
| ) | ||
| block_ids = all_block_ids[block_begin:block_end] | ||
| # Drop everything up to the last gap rather than compacting around it: | ||
| # a life cycle that keeps sink tokens resident has a sink prefix plus a | ||
| # window suffix, and returning both would misreport where the suffix | ||
| # starts. | ||
| gaps = np.flatnonzero(block_ids == BAD_PAGE_INDEX) | ||
| return block_ids[gaps[-1] + 1 :] if gaps.size else block_ids | ||
|
|
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- candidate files ---'
git ls-files | rg '(^|/)(cache_reuse|kv_cache_manager|disaggregation|.*test.*|CODING_GUIDELINES\.md)$|enable_pipelined_transfer|pipelined_transfer' | head -200
printf '%s\n' '--- symbols and call sites ---'
rg -n -S \
'get_block_ids_range|get_aggregated_page_indices|enable_pipelined_transfer|_build_prefill_chunk|KVCacheManagerV2|_CacheReuseAdapterV2|_CacheReuseAdapterV1' \
tensorrt_llm tests 2>/dev/null | head -300
printf '%s\n' '--- outline ---'
ast-grep outline tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py --view expanded
printf '%s\n' '--- implementation ---'
sed -n '1,280p' tensorrt_llm/_torch/disaggregation/resource/cache_reuse.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- cache reuse adapter ---'
sed -n '31,270p' tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py
printf '%s\n' '--- prefill chunk construction ---'
sed -n '680,845p' tensorrt_llm/_torch/disaggregation/transceiver.py
printf '%s\n' '--- transceiver selection and restrictions ---'
sed -n '110,235p' tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py
printf '%s\n' '--- V2 aggregation implementation ---'
sed -n '590,680p' tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
sed -n '330,360p' tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi
printf '%s\n' '--- adapter tests ---'
sed -n '100,220p' tests/unittest/disaggregated/test_cache_reuse_adapter.py
printf '%s\n' '--- V2 transfer test locations ---'
sed -n '1450,1535p' tests/unittest/disaggregated/test_kv_transfer.py
sed -n '2080,2135p' tests/integration/defs/accuracy/test_disaggregated_serving.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 39325
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- pipeline and V2 combinations ---'
rg -n -C 8 -S \
'enable_pipelined_transfer.{0,200}|use_kv_cache_manager_v2.{0,200}|KVCacheManagerV2 \+ NIXL|V2.*pipelined|pipelined.*V2' \
tests tensorrt_llm | head -400
printf '%s\n' '--- chunked transfer test setup ---'
sed -n '620,745p' tests/unittest/disaggregated/test_chunked_transfer.py
sed -n '1060,1210p' tests/unittest/disaggregated/test_chunked_transfer.py
printf '%s\n' '--- integration pipeline configurations ---'
sed -n '720,775p' tests/integration/defs/accuracy/test_disaggregated_serving.py
sed -n '1410,1465p' tests/integration/defs/accuracy/test_disaggregated_serving.py
sed -n '1580,1630p' tests/integration/defs/accuracy/test_disaggregated_serving.py
sed -n '2560,2660p' tests/integration/defs/accuracy/test_disaggregated_serving.py
printf '%s\n' '--- request allocation references ---'
rg -n -C 5 -S 'self\._blocks|_blocks =|num_blocks|allocate.*block|prompt_len|py_last_context_chunk' \
tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py \
tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py \
tensorrt_llm/_torch/disaggregation/transceiver.py | head -350Repository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -eu
python3 - <<'PY'
from dataclasses import dataclass
BAD_PAGE_INDEX = -1
`@dataclass`
class FakeCache:
blocks: list[int]
yielded: int = 0
def get_aggregated_page_indices(self, valid_only: bool = False):
for block in self.blocks:
self.yielded += 1
if block == BAD_PAGE_INDEX:
if not valid_only:
yield BAD_PAGE_INDEX
else:
yield block
def adapter_range(cache: FakeCache, begin: int, end: int) -> list[int]:
all_ids = list(cache.get_aggregated_page_indices(valid_only=False))
if end > len(all_ids):
raise ValueError
block_ids = all_ids[begin:end]
last_gap = max((i for i, value in enumerate(block_ids)
if value == BAD_PAGE_INDEX), default=None)
return block_ids[last_gap + 1:] if last_gap is not None else block_ids
for block_count, ranges in [
(8, [(0, 2), (2, 4), (4, 6), (6, 8)]),
(16, [(0, 4), (4, 8), (8, 12), (12, 16)]),
]:
cache = FakeCache(list(range(block_count)))
results = [adapter_range(cache, begin, end) for begin, end in ranges]
print({
"block_count": block_count,
"ranges": ranges,
"results": results,
"source_blocks_consumed": cache.yielded,
"expected_full_scan_count": block_count * len(ranges),
})
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 531
Add a bounded page-index query for KVCacheManagerV2.
get_block_ids_range materializes every allocated block before slicing [block_begin, block_end). _build_prefill_chunk calls it for each applicable layer group and chunk, causing O(number of chunks × allocated blocks) scanning per group. This affects the supported KVCacheManagerV2 with Python/NIXL pipelined-transfer path and defeats the bounded-retrieval optimization. Add a bounded V2 page-index API.
🤖 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/resource/cache_reuse.py` around lines 215
- 242, Update get_block_ids_range to use a bounded page-index retrieval API on
KVCacheManagerV2, requesting only the [block_begin, block_end) range instead of
materializing the full aggregated list and slicing it locally. Preserve
placeholder entries and the existing gap-trimming behavior, and add the
corresponding manager API implementation so the Python/NIXL pipelined-transfer
path performs bounded retrieval.
| kv_cache_transceiver = create_kv_cache_transceiver( | ||
| mapping, dist, kv_cache_manager, attention_type, | ||
| cache_transceiver_config, mamba_cache_manager) | ||
| mapping, | ||
| dist, | ||
| kv_cache_manager, | ||
| attention_type, | ||
| cache_transceiver_config, | ||
| mamba_cache_manager, | ||
| enable_chunked_prefill=llm_args.enable_chunked_prefill) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm chunked-prefill can be disabled locally without updating llm_args.
rg -n -C4 'enable_chunked_context' tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
rg -n -C2 'enable_chunked_prefill' tensorrt_llm/_torch/pyexecutor/_util.py tensorrt_llm/_torch/pyexecutor/py_executor_creator.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 5859
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- definitions and call sites ---'
rg -n -C8 'def create_kv_cache_transceiver|create_kv_cache_transceiver\(|ctx_chunk_config|enable_chunked_prefill|chunked_prefill' \
tensorrt_llm/_torch/pyexecutor tensorrt_llm/_torch | head -n 500
printf '%s\n' '--- creator structure ---'
ast-grep outline tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
printf '%s\n' '--- util structure ---'
ast-grep outline tensorrt_llm/_torch/pyexecutor/_util.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -e
python3 - <<'PY'
from pathlib import Path
paths = [
Path("tensorrt_llm/_torch/pyexecutor/_util.py"),
Path("tensorrt_llm/_torch/pyexecutor/py_executor_creator.py"),
]
needles = (
"def create_kv_cache_transceiver",
"enable_chunked_prefill=llm_args.enable_chunked_prefill",
"ctx_chunk_config",
"create_py_executor",
"create_py_executor_instance",
"chunked_prefill",
)
for path in paths:
lines = path.read_text().splitlines()
print(f"\n--- {path} ---")
starts = []
for i, line in enumerate(lines):
if any(n in line for n in needles):
starts.append(i)
shown = set()
for start in starts:
lo = max(0, start - 12)
hi = min(len(lines), start + 20)
for i in range(lo, hi):
if i not in shown:
print(f"{i+1}: {lines[i]}")
shown.add(i)
print()
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 35048
🏁 Script executed:
#!/bin/bash
set -e
python3 - <<'PY'
from pathlib import Path
import re
for name, patterns in {
"tensorrt_llm/_torch/pyexecutor/_util.py": [
r"def create_kv_cache_transceiver",
r"enable_chunked_prefill=llm_args\.enable_chunked_prefill",
],
"tensorrt_llm/_torch/pyexecutor/py_executor_creator.py": [
r"def create_py_executor",
r"enable_chunked_context",
r"ctx_chunk_config",
r"create_py_executor_instance",
],
}.items():
text = Path(name).read_text()
lines = text.splitlines()
print(f"\n--- {name} ({len(lines)} lines) ---")
for pattern in patterns:
print(f"PATTERN {pattern}")
for m in re.finditer(pattern, text):
line = text.count("\n", 0, m.start()) + 1
lo, hi = max(1, line - 15), min(len(lines), line + 30)
print(f"range {lo}-{hi}")
for n in range(lo, hi + 1):
print(f"{n}: {lines[n-1]}")
print()
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 50375
Pass the scheduler's effective chunking state.
create_py_executor can set ctx_chunk_config = None while llm_args.enable_chunked_prefill remains True. Pipelined transfer can then bypass its enable_chunked_prefill precondition. Pass ctx_chunk_config is not None instead.
🤖 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 3008 - 3015, Update the
enable_chunked_prefill argument in the create_kv_cache_transceiver call within
create_py_executor to pass whether ctx_chunk_config is not None, rather than
llm_args.enable_chunked_prefill, so the transceiver receives the scheduler’s
effective chunking state.
| @pytest.mark.skip_less_device(2) | ||
| @parametrize_with_ids("enable_block_reuse", [True]) | ||
| @parametrize_with_ids("disable_overlap_scheduler", [False]) | ||
| def test_pipelined_kv_transfer_nixl_python_accuracy( | ||
| self, enable_block_reuse: bool, disable_overlap_scheduler: bool): |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add @skip_pre_hopper to the new Gemma 3 pipelined test.
Every other test in TestGemma3_1BInstruct carries @skip_pre_hopper (test_auto_dtype at line 1501, test_kv_cache_v2_nixl_python at line 1548). The new test omits it, and the matching Llama test at line 741 has it. The CI entry targets l0_dgx_b200.yml, so CI is unaffected, but a manual QA run on pre-Hopper hardware would attempt this test.
🔧 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.
| @pytest.mark.skip_less_device(2) | |
| @parametrize_with_ids("enable_block_reuse", [True]) | |
| @parametrize_with_ids("disable_overlap_scheduler", [False]) | |
| def test_pipelined_kv_transfer_nixl_python_accuracy( | |
| self, enable_block_reuse: bool, disable_overlap_scheduler: bool): | |
| `@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( | |
| self, enable_block_reuse: bool, disable_overlap_scheduler: bool): |
🤖 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 - 1597, Add the `@skip_pre_hopper` decorator to
test_pipelined_kv_transfer_nixl_python, matching the other tests in
TestGemma3_1BInstruct and the corresponding Llama test while preserving the
existing decorators.
| # ------------- Disaggregated Serving: Pipelined KV Transfer (multi-GPU) --------------- | ||
| - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_pipelined_kv_transfer_nixl_python_accuracy[disable_overlap_scheduler=False-enable_block_reuse=False] | ||
| - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_pipelined_kv_transfer_nixl_python_accuracy[disable_overlap_scheduler=False-enable_block_reuse=True] | ||
| - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_pipelined_kv_transfer_nixl_python_accuracy[disable_overlap_scheduler=True-enable_block_reuse=False] | ||
| - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_pipelined_kv_transfer_nixl_python_accuracy[disable_overlap_scheduler=True-enable_block_reuse=True] | ||
| - accuracy/test_disaggregated_serving.py::TestGemma3_1BInstruct::test_pipelined_kv_transfer_nixl_python_accuracy[disable_overlap_scheduler=False-enable_block_reuse=True] |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Print the collected test ids for the new pipelined KV transfer tests.
set -euo pipefail
python - <<'PY'
import re, pathlib
p = pathlib.Path("tests/integration/defs/accuracy/test_disaggregated_serving.py")
src = p.read_text()
for m in re.finditer(
r"((?:^[ \t]*@[^\n]*\n)+)[ \t]*def (test_pipelined_kv_transfer_nixl_python_accuracy)\(",
src, re.M):
print("=== decorators (top to bottom) ===")
print(m.group(1).rstrip())
PY
echo "--- entries listed in the test-db ---"
rg -nF 'test_pipelined_kv_transfer_nixl_python_accuracy' tests/integration/test_lists/test-db/l0_dgx_b200.ymlRepository: NVIDIA/TensorRT-LLM
Length of output: 1496
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "--- parametrize_with_ids implementation ---"
rg -n -A35 -B10 'def parametrize_with_ids|parametrize_with_ids\s*=' tests/integration tests 2>/dev/null | head -n 220
echo "--- target test definitions and surrounding context ---"
rg -n -A28 -B12 'def test_pipelined_kv_transfer_nixl_python_accuracy' \
tests/integration/defs/accuracy/test_disaggregated_serving.py
echo "--- changed test-list files and repository status ---"
git status --short
git diff --name-status -- tests/integration/test_lists/test-dbRepository: NVIDIA/TensorRT-LLM
Length of output: 11246
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
try:
import pytest
print("pytest version:", pytest.__version__)
except Exception as exc:
print("pytest unavailable:", type(exc).__name__, str(exc))
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 222
🌐 Web query:
pytest stacked parametrize decorators generated test ID order outer inner
💡 Result:
When you stack multiple @pytest.mark.parametrize decorators on a test function, pytest executes the decorators from the bottom up (the innermost decorator is applied first) [1]. However, the order in which the parameters appear in the generated test IDs and the order in which test cases are generated follows the order in which the decorators are applied, starting from the outermost decorator [1][2]. Specifically: 1. Decorator Application: Because Python applies decorators from the innermost to the outermost, the decorator closest to the function definition is technically applied first to the function object [1]. 2. Test ID and Execution Order: Despite the order of application, pytest generates test IDs and Cartesian combinations based on the order of the decorators as they appear visually in the code—from top (outermost) to bottom (innermost) [1][2]. For example: @pytest.mark.parametrize("x", [1]) @pytest.mark.parametrize("y", [2][3]) def test_foo(x, y): pass In this case, the test IDs will be generated starting with the "x" parameter (the outermost decorator), resulting in test names like test_foo[0-2], test_foo[0-3], test_foo[1-2], and test_foo[1-3] [2]. The order of arguments in the test ID reflects the top-to-bottom order of the @pytest.mark.parametrize decorators, not the order of the arguments in the function definition [1].
Citations:
- 1: Order differences in IDs and parameters with multiple @parametrize decorators pytest-dev/pytest#4853
- 2: https://docs.pytest.org/en/stable/how-to/parametrize.html
- 3: https://doc.pytest.org/en/stable/example/parametrize.html
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
import itertools
import pathlib
import re
source = pathlib.Path("tests/integration/defs/accuracy/test_disaggregated_serving.py").read_text()
list_source = pathlib.Path("tests/integration/test_lists/test-db/l0_dgx_b200.yml").read_text()
pattern = re.compile(
r"(?P<decorators>(?:^[ \t]*`@parametrize_with_ids`[^\n]*\n)+)"
r"[ \t]*def test_pipelined_kv_transfer_nixl_python_accuracy"
)
matches = list(pattern.finditer(source))
listed = re.findall(
r"test_pipelined_kv_transfer_nixl_python_accuracy(\[[^]]+\])",
list_source,
)
for index, match in enumerate(matches, 1):
decorators = [
ast.parse(line.strip()[1:], mode="eval").body
for line in match.group("decorators").splitlines()
]
# Pytest composes stacked parametrization IDs in visual decorator order.
parameters = []
for decorator in decorators:
name = ast.literal_eval(decorator.args[0])
values = ast.literal_eval(decorator.args[1])
parameters.append((name, values))
expected = [
"[" + "-".join(f"{name}={value}" for name, value in combination) + "]"
for combination in itertools.product(
*[[(name, value) for value in values] for name, values in parameters]
)
]
print(f"test {index}: decorators top-to-bottom = {[name for name, _ in parameters]}")
print("expected IDs:", expected)
print("listed IDs: ", listed if index == 1 else listed[-1:])
print("matching expected IDs:", sum(item in expected for item in listed))
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
import itertools
import pathlib
import re
test_file = pathlib.Path("tests/integration/defs/accuracy/test_disaggregated_serving.py")
list_file = pathlib.Path("tests/integration/test_lists/test-db/l0_dgx_b200.yml")
tree = ast.parse(test_file.read_text())
listed = re.findall(
r"test_pipelined_kv_transfer_nixl_python_accuracy(\[[^]]+\])",
list_file.read_text(),
)
found = []
for node in ast.walk(tree):
if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
continue
if node.name != "test_pipelined_kv_transfer_nixl_python_accuracy":
continue
params = []
for decorator in node.decorator_list:
if not (
isinstance(decorator, ast.Call)
and isinstance(decorator.func, ast.Name)
and decorator.func.id == "parametrize_with_ids"
):
continue
name = ast.literal_eval(decorator.args[0])
values = ast.literal_eval(decorator.args[1])
params.append((name, values))
expected = [
"[" + "-".join(f"{name}={value}" for name, value in combination) + "]"
for combination in itertools.product(
*[[(name, value) for value in values] for name, values in params]
)
]
found.append((node.lineno, [name for name, _ in params], expected))
for lineno, names, expected in sorted(found):
print(f"line {lineno}: decorators top-to-bottom: {names}")
print("expected:", expected)
print("listed:", listed)
print("listed IDs matching any expected target ID:", [
item for item in listed
if any(item in expected for _, _, expected in found)
])
print("listed IDs not matching any expected target ID:", [
item for item in listed
if not any(item in expected for _, _, expected in found)
])
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 1376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "--- test-list files in the repository ---"
git ls-files 'tests/integration/test_lists/**' | sed -n '1,120p'
echo "--- CBTS artifacts ---"
git ls-files | rg '(^|/)(cbts_touchmap\.sqlite|.*CBTS.*|.*cbts.*coverage.*)' || true
echo "--- relevant list entries ---"
rg -n -C1 'test_pipelined_kv_transfer_nixl_python_accuracy' \
tests/integration/test_lists/test-db/l0_dgx_b200.ymlRepository: NVIDIA/TensorRT-LLM
Length of output: 6528
Reverse the parameter order in the five entries. Pytest generates IDs with enable_block_reuse before disable_overlap_scheduler, so the current IDs select no tests. Use [enable_block_reuse=<value>-disable_overlap_scheduler=<value>].
Test-list coverage: tests/integration/test_lists/test-db/l0_dgx_b200.yml adds five entries. No QA list changes. The coverage verdict is needs follow-up because no CBTS touchmap or 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,
Reverse the parameter order in all five parametrized test IDs under the
Disaggregated Serving section, placing enable_block_reuse before
disable_overlap_scheduler while preserving each value combination and test
target.
| with pytest.raises( | ||
| ValueError, | ||
| match="enable_chunked_prefill is required when enable_pipelined_transfer is set.", | ||
| ): | ||
| create_kv_cache_transceiver( |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Escape the regex metacharacters in the three pytest.raises(match=...) patterns.
Ruff reports RUF043 at lines 680, 706, and 795. Each pattern contains ., which re.search treats as "any character". The assertions therefore pass on messages that differ from the intended text. Use re.escape or a raw string with escaped dots.
🔧 Proposed fix
+import re
@@
with pytest.raises(
ValueError,
- match="enable_chunked_prefill is required when enable_pipelined_transfer is set.",
+ match=re.escape(
+ "enable_chunked_prefill is required when enable_pipelined_transfer is set."
+ ),
):
@@
with pytest.raises(
ValueError,
- match="pipeline_parallel_size=1 is required when enable_pipelined_transfer is set.",
+ match=re.escape(
+ "pipeline_parallel_size=1 is required when enable_pipelined_transfer is set."
+ ),
):
@@
with pytest.raises(
ValueError,
- match="schedule_style must be generation_first when enable_pipelined_transfer is set.",
+ match=re.escape(
+ "schedule_style must be generation_first when enable_pipelined_transfer is set."
+ ),
):Also applies to: 704-708, 793-797
🧰 Tools
🪛 Ruff (0.16.1)
[warning] 680-680: Pattern passed to match= contains metacharacters but is neither escaped nor raw
(RUF043)
🤖 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 678 -
682, Escape the literal periods in the pytest.raises match patterns at the three
affected assertions around create_kv_cache_transceiver, including the cases near
lines 680, 706, and 795. Use re.escape or regex-escaped raw strings so each
assertion matches the exact intended error message rather than treating dots as
wildcards.
Source: Linters/SAST tools
fredricz-20070104
left a comment
fredricz-20070104
left a comment
There was a problem hiding this comment.
Review summary - CONCERNS
Verdict: The change is careful and unusually well-tested, but it cannot merge as-is (GitHub reports mergeable_state: dirty — resolve conflicts first), and the unconditional wire-format widening is an operational regression worth an explicit decision.
Concerns
-
[MAJOR]
tensorrt_llm/_torch/disaggregation/native/transfer.py:~200-_KV_RESULT_PREFIXwidened unconditionally- What is wrong: The KV_AGENT_RESULT frame struct changed from
<qqq?Bqto<qqqq?Bq(newreceiver_slice_id). Your own comment notes there is no version negotiation and both servers must run matching builds. This new layout is emitted for every transfer, not only whenenable_pipelined_transferis set. - How it fails: In a rolling upgrade, an old-build generation server receiving a result from a new-build context server (or vice versa) unpacks a 4-quad frame with a 3-quad struct.
receiver_slice_id/is_last/status/transfer_sizeshift, so the receiver resolves the wrong task or trips the slice-count assert — a silent wrong-offset write or a hang, affecting all disagg users, not just those opting into pipelining. - Suggested fix: If lock-step ctx+gen upgrades are the intended contract, state that explicitly in the PR/release notes and deploy tooling. Better: gate the extra field behind a negotiated capability or add a frame-version byte so old/new peers interoperate.
- What is wrong: The KV_AGENT_RESULT frame struct changed from
-
[MAJOR-adjacent risk] V1 vs V2 gap-dropping semantics differ (
cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:4633,tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py:215)- What is wrong: V2 drops everything up to the last
BAD_PAGE_INDEXgap; the C++ V1 path drops only front-detached blocks (getNumFrontBlocksRemoved). The abstract contract says "blocks at or before any gap are dropped." - How it fails: If any V1 SWA lifecycle that pins sink tokens can leave an interior gap in the raw block table, V1 would return stale interior blocks, violating the contiguous-tail contract and writing KV to the wrong offsets silently. The added C++ test only covers front eviction.
- Suggested fix: Confirm V1 can never produce an interior gap for this query, and add a C++ test asserting it — or make V1 apply the same last-gap rule as V2.
- What is wrong: V2 drops everything up to the last
Minor notes (non-blocking)
tensorrt_llm/_torch/disaggregation/native/transfer.py:~1358-TxSession.statusiteratesself.kv_tasksoutsideself.lockwhilesend()appends under it; snapshot under the lock or document the eventual consistency.tensorrt_llm/_torch/disaggregation/transceiver.py:~720- thecompletedpath also setspy_kv_send_session_retired=Truevia_retire_send_session; the flag name/docstring imply failure, so a cleanly completed request ends up flagged 'retired'. Rename or restrict to the cancel/fail paths.
QA view
- Test coverage: adequate - session state machine, chunk projection,
_build_prefill_chunk(incremental source, SWA trim, prefix reuse), config validation, and the C++ range query are all covered. Uncovered: V1 interior-gap handling, and old/new wire-format interop. - SM coverage: architecture-independent (KV-transfer / cache-manager logic, no arch-guarded kernels or fp8/nvfp4). Integration tests run on Hopper (
skip_pre_hopper) and B200 (l0_dgx_b200). - Test code: generally strong (real Tx/Rx sessions with stubs). Minor: white-box construction via
object.__new__couples tests to private internals; accuracy tests hardcode caps; no interop-regression test for the frame change. - Test time: significant - five new 2-GPU disaggregated GSM8K accuracy cases added to
l0_dgx_b200.yml(each launches ctx+gen servers) plus ~2300 lines of new unit tests and 120s-timeout multi-config transfer-worker tests. - Needs
/qa-verify: yes - new multi-GPU disagg feature with a cluster-wide wire-format change; QA should confirm the B200 accuracy cases, validate the Gemma3 VSWA result, and verify/document that mixed old+new builds are never co-deployed.
What I could not verify
- Whether V1's SWA + sink-token lifecycle can create an interior gap in the raw block table (would make the C++/Python divergence a real corruption path).
- The definition of
token_rangefor the monolithic_create_kv_slice(defined above the diff window) — theslice_end = task._prompt_lensubstitution assumes it equals the oldtoken_range.endfor full transfers. - Runtime behaviour of the new
dispatch_lock/locknesting under real concurrent late-peer replay beyond the provided threaded unit test.
Automated review by NVCortex Lite, run by @fredricz-20070104.
fredricz-20070104
left a comment
fredricz-20070104
left a comment
There was a problem hiding this comment.
Review summary - Approve (non-blocking)
Approving so this is not blocked on me. The points raised in my review comment above are non-blocking — please read them and address what you agree with before merging.
Worth doing before this is relied on: New multi-GPU disagg feature with an unconditional ctx<->gen wire-format change: QA should confirm the B200 accuracy cases pass, validate the Gemma3 VSWA correctness result, and explicitly verify (or document as unsupported) that mixed old/new builds are not deployed together, since the frame layout changed with no negotiation.
Automated review by NVCortex Lite, run by @fredricz-20070104.
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Two whole-PR points before the inline comments:
The diff contains the base pipelined-transfer feature, not just the bounded query. This PR stacks on #15727 (unmerged), so as filed it ships enable_pipelined_transfer, the KV_AGENT_RESULT wire-format change, session retirement, and the executor gating — none of which the description or the [perf] title covers. If #15727 lands first this resolves itself; otherwise please either rebase or expand the description so the merged commit's description matches what it ships.
enable_pipelined_transfer is a new user-facing flag with no docs. docs/source/features/disagg-serving.md doesn't mention it, and its constraints are non-obvious (NIXL + Python transceiver only, requires enable_chunked_prefill, schedule_style: generation_first, sender pp=1, no Mamba/hybrid, incompatible with kv_cache_bounce_size_mb). A short section listing those, plus an example config, would save users a startup-error scavenger hunt. The golden manifest is regenerated — good — but note the new field will need telemetry/privacy CODEOWNER sign-off.
One coverage gap worth a follow-up rather than a blocker: both new accuracy tests pin use_kv_cache_manager_v2: False, so the V2 chunked source path (_CacheReuseAdapterV2.get_block_ids_range) has unit coverage only. A V2 variant of the Llama test would close that.
The range-contract work itself is in good shape — the contract is documented at every layer that enforces it, and the C++ test covering recycled front-detached SWA block IDs targets exactly the failure mode that would silently leak another request's KV.
| raise ValueError( | ||
| "enable_chunked_prefill is required when enable_pipelined_transfer is set." | ||
| ) | ||
| is_kv_cache_sender = getenv("TRTLLM_DISAGG_ROLE") != "generation" |
There was a problem hiding this comment.
This gate derives sender/receiver role from TRTLLM_DISAGG_ROLE, which is only set by the trtllm-serve disaggregated MPI launcher (serve.py sets it "for telemetry"). Workers launched standalone — separate trtllm-serve processes behind a disagg proxy, the common k8s/dynamo pattern — never see it, so a generation server with pp_size>1 and enable_pipelined_transfer inherited from a shared cache_transceiver_config is misclassified as a sender and fails startup, even though the pp=1 restriction only applies to the context side. Fail-closed is the right default, but the error message should name the escape hatch (set TRTLLM_DISAGG_ROLE=generation, or drop the flag from the gen server's config) so standalone deployments aren't stuck guessing.
| 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( |
There was a problem hiding this comment.
Nothing on the session records "the final slice has been submitted": between chunk k completing and chunk k+1 being sent, every registered task is TRANSFERRED, so status/is_completed() would read complete mid-prefill — and check_context_transfer_status would then retire the session, which permanently drops the peer registration. What actually prevents this is indirect: pipelining requires generation_first, which sets _need_aux=True, and the aux task isn't created until _finalize_send, so is_completed() stays false. But wait_complete(blocking=True) skips a None aux task entirely (line ~1489) and would return COMPLETED mid-pipeline — currently unreachable only because the sole production caller never passes at_least_request_num=None. Suggest making the invariant explicit: a _last_slice_submitted flag set in send() when slice.is_last_slice, checked in both is_completed() and the blocking wait path. That protects future relaxations (e.g. context_first pipelining) from silent mid-pipeline retirement.
| # | ||
| # This layout is the ctx<->gen wire format and has no version negotiation: the receiver unpacks | ||
| # whatever arrives against its compiled-in struct. Both servers must run matching builds. | ||
| _KV_RESULT_PREFIX = struct.Struct("<qqqq?Bq") |
There was a problem hiding this comment.
Widening _KV_RESULT_PREFIX from <qqq?Bq to <qqqq?Bq means a mixed-version ctx/gen pair fails as a struct.error swallowed by the listener's per-message try/except: the result is dropped, and the receiver's task hangs until the transfer timeout, with no hint that it's a version skew. Since the format is being broken anyway, this is the cheap moment to prepend a one-byte format version and log a clear "peer wire-format mismatch" error on mismatch — retrofitting versioning after another incompatible change is much harder.
| # Neither V2 backend has a bounded query, so read the whole aggregated | ||
| # list -- keeping the placeholders, which is what makes an entry's index | ||
| # its block ordinal -- and cut the range out of it here. | ||
| all_block_ids = np.fromiter( |
There was a problem hiding this comment.
The V2 path still materializes the full aggregated page list per chunk, so per-request block-ID work here remains O(num_chunks × prompt_blocks) per layer group — the O(prompt_blocks × L)-per-request bound in the PR description holds only for the C++ V1 manager. The comment is honest about the missing bounded query, but for the 128k-prompt/128-chunk case that motivated this PR the quadratic term survives on V2. Worth a tracked follow-up to add a bounded range query to KVCacheManagerV2/kv_cache_map (and worth noting the V1-only scope in the PR description).
|
|
||
| # Pipelined transfer records the transfer start time of the last slice. | ||
| # The records of the previous slices are overwritten. | ||
| req.set_kv_cache_transfer_start(tensorrt_llm.bindings.global_steady_clock_now()) |
There was a problem hiding this comment.
set_kv_cache_transfer_start is overwritten on every chunk, so the request-level transfer-time metric measures only the last slice's window rather than the whole pipelined transfer (session transfer_start_time keeps first-send semantics, so the two now disagree). If measuring the exposed (non-overlapped) tail is the intent, say so in the comment; otherwise guard with if req.get_kv_cache_transfer_start() is None-style logic so perf dashboards remain comparable between pipelined and monolithic runs.
| //! always the *contiguous* run ending at blockEnd, i.e. ordinals [blockEnd - result.size(), blockEnd), so a caller | ||
| //! can recover each id's block ordinal from blockEnd and the result size alone. Requesting a blockEnd past the | ||
| //! sequence's allocated blocks would break that guarantee and is rejected. | ||
| [[nodiscard]] std::vector<std::vector<SizeType32>> getCacheBlockIdsRange( |
There was a problem hiding this comment.
We are switching to KV-Cache Manager V2 very soon. Could you check if the same feature needs to be applied to KV-Cache Manager V2 as well?
5 participants
Description
Summary
Follow-up to the pipelined prefill-transfer PR. Building a chunk's
KVSliceused torequire materializing every block ID the request owns and then throwing away all but
the current chunk's slice of it. This PR gives the cache managers a bounded query so
_build_prefill_chunkfetches only the blocks the chunk actually sends.The cost of the old path scaled with the prompt, once per chunk. For a request cut into
Nchunks withLlayer groups the sender didO(N × prompt_blocks × L)block-ID work:with
tokens_per_block=32, a 128k-token prompt in 1024-token chunks rebuilt a4096-entry list per layer group, 128 times over, to send 32 blocks each time. It is now
O(prompt_blocks × L)for the whole request, with each chunk's work proportional to thechunk. This runs on the executor thread between forward steps, which is exactly the time
pipelining is trying to spend on compute.
Behavior on the wire is unchanged: the same block IDs are sent, in the same order, with
the same slice metadata.
Before / after
project_blocks_to_global_chunkis still used for the destination projection in_build_kv_write_meta; only the source side stops calling it._create_kv_slicenow hasone job — the whole-prompt monolithic slice — so its
resident_block_endparameter isgone.
The range contract
Every implementation returns the contiguous run of resident blocks ending at
block_end, i.e. ordinals[block_end - len(result), block_end). The result may beshorter than the requested range, never longer, and never has holes.
This is not a stylistic choice — it is what the rest of the transfer path already
assumes. The receiver reconstructs each layer group's starting token from
block_endminus the number of blocks it received, so a result that is not a contiguous tail would
be written to the wrong offsets, silently, with no error anywhere. Two consequences fall
out of it:
sink tokens leaves a sink prefix plus a window suffix; returning both would misreport
where the suffix starts.
block_endpast the request's allocated blocks is rejected rather than clamped,because clamping would return a run that does not end where the caller thinks it does.
_build_prefill_chunkenforces the same contract from the caller's side: a result longerthan the requested range raises, and a full-attention group that is not fully resident
raises rather than sending a short chunk.
Sliding-window layer groups
A windowed group's range starts at the request's final-window stale boundary, not at
the chunk start:
The boundary is a property of the whole request (it uses
prompt_len, and the+1covers the first generated token), so it is the same for every chunk. It has to be
applied here because of what the sender worker does downstream:
_build_kv_write_metaraisessrc_startto the stale boundary without dropping thesource blocks below it. If the slice carries blocks that reach further back than the
boundary, the raised
src_startpairs the run's head with the window tail's destinationblocks and the generation server decodes over KV belonging to other tokens. Under the old
code the trim happened implicitly, inside the full-slice construction that this PR
removes; it is now explicit.
A chunk that lies entirely below the final window contributes nothing, so the cache is not
queried at all — the common case for a short window and a long prompt.
Gemma3-1B exercises this directly: five of every six layers are windowed at 512 tokens,
so most chunks of a GSM8K prompt sit partly or entirely below the boundary. Without the
trim it scores 13.27 against a 25.52 reference; with it, 28.17.
Implementation
CacheReuseAdapter.get_block_ids_range(new abstract method)KVCacheManager.get_cache_indices_range→ nanobindget_cache_block_ids_range→BaseKVCacheManager::getCacheBlockIdsRangeget_block_idsdoes (the two diverge once host offload is enabled)get_aggregated_page_indices(valid_only=False), sliced and trimmed in Pythoncpp— the default — orpy) has a bounded query, so the range is cut out here. Keeping theBAD_PAGE_INDEXplaceholders is what makes an entry's index its block ordinalBound validation is aligned across backends: negative or reversed bounds raise
ValueErrorfrom both, and an out-of-rangeblock_endraises from both (ValueErrorfrom V2,
RuntimeErrorout of the C++ manager for V1).C++ (
BaseKVCacheManager::getCacheBlockIdsRange)A non-virtual convenience wrapper, so every V1 manager gets it. It reads the block table
and the front-eviction count off the same sequence object, so a manager that overrides
one but not the other cannot hand back a table and an eviction count that disagree — the
raw table keeps recycled IDs in front-detached SWA slots, which would otherwise look like
live blocks. The nanobind wrapper deliberately keeps the GIL, unlike its neighbours: it
reaches the virtual
getSequence, whose trampoline can call back into a Python subclass,and the copy it makes is bounded by the requested range.
KVCacheManager(V1 Python)get_cache_indices_rangeresolves the layer's window size through a new shared_resolve_cache_window_sizehelper and rejectsbeam_width > 1, which chunked transferdoes not support. The layer-offset → window-size mapping is now computed once at
construction (
_window_size_by_layer_offset) instead of being rebuilt on every call;get_batch_cache_indicesuses the same helper, so the two cannot drift.Test Coverage
TestBlockRangeAdaptersis intests/unittest/disaggregated/test_cache_reuse_adapter.py;the
test_build_prefill_chunk_*cases are intests/unittest/disaggregated/test_kv_transfer.py.TestBlockRangeAdapters::test_v1_adapter_requests_only_block_rangeTestBlockRangeAdapters::test_v1_adapter_skips_translation_for_empty_rangeTestBlockRangeAdapters::test_v2_range_honors_block_beginTestBlockRangeAdapters::test_v2_range_drops_leading_gap/..._drops_everything_before_an_interior_gap/..._is_empty_when_the_last_block_is_absent/..._ignores_gaps_outside_the_rangeTestBlockRangeAdapters::test_v2_range_rejects_invalid_bounds/..._rejects_block_end_past_allocationTestBlockRangeAdapters::test_v1_manager_rejects_invalid_bounds/..._resolves_layer_window_once_and_forwards_range/..._rejects_multiple_beamstest_build_prefill_chunk_rejects_short_full_attention_range/..._rejects_overlong_range/..._accepts_short_swa_rangetest_build_prefill_chunk_skips_windowed_group_before_final_windowtest_build_prefill_chunk_empty_range_skips_cache_queriestest_build_prefill_chunk_normalizes_swa_source_to_computed_prefix(block_begin, block_end)requestedKVCacheManagerTest.VSWAGetCacheBlockIdsRangeExcludesDetachedFrontBlocks_build_prefill_chunk_forintest_chunked_transfer.pywas driving a stubbed_create_kv_slice, so the prefix-reuse tests built on it had stopped exercising the realpath; it now goes through the range API. The V2 range-semantics tests moved from the
Python
_KVCacheto the adapter, so they cover whichever backendTLLM_KV_CACHE_MANAGER_V2_BACKENDselects.Validation
tests/unittest/disaggregated/test_kv_transfer.py(76 tests, includes the pipelinedGPU harness cases),
test_chunked_transfer.py+test_cache_reuse_adapter.py(136 tests),
test_bounce.py,test_mamba_cache_manager.py— all passTestGemma3_1BInstruct::test_pipelined_kv_transfer_nixl_python_accuracy— GSM8K 28.17vs 25.52 reference
pre-commitcleanPR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.Dev Engineer Review
get_block_ids_range.getCacheBlockIdsRange.CODING_GUIDELINES.mdand full CI results before merge.QA Engineer Review
tests/unittest/disaggregated/test_cache_reuse_adapter.py.tests/unittest/disaggregated/test_chunked_transfer.py.tests/unittest/disaggregated/test_kv_transfer.py.tests/unittest/disaggregated/test_disagg_utils.py.tests/unittest/disaggregated/test_bounce.py.tests/integration/test_lists/test-db/l0_dgx_b200.yml.