From e93714049e22bfcb1feb95aae9a5fa96e46f0498 Mon Sep 17 00:00:00 2001 From: Sangam Jindal Date: Tue, 28 Jul 2026 08:22:26 -0700 Subject: [PATCH] Implement asynchronous background worker thread and queue in KVCacheManager for non-blocking H2D/D2H dispatches. PiperOrigin-RevId: 955271371 --- tpu_raiden/kv_cache/BUILD | 1 + tpu_raiden/kv_cache/kv_cache_manager_base.cc | 117 ++++++++++++++++++- tpu_raiden/kv_cache/kv_cache_manager_base.h | 53 +++++++++ tpu_raiden/kv_cache/kv_cache_manager_test.cc | 79 +++++++++++++ 4 files changed, 244 insertions(+), 6 deletions(-) diff --git a/tpu_raiden/kv_cache/BUILD b/tpu_raiden/kv_cache/BUILD index 2ccf5f07..e13ea8ac 100644 --- a/tpu_raiden/kv_cache/BUILD +++ b/tpu_raiden/kv_cache/BUILD @@ -227,6 +227,7 @@ cc_library( "//tpu_raiden/transport:block_transport", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/functional:any_invocable", "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", diff --git a/tpu_raiden/kv_cache/kv_cache_manager_base.cc b/tpu_raiden/kv_cache/kv_cache_manager_base.cc index 5c4c874c..c9bbbb20 100644 --- a/tpu_raiden/kv_cache/kv_cache_manager_base.cc +++ b/tpu_raiden/kv_cache/kv_cache_manager_base.cc @@ -54,6 +54,7 @@ #include "tpu_raiden/core/status_macros.h" #include "tpu_raiden/core/tpu_utils.h" #include "tpu_raiden/kv_cache/logical_block_manager.h" +#include "tpu_raiden/kv_cache/pool_layout.h" #include "tpu_raiden/rpc/raiden_service.pb.h" #include "tpu_raiden/transport/block_transport.h" @@ -153,7 +154,7 @@ struct TransferPipelinedState { void SetError(const absl::Status& status) { if (status.ok()) return; { - absl::MutexLock lock(&err_mu); + absl::MutexLock lock(err_mu); if (first_error.ok()) { first_error = status; } @@ -163,7 +164,7 @@ struct TransferPipelinedState { if (promise_fulfilled.compare_exchange_strong(expected, true)) { absl::Status err; { - absl::MutexLock lock(&err_mu); + absl::MutexLock lock(err_mu); err = first_error; } promise.Set(err); @@ -178,7 +179,7 @@ struct TransferPipelinedState { if (has_failed.load(std::memory_order_acquire)) { absl::Status err; { - absl::MutexLock lock(&err_mu); + absl::MutexLock lock(err_mu); err = first_error; } promise.Set(err); @@ -339,6 +340,11 @@ KVCacheManagerBase::KVCacheManagerBase( dma_pool_ = std::make_unique(kPoolSize); push_pool_ = std::make_shared(kPoolSize); pull_pool_ = std::make_unique(kPoolSize); + const char* bg_env = std::getenv("RAIDEN_ENABLE_ASYNC_DISPATCH"); + enable_background_ = (bg_env != nullptr && std::string(bg_env) == "1"); + if (enable_background_) { + worker_thread_ = std::thread(&KVCacheManagerBase::WorkerLoop, this); + } } KVCacheManagerBase::KVCacheManagerBase( @@ -409,9 +415,21 @@ KVCacheManagerBase::KVCacheManagerBase( push_pool_ = std::make_shared(kPoolSize); pull_pool_ = std::make_unique(kPoolSize); InitTransportServer(); + const char* bg_env = std::getenv("RAIDEN_ENABLE_ASYNC_DISPATCH"); + enable_background_ = (bg_env != nullptr && std::string(bg_env) == "1"); + if (enable_background_) { + worker_thread_ = std::thread(&KVCacheManagerBase::WorkerLoop, this); + } } KVCacheManagerBase::~KVCacheManagerBase() { + if (worker_thread_.joinable()) { + { + absl::MutexLock lock(queue_mu_); + shutdown_ = true; + } + worker_thread_.join(); + } push_pool_.reset(); dma_pool_.reset(); pull_pool_.reset(); @@ -420,7 +438,40 @@ KVCacheManagerBase::~KVCacheManagerBase() { host_block_manager_.reset(); } -absl::StatusOr KVCacheManagerBase::H2d( +void KVCacheManagerBase::WorkerLoop() { + while (true) { + AsyncTask task; + { + absl::MutexLock lock(queue_mu_); + queue_mu_.Await( + absl::Condition(this, &KVCacheManagerBase::QueueNotEmptyOrShutdown)); + if (shutdown_ && task_queue_.empty()) { + break; + } + task = std::move(task_queue_.front()); + task_queue_.pop(); + } + auto status_or_future = std::move(task.work)(); + if (!status_or_future.ok()) { + task.promise.Set(status_or_future.status()); + } else { + status_or_future->OnReady( + [promise = std::move(task.promise)](auto status_or_holds) mutable { + if (status_or_holds.ok()) { + promise.Set(); + } else { + promise.Set(status_or_holds.status()); + } + }); + } + } +} + +bool KVCacheManagerBase::QueueNotEmptyOrShutdown() const { + return !task_queue_.empty() || shutdown_; +} + +absl::StatusOr KVCacheManagerBase::H2dImpl( const std::vector& src_offsets_major_dim, const std::vector& dst_offsets_major_dim, const std::vector& copy_sizes_major_dim, @@ -692,7 +743,7 @@ KVCacheManagerBase::DispatchD2hChunks(const std::vector& src_offsets, return logical_futures; } -absl::StatusOr KVCacheManagerBase::D2h( +absl::StatusOr KVCacheManagerBase::D2hImpl( const std::vector& src_offsets_major_dim, const std::vector& dst_offsets_major_dim, const std::vector& copy_sizes_major_dim, @@ -1988,7 +2039,7 @@ absl::Status KVCacheManagerBase::PushKVCacheResharded( } // 2. D2H to copy from device to host. - ASSIGN_OR_RETURN(raiden::PjRtCopyFuture d2h_future, D2h()); + ASSIGN_OR_RETURN(raiden::PjRtCopyFuture d2h_future, D2hImpl()); // 3. Group entries by dst_peer and collect unique block IDs std::map>> peer_transfers; @@ -2376,5 +2427,59 @@ uint8_t* KVCacheManagerBase::GetBlockHostPointer(size_t layer_idx, block_id); } +absl::StatusOr KVCacheManagerBase::H2d( + const std::vector& src_offsets_major_dim, + const std::vector& dst_offsets_major_dim, + const std::vector& copy_sizes_major_dim, + std::optional slot_idx, std::optional target_layer_idx, + std::optional target_shard_idx) { + if (enable_background_) { + auto [promise, future] = xla::MakePromise(); + AsyncTask task; + task.work = [this, src = src_offsets_major_dim, dst = dst_offsets_major_dim, + sizes = copy_sizes_major_dim, slot_idx, target_layer_idx, + target_shard_idx]() mutable { + return this->H2dImpl(src, dst, sizes, slot_idx, target_layer_idx, + target_shard_idx); + }; + task.promise = std::move(promise); + { + absl::MutexLock lock(queue_mu_); + task_queue_.push(std::move(task)); + } + return raiden::PjRtCopyFuture(future, {}); + } + return H2dImpl(src_offsets_major_dim, dst_offsets_major_dim, + copy_sizes_major_dim, slot_idx, target_layer_idx, + target_shard_idx); +} + +absl::StatusOr KVCacheManagerBase::D2h( + const std::vector& src_offsets_major_dim, + const std::vector& dst_offsets_major_dim, + const std::vector& copy_sizes_major_dim, + std::optional slot_idx, std::optional target_layer_idx, + std::optional target_shard_idx) { + if (enable_background_) { + auto [promise, future] = xla::MakePromise(); + AsyncTask task; + task.work = [this, src = src_offsets_major_dim, dst = dst_offsets_major_dim, + sizes = copy_sizes_major_dim, slot_idx, target_layer_idx, + target_shard_idx]() mutable { + return this->D2hImpl(src, dst, sizes, slot_idx, target_layer_idx, + target_shard_idx); + }; + task.promise = std::move(promise); + { + absl::MutexLock lock(queue_mu_); + task_queue_.push(std::move(task)); + } + return raiden::PjRtCopyFuture(future, {}); + } + return D2hImpl(src_offsets_major_dim, dst_offsets_major_dim, + copy_sizes_major_dim, slot_idx, target_layer_idx, + target_shard_idx); +} + } // namespace kv_cache } // namespace tpu_raiden diff --git a/tpu_raiden/kv_cache/kv_cache_manager_base.h b/tpu_raiden/kv_cache/kv_cache_manager_base.h index a403a71a..07e61136 100644 --- a/tpu_raiden/kv_cache/kv_cache_manager_base.h +++ b/tpu_raiden/kv_cache/kv_cache_manager_base.h @@ -19,12 +19,15 @@ #include #include #include +#include #include +#include #include #include #include "absl/base/thread_annotations.h" #include "absl/container/flat_hash_map.h" +#include "absl/functional/any_invocable.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/string_view.h" @@ -357,6 +360,22 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase { int block_id) override; protected: + virtual absl::StatusOr H2dImpl( + const std::vector& src_offsets_major_dim = {}, + const std::vector& dst_offsets_major_dim = {}, + const std::vector& copy_sizes_major_dim = {}, + std::optional slot_idx = std::nullopt, + std::optional layer_idx = std::nullopt, + std::optional shard_idx = std::nullopt); + + virtual absl::StatusOr D2hImpl( + const std::vector& src_offsets_major_dim = {}, + const std::vector& dst_offsets_major_dim = {}, + const std::vector& copy_sizes_major_dim = {}, + std::optional slot_idx = std::nullopt, + std::optional layer_idx = std::nullopt, + std::optional shard_idx = std::nullopt); + const PJRT_Api* c_api_ = nullptr; const PJRT_RawBuffer_Extension* extension_ = nullptr; size_t max_physical_size_ = 0; @@ -462,6 +481,40 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase { }; absl::flat_hash_map active_plans_ ABSL_GUARDED_BY(plans_mu_); + + // An asynchronous FFI task item representing a queued H2D or D2H copy + // request. Bundles the work lambda with the XLA promise that signals Python + // caller completion. + struct AsyncTask { + absl::AnyInvocable() &&> work; + xla::Promise promise; + }; + + // Background worker loop that dequeues tasks from task_queue_ and executes + // them in FIFO order, resolving promises when under-the-hood DMA completes. + void WorkerLoop(); + + // Condition variable predicate helper checking if tasks are available or if + // shutdown has been requested. + bool QueueNotEmptyOrShutdown() const ABSL_SHARED_LOCKS_REQUIRED(queue_mu_); + + // Mutex guarding access to the asynchronous background dispatch queue and + // shutdown state. + absl::Mutex queue_mu_; + + // FIFO task queue for sequential H2D/D2H transfer scheduling. + std::queue task_queue_ ABSL_GUARDED_BY(queue_mu_); + + // Background worker thread executing queued FFI transfers. + std::thread worker_thread_; + + // True if destruction has been initiated and the worker loop should exit + // after draining the queue. + bool shutdown_ ABSL_GUARDED_BY(queue_mu_) = false; + + // True if background FFI dispatching is enabled via the + // RAIDEN_ENABLE_ASYNC_DISPATCH environment variable. + bool enable_background_ = false; }; } // namespace kv_cache diff --git a/tpu_raiden/kv_cache/kv_cache_manager_test.cc b/tpu_raiden/kv_cache/kv_cache_manager_test.cc index e2fd5987..3f8dd4ae 100644 --- a/tpu_raiden/kv_cache/kv_cache_manager_test.cc +++ b/tpu_raiden/kv_cache/kv_cache_manager_test.cc @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -22,8 +23,10 @@ #include #include +#include "absl/base/thread_annotations.h" #include "absl/status/status.h" #include "absl/strings/str_cat.h" +#include "absl/synchronization/mutex.h" #include "tpu_raiden/kv_cache/kv_cache_manager_base.h" #include "tpu_raiden/rpc/raiden_service.pb.h" #include "tpu_raiden/transport/block_transport.h" @@ -1160,6 +1163,82 @@ TEST(KVCacheManagerTest, RemoteTwoStageApisRequireExplicitStaging) { EXPECT_THAT(h2d_write.status().message(), testing::HasSubstr("same length")); } +class TestBackgroundKVCacheManager : public TestKVCacheManager { + public: + TestBackgroundKVCacheManager(size_t num_layers, size_t num_shards, + size_t slice_byte_size, int host_blocks = 0) + : TestKVCacheManager(num_layers, num_shards, slice_byte_size, + host_blocks) {} + + absl::StatusOr H2dImpl( + const std::vector& src_offsets_major_dim, + const std::vector& dst_offsets_major_dim, + const std::vector& copy_sizes_major_dim, + std::optional slot_idx, std::optional layer_idx, + std::optional shard_idx) override { + absl::MutexLock lock(mu_); + execution_order_.push_back("H2dImpl"); + h2d_count_++; + return raiden::PjRtCopyFuture(std::vector{}); + } + + absl::StatusOr D2hImpl( + const std::vector& src_offsets_major_dim, + const std::vector& dst_offsets_major_dim, + const std::vector& copy_sizes_major_dim, + std::optional slot_idx, std::optional layer_idx, + std::optional shard_idx) override { + absl::MutexLock lock(mu_); + execution_order_.push_back("D2hImpl"); + d2h_count_++; + return raiden::PjRtCopyFuture(std::vector{}); + } + + absl::Mutex mu_; + std::vector execution_order_ ABSL_GUARDED_BY(mu_); + int h2d_count_ ABSL_GUARDED_BY(mu_) = 0; + int d2h_count_ ABSL_GUARDED_BY(mu_) = 0; +}; + +TEST(KVCacheManagerTest, BackgroundWorkerThreadExecutesInFifoOrder) { + setenv("RAIDEN_ENABLE_ASYNC_DISPATCH", "1", 1); + TestBackgroundKVCacheManager manager(/*num_layers=*/1, /*num_shards=*/1, + /*slice_byte_size=*/128, + /*host_blocks=*/2); + // Queue H2D, D2H, H2D sequentially + auto f1 = manager.H2d({0}, {0}, {1}); + auto f2 = manager.D2h({0}, {0}, {1}); + auto f3 = manager.H2d({0}, {0}, {1}); + ASSERT_TRUE(f1.ok()); + ASSERT_TRUE(f2.ok()); + ASSERT_TRUE(f3.ok()); + + // Await all futures + EXPECT_TRUE(f1->Await().ok()); + EXPECT_TRUE(f2->Await().ok()); + EXPECT_TRUE(f3->Await().ok()); + + absl::MutexLock lock(manager.mu_); + EXPECT_EQ(manager.h2d_count_, 2); + EXPECT_EQ(manager.d2h_count_, 1); + ASSERT_EQ(manager.execution_order_.size(), 3); + EXPECT_EQ(manager.execution_order_[0], "H2dImpl"); + EXPECT_EQ(manager.execution_order_[1], "D2hImpl"); + EXPECT_EQ(manager.execution_order_[2], "H2dImpl"); + unsetenv("RAIDEN_ENABLE_ASYNC_DISPATCH"); +} + +TEST(KVCacheManagerTest, BackgroundWorkerThreadDisabledByDefault) { + unsetenv("RAIDEN_ENABLE_ASYNC_DISPATCH"); + TestBackgroundKVCacheManager manager(/*num_layers=*/1, /*num_shards=*/1, + /*slice_byte_size=*/128, + /*host_blocks=*/2); + auto f1 = manager.H2d({0}, {0}, {1}); + ASSERT_TRUE(f1.ok()); + absl::MutexLock lock(manager.mu_); + EXPECT_EQ(manager.h2d_count_, 1); +} + } // namespace } // namespace kv_cache } // namespace tpu_raiden