diff --git a/tpu_sync/api/torch/kv_cache_manager.py b/tpu_sync/api/torch/kv_cache_manager.py index be8eb8c5f..129e8d0eb 100644 --- a/tpu_sync/api/torch/kv_cache_manager.py +++ b/tpu_sync/api/torch/kv_cache_manager.py @@ -118,6 +118,7 @@ def __init__( worker_id: Optional[str] = None, enable_shm: bool = False, backend_configs: Optional[Sequence[Any]] = None, + num_staging_layers: Optional[int] = None, ): """Instantiates the TransferEngine-based KVCacheManager. @@ -157,9 +158,25 @@ def __init__( and set `_rank` to this worker's rank in [0, size). Duplicate or missing ranks cause wrong or missed recalls. An invalid config aborts the process (LOG(FATAL)) at registration. + num_staging_layers: Optional number of host staging layer buffers to + allocate, in [1, num_layers]. Layer-major transfers lease these buffers + per layer, bounding host DRAM to num_staging_layers x num_slots x + max_blocks blocks instead of a full per-layer mirror. None keeps one + host buffer per layer. Only valid on the (max_blocks, num_slots) path. """ self._admission_summary = None impl = _torch_impl() + if num_staging_layers is not None: + if host_blocks_to_allocate is not None: + raise ValueError( + "num_staging_layers is not supported with host_blocks_to_allocate." + ) + # Checked here, before the C++ constructor pins the TPU buffers. + if not 1 <= num_staging_layers <= len(kv_caches): + raise ValueError( + f"num_staging_layers={num_staging_layers} must be in" + f" [1, {len(kv_caches)}] (number of kv_caches layers)." + ) if host_blocks_to_allocate is not None: self._impl = impl.KVCacheManager( kv_caches, @@ -195,6 +212,7 @@ def __init__( raiden_controller_address=raiden_controller_address, worker_id=worker_id, enable_shm=enable_shm, + num_staging_layers=num_staging_layers, ) if backend_configs: self._impl.register_kv_backends(list(backend_configs)) diff --git a/tpu_sync/frameworks/torch/kv_cache_manager.cc b/tpu_sync/frameworks/torch/kv_cache_manager.cc index 85e8c35e8..bb1b78725 100644 --- a/tpu_sync/frameworks/torch/kv_cache_manager.cc +++ b/tpu_sync/frameworks/torch/kv_cache_manager.cc @@ -119,7 +119,8 @@ TorchKVCacheManager::TorchKVCacheManager( std::optional host_blocks_to_allocate, bool unsafe_skip_buffer_lock, int parallelism, int64_t node_id, int64_t local_control_port, int64_t max_blocks, int64_t num_slots, double timeout_s, - std::vector kv_caches, bool enable_shm) + std::vector kv_caches, bool enable_shm, + std::optional num_staging_layers) : KVCacheManagerWithTransfer( unpacked.buffers, unpacked.has_logical_metadata ? unpacked.logical_slice_byte_size : 0, @@ -141,7 +142,8 @@ TorchKVCacheManager::TorchKVCacheManager( : unpacked.buffers[0][0] .buffer->GetOnDeviceSizeInBytes() .value_or(0)), - node_id, local_control_port, max_blocks, num_slots, timeout_s), + node_id, local_control_port, max_blocks, num_slots, timeout_s, + /*metrics_collector=*/nullptr, num_staging_layers), kv_caches_(std::move(kv_caches)), // Move the keep-alive refs in AFTER the base ctor has acquired the // buffers; they pin the materialized device buffers for our lifetime. @@ -153,14 +155,15 @@ TorchKVCacheManager::Create(const std::vector& kv_caches, int64_t max_blocks, int64_t num_slots, double timeout_s, bool unsafe_skip_buffer_lock, int parallelism, std::optional listener_port, - bool enable_shm) { + bool enable_shm, + std::optional num_staging_layers) { try { auto mgr = absl::WrapUnique(new TorchKVCacheManager( UnpackLayers(SingleShardLayers(kv_caches)), /*local_port=*/std::nullopt, /*host_blocks_to_allocate=*/std::nullopt, unsafe_skip_buffer_lock, parallelism, node_id, local_control_port, max_blocks, num_slots, - timeout_s, kv_caches, enable_shm)); + timeout_s, kv_caches, enable_shm, num_staging_layers)); if (listener_port) { mgr->listener_ = std::make_unique( mgr.get(), *listener_port); @@ -314,12 +317,14 @@ KVCacheManager::CreateWithNodeId( double timeout_s, bool unsafe_skip_buffer_lock, int parallelism, std::optional listener_port, int raiden_worker_port, std::optional raiden_controller_address, - std::optional worker_id, bool enable_shm) { - ABSL_ASSIGN_OR_RETURN(std::unique_ptr torch_manager, - TorchKVCacheManager::Create( - kv_caches, node_id, local_control_port, max_blocks, - num_slots, timeout_s, unsafe_skip_buffer_lock, - parallelism, listener_port, enable_shm)); + std::optional worker_id, bool enable_shm, + std::optional num_staging_layers) { + ABSL_ASSIGN_OR_RETURN( + std::unique_ptr torch_manager, + TorchKVCacheManager::Create( + kv_caches, node_id, local_control_port, max_blocks, num_slots, + timeout_s, unsafe_skip_buffer_lock, parallelism, listener_port, + enable_shm, num_staging_layers)); auto manager = absl::WrapUnique(new KVCacheManager(std::move(torch_manager))); ABSL_RETURN_IF_ERROR(manager->StartGrpcServer( raiden_worker_port, std::move(raiden_controller_address), diff --git a/tpu_sync/frameworks/torch/kv_cache_manager.h b/tpu_sync/frameworks/torch/kv_cache_manager.h index 658e284b0..40bb311bf 100644 --- a/tpu_sync/frameworks/torch/kv_cache_manager.h +++ b/tpu_sync/frameworks/torch/kv_cache_manager.h @@ -60,7 +60,8 @@ class TorchKVCacheManager : public KVCacheManagerWithTransfer { const std::vector& kv_caches, int64_t node_id, int64_t local_control_port, int64_t max_blocks, int64_t num_slots, double timeout_s, bool unsafe_skip_buffer_lock, int parallelism = 4, - std::optional listener_port = std::nullopt, bool enable_shm = false); + std::optional listener_port = std::nullopt, bool enable_shm = false, + std::optional num_staging_layers = std::nullopt); ~TorchKVCacheManager() override; @@ -129,7 +130,8 @@ class TorchKVCacheManager : public KVCacheManagerWithTransfer { bool unsafe_skip_buffer_lock, int parallelism, int64_t node_id, int64_t local_control_port, int64_t max_blocks, int64_t num_slots, double timeout_s, - std::vector kv_caches, bool enable_shm); + std::vector kv_caches, bool enable_shm, + std::optional num_staging_layers = std::nullopt); std::vector kv_caches_; // Keep-alives for the materialized device buffers backing the manager. @@ -157,8 +159,10 @@ class KVCacheManager { // Creates a transfer-enabled manager over the flat list of per-layer // tensors in |kv_caches| (single shard per layer). |enable_shm| opts this // manager's host buffers into the shared-memory segments named by - // RAIDEN_SHM_KEY; without it the env var is ignored. Failure modes as for - // the cache-only Create(). + // RAIDEN_SHM_KEY; without it the env var is ignored. |num_staging_layers| + // bounds the host mirror to that many host layer buffers (see + // KVCacheManagerBase). Failure modes as for the cache-only Create(), plus + // InvalidArgument for an out-of-range |num_staging_layers|. static absl::StatusOr> CreateWithNodeId( const std::vector& kv_caches, int64_t node_id, int64_t local_control_port, int64_t max_blocks, int64_t num_slots, @@ -167,7 +171,8 @@ class KVCacheManager { int raiden_worker_port = 0, std::optional raiden_controller_address = std::nullopt, std::optional worker_id = std::nullopt, - bool enable_shm = false); + bool enable_shm = false, + std::optional num_staging_layers = std::nullopt); ~KVCacheManager(); diff --git a/tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc b/tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc index 1cf90e5c4..60ecdcfb0 100644 --- a/tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc +++ b/tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc @@ -270,7 +270,8 @@ NB_MODULE(_tpu_raiden_torch, m) { nb::arg("parallelism") = 4, nb::arg("listener_port") = nb::none(), nb::arg("raiden_worker_port") = 0, nb::arg("raiden_controller_address") = nb::none(), - nb::arg("worker_id") = nb::none(), nb::arg("enable_shm") = false) + nb::arg("worker_id") = nb::none(), nb::arg("enable_shm") = false, + nb::arg("num_staging_layers") = nb::none()) .def("node_id", &KVCacheManager::node_id) .def("get_raiden_worker_port", &KVCacheManager::GetRaidenWorkerPort) .def( diff --git a/tpu_sync/kv_cache/kv_cache_manager_base.cc b/tpu_sync/kv_cache/kv_cache_manager_base.cc index 47c99f59a..cdb4bdd21 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_base.cc +++ b/tpu_sync/kv_cache/kv_cache_manager_base.cc @@ -45,7 +45,6 @@ #include "absl/status/statusor.h" #include "absl/strings/ascii.h" #include "absl/strings/match.h" -#include "absl/strings/numbers.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" #include "absl/strings/string_view.h" @@ -415,17 +414,24 @@ std::vector BuildDefaultShards(size_t num_shards) { return shards; } -// TPU_RAIDEN_NUM_STAGING_LAYERS: number of host layer buffers to allocate. -// Unset, unparsable, zero, or >= |num_layers| keeps the full per-layer mirror. -size_t ResolveEnvNumStagingLayers(size_t num_layers) { - const char* raw_env = std::getenv("TPU_RAIDEN_NUM_STAGING_LAYERS"); - size_t requested = 0; - if (raw_env == nullptr || !absl::SimpleAtoi(raw_env, &requested) || - requested == 0 || requested >= num_layers) { +// Number of host layer buffers to allocate: |num_staging_layers| when set, +// which must lie in [1, |num_layers|]; otherwise the full per-layer mirror. +size_t ResolveNumStagingLayers(size_t num_layers, + std::optional num_staging_layers) { + if (!num_staging_layers.has_value()) { return num_layers; } - LOG(INFO) << "KVCacheManagerBase: allocating " << requested << "/" - << num_layers << " host staging layer buffers"; + if (*num_staging_layers < 1 || + static_cast(*num_staging_layers) > num_layers) { + throw std::invalid_argument( + absl::StrCat("num_staging_layers (", *num_staging_layers, + ") must be in [1, ", num_layers, "]")); + } + const size_t requested = static_cast(*num_staging_layers); + if (requested < num_layers) { + LOG(INFO) << "KVCacheManagerBase: allocating " << requested << "/" + << num_layers << " host staging layer buffers"; + } return requested; } @@ -439,7 +445,7 @@ absl::Status CheckAllLayersHaveHostBuffers( return absl::FailedPreconditionError(absl::StrCat( direction, " of every layer needs a host buffer per layer, but only ", num_host_layers, " of ", num_layers, - " are allocated (TPU_RAIDEN_NUM_STAGING_LAYERS)")); + " are allocated (num_staging_layers)")); } return absl::OkStatus(); } @@ -453,7 +459,8 @@ KVCacheManagerBase::KVCacheManagerBase( HostBufferAllocator host_allocator, std::optional bind_ip, std::optional logical_slice_byte_size, std::vector logical_dimensions, - std::optional logical_physical_size) + std::optional logical_physical_size, + std::optional num_staging_layers) : RaidenManagerBase( layer_buffers.size(), raiden::DetectShards(layer_buffers), logical_slice_byte_size.has_value() @@ -464,6 +471,10 @@ KVCacheManagerBase::KVCacheManagerBase( local_port, parallelism, bind_ip, raiden::DetectNumaNodes(layer_buffers), GetLocalHostNicAddresses()), host_allocator_(host_allocator) { + // Validate before the empty-cache early return so an out-of-range value is + // rejected regardless of how many layers were registered. + const size_t num_host_layers_to_allocate = + ResolveNumStagingLayers(num_layers_, num_staging_layers); if (num_layers_ == 0 || num_shards_ == 0) { return; } @@ -538,9 +549,8 @@ KVCacheManagerBase::KVCacheManagerBase( buffer_holds_.push_back(std::move(device_info)); } absl::StatusOr host_layers = AllocateHostLayers( - num_layers_, num_shards_, - /*num_layers_to_allocate=*/ResolveEnvNumStagingLayers(num_layers_), - num_host_blocks, buffer_holds_, host_allocator); + num_layers_, num_shards_, num_host_layers_to_allocate, num_host_blocks, + buffer_holds_, host_allocator); if (!host_layers.ok()) { // Constructors cannot return a status; a failed host allocation surfaces // the way it always has. @@ -567,17 +577,18 @@ KVCacheManagerBase::KVCacheManagerBase( size_t num_layers, size_t num_shards, size_t slice_byte_size, std::optional local_port, std::optional host_blocks_to_allocate, int parallelism, HostBufferAllocator host_allocator, - std::optional bind_ip) + std::optional bind_ip, std::optional num_staging_layers) : KVCacheManagerBase(num_layers, num_shards, std::vector(num_layers, slice_byte_size), local_port, host_blocks_to_allocate, parallelism, - std::move(host_allocator), std::move(bind_ip)) {} + std::move(host_allocator), std::move(bind_ip), + num_staging_layers) {} KVCacheManagerBase::KVCacheManagerBase( size_t num_layers, size_t num_shards, std::vector slice_byte_sizes, std::optional local_port, std::optional host_blocks_to_allocate, int parallelism, HostBufferAllocator host_allocator, - std::optional bind_ip) + std::optional bind_ip, std::optional num_staging_layers) : RaidenManagerBase(num_layers, BuildDefaultShards(num_shards), slice_byte_sizes.empty() ? 0 : slice_byte_sizes[0], local_port, parallelism, bind_ip), @@ -600,7 +611,8 @@ KVCacheManagerBase::KVCacheManagerBase( absl::StatusOr host_layers = AllocateHostLayers( num_layers_, num_shards_, - /*num_layers_to_allocate=*/ResolveEnvNumStagingLayers(num_layers_), + /*num_layers_to_allocate=*/ + ResolveNumStagingLayers(num_layers_, num_staging_layers), /*num_host_blocks=*/total_blocks, /*device_layers=*/{}, host_allocator); if (!host_layers.ok()) { throw std::runtime_error(std::string(host_layers.status().message())); diff --git a/tpu_sync/kv_cache/kv_cache_manager_base.h b/tpu_sync/kv_cache/kv_cache_manager_base.h index 9f73f2fa4..596ee3fa7 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_base.h +++ b/tpu_sync/kv_cache/kv_cache_manager_base.h @@ -117,7 +117,9 @@ using HostBlockId = int64_t; class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase { public: // Core C++ Constructor wrapping raw PJRT buffers directly (used by JAX and - // PyTorch E2E) + // PyTorch E2E). |num_staging_layers| bounds the host mirror along the layer + // dimension to that many host layer buffers, in [1, num_layers]; unset keeps + // one host buffer per layer. Throws std::invalid_argument when out of range. KVCacheManagerBase( const std::vector>& layer_buffers, std::optional local_port = std::nullopt, @@ -127,26 +129,30 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase { std::optional bind_ip = std::nullopt, std::optional logical_slice_byte_size = std::nullopt, std::vector logical_dimensions = {}, - std::optional logical_physical_size = std::nullopt); + std::optional logical_physical_size = std::nullopt, + std::optional num_staging_layers = std::nullopt); - // Standard CPU-only constructor for remote workers E2E. + // Standard CPU-only constructor for remote workers E2E. |num_staging_layers| + // as in the device-buffer constructor. KVCacheManagerBase(size_t num_layers, size_t num_shards, size_t slice_byte_size, std::optional local_port = std::nullopt, std::optional host_blocks_to_allocate = std::nullopt, int parallelism = 1, HostBufferAllocator host_allocator = nullptr, - std::optional bind_ip = std::nullopt); + std::optional bind_ip = std::nullopt, + std::optional num_staging_layers = std::nullopt); // same as above, but with a vector of slice byte sizes for heterogeneous - // registration. + // registration. |num_staging_layers| as in the device-buffer constructor. KVCacheManagerBase(size_t num_layers, size_t num_shards, std::vector slice_byte_sizes, std::optional local_port = std::nullopt, std::optional host_blocks_to_allocate = std::nullopt, int parallelism = 1, HostBufferAllocator host_allocator = nullptr, - std::optional bind_ip = std::nullopt); + std::optional bind_ip = std::nullopt, + std::optional num_staging_layers = std::nullopt); ~KVCacheManagerBase() override; @@ -796,7 +802,7 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase { int block_id) override; // Returns the number of allocated host layer buffers (equals num_layers() - // unless bounded via TPU_RAIDEN_NUM_STAGING_LAYERS). + // unless bounded via the num_staging_layers constructor argument). size_t num_host_layers() const { return layers_.size(); } // Returns the total number of bytes allocated in host DRAM. diff --git a/tpu_sync/kv_cache/kv_cache_manager_test.cc b/tpu_sync/kv_cache/kv_cache_manager_test.cc index 70a4498b1..d8ea9c6a2 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_test.cc +++ b/tpu_sync/kv_cache/kv_cache_manager_test.cc @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -2406,20 +2407,22 @@ class ScopedEnvVar { struct BoundedStagingTestRig { static constexpr size_t kBlockBytes = 128; + static constexpr size_t kNumLayers = 2; std::unique_ptr client; std::vector> pjrt_buffers; + std::vector> layer_handles; std::unique_ptr manager; - void Init() { + void InitDeviceLayers() { TF_ASSERT_OK_AND_ASSIGN(client, xla::GetXlaPjrtCpuClient(xla::CpuClientOptions())); TF_ASSERT_OK_AND_ASSIGN( xla::PjRtMemorySpace * memory_space, client->addressable_devices()[0]->default_memory_space()); - std::vector> layer_handles(2); - for (size_t layer = 0; layer < 2; ++layer) { + layer_handles.resize(kNumLayers); + for (size_t layer = 0; layer < kNumLayers; ++layer) { std::vector init_data(kBlockBytes, static_cast(0xA0 + layer)); TF_ASSERT_OK_AND_ASSIGN( @@ -2435,22 +2438,31 @@ struct BoundedStagingTestRig { pjrt_buffers.push_back(std::move(buf)); layer_handles[layer].push_back(std::move(handle)); } + } - manager = std::make_unique( + std::unique_ptr CreateManager( + std::optional num_staging_layers) { + return std::make_unique( layer_handles, /*local_port=*/std::nullopt, /*host_blocks_to_allocate=*/1, /*unsafe_skip_buffer_lock=*/true, /*parallelism=*/1, /*host_allocator=*/nullptr, /*bind_ip=*/std::nullopt, - /*logical_slice_byte_size=*/kBlockBytes); + /*logical_slice_byte_size=*/kBlockBytes, + /*logical_dimensions=*/std::vector{}, + /*logical_physical_size=*/std::nullopt, num_staging_layers); + } + + void Init(int num_staging_layers) { + ASSERT_NO_FATAL_FAILURE(InitDeviceLayers()); + manager = CreateManager(num_staging_layers); } }; TEST(KVCacheManagerTest, AllLayerDispatchWithBoundedStagingFailsImmediately) { - ScopedEnvVar staging_env("TPU_RAIDEN_NUM_STAGING_LAYERS", "1"); ScopedEnvVar async_env("RAIDEN_ENABLE_ASYNC_DISPATCH", "1"); BoundedStagingTestRig rig; - ASSERT_NO_FATAL_FAILURE(rig.Init()); + ASSERT_NO_FATAL_FAILURE(rig.Init(/*num_staging_layers=*/1)); ASSERT_EQ(rig.manager->num_host_layers(), 1); uint8_t* host_ptr = @@ -2462,22 +2474,22 @@ TEST(KVCacheManagerTest, AllLayerDispatchWithBoundedStagingFailsImmediately) { // without issuing layer 0's copy or queueing a background task. EXPECT_THAT(rig.manager->H2dSyncDispatch({0}, {0}, {1}), StatusIs(absl::StatusCode::kFailedPrecondition, - HasSubstr("TPU_RAIDEN_NUM_STAGING_LAYERS"))); + HasSubstr("num_staging_layers"))); EXPECT_THAT(rig.manager->H2dSyncDispatch(), StatusIs(absl::StatusCode::kFailedPrecondition, - HasSubstr("TPU_RAIDEN_NUM_STAGING_LAYERS"))); + HasSubstr("num_staging_layers"))); EXPECT_THAT(rig.manager->D2hSyncDispatch({0}, {0}, {1}), StatusIs(absl::StatusCode::kFailedPrecondition, - HasSubstr("TPU_RAIDEN_NUM_STAGING_LAYERS"))); + HasSubstr("num_staging_layers"))); EXPECT_THAT(rig.manager->D2hSyncDispatch(), StatusIs(absl::StatusCode::kFailedPrecondition, - HasSubstr("TPU_RAIDEN_NUM_STAGING_LAYERS"))); + HasSubstr("num_staging_layers"))); EXPECT_THAT(rig.manager->H2d({0}, {0}, {1}), StatusIs(absl::StatusCode::kFailedPrecondition, - HasSubstr("TPU_RAIDEN_NUM_STAGING_LAYERS"))); + HasSubstr("num_staging_layers"))); EXPECT_THAT(rig.manager->D2h({0}, {0}, {1}), StatusIs(absl::StatusCode::kFailedPrecondition, - HasSubstr("TPU_RAIDEN_NUM_STAGING_LAYERS"))); + HasSubstr("num_staging_layers"))); // Host buffer 0 was never overwritten by a partial D2H of layer 0, and // device layer 0 was never overwritten by a partial H2D of 0x55. @@ -2496,11 +2508,10 @@ TEST(KVCacheManagerTest, AllLayerDispatchWithBoundedStagingFailsImmediately) { } TEST(KVCacheManagerTest, SingleLayerDispatchWithBoundedStagingSucceeds) { - ScopedEnvVar staging_env("TPU_RAIDEN_NUM_STAGING_LAYERS", "1"); ScopedEnvVar async_env("RAIDEN_ENABLE_ASYNC_DISPATCH", "1"); BoundedStagingTestRig rig; - ASSERT_NO_FATAL_FAILURE(rig.Init()); + ASSERT_NO_FATAL_FAILURE(rig.Init(/*num_staging_layers=*/1)); ASSERT_EQ(rig.manager->num_host_layers(), 1); uint8_t* host_ptr = @@ -2551,6 +2562,34 @@ TEST(KVCacheManagerTest, SingleLayerDispatchWithBoundedStagingSucceeds) { std::vector(BoundedStagingTestRig::kBlockBytes, 0xB1)); } +TEST(KVCacheManagerTest, NumStagingLayersBoundsHostLayerBuffers) { + BoundedStagingTestRig rig; + ASSERT_NO_FATAL_FAILURE(rig.InitDeviceLayers()); + + // Unset keeps one host buffer per layer; a value in [1, num_layers] + // allocates exactly that many. + EXPECT_EQ(rig.CreateManager(std::nullopt)->num_host_layers(), + BoundedStagingTestRig::kNumLayers); + EXPECT_EQ( + rig.CreateManager(BoundedStagingTestRig::kNumLayers)->num_host_layers(), + BoundedStagingTestRig::kNumLayers); + std::unique_ptr bounded = rig.CreateManager(1); + EXPECT_EQ(bounded->num_layers(), BoundedStagingTestRig::kNumLayers); + EXPECT_EQ(bounded->num_host_layers(), 1); + EXPECT_NE(bounded->GetHostPointer(/*layer_idx=*/0, /*shard_idx=*/0), nullptr); + EXPECT_EQ(bounded->GetHostPointer(/*layer_idx=*/1, /*shard_idx=*/0), nullptr); + + // Out-of-range values are rejected rather than silently clamped. + EXPECT_THROW(rig.CreateManager(0), std::invalid_argument); + EXPECT_THROW(rig.CreateManager(-1), std::invalid_argument); + EXPECT_THROW(rig.CreateManager(BoundedStagingTestRig::kNumLayers + 1), + std::invalid_argument); + + // Validation also runs for an empty cache, which otherwise returns early. + rig.layer_handles.clear(); + EXPECT_THROW(rig.CreateManager(1), std::invalid_argument); +} + } // namespace } // namespace kv_cache } // namespace tpu_raiden diff --git a/tpu_sync/kv_cache/kv_cache_manager_with_transfer.cc b/tpu_sync/kv_cache/kv_cache_manager_with_transfer.cc index 8501ef140..d87ae6783 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_with_transfer.cc +++ b/tpu_sync/kv_cache/kv_cache_manager_with_transfer.cc @@ -521,7 +521,8 @@ KVCacheManagerWithTransfer::KVCacheManagerWithTransfer( std::optional host_blocks_to_allocate, bool unsafe_skip_buffer_lock, int parallelism, HostBufferAllocator host_allocator, int64_t node_id, int64_t local_control_port, int64_t max_blocks, int64_t num_slots, - double timeout_s, std::shared_ptr metrics_collector) + double timeout_s, std::shared_ptr metrics_collector, + std::optional num_staging_layers) : KVCacheManagerWithTransfer( std::make_unique( layer_buffers, local_port, @@ -532,7 +533,8 @@ KVCacheManagerWithTransfer::KVCacheManagerWithTransfer( : std::nullopt, dimensions, physical_size > 0 ? std::make_optional(physical_size) - : std::nullopt), + : std::nullopt, + num_staging_layers), node_id, local_control_port, max_blocks, num_slots, timeout_s, std::move(metrics_collector)) {} diff --git a/tpu_sync/kv_cache/kv_cache_manager_with_transfer.h b/tpu_sync/kv_cache/kv_cache_manager_with_transfer.h index 840d01954..1c581138d 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_with_transfer.h +++ b/tpu_sync/kv_cache/kv_cache_manager_with_transfer.h @@ -143,7 +143,8 @@ class KVCacheManagerWithTransfer { int parallelism, HostBufferAllocator host_allocator, int64_t node_id = 0, int64_t local_control_port = -1, int64_t max_blocks = 0, int64_t num_slots = 0, double timeout_s = 120.0, - std::shared_ptr metrics_collector = nullptr); + std::shared_ptr metrics_collector = nullptr, + std::optional num_staging_layers = std::nullopt); // Metadata-based constructor for a heterogeneous registration: // slice_byte_sizes carries one stride per block array, in registration diff --git a/tpu_sync/kv_cache/kv_cache_manager_with_transfer_send_drain_test.cc b/tpu_sync/kv_cache/kv_cache_manager_with_transfer_send_drain_test.cc index 331de88b9..119f817fd 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_with_transfer_send_drain_test.cc +++ b/tpu_sync/kv_cache/kv_cache_manager_with_transfer_send_drain_test.cc @@ -20,7 +20,6 @@ #include #include #include -#include #include #include #include @@ -79,11 +78,13 @@ constexpr double kTimeoutS = 0.05; // A producer whose device-to-host copies complete when the test says so. class TestManager : public KVCacheManagerWithTransfer { public: - explicit TestManager(size_t num_layers, double timeout_s = kTimeoutS) - : KVCacheManagerWithTransfer(std::make_unique(num_layers, this), - /*node_id=*/0, - /*local_control_port=*/-1, /*max_blocks=*/1, - /*num_slots=*/kSlots, timeout_s) {} + explicit TestManager(size_t num_layers, double timeout_s = kTimeoutS, + std::optional num_staging_layers = std::nullopt) + : KVCacheManagerWithTransfer( + std::make_unique(num_layers, num_staging_layers, this), + /*node_id=*/0, + /*local_control_port=*/-1, /*max_blocks=*/1, + /*num_slots=*/kSlots, timeout_s) {} // Serves a pull for `uuid` the way ProcessPullStream does once the // consumer is acknowledged: the push runs on this thread and returns @@ -243,12 +244,15 @@ class TestManager : public KVCacheManagerWithTransfer { private: class TestBase : public kv_cache::KVCacheManagerBase { public: - TestBase(size_t num_layers, TestManager* owner) + TestBase(size_t num_layers, std::optional num_staging_layers, + TestManager* owner) : kv_cache::KVCacheManagerBase(num_layers, /*num_shards=*/1, std::vector(num_layers, 128), /*local_port=*/std::nullopt, /*host_blocks_to_allocate=*/kSlots, - /*parallelism=*/1, nullptr), + /*parallelism=*/1, nullptr, + /*bind_ip=*/std::nullopt, + num_staging_layers), owner_(owner) {} absl::StatusOr D2hSyncDispatch( @@ -317,11 +321,13 @@ class TestManager : public KVCacheManagerWithTransfer { // A consumer whose host-to-device copies complete when the test says so. class RecvTestManager : public KVCacheManagerWithTransfer { public: - explicit RecvTestManager(size_t num_layers, double timeout_s = 5.0) - : KVCacheManagerWithTransfer(std::make_unique(num_layers, this), - /*node_id=*/0, - /*local_control_port=*/-1, /*max_blocks=*/1, - /*num_slots=*/kSlots, timeout_s) {} + explicit RecvTestManager(size_t num_layers, double timeout_s = 5.0, + std::optional num_staging_layers = std::nullopt) + : KVCacheManagerWithTransfer( + std::make_unique(num_layers, num_staging_layers, this), + /*node_id=*/0, + /*local_control_port=*/-1, /*max_blocks=*/1, + /*num_slots=*/kSlots, timeout_s) {} void AddRecv(const std::string& req_id, uint64_t uuid, int32_t blocks_per_layer = 1, @@ -418,12 +424,15 @@ class RecvTestManager : public KVCacheManagerWithTransfer { private: class RecvBase : public kv_cache::KVCacheManagerBase { public: - RecvBase(size_t num_layers, RecvTestManager* owner) + RecvBase(size_t num_layers, std::optional num_staging_layers, + RecvTestManager* owner) : kv_cache::KVCacheManagerBase(num_layers, /*num_shards=*/1, std::vector(num_layers, 128), /*local_port=*/std::nullopt, /*host_blocks_to_allocate=*/kSlots, - /*parallelism=*/1, nullptr), + /*parallelism=*/1, nullptr, + /*bind_ip=*/std::nullopt, + num_staging_layers), owner_(owner) {} absl::StatusOr H2dSyncDispatch( @@ -458,20 +467,6 @@ class RecvTestManager : public KVCacheManagerWithTransfer { absl::Notification release_dispatch_; }; -// Sets environment variable |name| to |value| until destroyed. The staging -// geometry is read once when a manager is constructed, so the guard only has -// to outlive construction. -class ScopedEnv { - public: - ScopedEnv(const char* name, const char* value) : name_(name) { - setenv(name, value, /*overwrite=*/1); - } - ~ScopedEnv() { unsetenv(name_); } - - private: - const char* name_; -}; - class DispatchReleaseGuard { public: explicit DispatchReleaseGuard(RecvTestManager* manager) : manager_(manager) {} @@ -1278,9 +1273,9 @@ TEST(RecvLifecycleTest, EXPECT_EQ(receiver.base()->GetBlockHostPointer(0, 0, 3)[0], 0x3C); } -// StagingLayerManager: TPU_RAIDEN_NUM_STAGING_LAYERS=4 leaves four host layer -// buffers for a 64-layer cache. Layers lease them in arrival order, a fifth -// layer waits for a release, and a released buffer goes to the waiting layer +// StagingLayerManager: num_staging_layers=4 leaves four host layer buffers +// for a 64-layer cache. Layers lease them in arrival order, a fifth layer +// waits for a release, and a released buffer goes to the waiting layer // whichever layer freed it. TEST(RecvLifecycleTest, StagingLayerManagerCyclesHostLayerBuffers) { constexpr size_t kNumLayers = 64; @@ -1289,8 +1284,7 @@ TEST(RecvLifecycleTest, StagingLayerManagerCyclesHostLayerBuffers) { RecvTestManager(kNumLayers).base()->GetAllocatedHostDramBytes(); ASSERT_GT(full_dram_bytes, 0); - ScopedEnv staging_layers("TPU_RAIDEN_NUM_STAGING_LAYERS", "4"); - RecvTestManager consumer(kNumLayers, /*timeout_s=*/10.0); + RecvTestManager consumer(kNumLayers, /*timeout_s=*/10.0, kNumStagingLayers); kv_cache::KVCacheManagerBase* base = consumer.base(); EXPECT_EQ(base->num_host_layers(), kNumStagingLayers); EXPECT_EQ(base->GetAllocatedHostDramBytes(), @@ -1373,8 +1367,7 @@ TEST(RecvLifecycleTest, StagingLayerManagerReleasesAllResourcesAcrossFailureModes) { constexpr size_t kNumLayers = 16; constexpr size_t kNumStagingLayers = 4; - ScopedEnv staging_layers("TPU_RAIDEN_NUM_STAGING_LAYERS", "4"); - RecvTestManager consumer(kNumLayers, /*timeout_s=*/10.0); + RecvTestManager consumer(kNumLayers, /*timeout_s=*/10.0, kNumStagingLayers); kv_cache::KVCacheManagerBase* base = consumer.base(); const size_t block_bytes = base->block_bytes(0); @@ -1479,9 +1472,9 @@ TEST(RecvLifecycleTest, EXPECT_FALSE(consumer.has_recv(203)); } -// Sender-side StagingLayerManager: with TPU_RAIDEN_NUM_STAGING_LAYERS=4 on a -// 64-layer cache, StartPush leases the four host layer buffers for layers 0..3 -// and waits before issuing layer 4's D2H copy. Completing a layer's H2H push +// Sender-side StagingLayerManager: with num_staging_layers=4 on a 64-layer +// cache, StartPush leases the four host layer buffers for layers 0..3 and +// waits before issuing layer 4's D2H copy. Completing a layer's H2H push // returns its host layer buffer to the manager so the next waiting layer can // D2H into it and H2H out of it. TEST(SendLifecycleTest, StagingLayerManagerCyclesHostLayerBuffers) { @@ -1491,8 +1484,7 @@ TEST(SendLifecycleTest, StagingLayerManagerCyclesHostLayerBuffers) { TestManager(kNumLayers).base()->GetAllocatedHostDramBytes(); ASSERT_GT(full_dram_bytes, 0); - ScopedEnv staging_layers("TPU_RAIDEN_NUM_STAGING_LAYERS", "4"); - TestManager producer(kNumLayers, /*timeout_s=*/10.0); + TestManager producer(kNumLayers, /*timeout_s=*/10.0, kNumStagingLayers); producer.SetManualH2h(true); kv_cache::KVCacheManagerBase* base = producer.base(); EXPECT_EQ(base->num_host_layers(), kNumStagingLayers); @@ -1588,8 +1580,7 @@ TEST(SendLifecycleTest, StagingLayerManagerReleasesAllResourcesAcrossFailureModes) { constexpr size_t kNumLayers = 16; constexpr size_t kNumStagingLayers = 4; - ScopedEnv staging_layers("TPU_RAIDEN_NUM_STAGING_LAYERS", "4"); - TestManager producer(kNumLayers, /*timeout_s=*/10.0); + TestManager producer(kNumLayers, /*timeout_s=*/10.0, kNumStagingLayers); producer.SetManualH2h(true); kv_cache::KVCacheManagerBase* base = producer.base(); const size_t block_bytes = base->block_bytes(0); diff --git a/tpu_sync/kv_cache/pool_staging_test.cc b/tpu_sync/kv_cache/pool_staging_test.cc index 43206a967..5333b391a 100644 --- a/tpu_sync/kv_cache/pool_staging_test.cc +++ b/tpu_sync/kv_cache/pool_staging_test.cc @@ -19,7 +19,6 @@ #include #include -#include #include #include #include @@ -45,10 +44,13 @@ using ::absl_testing::StatusIs; class StagingTestManager : public KVCacheManagerBase { public: StagingTestManager(size_t num_layers, size_t num_shards, - size_t slice_byte_size, int host_blocks) + size_t slice_byte_size, int host_blocks, + std::optional num_staging_layers = std::nullopt) : KVCacheManagerBase(num_layers, num_shards, slice_byte_size, /*local_port=*/std::nullopt, - std::make_optional(host_blocks)) { + std::make_optional(host_blocks), /*parallelism=*/1, + /*host_allocator=*/nullptr, /*bind_ip=*/std::nullopt, + num_staging_layers) { buffer_holds_.resize(num_layers); for (size_t l = 0; l < num_layers; ++l) { buffer_holds_[l].holds.resize(num_shards); @@ -526,39 +528,13 @@ TEST(PoolStagingTest, PoolHostBaseAddrs) { EXPECT_TRUE(addrs.host_base_addrs().empty()); } -class ScopedEnvVar { - public: - ScopedEnvVar(const char* name, const char* value) : name_(name) { - const char* prev = std::getenv(name); - if (prev != nullptr) { - prev_ = prev; - } - setenv(name, value, /*overwrite=*/1); - } - - ~ScopedEnvVar() { - if (prev_.has_value()) { - setenv(name_, prev_->c_str(), /*overwrite=*/1); - } else { - unsetenv(name_); - } - } - - ScopedEnvVar(const ScopedEnvVar&) = delete; - ScopedEnvVar& operator=(const ScopedEnvVar&) = delete; - - private: - const char* name_; - std::optional prev_; -}; - TEST(PoolStagingTest, BoundedHostLayersPoolStaging) { - ScopedEnvVar staging_env("TPU_RAIDEN_NUM_STAGING_LAYERS", "2"); constexpr int64_t kStride = 64; constexpr int64_t kNumBlocks = 16; constexpr size_t kNumLayers = 4; StagingTestManager manager(/*num_layers=*/kNumLayers, /*num_shards=*/1, - /*slice_byte_size=*/kStride, /*host_blocks=*/1); + /*slice_byte_size=*/kStride, /*host_blocks=*/1, + /*num_staging_layers=*/2); EXPECT_EQ(manager.num_layers(), kNumLayers); EXPECT_EQ(manager.num_host_layers(), 2u); diff --git a/tpu_sync/kv_cache/reshard_receive_session_test.cc b/tpu_sync/kv_cache/reshard_receive_session_test.cc index fde4c7ae7..f79a3f2b1 100644 --- a/tpu_sync/kv_cache/reshard_receive_session_test.cc +++ b/tpu_sync/kv_cache/reshard_receive_session_test.cc @@ -16,7 +16,6 @@ #include // NOLINT(build/c++11) #include -#include #include #include #include @@ -60,32 +59,6 @@ using ::tpu_sync::rpc::MEMORY_TYPE_HBM; using ::tpu_sync::rpc::ShardPushEntryProto; using ::tpu_sync::rpc::StartTransferRequest; -class ScopedEnvVar { - public: - ScopedEnvVar(const char* name, const char* value) : name_(name) { - const char* prev = std::getenv(name); - if (prev != nullptr) { - prev_ = prev; - } - setenv(name, value, /*overwrite=*/1); - } - - ~ScopedEnvVar() { - if (prev_.has_value()) { - setenv(name_, prev_->c_str(), /*overwrite=*/1); - } else { - unsetenv(name_); - } - } - - ScopedEnvVar(const ScopedEnvVar&) = delete; - ScopedEnvVar& operator=(const ScopedEnvVar&) = delete; - - private: - const char* name_; - std::optional prev_; -}; - kv_cache::KVCacheManagerBase MakeTestBase() { return kv_cache::KVCacheManagerBase( /*num_layers=*/1, /*num_shards=*/1, @@ -481,12 +454,13 @@ TEST(ReshardReceiveSessionTest, GroupScopesDestinationBlocks) { } TEST(ReshardReceiveSessionTest, BoundedStagingSlotAcquisitionAndRelease) { - ScopedEnvVar staging_env("TPU_RAIDEN_NUM_STAGING_LAYERS", "1"); kv_cache::KVCacheManagerBase base( /*num_layers=*/2, /*num_shards=*/1, /*slice_byte_size=*/128, /*local_port=*/std::nullopt, - /*host_blocks_to_allocate=*/std::make_optional(4)); + /*host_blocks_to_allocate=*/std::make_optional(4), + /*parallelism=*/1, /*host_allocator=*/nullptr, + /*bind_ip=*/std::nullopt, /*num_staging_layers=*/1); base.AttachPlaceholderDeviceHoldForTest(); ASSERT_EQ(base.num_host_layers(), 1u); ASSERT_EQ(base.num_layers(), 2u);