Skip to content

Remote loads record nothing locally & local loads consume pins #697

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

Merged
copybara-service merged 1 commit into from
Aug 17, 2026
Merged

Remote loads record nothing locally & local loads consume pins #697

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
22 changes: 17 additions & 5 deletions tpu_sync/api/jax/kv_cache_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -379,11 +379,23 @@ def load(

`device_block_ids` is the destination and must name one device block per hash.

NOTE: The block_hashes must be pinned in the LRU cache before calling load
when loading from local host. Once the operation is complete (as reported by
poll_load_status), the caller must manually release/unpin them.
Blocks provided in `slices` must be already pinned externally, and remote
loads will re-resolve hashes at the peer, ignoring `slices`.
PIN CONTRACT:
local source -- every hash must be pinned on entry (lookup() is what
normally grants that pin), and a SUCCESSFUL load
consumes exactly one pin per hash. Do not release
afterwards. A FAILED load does not consume it: the entry
stays pinned so you can retry, or release it
deliberately. Giving up is your decision, not the
store's.
remote source -- no pin is required and none is consumed. A hash resolved
only through the registry never entered the local cache,
so there is nothing here to have pinned.

A load from a peer records NOTHING locally: no host copy is kept, so a
later lookup() of that hash is still a miss. Your own block manager is what
remembers you already own the device block.

Remote loads re-resolve hashes at the peer, ignoring the rest of `slices`.

Args:
block_hashes: List of block hashes to load.
Expand Down
12 changes: 4 additions & 8 deletions tpu_sync/api/jax/kv_cache_store_e2e_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -777,14 +777,10 @@ def _run_remote_read_to_hbm_test(self, enable_multi_numa: bool, use_slices: bool
self._await_terminal(
store_b.poll_load_status, len(hashes), "Job B peer-fetch load"
)
# The batch commits as a unit; the entry records where in HBM it landed.
after = store_b.lookup(hashes)

store_b.release(hashes)
self.assertLen(after, 2)
for i, (_, blk) in enumerate(after):
self.assertEqual(blk.status, kv_cache_store.BlockStatus.HBM)
self.assertEqual(blk.device_block_id, dst_device_blocks[i])
# A load from a peer records nothing locally: no host copy was kept, so
# there is no residency to describe. The bytes are in the device blocks
# the caller named and the cache is a miss for these hashes.
self.assertEmpty(store_b.lookup(hashes))
else:
# --- The thing under test: pull straight into HBM. ---------------------
# No insert first: the lookup answer IS the source coordinate, and the
Expand Down
9 changes: 3 additions & 6 deletions tpu_sync/api/jax/kv_cache_store_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -234,8 +234,7 @@ def test_pin_and_release(self):
controller.insert(hash_4, slice_4, True)

res_old = controller.lookup([b"7001", b"7002"])
self.assertLen(res_old, 2)
controller.release([b"7001", b"7002"])
self.assertEmpty(res_old)
res = controller.lookup([b"7002"])
self.assertLen(res, 1)
controller.release([b"7002"])
Expand Down Expand Up @@ -269,8 +268,7 @@ def test_partial_pin_rollback(self):
)

res_old = controller.lookup([b"8001", b"8002"])
self.assertLen(res_old, 2)
controller.release([b"8001", b"8002"])
self.assertEmpty(res_old)
res = controller.lookup([b"8004", b"8005"])
self.assertLen(res, 2)
controller.release([b"8004", b"8005"])
Expand Down Expand Up @@ -473,8 +471,7 @@ def test_insert_and_lock_release_and_delete(self):
success = controller.insert_and_lock(remote_hashes, remote_slices, True)
self.assertTrue(success)
res_local = controller.lookup([b"local_1"])
self.assertLen(res_local, 1)
controller.release([b"local_1"])
self.assertEmpty(res_local)

del_count = controller.release_and_delete(remote_hashes)
self.assertEqual(del_count, 2)
Expand Down
22 changes: 17 additions & 5 deletions tpu_sync/api/torch/kv_cache_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -438,11 +438,23 @@ def load(

`device_block_ids` is the destination and must name one device block per hash.

NOTE: The block_hashes must be pinned in the LRU cache before calling load
when loading from local host. Once the operation is complete (as reported by
poll_load_status), the caller must manually release/unpin them.
Blocks provided in `slices` must be already pinned externally, and remote
loads will re-resolve hashes at the peer, ignoring `slices`.
PIN CONTRACT:
local source -- every hash must be pinned on entry (lookup() is what
normally grants that pin), and a SUCCESSFUL load
consumes exactly one pin per hash. Do not release
afterwards. A FAILED load does not consume it: the entry
stays pinned so you can retry, or release it
deliberately. Giving up is your decision, not the
store's.
remote source -- no pin is required and none is consumed. A hash resolved
only through the registry never entered the local cache,
so there is nothing here to have pinned.

A load from a peer records NOTHING locally: no host copy is kept, so a
later lookup() of that hash is still a miss. Your own block manager is what
remembers you already own the device block.

Remote loads re-resolve hashes at the peer, ignoring the rest of `slices`.

Args:
block_hashes: List of block hashes to load.
Expand Down
9 changes: 3 additions & 6 deletions tpu_sync/api/torch/kv_cache_store_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -175,8 +175,7 @@ def test_pin_and_release(self):
controller.insert(hash_4, slice_4, True)

res_old = controller.lookup([b"7001", b"7002"])
self.assertLen(res_old, 2)
controller.release([b"7001", b"7002"])
self.assertEmpty(res_old)
res = controller.lookup([b"7002"])
self.assertLen(res, 1)
controller.release([b"7002"])
Expand Down Expand Up @@ -210,8 +209,7 @@ def test_partial_pin_rollback(self):
)

res_old = controller.lookup([b"8001", b"8002"])
self.assertLen(res_old, 2)
controller.release([b"8001", b"8002"])
self.assertEmpty(res_old)
res = controller.lookup([b"8004", b"8005"])
self.assertLen(res, 2)
controller.release([b"8004", b"8005"])
Expand Down Expand Up @@ -375,8 +373,7 @@ def test_insert_and_lock_release_and_delete(self):
success = controller.insert_and_lock(remote_hashes, remote_slices, True)
self.assertTrue(success)
res_local = controller.lookup([b"local_1"])
self.assertLen(res_local, 1)
controller.release([b"local_1"])
self.assertEmpty(res_local)

del_count = controller.release_and_delete(remote_hashes)
self.assertEqual(del_count, 2)
Expand Down
50 changes: 41 additions & 9 deletions tpu_sync/kv_cache/kv_cache_store.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1080,6 +1080,7 @@ absl::Status KVCacheStore::Load(absl::Span<const std::string> block_hashes,
}

RaidenId remote_id;
bool from_remote = false;
{
absl::MutexLock lock(mutex_);
auto lookup_or = backend()->Lookup(block_hashes);
Expand All @@ -1093,6 +1094,7 @@ absl::Status KVCacheStore::Load(absl::Span<const std::string> block_hashes,
BlockStatus first_status = slices[0].second.status;
if (first_status == BlockStatus::REMOTE) {
remote_id = slices[0].second.raiden_id;
from_remote = true;
}

for (size_t i = 0; i < slices.size(); ++i) {
Expand Down Expand Up @@ -1147,6 +1149,7 @@ absl::Status KVCacheStore::Load(absl::Span<const std::string> block_hashes,
std::vector<std::string>(block_hashes.begin(), block_hashes.end()),
.device_block_ids =
std::vector<int>(device_block_ids.begin(), device_block_ids.end()),
.from_remote = from_remote,
});
}

Expand All @@ -1169,12 +1172,14 @@ absl::Status KVCacheStore::Load(absl::Span<const std::string> block_hashes,
}

RaidenId remote_id;
bool from_remote = false;
{
absl::MutexLock lock(mutex_);

BlockStatus first_status = slices[0].status;
if (first_status == BlockStatus::REMOTE) {
remote_id = slices[0].raiden_id;
from_remote = true;
}

for (size_t i = 0; i < slices.size(); ++i) {
Expand All @@ -1195,6 +1200,13 @@ absl::Status KVCacheStore::Load(absl::Span<const std::string> block_hashes,
"Mixed remote node IDs in a single Load call");
}
} else {
// The caller's pin is what a successful local load consumes, so it has
// to exist. The no-slices form has always required it; this form did
// not, which left one signature hiding two different pin contracts.
if (backend()->GetPinCount(hash) <= 0) {
return absl::FailedPreconditionError(
absl::StrCat("Block is not pinned: ", hash));
}
if (existing.status != BlockStatus::HOST &&
existing.status != BlockStatus::HOST_AND_HBM) {
return absl::FailedPreconditionError(
Expand Down Expand Up @@ -1226,6 +1238,7 @@ absl::Status KVCacheStore::Load(absl::Span<const std::string> block_hashes,
std::vector<std::string>(block_hashes.begin(), block_hashes.end()),
.device_block_ids =
std::vector<int>(device_block_ids.begin(), device_block_ids.end()),
.from_remote = from_remote,
});
}

Expand Down Expand Up @@ -1967,8 +1980,25 @@ void KVCacheStore::PollLoadsInternal(std::vector<LoadState> ready_loads) {
for (auto& state : ready_loads) {
absl::Status status = state.future.Await();
absl::MutexLock lock(mutex_);
if (status.ok()) {
auto lookup_or = backend()->Lookup(state.block_hashes);
if (status.ok() && state.from_remote) {
// A load from a peer records NOTHING locally. The bytes went to the
// caller's device blocks and no local host copy was kept, so there is no
// residency to describe: an entry here would claim HBM with
// host_block_id -1, which eviction cannot reclaim (it only takes HOST and
// HOST_AND_HBM) and which nothing left in the API can delete.
//
// The consequence is deliberate: a later lookup() of the same hash is a
// miss, and a repeat request re-fetches unless the caller's own block
// manager remembers it already owns the device block.
for (const auto& hash : state.block_hashes) {
done_loads_.push_back(hash);
}
} else if (status.ok()) {
// Local source: the entry exists here by construction, so this lookup is
// purely local -- no registry fallback, which would otherwise put a
// blocking RPC inside the poller while it holds mutex_.
auto lookup_or = backend()->Lookup(state.block_hashes,
LookupOptions{.enable_global = false});
if (lookup_or.ok()) {
const auto& slices = lookup_or.value();
std::vector<std::string> update_hashes;
Expand All @@ -1978,20 +2008,22 @@ void KVCacheStore::PollLoadsInternal(std::vector<LoadState> ready_loads) {
if (i < slices.size()) {
RaidenBlockID block = slices[i].second;
block.device_block_id = state.device_block_ids[i];
if (block.status == BlockStatus::REMOTE) {
block.raiden_id = raiden_id_;
block.host_block_id = -1;
block.status = BlockStatus::HBM;
} else {
block.status = BlockStatus::HOST_AND_HBM;
}
block.status = BlockStatus::HOST_AND_HBM;
update_hashes.push_back(hash);
update_slices.push_back(block);
}
done_loads_.push_back(hash);
}
if (!update_hashes.empty()) {
backend()->Insert(update_hashes, update_slices, /*on_host=*/true);
// The load is done with the block, so the pin the caller acquired to
// keep it alive across the transfer is consumed here. Released AFTER
// the index update, so the entry cannot be evicted between the two.
//
// Only on success, and only for a local source: a failed load stays
// pinned so the caller can retry or release deliberately, and a
// remote load never had a caller pin to consume.
backend()->Release(update_hashes);
}
}
} else {
Expand Down
37 changes: 28 additions & 9 deletions tpu_sync/kv_cache/kv_cache_store.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -274,9 +274,13 @@ class KVCacheStore {
// `device_block_ids` is the destination and must name one device block per
// hash.
//
// NOTE: The block_hashes must be pinned in the LRU cache before calling Load.
// Once the operation is complete (as reported by PollLoadStatus), the caller
// must manually release/unpin them via Release.
// PIN CONTRACT: every hash must be pinned on entry -- Lookup() is what
// normally grants that pin -- and a SUCCESSFUL load consumes exactly one pin
// per hash. The caller does not release afterwards.
//
// A FAILED load does not: the entry stays pinned so the caller can retry, or
// release it deliberately. Deciding to give up is the caller's, not this
// store's.
absl::Status Load(absl::Span<const std::string> block_hashes,
absl::Span<const int> device_block_ids);

Expand All @@ -290,9 +294,16 @@ class KVCacheStore {
// hash.
//
// If `slices` is non-empty, the caller's pre-looked up RaidenBlockIDs are
// used directly. Note that blocks in `slices` must be already pinned
// externally (when Load from local host), and remote loads will re-resolve
// hashes at the peer, ignoring `slices`.
// used directly. Remote loads re-resolve hashes at the peer, ignoring the
// rest of `slices`.
//
// PIN CONTRACT, same as the overload above and now enforced the same way:
// local source -- every hash must be pinned on entry, and a successful
// load consumes one pin per hash.
// remote source -- no pin is required and none is consumed. A hash
// resolved only through the registry never entered the
// local index, so there is nothing here to have pinned,
// and a load from a peer records nothing either.
absl::Status Load(absl::Span<const std::string> block_hashes,
absl::Span<const RaidenBlockID> slices,
absl::Span<const int> device_block_ids);
Expand Down Expand Up @@ -353,10 +364,14 @@ class KVCacheStore {
// Polls the status of all active/inflight Load operations.
// Updates cache metadata upon successful H2D transfers:
// - Loaded from local host DRAM -> HOST_AND_HBM
// - Loaded from a peer -> HBM, with host_block_id -1.
// - Loaded from a peer -> nothing is recorded at all.
//
// Note: HBM-only entries hold a slot in the LRU but own no host block, and
// Evict only reclaims HOST and HOST_AND_HBM entries. They must be explicitly deleted.
// A peer load leaves no entry because there is nothing here to describe: no
// local host copy is kept, so the entry could only say HBM with
// host_block_id -1 -- which Evict cannot reclaim (it takes HOST and
// HOST_AND_HBM only) and which nothing would ever remove. A later lookup()
// of such a hash is therefore a miss, and the caller's own block manager is
// what remembers it already owns the device block.
//
// Returns:
// A tuple of {done_block_hashes, failed_block_hashes, pending_block_hashes}
Expand Down Expand Up @@ -507,6 +522,10 @@ class KVCacheStore {
tsl::Future<> future;
std::vector<std::string> block_hashes;
std::vector<int> device_block_ids;
// Whether the source was a peer. Decided at submit time and carried here
// because the poller cannot re-derive it: a remote load records nothing
// locally, so by completion there is no entry to read a status off.
bool from_remote = false;
};

struct RemoteReadState {
Expand Down
Loading
Loading