diff --git a/tpu_sync/api/torch/kv_cache_manager.py b/tpu_sync/api/torch/kv_cache_manager.py index f6f001b2..62bdc3d5 100644 --- a/tpu_sync/api/torch/kv_cache_manager.py +++ b/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)) diff --git a/tpu_sync/api/torch/kv_cache_manager_host_test.py b/tpu_sync/api/torch/kv_cache_manager_host_test.py index a62efaca..9b5633a9 100644 --- a/tpu_sync/api/torch/kv_cache_manager_host_test.py +++ b/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() diff --git a/tpu_sync/frameworks/torch/BUILD b/tpu_sync/frameworks/torch/BUILD index f331d5dd..56ba52d1 100644 --- a/tpu_sync/frameworks/torch/BUILD +++ b/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", diff --git a/tpu_sync/frameworks/torch/kv_cache_manager.cc b/tpu_sync/frameworks/torch/kv_cache_manager.cc index fd35025c..3798b330 100644 --- a/tpu_sync/frameworks/torch/kv_cache_manager.cc +++ b/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 TorchKVCacheManager::listener_port() const { if (listener_) { @@ -446,5 +459,100 @@ int KVCacheManager::GetRaidenWorkerPort() const { return controller::WorkerServiceServer::GetInstance().GetRaidenWorkerPort(); } +absl::StatusOr TorchKVCacheManager::H2d( + const std::vector& block_ids, + const std::vector& object_tensors, int64_t rank_id) { + return CopyObjectBlocks(block_ids, object_tensors, rank_id, /*is_h2d=*/true); +} + +absl::StatusOr TorchKVCacheManager::D2h( + const std::vector& block_ids, + const std::vector& object_tensors, int64_t rank_id) { + return CopyObjectBlocks(block_ids, object_tensors, rank_id, /*is_h2d=*/false); +} + +absl::StatusOr TorchKVCacheManager::CopyObjectBlocks( + const std::vector& block_ids, + const std::vector& 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 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(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(tensor.size(0)); + page_nbytes = static_cast(tensor.size(2)); + } else if (static_cast(tensor.size(0)) != num_ranks) { + return absl::InvalidArgumentError( + "all object_tensors must have the same num_ranks dimension"); + } else if (static_cast(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(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(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>(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 diff --git a/tpu_sync/frameworks/torch/kv_cache_manager.h b/tpu_sync/frameworks/torch/kv_cache_manager.h index d3f2e32b..ae438df2 100644 --- a/tpu_sync/frameworks/torch/kv_cache_manager.h +++ b/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 H2d( + const std::vector& block_ids, + const std::vector& object_tensors, int64_t rank_id); + + absl::StatusOr D2h( + const std::vector& block_ids, + const std::vector& object_tensors, int64_t rank_id); + private: + absl::StatusOr CopyObjectBlocks( + const std::vector& block_ids, + const std::vector& 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(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(mapped_address), pool_size_bytes); + } + + void ResetSharedMemoryMappedForTest() { + torch_manager_->base()->ResetSharedMemoryMappedForTest(); + } + + absl::StatusOr H2d( + const std::vector& block_ids, + const std::vector& object_tensors, int64_t rank_id) { + return torch_manager_->H2d(block_ids, object_tensors, rank_id); + } + + absl::StatusOr D2h( + const std::vector& block_ids, + const std::vector& object_tensors, int64_t rank_id) { + return torch_manager_->D2h(block_ids, object_tensors, rank_id); + } + absl::StatusOr, raiden::PjRtCopyFuture>> D2hAutoAllocate(const std::vector& src_offsets_major_dim = {}, const std::vector& copy_sizes_major_dim = {}) { diff --git a/tpu_sync/frameworks/torch/kv_cache_manager_test.cc b/tpu_sync/frameworks/torch/kv_cache_manager_test.cc index fc0de5a7..5cd0594f 100644 --- a/tpu_sync/frameworks/torch/kv_cache_manager_test.cc +++ b/tpu_sync/frameworks/torch/kv_cache_manager_test.cc @@ -14,12 +14,29 @@ #include "tpu_sync/frameworks/torch/kv_cache_manager.h" +#include + +#include +#include +#include +#include +#include +#include #include #include #include +#include // NOLINT(build/c++11) #include +#include "ATen/core/TensorBody.h" +#include "ATen/ops/from_blob.h" +#include "ATen/ops/zeros.h" +#include "c10/core/ScalarType.h" +#include "c10/core/TensorOptions.h" +#include "absl/status/status.h" #include "absl/strings/match.h" +#include "absl/time/clock.h" +#include "absl/time/time.h" #include "xla/pjrt/pjrt_client.h" #include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h" #include "xla/tsl/platform/statusor.h" @@ -27,7 +44,6 @@ #include "tpu_sync/core/controller/controller_service.h" #include "tpu_sync/core/controller/test_util.h" #include "tpu_sync/frameworks/torch/torch_tpu_utils_mock.h" -#include "torch/torch.h" namespace tpu_raiden { namespace torch { @@ -59,7 +75,7 @@ TEST_F(KVCacheManagerTorchTest, ConstructorSucceedsWithMocks) { /*device_layout=*/nullptr)); // Create a CPU PyTorch tensor - at::Tensor tensor = ::torch::zeros({8, 1024}, ::torch::kFloat32); + at::Tensor tensor = at::zeros({8, 1024}, at::kFloat); // Register the mapping RegisterMockTensor(tensor, pjrt_buffer.get()); @@ -170,7 +186,7 @@ TEST_F(KVCacheManagerTorchTest, kImmutableUntilTransferCompletes, /*on_done_with_host_buffer=*/nullptr, memory_space, /*device_layout=*/nullptr)); - at::Tensor tensor = ::torch::zeros({8, 1024}, ::torch::kFloat32); + at::Tensor tensor = at::zeros({8, 1024}, at::kFloat); RegisterMockTensor(tensor, pjrt_buffer.get()); std::vector> device_tensors = {{tensor}}; @@ -202,10 +218,438 @@ TEST_F(KVCacheManagerTorchTest, ASSERT_EQ(workers.size(), 1); EXPECT_EQ(workers[0].worker_id, "torch_worker_ctrl"); - // Verify registered transfer endpoints carry the DATA transport port (local_data_eps), NOT the control port (local_control_eps). + // Verify registered transfer endpoints carry the DATA transport port + // (local_data_eps), NOT the control port (local_control_eps). ASSERT_FALSE(workers[0].raiden_transfer_endpoints.empty()); - EXPECT_EQ(workers[0].raiden_transfer_endpoints[0].endpoint, local_data_eps[0].endpoint); - EXPECT_NE(workers[0].raiden_transfer_endpoints[0].endpoint, local_control_eps[0].endpoint); + EXPECT_EQ(workers[0].raiden_transfer_endpoints[0].endpoint, + local_data_eps[0].endpoint); + EXPECT_NE(workers[0].raiden_transfer_endpoints[0].endpoint, + local_control_eps[0].endpoint); +} + +TEST_F(KVCacheManagerTorchTest, MapSharedMemoryValidation) { + TF_ASSERT_OK_AND_ASSIGN( + xla::PjRtMemorySpace * memory_space, + client_->addressable_devices()[0]->default_memory_space()); + std::vector data(8 * 1024, 1.0f); + TF_ASSERT_OK_AND_ASSIGN( + auto pjrt_buffer, client_->BufferFromHostBuffer( + data.data(), xla::F32, {8, 1024}, + /*byte_strides=*/std::nullopt, + xla::PjRtClient::HostBufferSemantics:: + kImmutableUntilTransferCompletes, + /*on_done_with_host_buffer=*/nullptr, memory_space, + /*device_layout=*/nullptr)); + + at::Tensor tensor = at::zeros({8, 1024}, at::kFloat); + RegisterMockTensor(tensor, pjrt_buffer.get()); + std::vector> device_tensors = {{tensor}}; + + KVCacheManager manager(device_tensors, + /*local_port=*/std::nullopt, + /*host_blocks_to_allocate=*/8); + + EXPECT_FALSE(manager.is_shared_memory_mapped()); + + const int64_t page_size_val = sysconf(_SC_PAGESIZE); + ASSERT_GT(page_size_val, 0); + const size_t page_size = static_cast(page_size_val); + + void* aligned_ptr = nullptr; + ASSERT_EQ(posix_memalign(&aligned_ptr, page_size, page_size * 2), 0); + ASSERT_NE(aligned_ptr, nullptr); + + // Null pointer check + EXPECT_FALSE(manager.MapSharedMemory(0, page_size).ok()); + + // Zero size check + EXPECT_FALSE(manager.MapSharedMemory( + reinterpret_cast(aligned_ptr), 0).ok()); + + // Unaligned address check + EXPECT_FALSE(manager.MapSharedMemory( + reinterpret_cast(aligned_ptr) + 1, page_size).ok()); + + // Unaligned size check + EXPECT_FALSE(manager.MapSharedMemory( + reinterpret_cast(aligned_ptr), page_size + 1).ok()); + + // Address overflow check + EXPECT_FALSE(manager.MapSharedMemory( + std::numeric_limits::max() - page_size + 1, + page_size * 2).ok()); + + // Unmap when not mapped + EXPECT_FALSE(manager.UnmapSharedMemory().ok()); + + // Calling MapSharedMemory delegates to client->DmaMap, which fails with + // Unimplemented on CPU PJRT client, leaving mapping state unmapped. + absl::Status status = manager.MapSharedMemory( + reinterpret_cast(aligned_ptr), page_size); + EXPECT_FALSE(status.ok()); + EXPECT_FALSE(manager.is_shared_memory_mapped()); + + free(aligned_ptr); +} + +TEST_F(KVCacheManagerTorchTest, H2dD2hWithObjectTensorsValidation) { + TF_ASSERT_OK_AND_ASSIGN( + xla::PjRtMemorySpace * memory_space, + client_->addressable_devices()[0]->default_memory_space()); + std::vector data(8 * 1024, 1.0f); + TF_ASSERT_OK_AND_ASSIGN( + auto pjrt_buffer, client_->BufferFromHostBuffer( + data.data(), xla::F32, {8, 1024}, + /*byte_strides=*/std::nullopt, + xla::PjRtClient::HostBufferSemantics:: + kImmutableUntilTransferCompletes, + /*on_done_with_host_buffer=*/nullptr, memory_space, + /*device_layout=*/nullptr)); + + at::Tensor tensor = at::zeros({8, 1024}, at::kFloat); + RegisterMockTensor(tensor, pjrt_buffer.get()); + std::vector> device_tensors = {{tensor}}; + + KVCacheManager manager(device_tensors, + /*local_port=*/std::nullopt, + /*host_blocks_to_allocate=*/8); + + const int64_t page_nbytes = manager.slice_byte_size(); + EXPECT_EQ(page_nbytes, 1024 * sizeof(float)); + + // Valid 3D object tensor: [num_ranks=2, num_layers=1, page_nbytes=4096] + at::Tensor valid_obj = at::zeros({2, 1, page_nbytes}, at::kChar); + + // Empty block_ids + EXPECT_FALSE(manager.H2d({}, {}, 0).ok()); + + // Size mismatch + EXPECT_FALSE(manager.H2d({0}, {}, 0).ok()); + + // Calling H2d/D2h while unmapped fails with FailedPrecondition + EXPECT_FALSE(manager.H2d({0}, {valid_obj}, 0).ok()); + EXPECT_FALSE(manager.D2h({0}, {valid_obj}, 0).ok()); + + // Simulate mapped shared memory via test helper + manager.SetSharedMemoryMappedForTest(0x1000000, 1024 * 1024); + EXPECT_TRUE(manager.is_shared_memory_mapped()); + + // Duplicate mapping attempt fails + EXPECT_FALSE(manager.MapSharedMemory(0x1000000, 4096).ok()); + + // Out-of-bounds block_id (max_blocks is 8) + EXPECT_FALSE(manager.H2d({8}, {valid_obj}, 0).ok()); + EXPECT_FALSE(manager.H2d({-1}, {valid_obj}, 0).ok()); + + // Duplicate block_id + at::Tensor valid_obj2 = at::zeros({2, 1, page_nbytes}, at::kChar); + EXPECT_FALSE(manager.H2d({0, 0}, {valid_obj, valid_obj2}, 0).ok()); + + // Non-3D tensor + at::Tensor non_3d = at::zeros({2, page_nbytes}, at::kChar); + EXPECT_FALSE(manager.H2d({0}, {non_3d}, 0).ok()); + + // Non-1-byte dtype + at::Tensor float_obj = at::zeros({2, 1, page_nbytes}, at::kFloat); + EXPECT_FALSE(manager.H2d({0}, {float_obj}, 0).ok()); + + // num_layers mismatch (2 layers vs 1 layer) + at::Tensor wrong_layers = at::zeros({2, 2, page_nbytes}, at::kChar); + EXPECT_FALSE(manager.H2d({0}, {wrong_layers}, 0).ok()); + + // Slice size mismatch (2048 vs 4096) + at::Tensor wrong_slice = at::zeros({2, 1, 2048}, at::kChar); + EXPECT_FALSE(manager.H2d({0}, {wrong_slice}, 0).ok()); + + // rank_id out of range (rank_id=2 >= num_ranks=2) + EXPECT_FALSE(manager.H2d({0}, {valid_obj}, /*rank_id=*/2).ok()); + EXPECT_FALSE(manager.H2d({0}, {valid_obj}, /*rank_id=*/-1).ok()); + + // Reset test mapping state + manager.ResetSharedMemoryMappedForTest(); + EXPECT_FALSE(manager.is_shared_memory_mapped()); +} + +TEST_F(KVCacheManagerTorchTest, H2dD2hWithObjectTensorsEndToEnd) { + TF_ASSERT_OK_AND_ASSIGN( + xla::PjRtMemorySpace * memory_space, + client_->addressable_devices()[0]->default_memory_space()); + std::vector data(8 * 1024, 0.0f); + TF_ASSERT_OK_AND_ASSIGN( + auto pjrt_buffer, client_->BufferFromHostBuffer( + data.data(), xla::F32, {8, 1024}, + /*byte_strides=*/std::nullopt, + xla::PjRtClient::HostBufferSemantics:: + kImmutableUntilTransferCompletes, + /*on_done_with_host_buffer=*/nullptr, memory_space, + /*device_layout=*/nullptr)); + + at::Tensor tensor = at::zeros({8, 1024}, at::kFloat); + RegisterMockTensor(tensor, pjrt_buffer.get()); + std::vector> device_tensors = {{tensor}}; + + KVCacheManager manager(device_tensors, + /*local_port=*/std::nullopt, + /*host_blocks_to_allocate=*/8); + + const int64_t page_nbytes = manager.slice_byte_size(); + + // Model the production layout: the object tensors are carved out of the + // registered shared memory pool, so their addresses genuinely fall inside + // the mapped range. Registering an address the tensors do not live at + // would not exercise the DMA path at all. + const int64_t obj_nbytes = 2 * 1 * page_nbytes; + const size_t page_size = static_cast(sysconf(_SC_PAGESIZE)); + const size_t pool_size = + ((2 * static_cast(obj_nbytes) + page_size - 1) / page_size) * + page_size; + void* pool = nullptr; + ASSERT_EQ(posix_memalign(&pool, page_size, pool_size), 0); + std::memset(pool, 0, pool_size); + auto* pool_bytes = static_cast(pool); + + const c10::TensorOptions src_options = c10::TensorOptions().dtype(at::kChar); + const c10::TensorOptions dst_options = c10::TensorOptions().dtype(at::kByte); + at::Tensor src_obj = + at::from_blob(pool_bytes, {2, 1, page_nbytes}, src_options); + at::Tensor dst_obj = + at::from_blob(pool_bytes + obj_nbytes, {2, 1, page_nbytes}, dst_options); + + // Fill src_obj rank 0 with pattern 0x42 + uint8_t* src_ptr = static_cast(src_obj.data_ptr()); + std::memset(src_ptr, 0x42, page_nbytes); + + // Fill src_obj rank 1 with pattern 0x77 + std::memset(src_ptr + page_nbytes, 0x77, page_nbytes); + + manager.SetSharedMemoryMappedForTest(reinterpret_cast(pool), + pool_size); + + // H2D: copy src_obj rank 0 into block 3 + auto h2d_future_or = manager.H2d({3}, {src_obj}, /*rank_id=*/0); + TF_ASSERT_OK_AND_ASSIGN(auto h2d_future, h2d_future_or); + EXPECT_TRUE(h2d_future.Await().ok()); + + // D2H: copy block 3 back into dst_obj rank 0 + auto d2h_future_or = manager.D2h({3}, {dst_obj}, /*rank_id=*/0); + TF_ASSERT_OK_AND_ASSIGN(auto d2h_future, d2h_future_or); + EXPECT_TRUE(d2h_future.Await().ok()); + + // Verify dst_obj rank 0 matches pattern 0x42 + const uint8_t* dst_ptr = static_cast(dst_obj.data_ptr()); + for (int64_t i = 0; i < page_nbytes; ++i) { + ASSERT_EQ(dst_ptr[i], 0x42); + } + + // H2D: copy src_obj rank 1 into block 5 + auto h2d_rank1_or = manager.H2d({5}, {src_obj}, /*rank_id=*/1); + TF_ASSERT_OK_AND_ASSIGN(auto h2d_rank1, h2d_rank1_or); + EXPECT_TRUE(h2d_rank1.Await().ok()); + + // D2H: copy block 5 into dst_obj rank 1 + auto d2h_rank1_or = manager.D2h({5}, {dst_obj}, /*rank_id=*/1); + TF_ASSERT_OK_AND_ASSIGN(auto d2h_rank1, d2h_rank1_or); + EXPECT_TRUE(d2h_rank1.Await().ok()); + + // Verify dst_obj rank 1 matches pattern 0x77 + for (int64_t i = 0; i < page_nbytes; ++i) { + ASSERT_EQ(dst_ptr[page_nbytes + i], 0x77); + } + + // A tensor that lives outside the registered pool must be rejected instead + // of being handed to the DMA engine. + at::Tensor outside_obj = at::zeros({2, 1, page_nbytes}, at::kChar); + EXPECT_FALSE(manager.H2d({1}, {outside_obj}, /*rank_id=*/0).ok()); + EXPECT_FALSE(manager.D2h({1}, {outside_obj}, /*rank_id=*/0).ok()); + + manager.ResetSharedMemoryMappedForTest(); + free(pool); +} + +// Submissions racing an unmap must never tear down the registration out from +// under an in-flight copy: each call either succeeds or is refused outright. +TEST_F(KVCacheManagerTorchTest, ConcurrentCopiesRacingUnmapAreSafe) { + TF_ASSERT_OK_AND_ASSIGN( + xla::PjRtMemorySpace * memory_space, + client_->addressable_devices()[0]->default_memory_space()); + std::vector data(8 * 1024, 0.0f); + TF_ASSERT_OK_AND_ASSIGN( + auto pjrt_buffer, client_->BufferFromHostBuffer( + data.data(), xla::F32, {8, 1024}, + /*byte_strides=*/std::nullopt, + xla::PjRtClient::HostBufferSemantics:: + kImmutableUntilTransferCompletes, + /*on_done_with_host_buffer=*/nullptr, memory_space, + /*device_layout=*/nullptr)); + + at::Tensor tensor = at::zeros({8, 1024}, at::kFloat); + RegisterMockTensor(tensor, pjrt_buffer.get()); + std::vector> device_tensors = {{tensor}}; + + KVCacheManager manager(device_tensors, + /*local_port=*/std::nullopt, + /*host_blocks_to_allocate=*/8); + + const int64_t page_nbytes = manager.slice_byte_size(); + const int64_t obj_nbytes = 1 * 1 * page_nbytes; + const size_t page_size = static_cast(sysconf(_SC_PAGESIZE)); + constexpr int kNumThreads = 4; + const size_t pool_size = ((kNumThreads * static_cast(obj_nbytes) + + page_size - 1) / + page_size) * + page_size; + void* pool = nullptr; + ASSERT_EQ(posix_memalign(&pool, page_size, pool_size), 0); + std::memset(pool, 0, pool_size); + auto* pool_bytes = static_cast(pool); + + const c10::TensorOptions obj_options = c10::TensorOptions().dtype(at::kChar); + std::vector per_thread_obj; + per_thread_obj.reserve(kNumThreads); + for (int t = 0; t < kNumThreads; ++t) { + per_thread_obj.push_back(at::from_blob(pool_bytes + t * obj_nbytes, + {1, 1, page_nbytes}, obj_options)); + } + + manager.SetSharedMemoryMappedForTest(reinterpret_cast(pool), + pool_size); + + std::atomic stop = false; + std::vector threads; + threads.reserve(kNumThreads); + for (int t = 0; t < kNumThreads; ++t) { + threads.emplace_back([&manager, &per_thread_obj, &stop, t]() { + while (!stop.load(std::memory_order_relaxed)) { + // Each thread owns a distinct block so the copies never collide. + auto result = manager.H2d({t}, {per_thread_obj[t]}, /*rank_id=*/0); + if (!result.ok()) { + // The only legal refusal is "the pool is no longer mapped". + EXPECT_EQ(result.status().code(), + absl::StatusCode::kFailedPrecondition) + << result.status(); + } + } + }); + } + + // Racing unmap. DmaUnmap is unimplemented on the CPU client so this is + // expected to fail, but it still drives the drain path that must wait for + // every in-flight submission to register. + absl::SleepFor(absl::Milliseconds(20)); + manager.UnmapSharedMemory().IgnoreError(); + + stop.store(true, std::memory_order_relaxed); + for (std::thread& thread : threads) { + thread.join(); + } + + manager.ResetSharedMemoryMappedForTest(); + free(pool); +} + +// The pool is DMA mapped against exactly one PJRT client. A layer whose +// buffer belongs to a second client would be handed host memory that client +// never registered, so such a manager must refuse object tensor transfers +// instead of issuing the copy. +TEST_F(KVCacheManagerTorchTest, LayersOnDifferentPjRtClientsAreRejected) { + // A second, independent CPU client: its devices report a different + // PjRtClient* than client_, which is what the check keys on. + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr other_client, + xla::GetXlaPjrtCpuClient(xla::CpuClientOptions())); + ASSERT_NE(other_client.get(), client_.get()); + + TF_ASSERT_OK_AND_ASSIGN( + xla::PjRtMemorySpace * memory_space_a, + client_->addressable_devices()[0]->default_memory_space()); + TF_ASSERT_OK_AND_ASSIGN( + xla::PjRtMemorySpace * memory_space_b, + other_client->addressable_devices()[0]->default_memory_space()); + + std::vector data_a(8 * 1024, 1.0f); + std::vector data_b(8 * 1024, 2.0f); + TF_ASSERT_OK_AND_ASSIGN( + auto buffer_a, + client_->BufferFromHostBuffer( + data_a.data(), xla::F32, {8, 1024}, + /*byte_strides=*/std::nullopt, + xla::PjRtClient::HostBufferSemantics:: + kImmutableUntilTransferCompletes, + /*on_done_with_host_buffer=*/nullptr, memory_space_a, + /*device_layout=*/nullptr)); + TF_ASSERT_OK_AND_ASSIGN( + auto buffer_b, + other_client->BufferFromHostBuffer( + data_b.data(), xla::F32, {8, 1024}, + /*byte_strides=*/std::nullopt, + xla::PjRtClient::HostBufferSemantics:: + kImmutableUntilTransferCompletes, + /*on_done_with_host_buffer=*/nullptr, memory_space_b, + /*device_layout=*/nullptr)); + + // Layer 0 on client_, layer 1 on other_client. + std::vector> device_buffers = { + {buffer_a.get()}, {buffer_b.get()}}; + KVCacheManager manager(device_buffers, + /*local_port=*/std::nullopt, + /*host_blocks_to_allocate=*/8, + /*unsafe_skip_buffer_lock=*/true); + ASSERT_EQ(manager.num_layers(), 2); + + const int64_t page_nbytes = manager.slice_byte_size(); + const int64_t obj_nbytes = 2 * 2 * page_nbytes; + const size_t page_size = static_cast(sysconf(_SC_PAGESIZE)); + const size_t pool_size = + ((static_cast(obj_nbytes) + page_size - 1) / page_size) * + page_size; + void* pool = nullptr; + ASSERT_EQ(posix_memalign(&pool, page_size, pool_size), 0); + std::memset(pool, 0, pool_size); + + const c10::TensorOptions obj_options = c10::TensorOptions().dtype(at::kChar); + at::Tensor obj = at::from_blob(static_cast(pool), + {2, 2, page_nbytes}, obj_options); + + // Map the pool so the failure below is attributable to the split clients + // and not to a missing registration. + manager.SetSharedMemoryMappedForTest(reinterpret_cast(pool), + pool_size); + + const absl::Status h2d = manager.H2d({1}, {obj}, /*rank_id=*/0).status(); + EXPECT_EQ(h2d.code(), absl::StatusCode::kUnimplemented); + EXPECT_TRUE(absl::StrContains(h2d.message(), "share one PJRT client")) + << h2d.message(); + + const absl::Status d2h = manager.D2h({1}, {obj}, /*rank_id=*/0).status(); + EXPECT_EQ(d2h.code(), absl::StatusCode::kUnimplemented); + EXPECT_TRUE(absl::StrContains(d2h.message(), "share one PJRT client")) + << d2h.message(); + + manager.ResetSharedMemoryMappedForTest(); + free(pool); +} + +// A manager with no registered device KV cache must say exactly that, even +// for a well formed object tensor. num_layers() is 0 in that state, so +// validating tensor shapes first would report a bogus layer-count mismatch +// instead of the real problem. +TEST_F(KVCacheManagerTorchTest, ObjectTransferWithoutDeviceBuffersIsReported) { + const std::vector> no_device_buffers; + KVCacheManager manager(no_device_buffers, + /*local_port=*/std::nullopt, + /*host_blocks_to_allocate=*/0, + /*unsafe_skip_buffer_lock=*/true); + + at::Tensor obj = at::zeros({2, 1, 4096}, at::kChar); + + const absl::Status h2d = manager.H2d({0}, {obj}, /*rank_id=*/0).status(); + EXPECT_EQ(h2d.code(), absl::StatusCode::kFailedPrecondition); + EXPECT_TRUE(absl::StrContains(h2d.message(), "no registered device KV cache")) + << h2d.message(); + + const absl::Status d2h = manager.D2h({0}, {obj}, /*rank_id=*/0).status(); + EXPECT_EQ(d2h.code(), absl::StatusCode::kFailedPrecondition); + EXPECT_TRUE(absl::StrContains(d2h.message(), "no registered device KV cache")) + << d2h.message(); } } // namespace diff --git a/tpu_sync/frameworks/torch/tpu_raiden_host_module.cc b/tpu_sync/frameworks/torch/tpu_raiden_host_module.cc index 41cba48d..4de1a565 100644 --- a/tpu_sync/frameworks/torch/tpu_raiden_host_module.cc +++ b/tpu_sync/frameworks/torch/tpu_raiden_host_module.cc @@ -224,6 +224,26 @@ NB_MODULE(_tpu_raiden_host, m) { }, nb::arg("block_array_idx")) .def_prop_ro("transfer_address", &HostKVCacheManager::transfer_address) + .def( + "map_shared_memory", + [](HostKVCacheManager& self, uintptr_t mapped_address, + size_t pool_size_bytes) { + ThrowIfError(self.base()->MapSharedMemory( + reinterpret_cast(mapped_address), + pool_size_bytes), + "KVCacheManager map_shared_memory failed"); + }, + nb::arg("mapped_address"), nb::arg("pool_size_bytes")) + .def( + "unmap_shared_memory", + [](HostKVCacheManager& self) { + ThrowIfError(self.base()->UnmapSharedMemory(), + "KVCacheManager unmap_shared_memory failed"); + }) + .def_prop_ro("is_shared_memory_mapped", + [](const HostKVCacheManager& self) { + return self.base()->is_shared_memory_mapped(); + }) .def("get_local_endpoints", [](const HostKVCacheManager& self) { auto eps = self.get_local_endpoints(); diff --git a/tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc b/tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc index 3d4f93c3..fac8df50 100644 --- a/tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc +++ b/tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc @@ -419,6 +419,88 @@ NB_MODULE(_tpu_raiden_torch, m) { nb::arg("dst_offsets_major_dim") = std::vector{}, nb::arg("copy_sizes_major_dim") = std::vector{}, nb::call_guard()) + .def( + "map_shared_memory", + [](KVCacheManager& self, uintptr_t mapped_address, + size_t pool_size_bytes) { + absl::Status status = + self.MapSharedMemory(mapped_address, pool_size_bytes); + if (!status.ok()) { + if (status.code() == absl::StatusCode::kInvalidArgument) { + throw std::invalid_argument( + "KVCacheManager.map_shared_memory failed: " + + std::string(status.message())); + } + throw std::runtime_error( + "KVCacheManager.map_shared_memory failed: " + + std::string(status.message())); + } + }, + nb::arg("mapped_address"), nb::arg("pool_size_bytes"), + nb::call_guard()) + .def( + "unmap_shared_memory", + [](KVCacheManager& self) { + absl::Status status = self.UnmapSharedMemory(); + if (!status.ok()) { + throw std::runtime_error(absl::StrCat( + "KVCacheManager.unmap_shared_memory failed: ", + status.message())); + } + }, + nb::call_guard()) + .def_prop_ro("is_shared_memory_mapped", + &KVCacheManager::is_shared_memory_mapped) + .def( + "h2d", + [](KVCacheManager& self, const std::vector& block_ids, + const std::vector& object_tensors, int64_t rank_id) { + auto result = self.H2d(block_ids, object_tensors, rank_id); + if (!result.ok()) { + const auto code = result.status().code(); + if (code == absl::StatusCode::kInvalidArgument) { + throw std::invalid_argument( + "KVCacheManager.h2d failed: " + + std::string(result.status().message())); + } + if (code == absl::StatusCode::kOutOfRange) { + throw std::out_of_range( + "KVCacheManager.h2d failed: " + + std::string(result.status().message())); + } + throw std::runtime_error( + "KVCacheManager.h2d failed: " + + std::string(result.status().message())); + } + return tpu_raiden::RaidenFuture{std::move(result.value())}; + }, + nb::arg("block_ids"), nb::arg("object_tensors"), nb::arg("rank_id"), + nb::call_guard()) + .def( + "d2h", + [](KVCacheManager& self, const std::vector& block_ids, + const std::vector& object_tensors, int64_t rank_id) { + auto result = self.D2h(block_ids, object_tensors, rank_id); + if (!result.ok()) { + const auto code = result.status().code(); + if (code == absl::StatusCode::kInvalidArgument) { + throw std::invalid_argument( + "KVCacheManager.d2h failed: " + + std::string(result.status().message())); + } + if (code == absl::StatusCode::kOutOfRange) { + throw std::out_of_range( + "KVCacheManager.d2h failed: " + + std::string(result.status().message())); + } + throw std::runtime_error( + "KVCacheManager.d2h failed: " + + std::string(result.status().message())); + } + return tpu_raiden::RaidenFuture{std::move(result.value())}; + }, + nb::arg("block_ids"), nb::arg("object_tensors"), nb::arg("rank_id"), + nb::call_guard()) .def( "D2hAutoAllocate", [](KVCacheManager& self, diff --git a/tpu_sync/kv_cache/kv_cache_manager_base.cc b/tpu_sync/kv_cache/kv_cache_manager_base.cc index 03d47281..a85eb8b0 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_base.cc +++ b/tpu_sync/kv_cache/kv_cache_manager_base.cc @@ -14,6 +14,8 @@ #include "tpu_sync/kv_cache/kv_cache_manager_base.h" +#include + #include #include #include @@ -512,6 +514,12 @@ void KVCacheManagerBase::InitBackgroundWorker() { KVCacheManagerBase::~KVCacheManagerBase() { StopTransportServer(); + if (is_shared_memory_mapped()) { + absl::Status status = UnmapSharedMemory(); + if (!status.ok()) { + LOG(ERROR) << "KVCacheManagerBase unmap cleanup failed: " << status; + } + } if (worker_thread_.joinable()) { { absl::MutexLock lock(queue_mu_); @@ -3494,5 +3502,494 @@ KVCacheManagerBase::ResolveBlockSlices(int staging_block_id) const { return slices; } +xla::PjRtClient* KVCacheManagerBase::GetPjRtClient() const { + for (const auto& layer_device_info : buffer_holds_) { + for (const auto& hold : layer_device_info.holds) { + if (hold.device != nullptr && hold.device->client() != nullptr) { + return hold.device->client(); + } + if (hold.buffer != nullptr && hold.buffer->client() != nullptr) { + return hold.buffer->client(); + } + } + } + return nullptr; +} + +absl::Status KVCacheManagerBase::MapSharedMemory(void* mapped_address, + size_t pool_size_bytes) { + if (mapped_address == nullptr) { + return absl::InvalidArgumentError("mapped_address must be non-null"); + } + if (pool_size_bytes == 0) { + return absl::InvalidArgumentError( + "pool_size_bytes must be greater than zero"); + } + + const int64_t page_size_val = sysconf(_SC_PAGESIZE); + if (page_size_val <= 0) { + return absl::InternalError("sysconf(_SC_PAGESIZE) failed"); + } + const size_t page_size = static_cast(page_size_val); + const uintptr_t address_value = reinterpret_cast(mapped_address); + if (address_value % page_size != 0) { + return absl::InvalidArgumentError( + absl::StrCat("mapped_address=", address_value, + " must be aligned to system page size ", page_size)); + } + if (pool_size_bytes % page_size != 0) { + return absl::InvalidArgumentError( + absl::StrCat("pool_size_bytes=", pool_size_bytes, + " must be aligned to system page size ", page_size)); + } + if (pool_size_bytes > std::numeric_limits::max() - address_value) { + return absl::InvalidArgumentError( + "mapped_address + pool_size_bytes overflows the process address space"); + } + + xla::PjRtClient* client = GetPjRtClient(); + if (client == nullptr) { + return absl::FailedPreconditionError( + "KVCacheManagerBase has no active PJRT client for DMA mapping"); + } + + { + absl::MutexLock lock(external_mapping_mu_); + if (external_mapping_phase_ != MappingPhase::kUnmapped) { + return absl::FailedPreconditionError( + "shared memory is already mapped or registration is in progress"); + } + external_mapping_phase_ = MappingPhase::kMapping; + } + + const auto reset_mapping_reservation = [this]() { + absl::MutexLock lock(external_mapping_mu_); + external_mapping_phase_ = MappingPhase::kUnmapped; + }; + + const absl::Status status = client->DmaMap(mapped_address, pool_size_bytes); + if (!status.ok()) { + reset_mapping_reservation(); + return status; + } + + { + absl::MutexLock lock(external_mapping_mu_); + external_mapped_address_ = mapped_address; + external_mapped_size_ = pool_size_bytes; + external_mapping_phase_ = MappingPhase::kMapped; + } + return absl::OkStatus(); +} + +absl::Status KVCacheManagerBase::UnmapSharedMemory() { + void* mapped_address = nullptr; + std::vector copies_to_await; + absl::Status copy_status = absl::OkStatus(); + { + absl::MutexLock lock(external_mapping_mu_); + if (external_mapping_phase_ == MappingPhase::kUnmapped) { + return absl::FailedPreconditionError("shared memory is not mapped"); + } + if (external_mapping_phase_ == MappingPhase::kMapping) { + return absl::FailedPreconditionError( + "shared memory mapping is still in progress"); + } + if (external_mapping_phase_ == MappingPhase::kUnmapping) { + return absl::FailedPreconditionError( + "shared memory is already being unmapped"); + } + external_mapping_phase_ = MappingPhase::kUnmapping; + + // Moving to kUnmapping stops new leases, but submissions that acquired a + // lease before that are still racing to register their futures. Wait for + // them, otherwise the snapshot below would miss a live transfer and + // DmaUnmap would run underneath it. + external_mapping_mu_.Await(absl::Condition( + +[](int64_t* pending) { return *pending == 0; }, + &pending_external_copies_)); + + mapped_address = external_mapped_address_; + copies_to_await = std::move(in_flight_external_copies_); + in_flight_external_copies_.clear(); + external_copy_gc_watermark_ = kMinExternalCopyGcWatermark; + copy_status = std::move(deferred_external_copy_error_); + deferred_external_copy_error_ = absl::OkStatus(); + } + + auto first_error = [](absl::Status first, const absl::Status& next) { + return first.ok() ? next : first; + }; + + for (raiden::PjRtCopyFuture& future : copies_to_await) { + copy_status = first_error(std::move(copy_status), future.Await()); + } + + xla::PjRtClient* client = GetPjRtClient(); + absl::Status status = absl::OkStatus(); + if (client != nullptr) { + status = client->DmaUnmap(mapped_address); + } else { + status = absl::InternalError("PJRT client disappeared before DmaUnmap"); + } + + if (!status.ok()) { + absl::MutexLock lock(external_mapping_mu_); + deferred_external_copy_error_ = + first_error(std::move(deferred_external_copy_error_), copy_status); + external_mapping_phase_ = MappingPhase::kMapped; + return status; + } + + { + absl::MutexLock lock(external_mapping_mu_); + external_mapped_address_ = nullptr; + external_mapped_size_ = 0; + deferred_external_copy_error_ = absl::OkStatus(); + external_mapping_phase_ = MappingPhase::kUnmapped; + } + return copy_status; +} + +bool KVCacheManagerBase::is_shared_memory_mapped() const { + absl::MutexLock lock(external_mapping_mu_); + return external_mapping_phase_ == MappingPhase::kMapped; +} + +absl::StatusOr +KVCacheManagerBase::AcquireExternalCopyLease() { + absl::MutexLock lock(external_mapping_mu_); + if (external_mapping_phase_ != MappingPhase::kMapped) { + return absl::FailedPreconditionError( + "shared memory must be mapped before submitting a copy"); + } + ++pending_external_copies_; + return ExternalCopyLease(this); +} + +void KVCacheManagerBase::ReleaseExternalCopyLease( + std::optional future) { + absl::MutexLock lock(external_mapping_mu_); + if (future.has_value()) { + AddExternalCopyLocked(*std::move(future)); + } + --pending_external_copies_; +} + +absl::Status KVCacheManagerBase::ValidateExternalRange(const void* address, + size_t size) const { + absl::MutexLock lock(external_mapping_mu_); + return ValidateExternalRangeLocked(address, size); +} + +absl::Status KVCacheManagerBase::ValidateExternalRangeLocked( + const void* address, size_t size) const { + if (external_mapping_phase_ != MappingPhase::kMapped) { + return absl::FailedPreconditionError( + "shared memory must be mapped before submitting a copy"); + } + if (address == nullptr) { + return absl::InvalidArgumentError("transfer address must be non-null"); + } + const uintptr_t begin = reinterpret_cast(address); + if (size > std::numeric_limits::max() - begin) { + return absl::InvalidArgumentError( + "transfer address + size overflows the process address space"); + } + const uintptr_t pool_begin = + reinterpret_cast(external_mapped_address_); + const uintptr_t pool_end = pool_begin + external_mapped_size_; + if (begin < pool_begin || begin + size > pool_end) { + return absl::OutOfRangeError(absl::StrCat( + "transfer range [", begin, ", ", begin + size, + ") is outside the DMA mapped pool [", pool_begin, ", ", pool_end, + "); host memory must be allocated from the registered shared memory " + "pool")); + } + return absl::OkStatus(); +} + +void KVCacheManagerBase::AddExternalCopyLocked( + raiden::PjRtCopyFuture future) { + if (in_flight_external_copies_.size() >= external_copy_gc_watermark_) { + CollectFinishedExternalCopiesLocked(); + // Next sweep once the list has doubled, so the amortized cost of tracking + // a copy stays constant even when nothing completes. + external_copy_gc_watermark_ = std::max( + kMinExternalCopyGcWatermark, 2 * in_flight_external_copies_.size()); + } + in_flight_external_copies_.push_back(std::move(future)); +} + +void KVCacheManagerBase::CollectFinishedExternalCopiesLocked() { + size_t write_idx = 0; + for (size_t i = 0; i < in_flight_external_copies_.size(); ++i) { + if (in_flight_external_copies_[i].IsReady()) { + absl::Status err = in_flight_external_copies_[i].PollError(); + if (deferred_external_copy_error_.ok() && !err.ok()) { + deferred_external_copy_error_ = std::move(err); + } + } else { + if (write_idx != i) { + in_flight_external_copies_[write_idx] = + std::move(in_flight_external_copies_[i]); + } + ++write_idx; + } + } + in_flight_external_copies_.resize(write_idx); +} + +absl::Status KVCacheManagerBase::ValidateExternalObjectRequest( + const std::vector& block_ids, size_t num_objects, + int64_t rank_id) const { + if (buffer_holds_.empty()) { + return absl::FailedPreconditionError( + "KVCacheManager has no registered device KV cache"); + } + if (block_ids.empty()) { + return absl::InvalidArgumentError("block_ids must not be empty"); + } + if (block_ids.size() != num_objects) { + return absl::InvalidArgumentError(absl::StrCat( + "block_ids and object_tensors must have the same length; got ", + block_ids.size(), " and ", num_objects)); + } + if (rank_id < 0) { + return absl::InvalidArgumentError("rank_id must be non-negative"); + } + // buffer_holds_ is only populated on the device-backed path, and nothing + // guarantees it has one entry per layer; the submit loop indexes it with + // layer_id, so check the length up front. + if (buffer_holds_.size() < num_layers()) { + return absl::FailedPreconditionError(absl::StrCat( + "KVCacheManager has ", buffer_holds_.size(), + " registered layer buffers but num_layers=", num_layers())); + } + return absl::OkStatus(); +} + +absl::StatusOr +KVCacheManagerBase::CopyExternalObjectBlocks( + const std::vector& block_ids, + const std::vector& host_block_bases, size_t num_ranks, + size_t page_nbytes, int64_t rank_id, bool is_h2d, + std::shared_ptr keep_alive) { + absl::Status request_status = ValidateExternalObjectRequest( + block_ids, host_block_bases.size(), rank_id); + if (!request_status.ok()) { + return request_status; + } + if (page_nbytes == 0) { + return absl::InvalidArgumentError("page_nbytes must be greater than zero"); + } + if (static_cast(rank_id) >= num_ranks) { + return absl::OutOfRangeError(absl::StrCat( + "rank_id=", rank_id, " is outside the object first dimension [0, ", + num_ranks, ")")); + } + // This entry point is shared by the torch and JAX frontends, so it cannot + // lean on any one of them having screened its tensors. Reject null bases + // before the offset arithmetic below: adding an offset to a null pointer is + // undefined behaviour, and a nonzero rank offset would turn it into a + // non-null junk address that the null check in ValidateExternalRangeLocked + // no longer catches. + for (size_t obj_id = 0; obj_id < host_block_bases.size(); ++obj_id) { + if (host_block_bases[obj_id] == nullptr) { + return absl::InvalidArgumentError( + absl::StrCat("host_block_bases[", obj_id, "] is null")); + } + } + + const size_t num_layers_local = num_layers(); + const size_t slice_bytes = slice_byte_size(); + if (slice_bytes > 0 && page_nbytes != slice_bytes) { + return absl::InvalidArgumentError( + absl::StrCat("object page_nbytes=", page_nbytes, + " does not match manager slice_byte_size=", slice_bytes)); + } + + const size_t dev_physical_size = buffer_holds_[0].physical_size; + if (dev_physical_size == 0 || dev_physical_size % page_nbytes != 0) { + return absl::InvalidArgumentError(absl::StrCat( + "device physical size ", dev_physical_size, + " is not divisible by page_nbytes ", page_nbytes)); + } + const size_t num_blocks = dev_physical_size / page_nbytes; + + // MapSharedMemory registered the pool against a single client, the one + // GetPjRtClient() resolves to. DmaMap registration is per client, so a + // layer whose buffer belongs to a different client would hand the DMA + // engine host memory that client never registered. Devices may still + // differ: the registration covers the client, not one device. + xla::PjRtClient* const dma_client = GetPjRtClient(); + if (dma_client == nullptr) { + return absl::FailedPreconditionError( + "KVCacheManager has no active PJRT client for object tensor " + "transfers"); + } + + // The transfer loop below addresses every layer with the same page size and + // uses each layer's single shard. Reject the geometries that would break + // that assumption rather than silently reading the wrong bytes. + for (size_t layer_id = 0; layer_id < num_layers_local; ++layer_id) { + const auto& layer_info = buffer_holds_[layer_id]; + if (layer_info.holds.size() != 1) { + return absl::UnimplementedError(absl::StrCat( + "object tensor transfers require exactly one shard per layer; layer ", + layer_id, " has ", layer_info.holds.size(), " shards")); + } + if (layer_info.physical_size != dev_physical_size) { + return absl::UnimplementedError(absl::StrCat( + "object tensor transfers require a uniform per-layer device size; " + "layer ", + layer_id, " is ", layer_info.physical_size, " bytes but layer 0 is ", + dev_physical_size, " bytes")); + } + // Resolve this layer's client the same way GetPjRtClient() does, so the + // comparison is against the client the mapping was actually made on. + const auto& hold = layer_info.holds[0]; + xla::PjRtClient* layer_client = + hold.device != nullptr ? hold.device->client() : nullptr; + if (layer_client == nullptr && hold.buffer != nullptr) { + layer_client = hold.buffer->client(); + } + if (layer_client == nullptr) { + return absl::FailedPreconditionError(absl::StrCat( + "layer ", layer_id, " has no resolvable PJRT client")); + } + if (layer_client != dma_client) { + return absl::UnimplementedError(absl::StrCat( + "object tensor transfers require every layer to share one PJRT " + "client; layer ", + layer_id, " belongs to a different client than the one the shared " + "memory pool is DMA mapped on")); + } + } + + absl::flat_hash_set unique_blocks; + unique_blocks.reserve(block_ids.size()); + std::vector device_offsets; + device_offsets.reserve(block_ids.size()); + for (size_t i = 0; i < block_ids.size(); ++i) { + const int64_t b = block_ids[i]; + if (b < 0 || static_cast(b) >= num_blocks) { + return absl::OutOfRangeError( + absl::StrCat("block_ids[", i, "]=", b, + " is outside block range [0, ", num_blocks, ")")); + } + if (!unique_blocks.insert(b).second) { + return absl::InvalidArgumentError(absl::StrCat( + "block_ids contains duplicate block ", b, + "; concurrent writes to one block have undefined ordering")); + } + device_offsets.push_back(b * static_cast(page_nbytes)); + } + + const size_t rank_bytes = num_layers_local * page_nbytes; + const size_t host_rank_offset = static_cast(rank_id) * rank_bytes; + std::vector host_layer_offsets; + host_layer_offsets.reserve(num_layers_local); + for (size_t layer_id = 0; layer_id < num_layers_local; ++layer_id) { + host_layer_offsets.push_back(host_rank_offset + layer_id * page_nbytes); + } + + // Pin the DMA registration for the rest of this call. Checking + // is_shared_memory_mapped() here instead would be racy: the pool could be + // unmapped between the check and IssueH2dShard below. + absl::StatusOr lease = AcquireExternalCopyLease(); + if (!lease.ok()) { + return lease.status(); + } + + // Every byte handed to the DMA engine must lie inside the pool that was + // registered with DmaMap. Each object contributes one contiguous run: the + // `num_layers` pages belonging to `rank_id`. + // + // One lock for the whole loop rather than one per object: UnmapSharedMemory + // needs this same mutex to move to kUnmapping, and its Await releases the + // mutex while it drains. Locking per object would let an unmap land between + // two iterations and reject a later object after an earlier one already + // passed. The lock is scoped so it is dropped before `lease` is released, + // which reacquires it. + { + absl::MutexLock lock(external_mapping_mu_); + for (size_t obj_id = 0; obj_id < host_block_bases.size(); ++obj_id) { + const absl::Status range_status = ValidateExternalRangeLocked( + host_block_bases[obj_id] + host_rank_offset, rank_bytes); + if (!range_status.ok()) { + return absl::Status( + range_status.code(), + absl::StrCat("object_tensors[", obj_id, "] is not DMA mappable: ", + range_status.message())); + } + } + } + + std::vector layer_futures; + layer_futures.reserve(num_layers_local); + + const int64_t transfer_size = static_cast(page_nbytes); + for (size_t layer_id = 0; layer_id < num_layers_local; ++layer_id) { + const size_t host_offset = host_layer_offsets[layer_id]; + const auto& shard_hold = buffer_holds_[layer_id].holds[0]; + + absl::StatusOr future; + if (is_h2d) { + std::vector copies; + copies.reserve(host_block_bases.size()); + for (size_t obj_id = 0; obj_id < host_block_bases.size(); ++obj_id) { + copies.push_back(raiden::H2dCopy{ + .src = host_block_bases[obj_id] + host_offset, + .dst_off = device_offsets[obj_id], + .size = transfer_size, + }); + } + future = raiden::IssueH2dShard(shard_hold, copies); + } else { + std::vector copies; + copies.reserve(host_block_bases.size()); + for (size_t obj_id = 0; obj_id < host_block_bases.size(); ++obj_id) { + copies.push_back(raiden::D2hCopy{ + .dst = host_block_bases[obj_id] + host_offset, + .src_off = device_offsets[obj_id], + .size = transfer_size, + }); + } + future = raiden::IssueD2hShard(shard_hold, copies); + } + + if (!future.ok()) { + if (!layer_futures.empty()) { + // Layers [0, layer_id) are already in flight and still reference the + // caller's host storage; hand them to the lease so the mapping + // outlives them even though the overall transfer failed. + raiden::PjRtCopyFuture submitted = + raiden::JoinPjRtCopyFutures(layer_futures); + submitted.AddKeepAlive(keep_alive); + lease->Commit(std::move(submitted)); + } + return future.status(); + } + layer_futures.push_back(std::move(future.value())); + } + + raiden::PjRtCopyFuture joined = raiden::JoinPjRtCopyFutures(layer_futures); + joined.AddKeepAlive(std::move(keep_alive)); + lease->Commit(joined); + return joined; +} + +absl::Status KVCacheManagerBase::TrackExternalCopy( + raiden::PjRtCopyFuture future) { + absl::MutexLock lock(external_mapping_mu_); + if (external_mapping_phase_ != MappingPhase::kMapped) { + return absl::FailedPreconditionError( + "shared memory is not mapped; refusing to track the copy future"); + } + AddExternalCopyLocked(std::move(future)); + return absl::OkStatus(); +} + } // namespace kv_cache } // namespace tpu_raiden diff --git a/tpu_sync/kv_cache/kv_cache_manager_base.h b/tpu_sync/kv_cache/kv_cache_manager_base.h index b6f907c0..123f7817 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_base.h +++ b/tpu_sync/kv_cache/kv_cache_manager_base.h @@ -435,6 +435,148 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase { void SetExternalHostBuffer( const std::vector& buffer_holds); + // Registers an existing whole-pool host virtual memory mapping for TPU DMA. + // The caller owns the memory lifetime and must keep it page-locked until + // UnmapSharedMemory succeeds. + // + // At most one live mapping is supported at a time: calling this again while + // a pool is mapped fails with FailedPrecondition. Mapping a different pool + // requires an UnmapSharedMemory first, which is allowed and may be repeated. + // Registering several disjoint regions simultaneously is not supported; + // PJRT's DmaMap/DmaUnmap are keyed by address and would allow it, but the + // drain and validation bookkeeping here assumes a single pool. + absl::Status MapSharedMemory(void* mapped_address, size_t pool_size_bytes); + + // Drains in-flight copies and releases the shared-memory DMA registration. + // Blocks until every submission that already holds a lease has registered + // its future, then awaits all of them before calling DmaUnmap. + absl::Status UnmapSharedMemory(); + + // Returns true if an external shared memory pool is currently DMA mapped. + // + // This is a point-in-time observation only: a concurrent UnmapSharedMemory + // can invalidate it as soon as the lock is released, so it must NOT be used + // to guard a DMA submission. Use AcquireExternalCopyLease for that. + bool is_shared_memory_mapped() const; + + // Accessor for the PJRT client owning the device buffers. + xla::PjRtClient* GetPjRtClient() const; + + // Keeps the shared-memory DMA registration alive for the duration of a copy + // submission. UnmapSharedMemory will not call DmaUnmap while any lease is + // outstanding, which closes the window between "the pool is mapped" and + // "the resulting future has been registered as in-flight". + // + // Call Commit once the future exists. Destroying an uncommitted lease (an + // early-return error path) simply releases the hold. + class ExternalCopyLease { + public: + ExternalCopyLease() = default; + ~ExternalCopyLease() { Release(std::nullopt); } + + ExternalCopyLease(ExternalCopyLease&& other) noexcept + : manager_(std::exchange(other.manager_, nullptr)) {} + ExternalCopyLease& operator=(ExternalCopyLease&& other) noexcept { + if (this != &other) { + Release(std::nullopt); + manager_ = std::exchange(other.manager_, nullptr); + } + return *this; + } + ExternalCopyLease(const ExternalCopyLease&) = delete; + ExternalCopyLease& operator=(const ExternalCopyLease&) = delete; + + // Registers `future` as in-flight and releases the lease. + void Commit(raiden::PjRtCopyFuture future) { Release(std::move(future)); } + + private: + friend class KVCacheManagerBase; + + explicit ExternalCopyLease(KVCacheManagerBase* manager) + : manager_(manager) {} + + void Release(std::optional future) { + if (manager_ != nullptr) { + std::exchange(manager_, nullptr) + ->ReleaseExternalCopyLease(std::move(future)); + } + } + + KVCacheManagerBase* manager_ = nullptr; + }; + + // Acquires a lease on the shared-memory registration. Fails unless a pool + // is mapped and no unmap is in progress. + absl::StatusOr AcquireExternalCopyLease(); + + // Returns OK iff [address, address + size) lies entirely inside the pool + // that was handed to DmaMap. DMA must never be issued against host memory + // outside the registered pool. + absl::Status ValidateExternalRange(const void* address, size_t size) const; + + // Checks the parts of an object transfer request that do not depend on the + // frontend's tensor types: that device buffers are registered and cover + // every layer, and that block_ids/rank_id are structurally sane. + // + // CopyExternalObjectBlocks calls this itself. Frontends should also call it + // up front, before validating their own tensors, so that a manager with no + // device buffers reports that rather than a confusing shape mismatch. + absl::Status ValidateExternalObjectRequest( + const std::vector& block_ids, size_t num_objects, + int64_t rank_id) const; + + // Issues an object-tensor style transfer between a mapped shared-memory pool + // and the device KV cache. Framework agnostic: torch and JAX frontends + // validate their own tensor types, then delegate here. + // + // Each entry of `host_block_bases` points at the first byte of one + // caller-owned host object laid out as contiguous bytes + // [num_ranks, num_layers, page_nbytes], residing inside the pool registered + // by MapSharedMemory. Only the `rank_id` slice participates: for each layer + // `page_nbytes` bytes are copied to (H2D) or from (D2H) block + // `block_ids[i]` of that layer's device buffer, where `i` indexes + // `host_block_bases`. + // + // `keep_alive` is attached to every future issued, so the caller's host + // storage outlives the transfer even when a later layer fails to submit. + // + // Requires a uniform per-layer device size and exactly one shard per layer; + // hybrid (HMA) and multi-shard geometries are rejected rather than + // mis-addressed. + absl::StatusOr CopyExternalObjectBlocks( + const std::vector& block_ids, + const std::vector& host_block_bases, size_t num_ranks, + size_t page_nbytes, int64_t rank_id, bool is_h2d, + std::shared_ptr keep_alive); + + // Tracks an external in-flight copy future so it is drained before + // unmapping. Fails if the pool is no longer mapped, in which case the + // future is dropped rather than silently queued against a dead mapping. + // + // Prefer AcquireExternalCopyLease + Commit, which additionally keeps the + // registration alive across the submission itself. + absl::Status TrackExternalCopy(raiden::PjRtCopyFuture future); + + // Test helpers to simulate shared memory mapping state in unit tests without + // requiring hardware DMA support. + void SetSharedMemoryMappedForTest(void* mapped_address, + size_t pool_size_bytes) { + absl::MutexLock lock(external_mapping_mu_); + external_mapped_address_ = mapped_address; + external_mapped_size_ = pool_size_bytes; + external_mapping_phase_ = MappingPhase::kMapped; + } + + void ResetSharedMemoryMappedForTest() { + absl::MutexLock lock(external_mapping_mu_); + external_mapped_address_ = nullptr; + external_mapped_size_ = 0; + in_flight_external_copies_.clear(); + external_copy_gc_watermark_ = kMinExternalCopyGcWatermark; + deferred_external_copy_error_ = absl::OkStatus(); + external_mapping_phase_ = MappingPhase::kUnmapped; + } + // Returns the internal LogicalBlockManager. LogicalBlockManager* host_block_manager() const { return host_block_manager_.get(); @@ -633,6 +775,25 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase { std::optional layer_idx = std::nullopt, std::optional shard_idx = std::nullopt); + // Per-layer device buffer holds bundled with the layer's on-device size. + // For uniform models every layer has the same physical_size; for hybrid + // (HMA) models sizes may differ (e.g. mamba conv_state bf16 vs ssm f32). + struct LayerDeviceInfo { + std::vector holds; + // Total on-device bytes for this layer's buffer. Set by the + // device-backed constructor from PjRtBuffer. DMA functions use + // this for per-layer offset and copy-size calculations. + size_t physical_size = 0; + }; + + // Read-only view of the per-layer device geometry. Frontends that reach + // this manager by composition rather than inheritance (see + // KVCacheManagerWithTransfer::base()) need it to slice external host + // tensors against the device buffers. + const std::vector& buffer_holds() const { + return buffer_holds_; + } + bool has_device_buffers() const { return !buffer_holds_.empty(); } void AttachPlaceholderDeviceHoldForTest() { buffer_holds_.emplace_back(); } int parallelism() const { return parallelism_; } @@ -663,15 +824,7 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase { std::unique_ptr host_block_manager_; // Per-layer device buffer holds bundled with the layer's on-device size. - // For uniform models every layer has the same physical_size; for hybrid - // (HMA) models sizes may differ (e.g. mamba conv_state bf16 vs ssm f32). - struct LayerDeviceInfo { - std::vector holds; - // Total on-device bytes for this layer's buffer. Set by the - // device-backed constructor from PjRtBuffer. DMA functions use - // this for per-layer offset and copy-size calculations. - size_t physical_size = 0; - }; + // See the LayerDeviceInfo definition in the public section above. std::vector buffer_holds_; // Pool table. Explicit after RegisterPools; otherwise lazily materialized // implicit pools (one per storage, tag "opaque"). pools_mu_ guards the lazy @@ -906,6 +1059,54 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase { backends_ ABSL_GUARDED_BY(backends_mu_); bool InitializeSingleSecondaryBackend(const BackendConfig& config); + enum class MappingPhase { + kUnmapped, + kMapping, + kMapped, + kUnmapping, + }; + + mutable absl::Mutex external_mapping_mu_; + MappingPhase external_mapping_phase_ ABSL_GUARDED_BY(external_mapping_mu_) = + MappingPhase::kUnmapped; + void* external_mapped_address_ + ABSL_GUARDED_BY(external_mapping_mu_) = nullptr; + size_t external_mapped_size_ ABSL_GUARDED_BY(external_mapping_mu_) = 0; + std::vector in_flight_external_copies_ + ABSL_GUARDED_BY(external_mapping_mu_); + absl::Status deferred_external_copy_error_ + ABSL_GUARDED_BY(external_mapping_mu_) = absl::OkStatus(); + + // Number of submissions holding a lease that have not yet registered their + // future. UnmapSharedMemory waits for this to reach zero so that it cannot + // unmap underneath a copy that is still being issued. + int64_t pending_external_copies_ ABSL_GUARDED_BY(external_mapping_mu_) = 0; + + // in_flight_external_copies_ is swept only once it reaches this size. Each + // sweep polls every entry, so sweeping on every submission would make + // tracking quadratic in the number of concurrent copies; the watermark keeps + // it amortized constant. + static constexpr size_t kMinExternalCopyGcWatermark = 64; + size_t external_copy_gc_watermark_ ABSL_GUARDED_BY(external_mapping_mu_) = + kMinExternalCopyGcWatermark; + + // Releases a lease taken by AcquireExternalCopyLease, registering `future` + // as in-flight when the submission succeeded. + void ReleaseExternalCopyLease(std::optional future); + + absl::Status ValidateExternalRangeLocked(const void* address, + size_t size) const + ABSL_EXCLUSIVE_LOCKS_REQUIRED(external_mapping_mu_); + + // Adds `future` to the in-flight list, sweeping completed entries first if + // the list has grown past the watermark. + void AddExternalCopyLocked(raiden::PjRtCopyFuture future) + ABSL_EXCLUSIVE_LOCKS_REQUIRED(external_mapping_mu_); + + // Drops completed futures from in_flight_external_copies_, latching the + // first error into deferred_external_copy_error_. + void CollectFinishedExternalCopiesLocked() + ABSL_EXCLUSIVE_LOCKS_REQUIRED(external_mapping_mu_); }; } // namespace kv_cache diff --git a/tpu_sync/kv_cache/kv_cache_manager_test.cc b/tpu_sync/kv_cache/kv_cache_manager_test.cc index 71ceb660..6f091660 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_test.cc +++ b/tpu_sync/kv_cache/kv_cache_manager_test.cc @@ -529,6 +529,29 @@ TEST(KVCacheManagerTest, PoolBlockCopiesRejectHostOnlyManager) { EXPECT_THAT(h2d.status().message(), testing::HasSubstr("host-only")); } +// The torch frontend cannot produce a null base, but this entry point is +// shared with JAX, so the base must reject one itself: the offset arithmetic +// downstream is undefined behaviour on a null pointer, and a nonzero rank +// offset would carry it past the null check in ValidateExternalRangeLocked. +TEST(KVCacheManagerTest, CopyExternalObjectBlocksRejectsNullHostBase) { + KVCacheManagerBase manager(/*num_layers=*/1, /*num_shards=*/1, + /*slice_byte_size=*/64, + /*local_port=*/std::nullopt, + /*host_blocks_to_allocate=*/2); + manager.AttachPlaceholderDeviceHoldForTest(); + + const absl::Status status = manager + .CopyExternalObjectBlocks( + /*block_ids=*/{0}, + /*host_block_bases=*/{nullptr}, + /*num_ranks=*/1, /*page_nbytes=*/64, + /*rank_id=*/0, /*is_h2d=*/true, + /*keep_alive=*/nullptr) + .status(); + EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(status.message(), testing::HasSubstr("is null")); +} + // Without RegisterPools the manager exposes one implicit Opaque pool per // storage and reports no explicit pools. TEST(KVCacheManagerTest, ImplicitPoolsMirrorStorages) {