Reshard KV pools between pipeline-parallel engines

tpu_sync/core/reshard_send_session.cc

@@ -308,6 +308,14 @@ void ReshardSendSession::StartPoolPush(KVCacheManagerWithTransfer& manager,
     in_flight_ += static_cast<int>(transfers_by_peer.size());
   }
 
+  // The plan names this pool in the sender's own index space; the receiver
+  // resolves the wire index against the destination's pool table.
+  std::optional<int> wire_pool_idx;
+  auto wire_it = plan_.wire_pool_indices().find(static_cast<int32_t>(pool_idx));
+  if (wire_it != plan_.wire_pool_indices().end()) {
+    wire_pool_idx = wire_it->second;
+  }
+
   for (const auto& [peer, transfers] : transfers_by_peer) {
     std::vector<int> src_ids;
     std::vector<int> dst_ids;
@@ -325,7 +333,8 @@ void ReshardSendSession::StartPoolPush(KVCacheManagerWithTransfer& manager,
           absl::Cleanup end_op = [self]() { self->EndOp(); };
           self->RecordPushCompletion(
               manager, result.ok() ? absl::OkStatus() : result.status());
-        });
+        },
+        wire_pool_idx);
   }
 }
 

tpu_sync/core/transfer_program_reshard.cc

@@ -168,6 +168,9 @@ absl::StatusOr<::tpu_sync::proto::TransferProgramRequest> CompileStartTransfer(
   *binding->mutable_transfer_pool_indices() =
       request.transfer_pool_indices();
   *binding->mutable_pool_dtype_tags() = request.pool_dtype_tags();
+  for (const auto& [local, wire] : request.wire_pool_indices()) {
+    (*binding->mutable_wire_pool_indices())[local] = wire;
+  }
 
   ::tpu_sync::proto::FanIn* fan_in =
       program->mutable_completion()->mutable_fan_in();
@@ -241,6 +244,9 @@ absl::StatusOr<::tpu_sync::rpc::StartTransferRequest> LowerToStartTransfer(
   *out.mutable_dst_units() = binding.dst_units();
   *out.mutable_transfer_pool_indices() = binding.transfer_pool_indices();
   *out.mutable_pool_dtype_tags() = binding.pool_dtype_tags();
+  for (const auto& [local, wire] : binding.wire_pool_indices()) {
+    (*out.mutable_wire_pool_indices())[local] = wire;
+  }
 
   const ::tpu_sync::proto::ExecutionPolicy& policy = program.policy();
   out.set_parallelism(policy.parallelism());

tpu_sync/core/transfer_program_reshard_test.cc

@@ -160,6 +160,17 @@ TEST(TransferProgramReshard, SkipTilingMapRoundTrips) {
   EXPECT_EQ(Canonical(original), Canonical(*lowered));
 }
 
+TEST(TransferProgramReshard, WirePoolIndexMapRoundTrips) {
+  ::tpu_sync::rpc::StartTransferRequest original = MakePoolRequest(true);
+  (*original.mutable_wire_pool_indices())[0] = 9;
+  (*original.mutable_wire_pool_indices())[1] = 10;
+  auto program = CompileStartTransfer(original);
+  ASSERT_TRUE(program.ok()) << program.status();
+  auto lowered = LowerToStartTransfer(*program);
+  ASSERT_TRUE(lowered.ok()) << lowered.status();
+  EXPECT_EQ(Canonical(original), Canonical(*lowered));
+}
+
 TEST(TransferProgramReshard, LegacyDensePlanRefusesToCompile) {
   ::tpu_sync::rpc::StartTransferRequest legacy;
   legacy.set_uuid(1);

tpu_sync/core/utils.h

@@ -39,6 +39,7 @@
 #include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
+#include "absl/strings/str_join.h"
 #include "absl/strings/string_view.h"
 #include "absl/strings/strip.h"
 #include "absl/types/span.h"
@@ -409,10 +410,17 @@ inline absl::Status ValidateCommonPoolReshardPlan(
           absl::StrCat("transfer pool index out of range: ", pool_idx));
     }
     if (plan.pool_dtype_tags(pool_idx) != spec->dtype_tag) {
+      std::string local_pools;
+      for (size_t i = 0; i < base->num_pools(); ++i) {
+        absl::StrAppend(&local_pools, i ? "," : "", base->pool(i)->tag, ":",
+                        base->pool(i)->dtype_tag);
+      }
       return absl::InvalidArgumentError(
           absl::StrCat("plan dtype tag mismatch for pool ", pool_idx, " (",
                        spec->tag, "): plan=", plan.pool_dtype_tags(pool_idx),
-                       " local=", spec->dtype_tag));
+                       " local=", spec->dtype_tag, "; plan dtype tags=[",
+                       absl::StrJoin(plan.pool_dtype_tags(), ","),
+                       "] local pools=[", local_pools, "]"));
     }
   }
   if (plan.shard_push_schedules().empty()) {

tpu_sync/kv_cache/kv_cache_manager_base.cc

@@ -45,6 +45,7 @@
 #include "absl/status/statusor.h"
 #include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
 #include "absl/synchronization/notification.h"
@@ -2797,10 +2798,17 @@ absl::Status KVCacheManagerBase::RegisterActivePlan(
     }
     for (size_t pool_idx = 0; pool_idx < pools_.size(); ++pool_idx) {
       if (request.pool_dtype_tags(pool_idx) != pools_[pool_idx].dtype_tag) {
+        std::string local_pools;
+        for (size_t i = 0; i < pools_.size(); ++i) {
+          absl::StrAppend(&local_pools, i ? "," : "", pools_[i].tag, ":",
+                          pools_[i].dtype_tag);
+        }
         return absl::InvalidArgumentError(absl::StrCat(
             "plan dtype tag mismatch for pool ", pool_idx, " (",
             pools_[pool_idx].tag, "): plan=", request.pool_dtype_tags(pool_idx),
-            " local=", pools_[pool_idx].dtype_tag));
+            " local=", pools_[pool_idx].dtype_tag, "; plan dtype tags=[",
+            absl::StrJoin(request.pool_dtype_tags(), ","), "] local pools=[",
+            local_pools, "]"));
       }
     }
   }

