JAX: pool admission + control-plane listener parity with torch

[copybara-service]

JAX: pool admission + control-plane listener parity with torch

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 routes prefill→decode KV transfer through this reshard path so the two sides can run different TP degrees and page sizes, over #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<RaidenFuture> 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_<KVCacheManager> instantiates detail::wrap_move, which reaches the inline defaulted move constructor, which needs unique_ptr<KVCacheListener>'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), 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_tests (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.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.