Add the Python/nanobind pool surface and host-only construction path

experiments/BUILD

@@ -36,7 +36,7 @@ py_binary(
     srcs = ["reshard_stage1_e2e.py"],
     main = "reshard_stage1_e2e.py",
     deps = [
-        "//tpu_raiden/api/torch:kv_cache_manager_host_py",
+        "//tpu_raiden/experimental:kv_cache_manager_host_py",
         "//tpu_raiden/rpc:raiden_service_py_pb2",
     ],
 )
@@ -50,7 +50,7 @@ py_test(
         "notap",
     ],
     deps = [
-        "//tpu_raiden/api/torch:kv_cache_manager_host_py",
+        "//tpu_raiden/experimental:kv_cache_manager_host_py",
         "//tpu_raiden/rpc:raiden_service_py_pb2",
         "@com_google_absl_py//absl/testing:absltest",
     ],

experiments/reshard_stage1_e2e.py

@@ -46,7 +46,7 @@
 import time
 from typing import Any, Dict, List, Optional, Sequence, Tuple
 
-from tpu_raiden.api.torch import kv_cache_manager_host
+from tpu_raiden.experimental import kv_cache_manager_host
 from tpu_raiden.rpc import raiden_service_pb2
 
 NUM_FA_LAYERS = 15

experiments/reshard_stage1_local_primitive_test.py

@@ -38,7 +38,7 @@
 
 from absl.testing import absltest
 
-from tpu_raiden.api.torch import kv_cache_manager_host
+from tpu_raiden.experimental import kv_cache_manager_host
 from tpu_raiden.rpc import raiden_service_pb2
 
 NUM_FA_LAYERS = 15

tpu_raiden/api/torch/BUILD

@@ -26,11 +26,18 @@ py_library(
     ],
 )
 