tpu_sync/kv_cache/reshard/BUILD

@@ -103,6 +103,7 @@ cc_library(
         "//tpu_sync/kv_cache:pool_layout",
         "//tpu_sync/rpc:raiden_service_cc_proto",
         "@com_google_absl//absl/container:btree",
+        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",

tpu_sync/kv_cache/reshard/pool_reshard_planner.cc

@@ -30,6 +30,7 @@
 
 #include "absl/container/btree_map.h"
 #include "absl/container/btree_set.h"
+#include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
@@ -122,7 +123,7 @@ MetadataByUnit(
 }
 
 struct TagPrecheck {
-  std::vector<int32_t> selected;
+  std::vector<int32_t> selected;  // destination pool indices
   int64_t src_live = 0;
   int64_t dst_live = 0;
   std::vector<PoolLiveSegment> src_segments;
@@ -239,18 +240,40 @@ absl::StatusOr<PoolReshardPlan> BuildPoolReshardPlan(
     return absl::InvalidArgumentError(
         "Destination pool manifest must not be empty");
   }
+  // Pools pair up by tag. Source and destination units may each register a
+  // subset of the model's tags (a pipeline stage holds a layer subset);
+  // every tag both sides register must carry the same pool count and dtype
+  // for that tag, in manifest order.
+  absl::flat_hash_map<std::string, std::vector<int32_t>> dst_pools_by_tag;
+  for (int i = 0; i < dst_meta.pools_size(); ++i) {
+    dst_pools_by_tag[dst_meta.pools(i).tag()].push_back(i);
+  }
+  std::map<RaidenId, absl::flat_hash_map<std::string, std::vector<int32_t>>,
+           RequestBlockRegistry::RaidenIdLess>
+      src_pools_by_tag;
   for (const RaidenId& src_unit : request.src_units) {
     const auto& src_pools = src_by_unit.at(src_unit)->pools();
-    std::vector<std::pair<std::string, std::string>> src_identity;
-    src_identity.reserve(src_pools.size());
-    for (const auto& pool : src_pools) {
-      src_identity.emplace_back(pool.tag(), pool.dtype_tag());
+    auto& by_tag = src_pools_by_tag[src_unit];
+    for (int i = 0; i < src_pools.size(); ++i) {
+      by_tag[src_pools.Get(i).tag()].push_back(i);
     }
-    if (src_identity != dst_identity) {
-      return absl::InvalidArgumentError(absl::StrCat(
-          "Canonical pool manifest mismatch between source and destination "
-          "for ",
-          PythonRepr(src_unit)));
+    for (const auto& [tag, src_indices] : by_tag) {
+      auto dst_it = dst_pools_by_tag.find(tag);
+      // A source tag the destination lacks is simply not transferable
+      // (a pipeline-parallel destination stage holds a layer subset too);
+      // requesting it fails below with "do not match any registered pool".
+      if (dst_it == dst_pools_by_tag.end()) continue;
+      bool matches = dst_it->second.size() == src_indices.size();
+      for (size_t k = 0; matches && k < src_indices.size(); ++k) {
+        matches = src_pools.Get(src_indices[k]).dtype_tag() ==
+                  dst_meta.pools(dst_it->second[k]).dtype_tag();
+      }
+      if (!matches) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Canonical pool manifest mismatch between source and destination "
+            "for ",
+            PythonRepr(src_unit), " at tag ", PyStrRepr(tag)));
+      }
     }
   }
   {
@@ -285,23 +308,28 @@ absl::StatusOr<PoolReshardPlan> BuildPoolReshardPlan(
     }
   }
 
-  const tpu_sync::rpc::RegisterWorkUnitRequest& reference_src =
-      *src_by_unit.at(request.src_units[0]);
+  // Source geometry is compared per tag across the ranks that register the
+  // tag; the first such rank (in request order) is the tag's reference.
+  absl::flat_hash_map<std::string, RaidenId> tag_reference_unit;
   {
-    std::vector<std::string> reference_geometry;
-    for (const auto& pool : reference_src.pools()) {
-      reference_geometry.push_back(GeometrySignature(pool));
-    }
-    for (size_t i = 1; i < request.src_units.size(); ++i) {
-      const RaidenId& src_unit = request.src_units[i];
-      std::vector<std::string> geometry;
-      for (const auto& pool : src_by_unit.at(src_unit)->pools()) {
-        geometry.push_back(GeometrySignature(pool));
-      }
-      if (geometry != reference_geometry) {
-        return absl::InvalidArgumentError(
-            absl::StrCat("Source pool geometry differs across ranks at ",
-                         PythonRepr(src_unit)));
+    absl::flat_hash_map<std::string, std::vector<std::string>>
+        reference_geometry;
+    for (const RaidenId& src_unit : request.src_units) {
+      const auto& src_pools = src_by_unit.at(src_unit)->pools();
+      for (const auto& [tag, src_indices] : src_pools_by_tag.at(src_unit)) {
+        std::vector<std::string> geometry;
+        geometry.reserve(src_indices.size());
+        for (int32_t idx : src_indices) {
+          geometry.push_back(GeometrySignature(src_pools.Get(idx)));
+        }
+        auto [it, inserted] = reference_geometry.emplace(tag, geometry);
+        if (inserted) {
+          tag_reference_unit.emplace(tag, src_unit);
+        } else if (it->second != geometry) {
+          return absl::InvalidArgumentError(
+              absl::StrCat("Source pool geometry differs across ranks at ",
+                           PythonRepr(src_unit), " for tag ", PyStrRepr(tag)));
+        }
       }
     }
   }
@@ -450,13 +478,26 @@ absl::StatusOr<PoolReshardPlan> BuildPoolReshardPlan(
         precheck.selected.push_back(i);
       }
     }
+    // The source unit whose pools define the tag's source geometry, and its
+    // pool indices for the tag (aligned 1:1 with `selected`).
+    auto ref_it = tag_reference_unit.find(plan_tag);
+    if (ref_it == tag_reference_unit.end()) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "No source unit registers pools for tag ", PyStrRepr(plan_tag)));
+    }
+    const RaidenId& src_reference_unit = ref_it->second;
+    const std::vector<int32_t>& src_selected =
+        src_pools_by_tag.at(src_reference_unit).at(plan_tag);
+    const tpu_sync::rpc::RegisterWorkUnitRequest& reference_src =
+        *src_by_unit.at(src_reference_unit);
 
     std::set<int64_t> src_live_values;
     std::set<int64_t> dst_live_values;
     std::vector<std::vector<PoolLiveSegment>> src_segment_maps;
     std::vector<std::vector<PoolLiveSegment>> dst_segment_maps;
