From 25a82d50579fd26829419b8bcfbf846eea29a24c Mon Sep 17 00:00:00 2001 From: Googler Date: Thu, 13 Aug 2026 07:46:02 -0700 Subject: [PATCH] JAX: pool admission + control-plane listener parity with torch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # Problem Raiden's pool-addressed reshard pipeline — `RegisterWorkUnit` → `RegisterRequestBlocks` → `CoordinateTransfer` → `PoolReshardRegisterRecv` / `PoolReshardPush` — is fully implemented in the C++ core, but only the torch framework surface can reach it. A JAX worker cannot be admitted with a pool manifest, cannot be armed as a receiver, and cannot be fired as a sender. **Who this unblocks.** The pipeline has had no production caller in either framework. One is now proposed in vLLM's TPU backend — [vllm-project/tpu-inference#3379](https://github.com/vllm-project/tpu-inference/pull/3379) routes prefill→decode KV transfer through this reshard path so the two sides can run different TP degrees and page sizes, over [#3378](https://github.com/vllm-project/tpu-inference/pull/3378), which lowers a vLLM TPU KV cache into byte spans. That backend is JAX, so it reaches the manager through exactly the four entry points this PR adds: the `listener_port` constructor argument, `register_pools()`, `transfer_address` and `listener_address`. Without them the connector's controller path cannot construct a manager at all, and it stays gated off behind an environment variable for that reason. Nothing in this PR depends on those two in return — the surface is verified here on its own, by the tests below. # Approach Add the missing JAX-side surface: a way to construct a control-plane listener, bind the pool methods, and call them from Python. Nothing in the transport is reimplemented — `KVCacheManagerWithTransfer` already implements `RegisterActivePlan`, `RegisterRecv`, `PoolReshardPush` and `PoolReshardRegisterRecv`. The change is additive: **no torch file is moved and no torch file is changed.** All seven files are under `tpu_raiden/api/jax/` and `tpu_raiden/frameworks/jax/`, so the torch suite is untouched by construction and the diff carries near-zero rebase risk. # What changed **`frameworks/jax/kv_cache_manager.{h,cc}`** - An optional `listener_port` on all four `KVCacheManager` constructors. When set, a `kv_cache::KVCacheListener` binds to the sole sub-manager. Port 0 binds an ephemeral port; `listener_port()` reports the one actually bound. - `listener_port()` / `is_listener_active()` / `listener_address()` / `transfer_address()`, formatted with the same IPv6 bracketing rule the torch manager uses, so the controller sees identical address strings from either framework. - Thin forwarders for the 13 pool/plan methods, with signatures copied from `KVCacheManagerBase` so `BindPoolApi<>` instantiates unchanged. The JAX facade wraps `NumaAwareKVCacheManager` and doesn't derive from `KVCacheManagerBase`; `BindPoolApi`'s "must derive" comment turned out to be stricter than the code, which is a template that resolves names on the bound class. No refactor of the NumaAware hierarchy was needed. - `PoolTarget()` fails closed: `NumaAwareKVCacheManager::sole_sub_manager()` returns `nullptr` under `ENABLE_MULTI_NUMA`, so a multi-NUMA manager refuses pool operations rather than silently addressing only sub-manager 0. Multi-NUMA fan-out needs `global_shard_to_submanager_` remapping and stays out of scope. - `listener_` is declared after `numa_manager_` so it's destroyed first — it holds a raw pointer into a sub-manager. **`frameworks/jax/tpu_raiden_jax_module.cc`** decodes a `StartTransferRequest`, so it takes a dep on `//tpu_sync/rpc:raiden_service_cc_proto`. Note the asymmetry, which is not a typo: the proto lives under `tpu_sync/rpc/` but still declares `package tpu_raiden.rpc`, so the include path and the C++ namespace disagree by design. It also binds the listener accessors, `register_active_plan` / `unregister_active_plan` / `register_recv`, and applies `BindPoolApi` to the manager class. `pool_layout_nanobind.h` is framework-neutral despite living under `frameworks/torch/`: it includes no torch, and its `cc_library` is `hdrs`-only with public visibility and no torch dependency. Sharing it rather than promoting or copying it is what keeps the JAX and torch pool surfaces from drifting apart. **`api/jax/kv_cache_manager.py`** gains the `listener_port` constructor argument and the `register_pools` / `get_block_ref` / `pool_ids_with_tag` / `num_pools` / `has_explicit_pools` / `pool_spec` / `d2h_pool_blocks` / `h2d_pool_blocks` / `admission_summary` / `register_active_plan` / `unregister_active_plan` / `register_recv` surface, shaped exactly as `api/torch/kv_cache_manager.py`. It imports `api/torch/pool_layout` rather than forking it — pure dataclasses, no torch import, and `tpu_raiden` ships as a namespace package, so the import pulls in nothing else. The pre-existing `is_listener_active` property (the `WorkerService` gRPC port) keeps its meaning for backward compatibility; the new control-plane socket is exposed as `is_control_listener_active`. One compile detail worth knowing if you touch this header: `nanobind::class_` instantiates `detail::wrap_move`, which reaches the inline defaulted move constructor, which needs `unique_ptr`'s deleter — so the listener must be a **complete** type in the JAX header, not forward-declared. The torch manager escapes this because it isn't move-wrapped. # Tests and results **`api/jax/kv_cache_manager_pool_test.py`** (new, 10 tests) — run on a v6e-8 with `--device_type=tpu`: `Ran 10 tests … OK`. Covers pool admission round-trip, descriptor echo, block-ref stride arithmetic, admission rejections (empty table, out-of-range pool index, out-of-range `storage_index`), the ephemeral-port listener, and a byte oracle where pool D2H/H2D mirror host and device at the same offsets — a correct round trip is the identity on the device array, while a wrong base offset, block stride or region extent corrupts it instead of erroring. Runs a half-block-live pool alongside a dense one, since differing block stride and live extent is what the reshard path actually uses. **JAX↔JAX pool push, chip 0 → chip 1.** Three processes over a file rendezvous: a source worker on chip 0, a destination worker on chip 1, and a driver holding an in-process `RaidenController` with a real `WorkerRpcClient` and no controller server — `start_transfer` passes each worker's `control_plane_rpc_address` explicitly, the library-not-server shape a caller embeds. Geometry: 8 blocks × 8 tokens × 256 f32, blocks `[0,1,2,3]` → `[4,5,6,7]`. | run | pools | plan | result | |---|---|---|---| | single pool | 1 | 4 wire entries, 32768 B | 8192/8192 elements exact | | three pools | 3 | same schedule fanned over 3 storages | 24576/24576 elements exact | Both runs assert the negative too: destination blocks outside the plan still hold their sentinel, so a wrong-block write or a stride error can't pass silently. Both sides' `poll_stats()` report `done_sending` / `done_recving`. The same surface was later driven at a real reshard geometry by six single-chip workers — 4 source at page 128 into 2 destinations at page 64, 8 head groups, 2 layers, 256 tokens — landing 524,288 elements exactly, with a negative control that turns the check red. That evidence belongs to the connector PR that uses this surface ([tpu-inference#3379](https://github.com/vllm-project/tpu-inference/pull/3379)), but it's what this surface was built for. **Regression.** `bazel test //tpu_raiden/{core,kv_cache,api}/... //tpu_sync/...` under a hermetic py3.12 with `--define with_torch=false`, run twice with `--nocache_test_results --keep_going` — once on this branch, once on the base commit — so "pre-existing" is a measurement rather than an assertion: | | this branch | base | |---|---|---| | targets | 53 | 52 | | pass | 36 | 36 | | fail locally | 5 | 5 | | fail to build | 12 | 11 | The 36 passing targets are the same 36 on both sides, and the 5 that run and fail are the same 5: `core:host_memory_allocator_test`, `core:raiden_manager_base_test`, `core:raw_transfer_perf_test`, `core:tpu_utils_test`, `tpu_sync/rpc:coordination_helper_test`. The build-error set on this branch is the base's set plus exactly one target — `api/jax:kv_cache_manager_pool_test`, the test this PR adds, which lands in the same bucket as its pre-existing `api/jax` siblings `weight_synchronizer_test` and `kv_cache_store_recovery_e2e_test`. So this branch introduces no failure of any kind, and the one target it adds does not build on this host. The failures fall into four groups, none of them ours: the `torch_tpu` shim packages are absent because the host has no torch and the build is pointed at a stub module, BUILD rules reference `.cc` sources absent from the OSS tree, some external deps (`@@protobuf+//io`, `pyglib`) resolve to no package, and a set of C++ targets fail `CppCompile` under this host's gcc. Green and relevant: `kv_cache:pool_layout_test`, `api/torch:kv_cache_manager_host_test`, `core:kv_cache_manager_with_transfer_pool_reshard_test`, `core:kv_manager_holder_test`, `kv_cache:kv_cache_listener_test`, and both `raiden_controller_test`s (`core/controller` and `tpu_sync/rpc`). # Notes and follow-ups - **The torch extension targets could not be built on the test host at all** — no `torch_tpu` checkout and no torch, so `build.sh` falls back to a JAX-only build. "No torch regression" rests on the zero-torch-file diff plus the shared-code targets above, not on a torch suite run. Please run the torch suite where it builds. - **The new bazel target `//tpu_raiden/api/jax:kv_cache_manager_pool_test` also fails to build on that host**, for the same jaxlib/gcc reason as its pre-existing sibling `weight_synchronizer_test`, which fails identically at base. The 10 tests are verified by direct execution against the built `.so` on a pinned runtime, not under bazel — please re-run under bazel wherever the hermetic jaxlib build works. - **Multi-NUMA is out of scope and refused, not supported.** Under `ENABLE_MULTI_NUMA` the pool methods throw. Lifting that needs pool operations to fan out across sub-managers with `global_shard_to_submanager_` remapping. - The listener is bound to one sub-manager, matching the torch manager's own single-listener shape. - **Runtime note:** the JAX extension only loads against the XLA it was compiled against. On a runtime carrying a newer `jax`/`libtpu` pair, every device-attached manager construction fails with `Failed to acquire buffer handle: RawBuffer extension missing` — a `PJRT_Extension_Type` enum skew, not a missing feature, and it reproduces on pre-existing JAX tests at base. Nothing here changes that; noted so the failure is recognizable if you hit it. PiperOrigin-RevId: 964080221 --- tpu_raiden/api/jax/BUILD | 21 ++ tpu_raiden/api/jax/kv_cache_manager.py | 177 +++++++++++- .../api/jax/kv_cache_manager_pool_test.py | 260 ++++++++++++++++++ tpu_raiden/frameworks/jax/BUILD | 19 +- tpu_raiden/frameworks/jax/kv_cache_manager.cc | 73 ++++- tpu_raiden/frameworks/jax/kv_cache_manager.h | 124 ++++++++- .../frameworks/jax/tpu_raiden_jax_module.cc | 72 ++++- 7 files changed, 727 insertions(+), 19 deletions(-) create mode 100644 tpu_raiden/api/jax/kv_cache_manager_pool_test.py diff --git a/tpu_raiden/api/jax/BUILD b/tpu_raiden/api/jax/BUILD index f9f3b50a..c09e79c6 100644 --- a/tpu_raiden/api/jax/BUILD +++ b/tpu_raiden/api/jax/BUILD @@ -21,11 +21,32 @@ py_library( srcs = ["kv_cache_manager.py"], visibility = ["//visibility:public"], deps = [ + # Framework-neutral despite living under api/torch: pure dataclasses, + # no torch import. Shared rather than duplicated so the JAX and torch + # pool descriptors cannot drift. + "//tpu_raiden/api/torch:pool_layout", "//tpu_raiden/frameworks/jax:_tpu_raiden_jax", "//tpu_raiden/frameworks/jax:jax_test_utils", ], ) +py_test( + name = "kv_cache_manager_pool_test", + srcs = ["kv_cache_manager_pool_test.py"], + args = ["--device_type=cpu"], + tags = [ + "cpu:2", + ], + deps = [ + ":kv_cache_manager_jax_py", + "//tpu_raiden/api/torch:pool_layout", + "@com_google_absl_py//absl/flags", + "@com_google_absl_py//absl/testing:absltest", + "@jax//jax", + "@pypi//numpy", + ], +) + py_library( name = "kv_cache_store", srcs = ["kv_cache_store.py"], diff --git a/tpu_raiden/api/jax/kv_cache_manager.py b/tpu_raiden/api/jax/kv_cache_manager.py index 5a3b9e01..fa2cec82 100644 --- a/tpu_raiden/api/jax/kv_cache_manager.py +++ b/tpu_raiden/api/jax/kv_cache_manager.py @@ -14,9 +14,17 @@ """High-performance JAX KV Cache Manager (repurposed as TransferEngine).""" -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Sequence, Tuple, Union +# pool_layout is framework-neutral (pure dataclasses, no torch import, and +# tpu_raiden ships as a namespace package so importing it pulls nothing else +# from the torch API). It is shared rather than duplicated so the JAX and +# torch pool descriptors cannot drift. +from tpu_raiden.api.torch import pool_layout from tpu_raiden.frameworks.jax import _tpu_raiden_jax as _impl +PoolSpec = pool_layout.PoolSpec +RegionSpec = pool_layout.RegionSpec + class KVCacheManager: """Wrapper around compiled C++ TransferEngine. @@ -39,6 +47,7 @@ def __init__( raiden_worker_port: int = 0, raiden_controller_address: Optional[str] = None, worker_id: Optional[str] = None, + listener_port: Optional[int] = None, ): """Instantiates the TransferEngine-based KVCacheManager. @@ -58,7 +67,14 @@ def __init__( raiden_controller_address: Optional address of central RaidenController. If provided, the WorkerService gRPC server is enabled. worker_id: Optional identifier for this worker. + listener_port: Optional port for the control-plane KVCacheListener. A + RaidenController talks to this socket to arm a receiver + (PoolReshardRegisterRecv) and to fire a sender (PoolReshardPush) for a + pool-addressed reshard plan. Distinct from raiden_worker_port, which + is the buffer-oriented WorkerService gRPC server. Requires a manager + with exactly one NUMA sub-manager. """ + self._admission_summary: Optional[Dict[str, Any]] = None if host_blocks_to_allocate is not None: self._impl = _impl.KVCacheManager( kv_caches, @@ -72,6 +88,7 @@ def __init__( # Pass node_id so ReadRemote can always match src<->dst workers by # node_id (see the non-host_blocks branch, which already forwards it). node_id=node_id, + listener_port=listener_port, ) else: if max_blocks is None or num_slots is None: @@ -91,6 +108,7 @@ def __init__( raiden_worker_port=raiden_worker_port, raiden_controller_address=raiden_controller_address, worker_id=worker_id, + listener_port=listener_port, ) def get_raiden_worker_port(self) -> int: @@ -235,3 +253,160 @@ def dump_metrics_to_string(self) -> str: A JSON string representing the collected telemetry metrics. """ return self._impl.dump_metrics_to_string() + + # ========================================================================= + # POOL-ADDRESSED RESHARD APIs + # Mirrors api/torch/kv_cache_manager.py so a caller (e.g. the TPU vLLM + # connector) can drive either framework through the same names and shapes. + # ========================================================================= + + @property + def listener_port(self) -> Optional[int]: + """Returns the control-plane KVCacheListener port, or None if disabled.""" + return self._impl.listener_port + + @property + def is_control_listener_active(self) -> bool: + """Returns True if the control-plane KVCacheListener is running. + + Distinct from ``is_listener_active``, which reports on the buffer-oriented + WorkerService gRPC server. The two are separate sockets and either can be + up without the other. + """ + return self._impl.is_listener_active + + @property + def listener_address(self) -> str: + """Returns the formatted control listener endpoint string (host:port).""" + return self._impl.listener_address + + @property + def transfer_address(self) -> str: + """Returns the formatted data transfer endpoint string (host:port).""" + return self._impl.transfer_address + + def register_pools(self, pools: Sequence[Any]) -> Dict[str, Any]: + """Registers explicit block pools over the wrapped storages. + + Args: + pools: Sequence of ``pool_layout.PoolSpec`` (or equivalent mappings) in + the caller's canonical order. Pool indices travel on the wire, so both + transfer peers must agree on this order. + + Returns: + A generic admission summary (also served by ``admission_summary``). + """ + coerced = tuple(pool_layout.coerce_pool_spec(pool) for pool in pools) + if not coerced: + raise ValueError("pool table must be non-empty") + num_storages = int(self._impl.num_layers) + for pool_idx, pool in enumerate(coerced): + try: + pool.validate() + except Exception as exc: + raise ValueError(f"invalid pool {pool_idx}: {exc}") from exc + if pool.storage_index >= num_storages: + raise ValueError( + f"pool {pool_idx} ({pool.tag}) storage_index " + f"{pool.storage_index} out of range: manager wraps " + f"{num_storages} storages" + ) + self._impl.register_pools_native( + [pool.to_native_tuple() for pool in coerced] + ) + tags: Dict[str, int] = {} + for pool in coerced: + tags[pool.tag] = tags.get(pool.tag, 0) + 1 + storages = len({pool.storage_index for pool in coerced}) + summary = { + "admitted": True, + "pools": len(coerced), + "storages": storages, + "tags": tags, + } + self._admission_summary = dict(summary) + return dict(summary) + + def get_block_ref( + self, pool_idx: int, block_id: int, shard_idx: int = 0 + ) -> Dict[str, Any]: + """Returns a reference descriptor for one host-side pool block.""" + return dict( + self._impl.get_pool_block_ref_native( + pool_idx=pool_idx, shard_idx=shard_idx, block_id=block_id + ) + ) + + def pool_ids_with_tag(self, tag: str) -> List[int]: + """Returns the pool indices registered with the given opaque tag.""" + return [ + int(pool_idx) + for pool_idx in self._impl.pool_indices_with_tag_native(tag) + ] + + def num_pools(self) -> int: + """Returns the number of pools (implicit or explicit).""" + return int(self._impl.num_pools()) + + def has_explicit_pools(self) -> bool: + """Returns True if an explicit pool table was admitted.""" + return bool(self._impl.has_explicit_pools()) + + def pool_spec(self, pool_idx: int) -> Dict[str, Any]: + """Returns one pool's descriptor as a dict.""" + return dict(self._impl.pool_spec_native(pool_idx)) + + def d2h_pool_blocks( + self, + pool_idx: int, + block_ids: Sequence[int], + shard_idx: Optional[int] = None, + ) -> Any: + """Partial D2H of whole pool blocks into the host mirror.""" + return self._impl.d2h_pool_blocks(pool_idx, list(block_ids), shard_idx) + + def h2d_pool_blocks( + self, + pool_idx: int, + block_ids: Sequence[int], + shard_idx: Optional[int] = None, + ) -> Any: + """Partial H2D of whole pool blocks from the host mirror.""" + return self._impl.h2d_pool_blocks(pool_idx, list(block_ids), shard_idx) + + def admission_summary(self) -> Dict[str, Any]: + """Returns the last successful pool admission summary.""" + if self._admission_summary is None: + return {"admitted": False} + return dict(self._admission_summary) + + def register_active_plan( + self, uuid: int, request: Union[bytes, Any], is_sender: bool + ) -> None: + """Registers a serialized StartTransferRequest for strided push.""" + if hasattr(request, "SerializeToString"): + request = request.SerializeToString() + if not isinstance(request, (bytes, bytearray)): + raise TypeError("request must be bytes or a protobuf message") + self._impl.register_active_plan(uuid, bytes(request), is_sender) + + def unregister_active_plan(self, uuid: int) -> None: + """Removes a previously registered strided push plan.""" + self._impl.unregister_active_plan(uuid) + + def register_recv( + self, uuid: int, req_id: str, expected_block_count: int + ) -> None: + """[EXPERIMENTAL] Registers expected incoming blocks for push resharding. + + Allocates staging slots in the C++ receiver engine and sets the barrier for + the expected physical block-pushes; the engine triggers H2D into TPU HBM + once the count is reached. + + Args: + uuid: Unique identifier for the transfer transaction. + req_id: Request ID associated with the transfer. + expected_block_count: Total number of physical block-pushes expected from + all contributing source ranks. + """ + self._impl.register_recv(uuid, req_id, expected_block_count) diff --git a/tpu_raiden/api/jax/kv_cache_manager_pool_test.py b/tpu_raiden/api/jax/kv_cache_manager_pool_test.py new file mode 100644 index 00000000..1b61f1a1 --- /dev/null +++ b/tpu_raiden/api/jax/kv_cache_manager_pool_test.py @@ -0,0 +1,260 @@ +# 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. + +"""Pool-admission and control-plane parity tests for the JAX manager. + +These mirror api/torch/kv_cache_manager_host_test.py: the JAX manager must +expose the same pool surface, with the same shapes and the same errors, so a +caller (e.g. the TPU vLLM connector) can drive either framework identically. +""" + +from absl import flags +from absl.testing import absltest +import jax +import jax.numpy as jnp +import numpy as np + +from tpu_raiden.api.jax import kv_cache_manager as jax_kv +from tpu_raiden.api.torch import pool_layout + +flags.DEFINE_string( + "device_type", + "tpu", + "JAX backend to run against (e.g. 'tpu', 'cpu').", +) + +_NUM_BLOCKS = 4 +_BLOCK_SHAPE = (8, 4, 128) +_DTYPE = jnp.float32 + + +def _block_bytes() -> int: + return int(np.prod(_BLOCK_SHAPE)) * jnp.dtype(_DTYPE).itemsize + + +def _dense_pool(tag: str, storage_index: int) -> pool_layout.PoolSpec: + """One pool spanning every byte of every block of one storage.""" + block_bytes = _block_bytes() + return pool_layout.PoolSpec( + tag=tag, + storage_index=storage_index, + base_offset_bytes=0, + block_stride_bytes=block_bytes, + num_blocks=_NUM_BLOCKS, + regions=( + pool_layout.RegionSpec( + name="payload", + offset_bytes=0, + stride_bytes=block_bytes, + unit_bytes=block_bytes, + num_units=1, + ), + ), + dtype_tag="float32", + ) + + +def _half_block_pool(tag: str, storage_index: int) -> pool_layout.PoolSpec: + """A pool declaring only the first half of each block as live. + + Exercises the case the reshard path actually uses: block stride and live + extent differ, so a wrong offset shows up as corruption in the untouched + half rather than as an error. + """ + block_bytes = _block_bytes() + return pool_layout.PoolSpec( + tag=tag, + storage_index=storage_index, + base_offset_bytes=0, + block_stride_bytes=block_bytes, + num_blocks=_NUM_BLOCKS, + regions=( + pool_layout.RegionSpec( + name="payload", + offset_bytes=0, + stride_bytes=block_bytes, + unit_bytes=block_bytes // 2, + num_units=1, + ), + ), + dtype_tag="float32", + ) + + +class _ManagerTestBase(absltest.TestCase): + """Shared fixture: one single-device, single-shard JAX manager.""" + + def setUp(self): + super().setUp() + device_type = flags.FLAGS.device_type + try: + self.devices = jax.devices(device_type) + except RuntimeError as exc: + raise AssertionError(f"No {device_type} devices found") from exc + if not self.devices: + raise AssertionError(f"No {device_type} devices found") + + def _make_manager(self, num_layers=2, listener_port=None): + """Builds a single-device, single-shard manager over `num_layers` arrays.""" + key = jax.random.key(17) + refs = [] + arrays = [] + for layer in range(num_layers): + base = jax.random.uniform( + jax.random.fold_in(key, layer), + (_NUM_BLOCKS,) + _BLOCK_SHAPE, + dtype=_DTYPE, + ) + refs.append(np.asarray(base)) + arrays.append(jax.device_put(base, self.devices[0])) + jax.block_until_ready(arrays) + manager = jax_kv.KVCacheManager( + kv_caches=arrays, + local_control_port=0, + host_blocks_to_allocate=_NUM_BLOCKS, + unsafe_skip_buffer_lock=True, + listener_port=listener_port, + ) + return manager, arrays, refs + + +class JaxPoolApiTest(_ManagerTestBase): + + def test_register_pools_round_trip(self): + manager, _, _ = self._make_manager(num_layers=2) + pools = (_dense_pool("fa", 0), _dense_pool("fa", 1)) + + summary = manager.register_pools(pools) + + self.assertEqual( + summary, + {"admitted": True, "pools": 2, "storages": 2, "tags": {"fa": 2}}, + ) + self.assertEqual(manager.admission_summary(), summary) + self.assertTrue(manager.has_explicit_pools()) + self.assertEqual(manager.num_pools(), 2) + self.assertEqual(manager.pool_ids_with_tag("fa"), [0, 1]) + self.assertEmpty(manager.pool_ids_with_tag("state")) + + def test_pool_spec_echoes_the_admitted_descriptor(self): + manager, _, _ = self._make_manager(num_layers=1) + pool = _half_block_pool("fa", 0) + manager.register_pools((pool,)) + + spec = manager.pool_spec(0) + + self.assertEqual(spec["tag"], "fa") + self.assertEqual(spec["storage_index"], 0) + self.assertEqual(spec["base_offset_bytes"], 0) + self.assertEqual(spec["block_stride_bytes"], _block_bytes()) + self.assertEqual(spec["num_blocks"], _NUM_BLOCKS) + self.assertEqual(spec["dtype_tag"], "float32") + self.assertLen(spec["regions"], 1) + region = spec["regions"][0] + self.assertEqual(region["name"], "payload") + self.assertEqual(region["unit_bytes"], _block_bytes() // 2) + self.assertEqual(region["num_units"], 1) + self.assertEqual(region["units_per_stride"], 1) + + def test_block_refs_advance_by_the_declared_block_stride(self): + manager, _, _ = self._make_manager(num_layers=1) + manager.register_pools((_dense_pool("fa", 0),)) + + first = manager.get_block_ref(pool_idx=0, block_id=0) + second = manager.get_block_ref(pool_idx=0, block_id=1) + + self.assertEqual(first["tag"], "fa") + self.assertEqual(first["block_stride_bytes"], _block_bytes()) + self.assertEqual(second["ptr"] - first["ptr"], _block_bytes()) + + def test_register_pools_rejects_out_of_range_storage_index(self): + manager, _, _ = self._make_manager(num_layers=1) + + with self.assertRaisesRegex(ValueError, "out of range"): + manager.register_pools((_dense_pool("fa", 3),)) + + def test_register_pools_rejects_an_empty_table(self): + manager, _, _ = self._make_manager(num_layers=1) + + with self.assertRaisesRegex(ValueError, "non-empty"): + manager.register_pools(()) + + def test_pool_d2h_h2d_round_trip_leaves_the_device_untouched(self): + """The byte oracle: pool-addressed copies must be offset-exact. + + Pool D2H/H2D mirror host and device at the *same* byte offsets, so a + correct round trip is the identity on the device array. Any error in the + base offset, block stride, or region extent scrambles the array instead. + """ + manager, arrays, refs = self._make_manager(num_layers=2) + manager.register_pools((_half_block_pool("fa", 0), _dense_pool("fa", 1))) + + for pool_idx in (0, 1): + manager.d2h_pool_blocks(pool_idx, list(range(_NUM_BLOCKS))).Await() + for pool_idx in (0, 1): + manager.h2d_pool_blocks(pool_idx, list(range(_NUM_BLOCKS))).Await() + jax.block_until_ready(arrays) + + for layer, (array, ref) in enumerate(zip(arrays, refs)): + np.testing.assert_array_equal( + np.asarray(array), ref, err_msg=f"layer {layer} corrupted" + ) + + def test_pool_d2h_accepts_a_block_subset(self): + manager, arrays, refs = self._make_manager(num_layers=1) + manager.register_pools((_dense_pool("fa", 0),)) + + manager.d2h_pool_blocks(0, [1, 3]).Await() + manager.h2d_pool_blocks(0, [1, 3]).Await() + jax.block_until_ready(arrays) + + np.testing.assert_array_equal(np.asarray(arrays[0]), refs[0]) + + def test_pool_index_out_of_range_raises(self): + manager, _, _ = self._make_manager(num_layers=1) + manager.register_pools((_dense_pool("fa", 0),)) + + with self.assertRaises((IndexError, RuntimeError)): + manager.pool_spec(7) + + +class JaxControlListenerTest(_ManagerTestBase): + """The control plane the RaidenController drives, absent from JAX until now.""" + + def test_listener_is_off_by_default(self): + manager, _, _ = self._make_manager(num_layers=1) + + self.assertIsNone(manager.listener_port) + self.assertFalse(manager.is_control_listener_active) + self.assertEqual(manager.listener_address, "") + + def test_listener_binds_an_ephemeral_port(self): + manager, _, _ = self._make_manager(num_layers=1, listener_port=0) + + port = manager.listener_port + self.assertIsNotNone(port) + self.assertGreater(port, 0) + self.assertTrue(manager.is_control_listener_active) + self.assertTrue( + manager.listener_address.endswith(f":{port}"), + msg=f"listener_address={manager.listener_address!r} port={port}", + ) + # The control listener and the WorkerService gRPC server are separate + # sockets; the latter is off here, and that must not affect the former. + self.assertEqual(manager.get_raiden_worker_port(), 0) + self.assertFalse(manager.is_listener_active) + + +if __name__ == "__main__": + absltest.main() diff --git a/tpu_raiden/frameworks/jax/BUILD b/tpu_raiden/frameworks/jax/BUILD index ed10fda9..2dece308 100644 --- a/tpu_raiden/frameworks/jax/BUILD +++ b/tpu_raiden/frameworks/jax/BUILD @@ -48,6 +48,9 @@ cc_library( "//tpu_sync/core:utils", "//tpu_sync/core/controller:controller_client", "//tpu_sync/core/controller:worker_service_server", + "//tpu_sync/kv_cache:kv_cache_listener", + "//tpu_sync/kv_cache:pool_layout", + "//tpu_sync/rpc:raiden_service_cc_proto", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", @@ -142,10 +145,6 @@ nanobind_extension( ":nb_statusor", ":raw_transfer_jax", ":weight_synchronizer_jax", - "//tpu_sync/core:raw_transfer_core", - "//tpu_sync/kv_cache:kv_cache_store", - "//tpu_sync/kv_cache:kv_cache_store_wrapper", - "//tpu_sync/weight_sync:weight_synchronizer_base", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", @@ -153,6 +152,15 @@ nanobind_extension( "@nanobind", "@rules_python//python/cc:current_py_cc_headers", "@xla//xla/pjrt:status_casters", + # Framework-neutral despite living under frameworks/torch: the header + # includes no torch and its cc_library has no torch dep, so binding + # the same pool surface here keeps JAX and torch from diverging. + "//tpu_raiden/frameworks/torch:pool_layout_nanobind", + "//tpu_sync/core:raw_transfer_core", + "//tpu_sync/kv_cache:kv_cache_store", + "//tpu_sync/kv_cache:kv_cache_store_wrapper", + "//tpu_sync/rpc:raiden_service_cc_proto", + "//tpu_sync/weight_sync:weight_synchronizer_base", ], ) @@ -711,6 +719,9 @@ cc_library( "//tpu_sync/core:utils", "//tpu_sync/core/controller:controller_client", "//tpu_sync/core/controller:worker_service_server", + "//tpu_sync/kv_cache:kv_cache_listener", + "//tpu_sync/kv_cache:pool_layout", + "//tpu_sync/rpc:raiden_service_cc_proto", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", diff --git a/tpu_raiden/frameworks/jax/kv_cache_manager.cc b/tpu_raiden/frameworks/jax/kv_cache_manager.cc index e145ded3..c17abc87 100644 --- a/tpu_raiden/frameworks/jax/kv_cache_manager.cc +++ b/tpu_raiden/frameworks/jax/kv_cache_manager.cc @@ -53,6 +53,7 @@ #include "tpu_sync/core/utils.h" // IWYU pragma: keep #ifndef WITHOUT_PYTHON #include "tpu_raiden/frameworks/jax/utils.h" +#include "tpu_sync/kv_cache/kv_cache_listener.h" namespace nb = nanobind; #endif @@ -142,6 +143,17 @@ ::tpu_raiden::HostBufferAllocator LocalCreateHostMemoryAllocator( return alloc; }; } + +// Same bracketing rule as the torch manager uses, so the address strings the +// controller sees are identical whichever framework produced them. Defined +// outside the Python-only section because listener_address() and +// transfer_address() are declared unconditionally. +std::string FormatAddressWithPort(absl::string_view ip, int port) { + if (absl::StrContains(ip, ':')) { + return absl::StrCat("[", ip, "]:", port); + } + return absl::StrCat(ip, ":", port); +} } // namespace #ifndef WITHOUT_PYTHON @@ -929,10 +941,12 @@ KVCacheManager::KVCacheManager( std::optional host_blocks_to_allocate, bool unsafe_skip_buffer_lock, int parallelism, int raiden_worker_port, std::optional raiden_controller_address, - std::optional worker_id, int64_t node_id) + std::optional worker_id, int64_t node_id, + std::optional listener_port) : numa_manager_(std::make_unique( std::move(device_arrays), local_port, host_blocks_to_allocate, unsafe_skip_buffer_lock, parallelism, node_id)) { + StartListener(listener_port); StartGrpcServer(raiden_worker_port, raiden_controller_address, worker_id); } @@ -941,10 +955,11 @@ KVCacheManager::KVCacheManager( int64_t max_blocks, int64_t num_slots, double timeout_s, bool unsafe_skip_buffer_lock, int parallelism, int raiden_worker_port, std::optional raiden_controller_address, - std::optional worker_id) + std::optional worker_id, std::optional listener_port) : numa_manager_(std::make_unique( std::move(kv_caches), node_id, local_control_port, max_blocks, num_slots, timeout_s, unsafe_skip_buffer_lock, parallelism)) { + StartListener(listener_port); StartGrpcServer(raiden_worker_port, raiden_controller_address, worker_id); } #endif @@ -954,10 +969,11 @@ KVCacheManager::KVCacheManager( std::optional local_port, std::optional host_blocks_to_allocate, int parallelism, int raiden_worker_port, std::optional raiden_controller_address, - std::optional worker_id) + std::optional worker_id, std::optional listener_port) : numa_manager_(std::make_unique( num_layers, num_shards, slice_byte_size, local_port, host_blocks_to_allocate, parallelism)) { + StartListener(listener_port); StartGrpcServer(raiden_worker_port, raiden_controller_address, worker_id); } @@ -965,12 +981,61 @@ KVCacheManager::KVCacheManager( std::vector> sub_managers, int raiden_worker_port, std::optional raiden_controller_address, - std::optional worker_id) + std::optional worker_id, std::optional listener_port) : numa_manager_( std::make_unique(std::move(sub_managers))) { + StartListener(listener_port); StartGrpcServer(raiden_worker_port, raiden_controller_address, worker_id); } +KVCacheManagerWithTransfer& KVCacheManager::PoolTarget() const { + KVCacheManagerWithTransfer* target = numa_manager_->sole_sub_manager(); + if (target == nullptr) { + throw std::runtime_error(absl::StrCat( + "pool operations require exactly one NUMA sub-manager, but this " + "manager has ", + numa_manager_->num_sub_managers(), + ". Multi-NUMA pool fan-out is not implemented; construct the manager " + "without ENABLE_MULTI_NUMA.")); + } + return *target; +} + +void KVCacheManager::StartListener(std::optional listener_port) { + if (!listener_port.has_value()) { + return; + } + // Bind against the sub-manager, not the facade: KVCacheListener drives + // KVCacheManagerBase, and the facade is not one. PoolTarget() also enforces + // the single-sub-manager restriction, so a multi-NUMA manager refuses to + // start a listener instead of arming only shard group 0. + listener_ = std::make_unique(&PoolTarget(), + *listener_port); +} + +std::optional KVCacheManager::listener_port() const { + if (listener_) { + return listener_->listener_port(); + } + return std::nullopt; +} + +bool KVCacheManager::is_listener_active() const { + return listener_ ? listener_->is_active() : false; +} + +std::string KVCacheManager::listener_address() const { + auto port = listener_port(); + if (!port.has_value()) return ""; + return FormatAddressWithPort(PoolTarget().local_ip(), *port); +} + +std::string KVCacheManager::transfer_address() const { + auto port = numa_manager_->local_port(); + if (!port.has_value()) return ""; + return FormatAddressWithPort(PoolTarget().local_ip(), *port); +} + KVCacheManager::~KVCacheManager() { if (private_grpc_server_) { private_grpc_server_->SetTransferManager(nullptr); diff --git a/tpu_raiden/frameworks/jax/kv_cache_manager.h b/tpu_raiden/frameworks/jax/kv_cache_manager.h index bfcbccc9..548638d8 100644 --- a/tpu_raiden/frameworks/jax/kv_cache_manager.h +++ b/tpu_raiden/frameworks/jax/kv_cache_manager.h @@ -36,6 +36,13 @@ #include "tpu_sync/core/controller/worker_service_server.h" #include "tpu_sync/core/kv_cache_manager_with_transfer.h" #include "tpu_sync/core/tpu_utils.h" +// Included rather than forward declared: KVCacheManager owns the listener by +// unique_ptr and keeps an inline defaulted move constructor, which nanobind +// instantiates (detail::wrap_move) when it wraps the class. That instantiation +// needs unique_ptr's deleter, hence the complete type. Torch gets away with a +// forward declaration because its bound manager is not move-wrapped. +#include "tpu_sync/kv_cache/kv_cache_listener.h" +#include "tpu_sync/kv_cache/pool_layout.h" namespace xla { class PjRtBuffer; @@ -45,6 +52,7 @@ namespace tpu_raiden { class MetricsCollector; namespace kv_cache { + namespace jax { struct UnpackedCache { @@ -124,6 +132,18 @@ class NumaAwareKVCacheManager { std::vector get_local_endpoints() const; std::vector get_local_data_endpoints() const; + size_t num_sub_managers() const { return sub_managers_.size(); } + + // The sole sub-manager. Pool admission and the pool-addressed reshard path + // are byte-space operations over one manager's storages; under + // ENABLE_MULTI_NUMA they would have to fan out across sub-managers with + // global_shard_to_submanager_ remapping, which is deliberately not + // implemented yet. Returns nullptr in that case so callers fail loudly + // rather than silently addressing only sub-manager 0. + KVCacheManagerWithTransfer* sole_sub_manager() const { + return sub_managers_.size() == 1 ? sub_managers_[0].get() : nullptr; + } + void SetSubmanagerShardsForTesting( const std::vector>& assignment) { submanager_to_global_shards_ = assignment; @@ -255,7 +275,8 @@ class KVCacheManager { bool unsafe_skip_buffer_lock = false, int parallelism = 1, int raiden_worker_port = 0, std::optional raiden_controller_address = std::nullopt, - std::optional worker_id = std::nullopt, int64_t node_id = 0); + std::optional worker_id = std::nullopt, int64_t node_id = 0, + std::optional listener_port = std::nullopt); // New transfer-enabled constructor (flat list of arrays, single shard per // layer) @@ -264,7 +285,8 @@ class KVCacheManager { int64_t max_blocks, int64_t num_slots, double timeout_s, bool unsafe_skip_buffer_lock, int parallelism, int raiden_worker_port = 0, std::optional raiden_controller_address = std::nullopt, - std::optional worker_id = std::nullopt); + std::optional worker_id = std::nullopt, + std::optional listener_port = std::nullopt); #endif // FFI metadata constructor (cache-only by default) @@ -273,17 +295,32 @@ class KVCacheManager { std::optional local_port, std::optional host_blocks_to_allocate, int parallelism = 1, int raiden_worker_port = 0, std::optional raiden_controller_address = std::nullopt, - std::optional worker_id = std::nullopt); + std::optional worker_id = std::nullopt, + std::optional listener_port = std::nullopt); // Test-only constructor for sub-manager mock injection explicit KVCacheManager( std::vector> sub_managers, int raiden_worker_port = 0, std::optional raiden_controller_address = std::nullopt, - std::optional worker_id = std::nullopt); + std::optional worker_id = std::nullopt, + std::optional listener_port = std::nullopt); ~KVCacheManager(); + // --- Control-plane listener ------------------------------------------- + // The KVCacheListener decodes START_TRANSFER off a raw socket and drives + // PoolReshardRegisterRecv / PoolReshardPush. It is how a RaidenController + // arms a receiver and fires a sender. Torch has had this since + // frameworks/torch/kv_cache_manager.cc; these mirror that surface exactly. + // + // Note this is NOT the WorkerService gRPC port (get_raiden_worker_port()): + // that is the buffer-oriented worker service, a separate server. + std::optional listener_port() const; + bool is_listener_active() const; + std::string listener_address() const; + std::string transfer_address() const; + NumaAwareKVCacheManager* numa_manager() const { return numa_manager_.get(); } int GetRaidenWorkerPort() const; @@ -450,13 +487,92 @@ class KVCacheManager { dst_block_ids); } + // --- Pool admission + pool-addressed reshard -------------------------- + // Thin forwarders onto the sole sub-manager, in the same spirit as the + // NotifyForRead / StartRead forwarders above. They exist because + // BindPoolApi<> (frameworks/torch/pool_layout_nanobind.h) resolves these + // names on the bound class, and the JAX facade is a NumaAware* wrapper + // rather than a KVCacheManagerBase subclass -- so it cannot inherit them. + // Signatures are copied verbatim from KVCacheManagerBase so the template + // instantiates identically for JAX and for torch. + absl::Status RegisterPools(std::vector pools) { + return PoolTarget().RegisterPools(std::move(pools)); + } + + absl::StatusOr GetPoolBlockRef(size_t pool_idx, + size_t shard_idx, + int64_t block_id) const { + return PoolTarget().GetPoolBlockRef(pool_idx, shard_idx, block_id); + } + + const PoolSpec* pool(size_t pool_idx) const { + return PoolTarget().pool(pool_idx); + } + + size_t num_pools() const { return PoolTarget().num_pools(); } + + bool has_explicit_pools() const { + return PoolTarget().has_explicit_pools(); + } + + std::vector PoolIndicesWithTag(absl::string_view tag) const { + return PoolTarget().PoolIndicesWithTag(tag); + } + + int64_t LayerBlockByteSize(size_t layer_idx) const { + return PoolTarget().LayerBlockByteSize(layer_idx); + } + + absl::StatusOr GetBlockHostPointerValue(size_t layer_idx, + size_t shard_idx, + int block_id) { + return PoolTarget().GetBlockHostPointerValue(layer_idx, shard_idx, + block_id); + } + + absl::StatusOr D2hPoolBlocks( + size_t pool_idx, absl::Span block_ids, + std::optional shard_idx = std::nullopt) { + return PoolTarget().D2hPoolBlocks(pool_idx, block_ids, shard_idx); + } + + absl::StatusOr H2dPoolBlocks( + size_t pool_idx, absl::Span block_ids, + std::optional shard_idx = std::nullopt) { + return PoolTarget().H2dPoolBlocks(pool_idx, block_ids, shard_idx); + } + + absl::Status RegisterActivePlan( + uint64_t uuid, const tpu_raiden::rpc::StartTransferRequest& request, + bool is_sender) { + return PoolTarget().RegisterActivePlan(uuid, request, is_sender); + } + + absl::Status UnregisterActivePlan(uint64_t uuid) { + return PoolTarget().UnregisterActivePlan(uuid); + } + + absl::Status RegisterRecv(uint64_t uuid, const std::string& req_id, + int64_t expected_block_count) { + return PoolTarget().RegisterRecv(uuid, req_id, expected_block_count); + } + private: + // Resolves the sub-manager every pool op addresses, or throws. Throwing + // (rather than returning a Status) keeps the forwarders signature-identical + // to KVCacheManagerBase's, which is what lets BindPoolApi<> instantiate + // unchanged; nanobind turns it into a Python exception at the boundary. + KVCacheManagerWithTransfer& PoolTarget() const; + + void StartListener(std::optional listener_port); + void StartGrpcServer( int raiden_worker_port, std::optional raiden_controller_address = std::nullopt, std::optional worker_id = std::nullopt); std::unique_ptr numa_manager_; + std::unique_ptr listener_; std::unique_ptr private_grpc_server_; }; diff --git a/tpu_raiden/frameworks/jax/tpu_raiden_jax_module.cc b/tpu_raiden/frameworks/jax/tpu_raiden_jax_module.cc index f6120fd4..d76b7e8b 100644 --- a/tpu_raiden/frameworks/jax/tpu_raiden_jax_module.cc +++ b/tpu_raiden/frameworks/jax/tpu_raiden_jax_module.cc @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -37,10 +38,12 @@ #include "tpu_raiden/frameworks/jax/nb_statusor.h" // IWYU pragma: keep #include "tpu_raiden/frameworks/jax/raw_transfer_internal.h" #include "tpu_raiden/frameworks/jax/weight_synchronizer.h" +#include "tpu_raiden/frameworks/torch/pool_layout_nanobind.h" #include "tpu_sync/core/raiden_future.h" #include "tpu_sync/core/raw_transfer_core.h" #include "tpu_sync/kv_cache/kv_cache_store.h" #include "tpu_sync/kv_cache/kv_cache_store_wrapper.h" +#include "tpu_sync/rpc/raiden_service.pb.h" namespace nb = nanobind; @@ -95,26 +98,31 @@ NB_MODULE(_tpu_raiden_jax, m) { .def("IsReady", &tpu_raiden::RaidenFuture::IsReady) .def("is_ready", &tpu_raiden::RaidenFuture::IsReady); - nb::class_(m, "KVCacheManager") + auto manager_cls = + nb::class_(m, + "KVCacheManager"); + manager_cls .def(nb::init, std::optional, bool, int, int, std::optional, std::optional, - int64_t>(), + int64_t, std::optional>(), nb::arg("device_arrays"), nb::arg("local_port") = nb::none(), nb::arg("host_blocks_to_allocate") = nb::none(), nb::arg("unsafe_skip_buffer_lock") = false, nb::arg("parallelism") = 1, nb::arg("raiden_worker_port") = 0, nb::arg("raiden_controller_address") = nb::none(), - nb::arg("worker_id") = nb::none(), nb::arg("node_id") = 0) + nb::arg("worker_id") = nb::none(), nb::arg("node_id") = 0, + nb::arg("listener_port") = nb::none()) .def(nb::init, - std::optional>(), + std::optional, std::optional>(), nb::arg("kv_caches"), nb::arg("node_id") = 0, nb::arg("local_control_port"), nb::arg("max_blocks"), nb::arg("num_slots"), nb::arg("timeout_s") = 120.0, nb::arg("unsafe_skip_buffer_lock") = true, nb::arg("parallelism") = 4, nb::arg("raiden_worker_port") = 0, nb::arg("raiden_controller_address") = nb::none(), - nb::arg("worker_id") = nb::none()) + nb::arg("worker_id") = nb::none(), + nb::arg("listener_port") = nb::none()) // Use lambdas to wrap the returned raiden::PjRtCopyFuture into // KVCacheManagerFuture @@ -292,7 +300,59 @@ NB_MODULE(_tpu_raiden_jax, m) { &tpu_raiden::kv_cache::jax::KVCacheManager::UnlockBlocks, nb::arg("block_ids")) .def("dump_metrics_to_string", - &tpu_raiden::kv_cache::jax::KVCacheManager::DumpMetricsToString); + &tpu_raiden::kv_cache::jax::KVCacheManager::DumpMetricsToString) + + // Control-plane listener, mirroring the torch module's surface. This is + // the socket a RaidenController talks to when it arms a receiver or + // fires a sender for a pool-addressed reshard plan; distinct from + // get_raiden_worker_port(), which is the WorkerService gRPC port. + .def_prop_ro( + "listener_port", + &tpu_raiden::kv_cache::jax::KVCacheManager::listener_port) + .def_prop_ro( + "is_listener_active", + &tpu_raiden::kv_cache::jax::KVCacheManager::is_listener_active) + .def_prop_ro( + "listener_address", + &tpu_raiden::kv_cache::jax::KVCacheManager::listener_address) + .def_prop_ro( + "transfer_address", + &tpu_raiden::kv_cache::jax::KVCacheManager::transfer_address) + + .def( + "register_active_plan", + [](tpu_raiden::kv_cache::jax::KVCacheManager& self, uint64_t uuid, + nb::bytes serialized, bool is_sender) { + tpu_raiden::rpc::StartTransferRequest request; + if (!request.ParseFromArray(serialized.c_str(), + static_cast(serialized.size()))) { + throw std::runtime_error( + "register_active_plan: could not parse StartTransferRequest"); + } + tpu_raiden::torch_bindings::ThrowIfNotOk(self.RegisterActivePlan(uuid, request, is_sender), + "KVCacheManager register_active_plan failed"); + }, + nb::arg("uuid"), nb::arg("request"), nb::arg("is_sender")) + .def( + "unregister_active_plan", + [](tpu_raiden::kv_cache::jax::KVCacheManager& self, uint64_t uuid) { + tpu_raiden::torch_bindings::ThrowIfNotOk(self.UnregisterActivePlan(uuid), + "KVCacheManager unregister_active_plan failed"); + }, + nb::arg("uuid")) + .def( + "register_recv", + [](tpu_raiden::kv_cache::jax::KVCacheManager& self, uint64_t uuid, + const std::string& req_id, int64_t expected_block_count) { + tpu_raiden::torch_bindings::ThrowIfNotOk(self.RegisterRecv(uuid, req_id, expected_block_count), + "KVCacheManager register_recv failed"); + }, + nb::arg("uuid"), nb::arg("req_id"), nb::arg("expected_block_count")); + + // Pool admission + partial D2H/H2D, shared verbatim with the torch and host + // modules so the three surfaces cannot drift. + tpu_raiden::torch_bindings::BindPoolApi( + manager_cls); // ========================================================================= // 2. Bind WeightSynchronizer