[tpu-raiden] Support shared memory DMA mapping and external tensor H2D/D2H transfers in KVCacheManager

tpu_sync/api/torch/kv_cache_manager.py

@@ -518,3 +518,122 @@ def listener_port(self) -> Optional[int]:
   def is_listener_active(self) -> bool:
     """Returns whether the native C++ KVCacheListener is actively running."""
     return self._impl.is_listener_active
+
+  # =========================================================================
+  # EXPERIMENTAL SHARED-MEMORY DMA APIs
+  # The following APIs register a caller-owned host memory pool for TPU DMA
+  # and transfer directly to and from tensors inside that pool, bypassing the
+  # manager's own host staging blocks. They are experimental and subject to
+  # change in future releases.
+  # =========================================================================
+
+  def experimental_map_shared_memory(
+      self, mapped_address: int, pool_size_bytes: int
+  ) -> None:
+    """[EXPERIMENTAL] Registers an existing whole-pool mapping for TPU DMA.
+
+    WARNING: This API is experimental and subject to change in future releases.
+
+    The external pool owner must keep this process-local virtual address
+    mapped and page-locked until ``experimental_unmap_shared_memory``
+    succeeds.
+
+    At most one pool may be mapped at a time; mapping a different pool
+    requires unmapping the current one first.
+
+    Args:
+      mapped_address: Process-local virtual address of the shared pool.
+      pool_size_bytes: Total byte length of the shared pool.
+    """
+    self._impl.map_shared_memory(mapped_address, pool_size_bytes)
+
+  def experimental_unmap_shared_memory(self) -> None:
+    """[EXPERIMENTAL] Drains submitted copies and releases the registration.
+
+    WARNING: This API is experimental and subject to change in future releases.
+    """
+    self._impl.unmap_shared_memory()
+
+  @property
+  def experimental_is_shared_memory_mapped(self) -> bool:
+    """[EXPERIMENTAL] Returns whether shared memory is currently DMA mapped.
+
+    WARNING: This API is experimental and subject to change in future releases.
+    """
+    return self._impl.is_shared_memory_mapped
+
+  def experimental_d2h(
+      self,
+      block_ids: List[int],
+      object_tensors: List[Any],
+      rank_id: int,
+  ) -> Any:
+    """[EXPERIMENTAL] Copies device KV cache blocks into mapped host tensors.
+
+    WARNING: This API is experimental and subject to change in future releases.
+
+    This is not the same addressing scheme as ``d2h``. ``d2h`` takes block
+    offsets into the manager's own host staging buffers; this transfers
+    directly into caller-owned CPU tensors that live inside the shared memory
+    pool registered by ``experimental_map_shared_memory``. DMA is issued
+    against those tensors' memory, so the pool must still be mapped and every
+    tensor must lie entirely inside it.
+
+    Each object tensor is contiguous with a 1-byte dtype (e.g. ``uint8`` or
+    ``int8``) and shape ``[num_ranks, num_layers, page_nbytes]``. Only the
+    ``rank_id`` slice takes part: for layer ``l``, the ``page_nbytes`` bytes at
+    ``tensor[rank_id, l]`` are filled from block ``block_ids[i]`` of layer
+    ``l``'s device buffer, where ``i`` is that tensor's index.
+
+    Args:
+      block_ids: Device block ids, one per object tensor. These are block
+        indices, not byte offsets. Must be unique and within
+        ``[0, num_blocks)``.
+      object_tensors: Caller-owned contiguous CPU tensors with a 1-byte dtype
+        and shape ``[num_ranks, num_layers, page_nbytes]``, residing inside the
+        mapped pool. One per entry of ``block_ids``.
+      rank_id: Index into each tensor's first dimension, selecting which
+        rank's slice to transfer.
+
+    Returns:
+      A future representing the asynchronous copy transfer operation.
+    """
+    return self._impl.d2h(list(block_ids), list(object_tensors), int(rank_id))
+
+  def experimental_h2d(
+      self,
+      block_ids: List[int],
+      object_tensors: List[Any],
+      rank_id: int,
+  ) -> Any:
+    """[EXPERIMENTAL] Copies mapped host tensors into device KV cache blocks.
+
+    WARNING: This API is experimental and subject to change in future releases.
+
+    This is not the same addressing scheme as ``h2d``. ``h2d`` takes block
+    offsets into the manager's own host staging buffers; this transfers
+    directly from caller-owned CPU tensors that live inside the shared memory
+    pool registered by ``experimental_map_shared_memory``. DMA is issued
+    against those tensors' memory, so the pool must still be mapped and every
+    tensor must lie entirely inside it.
+
+    Each object tensor is contiguous with a 1-byte dtype (e.g. ``uint8`` or
+    ``int8``) and shape ``[num_ranks, num_layers, page_nbytes]``. Only the
+    ``rank_id`` slice takes part: for layer ``l``, the ``page_nbytes`` bytes at
+    ``tensor[rank_id, l]`` are written to block ``block_ids[i]`` of layer
+    ``l``'s device buffer, where ``i`` is that tensor's index.
+
+    Args:
+      block_ids: Device block ids, one per object tensor. These are block
+        indices, not byte offsets. Must be unique and within
+        ``[0, num_blocks)``.
+      object_tensors: Caller-owned contiguous CPU tensors with a 1-byte dtype
+        and shape ``[num_ranks, num_layers, page_nbytes]``, residing inside the
+        mapped pool. One per entry of ``block_ids``.
+      rank_id: Index into each tensor's first dimension, selecting which
+        rank's slice to transfer.
+
+    Returns:
+      A future representing the asynchronous copy transfer operation.
+    """
+    return self._impl.h2d(list(block_ids), list(object_tensors), int(rank_id))

