Repository navigation
[TRTLLM-9523][feat] Adapting the transceiver to manager v2 (step 6) - #11978
Merged
Merged
Conversation
Shixiaowei02
requested review from
Naveassaf,
brb-nv,
chuangz0,
schetlur-nv and
shaharmor98
March 6, 2026 06:41
Contributor
📝 WalkthroughWalkthroughThis pull request introduces a major refactoring of KV cache management and disaggregated serving infrastructure. Key changes include: (1) replacing legacy page table structures with a new KVCachePageTable model supporting per-layer-group pools; (2) introducing VRAM VMM descriptor splitting for NIXL transfers; (3) restructuring RankInfo to embed AttentionInfo and page_table; (4) adding new region mapper implementations for attention and SSM disaggregation; and (5) implementing multi-layer-group KV slice handling throughout the transfer and cache transceiver layers. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes 🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 16
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp (1)
2-3:⚠️ Potential issue | 🟡 MinorUpdate SPDX copyright year to include 2026.
Line 2 still ends at
2025, but this file is modified in 2026. Please update it to include 2026 (e.g.,2022-2026).As per coding guidelines: "All TensorRT-LLM source files should contain an NVIDIA copyright header with the year of the latest meaningful modification."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp` around lines 2 - 3, Update the SPDX copyright header year range in this file by changing the trailing year from 2025 to 2026 (e.g., update "2022-2025" to "2022-2026") so the header on the top lines that start with "SPDX-FileCopyrightText" and "SPDX-License-Identifier" reflects the 2026 modification.tensorrt_llm/_torch/pyexecutor/resource_manager.py (1)
2148-2156:⚠️ Potential issue | 🔴 Critical
KVCacheManagerV2.get_batch_cache_indicesdefault layer selection is unsafe across PP ranks.Line 2150 defaults
layer_idto0, but Line 2155 indexesself.layer_offsets[layer_id]using global layer IDs local to the current PP rank. On non-first PP ranks, global layer0may not exist, causing aKeyErrorin default call paths (e.g., Line 2000).💡 Proposed fix
- def get_batch_cache_indices(self, - request_ids: List[int], - layer_id: int = 0) -> List[List[int]]: + def get_batch_cache_indices( + self, + request_ids: List[int], + layer_id: int | None = None, + ) -> List[List[int]]: + + if layer_id is None: + if not self.pp_layers: + return [[] for _ in request_ids] + layer_id = self.pp_layers[0] + if layer_id not in self.layer_offsets: + raise ValueError( + f"layer_id {layer_id} is not local to this PP rank; local global layers: {self.pp_layers}" + ) + local_layer_id = self.layer_offsets[layer_id] return self._get_batch_cache_indices_by_pool_id( request_ids, - pool_id=self.layer_to_pool_mapping_dict[ - self.layer_offsets[layer_id]], + pool_id=self.layer_to_pool_mapping_dict[local_layer_id], is_kv_aggregate=True)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py` around lines 2148 - 2156, KVCacheManagerV2.get_batch_cache_indices currently defaults layer_id=0 but then treats it as a global layer ID when indexing self.layer_offsets and layer_to_pool_mapping_dict, which breaks on non-first PP ranks; change the signature to accept layer_id: Optional[int] = None (or keep int but treat 0 as local index) and when layer_id is None set it to the local-first-layer index (0 as a local index), then map to the global layer via self.layer_offsets[layer_id_local] before looking up pool_id; also add a bounds check that layer_id_local is within len(self.layer_offsets) and raise/handle a clear error if out of range; reference KVCacheManagerV2.get_batch_cache_indices, self.layer_offsets, self.layer_to_pool_mapping_dict, and _get_batch_cache_indices_by_pool_id when making the change.
🧹 Nitpick comments (6)
tensorrt_llm/_torch/attention_backend/trtllm.py (1)
630-631: Consider usingtorch.ones_like()for clarity and efficiency.The expression
0 * self.sparse_attn_indices + 1is unconventional. Usingtorch.ones_like(self.sparse_attn_indices)would be more readable and avoids creating two intermediate tensors.Additionally, the WAR comment could benefit from a reference to a tracking issue or more detail about what constitutes "illegal indices" for future maintainers.
♻️ Suggested improvement
# WAR on potential illegal sparse attention top_k indices -self.sparse_attn_indices = 0 * self.sparse_attn_indices + 1 if self.sparse_attn_indices is not None else None +self.sparse_attn_indices = torch.ones_like(self.sparse_attn_indices) if self.sparse_attn_indices is not None else None🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/attention_backend/trtllm.py` around lines 630 - 631, Replace the odd expression "0 * self.sparse_attn_indices + 1" with torch.ones_like(self.sparse_attn_indices) to make the intent clear and avoid extra temporaries; update the preceding WAR comment in trtllm.py near self.sparse_attn_indices to briefly state what "illegal sparse attention top_k indices" means (e.g., out-of-range or negative indices) and, if available, add a reference to the related tracking issue or ticket for future context.tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py (1)
278-284: Guard TP-ratio math for non-divisible TP sizes.The ratio currently uses floor division. If TP sizes are not divisible, offsets become silently incorrect.
💡 Proposed hardening
def _compute_tp_offsets( @@ ) -> tuple: @@ - ratio = max(self_tp_per_dp, peer_tp_per_dp) // min(self_tp_per_dp, peer_tp_per_dp) + larger = max(self_tp_per_dp, peer_tp_per_dp) + smaller = min(self_tp_per_dp, peer_tp_per_dp) + if larger % smaller != 0: + raise ValueError( + f"Incompatible TP sizes: self_tp_per_dp={self_tp_per_dp}, peer_tp_per_dp={peer_tp_per_dp}" + ) + ratio = larger // smaller🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py` around lines 278 - 284, The TP-ratio math using floor division produces incorrect offsets when TP sizes are not divisible; in the block that computes ratio and returns offsets (variables: ratio, self_tp_per_dp, peer_tp_per_dp, transfer_bytes, self_tp_rank, peer_tp_rank), replace the simple floor-division scheme with proper uneven-chunk math: treat the larger group as split into smaller_group-sized buckets with base_chunk = larger // smaller and extra = larger % smaller, compute the selected bucket index (peer_tp_rank % smaller or self_tp_rank % smaller as currently chosen), then compute offset_bytes = (bucket_index * base_chunk + min(bucket_index, extra)) * transfer_bytes so buckets 0..extra-1 get base_chunk+1 elements and the rest get base_chunk, and return offsets using that computed offset_bytes for the appropriate side instead of ratio * transfer_bytes.tests/unittest/disaggregated/test_peer.py (1)
49-52: Consider addingstrict=Truetozip()for safer iteration.The
zip()call pairspool_ptrsandblock_bytes, but if they have mismatched lengths, silent truncation could mask test setup bugs.♻️ Proposed fix
physical_pools = [ PhysicalPool(base_address=ptr, slot_bytes=bs, num_slots=128) - for ptr, bs in zip(pool_ptrs, block_bytes) + for ptr, bs in zip(pool_ptrs, block_bytes, strict=True) ]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/disaggregated/test_peer.py` around lines 49 - 52, The list comprehension that builds physical_pools using zip(pool_ptrs, block_bytes) should use zip(..., strict=True) to ensure lengths match and avoid silent truncation; update the comprehension that constructs PhysicalPool(base_address=ptr, slot_bytes=bs, num_slots=128) to iterate with zip(pool_ptrs, block_bytes, strict=True) so mismatched lengths raise immediately.tensorrt_llm/_torch/disaggregation/native/transfer.py (1)
1342-1343: Unusedmessageparameter.The
messageparameter is not used in_process_request_rank_info. Consider prefixing with underscore to indicate it's intentionally unused.♻️ Proposed fix
- def _process_request_rank_info(self, send_id: bytes, message: list[bytes]): + def _process_request_rank_info(self, send_id: bytes, _message: list[bytes]): self._messenger.send([send_id, self._rank_info.to_bytes()])🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/disaggregation/native/transfer.py` around lines 1342 - 1343, The parameter "message" in _process_request_rank_info is unused; rename it to "_message" (or prefix with an underscore) in the function signature to indicate it's intentionally unused and silence lint warnings. Update the signature for def _process_request_rank_info(self, send_id: bytes, message: list[bytes]): to use the underscored name (e.g., _message: list[bytes]) while leaving the body that calls self._messenger.send([send_id, self._rank_info.to_bytes()]) unchanged.tests/unittest/disaggregated/test_cache_transceiver_single_process.py (1)
531-531: Unused parameteruse_v2in_get_block_data_for_layer.The
use_v2parameter is declared but never used in the function body. Either remove it or use it for V1/V2-specific logic if needed.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/disaggregated/test_cache_transceiver_single_process.py` at line 531, The helper function _get_block_data_for_layer declares an unused parameter use_v2; remove use_v2 from the function signature and from all its call sites (or alternatively implement V1/V2 branching inside _get_block_data_for_layer using use_v2 to choose the correct manager method), ensuring you update references to request_id, layer_idx, expected_valid and any mgr calls accordingly so no unused parameter remains.tensorrt_llm/_torch/disaggregation/native/peer.py (1)
298-299: Consider exposing_tp_per_dpas a public method or property.The code accesses
self._attention_policy._tp_per_dp(...)which is a private method (prefixed with_). While this works, it couples this class to internal implementation details. Consider either making_tp_per_dppublic (rename totp_per_dp) or exposing a wrapper method onAttentionPolicy.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/disaggregation/native/peer.py` around lines 298 - 299, The code calls the private method self._attention_policy._tp_per_dp(...); change the API to avoid using a private member by either renaming AttentionPolicy._tp_per_dp to a public method tp_per_dp or adding a public wrapper on AttentionPolicy (e.g., def tp_per_dp(self, ri): return self._tp_per_dp(ri)), and then update this file to call self._attention_policy.tp_per_dp(peer_ri) (and any other call sites) so the class no longer depends on a private implementation detail.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tensorrt_llm/_torch/disaggregation/native/auxiliary.py`:
- Around line 131-143: free_slot currently removes the slot from _occupied_slots
and returns it to _free_slots but does not clear per-slot metadata, allowing
stale token-length/width reads; update free_slot to reset all per-slot metadata
(e.g., the per-slot token lengths/widths arrays and any related buffers
referenced elsewhere) for the given slot after validation, and also ensure
alloc_slot initializes those metadata entries (the code that defaults to full
widths at the reader path around the logic referenced at line ~177 should find
zeroed/initialized metadata rather than missing stale values). Target symbols:
free_slot, _occupied_slots, _free_slots, and the per-slot metadata structures
(token length/width arrays) so that freeing a slot zeroes or clears its metadata
and allocation reinitializes it.
In `@tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py`:
- Line 1: Add the required NVIDIA Apache-2.0 copyright/license header to the top
of the new Python file
tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py: insert the
standard NVIDIA Apache-2.0 header block (with the correct copyright year equal
to the latest meaningful modification year) before any imports (i.e., before the
existing "import numpy as np"), matching the header style used across other
TensorRT-LLM source files.
- Around line 151-152: The current computation for self_tp_per_dp and
peer_tp_per_dp in peer.py (variables self_tp_per_dp and peer_tp_per_dp)
incorrectly divides tp_size by dp_size and should instead use the
attention-aware per-DP TPU grouping value; replace those divisions with the
RankInfo field that encodes TP-per-DP for attention (e.g., use
self_ri.tp_size_per_dp_group and peer_ri.tp_size_per_dp_group or the existing
RankInfo.tp_size_per_dp_group accessor) so offsets and head-mismatch mapping
follow the attention DP semantics when attention DP is disabled or differs
between ranks.
In `@tensorrt_llm/_torch/disaggregation/native/mixers/attention/spec.py`:
- Around line 1-23: This new Python file lacks the required NVIDIA Apache-2.0
source header; add the standard NVIDIA copyright + Apache-2.0 license block at
the top of the file (using the year of the latest meaningful modification)
before any imports so it applies to the whole module; ensure the header appears
above the existing import and the AttentionInfo class and that formatting
matches other project source files' header block exactly.
In `@tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py`:
- Line 1: This file is missing the required NVIDIA Apache-2.0 license header;
add the standard NVIDIA Apache-2.0 copyright/license block (with the year of the
latest meaningful modification) at the top of the file before any imports (i.e.,
above the existing "from typing import List" line), ensuring it matches the
project's canonical header format used across other .py sources.
- Around line 111-120: After slicing src_group.ptrs and dst_group.ptrs using
self._src_layer_off/self._dst_layer_off and self._transfer_layers, validate that
the resulting src_layer_ptrs and dst_layer_ptrs have the same length before
building new_src_ptrs/new_dst_ptrs; if they differ, either adjust the effective
transfer count to min(len(src_layer_ptrs), len(dst_layer_ptrs)) or raise a clear
ValueError (or log and abort) so you don't construct mismatched mapped pointers
— update the logic around src_group.ptrs, dst_group.ptrs, self._src_layer_off,
self._dst_layer_off, self._transfer_layers, new_src_ptrs and new_dst_ptrs
accordingly.
In `@tensorrt_llm/_torch/disaggregation/native/py_cache_transceiver.py`:
- Around line 139-145: The sliding-window trimming mistakenly uses
block_ids[-expected_valid:] which is a no-op when expected_valid == 0; update
the logic where window_size is handled (the block trimming around variables
total_blocks, stale_end, expected_valid and block_ids) to handle the zero case
explicitly — only slice with negative index when expected_valid > 0, otherwise
set block_ids to an empty list (or otherwise clear stale entries) so no stale
pages are retained.
In `@tensorrt_llm/_torch/disaggregation/native/transfer.py`:
- Around line 1304-1306: The constructor for RankInfoServer has parameters addr
and port defaulting to None but typed as str and int; update the type hints on
RankInfoServer.__init__ to use Optional[str] for addr and Optional[int] for port
(or use typing.Union[str, None] / Union[int, None]) and ensure typing.Optional
is imported at the top of the module if not already present so the annotations
match the default None values.
In `@tensorrt_llm/_torch/disaggregation/resource/page.py`:
- Line 1: Add the required NVIDIA Apache-2.0 license header at the very top of
this source file (before the existing "from __future__ import annotations" line)
with the appropriate NVIDIA copyright notice and the year of the latest
meaningful modification; ensure the header matches the project's canonical
Apache-2.0 block used across TensorRT-LLM Python sources so the file (page.py)
complies with the repository coding guidelines.
In `@tensorrt_llm/_torch/disaggregation/resource/utils.py`:
- Around line 31-33: The code currently only checks the upper bound of slot_id
before computing pool.base_address + slot_id * pool.slot_bytes, allowing
negative slot_id to produce addresses below the base; add a lower-bound check
(if slot_id < 0: raise ValueError(...)) alongside the existing upper-bound check
so both out-of-range negatives and too-large indices are rejected—update the
same validation at both places where pool.base_address + slot_id *
pool.slot_bytes is computed (the occurrence around the first check and the
separate occurrence later) to raise a clear ValueError mentioning the invalid
slot_id and pool.num_slots.
- Line 1: Add the required NVIDIA Apache-2.0 copyright header block to the top
of this source file (place it before the existing "from __future__ import
annotations" line), using the NVIDIA-prescribed Apache-2.0 license text and the
year of the latest meaningful modification; ensure the header matches formatting
used across other TensorRT-LLM Python files and does not alter the existing
imports or code order.
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py`:
- Around line 155-156: The abstract base class method
KvCacheTransceiver.shutdown currently only contains a docstring which triggers
Ruff B027; give it an explicit body to silence the linter by adding a no-op
implementation (e.g., include a pass statement or simple return) inside the
shutdown method so it remains a default no-op, or alternatively decorate it with
`@abstractmethod` if you intend subclasses to implement it—update the shutdown
method in class KvCacheTransceiver accordingly.
In `@tensorrt_llm/_torch/pyexecutor/py_executor_creator.py`:
- Around line 787-792: Wrap the kv_cache_transceiver.shutdown() call in a
try/except (or try/except/finally) so any exception from
py_executor.kv_cache_transceiver.shutdown() is caught and does not prevent
kv_cache_creator.teardown_managers(resources) from running; ensure teardown is
invoked in a finally block (or after the except) and that del py_executor and
gc.collect() still run regardless of shutdown errors, optionally logging the
shutdown exception for debugging purposes.
In `@tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py`:
- Around line 269-273: The warning currently logs whenever nothing was scheduled
(based only on scheduled_requests and scheduled_disagg_gen_init_requests) which
can be noisy; change the condition so the warning is emitted only when there are
actual pending candidates but none could be scheduled (e.g. compute
pending_candidates_count from whatever list/iterator of candidates used earlier,
or add a resource_failure flag) and require pending_candidates_count > 0 (or
resource_failure == True) in addition to len(scheduled_requests) +
len(scheduled_disagg_gen_init_requests) == 0 before calling logger.warning;
update the block around scheduled_requests / scheduled_disagg_gen_init_requests
accordingly.
In `@tests/unittest/disaggregated/test_cache_transceiver_single_process.py`:
- Around line 725-731: The assertion's msg lambda inside
torch.testing.assert_close captures loop variables req_idx and layer_idx by
reference, causing incorrect values if the lambda is evaluated later; fix it by
binding those variables as default parameters in the lambda (e.g., change
msg=lambda m: ... to msg=lambda m, req_idx=req_idx, layer_idx=layer_idx: ...) at
the torch.testing.assert_close call so the message reflects the correct req_idx
and layer_idx for that assertion.
In `@tests/unittest/disaggregated/test_kv_transfer.py`:
- Around line 404-410: The get_block_data function has request_id defaulting to
None but typed as int, which creates an implicit Optional; update the signature
of get_block_data to annotate request_id as Optional[int] (or Union[int, None])
and import typing.Optional if not already imported, and adjust any internal
usage that assumes a non-None value to handle None where necessary; reference:
function get_block_data and parameter request_id.
---
Outside diff comments:
In `@cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp`:
- Around line 2-3: Update the SPDX copyright header year range in this file by
changing the trailing year from 2025 to 2026 (e.g., update "2022-2025" to
"2022-2026") so the header on the top lines that start with
"SPDX-FileCopyrightText" and "SPDX-License-Identifier" reflects the 2026
modification.
In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py`:
- Around line 2148-2156: KVCacheManagerV2.get_batch_cache_indices currently
defaults layer_id=0 but then treats it as a global layer ID when indexing
self.layer_offsets and layer_to_pool_mapping_dict, which breaks on non-first PP
ranks; change the signature to accept layer_id: Optional[int] = None (or keep
int but treat 0 as local index) and when layer_id is None set it to the
local-first-layer index (0 as a local index), then map to the global layer via
self.layer_offsets[layer_id_local] before looking up pool_id; also add a bounds
check that layer_id_local is within len(self.layer_offsets) and raise/handle a
clear error if out of range; reference KVCacheManagerV2.get_batch_cache_indices,
self.layer_offsets, self.layer_to_pool_mapping_dict, and
_get_batch_cache_indices_by_pool_id when making the change.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/trtllm.py`:
- Around line 630-631: Replace the odd expression "0 * self.sparse_attn_indices
+ 1" with torch.ones_like(self.sparse_attn_indices) to make the intent clear and
avoid extra temporaries; update the preceding WAR comment in trtllm.py near
self.sparse_attn_indices to briefly state what "illegal sparse attention top_k
indices" means (e.g., out-of-range or negative indices) and, if available, add a
reference to the related tracking issue or ticket for future context.
In `@tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py`:
- Around line 278-284: The TP-ratio math using floor division produces incorrect
offsets when TP sizes are not divisible; in the block that computes ratio and
returns offsets (variables: ratio, self_tp_per_dp, peer_tp_per_dp,
transfer_bytes, self_tp_rank, peer_tp_rank), replace the simple floor-division
scheme with proper uneven-chunk math: treat the larger group as split into
smaller_group-sized buckets with base_chunk = larger // smaller and extra =
larger % smaller, compute the selected bucket index (peer_tp_rank % smaller or
self_tp_rank % smaller as currently chosen), then compute offset_bytes =
(bucket_index * base_chunk + min(bucket_index, extra)) * transfer_bytes so
buckets 0..extra-1 get base_chunk+1 elements and the rest get base_chunk, and
return offsets using that computed offset_bytes for the appropriate side instead
of ratio * transfer_bytes.
In `@tensorrt_llm/_torch/disaggregation/native/peer.py`:
- Around line 298-299: The code calls the private method
self._attention_policy._tp_per_dp(...); change the API to avoid using a private
member by either renaming AttentionPolicy._tp_per_dp to a public method
tp_per_dp or adding a public wrapper on AttentionPolicy (e.g., def
tp_per_dp(self, ri): return self._tp_per_dp(ri)), and then update this file to
call self._attention_policy.tp_per_dp(peer_ri) (and any other call sites) so the
class no longer depends on a private implementation detail.
In `@tensorrt_llm/_torch/disaggregation/native/transfer.py`:
- Around line 1342-1343: The parameter "message" in _process_request_rank_info
is unused; rename it to "_message" (or prefix with an underscore) in the
function signature to indicate it's intentionally unused and silence lint
warnings. Update the signature for def _process_request_rank_info(self, send_id:
bytes, message: list[bytes]): to use the underscored name (e.g., _message:
list[bytes]) while leaving the body that calls self._messenger.send([send_id,
self._rank_info.to_bytes()]) unchanged.
In `@tests/unittest/disaggregated/test_cache_transceiver_single_process.py`:
- Line 531: The helper function _get_block_data_for_layer declares an unused
parameter use_v2; remove use_v2 from the function signature and from all its
call sites (or alternatively implement V1/V2 branching inside
_get_block_data_for_layer using use_v2 to choose the correct manager method),
ensuring you update references to request_id, layer_idx, expected_valid and any
mgr calls accordingly so no unused parameter remains.
In `@tests/unittest/disaggregated/test_peer.py`:
- Around line 49-52: The list comprehension that builds physical_pools using
zip(pool_ptrs, block_bytes) should use zip(..., strict=True) to ensure lengths
match and avoid silent truncation; update the comprehension that constructs
PhysicalPool(base_address=ptr, slot_bytes=bs, num_slots=128) to iterate with
zip(pool_ptrs, block_bytes, strict=True) so mismatched lengths raise
immediately.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 01336ef2-57ee-4bd6-b968-4927c23a56cf
📒 Files selected for processing (38)
cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txtcpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cppcpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.hcpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpptensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/disaggregation/base/region.pytensorrt_llm/_torch/disaggregation/base/transfer.pytensorrt_llm/_torch/disaggregation/native/auxiliary.pytensorrt_llm/_torch/disaggregation/native/mixers/__init__.pytensorrt_llm/_torch/disaggregation/native/mixers/attention/__init__.pytensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.pytensorrt_llm/_torch/disaggregation/native/mixers/attention/spec.pytensorrt_llm/_torch/disaggregation/native/mixers/ssm/__init__.pytensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.pytensorrt_llm/_torch/disaggregation/native/peer.pytensorrt_llm/_torch/disaggregation/native/py_cache_transceiver.pytensorrt_llm/_torch/disaggregation/native/rank_info.pytensorrt_llm/_torch/disaggregation/native/region/block.pytensorrt_llm/_torch/disaggregation/native/region/page.pytensorrt_llm/_torch/disaggregation/native/transfer.pytensorrt_llm/_torch/disaggregation/resource/kv_extractor.pytensorrt_llm/_torch/disaggregation/resource/kv_extractor_v2.pytensorrt_llm/_torch/disaggregation/resource/page.pytensorrt_llm/_torch/disaggregation/resource/utils.pytensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.pytensorrt_llm/_torch/pyexecutor/py_executor_creator.pytensorrt_llm/_torch/pyexecutor/resource_manager.pytensorrt_llm/_torch/pyexecutor/scheduler/scheduler.pytests/integration/defs/accuracy/test_disaggregated_serving.pytests/integration/test_lists/test-db/l0_dgx_b200.ymltests/integration/test_lists/test-db/l0_dgx_h100.ymltests/unittest/disaggregated/region/test_block.pytests/unittest/disaggregated/test_cache_transceiver_single_process.pytests/unittest/disaggregated/test_extractor.pytests/unittest/disaggregated/test_extractor_v2.pytests/unittest/disaggregated/test_kv_transfer.pytests/unittest/disaggregated/test_kv_transfer_mp.pytests/unittest/disaggregated/test_peer.py
💤 Files with no reviewable changes (4)
- tests/unittest/disaggregated/test_extractor_v2.py
- tensorrt_llm/_torch/disaggregation/native/region/block.py
- tensorrt_llm/_torch/disaggregation/resource/kv_extractor_v2.py
- tensorrt_llm/_torch/disaggregation/native/region/page.py
chuangz0
reviewed
Mar 6, 2026
Shixiaowei02
changed the title
Shixiaowei02
force-pushed
the
feat/v2_adapt_2
branch
from
March 6, 2026 09:01
27d709d to
0cbedd8
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Shixiaowei02
requested review from
bo-nv,
lowsfer,
pcastonguay,
reasonsolo and
yizhang-nv
March 6, 2026 09:09
Collaborator
|
PR_Github #38001 [ run ] triggered by Bot. Commit: |
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #38004 [ run ] triggered by Bot. Commit: |
Shixiaowei02
force-pushed
the
feat/v2_adapt_2
branch
from
March 6, 2026 09:27
ea47557 to
67a8a3d
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #38006 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #38258 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #38259 [ run ] triggered by Bot. Commit: |
Shixiaowei02
force-pushed
the
feat/v2_adapt_2
branch
from
March 9, 2026 12:18
790123e to
a6a6767
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #38264 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #38264 [ run ] completed with state
|
Shixiaowei02
force-pushed
the
feat/v2_adapt_2
branch
from
March 10, 2026 02:28
a6a6767 to
65afb50
Compare
Co-authored-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com>
Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com>
Fix VMM virtual address alignment in cuMemAddressReserve (pass phys_mem_size as alignment to ensure chunk-aligned pool base), fix NIXL transfer agent init/setup, peer overlap detection, and dst split for kv_cache_manager_v2 compatibility. Co-Authored-By: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com>
Shixiaowei02
force-pushed
the
feat/v2_adapt_2
branch
from
March 10, 2026 02:36
65afb50 to
239eab4
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #38354 [ run ] triggered by Bot. Commit: |
lowsfer
approved these changes
Mar 10, 2026
lowsfer
left a comment
lowsfer
left a comment
Member
There was a problem hiding this comment.
LGTM for the part related to kv cache manager v2.
Collaborator
|
PR_Github #38354 [ run ] completed with state
|
Collaborator
Author
|
/bot run --stage-list "A100X-PyTorch-1, A30-CPP-3" |
Collaborator
|
PR_Github #38441 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #38441 [ run ] completed with state |
Collaborator
Author
|
/bot skip --comment "CI passed at 38354 and 38441 in the assembled form." |
Collaborator
|
PR_Github #38511 [ skip ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #38511 [ skip ] completed with state |
Shixiaowei02
enabled auto-merge (squash)
March 11, 2026 08:08
pcastonguay
approved these changes
Mar 11, 2026
limin2021
pushed a commit
to limin2021/TensorRT-LLM
that referenced
this pull request
Mar 19, 2026
…VIDIA#11978) Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com> Co-authored-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
longcheng-nv
pushed a commit
to longcheng-nv/TensorRT-LLM
that referenced
this pull request
Mar 31, 2026
…VIDIA#11978) Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com> Co-authored-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
1 task done
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
Release Notes
New Features
Improvements
Description
Test Coverage
PR 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)
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.