Stage pool reshard transfers in place on a full host mirror

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",

tpu_sync/kv_cache/kv_cache_manager_with_transfer_pool_reshard_test.cc

@@ -124,7 +124,9 @@ class TestManager : public KVCacheManagerWithTransfer {
     absl::StatusOr<std::shared_ptr<TransferSession>> session = GetSession(uuid);
     ABSL_ASSERT_OK(session);
     static_cast<FakeH2dBase&>(*base_).ArmPoolUpload(pool_idx, status);
-    ABSL_ASSERT_OK((*session)->ResolveStagingSlot(pool_idx));
+    // kNotFound: a full host mirror stages the pool in place.
+    absl::StatusOr<size_t> 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=*/{}));
   }

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<size_t> 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<size_t> 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<size_t> LookupStagingSlot(size_t pool_idx) const;
@@ -561,7 +570,8 @@ ReshardReceiveSession::ReshardReceiveSession(
       start_time_(std::chrono::steady_clock::now()),
       transfer_plan_(
           std::make_shared<const TransferPlan>(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<size_t> 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<size_t> 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<size_t> host_layer_idx;
+    if (HostLayersBounded()) {
+      absl::StatusOr<size_t> 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<raiden::PjRtCopyFuture> copy_future = base_->H2dPoolBlocks(
         pool_idx, dst_chip_block_ids,
         /*shard_idx=*/std::nullopt, uuid(), host_layer_idx);

tpu_sync/kv_cache/reshard_receive_session_test.cc

@@ -26,9 +26,11 @@
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#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<const int64_t> block_ids,
       std::optional<size_t> shard_idx, std::optional<uint64_t> uuid,
       std::optional<size_t> 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<std::pair<size_t, std::optional<size_t>>> uploads() const {
+    absl::MutexLock lock(fake_mu_);
+    return uploads_;
+  }
+
+ private:
+  mutable absl::Mutex fake_mu_;
+  std::vector<std::pair<size_t, std::optional<size_t>>> 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<size_t>(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<kv_cache::PoolSpec> pools = {
+      DensePool("fa", 128),
+      DensePool("state", 128),
+  };
+  pools[1].storage_index = 1;
+  ABSL_ASSERT_OK(base.RegisterPools(pools));
+
+  std::unique_ptr<StagingBlockAllocator> 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<TransferSession> session,
+      CreateReshardReceiveSession(
+          &base, allocator.get(), plan, std::vector<int64_t>{0},
+          std::chrono::steady_clock::now() + std::chrono::seconds(10)));
+  kv_cache::KVCacheManagerBase::TransferEventHooks hooks;
+  hooks.get_session =
+      [weak = std::weak_ptr<TransferSession>(session)](
+          uint64_t uuid) -> absl::StatusOr<std::shared_ptr<TransferSession>> {
+    if (std::shared_ptr<TransferSession> 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<transport::BlockChunk> 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<size_t, std::optional<size_t>>;
+  EXPECT_THAT(base.uploads(),
+              ::testing::UnorderedElementsAre(Upload{0, std::nullopt},
+                                              Upload{1, std::nullopt}));
+  ABSL_EXPECT_OK(session->GetStatus());
+}
+
 }  // namespace
 }  // namespace tpu_raiden

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<size_t> 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<size_t> 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,14 +338,20 @@ ReshardSendSession::ReshardSendSession(
       transfer_plan_(
           std::make_shared<const TransferPlan>(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_);
   ReleaseLeasesLocked();
 }
 
 absl::StatusOr<size_t> 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<size_t> ReshardSendSession::AcquireStagingSlot(size_t pool_idx) {
 
 absl::StatusOr<size_t> 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<size_t> 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<size_t> host_slot;
+    if (absl::StatusOr<size_t> 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<raiden::PjRtCopyFuture> 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);