Stage resharding transfers in plan-allocated host blocks

Resharding resolved every block on the wire by identity into the host
mirror: a sender staged each device block into the host block of the same
id, and a receiver landed each incoming block at its destination id and
uploaded from there. That only works when the host mirror spans the whole
device block space, so both sides of a reshard pinned a host copy of their
entire device KV pool.

Under TPU_RAIDEN_DYNAMIC_HOST_STAGING=1 a plan now allocates host blocks
for the device blocks it names -- its source blocks on a sender, its
destination blocks on a receiver -- and records the mapping. Chunk
resolution consults the mapping, the receiver's upload reads from the
mapped blocks, and the blocks return to the pool when the receive completes
or the plan is unregistered. The wire is unchanged: it still names
device blocks. Senders obtain their staging blocks through
plan_host_blocks() and copy into those instead of into the mirror.

A demand-staged receiver plan is dropped when its transfer settles,
whether or not it was unregistered first, so its mapping never outlives
the blocks it names; until then, pushes already accepted keep landing in
the plan's blocks, and unregistering a plan that already settled is a
no-op. plan_host_blocks() refuses a plan that is not registered, a
pool-addressed plan, and a block the plan's schedules do not name.

Registering and unregistering a plan serialize with each other, and a
plan becomes visible only after the staging and receive state that belong
to it exist, so neither can observe the other half-done. Each registration
carries a generation, and a receive's settlement cleanup only removes the
registration it belongs to, never a newer one reusing the same uuid.

Off by default. Pool-addressed plans keep their existing addressing.

**Validation**

8->4 reshard (prefill TP8 -> decode attention TP4 x DP2), Qwen3-1.7B on
v7x, 128-token pages, 8k context, 64 x 512-token requests at 64-way
concurrency. Flag off vs on, same build and load, two runs per arm. The
build includes the multi-sender completion fix for pushed transfers.

| | off | on (512-block pool) |
|---|---|---|
| host staging per rank | whole device pool mirrored: 1025 blocks, 1.8 GB prefill / 3.5 GB decode | 512 blocks, 0.9 GB / 1.8 GB |
| pinned host memory, both hosts | 42 GB | 21 GB |
| gsm8k, 200 questions (symmetric serve: 0.675, 0.705) | 0.690, 0.670 | 0.685, 0.665 |
| greedy probe, 5 sequential + 8 concurrent repeats | repeatable, equal to the symmetric serve | same |
| req/s | 12.1, 12.8 | 13.0, 12.4 |
| mean TTFT / ITL, ms | 1187 / 28.2, 1225 / 27.0 | 1571 / 23.3, 1208 / 26.8 |
| failed transfers | 0 | 0 |

- The pool no longer tracks the device pool, so the saving grows with the
  model: on gpt-oss-120b TP8 the mirror is 56.5 GiB per rank (452 GiB per
  host); the same 64 in-flight requests fit in about 2.3 GiB per host.
- One probe prompt alternates between two sane continuations across its
  repeats, in both arms.

Author: yinlin-google

tpu_sync/api/torch/kv_cache_manager.py

@@ -207,6 +207,25 @@ def unregister_active_plan(self, uuid: int) -> None:
     """Removes a previously registered strided push plan."""
     self._impl.unregister_active_plan(uuid)
 
