From bd56349335136eb4b49f6b5d326d4b09562b0d95 Mon Sep 17 00:00:00 2001 From: Googler Date: Thu, 29 Oct 2026 17:11:02 -0700 Subject: [PATCH] =?UTF-8?q?Index=20active=20transfer=20schedules=20by=20bl?= =?UTF-8?q?ock=20ID=20in=20`RegisterActivePlan`=20so=20`KVCacheManagerBase?= =?UTF-8?q?::GetBlockChunks`=20scales=20linearly=20`O(N)`=20instead=20of?= =?UTF-8?q?=20quadratically=20`O(N=C2=B2)`.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PiperOrigin-RevId: 1010081939 --- tpu_sync/kv_cache/kv_cache_manager_base.cc | 52 ++++++++++--- tpu_sync/kv_cache/kv_cache_manager_base.h | 6 ++ tpu_sync/kv_cache/pool_staging_test.cc | 85 ++++++++++++++++++++++ 3 files changed, 131 insertions(+), 12 deletions(-) diff --git a/tpu_sync/kv_cache/kv_cache_manager_base.cc b/tpu_sync/kv_cache/kv_cache_manager_base.cc index 7163bb29d..c6fef4c8f 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_base.cc +++ b/tpu_sync/kv_cache/kv_cache_manager_base.cc @@ -3240,11 +3240,20 @@ absl::Status KVCacheManagerBase::RegisterActivePlan( } } } + auto plan_ptr = std::make_shared( + RegisteredPlan{request, is_sender, std::move(host_block_of), + std::move(staged_device_blocks), generation}); + for (const auto& [src_shard, schedule] : + plan_ptr->request.shard_push_schedules()) { + auto& block_map = plan_ptr->schedule_by_block[src_shard]; + for (const auto& entry : schedule.entries()) { + block_map[is_sender ? entry.src_block_id() : entry.dst_block_id()] + .push_back(&entry); + } + } absl::MutexLock l(plans_mu_); - if (auto [it, inserted] = active_plans_.try_emplace( - uuid, std::make_shared(RegisteredPlan{ - request, is_sender, std::move(host_block_of), - std::move(staged_device_blocks), generation})); + if (auto [it, inserted] = + active_plans_.try_emplace(uuid, std::move(plan_ptr)); !inserted) { return absl::AlreadyExistsError( absl::StrCat("Plan with UUID ", uuid, " is already registered!")); @@ -3476,10 +3485,20 @@ KVCacheManagerBase::GetBlockChunks(size_t layer_idx, size_t shard_idx, } const auto& request = plan_snapshot->request; - const auto& schedules = request.shard_push_schedules(); - auto schedule_it = schedules.find(static_cast(shard_idx)); - + const auto& schedules = plan_snapshot->schedule_by_block; bool is_sender = plan_snapshot->is_sender; + const int32_t fixed_src_shard = is_sender + ? static_cast(shard_idx) + : static_cast(sender_node_id); + auto schedule_it = + fixed_src_shard != -1 ? schedules.find(fixed_src_shard) : schedules.end(); + const auto block_entries = [](const RegisteredPlan::BlockEntryMap& block_map, + int64_t block_id) + -> absl::Span { + auto it = block_map.find(block_id); + if (it == block_map.end()) return {}; + return it->second; + }; // Pool reshard plans scope every entry to one pool group; an entry only // resolves chunks for pools of its own group. Only LEGACY (non-pool) @@ -3506,6 +3525,7 @@ KVCacheManagerBase::GetBlockChunks(size_t layer_idx, size_t shard_idx, } std::vector chunks; + std::vector block_resolved_chunks; size_t accumulated_bytes = 0; for (int64_t block_id : block_ids) { @@ -3527,12 +3547,14 @@ KVCacheManagerBase::GetBlockChunks(size_t layer_idx, size_t shard_idx, // Bounded storage: this uuid's lease does not cover the block. return {}; } - std::vector block_resolved_chunks; + block_resolved_chunks.clear(); if (is_sender) { if (schedule_it != schedules.end()) { const auto& schedule = schedule_it->second; - for (const auto& entry : schedule.entries()) { + for (const ::tpu_sync::rpc::ShardPushEntryProto* entry_ptr : + block_entries(schedule, block_id)) { + const ::tpu_sync::rpc::ShardPushEntryProto& entry = *entry_ptr; if (!entry_targets_pool(entry)) { continue; } @@ -3573,7 +3595,9 @@ KVCacheManagerBase::GetBlockChunks(size_t layer_idx, size_t shard_idx, // (merging) where multiple source blocks target the same dst_block_id. // If src_block_id is -1, we fall back to matching only on dst_block_id. for (const auto& [src_shard, src_schedule] : schedules) { - for (const auto& entry : src_schedule.entries()) { + for (const ::tpu_sync::rpc::ShardPushEntryProto* entry_ptr : + block_entries(src_schedule, block_id)) { + const ::tpu_sync::rpc::ShardPushEntryProto& entry = *entry_ptr; if (entry_targets_pool(entry) && static_cast(entry.dst_shard_idx()) == shard_idx && static_cast(entry.dst_block_id()) == block_id && @@ -3588,10 +3612,14 @@ KVCacheManagerBase::GetBlockChunks(size_t layer_idx, size_t shard_idx, } if (found_src_shard != -1) { - auto schedule_found_it = schedules.find(found_src_shard); + auto schedule_found_it = sender_node_id != -1 + ? schedule_it + : schedules.find(found_src_shard); if (schedule_found_it != schedules.end()) { const auto& schedule = schedule_found_it->second; - for (const auto& entry : schedule.entries()) { + for (const ::tpu_sync::rpc::ShardPushEntryProto* entry_ptr : + block_entries(schedule, block_id)) { + const ::tpu_sync::rpc::ShardPushEntryProto& entry = *entry_ptr; if (!entry_targets_pool(entry)) { continue; } diff --git a/tpu_sync/kv_cache/kv_cache_manager_base.h b/tpu_sync/kv_cache/kv_cache_manager_base.h index 94f5faddd..177d2973b 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_base.h +++ b/tpu_sync/kv_cache/kv_cache_manager_base.h @@ -1048,6 +1048,9 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase { mutable absl::Mutex plans_mu_; struct RegisteredPlan { + using EntryList = std::vector; + using BlockEntryMap = absl::flat_hash_map; + ::tpu_sync::rpc::StartTransferRequest request; bool is_sender = false; // Host block holding each device block the plan names on this side @@ -1063,6 +1066,9 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase { // registration cannot touch a newer one under the same uuid. 0 when // the caller keeps no generations; cleanup then matches by uuid alone. uint64_t generation = 0; + // Per-source-shard schedule entries indexed by src_block_id (sender) or + // dst_block_id (receiver). Pointers alias `request.shard_push_schedules()`. + absl::flat_hash_map schedule_by_block; }; // Plans are stored behind shared_ptr so the per-push readers // (GetBlockChunks / GetPoolPushProgressSpec) snapshot with a refcount bump diff --git a/tpu_sync/kv_cache/pool_staging_test.cc b/tpu_sync/kv_cache/pool_staging_test.cc index 5333b391a..ee60eef65 100644 --- a/tpu_sync/kv_cache/pool_staging_test.cc +++ b/tpu_sync/kv_cache/pool_staging_test.cc @@ -604,6 +604,91 @@ TEST(PoolStagingTest, BoundedHostLayersPoolStaging) { manager.ReleasePoolStagingLeases(42); } +TEST(PoolStagingTest, GetBlockChunksLargePlanLinearScaling) { + constexpr size_t kNumShards = 8; + constexpr size_t kNumFaPools = 6; + constexpr size_t kTotalPools = 93; + constexpr int64_t kStride = 64; + constexpr int64_t kMaxBlocks = 4096; + constexpr int64_t kNumBlocks = 3072; + constexpr uint64_t kUuid = 101; + constexpr char kPeer[] = "127.0.0.1:1"; + + StagingTestManager receiver(1, kNumShards, kStride, /*host_blocks=*/1); + StagingTestManager sender(1, kNumShards, kStride, /*host_blocks=*/1); + const size_t storage_bytes = kTotalPools * kStride * kMaxBlocks; + receiver.SetDeviceBacked(0, storage_bytes); + sender.SetDeviceBacked(0, storage_bytes); + + std::vector specs; + specs.reserve(kTotalPools); + for (size_t pool_idx = 0; pool_idx < kTotalPools; ++pool_idx) { + specs.push_back(DensePool("pool_" + std::to_string(pool_idx), 0, + pool_idx * kStride * kMaxBlocks, kStride, + kMaxBlocks, /*staging_blocks_per_request=*/0)); + } + ABSL_ASSERT_OK(receiver.RegisterPools(specs, /*staging_leases=*/0)); + ABSL_ASSERT_OK(sender.RegisterPools(specs, /*staging_leases=*/0)); + + tpu_sync::rpc::StartTransferRequest request; + request.set_uuid(kUuid); + request.set_req_id("req"); + tpu_sync::rpc::PoolGroupProto* fa_group = request.add_pool_groups(); + tpu_sync::rpc::PoolGroupProto* swa_group = request.add_pool_groups(); + fa_group->set_expected_pushes(1); + swa_group->set_expected_pushes(1); + for (size_t pool_idx = 0; pool_idx < kTotalPools; ++pool_idx) { + request.add_transfer_pool_indices(pool_idx); + (pool_idx < kNumFaPools ? fa_group : swa_group)->add_pool_indices(pool_idx); + } + + std::vector dst_blocks(kNumBlocks); + for (size_t shard_idx = 0; shard_idx < kNumShards; ++shard_idx) { + tpu_sync::rpc::ShardPushScheduleProto& schedule = + (*request.mutable_shard_push_schedules())[shard_idx]; + for (int64_t block_id = 0; block_id < kNumBlocks; ++block_id) { + dst_blocks[block_id] = block_id + 100; + tpu_sync::rpc::ShardPushEntryProto* entry = schedule.add_entries(); + entry->set_dst_peer(kPeer); + entry->set_dst_shard_idx(shard_idx); + entry->set_src_block_id(block_id); + entry->set_dst_block_id(dst_blocks[block_id]); + entry->set_size_bytes(kStride); + entry->set_pool_group(0); + } + } + ABSL_ASSERT_OK( + receiver.RegisterActivePlan(kUuid, request, /*is_sender=*/false)); + for (size_t pool_idx = 0; pool_idx < kNumFaPools; ++pool_idx) { + TF_ASSERT_OK_AND_ASSIGN(tpu_sync::rpc::PoolHostAddrsProto pool_addrs, + receiver.PoolHostBaseAddrs(kUuid, pool_idx)); + (*(*request.mutable_receiver_addrs())[kPeer].mutable_pools())[pool_idx] = + pool_addrs; + } + ABSL_ASSERT_OK(sender.RegisterActivePlan(kUuid, request, /*is_sender=*/true)); + + for (size_t shard_idx = 0; shard_idx < kNumShards; ++shard_idx) { + for (size_t pool_idx = 0; pool_idx < kNumFaPools; ++pool_idx) { + std::vector recv_chunks = receiver.GetBlockChunks( + pool_idx, shard_idx, dst_blocks, kNumBlocks * kStride, kUuid, + /*sender_node_id=*/shard_idx); + ASSERT_EQ(recv_chunks.size(), static_cast(kNumBlocks)); + for (int64_t block_id = 0; block_id < kNumBlocks; ++block_id) { + std::vector send_chunks = sender.GetBlockChunks( + pool_idx, shard_idx, {block_id}, kStride, kUuid, + /*sender_node_id=*/-1, kPeer, block_id, dst_blocks[block_id]); + ASSERT_EQ(send_chunks.size(), 1u); + EXPECT_EQ(send_chunks[0].raddr, recv_chunks[block_id].ptr); + } + } + } + EXPECT_TRUE(receiver + .GetBlockChunks(kNumFaPools, /*shard_idx=*/0, dst_blocks, + kNumBlocks * kStride, kUuid, + /*sender_node_id=*/0) + .empty()); +} + } // namespace } // namespace kv_cache } // namespace tpu_raiden