Merge branch 'main' into initial-stats-sweep

Author: BenjaminBraunDev

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

@@ -26,6 +26,7 @@
 #endif
 
 #include <nanobind/nanobind.h>
+#include <nanobind/ndarray.h>
 #include <nanobind/stl/optional.h>
 #include <nanobind/stl/string.h>
 #include <nanobind/stl/tuple.h>
@@ -83,6 +84,47 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m)
                 new (self) kvc::MemoryDescs(type, std::move(descs));
             },
             nb::arg("type"), nb::arg("tuples"))
+        // 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)
+            {
+                size_t n = addrs.shape(0);
+                auto const* a = addrs.data();
+                auto const* s = sizes.data();
+                auto const* d = deviceIds.data();
+                std::vector<kvc::MemoryDesc> descs;
+                descs.reserve(n);
+                for (size_t i = 0; i < n; ++i)
+                {
+                    descs.emplace_back(
+                        static_cast<uintptr_t>(a[i]), static_cast<size_t>(s[i]), static_cast<uint32_t>(d[i]));
+                }
+                return kvc::MemoryDescs(type, std::move(descs));
+            },
+            nb::arg("type"), nb::arg("addrs"), nb::arg("sizes"), nb::arg("device_ids"),
+            nb::call_guard<nb::gil_scoped_release>())
+        // 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)
+            {
+                size_t n = addrs.shape(0);
+                auto const* a = addrs.data();
+                auto const* s = sizes.data();
+                std::vector<kvc::MemoryDesc> descs;
+                descs.reserve(n);
+                for (size_t i = 0; i < n; ++i)
+                {
+                    descs.emplace_back(static_cast<uintptr_t>(a[i]), static_cast<size_t>(s[i]), deviceId);
+                }
+                return kvc::MemoryDescs(type, std::move(descs));
+            },
+            nb::arg("type"), nb::arg("addrs"), nb::arg("sizes"), nb::arg("device_id"),
+            nb::call_guard<nb::gil_scoped_release>())
         .def_prop_ro("type", &kvc::MemoryDescs::getType)
         .def_prop_ro("descs", &kvc::MemoryDescs::getDescs);
 
@@ -105,9 +147,24 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m)
             });
 
     // TransferRequest class
+    //
+    // NOTE: The constructor uses std::move to transfer ownership of src_descs / dst_descs
+    // into the TransferRequest.  This avoids an O(n) copy of the internal
+    // std::vector<MemoryDesc> (24 bytes * n).  For 40k descriptors this saves ~937 KB
+    // of memcpy and turns a ~58 us copy into an O(1) pointer swap (~0.4 us).
+    //
+    // IMPORTANT: After construction, the Python MemoryDescs objects passed as src_descs
+    // and dst_descs are left in a moved-from state — their internal descriptor list
+    // becomes empty.  Do NOT access them after passing to TransferRequest.
     nb::class_<kvc::TransferRequest>(m, "TransferRequest")
-        .def(nb::init<kvc::TransferOp, kvc::TransferDescs, kvc::TransferDescs, std::string const&,
-                 std::optional<kvc::SyncMessage>>(),
+        .def(
+            "__init__",
+            [](kvc::TransferRequest* self, kvc::TransferOp op, kvc::TransferDescs& srcDescs,
+                kvc::TransferDescs& dstDescs, std::string const& remoteName,
+                std::optional<kvc::SyncMessage> syncMessage) {
+                new (self) kvc::TransferRequest(
+                    op, std::move(srcDescs), std::move(dstDescs), remoteName, std::move(syncMessage));
+            },
             nb::arg("op"), nb::arg("src_descs"), nb::arg("dst_descs"), nb::arg("remote_name"),
             nb::arg("sync_message") = std::nullopt)
         .def_prop_ro("op", &kvc::TransferRequest::getOp)

cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp

@@ -18,6 +18,7 @@
 #include "tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h"
 #include "tensorrt_llm/common/envUtils.h"
 #include "tensorrt_llm/common/logger.h"
+#include "tensorrt_llm/common/nvtxUtils.h"
 #include "tensorrt_llm/executor/transferAgent.h"
 #include "tensorrt_llm/runtime/utils/mpiUtils.h"
 
@@ -789,13 +790,15 @@ void NixlTransferAgent::invalidateRemoteAgent(std::string const& name)
     // Set TRTLLM_NIXL_ENABLE_COALESCE=1 to enable this optimization
     if (common::getEnvNixlEnableCoalesce())
     {
+        NVTX3_SCOPED_RANGE(coalesceTransferDescs_CreateXferReq);
         auto [coalescedSrc, coalescedDst] = NixlHelper::coalesceTransferDescs(splitSrc, splitDst);
         status
             = mRawAgent->createXferReq(NixlHelper::convert(request.getOp()), NixlHelper::convertXferDist(coalescedSrc),
                 NixlHelper::convertXferDist(coalescedDst), request.getRemoteName(), handle, &mExtraParams);
     }
     else
     {
+        NVTX3_SCOPED_RANGE(createXferReq);
         status = mRawAgent->createXferReq(NixlHelper::convert(request.getOp()), NixlHelper::convertXferDist(splitSrc),
             NixlHelper::convertXferDist(splitDst), request.getRemoteName(), handle, &mExtraParams);
     }
@@ -804,8 +807,10 @@ void NixlTransferAgent::invalidateRemoteAgent(std::string const& name)
         " rank: %d createXferReq failed with status: %s selfname: %s remoteAgent name: %s",
         mpi::MpiComm::world().getRank(), nixlEnumStrings::statusStr(status).c_str(), mName.c_str(),
         request.getRemoteName().c_str());
