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