tpu_sync/api/torch/kv_cache_manager_host_test.py

@@ -148,6 +148,31 @@ def test_register_refs_tags_and_admission_summary(self):
     with self.assertRaisesRegex(RuntimeError, "host-only manager"):
       manager.h2d_pool_blocks(pool_idx=0, block_ids=[0])
 
+  def test_shared_memory_mapping_api(self):
+    manager = kv_cache_manager.KVCacheManager.create_host_only_for_testing(
+        num_layers=1,
+        num_shards=1,
+        slice_byte_size=128,
+        node_id=7,
+        host_blocks=2,
+        parallelism=1,
+    )
+    self.assertFalse(manager.experimental_is_shared_memory_mapped)
+
+    # Map with null address or 0 size raises RuntimeError
+    with self.assertRaises(RuntimeError):
+      manager.experimental_map_shared_memory(0, 4096)
+    with self.assertRaises(RuntimeError):
+      manager.experimental_map_shared_memory(4096, 0)
+
+    # Unmap when not mapped raises RuntimeError
+    with self.assertRaises(RuntimeError):
+      manager.experimental_unmap_shared_memory()
+
+    # Host-only manager has no active PJRT client for DMA mapping
+    with self.assertRaisesRegex(RuntimeError, "no active PJRT client"):
+      manager.experimental_map_shared_memory(4096, 4096)
+
 
 if __name__ == "__main__":
   unittest.main()

tpu_sync/frameworks/torch/BUILD

