From 3d98f15362b6d288678ffb65ad0dc02fee0ba052 Mon Sep 17 00:00:00 2001 From: justinlu Date: Thu, 29 Oct 2026 14:15:43 -0700 Subject: [PATCH] Support bounded host layer staging for pool reshard sends PiperOrigin-RevId: 1009957389 --- tpu_sync/kv_cache/BUILD | 1 + tpu_sync/kv_cache/kv_cache_manager_base.cc | 5 +- tpu_sync/kv_cache/kv_cache_manager_base.h | 3 +- .../kv_cache_manager_with_transfer.cc | 20 +- tpu_sync/kv_cache/reshard_send_session.cc | 187 ++++++++++++++++-- tpu_sync/kv_cache/reshard_send_session.h | 23 ++- .../kv_cache/reshard_send_session_test.cc | 79 ++++++++ 7 files changed, 290 insertions(+), 28 deletions(-) diff --git a/tpu_sync/kv_cache/BUILD b/tpu_sync/kv_cache/BUILD index 489f83d52..ba3287c9f 100644 --- a/tpu_sync/kv_cache/BUILD +++ b/tpu_sync/kv_cache/BUILD @@ -484,6 +484,7 @@ cc_test( "@com_google_absl//absl/synchronization", "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest_main", + "@xla//xla/tsl/platform:statusor", ], ) diff --git a/tpu_sync/kv_cache/kv_cache_manager_base.cc b/tpu_sync/kv_cache/kv_cache_manager_base.cc index cdb4bdd21..7163bb29d 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_base.cc +++ b/tpu_sync/kv_cache/kv_cache_manager_base.cc @@ -2686,9 +2686,10 @@ absl::StatusOr KVCacheManagerBase::CopyPoolBlocks( absl::StatusOr KVCacheManagerBase::D2hPoolBlocks( size_t pool_idx, absl::Span block_ids, - std::optional shard_idx, std::optional uuid) { + std::optional shard_idx, std::optional uuid, + std::optional host_layer_idx) { return CopyPoolBlocks(pool_idx, block_ids, shard_idx, - /*device_to_host=*/true, uuid); + /*device_to_host=*/true, uuid, host_layer_idx); } absl::StatusOr KVCacheManagerBase::H2dPoolBlocks( diff --git a/tpu_sync/kv_cache/kv_cache_manager_base.h b/tpu_sync/kv_cache/kv_cache_manager_base.h index 596ee3fa7..94f5faddd 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_base.h +++ b/tpu_sync/kv_cache/kv_cache_manager_base.h @@ -705,7 +705,8 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase { absl::StatusOr D2hPoolBlocks( size_t pool_idx, absl::Span block_ids, std::optional shard_idx = std::nullopt, - std::optional uuid = std::nullopt); + std::optional uuid = std::nullopt, + std::optional host_layer_idx = std::nullopt); // Inverse of D2hPoolBlocks: host mirror -> device storage. absl::StatusOr H2dPoolBlocks( 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 d87ae6783..37407e8a4 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_with_transfer.cc +++ b/tpu_sync/kv_cache/kv_cache_manager_with_transfer.cc @@ -485,16 +485,26 @@ void KVCacheManagerWithTransfer::InitializeBaseHooks() { hooks.lookup_send_staging = [this](uint64_t uuid, size_t layer_idx) -> absl::StatusOr { std::shared_ptr send_session; + std::shared_ptr reshard_session; { absl::MutexLock lock(mu_); auto it = send_sessions_.find(uuid); - if (it == send_sessions_.end()) { - return absl::NotFoundError( - absl::StrCat("Send session ", uuid, " not found")); + if (it != send_sessions_.end()) { + send_session = it->second; + } else { + auto reshard_it = active_pool_reshard_sends_.find(uuid); + if (reshard_it != active_pool_reshard_sends_.end()) { + reshard_session = reshard_it->second; + } else { + return absl::NotFoundError( + absl::StrCat("Send session ", uuid, " not found")); + } } - send_session = it->second; } - return send_session->LookupStagingSlot(layer_idx); + if (send_session != nullptr) { + return send_session->LookupStagingSlot(layer_idx); + } + return reshard_session->LookupStagingSlot(layer_idx); }; base_->SetTransferEventHooks(std::move(hooks)); } diff --git a/tpu_sync/kv_cache/reshard_send_session.cc b/tpu_sync/kv_cache/reshard_send_session.cc index 0797347a4..31e1be8c3 100644 --- a/tpu_sync/kv_cache/reshard_send_session.cc +++ b/tpu_sync/kv_cache/reshard_send_session.cc @@ -15,6 +15,7 @@ #include "tpu_sync/kv_cache/reshard_send_session.h" #include +#include #include // NOLINT(build/c++11) #include #include @@ -23,6 +24,7 @@ #include #include #include +#include // NOLINT(build/c++11) #include #include @@ -167,21 +169,107 @@ ReshardSendSession::ReshardSendSession( plan_(std::move(plan)), dst_peers_(std::move(dst_peers)), data_connection_reuse_(data_connection_reuse), - remaining_pool_peer_pushes_(remaining_pool_peer_pushes) {} + remaining_pool_peer_pushes_(remaining_pool_peer_pushes), + staging_layer_manager_(base != nullptr ? base->num_host_layers() : 0) {} ReshardSendSession::~ReshardSendSession() { absl::MutexLock lock(mu_); ReleaseLeasesLocked(); } +absl::StatusOr ReshardSendSession::AcquireStagingSlot(size_t pool_idx) { + { + absl::MutexLock lock(mu_); + if (done_ || finalizing_) { + return !status_.ok() ? status_ + : absl::CancelledError("Send session is draining"); + } + } + absl::StatusOr slot = + staging_layer_manager_.Acquire(pool_idx, deadline_, &cancelled_); + if (!slot.ok()) { + return slot.status(); + } + { + absl::MutexLock lock(mu_); + if (done_ || finalizing_) { + staging_layer_manager_.Release(pool_idx); + return !status_.ok() ? status_ + : absl::CancelledError("Send session is draining"); + } + } + return *slot; +} + +absl::StatusOr ReshardSendSession::LookupStagingSlot( + size_t pool_idx) const { + { + absl::MutexLock lock(mu_); + if (done_ || finalizing_) { + return !status_.ok() ? status_ + : absl::CancelledError("Send session is draining"); + } + } + std::optional slot = staging_layer_manager_.Lookup(pool_idx); + if (!slot.has_value()) { + return absl::FailedPreconditionError( + absl::StrCat("Pool ", pool_idx, " has no leased staging slot")); + } + return *slot; +} + +absl::Status ReshardSendSession::ValidatePushPreconditions() const { + if (base_->transport_server() == nullptr) { + return absl::FailedPreconditionError("transport server is not running"); + } + if (base_->num_host_layers() == 0) { + return absl::FailedPreconditionError( + "pool reshard push requires at least one host staging layer"); + } + for (int32_t encoded_pool_idx : plan_.transfer_pool_indices()) { + const size_t pool_idx = static_cast(encoded_pool_idx); + if (base_->pool(pool_idx) == nullptr) { + return absl::InvalidArgumentError( + absl::StrCat("transfer pool index out of range: ", pool_idx)); + } + } + return absl::OkStatus(); +} + absl::Status ReshardSendSession::ExecutePush( KVCacheManagerWithTransfer& manager, absl::Span src_block_ids) { { absl::MutexLock lock(mu_); + if (done_ || finalizing_) { + return !status_.ok() ? status_ + : absl::CancelledError("Send session is draining"); + } ++in_flight_; manager_ = &manager; } + absl::Cleanup end_op = [this]() { EndOp(); }; + // Fail fast on the plan-shape and transport preconditions that need no + // waiting, so the START_TRANSFER reply still carries them. Staging slot + // and arena leases can block on earlier pools' pushes and stay on the + // pipeline thread below. + if (absl::Status precondition = ValidatePushPreconditions(); + !precondition.ok()) { + RecordPushCompletion(manager, precondition); + return precondition; + } + std::vector src_ids(src_block_ids.begin(), src_block_ids.end()); + std::thread([self = shared_from_this(), &manager, + src_ids = std::move(src_ids)]() mutable { + self->ExecutePushPipeline(manager, std::move(src_ids)); + self->EndOp(); + }).detach(); + std::move(end_op).Cancel(); + return absl::OkStatus(); +} + +void ReshardSendSession::ExecutePushPipeline( + KVCacheManagerWithTransfer& manager, std::vector src_block_ids) { // Multi-tag plans scope each pool's staging and pushes to its group's // entries; the flat src_block_ids argument is the legacy single-tag // whole-plan block list. @@ -201,8 +289,25 @@ absl::Status ReshardSendSession::ExecutePush( local_schedule_it = plan_.shard_push_schedules().begin(); } + std::vector ordered_pools; + ordered_pools.reserve(plan_.transfer_pool_indices_size()); for (int32_t encoded_pool_idx : plan_.transfer_pool_indices()) { - const size_t pool_idx = static_cast(encoded_pool_idx); + ordered_pools.push_back(static_cast(encoded_pool_idx)); + } + std::stable_sort( + ordered_pools.begin(), ordered_pools.end(), [&](size_t a, size_t b) { + const int ga = pool_group_index(a); + const int gb = pool_group_index(b); + const int rank_a = (ga >= 0 && ga < plan_.pool_groups_size()) + ? plan_.pool_groups(ga).order_rank() + : 0; + const int rank_b = (gb >= 0 && gb < plan_.pool_groups_size()) + ? plan_.pool_groups(gb).order_rank() + : 0; + return rank_a < rank_b; + }); + + for (size_t pool_idx : ordered_pools) { std::vector pool_src_block_ids(src_block_ids.begin(), src_block_ids.end()); if (local_schedule_it != plan_.shard_push_schedules().end()) { @@ -223,6 +328,20 @@ absl::Status ReshardSendSession::ExecutePush( continue; } } + + absl::StatusOr host_slot = AcquireStagingSlot(pool_idx); + if (!host_slot.ok()) { + RecordPushCompletion(manager, host_slot.status()); + return; + } + { + absl::MutexLock lock(mu_); + if (done_ || finalizing_) { + staging_layer_manager_.Release(pool_idx); + return; + } + } + // Bounded host staging: lease one arena slot per source page of this // pool's storage for the transfer before staging its bytes (no-op on // full-mirror storages). Released in SettleLocked. @@ -233,16 +352,17 @@ absl::Status ReshardSendSession::ExecutePush( base_->pool_staging_lease_timeout()); if (!lease_status.ok()) { RecordPushCompletion(manager, lease_status); - EndOp(); - return lease_status; + staging_layer_manager_.Release(pool_idx); + return; } } - absl::StatusOr future = base_->D2hPoolBlocks( - pool_idx, pool_src_block_ids, /*shard_idx=*/std::nullopt, uuid_); + absl::StatusOr future = + base_->D2hPoolBlocks(pool_idx, pool_src_block_ids, + /*shard_idx=*/std::nullopt, uuid_, *host_slot); if (!future.ok()) { RecordPushCompletion(manager, future.status()); - EndOp(); - return future.status(); + staging_layer_manager_.Release(pool_idx); + return; } raiden::PjRtCopyFuture pool_future = *std::move(future); { @@ -251,17 +371,16 @@ absl::Status ReshardSendSession::ExecutePush( ++in_flight_; } pool_future.OnReady( - [self = shared_from_this(), &manager, pool_idx](auto status_or) { + [self = shared_from_this(), &manager, pool_idx](auto copy_status) { absl::Cleanup end_op = [self]() { self->EndOp(); }; - if (!status_or.ok()) { - self->RecordPushCompletion(manager, status_or.status()); + if (!copy_status.ok()) { + self->RecordPushCompletion(manager, copy_status.status()); + self->staging_layer_manager_.Release(pool_idx); return; } self->StartPoolPush(manager, pool_idx); }); } - EndOp(); - return absl::OkStatus(); } void ReshardSendSession::StartPoolPush(KVCacheManagerWithTransfer& manager, @@ -272,6 +391,7 @@ void ReshardSendSession::StartPoolPush(KVCacheManagerWithTransfer& manager, { absl::MutexLock lock(mu_); if (done_ || finalizing_) { + staging_layer_manager_.Release(pool_idx); return; } } @@ -303,16 +423,25 @@ void ReshardSendSession::StartPoolPush(KVCacheManagerWithTransfer& manager, } } + auto remaining_pushes = + std::make_shared>(transfers_by_peer.size()); + if (transfers_by_peer.empty()) { + staging_layer_manager_.Release(pool_idx); + return; + } + transport::BlockTransport* transport_srv = base_->transport_server(); if (transport_srv == nullptr) { RecordPushCompletion(manager, absl::FailedPreconditionError( "transport server is not running")); + staging_layer_manager_.Release(pool_idx); return; } { absl::MutexLock lock(mu_); if (done_ || finalizing_) { + staging_layer_manager_.Release(pool_idx); return; } in_flight_ += static_cast(transfers_by_peer.size()); @@ -339,12 +468,23 @@ void ReshardSendSession::StartPoolPush(KVCacheManagerWithTransfer& manager, ->AsyncPush({peer}, src_ids, dst_ids, parallelism_, transport::MajorOrder::kLayerMajor, uuid_, static_cast(pool_idx), wire_pool_idx) - .OnReady([self = shared_from_this(), - &manager](const absl::StatusOr>& result) { - absl::Cleanup end_op = [self]() { self->EndOp(); }; - self->RecordPushCompletion( - manager, result.ok() ? absl::OkStatus() : result.status()); - }); + .OnReady( + [self = shared_from_this(), &manager, pool_idx, + remaining_pushes](const absl::StatusOr>& result) { + absl::Cleanup end_op = [self]() { self->EndOp(); }; + const bool last_for_pool = remaining_pushes->fetch_sub(1) == 1; + if (!result.ok()) { + self->RecordPushCompletion(manager, result.status()); + if (last_for_pool) { + self->staging_layer_manager_.Release(pool_idx); + } + return; + } + if (last_for_pool) { + self->staging_layer_manager_.Release(pool_idx); + } + self->RecordPushCompletion(manager, absl::OkStatus()); + }); } } @@ -364,9 +504,13 @@ void ReshardSendSession::RecordPushCompletion( status_ = status; } finalizing_ = true; + cancelled_.store(true, std::memory_order_relaxed); + staging_layer_manager_.Shutdown(); should_unregister = !plan_unregistered_; } else if (--remaining_pool_peer_pushes_ == 0) { finalizing_ = true; + cancelled_.store(true, std::memory_order_relaxed); + staging_layer_manager_.Shutdown(); should_unregister = !plan_unregistered_; } if (should_unregister) plan_unregistered_ = true; @@ -398,6 +542,8 @@ void ReshardSendSession::Finish(const absl::Status& status) { } if (done_ || finalizing_) return; finalizing_ = true; + cancelled_.store(true, std::memory_order_relaxed); + staging_layer_manager_.Shutdown(); // Settling here, with no op in flight, means no push was ever issued: // ExecutePush holds an op while it schedules, StartPoolPush counts every // push before issuing it, and each push stays counted until its completion @@ -468,6 +614,7 @@ void ReshardSendSession::EndOp() { } void ReshardSendSession::ReleaseLeasesLocked() { + staging_layer_manager_.ReleaseAll(); if (leases_released_) return; leases_released_ = true; base_->ReleasePoolStagingLeases(uuid_); @@ -477,6 +624,8 @@ void ReshardSendSession::SettleLocked() { // Every push of every pool has completed (or the send failed): release the // host staging arena slots and mark done atomically under |mu_|. ReleaseLeasesLocked(); + cancelled_.store(true, std::memory_order_relaxed); + staging_layer_manager_.Shutdown(); done_ = true; } diff --git a/tpu_sync/kv_cache/reshard_send_session.h b/tpu_sync/kv_cache/reshard_send_session.h index 978b7ee85..8fc1e1d2b 100644 --- a/tpu_sync/kv_cache/reshard_send_session.h +++ b/tpu_sync/kv_cache/reshard_send_session.h @@ -92,13 +92,32 @@ class ReshardSendSession int32_t sub_plan_count() const { return plan_.sub_plan_count(); } // Leases host staging per pool, issues D2H copies for each pool in |plan_|, - // and chains async H2H peer pushes upon D2H completion. + // and chains async H2H peer pushes upon D2H completion. Non-blocking + // preconditions (transport, host staging layers, pool indices) are checked + // synchronously and returned; the staging waits, copies and pushes run on + // a pipeline thread and report through GetStatus() / Done(). absl::Status ExecutePush(KVCacheManagerWithTransfer& manager, absl::Span src_block_ids); + // Returns the host staging buffer leased to |pool_idx|, leasing one when it + // has none and waiting until deadline when every slot is leased. + absl::StatusOr AcquireStagingSlot(size_t pool_idx); + + // Returns the host staging buffer currently leased to |pool_idx|, or + // kFailedPrecondition if |pool_idx| holds no lease. + absl::StatusOr LookupStagingSlot(size_t pool_idx) const; + private: friend struct ReshardSendSessionTestPeer; + // Checks the preconditions of ExecutePush that need no waiting: a running + // transport server, at least one host staging layer, and every transfer + // pool index resolving on |base_|. + absl::Status ValidatePushPreconditions() const; + + void ExecutePushPipeline(KVCacheManagerWithTransfer& manager, + std::vector src_block_ids); + ReshardSendSession(kv_cache::KVCacheManagerBase* base, StagingBlockAllocator* staging_allocator, std::string req_id, uint64_t uuid, int parallelism, @@ -153,6 +172,8 @@ class ReshardSendSession bool plan_unregistered_ ABSL_GUARDED_BY(mu_) = false; // The last op already closed the idle data connections (kPerTransfer). bool connections_released_ ABSL_GUARDED_BY(mu_) = false; + StagingLayerManager staging_layer_manager_; + std::atomic cancelled_{false}; }; } // namespace tpu_raiden diff --git a/tpu_sync/kv_cache/reshard_send_session_test.cc b/tpu_sync/kv_cache/reshard_send_session_test.cc index 1e7310a37..4d46b4273 100644 --- a/tpu_sync/kv_cache/reshard_send_session_test.cc +++ b/tpu_sync/kv_cache/reshard_send_session_test.cc @@ -33,6 +33,7 @@ #include "absl/strings/str_cat.h" #include "absl/synchronization/mutex.h" #include "absl/types/span.h" +#include "xla/tsl/platform/statusor.h" #include "tpu_sync/core/staging_arena.h" #include "tpu_sync/kv_cache/kv_cache_manager_base.h" #include "tpu_sync/kv_cache/kv_cache_manager_with_transfer.h" @@ -50,6 +51,11 @@ struct ReshardSendSessionTestPeer { return ReshardSendSession::ValidatePlan(base, plan, src_block_ids); } + static std::optional ReleaseStagingSlot(ReshardSendSession& session, + size_t pool_idx) { + return session.staging_layer_manager_.Release(pool_idx); + } + // Counts one (pool, peer) push of `session` in flight on `manager`, as // ExecutePush and StartPoolPush do before issuing it. The executor's device // copies are device-only, so deviceless tests drive the op accounting here. @@ -71,6 +77,7 @@ struct ReshardSendSessionTestPeer { namespace { +using ::absl_testing::IsOkAndHolds; using ::absl_testing::StatusIs; using ::testing::AllOf; using ::testing::HasSubstr; @@ -452,5 +459,77 @@ TEST(ReshardSendSessionTest, CreateRejectsUnknownDataConnReuseValue) { EXPECT_FALSE(base.HasActivePlan(plan.uuid())); } +TEST(ReshardSendSessionTest, BoundedStagingSlotAcquisitionAndRelease) { + 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), + /*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); + + std::vector pools = { + DensePool("fa", 128), + DensePool("state", 128), + }; + pools[1].storage_index = 1; + ABSL_ASSERT_OK(base.RegisterPools(pools)); + + std::unique_ptr allocator = + CreateStagingBlockAllocator(&base, /*num_slots=*/0, /*max_blocks=*/0); + StartTransferRequest plan = ValidPlan( + /*uuid=*/1301, /*dtype_tags=*/{"bf16", "bf16"}, + /*transferred_pools=*/{0, 1}); + TF_ASSERT_OK_AND_ASSIGN( + std::shared_ptr session, + ReshardSendSession::Create( + &base, allocator.get(), std::vector{0}, /*parallelism=*/1, + std::chrono::steady_clock::now() + std::chrono::seconds(2), plan)); + + // Pool 0 leases slot 0 (the only allocated host layer). + EXPECT_THAT(session->AcquireStagingSlot(0), IsOkAndHolds(0u)); + EXPECT_THAT(session->LookupStagingSlot(0), IsOkAndHolds(0u)); + + // Pool 1 blocks waiting for slot 0 until pool 0 releases its lease. + absl::StatusOr slot1; + std::thread release_waiter([&]() { slot1 = session->AcquireStagingSlot(1); }); + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + EXPECT_EQ(ReshardSendSessionTestPeer::ReleaseStagingSlot(*session, 0), + std::make_optional(0u)); + release_waiter.join(); + EXPECT_THAT(slot1, IsOkAndHolds(0u)); + EXPECT_THAT(session->LookupStagingSlot(1), IsOkAndHolds(0u)); + + // Cancelling the session while pool 0 is blocked in AcquireStagingSlot wakes + // the waiter with CancelledError and releases all staging slots on settle. + absl::StatusOr slot0_cancelled; + std::thread cancel_waiter( + [&]() { slot0_cancelled = session->AcquireStagingSlot(0); }); + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + session->Finish(absl::CancelledError("cancelled")); + cancel_waiter.join(); + EXPECT_THAT(slot0_cancelled, StatusIs(absl::StatusCode::kCancelled)); + EXPECT_TRUE(session->Done()); + EXPECT_THAT(session->LookupStagingSlot(1), + StatusIs(absl::StatusCode::kCancelled)); + + // Timeout when no slot is released before the session deadline. + StartTransferRequest timeout_plan = ValidPlan( + /*uuid=*/1302, /*dtype_tags=*/{"bf16", "bf16"}, + /*transferred_pools=*/{0, 1}); + TF_ASSERT_OK_AND_ASSIGN( + std::shared_ptr timeout_session, + ReshardSendSession::Create( + &base, allocator.get(), std::vector{0}, /*parallelism=*/1, + std::chrono::steady_clock::now() + std::chrono::milliseconds(30), + timeout_plan)); + EXPECT_THAT(timeout_session->AcquireStagingSlot(0), IsOkAndHolds(0u)); + EXPECT_THAT(timeout_session->AcquireStagingSlot(1), + StatusIs(absl::StatusCode::kDeadlineExceeded)); +} + } // namespace } // namespace tpu_raiden