Index active transfer schedules by block ID in RegisterActivePlan so KVCacheManagerBase::GetBlockChunks scales linearly O(N) instead of quadratically O(N²).

tpu_sync/kv_cache/kv_cache_manager_base.cc

@@ -3240,11 +3240,20 @@ absl::Status KVCacheManagerBase::RegisterActivePlan(
       }
     }
   }
+  auto plan_ptr = std::make_shared<RegisteredPlan>(
+      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<const RegisteredPlan>(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<int32_t>(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<int32_t>(shard_idx)
+                                      : static_cast<int32_t>(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<const ::tpu_sync::rpc::ShardPushEntryProto* const> {
+    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<tpu_raiden::transport::BlockChunk> chunks;
+  std::vector<tpu_raiden::transport::BlockChunk> 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<tpu_raiden::transport::BlockChunk> 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<size_t>(entry.dst_shard_idx()) == shard_idx &&
                 static_cast<size_t>(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;
             }

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<const ::tpu_sync::rpc::ShardPushEntryProto*>;
+    using BlockEntryMap = absl::flat_hash_map<int64_t, EntryList>;
+
     ::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<int32_t, BlockEntryMap> schedule_by_block;
   };
   // Plans are stored behind shared_ptr so the per-push readers
   // (GetBlockChunks / GetPoolPushProgressSpec) snapshot with a refcount bump

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<PoolSpec> 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<int64_t> 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<transport::BlockChunk> 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<size_t>(kNumBlocks));
+      for (int64_t block_id = 0; block_id < kNumBlocks; ++block_id) {
+        std::vector<transport::BlockChunk> 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