[None][fix] Filter empty aux buffers from NIXL transfers

cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp

@@ -83,9 +83,9 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m)
         // Classmethod: batch construction from numpy arrays
         .def_static(
             "from_arrays",
-            [](kvc::MemoryType type, nb::ndarray<int64_t, nb::ndim<1>, nb::c_contig, nb::device::cpu> addrs,
-                nb::ndarray<int64_t, nb::ndim<1>, nb::c_contig, nb::device::cpu> sizes,
-                nb::ndarray<int32_t, nb::ndim<1>, nb::c_contig, nb::device::cpu> deviceIds)
+            [](kvc::MemoryType type, nb::ndarray<int64_t const, nb::ndim<1>, nb::c_contig, nb::device::cpu> addrs,
+                nb::ndarray<int64_t const, nb::ndim<1>, nb::c_contig, nb::device::cpu> sizes,
+                nb::ndarray<int32_t const, nb::ndim<1>, nb::c_contig, nb::device::cpu> deviceIds)
             {
                 size_t n = addrs.shape(0);
                 auto const* a = addrs.data();
@@ -105,8 +105,8 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m)
         // Classmethod: batch construction with uniform device_id (avoids np.full allocation)
         .def_static(
             "from_arrays_uniform_device",
-            [](kvc::MemoryType type, nb::ndarray<int64_t, nb::ndim<1>, nb::c_contig, nb::device::cpu> addrs,
-                nb::ndarray<int64_t, nb::ndim<1>, nb::c_contig, nb::device::cpu> sizes, uint32_t deviceId)
+            [](kvc::MemoryType type, nb::ndarray<int64_t const, nb::ndim<1>, nb::c_contig, nb::device::cpu> addrs,
+                nb::ndarray<int64_t const, nb::ndim<1>, nb::c_contig, nb::device::cpu> sizes, uint32_t deviceId)
             {
                 size_t n = addrs.shape(0);
                 auto const* a = addrs.data();

tensorrt_llm/_torch/disaggregation/native/auxiliary.py

@@ -34,6 +34,81 @@ def from_dict(cls, data: dict[str, Any]) -> "AuxBufferMeta":
         )
 
 
+@dataclass(frozen=True)
+class AuxTransferLayout:
+    """Static auxiliary memory layout shared by all transfers to one peer."""
+
+    src_base_ptrs: np.ndarray
+    dst_base_ptrs: np.ndarray
+    src_item_sizes: np.ndarray
+    dst_item_sizes: np.ndarray
+
+
+def _readonly(array: np.ndarray) -> np.ndarray:
+    """Return an array protected from accidental in-place updates."""
+    array.flags.writeable = False
+    return array
+
+
+def get_non_empty_aux_indices(ptrs: np.ndarray, sizes: np.ndarray, context: str) -> np.ndarray:
+    """Validate auxiliary memory descriptors and return their non-empty indices."""
+    if ptrs.shape != sizes.shape:
+        raise ValueError(f"{context}: pointer/size count mismatch: {ptrs.shape=} != {sizes.shape=}")
+
+    negative_sizes = sizes < 0
+    if negative_sizes.any():
+        indices = np.flatnonzero(negative_sizes).tolist()
+        raise ValueError(f"{context}: negative sizes at indices {indices}")
+
+    null_non_empty = (ptrs == 0) & (sizes > 0)
+    if null_non_empty.any():
+        indices = np.flatnonzero(null_non_empty).tolist()
+        raise ValueError(f"{context}: null pointers with non-zero sizes at indices {indices}")
+
+    return np.flatnonzero(sizes > 0)
+
+
+def build_aux_transfer_layout(
+    src_meta: AuxBufferMeta, dst_meta: AuxBufferMeta
+) -> AuxTransferLayout:
+    """Validate and build the static auxiliary transfer layout for one peer."""
+    src_item_sizes = src_meta.item_sizes.astype(np.int64, copy=False)
+    dst_item_sizes = dst_meta.item_sizes.astype(np.int64, copy=False)
+    src_indices = get_non_empty_aux_indices(
+        src_meta.ptrs, src_item_sizes, "source auxiliary transfer"
+    )
+    dst_indices = get_non_empty_aux_indices(
+        dst_meta.ptrs, dst_item_sizes, "destination auxiliary transfer"
+    )
+    if src_meta.ptrs.shape != dst_meta.ptrs.shape:
+        raise ValueError(
+            "Source and destination auxiliary layouts do not match: "
+            f"{src_meta.ptrs.shape=} != {dst_meta.ptrs.shape=}"
+        )
+
+    dst_non_empty = np.zeros(dst_meta.ptrs.shape, dtype=bool)
+    dst_non_empty[dst_indices] = True
+    missing_dst = src_indices[~dst_non_empty[src_indices]]
+    if missing_dst.size > 0:
+        raise ValueError(
+            "Destination auxiliary buffers are empty for non-empty source "
+            f"indices {missing_dst.tolist()}"
+        )
+
+    too_small = src_indices[dst_item_sizes[src_indices] < src_item_sizes[src_indices]]
+    if too_small.size > 0:
+        raise ValueError(
+            f"Destination auxiliary buffers are too small at indices {too_small.tolist()}"
+        )
+
+    return AuxTransferLayout(
+        src_base_ptrs=_readonly(src_meta.ptrs[src_indices]),
+        dst_base_ptrs=_readonly(dst_meta.ptrs[src_indices]),
+        src_item_sizes=_readonly(src_item_sizes[src_indices]),
+        dst_item_sizes=_readonly(dst_item_sizes[src_indices]),
+    )
+
+
 AuxSlot = namedtuple("AuxSlot", ["id", "buffer"])
 
 

tensorrt_llm/_torch/disaggregation/native/peer.py

@@ -15,12 +15,13 @@
 
 from collections import Counter
 from dataclasses import dataclass, field
-from typing import Dict, List, Tuple
+from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 
 from tensorrt_llm import logger
 from tensorrt_llm._torch.disaggregation.base.region import RegionMapperBase
+from tensorrt_llm._torch.disaggregation.native.auxiliary import AuxTransferLayout
 from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import AttentionPolicy
 from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo
 from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1
@@ -60,6 +61,7 @@ def __init__(self, self_rank_info: RankInfo, self_extractor: KVRegionExtractorV1
         self._self_ext_cache = self_extractor
         self._peer_ext_cache: Dict[str, KVRegionExtractorV1] = {}
         self._overlap_cache: Dict[str, PeerOverlap] = {}
+        self._aux_transfer_layout_cache: Dict[str, AuxTransferLayout] = {}
         self._lg_pool_mapping_cache: Dict[
             str, Dict[LGPoolKey, LGPoolKey]
         ] = {}  # peer_key -> {(self_lg, self_pi) -> (peer_lg, peer_pi)}
@@ -72,6 +74,7 @@ def register(self, peer_name: str, peer_rank: int, peer_ri: RankInfo):
             )
         key = self._unique_key(peer_name, peer_rank)
         self._peer_ri_cache[key] = peer_ri
+        self._aux_transfer_layout_cache.pop(key, None)
         peer_ri = self.get_peer_rank_info(peer_name, peer_rank)
         extractor = KVRegionExtractorV1(peer_ri.page_table)
         self._peer_ext_cache[key] = extractor
@@ -117,6 +120,7 @@ def unregister(self, peer_name: str, peer_rank: int):
             del self._peer_ri_cache[key]
         if key in self._peer_ext_cache:
             del self._peer_ext_cache[key]
+        self._aux_transfer_layout_cache.pop(key, None)
         # Clean up kv_map_cache entries for this peer
         keys_to_remove = [k for k in self._kv_map_cache if k[0] == key]
         for k in keys_to_remove:
@@ -127,6 +131,16 @@ def unregister(self, peer_name: str, peer_rank: int):
     def get_peer_rank_info(self, peer_name: str, peer_rank: int):
         return self._peer_ri_cache[self._unique_key(peer_name, peer_rank)]
 
+    def get_aux_transfer_layout(
+        self, peer_name: str, peer_rank: int
+    ) -> Optional[AuxTransferLayout]:
+        return self._aux_transfer_layout_cache.get(self._unique_key(peer_name, peer_rank))
+
+    def cache_aux_transfer_layout(
+        self, peer_name: str, peer_rank: int, layout: AuxTransferLayout
+    ) -> None:
+        self._aux_transfer_layout_cache[self._unique_key(peer_name, peer_rank)] = layout
+
     @property
     def self_rank_info(self) -> RankInfo:
         return self._ri

tensorrt_llm/_torch/disaggregation/native/transfer.py

@@ -54,7 +54,11 @@
     TxSessionBase,
     WaitResult,
 )
-from tensorrt_llm._torch.disaggregation.native.auxiliary import AuxBuffer
+from tensorrt_llm._torch.disaggregation.native.auxiliary import (
+    AuxBuffer,
+    build_aux_transfer_layout,
+    get_non_empty_aux_indices,
+)
 from tensorrt_llm._torch.disaggregation.native.messenger import ZMQMessenger, decode_message
 from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import MambaPolicy
 from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar
@@ -934,9 +938,17 @@ def _build_aux_write_meta(self, task: AuxSendTask, req_info: RecvReqInfo) -> Wri
             peer_slot = req_info.aux_slot
             assert peer_slot is not None, f"aux_slot is None for request {req_info.unique_rid}"
             assert task._slot is not None
-            src_ptrs = src_aux_meta.ptrs + src_aux_meta.item_sizes * task._slot
-            dst_ptrs = peer_aux_meta.ptrs + peer_aux_meta.item_sizes * peer_slot
-            sizes = src_aux_meta.item_sizes.astype(np.int64, copy=False)
+            layout = self._registrar.get_aux_transfer_layout(
+                peer_ri.instance_name, peer_ri.instance_rank
+            )
+            if layout is None:
+                layout = build_aux_transfer_layout(src_aux_meta, peer_aux_meta)
+                self._registrar.cache_aux_transfer_layout(
+                    peer_ri.instance_name, peer_ri.instance_rank, layout
+                )
+            src_ptrs = layout.src_base_ptrs + layout.src_item_sizes * task._slot
+            dst_ptrs = layout.dst_base_ptrs + layout.dst_item_sizes * peer_slot
+            sizes = layout.src_item_sizes
 
         if timer:
             timer.record_prepare_args_end(peer_ri.instance_rank)
@@ -2353,10 +2365,14 @@ def _register_kv_cache(self):
     def _register_aux_buffer(self):
         assert self._aux_buffer is not None
         aux_meta = self._aux_buffer.meta
-        ptr_num = len(aux_meta.ptrs)
+        non_empty = get_non_empty_aux_indices(
+            aux_meta.ptrs, aux_meta.size, "auxiliary registration"
+        )
         ptr_descs = []
-        for i in range(ptr_num):
+        for i in non_empty:
             ptr_descs.append((aux_meta.ptrs[i], aux_meta.size[i], 0, f"aux_buffer_ptr_{i}"))
+        if not ptr_descs:
+            return
         reg_memory_desc = RegMemoryDescs("DRAM", ptr_descs)
         self._agent.register_memory(reg_memory_desc)
         logger.debug(f"Registered auxiliary buffer memory with transfer agent: {reg_memory_desc}")

tests/integration/test_lists/test-db/l0_a10.yml

@@ -73,6 +73,7 @@ l0_a10:
   - unittest/disaggregated/test_peer.py
   - unittest/disaggregated/test_cache_reuse_adapter.py
   - unittest/disaggregated/test_bounce.py
+  - unittest/disaggregated/test_aux_buffer_registration.py
   - unittest/disaggregated/region/test_block.py
   - unittest/disaggregated/test_mamba_transfer.py
   - unittest/tools

tests/integration/test_lists/test-db/l0_h100.yml

@@ -91,6 +91,7 @@ l0_h100:
   - unittest/disaggregated/test_kv_transfer.py
   - unittest/disaggregated/test_kv_transfer_mp.py
   - unittest/disaggregated/test_transceiver_bounded_polling.py
+  - unittest/disaggregated/test_aux_buffer_registration.py
   - unittest/disaggregated/test_pool_matching.py
   - unittest/disaggregated/test_deepseek_v4_kv_transfer.py
   - unittest/disaggregated/test_minimax_m3_kv_transfer.py

tests/unittest/bindings/test_transfer_agent_bindings.py

@@ -12,6 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import numpy as np
 import pytest
 
 # Try to import the transfer agent binding module
@@ -158,6 +159,38 @@ def test_memory_descs():
     assert descs.descs[1].addr == 0x2000
 
 
+def test_memory_descs_from_readonly_arrays():
+    """Test MemoryDescs construction from read-only arrays."""
+    addrs = np.array([0x1000, 0x2000], dtype=np.int64)
+    sizes = np.array([4096, 8192], dtype=np.int64)
+    device_ids = np.array([0, 1], dtype=np.int32)
+    addrs.flags.writeable = False
+    sizes.flags.writeable = False
+    device_ids.flags.writeable = False
+
+    descs = tab.MemoryDescs.from_arrays(tab.MemoryType.VRAM, addrs, sizes, device_ids)
+
+    assert [(desc.addr, desc.len, desc.device_id) for desc in descs.descs] == [
+        (0x1000, 4096, 0),
+        (0x2000, 8192, 1),
+    ]
+
+
+def test_memory_descs_from_readonly_arrays_uniform_device():
+    """Test uniform-device MemoryDescs construction from read-only arrays."""
+    addrs = np.array([0x1000, 0x2000], dtype=np.int64)
+    sizes = np.array([4096, 8192], dtype=np.int64)
+    addrs.flags.writeable = False
+    sizes.flags.writeable = False
+
+    descs = tab.MemoryDescs.from_arrays_uniform_device(tab.MemoryType.VRAM, addrs, sizes, 2)
+
+    assert [(desc.addr, desc.len, desc.device_id) for desc in descs.descs] == [
+        (0x1000, 4096, 2),
+        (0x2000, 8192, 2),
+    ]
+
+
 def test_memory_descs_empty():
     """Test MemoryDescs with empty list."""
     descs = tab.MemoryDescs(tab.MemoryType.DRAM, [])