diff --git a/tpu_raiden/api/jax/kv_cache_store.py b/tpu_raiden/api/jax/kv_cache_store.py index 7b463a45..51d91620 100644 --- a/tpu_raiden/api/jax/kv_cache_store.py +++ b/tpu_raiden/api/jax/kv_cache_store.py @@ -62,23 +62,27 @@ def __repr__(self) -> str: class KVCacheStore: """Wrapper around compiled C++ KVCacheStore.""" - def __init__(self, capacity: int): - self._impl = _impl.KVCacheStore(capacity=capacity) + def __init__(self, capacity: int, global_registry_address: str = ""): + self._impl = _impl.KVCacheStore( + capacity=capacity, global_registry_address=global_registry_address + ) def lookup( self, block_hashes: list[bytes], + enable_global: bool = False, ) -> list[tuple[bytes, list[RaidenId]]]: """Checks the LRU directory for cached block hashes. Args: block_hashes: Incoming block hashes to check. + enable_global: Whether to fallback to global registry on miss. Returns: A list of tuples containing the block hash and a list of matching RaidenId replicas, halting immediately upon the first cache miss. """ - raw_res = self._impl.lookup(block_hashes) + raw_res = self._impl.lookup(block_hashes, enable_global) final_res = [] for hash_val, raw_slices in raw_res: wrapped_slices = [RaidenId(impl=rs) for rs in raw_slices] diff --git a/tpu_raiden/api/jax/kv_cache_store.pyi b/tpu_raiden/api/jax/kv_cache_store.pyi index 495fbc36..a89602dc 100644 --- a/tpu_raiden/api/jax/kv_cache_store.pyi +++ b/tpu_raiden/api/jax/kv_cache_store.pyi @@ -15,10 +15,11 @@ class RaidenId: ) -> None: ... class KVCacheStore: - def __init__(self, capacity: int) -> None: ... + def __init__(self, capacity: int, global_registry_address: str = ...) -> None: ... def lookup( self, block_hashes: list[bytes], + enable_global: bool = ..., ) -> list[tuple[bytes, list[RaidenId]]]: """Checks the LRU directory for cached block hashes. diff --git a/tpu_raiden/api/jax/kv_cache_store_test.py b/tpu_raiden/api/jax/kv_cache_store_test.py index ef7601d4..9c468d8c 100644 --- a/tpu_raiden/api/jax/kv_cache_store_test.py +++ b/tpu_raiden/api/jax/kv_cache_store_test.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import unittest + from absl.testing import absltest from tpu_raiden.api.jax import kv_cache_store @@ -147,6 +149,73 @@ def test_large_and_arbitrary_length_hashes(self): self.assertEqual(lookup_res[0][0], large_hash) self.assertEqual(lookup_res[1][0], long_hash) + def test_global_lookup_case1_local_hit(self): + # Case 1: Full local hit, no global hit. + # We don't need a registry server for this because it shouldn't be queried. + controller = kv_cache_store.KVCacheStore(capacity=20) + hashes = [b"local_only"] + slices = [ + [kv_cache_store.RaidenId("local_job", "0", "kv_cache", 0)], + ] + self.assertTrue(controller.insert(hashes, slices, True)[0]) + + res = controller.lookup(hashes, enable_global=True) + self.assertLen(res, 1) + self.assertEqual(res[0][0], b"local_only") + self.assertEqual(res[0][1][0].job_name, "local_job") + self.assertEqual(res[0][1][0].data_replica_idx, 0) + + def test_global_lookup_case2_and_3_mocked(self): + # We mock _impl to simulate Case 2 and Case 3 because we don't + # have a running registry server in Python tests. + controller = kv_cache_store.KVCacheStore(capacity=20) + + # Create a mock for the C++ impl + mock_impl = unittest.mock.MagicMock() + controller._impl = mock_impl + + # Case 2: Both local and global have the same hit, but we return local. + local_id = kv_cache_store._impl.RaidenId("local_job", "0", "kv_cache", 1) + mock_impl.lookup.return_value = [(b"shared_hash", [local_id])] + + res = controller.lookup([b"shared_hash"], enable_global=True) + self.assertLen(res, 1) + self.assertEqual(res[0][0], b"shared_hash") + self.assertEqual(res[0][1][0].job_name, "local_job") + self.assertEqual(res[0][1][0].data_replica_idx, 1) + mock_impl.lookup.assert_called_with([b"shared_hash"], True) + + # Case 3: No local hit, only global hits. + remote_id1 = kv_cache_store._impl.RaidenId( + "10.0.0.1:1234", "0", "kv_cache", 42 + ) + remote_id2 = kv_cache_store._impl.RaidenId( + "10.0.0.2:1234", "0", "kv_cache", 43 + ) + mock_impl.lookup.return_value = [ + (b"global_1", [remote_id1]), + (b"global_2", [remote_id2]), + ] + + res = controller.lookup([b"global_1", b"global_2"], enable_global=True) + self.assertLen(res, 2) + self.assertEqual(res[0][0], b"global_1") + self.assertEqual(res[0][1][0].job_name, "10.0.0.1:1234") + self.assertEqual(res[0][1][0].data_replica_idx, 42) + self.assertEqual(res[1][0], b"global_2") + self.assertEqual(res[1][1][0].job_name, "10.0.0.2:1234") + self.assertEqual(res[1][1][0].data_replica_idx, 43) + mock_impl.lookup.assert_called_with([b"global_1", b"global_2"], True) + + def test_global_lookup_error_ignored(self): + controller = kv_cache_store.KVCacheStore( + capacity=20, global_registry_address="invalid.address:12345" + ) + hashes = [b"9001"] + # Should not fail, just return empty because the registry is down + res = controller.lookup(hashes, enable_global=True) + self.assertEmpty(res) + if __name__ == "__main__": absltest.main() diff --git a/tpu_raiden/api/torch/BUILD b/tpu_raiden/api/torch/BUILD index 4d81ed24..f78f734c 100644 --- a/tpu_raiden/api/torch/BUILD +++ b/tpu_raiden/api/torch/BUILD @@ -44,6 +44,17 @@ py_library( visibility = ["//visibility:public"], deps = [ "//tpu_raiden/frameworks/torch:_tpu_raiden_torch", + "@torch_tpu//torch_tpu", + "@torch_tpu//torch_tpu:_loader", + ], +) + +py_test( + name = "kv_cache_store_test", + srcs = ["kv_cache_store_test.py"], + deps = [ + ":kv_cache_store", + "@com_google_absl_py//absl/testing:absltest", ], ) diff --git a/tpu_raiden/api/torch/kv_cache_store.py b/tpu_raiden/api/torch/kv_cache_store.py index 95d31d1f..7ffc5486 100644 --- a/tpu_raiden/api/torch/kv_cache_store.py +++ b/tpu_raiden/api/torch/kv_cache_store.py @@ -12,9 +12,27 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Python wrapper for the compiled C++ KVCacheStore.""" - +import ctypes +import os +import pathlib from typing import Any + +import torch_tpu +from torch_tpu import _loader as _torch_tpu_loader + +Path = pathlib.Path + + +def _load_torch_tpu_common() -> None: + _torch_tpu_loader.load() + common = Path(torch_tpu.__file__).resolve().parent / "common" + lib = common / "libpywrap_torch_tpu_common.so" + if lib.exists(): + ctypes.CDLL(str(lib), mode=os.RTLD_GLOBAL | os.RTLD_NOW) + + +_load_torch_tpu_common() + from tpu_raiden.frameworks.torch import _tpu_raiden_torch as _impl @@ -62,23 +80,27 @@ def __repr__(self) -> str: class KVCacheStore: """Wrapper around compiled C++ KVCacheStore.""" - def __init__(self, capacity: int): - self._impl = _impl.KVCacheStore(capacity=capacity) + def __init__(self, capacity: int, global_registry_address: str = ""): + self._impl = _impl.KVCacheStore( + capacity=capacity, global_registry_address=global_registry_address + ) def lookup( self, block_hashes: list[bytes], + enable_global: bool = False, ) -> list[tuple[bytes, list[RaidenId]]]: """Checks the LRU directory for cached block hashes. Args: block_hashes: Incoming block hashes to check. + enable_global: Whether to fallback to global registry on miss. Returns: A list of tuples containing the block hash and a list of matching RaidenId replicas, halting immediately upon the first cache miss. """ - raw_res = self._impl.lookup(block_hashes) + raw_res = self._impl.lookup(block_hashes, enable_global) final_res = [] for hash_val, raw_slices in raw_res: wrapped_slices = [RaidenId(impl=rs) for rs in raw_slices] diff --git a/tpu_raiden/api/torch/kv_cache_store_test.py b/tpu_raiden/api/torch/kv_cache_store_test.py new file mode 100644 index 00000000..da7e9da2 --- /dev/null +++ b/tpu_raiden/api/torch/kv_cache_store_test.py @@ -0,0 +1,221 @@ +# Copyright 2026 Google LLC. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +from absl.testing import absltest + +from tpu_raiden.api.torch import kv_cache_store + + +class KVCacheStoreTest(absltest.TestCase): + + def test_basic_tests(self): + controller = kv_cache_store.KVCacheStore(capacity=20) + self.assertEqual(controller.capacity(), 20) + + hashes = [b"6001", b"6002"] + slices = [ + [kv_cache_store.RaidenId("inference_server", "0", "kv_cache", 0)], + [kv_cache_store.RaidenId("inference_server", "1", "kv_cache", 0)], + ] + + # 1. Insert + self.assertTrue(controller.insert(hashes, slices, True)[0]) + self.assertFalse( + controller.insert(hashes, slices, True)[0] + ) # Already exists + + # 2. Lookup with a partial miss at the end + hashes_with_miss = [b"6001", b"6002", b"6003"] + lookup_res = controller.lookup(hashes_with_miss) + self.assertLen(lookup_res, 2) + self.assertEqual(lookup_res[0][0], b"6001") + self.assertLen(lookup_res[0][1], 1) + self.assertEqual(lookup_res[0][1][0].job_name, "inference_server") + self.assertEqual(lookup_res[0][1][0].job_replica_id, "0") + + # Lookup with an early miss + hashes_early_miss = [b"6001", b"6003", b"6002"] + lookup_res_early = controller.lookup(hashes_early_miss) + self.assertLen(lookup_res_early, 1) + self.assertEqual(lookup_res_early[0][0], b"6001") + + # 3. Delete + controller.delete(hashes, slices) + self.assertTrue( + controller.insert(hashes, slices, True)[0] + ) # Successful again + + def test_pin_and_release(self): + controller = kv_cache_store.KVCacheStore(capacity=2) + + hashes = [b"7001", b"7002"] + slices = [ + [kv_cache_store.RaidenId("inference_server", "0", "kv_cache", 0)], + [kv_cache_store.RaidenId("inference_server", "1", "kv_cache", 0)], + ] + + self.assertTrue(controller.insert(hashes, slices, True)[0]) + + # Pin both + self.assertTrue(controller.pin(hashes)) + + # Inserting a third element should fail to evict because both items are + # pinned. + hash_3 = [b"7003"] + slice_3 = [ + [kv_cache_store.RaidenId("inference_server", "2", "kv_cache", 0)] + ] + controller.insert(hash_3, slice_3, True) + + # Release 7001 + controller.release([b"7001"]) + + # Now inserting a fourth element (7004) should successfully evict 7001 + hash_4 = [b"7004"] + slice_4 = [ + [kv_cache_store.RaidenId("inference_server", "3", "kv_cache", 0)] + ] + controller.insert(hash_4, slice_4, True) + + self.assertEmpty(controller.lookup([b"7001", b"7002"])) + self.assertLen(controller.lookup([b"7002"]), 1) + + def test_partial_pin_rollback(self): + controller = kv_cache_store.KVCacheStore(capacity=2) + + hashes = [b"8001", b"8002"] + slices = [ + [kv_cache_store.RaidenId("inference_server", "0", "kv_cache", 0)], + [kv_cache_store.RaidenId("inference_server", "1", "kv_cache", 0)], + ] + self.assertTrue(controller.insert(hashes, slices, True)[0]) + + # Attempt to pin a sequence with a missing hash (8003). + self.assertFalse(controller.pin([b"8001", b"8002", b"8003"])) + + # Now inserting two new items (8004, 8005) should successfully evict 8001 + # and 8002 because their pins were completely rolled back! + self.assertTrue( + controller.insert( + [b"8004", b"8005"], + [ + [ + kv_cache_store.RaidenId( + "inference_server", "2", "kv_cache", 0 + ) + ], + [ + kv_cache_store.RaidenId( + "inference_server", "3", "kv_cache", 0 + ) + ], + ], + True, + )[0] + ) + + self.assertEmpty(controller.lookup([b"8001", b"8002"])) + self.assertLen(controller.lookup([b"8004", b"8005"]), 2) + + def test_large_and_arbitrary_length_hashes(self): + controller = kv_cache_store.KVCacheStore(capacity=5) + + # Test both high-bit 8-byte hash and a very long arbitrary length hash + large_hash = b"\xff" * 8 + long_hash = b"a" * 100 + hashes = [large_hash, long_hash] + slices = [ + [kv_cache_store.RaidenId("inference_server", "0", "kv_cache", 0)], + [kv_cache_store.RaidenId("inference_server", "1", "kv_cache", 0)], + ] + + self.assertTrue(controller.insert(hashes, slices, True)[0]) + + lookup_res = controller.lookup(hashes) + self.assertLen(lookup_res, 2) + self.assertEqual(lookup_res[0][0], large_hash) + self.assertEqual(lookup_res[1][0], long_hash) + + def test_global_lookup_case1_local_hit(self): + # Case 1: Full local hit, no global hit. + # We don't need a registry server for this because it shouldn't be queried. + controller = kv_cache_store.KVCacheStore(capacity=20) + hashes = [b"local_only"] + slices = [ + [kv_cache_store.RaidenId("local_job", "0", "kv_cache", 0)], + ] + self.assertTrue(controller.insert(hashes, slices, True)[0]) + + res = controller.lookup(hashes, enable_global=True) + self.assertLen(res, 1) + self.assertEqual(res[0][0], b"local_only") + self.assertEqual(res[0][1][0].job_name, "local_job") + self.assertEqual(res[0][1][0].data_replica_idx, 0) + + def test_global_lookup_case2_and_3_mocked(self): + # We mock _impl to simulate Case 2 and Case 3 because we don't + # have a running registry server in Python tests. + controller = kv_cache_store.KVCacheStore(capacity=20) + + # Create a mock for the C++ impl + mock_impl = unittest.mock.MagicMock() + controller._impl = mock_impl + + # Case 2: Both local and global have the same hit, but we return local. + local_id = kv_cache_store._impl.RaidenId("local_job", "0", "kv_cache", 1) + mock_impl.lookup.return_value = [(b"shared_hash", [local_id])] + + res = controller.lookup([b"shared_hash"], enable_global=True) + self.assertLen(res, 1) + self.assertEqual(res[0][0], b"shared_hash") + self.assertEqual(res[0][1][0].job_name, "local_job") + self.assertEqual(res[0][1][0].data_replica_idx, 1) + mock_impl.lookup.assert_called_with([b"shared_hash"], True) + + # Case 3: No local hit, only global hits. + remote_id1 = kv_cache_store._impl.RaidenId( + "10.0.0.1:1234", "0", "kv_cache", 42 + ) + remote_id2 = kv_cache_store._impl.RaidenId( + "10.0.0.2:1234", "0", "kv_cache", 43 + ) + mock_impl.lookup.return_value = [ + (b"global_1", [remote_id1]), + (b"global_2", [remote_id2]), + ] + + res = controller.lookup([b"global_1", b"global_2"], enable_global=True) + self.assertLen(res, 2) + self.assertEqual(res[0][0], b"global_1") + self.assertEqual(res[0][1][0].job_name, "10.0.0.1:1234") + self.assertEqual(res[0][1][0].data_replica_idx, 42) + self.assertEqual(res[1][0], b"global_2") + self.assertEqual(res[1][1][0].job_name, "10.0.0.2:1234") + self.assertEqual(res[1][1][0].data_replica_idx, 43) + mock_impl.lookup.assert_called_with([b"global_1", b"global_2"], True) + + def test_global_lookup_error_ignored(self): + controller = kv_cache_store.KVCacheStore( + capacity=20, global_registry_address="invalid.address:12345" + ) + hashes = [b"9001"] + # Should not fail, just return empty because the registry is down + res = controller.lookup(hashes, enable_global=True) + self.assertEmpty(res) + + +if __name__ == "__main__": + absltest.main() diff --git a/tpu_raiden/frameworks/jax/kv_cache_store.pyi b/tpu_raiden/frameworks/jax/kv_cache_store.pyi index 074799c5..275897f2 100644 --- a/tpu_raiden/frameworks/jax/kv_cache_store.pyi +++ b/tpu_raiden/frameworks/jax/kv_cache_store.pyi @@ -17,10 +17,12 @@ class KVCacheStore: def __init__( self, capacity: int, + global_registry_address: str = '', ) -> None: ... def lookup( self, block_hashes: list[bytes], + enable_global: bool = False, ) -> list[tuple[bytes, list[RaidenId]]]: """Checks the LRU directory for cached block hashes. Returns a list of all matched replica pairs prior to the first miss.""" ... diff --git a/tpu_raiden/frameworks/jax/tpu_raiden_jax_module.cc b/tpu_raiden/frameworks/jax/tpu_raiden_jax_module.cc index a5738e11..843e84eb 100644 --- a/tpu_raiden/frameworks/jax/tpu_raiden_jax_module.cc +++ b/tpu_raiden/frameworks/jax/tpu_raiden_jax_module.cc @@ -58,8 +58,10 @@ std::vector ToStdStringVector( class KVCacheStoreWrapper { public: - explicit KVCacheStoreWrapper(size_t lru_capacity) { - controller_ = std::make_unique(lru_capacity); + explicit KVCacheStoreWrapper(size_t lru_capacity, + std::string global_registry_address = "") { + controller_ = + std::make_unique(lru_capacity, global_registry_address); } KVCacheStore* operator->() { return controller_.get(); } KVCacheStore& operator*() { return *controller_; } @@ -405,13 +407,14 @@ NB_MODULE(_tpu_raiden_jax, m) { &tpu_raiden::kv_cache::RaidenId::data_replica_idx); nb::class_(m, "KVCacheStore") - .def(nb::init(), nb::arg("capacity")) + .def(nb::init(), nb::arg("capacity"), + nb::arg("global_registry_address") = "") .def( "lookup", [](tpu_raiden::kv_cache::KVCacheStoreWrapper& self, - const std::vector& block_hashes) { + const std::vector& block_hashes, bool enable_global) { auto hashes = ToStdStringVector(block_hashes); - auto res = self->Lookup(hashes); + auto res = self->Lookup(hashes, enable_global); if (!res.ok()) { throw std::runtime_error("KVCacheStore lookup failed: " + std::string(res.status().message())); @@ -427,7 +430,7 @@ NB_MODULE(_tpu_raiden_jax, m) { } return py_res; }, - nb::arg("block_hashes"), + nb::arg("block_hashes"), nb::arg("enable_global") = false, "Checks the LRU directory for cached block hashes. Returns a list of " "all matched replica pairs prior to the first miss.") .def( diff --git a/tpu_raiden/frameworks/torch/tpu_raiden_torch_module.cc b/tpu_raiden/frameworks/torch/tpu_raiden_torch_module.cc index 175208ec..6f7b5dbc 100644 --- a/tpu_raiden/frameworks/torch/tpu_raiden_torch_module.cc +++ b/tpu_raiden/frameworks/torch/tpu_raiden_torch_module.cc @@ -58,8 +58,10 @@ std::vector ToStdStringVector( class KVCacheStoreWrapper { public: - explicit KVCacheStoreWrapper(size_t lru_capacity) { - controller_ = std::make_unique(lru_capacity); + explicit KVCacheStoreWrapper(size_t lru_capacity, + std::string global_registry_address = "") { + controller_ = + std::make_unique(lru_capacity, global_registry_address); } KVCacheStore* operator->() { return controller_.get(); } KVCacheStore& operator*() { return *controller_; } @@ -393,13 +395,14 @@ NB_MODULE(_tpu_raiden_torch, m) { &tpu_raiden::kv_cache::RaidenId::data_replica_idx); nb::class_(m, "KVCacheStore") - .def(nb::init(), nb::arg("capacity")) + .def(nb::init(), nb::arg("capacity"), + nb::arg("global_registry_address") = "") .def( "lookup", [](tpu_raiden::kv_cache::KVCacheStoreWrapper& self, - const std::vector& block_hashes) { + const std::vector& block_hashes, bool enable_global) { auto hashes = ToStdStringVector(block_hashes); - auto res = self->Lookup(hashes); + auto res = self->Lookup(hashes, enable_global); if (!res.ok()) { throw std::runtime_error("KVCacheStore lookup failed: " + std::string(res.status().message())); @@ -415,7 +418,7 @@ NB_MODULE(_tpu_raiden_torch, m) { } return py_res; }, - nb::arg("block_hashes"), + nb::arg("block_hashes"), nb::arg("enable_global") = false, "Checks the LRU directory for cached block hashes. Returns a list of " "all matched replica pairs prior to the first miss.") .def( diff --git a/tpu_raiden/kv_cache/BUILD b/tpu_raiden/kv_cache/BUILD index d66da5e8..000e475b 100644 --- a/tpu_raiden/kv_cache/BUILD +++ b/tpu_raiden/kv_cache/BUILD @@ -173,8 +173,11 @@ cc_library( deps = [ ":lru_cache", "//tpu_raiden/core:raw_transfer_core", + "//tpu_raiden/kv_cache/global_registry:global_registry_client_cc", + "@com_github_grpc_grpc//:grpc++", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/synchronization", @@ -193,6 +196,9 @@ cc_test( deps = [ ":kv_cache_store", "//tpu_raiden/core:raw_transfer_core", + "//tpu_raiden/kv_cache/global_registry:global_registry_client_cc", + "//tpu_raiden/kv_cache/global_registry:global_registry_server_lib", + "@com_github_grpc_grpc//:grpc++", "@com_google_absl//absl/status:statusor", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", diff --git a/tpu_raiden/kv_cache/kv_cache_store.cc b/tpu_raiden/kv_cache/kv_cache_store.cc index a410ffd6..f1cff4eb 100644 --- a/tpu_raiden/kv_cache/kv_cache_store.cc +++ b/tpu_raiden/kv_cache/kv_cache_store.cc @@ -15,35 +15,77 @@ #include "tpu_raiden/kv_cache/kv_cache_store.h" #include +#include #include #include #include #include +#include "grpcpp/grpcpp.h" +#include "absl/log/log.h" #include "absl/status/statusor.h" #include "absl/synchronization/mutex.h" +#include "third_party/grpc/include/grpcpp/security/credentials.h" +#include "tpu_raiden/kv_cache/global_registry/global_registry_client.h" #include "tpu_raiden/kv_cache/lru_cache.h" namespace tpu_raiden { namespace kv_cache { -KVCacheStore::KVCacheStore(size_t capacity) : lru_cache_(capacity) {} +KVCacheStore::KVCacheStore(size_t capacity, std::string global_registry_address) + : lru_cache_(capacity) { + if (!global_registry_address.empty()) { + auto channel = grpc::CreateChannel(global_registry_address, + grpc::InsecureChannelCredentials()); + registry_client_ = + std::make_unique(channel); + } +} KVCacheStore::~KVCacheStore() = default; absl::StatusOr>>> -KVCacheStore::Lookup(const std::vector& block_hashes) { - absl::MutexLock lock(mutex_); - +KVCacheStore::Lookup(const std::vector& block_hashes, + bool enable_global) { std::vector>> results; - results.reserve(block_hashes.size()); - for (const std::string& hash : block_hashes) { - std::vector* existing = lru_cache_.Get(hash); - if (!existing || existing->empty()) { - break; + size_t local_hits = 0; + size_t limit = 0; + { + absl::MutexLock lock(mutex_); + limit = std::min(block_hashes.size(), lru_cache_.available_space()); + results.reserve(limit); + for (size_t i = 0; i < limit; ++i) { + const std::string& hash = block_hashes[i]; + std::vector* existing = lru_cache_.Get(hash); + if (!existing || existing->empty()) { + break; + } + results.push_back(std::make_pair(hash, *existing)); + local_hits++; + } + } + + if (enable_global && local_hits < limit && registry_client_) { + std::vector remaining_hashes(block_hashes.begin() + local_hits, + block_hashes.begin() + limit); + auto global_results_or = registry_client_->Lookup(remaining_hashes); + if (global_results_or.ok()) { + const auto& global_results = global_results_or.value(); + for (size_t i = 0; i < global_results.size(); ++i) { + const auto& metadata = global_results[i]; + RaidenId remote_id; + remote_id.job_name = metadata.host_address(); + remote_id.job_replica_id = "0"; + remote_id.data_name = "kv_cache"; + remote_id.data_replica_idx = metadata.block_id(); + results.push_back(std::make_pair(remaining_hashes[i], + std::vector{remote_id})); + } + } else { + LOG(WARNING) << "Global registry lookup failed: " + << global_results_or.status().message(); } - results.push_back(std::make_pair(hash, *existing)); } return results; diff --git a/tpu_raiden/kv_cache/kv_cache_store.h b/tpu_raiden/kv_cache/kv_cache_store.h index b4051afc..7156eaac 100644 --- a/tpu_raiden/kv_cache/kv_cache_store.h +++ b/tpu_raiden/kv_cache/kv_cache_store.h @@ -33,6 +33,10 @@ namespace tpu_raiden { namespace kv_cache { +namespace global_registry { +class GlobalRegistryClient; +} + // Represents a microservice slice identifier / entity address hosting a replica // of a Key-Value cache block. struct RaidenId { @@ -53,7 +57,8 @@ struct RaidenId { // nodes and microservice slices. class KVCacheStore { public: - explicit KVCacheStore(size_t capacity); + explicit KVCacheStore(size_t capacity, + std::string global_registry_address = ""); ~KVCacheStore(); @@ -65,8 +70,11 @@ class KVCacheStore { // Checks the LRU directory for cached block hashes. Returns a list of all // matched replica pairs (block hash and vector of RaidenIds) encountered // in sequence prior to the first miss. + // If enable_global is true, it will query the global registry for any + // misses after the local lookup. absl::StatusOr>>> - Lookup(const std::vector& block_hashes); + Lookup(const std::vector& block_hashes, + bool enable_global = false); // Caches sharded buffers into host-RAM/HBM backing store. // Returns: @@ -99,6 +107,7 @@ class KVCacheStore { mutable absl::Mutex mutex_; mutable LRUCache> lru_cache_ ABSL_GUARDED_BY(mutex_); + std::unique_ptr registry_client_; }; } // namespace kv_cache diff --git a/tpu_raiden/kv_cache/kv_cache_store_test.cc b/tpu_raiden/kv_cache/kv_cache_store_test.cc index ea91d792..f88f679c 100644 --- a/tpu_raiden/kv_cache/kv_cache_store_test.cc +++ b/tpu_raiden/kv_cache/kv_cache_store_test.cc @@ -14,11 +14,18 @@ #include "tpu_raiden/kv_cache/kv_cache_store.h" +#include #include #include +#include "grpcpp/grpcpp.h" #include #include "absl/status/statusor.h" +#include "third_party/grpc/include/grpcpp/security/server_credentials.h" +#include "third_party/grpc/include/grpcpp/server.h" +#include "third_party/grpc/include/grpcpp/server_builder.h" +#include "tpu_raiden/kv_cache/global_registry/global_registry_client.h" +#include "tpu_raiden/kv_cache/global_registry/global_registry_server.h" namespace tpu_raiden { namespace kv_cache { @@ -146,6 +153,302 @@ TEST(KVCacheStoreTest, EvictionTracking) { EXPECT_EQ(controller.Lookup({"103"})->size(), 1); } +TEST(KVCacheStoreTest, GlobalLookupFallback) { + // 1. Start a local registry server + auto service = std::make_unique(); + grpc::ServerBuilder builder; + int port = 0; + builder.AddListeningPort("localhost:0", grpc::InsecureServerCredentials(), + &port); + builder.RegisterService(service.get()); + auto server = builder.BuildAndStart(); + std::string server_address = "localhost:" + std::to_string(port); + + // 2. Register some blocks in the registry + auto channel = + grpc::CreateChannel(server_address, grpc::InsecureChannelCredentials()); + global_registry::GlobalRegistryClient registry_client(channel); + + std::string hash1 = "global_hash_1"; + std::string host1 = "10.0.0.1:1234"; + int32_t block1 = 42; + + std::string hash2 = "global_hash_2"; + std::string host2 = "10.0.0.2:1234"; + int32_t block2 = 43; + + // For Case 2: Register a hash that will also be present locally, but with + // a different remote address in the registry. + std::string hash_shared = "shared_hash"; + std::string host_shared_remote = "10.0.0.9:1234"; + int32_t block_shared_remote = 99; + + ASSERT_TRUE( + registry_client + .Register({{hash1, host1, block1}, + {hash2, host2, block2}, + {hash_shared, host_shared_remote, block_shared_remote}}) + .ok()); + + // 3. Create KVCacheStore with the registry address + KVCacheStore store(50, server_address); + + // Insert blocks locally + std::vector local_hashes = {"local_only_hash", "shared_hash"}; + std::vector> local_slices = { + {RaidenId{"local_job", "0", "kv_cache", 0}}, + {RaidenId{"local_job", "0", "kv_cache", 1}}}; + ASSERT_TRUE(store.Insert(local_hashes, local_slices, true).first); + + // Case 1: Full local hit, no global hit + { + auto lookup_res = store.Lookup({"local_only_hash"}, /*enable_global=*/true); + ASSERT_TRUE(lookup_res.ok()); + ASSERT_EQ(lookup_res->size(), 1); + EXPECT_EQ((*lookup_res)[0].first, "local_only_hash"); + EXPECT_EQ((*lookup_res)[0].second[0].job_name, "local_job"); + EXPECT_EQ((*lookup_res)[0].second[0].data_replica_idx, 0); + } + + // Case 2: Both local and global has the same hit, but we return local hit + // results + { + auto lookup_res = store.Lookup({"shared_hash"}, /*enable_global=*/true); + ASSERT_TRUE(lookup_res.ok()); + ASSERT_EQ(lookup_res->size(), 1); + EXPECT_EQ((*lookup_res)[0].first, "shared_hash"); + // Should return local info, not remote info from registry + EXPECT_EQ((*lookup_res)[0].second[0].job_name, "local_job"); + EXPECT_EQ((*lookup_res)[0].second[0].data_replica_idx, 1); + } + + // Case 3: No local hit, only global hits + { + auto lookup_res = store.Lookup({"global_hash_1", "global_hash_2"}, + /*enable_global=*/true); + ASSERT_TRUE(lookup_res.ok()); + ASSERT_EQ(lookup_res->size(), 2); + + EXPECT_EQ((*lookup_res)[0].first, "global_hash_1"); + EXPECT_EQ((*lookup_res)[0].second[0].job_name, host1); + EXPECT_EQ((*lookup_res)[0].second[0].data_replica_idx, block1); + + EXPECT_EQ((*lookup_res)[1].first, "global_hash_2"); + EXPECT_EQ((*lookup_res)[1].second[0].job_name, host2); + EXPECT_EQ((*lookup_res)[1].second[0].data_replica_idx, block2); + } + + // 4. Lookup with enable_global = false + // It should stop at the first miss (which is the global hash if we query it) + // If we query {"local_only_hash", "global_hash_1"}, it should return + // local_only_hash and stop. + { + auto lookup_res = store.Lookup({"local_only_hash", "global_hash_1"}, + /*enable_global=*/false); + ASSERT_TRUE(lookup_res.ok()); + EXPECT_EQ(lookup_res->size(), 1); + EXPECT_EQ((*lookup_res)[0].first, "local_only_hash"); + } + + // 5. Lookup with enable_global = true + // It should return both local and global + { + auto lookup_res = + store.Lookup({"local_only_hash", "global_hash_1", "global_hash_2"}, + /*enable_global=*/true); + ASSERT_TRUE(lookup_res.ok()); + ASSERT_EQ(lookup_res->size(), 3); + + EXPECT_EQ((*lookup_res)[0].first, "local_only_hash"); + EXPECT_EQ((*lookup_res)[0].second[0].job_name, "local_job"); + + EXPECT_EQ((*lookup_res)[1].first, "global_hash_1"); + EXPECT_EQ((*lookup_res)[1].second[0].job_name, host1); + EXPECT_EQ((*lookup_res)[1].second[0].data_replica_idx, block1); + + EXPECT_EQ((*lookup_res)[2].first, "global_hash_2"); + EXPECT_EQ((*lookup_res)[2].second[0].job_name, host2); + EXPECT_EQ((*lookup_res)[2].second[0].data_replica_idx, block2); + } + + // 6. Lookup with enable_global = true, but registry has a miss + // It should stop at the first miss in registry + { + auto lookup_res = store.Lookup( + {"local_only_hash", "global_hash_1", "missing_hash", "global_hash_2"}, + /*enable_global=*/true); + ASSERT_TRUE(lookup_res.ok()); + ASSERT_EQ(lookup_res->size(), 2); // local_only_hash, global_hash_1 + EXPECT_EQ((*lookup_res)[0].first, "local_only_hash"); + EXPECT_EQ((*lookup_res)[1].first, "global_hash_1"); + } + + server->Shutdown(); +} + +TEST(KVCacheStoreTest, GlobalLookupRegistryDown) { + // Create KVCacheStore with an unreachable registry address + KVCacheStore store(50, "invalid.address:12345"); + + // Insert one block locally + std::vector local_hashes = {"local_hash"}; + std::vector> local_slices = { + {RaidenId{"local_job", "0", "kv_cache", 0}}}; + ASSERT_TRUE(store.Insert(local_hashes, local_slices, true).first); + + // Lookup with enable_global = true. + // It should NOT fail even though the registry is down. It should return the + // local hit. + auto lookup_res = store.Lookup({"local_hash", "missing_hash"}, + /*enable_global=*/true); + ASSERT_TRUE(lookup_res.ok()); + EXPECT_EQ(lookup_res->size(), 1); + EXPECT_EQ((*lookup_res)[0].first, "local_hash"); + EXPECT_EQ((*lookup_res)[0].second[0].job_name, "local_job"); +} + +TEST(KVCacheStoreTest, LookupCapLimit) { + KVCacheStore store(2); + + std::vector hashes = {"101", "102"}; + std::vector> slices = { + {RaidenId{"inference_server", "0", "kv_cache", 0}}, + {RaidenId{"inference_server", "1", "kv_cache", 1}}}; + + ASSERT_TRUE(store.Insert(hashes, slices, true).first); + + // Lookup 3 hashes, but capacity is 2. It should only return 2. + std::vector lookup_hashes = {"101", "102", "103"}; + auto lookup_res = store.Lookup(lookup_hashes); + ASSERT_TRUE(lookup_res.ok()); + EXPECT_EQ(lookup_res->size(), 2); + EXPECT_EQ((*lookup_res)[0].first, "101"); + EXPECT_EQ((*lookup_res)[1].first, "102"); +} + +TEST(KVCacheStoreTest, LookupCapLimitWithGlobal) { + // 1. Start a local registry server + auto service = std::make_unique(); + grpc::ServerBuilder builder; + int port = 0; + builder.AddListeningPort("localhost:0", grpc::InsecureServerCredentials(), + &port); + builder.RegisterService(service.get()); + auto server = builder.BuildAndStart(); + std::string server_address = "localhost:" + std::to_string(port); + + // 2. Register some blocks in the registry + auto channel = + grpc::CreateChannel(server_address, grpc::InsecureChannelCredentials()); + global_registry::GlobalRegistryClient registry_client(channel); + + std::string hash1 = "global_hash_1"; + std::string host1 = "10.0.0.1:1234"; + int32_t block1 = 42; + + std::string hash2 = "global_hash_2"; + std::string host2 = "10.0.0.2:1234"; + int32_t block2 = 43; + + std::string hash3 = "global_hash_3"; + std::string host3 = "10.0.0.3:1234"; + int32_t block3 = 44; + + ASSERT_TRUE(registry_client + .Register({{hash1, host1, block1}, + {hash2, host2, block2}, + {hash3, host3, block3}}) + .ok()); + + // 3. Create KVCacheStore with capacity 2 + KVCacheStore store(2, server_address); + + // Lookup 3 hashes, but capacity is 2. It should only return 2. + std::vector lookup_hashes = {"global_hash_1", "global_hash_2", + "global_hash_3"}; + auto lookup_res = store.Lookup(lookup_hashes, /*enable_global=*/true); + ASSERT_TRUE(lookup_res.ok()); + EXPECT_EQ(lookup_res->size(), 2); + EXPECT_EQ((*lookup_res)[0].first, "global_hash_1"); + EXPECT_EQ((*lookup_res)[1].first, "global_hash_2"); + + server->Shutdown(); +} + +TEST(KVCacheStoreTest, LookupCapLimitMixed) { + // 1. Start a local registry server + auto service = std::make_unique(); + grpc::ServerBuilder builder; + int port = 0; + builder.AddListeningPort("localhost:0", grpc::InsecureServerCredentials(), + &port); + builder.RegisterService(service.get()); + auto server = builder.BuildAndStart(); + std::string server_address = "localhost:" + std::to_string(port); + + // 2. Register some blocks in the registry + auto channel = + grpc::CreateChannel(server_address, grpc::InsecureChannelCredentials()); + global_registry::GlobalRegistryClient registry_client(channel); + + std::string hash2 = "global_hash_2"; + std::string host2 = "10.0.0.2:1234"; + int32_t block2 = 43; + + std::string hash3 = "global_hash_3"; + std::string host3 = "10.0.0.3:1234"; + int32_t block3 = 44; + + ASSERT_TRUE( + registry_client.Register({{hash2, host2, block2}, {hash3, host3, block3}}) + .ok()); + + // 3. Create KVCacheStore with capacity 2 + KVCacheStore store(2, server_address); + + // Insert 1 block locally + std::vector local_hashes = {"local_hash_1"}; + std::vector> local_slices = { + {RaidenId{"local_job", "0", "kv_cache", 0}}}; + ASSERT_TRUE(store.Insert(local_hashes, local_slices, true).first); + + // Lookup 3 hashes, but capacity is 2. It should only return 2 (1 local, 1 + // global). + std::vector lookup_hashes = {"local_hash_1", "global_hash_2", + "global_hash_3"}; + auto lookup_res = store.Lookup(lookup_hashes, /*enable_global=*/true); + ASSERT_TRUE(lookup_res.ok()); + EXPECT_EQ(lookup_res->size(), 2); + EXPECT_EQ((*lookup_res)[0].first, "local_hash_1"); + EXPECT_EQ((*lookup_res)[1].first, "global_hash_2"); + + server->Shutdown(); +} + +TEST(KVCacheStoreTest, LookupAvailableSpaceLimit) { + KVCacheStore store(3); + + std::vector hashes = {"101", "102", "103"}; + std::vector> slices = { + {RaidenId{"inference_server", "0", "kv_cache", 0}}, + {RaidenId{"inference_server", "1", "kv_cache", 1}}, + {RaidenId{"inference_server", "2", "kv_cache", 2}}}; + + ASSERT_TRUE(store.Insert(hashes, slices, true).first); + + // Pin 101. Pinned count = 1. Available space = 3 - 1 = 2. + EXPECT_TRUE(store.Pin({"101"})); + + // Lookup 4 hashes. Since available space is 2, it should only return the + // first 2. + std::vector lookup_hashes = {"101", "102", "103", "104"}; + auto lookup_res = store.Lookup(lookup_hashes); + ASSERT_TRUE(lookup_res.ok()); + EXPECT_EQ(lookup_res->size(), 2); + EXPECT_EQ((*lookup_res)[0].first, "101"); + EXPECT_EQ((*lookup_res)[1].first, "102"); +} + } // namespace } // namespace kv_cache } // namespace tpu_raiden diff --git a/tpu_raiden/kv_cache/lru_cache.h b/tpu_raiden/kv_cache/lru_cache.h index 35414286..28f82870 100644 --- a/tpu_raiden/kv_cache/lru_cache.h +++ b/tpu_raiden/kv_cache/lru_cache.h @@ -53,6 +53,9 @@ class LRUCache { // Returns the number of elements currently stored in the cache. size_t size() const { return map_.size(); } + // Returns the available space (free + evictable capacity). + size_t available_space() const { return capacity_ - pinned_list_.size(); } + // Returns true if the cache is completely empty. bool empty() const { return map_.empty(); }