+py_library(
+    name = "pool_layout",
+    srcs = ["pool_layout.py"],
+    visibility = ["//visibility:public"],
+)
+
 py_library(
     name = "kv_cache_manager_torch_py",
     srcs = ["kv_cache_manager.py"],
     visibility = ["//visibility:public"],
     deps = [
+        ":pool_layout",
         ":torch_tpu_common_loader",
         "//tpu_raiden/frameworks/torch:_tpu_raiden_torch",
         "//tpu_raiden/rpc:raiden_service_py_pb2",
@@ -39,10 +46,21 @@ py_library(
 
 py_library(
     name = "kv_cache_manager_host_py",
-    srcs = ["kv_cache_manager_host.py"],
+    srcs = ["kv_cache_manager.py"],
     visibility = ["//visibility:public"],
     deps = [
+        ":pool_layout",
         "//tpu_raiden/frameworks/torch:_tpu_raiden_host",
+        "//tpu_raiden/rpc:raiden_service_py_pb2",
+    ],
+)
+
+py_test(
+    name = "kv_cache_manager_host_test",
+    srcs = ["kv_cache_manager_host_test.py"],
+    deps = [
+        ":kv_cache_manager_host_py",
+        ":pool_layout",
     ],
 )
 

tpu_raiden/api/torch/kv_cache_manager.py

@@ -30,9 +30,21 @@
 
 from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 
+_HOST_IMPL = None
 _TORCH_IMPL = None
 
 
+def _host_impl():
+  """Returns the host-only extension without importing torch_tpu."""
+  global _HOST_IMPL
+  if _HOST_IMPL is None:
+    # pylint: disable=g-import-not-at-top
+    from tpu_raiden.frameworks.torch import _tpu_raiden_host as impl
+    # pylint: enable=g-import-not-at-top
+    _HOST_IMPL = impl
+  return _HOST_IMPL
+
+
 def _torch_impl():
   """Returns the torch-backed extension, loading torch_tpu runtime first."""
   global _TORCH_IMPL
@@ -84,6 +96,7 @@ def __init__(
       listener_port: Sockets server port for incoming C++ KVCacheListener
         commands.
     """
+    self._admission_summary = None
     impl = _torch_impl()
     if host_blocks_to_allocate is not None:
       self._impl = impl.KVCacheManager(
@@ -111,6 +124,32 @@ def __init__(
           parallelism=parallelism,
       )
 
+  @classmethod
+  def create_host_only_for_testing(
+      cls,
+      *,
+      num_layers: int,
+      num_shards: int = 1,
+      slice_byte_size: int,
+      node_id: int,
+      local_port: Optional[int] = None,
+      host_blocks: int,
+      parallelism: int = 4,
+  ) -> "KVCacheManager":
+    """Creates a CPU-only manager backed by Raiden-owned host memory."""
+    obj = cls.__new__(cls)
+    obj._admission_summary = None
+    obj._impl = _host_impl().KVCacheManager(
+        num_layers=num_layers,
+        num_shards=num_shards,
+        slice_byte_size=slice_byte_size,
+        node_id=node_id,
+        local_port=local_port,
+        host_blocks_to_allocate=host_blocks,
+        parallelism=parallelism,
+    )
+    return obj
+
   @property
   def node_id(self) -> int:
     """Returns the active Worker or Shard ID."""
@@ -183,6 +222,101 @@ def write_block_bytes(
     """Overwrites one host block with bytes."""
     self._impl.write_block_bytes(layer_idx, block_id, payload)
 
+  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``).
+    """
+    # pylint: disable=g-import-not-at-top
+    from tpu_raiden.api.torch import pool_layout
+    # pylint: enable=g-import-not-at-top
+
+    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 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_recv(
       self, uuid: int, req_id: str, expected_block_count: int
   ) -> None:

tpu_raiden/api/torch/kv_cache_manager_host_test.py

@@ -0,0 +1,119 @@
+# 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.
+
+# 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.
+
+"""Host-only tests for the Python pool-admission surface."""
+
+import unittest
+
+from tpu_raiden.api.torch import kv_cache_manager
+from tpu_raiden.api.torch import pool_layout
+
+
+class KVCacheManagerHostTest(unittest.TestCase):
+
+  def test_register_refs_tags_and_admission_summary(self):
+    manager = kv_cache_manager.KVCacheManager.create_host_only_for_testing(
+        num_layers=1,
+        num_shards=1,
+        slice_byte_size=256,
+        node_id=7,
+        host_blocks=4,
+        parallelism=1,
+    )
+    first_pool = pool_layout.PoolSpec(
+        tag="kind_a",
+        storage_index=0,
+        base_offset_bytes=32,
+        block_stride_bytes=128,
+        num_blocks=4,
+        regions=(
+            pool_layout.RegionSpec(
+                name="payload",
+                offset_bytes=0,
+                stride_bytes=128,
+                unit_bytes=128,
+                num_units=1,
+            ),
+        ),
+        dtype_tag="dtype_a",
+    )
+    second_pool = pool_layout.PoolSpec(
+        tag="kind_b",
+        storage_index=0,
+        base_offset_bytes=640,
+        block_stride_bytes=64,
+        num_blocks=4,
+        regions=(
+            pool_layout.RegionSpec(
+                name="payload",
+                offset_bytes=0,
+                stride_bytes=64,
+                unit_bytes=64,
+                num_units=1,
+            ),
+        ),
+        dtype_tag="dtype_b",
+    )
+
+    summary = manager.register_pools((first_pool, second_pool))
+
+    self.assertEqual(
+        summary,
+        {
+            "admitted": True,
+            "pools": 2,
+            "storages": 1,
+            "tags": {"kind_a": 1, "kind_b": 1},
+        },
+    )
+    self.assertEqual(manager.admission_summary(), summary)
+    self.assertEqual(manager.num_pools(), 2)
+    self.assertEqual(manager.pool_ids_with_tag("kind_a"), [0])
+    self.assertEqual(manager.pool_ids_with_tag("kind_b"), [1])
+    self.assertEqual(manager.pool_ids_with_tag("missing"), [])
+
+    first_ref = manager.get_block_ref(pool_idx=0, block_id=0)
+    first_ref_last = manager.get_block_ref(pool_idx=0, block_id=3)
+    second_ref = manager.get_block_ref(pool_idx=1, block_id=2)
+    self.assertEqual(first_ref_last["ptr"], first_ref["ptr"] + 3 * 128)
+    self.assertEqual(second_ref["ptr"], first_ref["ptr"] + (640 - 32) + 2 * 64)
+    self.assertEqual(first_ref["block_stride_bytes"], 128)
+    self.assertEqual(second_ref["block_stride_bytes"], 64)
+    self.assertEqual(first_ref["tag"], "kind_a")
+    self.assertEqual(second_ref["tag"], "kind_b")
+    self.assertEqual(second_ref["dtype_tag"], "dtype_b")
+    self.assertEqual(manager.pool_spec(1)["base_offset_bytes"], 640)
+
+    with self.assertRaisesRegex(RuntimeError, "host-only manager"):
+      manager.d2h_pool_blocks(pool_idx=0, block_ids=[0])
+    with self.assertRaisesRegex(RuntimeError, "host-only manager"):
+      manager.h2d_pool_blocks(pool_idx=0, block_ids=[0])
+
+
+if __name__ == "__main__":
+  unittest.main()