From 0976a5fb5c3d6d661ce831abecd1236b25b8b34d Mon Sep 17 00:00:00 2001 From: yinlin Date: Fri, 25 Sep 2026 23:55:04 -0700 Subject: [PATCH] Reshard KV pools between pipeline-parallel engines A pipeline-parallel engine splits its layers over its stages: each stage registers only the pools of its own layers, every decode stage plans the same request once, and a decode stage's plan names a single prefill stage. Pool resharding assumed the opposite in four places: 1. Source and destination pools paired by position and the two pool tables had to be equal, so planning failed with "Canonical pool manifest mismatch" as soon as a producer stage registered a layer subset. 2. A pushed pool was named by the sender's own pool index, which the receiver resolved against its own table, so every stage but the first had its pushes rejected with "No transfer chunks found for block N". 3. A request's block snapshot could be claimed by one planning attempt, so the first decode stage claimed it and the other stages were refused with "already claimed by another planning attempt" and timed out. 4. A source was keyed by its position in the plan's source list instead of its transfer rank, so a plan naming one source ran on worker_0 and the receiver found no schedule for the sender's node id. Changes: - Pool pairing: pools pair by tag. A source may register a subset of the destination's tags and a destination a subset of the source's; tags on both sides must agree on pool count and dtype. A sender's start-transfer request is rewritten into its own pool index space. - Wire pool index: StartTransferRequest carries wire_pool_indices (sender index -> destination index, the identity for equal pool tables), carried through the transfer-program reshard binding; a sender's push names the destination's block array. - Shared claims: planning attempts that name the same source unit set share one request claim; a different unit set is still refused; the claim is freed when the last attempt leaves. Dtype-tag mismatch errors name the plan's dtype list and the local pool table. - Rank keys: each source's schedule and each sender's worker are keyed by the source's registered transfer rank. Validation: the OSS build passes all 7 reshard, session and transport test targets (reshard_service_test 29 cases incl. SubsetSourceManifestsPairPoolsByTag, SourceTagsAbsentOnTheDestinationAreIgnored, PipelinedDestinationStagesShareOneRequestClaim, AbandonedStageLeavesTheSiblingClaimInPlace and AbandoningTheLastStageFreesTheClaim; block_transport_test 30 cases incl. PushNamesTheReceiversArrayByWireLayerIndex; transfer_program_reshard_test incl. WirePoolIndexMapRoundTrips). End to end, PP8 -> DP8 and PP8 -> PP8 pairs on Qwen3.5-35B-A3B-FP8 complete every request and answer every greedy probe whose KV crosses the reshard (vllm-torchtpu PR 871). PiperOrigin-RevId: 988715760 --- tpu_sync/core/reshard_send_session.cc | 11 +- tpu_sync/core/transfer_program_reshard.cc | 6 + .../core/transfer_program_reshard_test.cc | 11 + tpu_sync/core/utils.h | 10 +- tpu_sync/kv_cache/kv_cache_manager_base.cc | 10 +- tpu_sync/kv_cache/reshard/BUILD | 1 + .../kv_cache/reshard/pool_reshard_planner.cc | 132 +++-- .../kv_cache/reshard/pool_reshard_planner.h | 13 + .../reshard/request_block_registry.cc | 34 +- .../kv_cache/reshard/request_block_registry.h | 6 +- .../kv_cache/reshard/reshard_coordinator.cc | 48 +- .../kv_cache/reshard/reshard_service_test.cc | 466 ++++++++++++++++++ tpu_sync/proto/transfer_program.proto | 3 + tpu_sync/rpc/raiden_service.proto | 5 + tpu_sync/transport/block_transport.cc | 19 +- tpu_sync/transport/block_transport.h | 11 +- tpu_sync/transport/block_transport_test.cc | 67 +++ 17 files changed, 778 insertions(+), 75 deletions(-) diff --git a/tpu_sync/core/reshard_send_session.cc b/tpu_sync/core/reshard_send_session.cc index 9a72aa000..961bccba9 100644 --- a/tpu_sync/core/reshard_send_session.cc +++ b/tpu_sync/core/reshard_send_session.cc @@ -308,6 +308,14 @@ void ReshardSendSession::StartPoolPush(KVCacheManagerWithTransfer& manager, in_flight_ += static_cast(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 wire_pool_idx; + auto wire_it = plan_.wire_pool_indices().find(static_cast(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 src_ids; std::vector 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); } } diff --git a/tpu_sync/core/transfer_program_reshard.cc b/tpu_sync/core/transfer_program_reshard.cc index dff216aa7..b42cf7ab8 100644 --- a/tpu_sync/core/transfer_program_reshard.cc +++ b/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()); diff --git a/tpu_sync/core/transfer_program_reshard_test.cc b/tpu_sync/core/transfer_program_reshard_test.cc index ee44e0abf..ba92e051b 100644 --- a/tpu_sync/core/transfer_program_reshard_test.cc +++ b/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); diff --git a/tpu_sync/core/utils.h b/tpu_sync/core/utils.h index 3e6e71207..8c929d6c7 100644 --- a/tpu_sync/core/utils.h +++ b/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()) { diff --git a/tpu_sync/kv_cache/kv_cache_manager_base.cc b/tpu_sync/kv_cache/kv_cache_manager_base.cc index a85eb8b0f..da3555a72 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_base.cc +++ b/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, "]")); } } } diff --git a/tpu_sync/kv_cache/reshard/BUILD b/tpu_sync/kv_cache/reshard/BUILD index 848de344b..9cb74a9d3 100644 --- a/tpu_sync/kv_cache/reshard/BUILD +++ b/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", diff --git a/tpu_sync/kv_cache/reshard/pool_reshard_planner.cc b/tpu_sync/kv_cache/reshard/pool_reshard_planner.cc index cbcfa8a3c..677b48b21 100644 --- a/tpu_sync/kv_cache/reshard/pool_reshard_planner.cc +++ b/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 selected; + std::vector selected; // destination pool indices int64_t src_live = 0; int64_t dst_live = 0; std::vector src_segments; @@ -239,18 +240,40 @@ absl::StatusOr 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> 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>, + 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> 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 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 tag_reference_unit; { - std::vector 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 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> + 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 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 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& 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 src_live_values; std::set dst_live_values; std::vector> src_segment_maps; std::vector> 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 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 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 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& dtype_tags = plan.src_pool_dtype_tags[unit]; + for (const auto& pool : src_pools) { + dtype_tags.push_back(pool.dtype_tag()); + } + std::map& 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(ordinal); + plan.src_schedule_keys[plan.src_units[ordinal]] = static_cast( + src_by_unit.at(plan.src_units[ordinal])->transfer_rank()); } plan.parallelism = static_cast(requested_parallelism); plan.num_tokens = std::max(request.num_tokens, 0); diff --git a/tpu_sync/kv_cache/reshard/pool_reshard_planner.h b/tpu_sync/kv_cache/reshard/pool_reshard_planner.h index ee5d1c35b..2d619d8ed 100644 --- a/tpu_sync/kv_cache/reshard/pool_reshard_planner.h +++ b/tpu_sync/kv_cache/reshard/pool_reshard_planner.h @@ -84,7 +84,20 @@ struct PoolReshardPlan { int32_t expected_pushes_per_pool = 0; std::vector transfer_pool_indices; std::vector 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, + RequestBlockRegistry::RaidenIdLess> + src_pool_indices; + std::map, + RequestBlockRegistry::RaidenIdLess> + src_pool_dtype_tags; std::vector 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 src_schedule_keys; int32_t parallelism = 1; diff --git a/tpu_sync/kv_cache/reshard/request_block_registry.cc b/tpu_sync/kv_cache/reshard/request_block_registry.cc index c6d31934c..cc0b86257 100644 --- a/tpu_sync/kv_cache/reshard/request_block_registry.cc +++ b/tpu_sync/kv_cache/reshard/request_block_registry.cc @@ -452,20 +452,11 @@ RequestBlockRegistry::LookupAndClaim(const std::string& req_id, int64_t uuid, } std::set claimed_units(units.begin(), units.end()); auto existing_units_it = claimed_units_.find(lifecycle_key); - if (existing_units_it != claimed_units_.end()) { - auto owner_it = claimed_owners_.find(lifecycle_key); - const void* existing_owner = - owner_it == claimed_owners_.end() ? nullptr : owner_it->second; - if (existing_owner != claim_owner) { - return absl::InvalidArgumentError( - "Request block snapshot is already claimed by another planning " - "attempt"); - } - if (existing_units_it->second != claimed_units) { - return absl::InvalidArgumentError( - "Request block snapshot was already claimed for a different " - "source unit set"); - } + if (existing_units_it != claimed_units_.end() && + existing_units_it->second != claimed_units) { + return absl::InvalidArgumentError( + "Request block snapshot was already claimed for a different " + "source unit set"); } std::map result; for (const RaidenId& unit : units) { @@ -481,7 +472,7 @@ RequestBlockRegistry::LookupAndClaim(const std::string& req_id, int64_t uuid, // validated and copied while cancellation is excluded by the shared lock. claimed_[lifecycle_key] = now + ttl_s_; claimed_units_[lifecycle_key] = claimed_units; - claimed_owners_[lifecycle_key] = claim_owner; + claimed_owners_[lifecycle_key].insert(claim_owner); auto completion_it = completed_units_.find(lifecycle_key); if (completion_it != completed_units_.end()) { completion_it->second.expires_at = now + ttl_s_; @@ -505,14 +496,15 @@ bool RequestBlockRegistry::AbandonClaim(const std::string& req_id, int64_t uuid, absl::MutexLock lock(*mu_); auto claimed_it = claimed_.find(lifecycle_key); auto owner_it = claimed_owners_.find(lifecycle_key); - const void* existing_owner = - owner_it == claimed_owners_.end() ? nullptr : owner_it->second; - if (claimed_it == claimed_.end() || existing_owner != claim_owner) { + if (claimed_it == claimed_.end() || owner_it == claimed_owners_.end() || + owner_it->second.erase(claim_owner) == 0) { return false; } - claimed_.erase(lifecycle_key); - claimed_units_.erase(lifecycle_key); - claimed_owners_.erase(lifecycle_key); + if (owner_it->second.empty()) { + claimed_.erase(lifecycle_key); + claimed_units_.erase(lifecycle_key); + claimed_owners_.erase(lifecycle_key); + } return true; } diff --git a/tpu_sync/kv_cache/reshard/request_block_registry.h b/tpu_sync/kv_cache/reshard/request_block_registry.h index e27ed88b6..4ccdb5f95 100644 --- a/tpu_sync/kv_cache/reshard/request_block_registry.h +++ b/tpu_sync/kv_cache/reshard/request_block_registry.h @@ -141,7 +141,11 @@ class RequestBlockRegistry { std::map claimed_ ABSL_GUARDED_BY(mu_); std::map> claimed_units_ ABSL_GUARDED_BY(mu_); - std::map claimed_owners_ ABSL_GUARDED_BY(mu_); + // Every planning attempt holding the claim; a pipelined consumer plans + // one request once per destination stage, all against the same source + // unit set, and the claim outlives the last of them. + std::map> claimed_owners_ + ABSL_GUARDED_BY(mu_); std::map completed_units_ ABSL_GUARDED_BY(mu_); std::map cancelled_ ABSL_GUARDED_BY(mu_); }; diff --git a/tpu_sync/kv_cache/reshard/reshard_coordinator.cc b/tpu_sync/kv_cache/reshard/reshard_coordinator.cc index c4b43a951..2619d5369 100644 --- a/tpu_sync/kv_cache/reshard/reshard_coordinator.cc +++ b/tpu_sync/kv_cache/reshard/reshard_coordinator.cc @@ -19,6 +19,7 @@ #include #include #include +#include #include #include // NOLINT(build/c++11) #include @@ -113,11 +114,43 @@ tpu_sync::rpc::StartTransferRequest BuildStartTransferForTarget( start_req.set_use_block_chunks(true); start_req.set_expected_block_count(plan.expected_block_count); start_req.set_req_id(plan.req_id); + // The plan addresses pools in the destination's index space. A sender + // executes against its own pool table, which may register only a subset + // of the destination's tags, so its request is rewritten into its own + // index space: transferred pools it does not register are dropped (their + // groups keep their position so entry pool_group references hold), and + // the dtype tag list is the sender's own. + const std::map* sender_remap = nullptr; + if (is_sender && !is_receiver) { + auto remap_it = plan.src_pool_indices.find(target); + if (remap_it != plan.src_pool_indices.end()) { + sender_remap = &remap_it->second; + } + } + auto local_pool_index = [sender_remap](int32_t index) -> int32_t { + if (sender_remap == nullptr) return index; + auto it = sender_remap->find(index); + return it == sender_remap->end() ? -1 : it->second; + }; for (int32_t index : plan.transfer_pool_indices) { - start_req.add_transfer_pool_indices(index); + const int32_t local = local_pool_index(index); + if (local >= 0) start_req.add_transfer_pool_indices(local); + } + // The receiver resolves each push against its own pool table, so a + // rewritten sender still names pools on the wire by destination index. + if (sender_remap != nullptr) { + for (const auto& [dst_index, local] : *sender_remap) { + (*start_req.mutable_wire_pool_indices())[local] = dst_index; + } } - for (const std::string& tag : plan.pool_dtype_tags) { - start_req.add_pool_dtype_tags(tag); + if (sender_remap != nullptr) { + for (const std::string& tag : plan.src_pool_dtype_tags.at(target)) { + start_req.add_pool_dtype_tags(tag); + } + } else { + for (const std::string& tag : plan.pool_dtype_tags) { + start_req.add_pool_dtype_tags(tag); + } } start_req.set_parallelism(plan.parallelism); // Python assigns transfer_plan.skip_d2h unconditionally, which marks the @@ -128,7 +161,8 @@ tpu_sync::rpc::StartTransferRequest BuildStartTransferForTarget( for (const PlanPoolGroup& group : plan.pool_groups) { tpu_sync::rpc::PoolGroupProto* group_proto = start_req.add_pool_groups(); for (int32_t index : group.pool_indices) { - group_proto->add_pool_indices(index); + const int32_t local = local_pool_index(index); + if (local >= 0) group_proto->add_pool_indices(local); } for (int64_t block_id : group.dst_device_block_ids) { group_proto->add_dst_device_block_ids(block_id); @@ -178,9 +212,9 @@ tpu_sync::rpc::StartTransferRequest BuildStartTransferForTarget( }; if (is_receiver) { - // Receiver path: every source's schedule, keyed by source ordinal, - // filtered to entries targeting this receiver (single-endpoint pool - // plans always match). + // Receiver path: every source's schedule, keyed by the source's transfer + // rank (the node id its pushes carry), filtered to entries targeting + // this receiver (single-endpoint pool plans always match). const std::string& target_peer = plan.dst_peers.at(target); for (const auto& [src_unit, entries] : plan.schedules) { auto key_it = plan.src_schedule_keys.find(src_unit); diff --git a/tpu_sync/kv_cache/reshard/reshard_service_test.cc b/tpu_sync/kv_cache/reshard/reshard_service_test.cc index 6faeff93e..c9f02093c 100644 --- a/tpu_sync/kv_cache/reshard/reshard_service_test.cc +++ b/tpu_sync/kv_cache/reshard/reshard_service_test.cc @@ -315,6 +315,136 @@ class ReshardStackTest : public ::testing::Test { /*num_blocks=*/16); } + // Pipeline-parallel on both sides: source rank r and destination stage r + // each register only layer r's pool. Stage r's control address is + // 10.0.0.2:<9600 + r>. + static RaidenId PipelineStage(int stage) { + RaidenId unit = DstUnit(); + unit.job_replica_id = absl::StrCat(unit.job_replica_id, "-rank", stage); + return unit; + } + + void RegisterPipelinedUnits(int num_stages) { + for (int rank = 0; rank < num_stages; ++rank) { + tpu_sync::rpc::ControlRequest req; + req.set_command( + tpu_sync::rpc::ControlRequest::COMMAND_REGISTER_WORK_UNIT); + auto* reg = req.mutable_register_work_unit_request(); + *reg->mutable_unit() = RaidenIdToProto(Unit(rank)); + reg->add_shards(absl::StrCat("10.0.0.1:", 9000 + rank)); + reg->set_control_plane_rpc_address( + absl::StrCat("10.0.0.1:", 9100 + rank)); + *reg->add_pools() = MakePool(absl::StrCat("fa.l", rank), 1024, 1024, 16); + reg->set_layout_fingerprint("fp1"); + reg->set_page_tokens(512); + reg->set_transfer_parallelism(num_stages); + reg->set_transfer_rank(rank); + tpu_sync::rpc::ControlResponse resp = Handle(req.SerializeAsString()); + ASSERT_TRUE(resp.success()) << resp.message(); + } + for (int stage = 0; stage < num_stages; ++stage) { + tpu_sync::rpc::ControlRequest req; + req.set_command( + tpu_sync::rpc::ControlRequest::COMMAND_REGISTER_WORK_UNIT); + auto* reg = req.mutable_register_work_unit_request(); + *reg->mutable_unit() = RaidenIdToProto(PipelineStage(stage)); + reg->add_shards(absl::StrCat("10.0.0.2:", 9400 + stage)); + reg->set_control_plane_rpc_address( + absl::StrCat("10.0.0.2:", 9600 + stage)); + *reg->add_pools() = MakePool(absl::StrCat("fa.l", stage), 1024, 1024, 16); + reg->set_layout_fingerprint("fp1"); + reg->set_page_tokens(512); + reg->set_transfer_parallelism(num_stages); + reg->set_transfer_rank(stage); + tpu_sync::rpc::ControlResponse resp = Handle(req.SerializeAsString()); + ASSERT_TRUE(resp.success()) << resp.message(); + } + } + + // Every source rank registers one block of its layer's pool for the + // request. + void RegisterPipelinedSpans(const std::string& req_id, int64_t uuid, + int num_stages) { + for (int rank = 0; rank < num_stages; ++rank) { + tpu_sync::rpc::ControllerRequest req; + req.set_command( + tpu_sync::rpc::ControllerRequest::COMMAND_REGISTER_REQUEST_BLOCKS); + auto* block_req = req.mutable_register_request_blocks_request(); + block_req->set_req_id(req_id); + block_req->set_uuid(uuid); + *block_req->mutable_unit() = RaidenIdToProto(Unit(rank)); + block_req->add_block_ids(3 + rank); + auto* entry = block_req->add_pool_spans(); + entry->set_tag(absl::StrCat("fa.l", rank)); + entry->add_block_ids(3 + rank); + auto* span = entry->add_spans(); + span->set_src_block_ordinal(0); + span->set_src_offset_bytes(0); + span->set_dst_block_index(0); + span->set_dst_offset_bytes(0); + span->set_size_bytes(1024); + span->set_count(1); + entry->set_declared_bytes(1024); + entry->set_dst_space_version(1); + tpu_sync::rpc::ControllerResponse resp = + HandleController(req.SerializeAsString()); + ASSERT_TRUE(resp.success()) << resp.message(); + } + } + + // One destination stage plans the request against every source rank, + // pulling its own layer into block 7. + tpu_sync::rpc::ControllerResponse CoordinateStage(const std::string& req_id, + int64_t uuid, + int num_stages, int stage) { + tpu_sync::rpc::ControllerRequest req; + req.set_command( + tpu_sync::rpc::ControllerRequest::COMMAND_COORDINATE_TRANSFER); + auto* coord = req.mutable_coordinate_transfer_request(); + for (int rank = 0; rank < num_stages; ++rank) { + *coord->add_src_units() = RaidenIdToProto(Unit(rank)); + } + *coord->add_dst_units() = RaidenIdToProto(PipelineStage(stage)); + coord->set_uuid(uuid); + coord->set_is_sender(true); + coord->set_dst_mem_type(tpu_sync::rpc::MEMORY_TYPE_HBM); + coord->set_use_block_chunks(true); + coord->set_req_id(req_id); + coord->add_dst_device_block_ids(7); + coord->add_transfer_pool_tags(absl::StrCat("fa.l", stage)); + coord->add_dst_block_counts(1); + return HandleController(req.SerializeAsString()); + } + + tpu_sync::rpc::ControllerResponse CancelIfUnclaimed(const std::string& req_id, + int64_t uuid) { + tpu_sync::rpc::ControllerRequest req; + req.set_command(tpu_sync::rpc::ControllerRequest:: + COMMAND_CANCEL_REQUEST_BLOCKS_IF_UNCLAIMED); + req.mutable_cancel_request_blocks_if_unclaimed_request()->set_req_id( + req_id); + req.mutable_cancel_request_blocks_if_unclaimed_request()->set_uuid(uuid); + return HandleController(req.SerializeAsString()); + } + + std::vector RequestBlockStatuses(const std::string& req_id, + int64_t uuid) { + tpu_sync::rpc::ControllerRequest req; + req.set_command( + tpu_sync::rpc::ControllerRequest::COMMAND_GET_REQUEST_BLOCK_STATUS); + auto* key = req.mutable_get_request_block_status_request()->add_keys(); + key->set_req_id(req_id); + key->set_uuid(uuid); + tpu_sync::rpc::ControllerResponse resp = + HandleController(req.SerializeAsString()); + EXPECT_TRUE(resp.success()) << resp.message(); + std::vector statuses; + for (int status : resp.get_request_block_status_response().statuses()) { + statuses.push_back(status); + } + return statuses; + } + tpu_sync::rpc::ControlResponse Handle(const std::string& bytes) { tpu_sync::rpc::ControlResponse resp; resp.ParseFromString(service_->HandleFrame(bytes)); @@ -867,6 +997,342 @@ TEST_F(ReshardStackTest, TwoDestinationsArmEachThenDispatchSendersOnce) { } } +TEST_F(ReshardStackTest, SubsetSourceManifestsPairPoolsByTag) { + // Pipeline-parallel source: rank r registers only its own layer's pool + // (tag fa.l), and the destination registers both layers. Each rank + // covers the whole request for its layer; the plan pairs pools by tag, + // hands the receiver destination-canonical pool indices, and rewrites + // every sender's request into that sender's own pool index space. + for (int rank = 0; rank < 2; ++rank) { + tpu_sync::rpc::ControlRequest req; + req.set_command(tpu_sync::rpc::ControlRequest::COMMAND_REGISTER_WORK_UNIT); + auto* reg = req.mutable_register_work_unit_request(); + *reg->mutable_unit() = RaidenIdToProto(Unit(rank)); + reg->add_shards(absl::StrCat("10.0.0.1:", 9000 + rank)); + reg->set_control_plane_rpc_address(absl::StrCat("10.0.0.1:", 9100 + rank)); + *reg->add_pools() = MakePool(absl::StrCat("fa.l", rank), 1024, 1024, 16); + reg->set_layout_fingerprint("fp1"); + reg->set_page_tokens(512); + reg->set_transfer_parallelism(2); + reg->set_transfer_rank(rank); + tpu_sync::rpc::ControlResponse resp = Handle(req.SerializeAsString()); + ASSERT_TRUE(resp.success()) << resp.message(); + } + { + tpu_sync::rpc::ControlRequest req; + req.set_command(tpu_sync::rpc::ControlRequest::COMMAND_REGISTER_WORK_UNIT); + auto* reg = req.mutable_register_work_unit_request(); + *reg->mutable_unit() = RaidenIdToProto(DstUnit(0)); + reg->add_shards("10.0.0.2:9400"); + reg->set_control_plane_rpc_address("10.0.0.2:9600"); + *reg->add_pools() = MakePool("fa.l0", 1024, 1024, 16); + *reg->add_pools() = MakePool("fa.l1", 1024, 1024, 16); + reg->set_layout_fingerprint("fp1"); + reg->set_page_tokens(512); + reg->set_transfer_parallelism(2); + reg->set_transfer_rank(0); + tpu_sync::rpc::ControlResponse resp = Handle(req.SerializeAsString()); + ASSERT_TRUE(resp.success()) << resp.message(); + } + for (int rank = 0; rank < 2; ++rank) { + tpu_sync::rpc::ControllerRequest req; + req.set_command( + tpu_sync::rpc::ControllerRequest::COMMAND_REGISTER_REQUEST_BLOCKS); + auto* block_req = req.mutable_register_request_blocks_request(); + block_req->set_req_id("req-pp"); + block_req->set_uuid(77); + *block_req->mutable_unit() = RaidenIdToProto(Unit(rank)); + block_req->add_block_ids(3 + rank); + auto* entry = block_req->add_pool_spans(); + entry->set_tag(absl::StrCat("fa.l", rank)); + entry->add_block_ids(3 + rank); + auto* span = entry->add_spans(); + span->set_src_block_ordinal(0); + span->set_src_offset_bytes(0); + span->set_dst_block_index(0); + span->set_dst_offset_bytes(0); + span->set_size_bytes(1024); + span->set_count(1); + entry->set_declared_bytes(1024); + entry->set_dst_space_version(1); + tpu_sync::rpc::ControllerResponse resp = + HandleController(req.SerializeAsString()); + ASSERT_TRUE(resp.success()) << resp.message(); + } + + tpu_sync::rpc::ControllerRequest req; + req.set_command( + tpu_sync::rpc::ControllerRequest::COMMAND_COORDINATE_TRANSFER); + auto* coord = req.mutable_coordinate_transfer_request(); + *coord->add_src_units() = RaidenIdToProto(Unit(0)); + *coord->add_src_units() = RaidenIdToProto(Unit(1)); + *coord->add_dst_units() = RaidenIdToProto(DstUnit(0)); + coord->set_uuid(77); + coord->set_is_sender(true); + coord->set_dst_mem_type(tpu_sync::rpc::MEMORY_TYPE_HBM); + coord->set_use_block_chunks(true); + coord->set_req_id("req-pp"); + // Both tags land on the same destination page (one page per layer pool). + coord->add_dst_device_block_ids(7); + coord->add_dst_device_block_ids(7); + coord->add_transfer_pool_tags("fa.l0"); + coord->add_transfer_pool_tags("fa.l1"); + coord->add_dst_block_counts(1); + coord->add_dst_block_counts(1); + tpu_sync::rpc::ControllerResponse resp = + HandleController(req.SerializeAsString()); + ASSERT_TRUE(resp.success()) << resp.message(); + + // One arm, two sender dispatches. + ASSERT_EQ(transport_.calls_.size(), 3u); + tpu_sync::rpc::ControlRequest arm; + ASSERT_TRUE(arm.ParseFromString(transport_.calls_[0].second)); + const auto& arm_req = arm.start_transfer_request(); + EXPECT_FALSE(arm_req.is_sender()); + // Receiver: destination-canonical pools 0 and 1, one group per tag, each + // fed by exactly one sender. + ASSERT_EQ(arm_req.transfer_pool_indices_size(), 2); + EXPECT_EQ(arm_req.transfer_pool_indices(0), 0); + EXPECT_EQ(arm_req.transfer_pool_indices(1), 1); + ASSERT_EQ(arm_req.pool_dtype_tags_size(), 2); + ASSERT_EQ(arm_req.pool_groups_size(), 2); + for (int g = 0; g < 2; ++g) { + ASSERT_EQ(arm_req.pool_groups(g).pool_indices_size(), 1); + EXPECT_EQ(arm_req.pool_groups(g).pool_indices(0), g); + EXPECT_EQ(arm_req.pool_groups(g).expected_pushes(), 1); + } + EXPECT_EQ(arm_req.shard_push_schedules_size(), 2); + EXPECT_EQ(arm_req.wire_pool_indices_size(), 0); + + // Senders: each request is rewritten into the sender's own single-pool + // index space; the group it does not feed is kept (positionally) but + // names no local pool. + for (int i = 1; i <= 2; ++i) { + tpu_sync::rpc::ControlRequest dispatch; + ASSERT_TRUE(dispatch.ParseFromString(transport_.calls_[i].second)); + const auto& send_req = dispatch.start_transfer_request(); + EXPECT_TRUE(send_req.is_sender()); + ASSERT_EQ(send_req.transfer_pool_indices_size(), 1); + EXPECT_EQ(send_req.transfer_pool_indices(0), 0); + ASSERT_EQ(send_req.pool_dtype_tags_size(), 1); + ASSERT_EQ(send_req.pool_groups_size(), 2); + const int rank = transport_.calls_[i].first == "10.0.0.1:9100" ? 0 : 1; + for (int g = 0; g < 2; ++g) { + if (g == rank) { + ASSERT_EQ(send_req.pool_groups(g).pool_indices_size(), 1); + EXPECT_EQ(send_req.pool_groups(g).pool_indices(0), 0); + } else { + EXPECT_EQ(send_req.pool_groups(g).pool_indices_size(), 0); + } + } + ASSERT_EQ(send_req.shard_push_schedules_size(), 1); + const auto& schedule = send_req.shard_push_schedules().at(0); + ASSERT_EQ(schedule.entries_size(), 1); + EXPECT_EQ(schedule.entries(0).pool_group(), rank); + EXPECT_EQ(schedule.entries(0).src_block_id(), 3 + rank); + EXPECT_EQ(schedule.entries(0).dst_block_id(), 7); + // The sender's single local pool is named on the wire by the + // destination index it feeds. + ASSERT_EQ(send_req.wire_pool_indices_size(), 1); + EXPECT_EQ(send_req.wire_pool_indices().at(0), rank); + } +} + +TEST_F(ReshardStackTest, PipelinedDestinationStagesShareOneRequestClaim) { + // Every destination stage plans the same request against the whole source + // rank set; the planner keeps the source with its layer, and the request + // claim is shared by the stages. + RegisterPipelinedUnits(/*num_stages=*/2); + RegisterPipelinedSpans("req-pp2", 78, /*num_stages=*/2); + for (int stage = 0; stage < 2; ++stage) { + tpu_sync::rpc::ControllerResponse resp = + CoordinateStage("req-pp2", 78, /*num_stages=*/2, stage); + ASSERT_TRUE(resp.success()) << "stage " << stage << ": " << resp.message(); + } + + // Per stage: one arm on that stage, one sender on the rank with its layer. + ASSERT_EQ(transport_.calls_.size(), 4u); + EXPECT_EQ(transport_.calls_[0].first, "10.0.0.2:9600"); + EXPECT_EQ(transport_.calls_[1].first, "10.0.0.1:9100"); + EXPECT_EQ(transport_.calls_[2].first, "10.0.0.2:9601"); + EXPECT_EQ(transport_.calls_[3].first, "10.0.0.1:9101"); + // Stage 1's arm carries its single source under that source's transfer + // rank, which is the node id the receiver resolves the pushes by. + tpu_sync::rpc::ControlRequest arm; + ASSERT_TRUE(arm.ParseFromString(transport_.calls_[2].second)); + const auto& arm_req = arm.start_transfer_request(); + ASSERT_EQ(arm_req.shard_push_schedules_size(), 1); + EXPECT_EQ(arm_req.shard_push_schedules().count(1), 1u); +} + +TEST_F(ReshardStackTest, AbandonedStageLeavesTheSiblingClaimInPlace) { + using ProtoStatus = tpu_sync::rpc::GetRequestBlockStatusResponse; + RegisterPipelinedUnits(/*num_stages=*/2); + RegisterPipelinedSpans("req-pp3", 79, /*num_stages=*/2); + + // Stage 1's receiver refuses its arm, so stage 1's planning attempt + // abandons the claim it shares with stage 0. + transport_.FailFor("10.0.0.2:9601"); + tpu_sync::rpc::ControllerResponse first = + CoordinateStage("req-pp3", 79, /*num_stages=*/2, /*stage=*/0); + ASSERT_TRUE(first.success()) << first.message(); + tpu_sync::rpc::ControllerResponse refused = + CoordinateStage("req-pp3", 79, /*num_stages=*/2, /*stage=*/1); + ASSERT_FALSE(refused.success()); + EXPECT_THAT(refused.message(), HasSubstr("injected arm refusal")); + + // Stage 0 still holds the claim: the request cannot be cancelled and reads + // as claimed. + EXPECT_EQ(CancelIfUnclaimed("req-pp3", 79).response_data(), "false"); + EXPECT_EQ(RequestBlockStatuses("req-pp3", 79), + std::vector{ProtoStatus::STATUS_CLAIMED}); + + // Stage 1 rejoins the claim once its receiver accepts the arm. + transport_.FailFor(""); + tpu_sync::rpc::ControllerResponse retry = + CoordinateStage("req-pp3", 79, /*num_stages=*/2, /*stage=*/1); + ASSERT_TRUE(retry.success()) << retry.message(); +} + +TEST_F(ReshardStackTest, AbandoningTheLastStageFreesTheClaim) { + using ProtoStatus = tpu_sync::rpc::GetRequestBlockStatusResponse; + RegisterPipelinedUnits(/*num_stages=*/2); + RegisterPipelinedSpans("req-pp4", 80, /*num_stages=*/2); + + // Both stages' receivers refuse their arms: every planning attempt + // abandons the claim, and the last one leaving frees it. + for (int stage = 0; stage < 2; ++stage) { + transport_.FailFor(absl::StrCat("10.0.0.2:", 9600 + stage)); + tpu_sync::rpc::ControllerResponse resp = + CoordinateStage("req-pp4", 80, /*num_stages=*/2, stage); + ASSERT_FALSE(resp.success()) << "stage " << stage; + EXPECT_THAT(resp.message(), HasSubstr("injected arm refusal")); + } + EXPECT_EQ(RequestBlockStatuses("req-pp4", 80), + std::vector{ProtoStatus::STATUS_REGISTERED}); + + // Nothing holds the claim: the request can be cancelled, and a stage that + // plans afterwards is refused. + EXPECT_EQ(CancelIfUnclaimed("req-pp4", 80).response_data(), "true"); + EXPECT_EQ(RequestBlockStatuses("req-pp4", 80), + std::vector{ProtoStatus::STATUS_CANCELLED}); + transport_.FailFor(""); + tpu_sync::rpc::ControllerResponse late = + CoordinateStage("req-pp4", 80, /*num_stages=*/2, /*stage=*/0); + ASSERT_FALSE(late.success()); + EXPECT_THAT(late.message(), HasSubstr("was cancelled")); +} + +TEST_F(ReshardStackTest, SpansForAnUnregisteredTagAreRefusedAtRegistration) { + // A rank that registers no pool for a tag cannot declare spans for it: + // the registry refuses the declaration before it can reach planning. + RegisterAllUnits(/*num_src=*/2, /*live=*/1024, /*stride=*/1024, + /*num_blocks=*/16); + tpu_sync::rpc::ControllerRequest req; + req.set_command( + tpu_sync::rpc::ControllerRequest::COMMAND_REGISTER_REQUEST_BLOCKS); + auto* block_req = req.mutable_register_request_blocks_request(); + block_req->set_req_id("req-x"); + block_req->set_uuid(43); + *block_req->mutable_unit() = RaidenIdToProto(Unit(1)); + block_req->add_block_ids(5); + auto* entry = block_req->add_pool_spans(); + entry->set_tag("fa.other"); + entry->add_block_ids(5); + auto* span = entry->add_spans(); + span->set_dst_block_index(1); + span->set_size_bytes(1024); + span->set_count(1); + entry->set_declared_bytes(1024); + tpu_sync::rpc::ControllerResponse resp = + HandleController(req.SerializeAsString()); + ASSERT_FALSE(resp.success()); + EXPECT_THAT(resp.message(), HasSubstr("does not match any registered pool")); +} + +TEST_F(ReshardStackTest, SourceTagsAbsentOnTheDestinationAreIgnored) { + // Pipeline-parallel on both sides: prefill stage 0 registers layers 0 + // and 1, decode stage 0 registers layer 0 only. Planning layer 0 pairs + // the shared tag and ignores the source's extra layer. + { + tpu_sync::rpc::ControlRequest req; + req.set_command(tpu_sync::rpc::ControlRequest::COMMAND_REGISTER_WORK_UNIT); + auto* reg = req.mutable_register_work_unit_request(); + *reg->mutable_unit() = RaidenIdToProto(Unit(0)); + reg->add_shards("10.0.0.1:9000"); + reg->set_control_plane_rpc_address("10.0.0.1:9100"); + *reg->add_pools() = MakePool("fa.l0", 1024, 1024, 16); + *reg->add_pools() = MakePool("fa.l1", 1024, 1024, 16); + reg->set_layout_fingerprint("fp1"); + reg->set_page_tokens(512); + reg->set_transfer_parallelism(1); + reg->set_transfer_rank(0); + tpu_sync::rpc::ControlResponse resp = Handle(req.SerializeAsString()); + ASSERT_TRUE(resp.success()) << resp.message(); + } + { + tpu_sync::rpc::ControlRequest req; + req.set_command(tpu_sync::rpc::ControlRequest::COMMAND_REGISTER_WORK_UNIT); + auto* reg = req.mutable_register_work_unit_request(); + *reg->mutable_unit() = RaidenIdToProto(DstUnit()); + reg->add_shards("10.0.0.2:9400"); + reg->set_control_plane_rpc_address("10.0.0.2:9600"); + *reg->add_pools() = MakePool("fa.l0", 1024, 1024, 16); + reg->set_layout_fingerprint("fp1"); + reg->set_page_tokens(512); + reg->set_transfer_parallelism(1); + reg->set_transfer_rank(0); + tpu_sync::rpc::ControlResponse resp = Handle(req.SerializeAsString()); + ASSERT_TRUE(resp.success()) << resp.message(); + } + { + tpu_sync::rpc::ControllerRequest req; + req.set_command( + tpu_sync::rpc::ControllerRequest::COMMAND_REGISTER_REQUEST_BLOCKS); + auto* block_req = req.mutable_register_request_blocks_request(); + block_req->set_req_id("req-pp2pp"); + block_req->set_uuid(78); + *block_req->mutable_unit() = RaidenIdToProto(Unit(0)); + block_req->add_block_ids(3); + for (const char* tag : {"fa.l0", "fa.l1"}) { + auto* entry = block_req->add_pool_spans(); + entry->set_tag(tag); + entry->add_block_ids(3); + auto* span = entry->add_spans(); + span->set_size_bytes(1024); + span->set_count(1); + entry->set_declared_bytes(1024); + entry->set_dst_space_version(1); + } + tpu_sync::rpc::ControllerResponse resp = + HandleController(req.SerializeAsString()); + ASSERT_TRUE(resp.success()) << resp.message(); + } + tpu_sync::rpc::ControllerRequest req; + req.set_command( + tpu_sync::rpc::ControllerRequest::COMMAND_COORDINATE_TRANSFER); + auto* coord = req.mutable_coordinate_transfer_request(); + *coord->add_src_units() = RaidenIdToProto(Unit(0)); + *coord->add_dst_units() = RaidenIdToProto(DstUnit()); + coord->set_uuid(78); + coord->set_is_sender(true); + coord->set_dst_mem_type(tpu_sync::rpc::MEMORY_TYPE_HBM); + coord->set_use_block_chunks(true); + coord->set_req_id("req-pp2pp"); + coord->add_dst_device_block_ids(7); + coord->add_transfer_pool_tags("fa.l0"); + tpu_sync::rpc::ControllerResponse resp = + HandleController(req.SerializeAsString()); + ASSERT_TRUE(resp.success()) << resp.message(); + ASSERT_EQ(transport_.calls_.size(), 2u); + tpu_sync::rpc::ControlRequest dispatch; + ASSERT_TRUE(dispatch.ParseFromString(transport_.calls_[1].second)); + const auto& send_req = dispatch.start_transfer_request(); + ASSERT_EQ(send_req.transfer_pool_indices_size(), 1); + EXPECT_EQ(send_req.transfer_pool_indices(0), 0); + ASSERT_EQ(send_req.pool_dtype_tags_size(), 2); +} + TEST_F(ReshardStackTest, MismatchedDestinationsFailClosed) { RegisterAllUnits(/*num_src=*/1, /*live=*/1024, /*stride=*/1024, /*num_blocks=*/16, /*num_dst=*/2); diff --git a/tpu_sync/proto/transfer_program.proto b/tpu_sync/proto/transfer_program.proto index 2ac317eb5..389393226 100644 --- a/tpu_sync/proto/transfer_program.proto +++ b/tpu_sync/proto/transfer_program.proto @@ -160,6 +160,9 @@ message ReshardBinding { repeated tpu_sync.rpc.RaidenIdProto dst_units = 2; repeated int32 transfer_pool_indices = 3; repeated string pool_dtype_tags = 4; + // Local pool index -> the destination index that names the pool on the + // wire; a pool absent from the map is named by its local index. + map wire_pool_indices = 5; } message TransferProgram { diff --git a/tpu_sync/rpc/raiden_service.proto b/tpu_sync/rpc/raiden_service.proto index f460fa108..40d2325fe 100644 --- a/tpu_sync/rpc/raiden_service.proto +++ b/tpu_sync/rpc/raiden_service.proto @@ -215,6 +215,11 @@ message StartTransferRequest { // destined for this receiver for each layer/tensor. Used to trigger pipelined // per-layer H2D upon completion of all chunks for a layer. map expected_layer_chunk_counts = 21; + // A sender addresses its pools by local index; this maps each local index + // to the destination index that names the pool on the wire. The coordinator + // fills it for every sender from the plan's pool pairing, identity entries + // included; a pool absent from the map is named by its local index. + map wire_pool_indices = 22; } // One tag class's slice of a multi-tag byte-span transfer. diff --git a/tpu_sync/transport/block_transport.cc b/tpu_sync/transport/block_transport.cc index 452bd23ec..b7dc10197 100644 --- a/tpu_sync/transport/block_transport.cc +++ b/tpu_sync/transport/block_transport.cc @@ -756,7 +756,8 @@ void BlockTransport::AsyncPush( const std::vector& src_block_ids, const std::vector& dst_block_ids, int parallelism, MajorOrder major_order, uint64_t uuid, int layer_idx, - std::function>)> raw_on_complete) { + std::function>)> raw_on_complete, + std::optional wire_layer_idx) { auto on_complete = [raw_on_complete](absl::StatusOr> res) { if (!res.ok()) { RecordTransferFailure(res.status(), metric_labels::kDirectionPush); @@ -783,8 +784,9 @@ void BlockTransport::AsyncPush( // In multi-NIC setups, `peers` contains all NIC rail endpoints for the // destination host. Request chunk resolution is identical across NICs, so // we pass `peers[0]` as the destination peer to build requests. - auto requests = BuildBlockRequests(peers[0], src_block_ids, dst_block_ids, - major_order, uuid, layer_idx, P); + auto requests = + BuildBlockRequests(peers[0], src_block_ids, dst_block_ids, major_order, + uuid, layer_idx, P, wire_layer_idx); if (!requests.ok()) { on_complete(requests.status()); return; @@ -898,13 +900,18 @@ lib::Request BlockTransport::BuildBlockRequest( absl::StatusOr> BlockTransport::BuildBlockRequests( absl::string_view peer, const std::vector& src_block_ids, const std::vector& dst_block_ids, MajorOrder major_order, - uint64_t uuid, int layer_idx, int parallelism) { + uint64_t uuid, int layer_idx, int parallelism, + std::optional wire_layer_idx) { const size_t num_blocks = src_block_ids.size(); const uint8_t socket_opcode = static_cast(dst_block_ids.empty() ? 1 : 6); const uint32_t remote_id = static_cast(block_delegate_->node_id()); - const uint32_t local_id = - layer_idx == -1 ? 0xFFFF'FFFF : static_cast(layer_idx); + // Chunks are read from the local block array `layer_idx`; the request + // header names the array the receiver writes into. + const int header_layer_idx = wire_layer_idx.value_or(layer_idx); + const uint32_t local_id = header_layer_idx == -1 + ? 0xFFFF'FFFF + : static_cast(header_layer_idx); if (num_blocks == 0) { return std::vector{BuildBlockRequest( socket_opcode, /*laddr=*/nullptr, /*raddr=*/nullptr, /*len=*/0, diff --git a/tpu_sync/transport/block_transport.h b/tpu_sync/transport/block_transport.h index 776a9068e..691f86dfc 100644 --- a/tpu_sync/transport/block_transport.h +++ b/tpu_sync/transport/block_transport.h @@ -73,13 +73,17 @@ class BlockTransport final { return peregrine_control_.get(); } - // Asynchronous Scatter-Gather Push + // Asynchronous Scatter-Gather Push. `layer_idx` selects the local block + // array; `wire_layer_idx`, when set, is the index the receiver resolves the + // pushed blocks against (a sender whose pool table is a subset of the + // receiver's). void AsyncPush( const std::vector& peers, const std::vector& src_block_ids, const std::vector& dst_block_ids, int parallelism, MajorOrder major_order, uint64_t uuid, int layer_idx, - std::function>)> raw_on_complete); + std::function>)> raw_on_complete, + std::optional wire_layer_idx = std::nullopt); // Synchronous Scatter-Gather Push (op = 1 / op = 6) absl::StatusOr> SyncPush( @@ -152,7 +156,8 @@ class BlockTransport final { absl::StatusOr> BuildBlockRequests( absl::string_view peer, const std::vector& src_block_ids, const std::vector& dst_block_ids, MajorOrder major_order, - uint64_t uuid = 0, int layer_idx = -1, int parallelism = 1); + uint64_t uuid = 0, int layer_idx = -1, int parallelism = 1, + std::optional wire_layer_idx = std::nullopt); // Builds a batch of Requests for block pull transfer. absl::StatusOr> BuildBlockPullRequests( diff --git a/tpu_sync/transport/block_transport_test.cc b/tpu_sync/transport/block_transport_test.cc index bfe5b9574..b7efcc565 100644 --- a/tpu_sync/transport/block_transport_test.cc +++ b/tpu_sync/transport/block_transport_test.cc @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -476,6 +477,72 @@ TEST_P(BlockTransportTest, PushAndPullCorrectness) { EXPECT_EQ(delegate2.data()[size - 1], 0xAB); } +// Resolves chunks inside the block array named by `layer_idx`, so a test can +// tell which of the receiver's arrays a push landed in. +class LayerRoutedDelegate : public MockDelegate { + public: + using MockDelegate::MockDelegate; + + std::vector GetBlockChunks(size_t layer_idx, size_t shard_idx, + absl::Span block_ids, + size_t total_bytes, uint64_t uuid, + int64_t sender_node_id = -1, + absl::string_view peer = "", + int64_t src_block_id = -1, + int64_t dst_block_id = -1) override { + (void)total_bytes; + (void)uuid; + (void)sender_node_id; + (void)peer; + (void)src_block_id; + // A sender names the block by `dst_block_id`; a receiver passes it in + // `block_ids`. + const int64_t block = dst_block_id >= 0 + ? dst_block_id + : (block_ids.empty() ? -1 : block_ids[0]); + if (block < 0 || block >= 4) { + return {}; + } + return {{.ptr = data(layer_idx, shard_idx) + block * 64, .size = 64}}; + } +}; + +TEST_P(BlockTransportTest, PushNamesTheReceiversArrayByWireLayerIndex) { + const size_t size = 1024; + // The sender holds one block array; the receiver holds two. + LayerRoutedDelegate sender_delegate(size, /*max_blocks=*/4, + /*num_layers=*/1); + LayerRoutedDelegate receiver_delegate(size, /*max_blocks=*/4, + /*num_layers=*/2); + std::memset(sender_delegate.data(0), 0xAB, size); + std::memset(receiver_delegate.data(0), 0x00, size); + std::memset(receiver_delegate.data(1), 0x00, size); + + BlockTransport sender(&sender_delegate, 0); + BlockTransport receiver(&receiver_delegate, 0); + BindControlChannels(&sender, &sender_delegate, &receiver, &receiver_delegate); + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + + // Chunks come from the sender's array 0; the wire names the receiver's + // array 1. + std::promise>> promise; + auto future = promise.get_future(); + sender.AsyncPush( + {absl::StrCat("localhost:", receiver.local_port())}, + /*src_block_ids=*/{0}, /*dst_block_ids=*/{0}, /*parallelism=*/1, + MajorOrder::kLayerMajor, /*uuid=*/0, /*layer_idx=*/0, + [&promise](absl::StatusOr> res) { + promise.set_value(std::move(res)); + }, + /*wire_layer_idx=*/1); + auto pushed = future.get(); + ASSERT_TRUE(pushed.ok()) << pushed.status().message(); + + EXPECT_EQ(receiver_delegate.data(1)[0], 0xAB); + EXPECT_EQ(receiver_delegate.data(1)[63], 0xAB); + EXPECT_EQ(receiver_delegate.data(0)[0], 0x00); +} + TEST_P(BlockTransportTest, PullNonContiguous) { size_t size = 1024; // Delegate 1 has 3 blocks capacity