Skip to content

Implement asynchronous background worker thread and queue in KVCacheManager for non-blocking H2D/D2H dispatches. #471

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Implement asynchronous background worker thread and queue in KVCacheManager for non-blocking H2D/D2H dispatches. #471

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions tpu_raiden/kv_cache/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
117 changes: 111 additions & 6 deletions tpu_raiden/kv_cache/kv_cache_manager_base.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down Expand Up @@ -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;
}
Expand All @@ -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);
Expand All @@ -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);
Expand Down Expand Up @@ -339,6 +340,11 @@ KVCacheManagerBase::KVCacheManagerBase(
dma_pool_ = std::make_unique<NumaThreadPool>(kPoolSize);
push_pool_ = std::make_shared<NumaThreadPool>(kPoolSize);
pull_pool_ = std::make_unique<NumaThreadPool>(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(
Expand Down Expand Up @@ -409,9 +415,21 @@ KVCacheManagerBase::KVCacheManagerBase(
push_pool_ = std::make_shared<NumaThreadPool>(kPoolSize);
pull_pool_ = std::make_unique<NumaThreadPool>(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();
Expand All @@ -420,7 +438,40 @@ KVCacheManagerBase::~KVCacheManagerBase() {
host_block_manager_.reset();
}

absl::StatusOr<raiden::PjRtCopyFuture> 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<raiden::PjRtCopyFuture> KVCacheManagerBase::H2dImpl(
const std::vector<int64_t>& src_offsets_major_dim,
const std::vector<int64_t>& dst_offsets_major_dim,
const std::vector<int64_t>& copy_sizes_major_dim,
Expand Down Expand Up @@ -692,7 +743,7 @@ KVCacheManagerBase::DispatchD2hChunks(const std::vector<int64_t>& src_offsets,
return logical_futures;
}

absl::StatusOr<raiden::PjRtCopyFuture> KVCacheManagerBase::D2h(
absl::StatusOr<raiden::PjRtCopyFuture> KVCacheManagerBase::D2hImpl(
const std::vector<int64_t>& src_offsets_major_dim,
const std::vector<int64_t>& dst_offsets_major_dim,
const std::vector<int64_t>& copy_sizes_major_dim,
Expand Down Expand Up @@ -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<std::string, std::vector<std::pair<int, int>>> peer_transfers;
Expand Down Expand Up @@ -2376,5 +2427,59 @@ uint8_t* KVCacheManagerBase::GetBlockHostPointer(size_t layer_idx,
block_id);
}

absl::StatusOr<raiden::PjRtCopyFuture> KVCacheManagerBase::H2d(
const std::vector<int64_t>& src_offsets_major_dim,
const std::vector<int64_t>& dst_offsets_major_dim,
const std::vector<int64_t>& copy_sizes_major_dim,
std::optional<int64_t> slot_idx, std::optional<size_t> target_layer_idx,
std::optional<size_t> 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<raiden::PjRtCopyFuture> KVCacheManagerBase::D2h(
const std::vector<int64_t>& src_offsets_major_dim,
const std::vector<int64_t>& dst_offsets_major_dim,
const std::vector<int64_t>& copy_sizes_major_dim,
std::optional<int64_t> slot_idx, std::optional<size_t> target_layer_idx,
std::optional<size_t> 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
53 changes: 53 additions & 0 deletions tpu_raiden/kv_cache/kv_cache_manager_base.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,12 +19,15 @@
#include <cstdint>
#include <memory>
#include <optional>
#include <queue>
#include <string>
#include <thread>
#include <utility>
#include <vector>

#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"
Expand Down Expand Up @@ -357,6 +360,22 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase {
int block_id) override;

protected:
virtual absl::StatusOr<raiden::PjRtCopyFuture> H2dImpl(
const std::vector<int64_t>& src_offsets_major_dim = {},
const std::vector<int64_t>& dst_offsets_major_dim = {},
const std::vector<int64_t>& copy_sizes_major_dim = {},
std::optional<int64_t> slot_idx = std::nullopt,
std::optional<size_t> layer_idx = std::nullopt,
std::optional<size_t> shard_idx = std::nullopt);

virtual absl::StatusOr<raiden::PjRtCopyFuture> D2hImpl(
const std::vector<int64_t>& src_offsets_major_dim = {},
const std::vector<int64_t>& dst_offsets_major_dim = {},
const std::vector<int64_t>& copy_sizes_major_dim = {},
std::optional<int64_t> slot_idx = std::nullopt,
std::optional<size_t> layer_idx = std::nullopt,
std::optional<size_t> shard_idx = std::nullopt);

const PJRT_Api* c_api_ = nullptr;
const PJRT_RawBuffer_Extension* extension_ = nullptr;
size_t max_physical_size_ = 0;
Expand Down Expand Up @@ -462,6 +481,40 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase {
};
absl::flat_hash_map<uint64_t, RegisteredPlan> 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<absl::StatusOr<raiden::PjRtCopyFuture>() &&> work;
xla::Promise<void> 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<AsyncTask> 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
Expand Down
79 changes: 79 additions & 0 deletions tpu_raiden/kv_cache/kv_cache_manager_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,15 +15,18 @@
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <optional>
#include <string>
#include <vector>

#include <gmock/gmock.h>
#include <gtest/gtest.h>
#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"
Expand Down Expand Up @@ -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<raiden::PjRtCopyFuture> H2dImpl(
const std::vector<int64_t>& src_offsets_major_dim,
const std::vector<int64_t>& dst_offsets_major_dim,
const std::vector<int64_t>& copy_sizes_major_dim,
std::optional<int64_t> slot_idx, std::optional<size_t> layer_idx,
std::optional<size_t> shard_idx) override {
absl::MutexLock lock(mu_);
execution_order_.push_back("H2dImpl");
h2d_count_++;
return raiden::PjRtCopyFuture(std::vector<raiden::BufferHolder>{});
}

absl::StatusOr<raiden::PjRtCopyFuture> D2hImpl(
const std::vector<int64_t>& src_offsets_major_dim,
const std::vector<int64_t>& dst_offsets_major_dim,
const std::vector<int64_t>& copy_sizes_major_dim,
std::optional<int64_t> slot_idx, std::optional<size_t> layer_idx,
std::optional<size_t> shard_idx) override {
absl::MutexLock lock(mu_);
execution_order_.push_back("D2hImpl");
d2h_count_++;
return raiden::PjRtCopyFuture(std::vector<raiden::BufferHolder>{});
}

absl::Mutex mu_;
std::vector<std::string> 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
Loading