-
-    status = mRawAgent->postXferReq(handle, &mExtraParams);
+    {
+        NVTX3_SCOPED_RANGE(postXferReq);
+        status = mRawAgent->postXferReq(handle, &mExtraParams);
+    }
     return std::make_unique<NixlTransferStatus>(mRawAgent.get(), handle);
 }
 
@@ -932,6 +937,7 @@ MemoryDescs NixlTransferAgent::splitDescsFromRegistry(MemoryDescs const& descs)
 std::pair<MemoryDescs, MemoryDescs> NixlTransferAgent::splitTransferDescsFromRegistry(
     MemoryDescs const& srcDescs, MemoryDescs const& dstDescs) const
 {
+    NVTX3_SCOPED_RANGE(splitTransferDescsFromRegistry);
     if (srcDescs.getType() != MemoryType::kVRAM)
         return {srcDescs, dstDescs};
 

tensorrt_llm/_torch/disaggregation/base/agent.py

@@ -3,6 +3,8 @@
 from dataclasses import dataclass
 from typing import List, NamedTuple, Optional, Tuple
 
+import numpy as np
+
 from tensorrt_llm import logger
 
 
@@ -27,10 +29,42 @@ class MemoryDesc(NamedTuple):
     device_id: int
 
 
-@dataclass
 class MemoryDescs:
-    type: str
-    descs: List[MemoryDesc]
+    """Describes a set of memory regions with a common type.
+
+    descs: List of (ptr, size, device_id) tuples.
+    """
+
+    __slots__ = ("type", "descs")
+
+    def __init__(self, type: str, descs: List[tuple[int, int, int]]):
+        self.type = type
+        self.descs = descs
+
+    @classmethod
+    def from_arrays(
+        cls, type: str, addrs: np.ndarray, sizes: np.ndarray, device_ids: np.ndarray
+    ) -> "MemoryDescs":
+        """Batch-construct from numpy arrays of addrs, sizes, device_ids.
+
+        Pure-Python fallback; the C++ binding overrides this with a version
+        that reads numpy raw pointers directly.
+        """
+        descs = np.stack([addrs, sizes, device_ids], axis=1).tolist()
+        return cls(type, [tuple(d) for d in descs])
+
+    @classmethod
+    def from_arrays_uniform_device(
+        cls, type: str, addrs: np.ndarray, sizes: np.ndarray, device_id: int
+    ) -> "MemoryDescs":
+        """Batch-construct from numpy arrays with a single device_id for all entries.
+
+        Pure-Python fallback; the C++ binding overrides this with a version
+        that reads numpy raw pointers directly.
+        """
+        dev_ids = np.full(addrs.size, device_id, dtype=np.int32)
+        descs = np.stack([addrs, sizes, dev_ids], axis=1).tolist()
+        return cls(type, [tuple(d) for d in descs])
 
 
 @dataclass

tensorrt_llm/_torch/disaggregation/base/region.py

@@ -3,6 +3,8 @@
 from enum import IntFlag, auto
 from typing import List, NamedTuple, Optional
 
+import numpy as np
+
 
 @dataclass(frozen=True)
 class IndexRange:
@@ -33,7 +35,7 @@ class MemRegion(NamedTuple):
 class MemRegionGroup(NamedTuple):
     """Describes a block of memory by starting pointer and size in bytes."""
 
-    ptrs: List[int]
+    ptrs: np.ndarray  # dtype=np.int64
     bytes_per_region: int
 
 
@@ -89,10 +91,10 @@ class RegionExtractorBase(ABC):
     """
 
     @abstractmethod
-    def extract(self, region_ids: Optional[List[int]] = None) -> List[SpecRegion]:
+    def extract(self, region_ids: Optional[np.ndarray] = None) -> List[SpecRegion]:
         """
         Args:
-            region_ids: (Optional) List of integer region identifiers to extract.
+            region_ids: (Optional) np.ndarray of integer region identifiers to extract.
         Returns:
             List of Regions for corresponding regions.
         """

tensorrt_llm/_torch/disaggregation/base/transfer.py

@@ -6,6 +6,8 @@
 from enum import Enum
 from typing import List, Optional, cast
 
+import numpy as np
+
 from tensorrt_llm import DisaggregatedParams
 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest
 
@@ -46,9 +48,9 @@ class KVSlice:
 
     token_range: Optional[TokenRange] = None
     layer_range: Optional[LayerRange] = None
-    block_ids_per_layer_groups: List[List[int]] = field(
+    block_ids_per_layer_groups: List[np.ndarray] = field(
         default_factory=list
-    )  # Physical block IDs per layer group
+    )  # Physical block IDs per layer group, each np.ndarray(dtype=np.int64)
     is_last_slice: bool = False
 
 

tensorrt_llm/_torch/disaggregation/native/auxiliary.py

@@ -3,32 +3,33 @@
 from dataclasses import dataclass, field
 from typing import Any
 
+import numpy as np
 import torch
 
 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest
 
 
 @dataclass
 class AuxBufferMeta:
-    ptrs: list[int]
-    size: list[int]
-    item_sizes: list[int] = field(default_factory=list)
+    ptrs: np.ndarray  # dtype=np.int64
+    size: np.ndarray  # dtype=np.int64
+    item_sizes: np.ndarray = field(default_factory=lambda: np.array([], dtype=np.int64))
     device: str = "cpu"
 
     def to_dict(self) -> dict[str, Any]:
         return {
-            "ptrs": self.ptrs,
-            "size": self.size,
-            "item_sizes": self.item_sizes,
+            "ptrs": self.ptrs.tolist(),
+            "size": self.size.tolist(),
+            "item_sizes": self.item_sizes.tolist(),
             "device": self.device,
         }
 
     @classmethod
     def from_dict(cls, data: dict[str, Any]) -> "AuxBufferMeta":
         return cls(
-            ptrs=data["ptrs"],
-            size=data["size"],
-            item_sizes=data.get("item_sizes", []),
+            ptrs=np.array(data["ptrs"], dtype=np.int64),
+            size=np.array(data["size"], dtype=np.int64),
+            item_sizes=np.array(data.get("item_sizes", []), dtype=np.int64),
             device=data.get("device", "cpu"),
         )
 
@@ -109,21 +110,30 @@ def __init__(self, max_slot_num: int, beam_width: int, max_draft_len: int, devic
         )
 
         self._meta = AuxBufferMeta(
-            ptrs=[
-                self._first_tokens_buffer.data_ptr(),
-                self._draft_tokens_buffer.data_ptr(),
-                self._token_counts_buffer.data_ptr(),
-            ],
-            size=[
-                self._first_tokens_buffer.numel() * self._first_tokens_buffer.element_size(),
-                self._draft_tokens_buffer.numel() * self._draft_tokens_buffer.element_size(),
-                self._token_counts_buffer.numel() * self._token_counts_buffer.element_size(),
-            ],
-            item_sizes=[
-                self._first_tokens_buffer[0].numel() * self._first_tokens_buffer.element_size(),
-                self._draft_tokens_buffer[0].numel() * self._draft_tokens_buffer.element_size(),
-                self._token_counts_buffer[0].numel() * self._token_counts_buffer.element_size(),
-            ],
+            ptrs=np.array(
+                [
+                    self._first_tokens_buffer.data_ptr(),
+                    self._draft_tokens_buffer.data_ptr(),
+                    self._token_counts_buffer.data_ptr(),
+                ],
+                dtype=np.int64,
+            ),
+            size=np.array(
+                [
+                    self._first_tokens_buffer.numel() * self._first_tokens_buffer.element_size(),
+                    self._draft_tokens_buffer.numel() * self._draft_tokens_buffer.element_size(),
+                    self._token_counts_buffer.numel() * self._token_counts_buffer.element_size(),
+                ],
+                dtype=np.int64,
+            ),
+            item_sizes=np.array(
+                [
+                    self._first_tokens_buffer[0].numel() * self._first_tokens_buffer.element_size(),
+                    self._draft_tokens_buffer[0].numel() * self._draft_tokens_buffer.element_size(),
+                    self._token_counts_buffer[0].numel() * self._token_counts_buffer.element_size(),
+                ],
+                dtype=np.int64,
+            ),
             device=self._device,
         )