-    for (int32_t pool_idx : precheck.selected) {
-      const auto& src_pool = reference_src.pools(pool_idx);
+    for (size_t k = 0; k < precheck.selected.size(); ++k) {
+      const int32_t pool_idx = precheck.selected[k];
+      const auto& src_pool = reference_src.pools(src_selected[k]);
       const auto& dst_pool = dst_meta.pools(pool_idx);
       auto src_segments = LiveSegments(src_pool);
       if (!src_segments.ok()) return src_segments.status();
@@ -561,9 +602,14 @@ absl::StatusOr<PoolReshardPlan> BuildPoolReshardPlan(
       const RequestBlockRegistration& registration = registrations.at(unit);
       for (const PoolSpanRegistration& entry : registration.pool_spans) {
         if (entry.tag != plan_tag) continue;
-        if (!entry.spans.empty()) {
-          declared.emplace_back(unit, &entry);
+        if (entry.spans.empty()) continue;
+        if (src_pools_by_tag.at(unit).find(plan_tag) ==
+            src_pools_by_tag.at(unit).end()) {
+          return absl::InvalidArgumentError(absl::StrCat(
+              "Byte spans are declared for tag ", PyStrRepr(plan_tag), " by ",
+              PythonRepr(unit), ", which registers no pool with that tag"));
         }
+        declared.emplace_back(unit, &entry);
       }
     }
     if (declared.empty()) {
@@ -655,11 +701,13 @@ absl::StatusOr<PoolReshardPlan> BuildPoolReshardPlan(
     }
     declared = std::move(converted);
 
-    for (int32_t pool_idx : precheck.selected) {
+    for (size_t k = 0; k < precheck.selected.size(); ++k) {
+      const int32_t pool_idx = precheck.selected[k];
       const int64_t dst_num_blocks = dst_meta.pools(pool_idx).num_blocks();
       for (const auto& [unit, entry] : declared) {
+        const int32_t src_pool_idx = src_pools_by_tag.at(unit).at(plan_tag)[k];
         const int64_t limit =
-            src_by_unit.at(unit)->pools(pool_idx).num_blocks();
+            src_by_unit.at(unit)->pools(src_pool_idx).num_blocks();
         for (int64_t block_id : entry->block_ids) {
           if (block_id >= limit) {
             return absl::InvalidArgumentError(
@@ -980,10 +1028,26 @@ absl::StatusOr<PoolReshardPlan> BuildPoolReshardPlan(
   for (const auto& pool : dst_meta.pools()) {
     plan.pool_dtype_tags.push_back(pool.dtype_tag());
   }
+  for (const RaidenId& unit : plan.src_units) {
+    const auto& src_pools = src_by_unit.at(unit)->pools();
+    std::vector<std::string>& dtype_tags = plan.src_pool_dtype_tags[unit];
+    for (const auto& pool : src_pools) {
+      dtype_tags.push_back(pool.dtype_tag());
+    }
+    std::map<int32_t, int32_t>& remap = plan.src_pool_indices[unit];
+    const auto& by_tag = src_pools_by_tag.at(unit);
+    for (const TagPrecheck& precheck : tag_precheck) {
+      auto tag_it = by_tag.find(dst_meta.pools(precheck.selected[0]).tag());
+      if (tag_it == by_tag.end()) continue;
+      for (size_t k = 0; k < precheck.selected.size(); ++k) {
+        remap[precheck.selected[k]] = tag_it->second[k];
+      }
+    }
+  }
   plan.dst_device_block_ids = dst_ids;
   for (size_t ordinal = 0; ordinal < plan.src_units.size(); ++ordinal) {
-    plan.src_schedule_keys[plan.src_units[ordinal]] =
-        static_cast<int32_t>(ordinal);
+    plan.src_schedule_keys[plan.src_units[ordinal]] = static_cast<int32_t>(
+        src_by_unit.at(plan.src_units[ordinal])->transfer_rank());
   }
   plan.parallelism = static_cast<int32_t>(requested_parallelism);
   plan.num_tokens = std::max<int64_t>(request.num_tokens, 0);

tpu_sync/kv_cache/reshard/pool_reshard_planner.h

@@ -84,7 +84,20 @@ struct PoolReshardPlan {
   int32_t expected_pushes_per_pool = 0;
   std::vector<int32_t> transfer_pool_indices;
   std::vector<std::string> pool_dtype_tags;
+  // Pools pair up by tag, and a source unit may register only a subset of
+  // the destination's tags. Per source unit: destination pool index ->
+  // that unit's own pool index for every transferred pool it registers,
+  // and its complete per-pool dtype tag list. The coordinator rewrites a
+  // sender's request into the sender's pool index space with these.
+  std::map<RaidenId, std::map<int32_t, int32_t>,
+           RequestBlockRegistry::RaidenIdLess>
+      src_pool_indices;
+  std::map<RaidenId, std::vector<std::string>,
+           RequestBlockRegistry::RaidenIdLess>
+      src_pool_dtype_tags;
   std::vector<int64_t> dst_device_block_ids;
+  // A source's schedule key is its registered transfer rank: the worker id
+  // it dispatches under and the node id a receiver resolves its pushes by.
   std::map<RaidenId, int32_t, RequestBlockRegistry::RaidenIdLess>
       src_schedule_keys;
   int32_t parallelism = 1;