@@ -189,6 +189,7 @@ cc_library(
         "//tpu_sync/core:buffer_utils",
         "//tpu_sync/core:kv_cache_manager_with_transfer",
         "//tpu_sync/core:kv_manager_holder",
+        "//tpu_sync/core:raw_transfer_core",
         "//tpu_sync/core:tpu_utils",
         "//tpu_sync/core:utils",
         "//tpu_sync/core:xla_raw_transfer_headers",
@@ -427,6 +428,7 @@ cc_library(
         "//tpu_sync/core:buffer_utils",
         "//tpu_sync/core:kv_cache_manager_with_transfer",
         "//tpu_sync/core:kv_manager_holder",
+        "//tpu_sync/core:raw_transfer_core",
         "//tpu_sync/core:tpu_utils",
         "//tpu_sync/core:utils",
         "//tpu_sync/core:xla_raw_transfer_headers",
@@ -544,9 +546,13 @@ cc_test(
         ":torch_tpu_utils_mock",
         "//tpu_sync/core/controller:controller_service",
         "//tpu_sync/core/controller:test_util",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/time",
         "@torch_tpu//shims/torch",
+        "@torch_tpu//shims/torch:aten_headers",
         "@torch_tpu//shims/torch:torch_headers",
+        "@torch_tpu//shims/torch/c10/core:base_headers",
         "@xla//xla/pjrt:pjrt_client",
         "@xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client",
         "@xla//xla/tsl/platform:statusor",

tpu_sync/frameworks/torch/kv_cache_manager.cc

@@ -38,6 +38,7 @@
 #include "tpu_sync/core/controller/worker_service_server.h"
 #include "tpu_sync/core/kv_cache_manager_with_transfer.h"
 #include "tpu_sync/core/kv_manager_holder.h"
+#include "tpu_sync/core/raw_transfer_core.h"
 #include "tpu_sync/core/tpu_utils.h"
 #include "tpu_sync/core/utils.h"
 #include "tpu_sync/frameworks/torch/torch_utils.h"
@@ -193,7 +194,19 @@ TorchKVCacheManager::TorchKVCacheManager(
           /*num_slots=*/0, /*timeout_s=*/120.0),
       kv_caches_({}) {}
 
-TorchKVCacheManager::~TorchKVCacheManager() = default;
+TorchKVCacheManager::~TorchKVCacheManager() {
+  // Drain external copies here rather than leaving it to ~KVCacheManagerBase.
+  // base_ is owned by KVCacheManagerWithTransfer and so is torn down after
+  // this subobject; any completion that reaches back into the derived manager
+  // at that point would touch freed memory.
+  if (base()->is_shared_memory_mapped()) {
+    const absl::Status status = base()->UnmapSharedMemory();
+    if (!status.ok()) {
+      LOG(ERROR) << "TorchKVCacheManager shared memory unmap failed: "
+                 << status;
+    }
+  }
+}
 
 std::optional<int> TorchKVCacheManager::listener_port() const {
   if (listener_) {
@@ -446,5 +459,100 @@ int KVCacheManager::GetRaidenWorkerPort() const {
   return controller::WorkerServiceServer::GetInstance().GetRaidenWorkerPort();
 }
 
+absl::StatusOr<raiden::PjRtCopyFuture> TorchKVCacheManager::H2d(
+    const std::vector<int64_t>& block_ids,
+    const std::vector<at::Tensor>& object_tensors, int64_t rank_id) {
+  return CopyObjectBlocks(block_ids, object_tensors, rank_id, /*is_h2d=*/true);
+}
+
+absl::StatusOr<raiden::PjRtCopyFuture> TorchKVCacheManager::D2h(
+    const std::vector<int64_t>& block_ids,
+    const std::vector<at::Tensor>& object_tensors, int64_t rank_id) {
+  return CopyObjectBlocks(block_ids, object_tensors, rank_id, /*is_h2d=*/false);
+}
+
+absl::StatusOr<raiden::PjRtCopyFuture> TorchKVCacheManager::CopyObjectBlocks(
+    const std::vector<int64_t>& block_ids,
+    const std::vector<at::Tensor>& object_tensors, int64_t rank_id,
+    bool is_h2d) {
+  // Only at::Tensor-shaped validation belongs here.  The transfer itself is
+  // framework agnostic and lives on KVCacheManagerBase so the JAX frontend can
+  // reuse it; this function reduces the tensors to raw host pointers plus the
+  // [num_ranks, num_layers, page_nbytes] geometry the base needs.
+  //
+  // Check the request against the manager before the tensors: num_layers is
+  // meaningless without registered device buffers, so validating shapes first
+  // would report a layer-count mismatch for a manager that simply has no
+  // device KV cache.  CopyExternalObjectBlocks repeats this check.
+  absl::Status request_status = base()->ValidateExternalObjectRequest(
+      block_ids, object_tensors.size(), rank_id);
+  if (!request_status.ok()) {
+    return request_status;
+  }
+
+  const size_t num_layers = base()->num_layers();
+
+  size_t page_nbytes = 0;
+  size_t num_ranks = 0;
+  std::vector<uint8_t*> host_bases;
+  host_bases.reserve(object_tensors.size());
+
+  for (size_t obj_id = 0; obj_id < object_tensors.size(); ++obj_id) {
+    const at::Tensor& tensor = object_tensors[obj_id];
+    if (!tensor.device().is_cpu()) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("object_tensors[", obj_id, "] must be a CPU tensor"));
+    }
+    if (!tensor.is_contiguous()) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "object_tensors[", obj_id,
+          "] must be contiguous; a nonzero storage offset is allowed"));
+    }
+    if (tensor.dim() != 3 || tensor.element_size() != 1) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "object_tensors[", obj_id,
+          "] must be a rank-3 CPU tensor with a 1-byte dtype and shape "
+          "[num_ranks, ",
+          num_layers, ", page_nbytes]"));
+    }
+    if (tensor.size(0) <= 0 ||
+        tensor.size(1) != static_cast<int64_t>(num_layers) ||
+        tensor.size(2) <= 0) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "object_tensors[", obj_id, "] must have shape [num_ranks, ",
+          num_layers,
+          ", page_nbytes] with num_ranks > 0 and page_nbytes > 0"));
+    }
+
+    if (obj_id == 0) {
+      num_ranks = static_cast<size_t>(tensor.size(0));
+      page_nbytes = static_cast<size_t>(tensor.size(2));
+    } else if (static_cast<size_t>(tensor.size(0)) != num_ranks) {
+      return absl::InvalidArgumentError(
+          "all object_tensors must have the same num_ranks dimension");
+    } else if (static_cast<size_t>(tensor.size(2)) != page_nbytes) {
+      return absl::InvalidArgumentError(
+          "all object_tensors must have the same page_nbytes dimension");
+    }
+
+    const size_t expected_object_bytes = num_ranks * num_layers * page_nbytes;
+    if (static_cast<size_t>(tensor.nbytes()) != expected_object_bytes) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("object_tensors[", obj_id, "] has ", tensor.nbytes(),
+                       " bytes, expected ", expected_object_bytes));
+    }
+
+    host_bases.push_back(static_cast<uint8_t*>(tensor.data_ptr()));
+  }
+
+  // Keeps the caller's storage alive for as long as any issued copy still
+  // references it, including on a partial submit failure.
+  auto tensor_holds = std::make_shared<std::vector<at::Tensor>>(object_tensors);
+
+  return base()->CopyExternalObjectBlocks(block_ids, host_bases, num_ranks,
+                                          page_nbytes, rank_id, is_h2d,
+                                          std::move(tensor_holds));
+}
+
 }  // namespace torch
 }  // namespace tpu_raiden