+  def plan_host_blocks(self, uuid: int, block_ids: Sequence[int]) -> List[int]:
+    """Host blocks staging ``block_ids`` under a registered plan.
+
+    A sender stages its device blocks into these before pushing.
+
+    Args:
+      uuid: The plan's identifier, as passed to ``register_active_plan``.
+      block_ids: Device block ids the plan names.
+
+    Returns:
+      One host block id per entry of ``block_ids``; the ids themselves when
+      the plan stages blocks at their own ids.
+
+    Raises:
+      RuntimeError: The plan is not registered, or a block is not staged by
+        the plan.
+    """
+    return [int(b) for b in self._impl.plan_host_blocks(uuid, list(block_ids))]
+
   def push_registered_plan(
       self,
       uuid: int,

tpu_sync/core/kv_cache_manager_with_transfer.cc

@@ -204,6 +204,10 @@ class KVCacheManagerWithTransfer : public kv_cache::KVCacheManagerBase {
   absl::Status RegisterActivePlan(
       uint64_t uuid, const ::tpu_sync::rpc::StartTransferRequest& request,
       bool is_sender) override;
+  absl::Status UnregisterActivePlan(uint64_t uuid) override;
+  // Drops the plan of a receive that has settled; a plan already gone, or
+  // a newer registration reusing the uuid, is left alone.
+  void UnregisterSettledPlan(uint64_t uuid, uint64_t generation);
 
   absl::Status RegisterRecv(uint64_t uuid, const std::string& req_id,
                             int64_t expected_block_count) override;
@@ -341,6 +345,18 @@ class KVCacheManagerWithTransfer : public kv_cache::KVCacheManagerBase {
   Slot AcquireSlotLocked();
   void ReleaseSlotLocked(int64_t slot_idx);
   struct RecvEntry;  // defined below; staging helpers take it by pointer
+  // Serializes plan registration and unregistration, so a plan is never
+  // published without its staging owner or torn down against a half-built
+  // registration.
+  absl::Mutex plan_lifecycle_mu_;
+  // Source of plan generations: the uuid names a transfer, a generation
+  // names one registration of it (the same uuid may come back, e.g. on a
+  // retry of the same transfer).
+  uint64_t plan_generation_counter_ ABSL_GUARDED_BY(plan_lifecycle_mu_) = 0;
+  // Host staging held by a plan: a sender's, or a receiver's whose
+  // destination is host memory. Released when the plan is unregistered.
+  absl::flat_hash_map<uint64_t, std::vector<int>> plan_staging_
+      ABSL_GUARDED_BY(mu_);
   // Host staging for one incoming read: exactly `num_blocks` blocks under
   // demand staging, a whole fixed slot otherwise. Returns nullopt when the
   // staging pool cannot seat the request.
@@ -399,6 +415,15 @@ class KVCacheManagerWithTransfer : public kv_cache::KVCacheManagerBase {
     std::chrono::steady_clock::time_point start_time;
     std::vector<raiden::PjRtCopyFuture> h2d_futures;
     bool is_pool_reshard = false;
+    // The plan is dropped when this receive settles: set for every
+    // demand-staged receiver plan (whose mapping would otherwise outlive its
+    // freed blocks) and when an unregister arrives while the receive is in
+    // flight (the plan stays mapped until then so late pushes resolve
+    // through its blocks).
+    bool unregister_on_settle = false;
+    // Generation of the plan this receive belongs to; settlement cleanup
+    // only touches that registration.
+    uint64_t plan_generation = 0;
     std::set<size_t> expected_pool_indices;
     std::set<size_t> started_pool_indices;
     std::set<size_t> completed_pool_indices;

tpu_sync/core/kv_cache_manager_with_transfer.h

@@ -24,7 +24,10 @@
 #include <limits>
 #include <memory>
 #include <optional>
+#include <chrono>
 #include <string>
+#include <thread>
+#include <tuple>
 #include <utility>
 #include <vector>
 
@@ -43,7 +46,9 @@ namespace {
 
 using ::testing::_;
 using ::testing::Ge;
+using ::testing::Contains;
 using ::testing::IsEmpty;
+using ::tpu_sync::rpc::MEMORY_TYPE_DRAM;
 using ::tpu_sync::rpc::MEMORY_TYPE_HBM;
 using ::tpu_sync::rpc::ShardPushEntryProto;
 using ::tpu_sync::rpc::StartTransferRequest;
@@ -68,6 +73,9 @@ class TestManager : public KVCacheManagerWithTransfer {
   // that never touch the holds (validation and the no-bytes-owned sender
   // completion) may rely on it.
   void AttachPlaceholderDeviceHold() { buffer_holds_.emplace_back(); }
+  // Stages plans in per-transfer host blocks, as
+  // TPU_RAIDEN_DYNAMIC_HOST_STAGING=1 does for device-attached managers.
+  void EnableDemandStaging() { dynamic_host_staging_ = true; }
 };
 
 kv_cache::PoolSpec DensePool(std::string tag, int64_t block_stride = 128,
@@ -126,6 +134,34 @@ StartTransferRequest ValidPlan(
   return plan;
 }
 
+// A block-addressed plan (no pool groups) that moves `src` to `dst` on one
+// shard, destined for HBM or for host memory.
+StartTransferRequest BlockPlan(int64_t uuid, const std::vector<int32_t>& src,
+                               const std::vector<int32_t>& dst,
+                               ::tpu_sync::rpc::MemoryType dst_mem_type) {
+  StartTransferRequest plan;
+  plan.set_uuid(uuid);
+  plan.set_req_id("block_plan_req_" + std::to_string(uuid));
+  plan.set_dst_mem_type(dst_mem_type);
+  plan.set_use_block_chunks(true);
+  plan.set_parallelism(1);
+  auto* schedule = &(*plan.mutable_shard_push_schedules())[0];
+  for (size_t i = 0; i < src.size(); ++i) {
+    auto* entry = schedule->add_entries();
+    entry->set_dst_peer("127.0.0.1:1");
+    entry->set_dst_shard_idx(0);
+    entry->set_src_block_id(src[i]);
+    entry->set_dst_block_id(dst[i]);
+    entry->set_src_offset_bytes(0);
+    entry->set_dst_offset_bytes(0);
+    entry->set_size_bytes(16);
+    entry->set_src_stride_bytes(0);
+    entry->set_dst_stride_bytes(0);
+    entry->set_count(1);
+  }
+  return plan;
+}
+
 void ExpectInvalid(const absl::Status& status, const std::string& fragment) {
   EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument)
       << status.ToString();
@@ -686,5 +722,198 @@ TEST(PoolReshardRecvTest, FinishPoolReshardRecvDoesNotRecordMetricOnFailure) {
                                     absl::InternalError("simulated failure"));
 }
 
+
+TEST(SendDeadlineTest, ExpiredSendEntryFailsInsteadOfReportingDone) {
+  TestManager manager(/*timeout_s=*/0.05);
+  ASSERT_GT(manager.NotifyForRead("expired_send_req", 31, {0, 1}), 0);
+
+  std::this_thread::sleep_for(std::chrono::milliseconds(120));
+  const auto [done_sending, done_recving, failed_recving] =
+      manager.CompleteReadRaw();
+  EXPECT_THAT(done_sending, IsEmpty());
+  EXPECT_THAT(failed_recving, Contains("expired_send_req"));
+}
+
+TEST(DemandStagingTest, SenderPlanReturnsStagingOnUnregister) {
+  TestManager manager;
+  manager.EnableDemandStaging();
+  auto* pool = manager.host_block_manager();
+  const int free_before = pool->num_free_blocks();
+  ASSERT_TRUE(manager
+                  .RegisterActivePlan(21, BlockPlan(21, {0, 1}, {2, 3},
+                                                    MEMORY_TYPE_HBM),
+                                      /*is_sender=*/true)
+                  .ok());
+  EXPECT_EQ(pool->num_free_blocks(), free_before - 2);
+  EXPECT_EQ(pool->num_locked_blocks(), 2);
+  ASSERT_TRUE(manager.UnregisterActivePlan(21).ok());
+  EXPECT_EQ(pool->num_free_blocks(), free_before);
+  EXPECT_EQ(pool->num_locked_blocks(), 0);
+}
+
+TEST(DemandStagingTest, HostReceiverPlanReturnsStagingOnUnregister) {
+  TestManager manager;
+  manager.EnableDemandStaging();
+  auto* pool = manager.host_block_manager();
+  const int free_before = pool->num_free_blocks();
+  ASSERT_TRUE(manager
+                  .RegisterActivePlan(22, BlockPlan(22, {0, 1}, {2, 3},
+                                                    MEMORY_TYPE_DRAM),
+                                      /*is_sender=*/false)
+                  .ok());
+  EXPECT_EQ(pool->num_free_blocks(), free_before - 2);
+  EXPECT_EQ(pool->num_locked_blocks(), 2);
+  ASSERT_TRUE(manager.UnregisterActivePlan(22).ok());
+  EXPECT_EQ(pool->num_free_blocks(), free_before);
+  EXPECT_EQ(pool->num_locked_blocks(), 0);
+}
+
+TEST(DemandStagingTest, DuplicateSenderRegistrationKeepsOriginalStaging) {
+  TestManager manager;
+  manager.EnableDemandStaging();
+  auto* pool = manager.host_block_manager();
+  const int free_before = pool->num_free_blocks();
+  const StartTransferRequest plan =
+      BlockPlan(23, {0, 1}, {2, 3}, MEMORY_TYPE_HBM);
+  ASSERT_TRUE(manager.RegisterActivePlan(23, plan, /*is_sender=*/true).ok());
+  const int free_registered = pool->num_free_blocks();
+  EXPECT_EQ(free_registered, free_before - 2);
+
+  // The second registration is refused and must neither keep its own
+  // staging nor disturb the first plan's.
+  const absl::Status duplicate =
+      manager.RegisterActivePlan(23, plan, /*is_sender=*/true);
+  EXPECT_EQ(duplicate.code(), absl::StatusCode::kAlreadyExists)
+      << duplicate.ToString();
+  EXPECT_EQ(pool->num_free_blocks(), free_registered);
+  EXPECT_EQ(pool->num_locked_blocks(), 2);
+
+  ASSERT_TRUE(manager.UnregisterActivePlan(23).ok());
+  EXPECT_EQ(pool->num_free_blocks(), free_before);
+  EXPECT_EQ(pool->num_locked_blocks(), 0);
+}
+
+
+TEST(DemandStagingTest, UnregisteringInFlightReceiverDefersUntilItSettles) {
+  TestManager manager(/*timeout_s=*/0.05);
+  manager.EnableDemandStaging();
+  manager.AttachPlaceholderDeviceHold();
+  auto* pool = manager.host_block_manager();
+  const int free_before = pool->num_free_blocks();
+  ASSERT_TRUE(manager
+                  .RegisterActivePlan(24, BlockPlan(24, {0, 1}, {2, 3},
+                                                    MEMORY_TYPE_HBM),
+                                      /*is_sender=*/false)
+                  .ok());
+  EXPECT_EQ(pool->num_free_blocks(), free_before - 2);
+  const auto staged = manager.PlanHostBlocks(24, {2, 3});
+  ASSERT_TRUE(staged.ok()) << staged.status();
+
+  // Unregistering while the receive is in flight keeps the plan's mapping
+  // and its staging, so pushes already accepted still land in the plan's
+  // blocks ...
+  ASSERT_TRUE(manager.UnregisterActivePlan(24).ok());
+  const auto still_staged = manager.PlanHostBlocks(24, {2, 3});
+  ASSERT_TRUE(still_staged.ok()) << still_staged.status();
+  EXPECT_EQ(*still_staged, *staged);
+  EXPECT_EQ(pool->num_free_blocks(), free_before - 2);
+
+  // ... until the receive settles; here it times out. Then the plan, its
+  // staging and the receive entry go together.
+  std::this_thread::sleep_for(std::chrono::milliseconds(120));
+  const auto [done_sending, done_recving, failed_recving] =
+      manager.CompleteReadRaw();
+  EXPECT_THAT(failed_recving, Contains("block_plan_req_24"));
+  EXPECT_EQ(manager.PlanHostBlocks(24, {2, 3}).status().code(),
+            absl::StatusCode::kNotFound);
+  EXPECT_EQ(pool->num_free_blocks(), free_before);
+  EXPECT_EQ(pool->num_locked_blocks(), 0);
+}
+
+TEST(DemandStagingTest, PlanHostBlocksFailsClosed) {
+  TestManager manager;
+  manager.EnableDemandStaging();
+  auto* pool = manager.host_block_manager();
+  EXPECT_EQ(manager.PlanHostBlocks(25, {0}).status().code(),
+            absl::StatusCode::kNotFound);
+
+  // Occupy the identity blocks first, so the plan's host blocks provably
+  // differ from its device blocks.
+  const auto occupied = pool->Allocate(2, /*lock=*/true);
+  ASSERT_TRUE(occupied.ok());
+  ASSERT_TRUE(manager
+                  .RegisterActivePlan(25, BlockPlan(25, {0, 1}, {2, 3},
+                                                    MEMORY_TYPE_HBM),
+                                      /*is_sender=*/true)
+                  .ok());
+  const auto staged = manager.PlanHostBlocks(25, {0, 1});
+  ASSERT_TRUE(staged.ok()) << staged.status();
+  ASSERT_EQ(staged->size(), 2u);
+  EXPECT_NE(*staged, (std::vector<int64_t>{0, 1}));
+  EXPECT_TRUE(pool->IsLocked(static_cast<int>((*staged)[0])));
+  EXPECT_TRUE(pool->IsLocked(static_cast<int>((*staged)[1])));
+  // A block the plan does not stage is refused rather than passed through.
+  EXPECT_EQ(manager.PlanHostBlocks(25, {0, 7}).status().code(),
+            absl::StatusCode::kInvalidArgument);
+
+  ASSERT_TRUE(manager.UnregisterActivePlan(25).ok());
+  EXPECT_EQ(manager.PlanHostBlocks(25, {0}).status().code(),
+            absl::StatusCode::kNotFound);
+}
+
+TEST(DemandStagingTest, PlanHostBlocksIsIdentityForFixedStaging) {
+  TestManager manager;
+  ASSERT_TRUE(manager
+                  .RegisterActivePlan(26, BlockPlan(26, {0, 1}, {2, 3},
+                                                    MEMORY_TYPE_HBM),
+                                      /*is_sender=*/true)
+                  .ok());
+  const auto staged = manager.PlanHostBlocks(26, {0, 1});
+  ASSERT_TRUE(staged.ok()) << staged.status();
+  EXPECT_EQ(*staged, (std::vector<int64_t>{0, 1}));
+  // Identity covers only the blocks the plan names.
+  EXPECT_EQ(manager.PlanHostBlocks(26, {7}).status().code(),
+            absl::StatusCode::kInvalidArgument);
+  ASSERT_TRUE(manager.UnregisterActivePlan(26).ok());
+}
+
+TEST(DemandStagingTest, PlanHostBlocksRejectsBlocksOfAnEmptyPlan) {
+  TestManager manager;
+  manager.EnableDemandStaging();
+  ASSERT_TRUE(manager
+                  .RegisterActivePlan(28, BlockPlan(28, {}, {},
+                                                    MEMORY_TYPE_HBM),
+                                      /*is_sender=*/true)
+                  .ok());
+  EXPECT_EQ(manager.PlanHostBlocks(28, {0}).status().code(),
+            absl::StatusCode::kInvalidArgument);
+  ASSERT_TRUE(manager.UnregisterActivePlan(28).ok());
+}
+
+TEST(DemandStagingTest, DemandStagedReceiverPlanUnregistersWhenItSettles) {
+  TestManager manager(/*timeout_s=*/0.05);
+  manager.EnableDemandStaging();
+  manager.AttachPlaceholderDeviceHold();
+  auto* pool = manager.host_block_manager();
+  const int free_before = pool->num_free_blocks();
+  ASSERT_TRUE(manager
+                  .RegisterActivePlan(27, BlockPlan(27, {0, 1}, {2, 3},
+                                                    MEMORY_TYPE_HBM),
+                                      /*is_sender=*/false)
+                  .ok());
+  ASSERT_TRUE(manager.PlanHostBlocks(27, {2, 3}).ok());
+
+  // Nobody unregisters; the plan still goes when the receive settles (here
+  // by timeout), leaving neither a stale mapping nor held blocks behind.
+  std::this_thread::sleep_for(std::chrono::milliseconds(120));
+  const auto [done_sending, done_recving, failed_recving] =
+      manager.CompleteReadRaw();
+  EXPECT_THAT(failed_recving, Contains("block_plan_req_27"));
+  EXPECT_EQ(manager.PlanHostBlocks(27, {2, 3}).status().code(),
+            absl::StatusCode::kNotFound);
+  EXPECT_EQ(pool->num_free_blocks(), free_before);
+  EXPECT_EQ(pool->num_locked_blocks(), 0);
+}
+
 }  // namespace
 }  // namespace tpu_raiden

tpu_sync/core/kv_cache_manager_with_transfer_pool_reshard_test.cc

@@ -238,6 +238,11 @@ class KVCacheManager {
     return torch_manager_->UnregisterActivePlan(uuid);
   }
 
+  absl::StatusOr<std::vector<int64_t>> PlanHostBlocks(
+      uint64_t uuid, const std::vector<int64_t>& block_ids) {
+    return torch_manager_->PlanHostBlocks(uuid, block_ids);
+  }
+
   absl::Status RegisterRecv(uint64_t uuid, const std::string& req_id,
                             int64_t expected_block_count) {
     return torch_manager_->RegisterRecv(uuid, req_id, expected_block_count);

tpu_sync/frameworks/torch/kv_cache_manager.h

@@ -242,7 +242,11 @@ NB_MODULE(_tpu_raiden_host, m) {
       .def(
           "unregister_active_plan",
           [](HostKVCacheManager& self, uint64_t uuid) {
-            ThrowIfError(self.UnregisterActivePlan(uuid),
+            absl::Status status = self.UnregisterActivePlan(uuid);
+            // A plan that already settled is gone; unregistering it again is
+            // a no-op rather than an error.
+            if (absl::IsNotFound(status)) return;
+            ThrowIfError(status,
                          "KVCacheManager unregister_active_plan failed");
           },
           nb::arg("uuid"))

tpu_sync/frameworks/torch/tpu_raiden_host_module.cc

@@ -272,13 +272,27 @@ NB_MODULE(_tpu_raiden_torch, m) {
           "unregister_active_plan",
           [](KVCacheManager& self, uint64_t uuid) {
             absl::Status status = self.UnregisterActivePlan(uuid);
+            // A plan that already settled is gone; unregistering it again is
+            // a no-op rather than an error.
+            if (absl::IsNotFound(status)) return;
             if (!status.ok()) {
               throw std::runtime_error(
                   "KVCacheManager unregister_active_plan failed: " +
                   std::string(status.message()));
             }
           },
           nb::arg("uuid"))
+      .def(
+          "plan_host_blocks",
+          [](KVCacheManager& self, uint64_t uuid,
+             const std::vector<int64_t>& block_ids) {
+            auto blocks = self.PlanHostBlocks(uuid, block_ids);
+            if (!blocks.ok()) {
+              throw std::runtime_error(std::string(blocks.status().message()));
+            }
+            return *std::move(blocks);
+          },
+          nb::arg("uuid"), nb::arg("block_ids"))
       .def(
           "push_registered_plan",
           [](KVCacheManager& self, uint64_t uuid, const std::string& peer,