diff --git a/tpu_sync/kv_cache/BUILD b/tpu_sync/kv_cache/BUILD index a74c2afc1..6d4abd415 100644 --- a/tpu_sync/kv_cache/BUILD +++ b/tpu_sync/kv_cache/BUILD @@ -519,9 +519,13 @@ cc_test( "//tpu_sync/core:raw_transfer_core", "//tpu_sync/core:staging_arena", "//tpu_sync/rpc:raiden_service_cc_proto", + "//tpu_sync/transport:block_transport_delegate", + "//tpu_sync/transport:transfer_session", + "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/synchronization", "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest_main", "@xla//xla:future", diff --git a/tpu_sync/kv_cache/kv_cache_manager_with_transfer_pool_reshard_test.cc b/tpu_sync/kv_cache/kv_cache_manager_with_transfer_pool_reshard_test.cc index 4596746a1..3755276a1 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_with_transfer_pool_reshard_test.cc +++ b/tpu_sync/kv_cache/kv_cache_manager_with_transfer_pool_reshard_test.cc @@ -124,7 +124,9 @@ class TestManager : public KVCacheManagerWithTransfer { absl::StatusOr> session = GetSession(uuid); ABSL_ASSERT_OK(session); static_cast(*base_).ArmPoolUpload(pool_idx, status); - ABSL_ASSERT_OK((*session)->ResolveStagingSlot(pool_idx)); + // kNotFound: a full host mirror stages the pool in place. + absl::StatusOr slot = (*session)->ResolveStagingSlot(pool_idx); + ASSERT_TRUE(slot.ok() || absl::IsNotFound(slot.status())) << slot.status(); ABSL_ASSERT_OK((*session)->OnDataReceived( pool_idx, /*minor_dimension_indices=*/{}, /*shard_ids=*/{})); } diff --git a/tpu_sync/kv_cache/reshard_receive_session.cc b/tpu_sync/kv_cache/reshard_receive_session.cc index af1d32efe..dca5691a5 100644 --- a/tpu_sync/kv_cache/reshard_receive_session.cc +++ b/tpu_sync/kv_cache/reshard_receive_session.cc @@ -138,7 +138,8 @@ class ReshardReceiveSession final return tsl::Future<>(absl::OkStatus()); } - // Leases a buffer to |pool_idx| on first use. + // Leases a buffer to |pool_idx| on first use when host layers are bounded; + // kNotFound on a full host mirror (the pool is addressed in place). absl::StatusOr ResolveStagingSlot(size_t pool_idx) override { return AcquireStagingSlot(pool_idx); } @@ -156,9 +157,17 @@ class ReshardReceiveSession final absl::Status OnPoolReceived(size_t pool_idx); // Returns the host staging buffer leased to |pool_idx|, leasing one when it - // has none and waiting until deadline when every slot is leased. + // has none and waiting until deadline when every slot is leased. Returns + // kNotFound on a full host mirror, where every pool stages in its own + // storage's host buffer and there is no slot to lease. absl::StatusOr AcquireStagingSlot(size_t pool_idx); + // True when fewer host layer buffers are allocated than storages, so pools + // stage through leased slots instead of their own storage's host buffer. + bool HostLayersBounded() const { + return base_ != nullptr && base_->num_host_layers() < base_->num_layers(); + } + // 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; @@ -561,7 +570,8 @@ ReshardReceiveSession::ReshardReceiveSession( start_time_(std::chrono::steady_clock::now()), transfer_plan_( std::make_shared(std::move(transfer_plan))), - staging_layer_manager_(base != nullptr ? base->num_host_layers() : 0) { + staging_layer_manager_(HostLayersBounded() ? base->num_host_layers() + : 0) { absl::MutexLock lock(mu_); chip_block_ids_.assign(chip_blocks.begin(), chip_blocks.end()); for (int32_t pool_idx : request().transfer_pool_indices()) { @@ -647,6 +657,13 @@ void ReshardReceiveSession::FinishLocked(const absl::Status& status) { absl::StatusOr ReshardReceiveSession::AcquireStagingSlot( size_t pool_idx) { + // A full host mirror stages every pool in place: GetBlockChunks, the + // transport's chunk validation and H2dPoolBlocks all address the pool's own + // storage. A leased slot would land the bytes in another storage's host + // buffer, outside the span ValidateChunks accepts for the pool. + if (!HostLayersBounded()) { + return absl::NotFoundError("no bounded host staging"); + } std::chrono::steady_clock::time_point deadline; { absl::MutexLock lock(mu_); @@ -823,13 +840,18 @@ void ReshardReceiveSession::ExecuteEligiblePoolH2ds() { EndOp(); continue; } - absl::StatusOr slot = LookupStagingSlot(pool_idx); - if (!slot.ok()) { - FinishPoolH2d(pool_idx, slot.status()); - EndOp(); - continue; + // Host source of the upload: the pool's own storage on a full mirror, or + // the slot leased when its bytes were addressed in GetBlockChunks. + std::optional host_layer_idx; + if (HostLayersBounded()) { + absl::StatusOr slot = LookupStagingSlot(pool_idx); + if (!slot.ok()) { + FinishPoolH2d(pool_idx, slot.status()); + EndOp(); + continue; + } + host_layer_idx = *slot; } - const size_t host_layer_idx = *slot; absl::StatusOr copy_future = base_->H2dPoolBlocks( pool_idx, dst_chip_block_ids, /*shard_idx=*/std::nullopt, uuid(), host_layer_idx); diff --git a/tpu_sync/kv_cache/reshard_receive_session_test.cc b/tpu_sync/kv_cache/reshard_receive_session_test.cc index 376c306a3..5fd5c2871 100644 --- a/tpu_sync/kv_cache/reshard_receive_session_test.cc +++ b/tpu_sync/kv_cache/reshard_receive_session_test.cc @@ -26,9 +26,11 @@ #include #include +#include "absl/base/thread_annotations.h" #include "absl/status/status.h" #include "absl/status/status_matchers.h" #include "absl/status/statusor.h" +#include "absl/synchronization/mutex.h" #include "absl/types/span.h" #include "xla/future.h" #include "xla/tsl/platform/statusor.h" @@ -38,6 +40,8 @@ #include "tpu_sync/kv_cache/kv_cache_manager_with_transfer.h" #include "tpu_sync/kv_cache/pool_layout.h" #include "tpu_sync/rpc/raiden_service.pb.h" +#include "tpu_sync/transport/block_transport_delegate.h" +#include "tpu_sync/transport/transfer_session.h" namespace tpu_raiden { @@ -165,10 +169,26 @@ class FakeH2dBase : public kv_cache::KVCacheManagerBase { size_t pool_idx, absl::Span block_ids, std::optional shard_idx, std::optional uuid, std::optional host_layer_idx) override { + { + absl::MutexLock lock(fake_mu_); + uploads_.emplace_back(pool_idx, host_layer_idx); + } auto [promise, future] = xla::MakePromise(); promise.Set(); return raiden::PjRtCopyFuture(std::move(future), raiden::BufferHolders{}); } + + // (pool index, host staging buffer) of every upload issued, in order; + // a nullopt buffer means the pool's own storage. + std::vector>> uploads() const { + absl::MutexLock lock(fake_mu_); + return uploads_; + } + + private: + mutable absl::Mutex fake_mu_; + std::vector>> uploads_ + ABSL_GUARDED_BY(fake_mu_); }; TEST(ReshardReceiveSessionTest, AcceptsCanonicalPlanOnExplicitPools) { @@ -517,5 +537,87 @@ TEST(ReshardReceiveSessionTest, BoundedStagingSlotAcquisitionAndRelease) { StatusIs(absl::StatusCode::kCancelled)); } +// Mirrors the bounds check the transport's ValidateChunks applies to every +// chunk of an incoming push: the chunk must lie inside the pool's flat host +// span when the manager exposes one. +bool ChunkPassesTransportValidation(kv_cache::KVCacheManagerBase& base, + size_t pool_idx, + const transport::BlockChunk& chunk) { + uint8_t* array_base = base.GetBlockArrayHostPointer(pool_idx, 0); + if (array_base == nullptr) return true; + const size_t array_size = base.GetBlockArrayHostSize(pool_idx, 0); + if (chunk.ptr < array_base) return false; + const size_t offset = static_cast(chunk.ptr - array_base); + return offset <= array_size && chunk.size <= array_size - offset; +} + +// On a full host mirror (no num_staging_layers) every pool stages in its own +// storage's host buffer. Regression test: the receive used to lease a staging +// slot even here, so the pool whose bytes arrived first (pool 1, storage 1) +// was addressed in host buffer 0 while the transport validated its chunks +// against storage 1, failing the push with "Chunk out of bounds". +TEST(ReshardReceiveSessionTest, FullHostMirrorAddressesPoolsInPlace) { + FakeH2dBase base( + /*num_layers=*/2, /*num_shards=*/1, + /*slice_byte_size=*/128, + /*local_port=*/std::nullopt, + /*host_blocks_to_allocate=*/std::make_optional(4)); + base.AttachPlaceholderDeviceHoldForTest(); + ASSERT_EQ(base.num_host_layers(), base.num_layers()); + + 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=*/1303, /*dtype_tags=*/{"bf16", "bf16"}, + /*transferred_pools=*/{0, 1}); + TF_ASSERT_OK_AND_ASSIGN( + std::shared_ptr session, + CreateReshardReceiveSession( + &base, allocator.get(), plan, std::vector{0}, + std::chrono::steady_clock::now() + std::chrono::seconds(10))); + kv_cache::KVCacheManagerBase::TransferEventHooks hooks; + hooks.get_session = + [weak = std::weak_ptr(session)]( + uint64_t uuid) -> absl::StatusOr> { + if (std::shared_ptr live = weak.lock()) return live; + return absl::NotFoundError("session gone"); + }; + base.SetTransferEventHooks(std::move(hooks)); + + // Pushes land in any order; pool 1 arrives first. There is no slot to lease. + EXPECT_THAT(session->ResolveStagingSlot(1), + StatusIs(absl::StatusCode::kNotFound)); + for (size_t pool_idx : {size_t{1}, size_t{0}}) { + std::vector chunks = base.GetBlockChunks( + pool_idx, /*shard_idx=*/0, {0}, /*total_bytes=*/16, plan.uuid(), + /*sender_node_id=*/0); + ASSERT_EQ(chunks.size(), 1u) << "pool " << pool_idx; + EXPECT_EQ(chunks[0].ptr, base.GetBlockArrayHostPointer(pool_idx, 0)) + << "pool " << pool_idx; + EXPECT_TRUE(ChunkPassesTransportValidation(base, pool_idx, chunks[0])) + << "pool " << pool_idx; + } + + // Each upload reads the pool's own storage. + ABSL_EXPECT_OK(session->OnDataReceived(/*major_dimension_idx=*/1, + /*minor_dimension_indices=*/{}, + /*shard_ids=*/{})); + ABSL_EXPECT_OK(session->OnDataReceived(/*major_dimension_idx=*/0, + /*minor_dimension_indices=*/{}, + /*shard_ids=*/{})); + using Upload = std::pair>; + EXPECT_THAT(base.uploads(), + ::testing::UnorderedElementsAre(Upload{0, std::nullopt}, + Upload{1, std::nullopt})); + ABSL_EXPECT_OK(session->GetStatus()); +} + } // namespace } // namespace tpu_raiden diff --git a/tpu_sync/kv_cache/reshard_send_session.cc b/tpu_sync/kv_cache/reshard_send_session.cc index bb7f08c11..6e70b16f6 100644 --- a/tpu_sync/kv_cache/reshard_send_session.cc +++ b/tpu_sync/kv_cache/reshard_send_session.cc @@ -154,15 +154,24 @@ class ReshardSendSession final private: // Returns the host staging buffer leased to |pool_idx|, leasing one when it - // has none and waiting until deadline when every slot is leased. + // has none and waiting until deadline when every slot is leased. Returns + // kNotFound on a full host mirror, where every pool stages in its own + // storage's host buffer and there is no slot to lease. absl::StatusOr AcquireStagingSlot(size_t pool_idx); - // Returns the host staging buffer currently leased to |pool_idx|. Fails + // Returns the host staging buffer currently leased to |pool_idx|. Returns + // kNotFound on a full host mirror (the pool is addressed in place). Fails // with the session's status (or kCancelled) once the session is draining // and holds no lease for |pool_idx|, and kFailedPrecondition when a live // session holds none. absl::StatusOr LookupStagingSlot(size_t pool_idx) const; + // True when fewer host layer buffers are allocated than storages, so pools + // stage through leased slots instead of their own storage's host buffer. + bool HostLayersBounded() const { + return base_ != nullptr && base_->num_host_layers() < base_->num_layers(); + } + // Checks the preconditions of Start that need no waiting: a running // transport server, at least one host staging layer, and every transfer // pool index resolving on |base_|. @@ -329,7 +338,8 @@ ReshardSendSession::ReshardSendSession( transfer_plan_( std::make_shared(std::move(transfer_plan))), remaining_pool_peer_pushes_(remaining_pool_peer_pushes), - staging_layer_manager_(base != nullptr ? base->num_host_layers() : 0) {} + staging_layer_manager_(HostLayersBounded() ? base->num_host_layers() + : 0) {} ReshardSendSession::~ReshardSendSession() { absl::MutexLock lock(mu_); @@ -337,6 +347,11 @@ ReshardSendSession::~ReshardSendSession() { } absl::StatusOr ReshardSendSession::AcquireStagingSlot(size_t pool_idx) { + // A full host mirror stages every pool in place: the D2H copy, GetBlockChunks + // and the transport's chunk validation all address the pool's own storage. + if (!HostLayersBounded()) { + return absl::NotFoundError("no bounded host staging"); + } { absl::MutexLock lock(mu_); if (done_ || finalizing_) { @@ -362,6 +377,9 @@ absl::StatusOr ReshardSendSession::AcquireStagingSlot(size_t pool_idx) { absl::StatusOr ReshardSendSession::LookupStagingSlot( size_t pool_idx) const { + if (!HostLayersBounded()) { + return absl::NotFoundError("no bounded host staging"); + } // A lease outlives the session's finalizing_ flip: a push whose completion // started finalizing may still be addressing another pool's host bytes. // SettleLocked releases every lease, so a done session never resolves one. @@ -499,9 +517,13 @@ void ReshardSendSession::ExecutePushPipeline( } } - absl::StatusOr host_slot = AcquireStagingSlot(pool_idx); - if (!host_slot.ok()) { - RecordPushCompletion(host_slot.status()); + // Host staging buffer of the D2H copy: a leased slot when host layers are + // bounded, else (kNotFound) the pool's own storage. + std::optional host_slot; + if (absl::StatusOr slot = AcquireStagingSlot(pool_idx); slot.ok()) { + host_slot = *slot; + } else if (!absl::IsNotFound(slot.status())) { + RecordPushCompletion(slot.status()); return; } { @@ -528,7 +550,7 @@ void ReshardSendSession::ExecutePushPipeline( } absl::StatusOr future = base_->D2hPoolBlocks(pool_idx, pool_src_block_ids, - /*shard_idx=*/std::nullopt, uuid(), *host_slot); + /*shard_idx=*/std::nullopt, uuid(), host_slot); if (!future.ok()) { RecordPushCompletion(future.status()); staging_layer_manager_.Release(pool_idx); diff --git a/tpu_sync/kv_cache/reshard_send_session_test.cc b/tpu_sync/kv_cache/reshard_send_session_test.cc index 9e7672f12..5d074a654 100644 --- a/tpu_sync/kv_cache/reshard_send_session_test.cc +++ b/tpu_sync/kv_cache/reshard_send_session_test.cc @@ -372,6 +372,14 @@ class FakeD2hBase : public kv_cache::KVCacheManagerBase { return true; } + // Returns the host staging buffer the last copy of `pool_idx` targeted + // (nullopt: the pool's own storage). + std::optional LastHostLayerIdx(size_t pool_idx) const { + absl::MutexLock lock(fake_mu_); + auto it = last_host_layer_idx_.find(pool_idx); + return it == last_host_layer_idx_.end() ? std::nullopt : it->second; + } + absl::StatusOr D2hPoolBlocks( size_t pool_idx, absl::Span block_ids, std::optional shard_idx, std::optional uuid, @@ -379,6 +387,7 @@ class FakeD2hBase : public kv_cache::KVCacheManagerBase { auto [promise, future] = xla::MakePromise(); absl::MutexLock lock(fake_mu_); ++copies_issued_[pool_idx]; + last_host_layer_idx_[pool_idx] = host_layer_idx; if (hold_next_[pool_idx]) { hold_next_[pool_idx] = false; pending_[pool_idx] = std::move(promise); @@ -394,6 +403,8 @@ class FakeD2hBase : public kv_cache::KVCacheManagerBase { absl::flat_hash_map> pending_ ABSL_GUARDED_BY(fake_mu_); absl::flat_hash_map copies_issued_ ABSL_GUARDED_BY(fake_mu_); + absl::flat_hash_map> last_host_layer_idx_ + ABSL_GUARDED_BY(fake_mu_); }; // Routes `sender`'s transport session lookups to `session`, as the manager @@ -763,11 +774,16 @@ TEST(ReshardSendSessionTest, PipelineSlotWaitTimeoutSettlesSendOutOfBand) { } // A pool pushing to several peers keeps its staging slot until the last -// peer push completes, then the send settles clean. +// peer push completes, then the send settles clean. Slots are leased only when +// host layers are bounded (here one host layer buffer for two storages). TEST(ReshardSendSessionTest, StagingSlotReleasedAfterLastPeerPush) { - FakeD2hBase sender(/*num_layers=*/1, /*num_shards=*/1, + FakeD2hBase sender(/*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); + ASSERT_EQ(sender.num_host_layers(), 1u); + ABSL_ASSERT_OK(sender.RegisterPools({DensePool("fa", 128)})); kv_cache::KVCacheManagerBase receiver_a = MakeTestBase(); kv_cache::KVCacheManagerBase receiver_b = MakeTestBase(); auto recv_session = std::make_shared(); @@ -780,7 +796,7 @@ TEST(ReshardSendSessionTest, StagingSlotReleasedAfterLastPeerPush) { sender.AttachPlaceholderDeviceHoldForTest(); std::unique_ptr allocator = CreateStagingBlockAllocator(&sender, /*num_slots=*/0, /*max_blocks=*/0); - StartTransferRequest plan = ValidPlan(/*uuid=*/1403, /*dtype_tags=*/{""}); + StartTransferRequest plan = ValidPlan(/*uuid=*/1403); auto* schedule = &(*plan.mutable_shard_push_schedules())[0]; schedule->mutable_entries(0)->set_dst_peer(peer_a); *schedule->add_entries() = schedule->entries(0); @@ -808,5 +824,56 @@ TEST(ReshardSendSessionTest, StagingSlotReleasedAfterLastPeerPush) { StatusIs(absl::StatusCode::kCancelled)); } +// On a full host mirror (no num_staging_layers) every pool stages in its own +// storage's host buffer. Regression test: the send used to lease a staging +// slot even here, so pool 0 (on storage 1) was copied to host buffer 0 and its +// push chunks failed the transport's ValidateChunks, which checks them against +// pool 0's own storage ("Chunk out of bounds"). +TEST(ReshardSendSessionTest, FullHostMirrorStagesPoolsInPlace) { + FakeD2hBase sender(/*num_layers=*/2, /*num_shards=*/1, + /*slice_byte_size=*/128, /*local_port=*/std::nullopt, + /*host_blocks_to_allocate=*/std::make_optional(4)); + ASSERT_EQ(sender.num_host_layers(), sender.num_layers()); + std::vector pools = { + DensePool("fa", 128), + DensePool("state", 128), + }; + // Pool order differs from storage order, so slot order != storage order. + pools[0].storage_index = 1; + pools[1].storage_index = 0; + ABSL_ASSERT_OK(sender.RegisterPools(pools)); + kv_cache::KVCacheManagerBase receiver( + /*num_layers=*/2, /*num_shards=*/1, /*slice_byte_size=*/128, + /*local_port=*/std::nullopt, + /*host_blocks_to_allocate=*/std::make_optional(4)); + AcceptPushes(receiver, std::make_shared()); + const std::string peer = + absl::StrCat("127.0.0.1:", receiver.transport_server()->local_port()); + sender.AttachPlaceholderDeviceHoldForTest(); + std::unique_ptr allocator = + CreateStagingBlockAllocator(&sender, /*num_slots=*/0, /*max_blocks=*/0); + StartTransferRequest plan = ValidPlan( + /*uuid=*/1404, /*dtype_tags=*/{"bf16", "bf16"}, + /*transferred_pools=*/{0, 1}); + (*plan.mutable_shard_push_schedules())[0].mutable_entries(0)->set_dst_peer( + peer); + TF_ASSERT_OK_AND_ASSIGN( + std::shared_ptr session, + CreateReshardSendSession( + &sender, allocator.get(), std::vector{0}, + /*parallelism=*/1, + std::chrono::steady_clock::now() + std::chrono::seconds(10), plan)); + RouteSessionLookups(sender, session); + EXPECT_THAT(session->ResolveStagingSlot(0), + StatusIs(absl::StatusCode::kNotFound)); + + ABSL_ASSERT_OK(session->Start()); + ABSL_EXPECT_OK(session->AwaitForDone()); + EXPECT_EQ(sender.CopiesIssued(0), 1); + EXPECT_EQ(sender.CopiesIssued(1), 1); + EXPECT_EQ(sender.LastHostLayerIdx(0), std::nullopt); + EXPECT_EQ(sender.LastHostLayerIdx(1), std::nullopt); +} + } // namespace } // namespace tpu_raiden