tpu_sync/frameworks/torch/kv_cache_manager.h

@@ -105,7 +105,22 @@ class TorchKVCacheManager : public KVCacheManagerWithTransfer {
   absl::Status WriteBlockBytes(size_t layer_idx, int block_id,
                                absl::string_view payload, size_t shard_idx = 0);
 
+  // The legacy offset-based H2d/D2h live on KVCacheManagerBase, reached via
+  // base(); they are not hidden by the object-tensor overloads below.
+  // Overloaded H2D and D2H operating on external PyTorch object tensor views
+  absl::StatusOr<raiden::PjRtCopyFuture> H2d(
+      const std::vector<int64_t>& block_ids,
+      const std::vector<at::Tensor>& object_tensors, int64_t rank_id);
+
+  absl::StatusOr<raiden::PjRtCopyFuture> D2h(
+      const std::vector<int64_t>& block_ids,
+      const std::vector<at::Tensor>& object_tensors, int64_t rank_id);
+
  private:
+  absl::StatusOr<raiden::PjRtCopyFuture> CopyObjectBlocks(
+      const std::vector<int64_t>& block_ids,
+      const std::vector<at::Tensor>& object_tensors, int64_t rank_id,
+      bool is_h2d);
   // Buffers unpacked from a 2D tensor list, together with the owning
   // TensorBufferHandles that must outlive their use (see UnpackTorchTensor).
   struct UnpackedLayers {
@@ -336,6 +351,42 @@ class KVCacheManager {
         slot_idx, layer_idx, shard_idx);
   }
 
+  absl::Status MapSharedMemory(uintptr_t mapped_address,
+                               size_t pool_size_bytes) {
+    return torch_manager_->base()->MapSharedMemory(
+        reinterpret_cast<void*>(mapped_address), pool_size_bytes);
+  }
+
+  absl::Status UnmapSharedMemory() {
+    return torch_manager_->base()->UnmapSharedMemory();
+  }
+
+  bool is_shared_memory_mapped() const {
+    return torch_manager_->base()->is_shared_memory_mapped();
+  }
+
+  void SetSharedMemoryMappedForTest(uintptr_t mapped_address,
+                                    size_t pool_size_bytes) {
+    torch_manager_->base()->SetSharedMemoryMappedForTest(
+        reinterpret_cast<void*>(mapped_address), pool_size_bytes);
+  }
+
+  void ResetSharedMemoryMappedForTest() {
+    torch_manager_->base()->ResetSharedMemoryMappedForTest();
+  }
+
+  absl::StatusOr<raiden::PjRtCopyFuture> H2d(
+      const std::vector<int64_t>& block_ids,
+      const std::vector<at::Tensor>& object_tensors, int64_t rank_id) {
+    return torch_manager_->H2d(block_ids, object_tensors, rank_id);
+  }
+
+  absl::StatusOr<raiden::PjRtCopyFuture> D2h(
+      const std::vector<int64_t>& block_ids,
+      const std::vector<at::Tensor>& object_tensors, int64_t rank_id) {
+    return torch_manager_->D2h(block_ids, object_tensors, rank_id);
+  }
+
   absl::StatusOr<std::pair<std::vector<int>, raiden::PjRtCopyFuture>>
   D2hAutoAllocate(const std::vector<int64_t>& src_offsets_major_dim = {},
                   const std::vector<int64_t>& copy_sizes_major_dim = {}) {