From 37b1696566228458c5435a25d0b13125698aa380 Mon Sep 17 00:00:00 2001 From: Googler Date: Wed, 23 Sep 2026 13:55:26 -0700 Subject: [PATCH] [tpu-raiden] Add real-TPU e2e test for external-tensor H2D/D2H over mapped shared memory PiperOrigin-RevId: 986985408 --- .../api/torch/kv_cache_manager_dma_test.py | 311 ++++++++++++++++++ 1 file changed, 311 insertions(+) create mode 100644 tpu_sync/api/torch/kv_cache_manager_dma_test.py diff --git a/tpu_sync/api/torch/kv_cache_manager_dma_test.py b/tpu_sync/api/torch/kv_cache_manager_dma_test.py new file mode 100644 index 000000000..71197cbf4 --- /dev/null +++ b/tpu_sync/api/torch/kv_cache_manager_dma_test.py @@ -0,0 +1,311 @@ +# 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. + +"""Real-TPU tests for shared memory DMA mapping and object tensor transfers. + +These exercise KVCacheManager.experimental_map_shared_memory plus the +object-tensor H2D/D2H overloads against a live PJRT client. + +test_object_tensor_bandwidth additionally measures throughput. Its threshold is +deliberately loose: it only catches a silent fall-off to staged pageable copies, +an order of magnitude below DMA speed. It is not a regression gate, and +scheduling jitter must not flake it. +""" + +import os +import statistics +import time + +from absl import flags +from absl.testing import absltest +import torch + +from tpu_sync.api.torch import kv_cache_manager +from tpu_sync.api.torch import torch_tpu_common_loader + +KVCacheManager = kv_cache_manager.KVCacheManager + +_NUM_LAYERS = 2 +_NUM_BLOCKS = 4 +_NUM_RANKS = 2 +# Per-layer device buffer: [num_blocks, ...] float32, so one block per layer is +# 8 * 8 * 128 * 4 = 32 KiB. +_BLOCK_SHAPE = (_NUM_BLOCKS, 8, 8, 128) + +# The correctness tests above move 32 KiB per layer, where per-call overhead +# dominates and the measured rate says nothing about DMA bandwidth. The +# bandwidth test uses its own, much larger geometry. +_PERF_LAYERS = flags.DEFINE_integer( + "perf_layers", 8, "Layers (device buffers) for the bandwidth test." +) +_PERF_BLOCKS = flags.DEFINE_integer( + "perf_blocks", 8, "Blocks per h2d/d2h call in the bandwidth test." +) +_PERF_BLOCK_MIB = flags.DEFINE_integer( + "perf_block_mib", 4, "MiB per block per layer in the bandwidth test." +) +_PERF_ITERS = flags.DEFINE_integer( + "perf_iters", 8, "Timed iterations per direction." +) +_PERF_WARMUP = flags.DEFINE_integer( + "perf_warmup", 2, "Untimed warmup iterations." +) +_MIN_GBPS = flags.DEFINE_float( + "min_gbps", + 2.0, + "Fail below this bandwidth (GB/s, decimal). Loose on purpose: it catches a " + "fall-off to staged pageable copies, not a drift regression.", +) + + +class SharedMemoryDmaTest(absltest.TestCase): + + def setUp(self): + super().setUp() + # Registers the "tpu" device name. Importing torch_tpu is not enough: the + # privateuse1 rename happens inside torch_tpu's loader, which + # kv_cache_manager only calls lazily from the manager constructor, i.e. + # after setUp needs the device. + torch_tpu_common_loader.load_torch_tpu_common() + self.device = torch.device("tpu") + self.page_size = os.sysconf("SC_PAGE_SIZE") + + def _make_manager(self) -> KVCacheManager: + kv_caches = [ + torch.zeros(_BLOCK_SHAPE, dtype=torch.float32, device=self.device) + for _ in range(_NUM_LAYERS) + ] + return KVCacheManager( + kv_caches=kv_caches, + node_id=0, + local_control_port=0, + max_blocks=_NUM_BLOCKS, + num_slots=2, + ) + + def _alloc_pool(self, min_bytes: int): + """Returns (backing_tensor, pad, pool_addr, pool_size) for a page-aligned pool. + + experimental_map_shared_memory requires both the address and the size to + be page aligned, so over-allocate by one page and start the pool at the + first aligned offset inside it. + + Args: + min_bytes: smallest usable pool size; rounded up to a page multiple. + """ + pool_size = ( + (min_bytes + self.page_size - 1) // self.page_size + ) * self.page_size + backing = torch.empty(pool_size + self.page_size, dtype=torch.int8) + base = backing.data_ptr() + pad = (-base) % self.page_size + return backing, pad, base + pad, pool_size + + def _bench(self, fn) -> float: + """Returns the median wall time of fn over _PERF_ITERS timed calls.""" + times = [] + for i in range(_PERF_WARMUP.value + _PERF_ITERS.value): + start = time.perf_counter() + fn() + elapsed = time.perf_counter() - start + if i >= _PERF_WARMUP.value: + times.append(elapsed) + return statistics.median(times) + + def test_h2d_d2h_roundtrip_through_mapped_pool(self): + manager = self._make_manager() + page_nbytes = manager._impl.slice_byte_size # pylint: disable=protected-access + obj_nbytes = _NUM_RANKS * _NUM_LAYERS * page_nbytes + + # src and dst object tensors both live inside the registered pool, which is + # how a caller is expected to use this API. + backing, pad, pool_addr, pool_size = self._alloc_pool(2 * obj_nbytes) + src = backing[pad : pad + obj_nbytes].view( + _NUM_RANKS, _NUM_LAYERS, page_nbytes + ) + dst = backing[pad + obj_nbytes : pad + 2 * obj_nbytes].view( + _NUM_RANKS, _NUM_LAYERS, page_nbytes + ) + src[0].fill_(0x42) + src[1].fill_(0x77) + dst.fill_(0) + + manager.experimental_map_shared_memory(pool_addr, pool_size) + self.assertTrue(manager.experimental_is_shared_memory_mapped) + + try: + # Rank 0 -> block 3 -> back again. + manager.experimental_h2d([3], [src], 0).wait() + manager.experimental_d2h([3], [dst], 0).wait() + self.assertTrue( + torch.equal(dst[0], src[0]), + "rank 0 did not round-trip through the device", + ) + + # Rank 1 lands in a different block and must not disturb rank 0. + manager.experimental_h2d([1], [src], 1).wait() + manager.experimental_d2h([1], [dst], 1).wait() + self.assertTrue( + torch.equal(dst[1], src[1]), + "rank 1 did not round-trip through the device", + ) + self.assertTrue(torch.equal(dst[0], src[0])) + finally: + manager.experimental_unmap_shared_memory() + + self.assertFalse(manager.experimental_is_shared_memory_mapped) + del backing + + def test_tensor_outside_mapped_pool_is_rejected(self): + manager = self._make_manager() + page_nbytes = manager._impl.slice_byte_size # pylint: disable=protected-access + obj_nbytes = _NUM_RANKS * _NUM_LAYERS * page_nbytes + + backing, pad, pool_addr, pool_size = self._alloc_pool(obj_nbytes) + inside = backing[pad : pad + obj_nbytes].view( + _NUM_RANKS, _NUM_LAYERS, page_nbytes + ) + # Separate allocation, so it cannot overlap the pool. + outside = torch.zeros( + (_NUM_RANKS, _NUM_LAYERS, page_nbytes), dtype=torch.int8 + ) + + manager.experimental_map_shared_memory(pool_addr, pool_size) + try: + # Sanity: the in-pool tensor is accepted, so a rejection below is really + # about the address and not about some other argument. + manager.experimental_h2d([0], [inside], 0).wait() + + with self.assertRaisesRegex(Exception, "outside the DMA mapped pool"): + manager.experimental_h2d([0], [outside], 0) + with self.assertRaisesRegex(Exception, "outside the DMA mapped pool"): + manager.experimental_d2h([0], [outside], 0) + finally: + manager.experimental_unmap_shared_memory() + del backing + + def test_transfers_refused_when_not_mapped(self): + manager = self._make_manager() + page_nbytes = manager._impl.slice_byte_size # pylint: disable=protected-access + obj_nbytes = _NUM_RANKS * _NUM_LAYERS * page_nbytes + + backing, pad, pool_addr, pool_size = self._alloc_pool(obj_nbytes) + obj = backing[pad : pad + obj_nbytes].view( + _NUM_RANKS, _NUM_LAYERS, page_nbytes + ) + + # Before mapping. + self.assertFalse(manager.experimental_is_shared_memory_mapped) + with self.assertRaisesRegex(Exception, "must be mapped"): + manager.experimental_h2d([0], [obj], 0) + + # And again after the pool is released. + manager.experimental_map_shared_memory(pool_addr, pool_size) + manager.experimental_h2d([0], [obj], 0).wait() + manager.experimental_unmap_shared_memory() + + self.assertFalse(manager.experimental_is_shared_memory_mapped) + with self.assertRaisesRegex(Exception, "must be mapped"): + manager.experimental_d2h([0], [obj], 0) + del backing + + def test_remap_after_unmap(self): + manager = self._make_manager() + page_nbytes = manager._impl.slice_byte_size # pylint: disable=protected-access + obj_nbytes = _NUM_RANKS * _NUM_LAYERS * page_nbytes + + backing, pad, pool_addr, pool_size = self._alloc_pool(obj_nbytes) + obj = backing[pad : pad + obj_nbytes].view( + _NUM_RANKS, _NUM_LAYERS, page_nbytes + ) + obj[0].fill_(0x5A) + + for _ in range(2): + manager.experimental_map_shared_memory(pool_addr, pool_size) + self.assertTrue(manager.experimental_is_shared_memory_mapped) + manager.experimental_h2d([2], [obj], 0).wait() + manager.experimental_d2h([2], [obj], 0).wait() + manager.experimental_unmap_shared_memory() + self.assertFalse(manager.experimental_is_shared_memory_mapped) + del backing + + def test_object_tensor_bandwidth(self): + num_layers = _PERF_LAYERS.value + num_blocks = _PERF_BLOCKS.value + # float32, so 1024 * 256 * 4 == 1 MiB per unit of dim 1. + block_shape = (num_blocks, _PERF_BLOCK_MIB.value, 1024, 256) + + kv_caches = [ + torch.zeros(block_shape, dtype=torch.float32, device=self.device) + for _ in range(num_layers) + ] + torch.tpu.synchronize() + manager = KVCacheManager( + kv_caches=kv_caches, + node_id=0, + local_control_port=0, + max_blocks=num_blocks, + num_slots=2, + ) + page_nbytes = manager._impl.slice_byte_size # pylint: disable=protected-access + + # One object tensor per block, each [num_ranks=1, num_layers, page_nbytes]. + # Only rank 0's slice is transferred, so that is the whole tensor here. + obj_nbytes = num_layers * page_nbytes + backing, pad, pool_addr, pool_size = self._alloc_pool( + num_blocks * obj_nbytes + ) + tensors = [ + backing[pad + i * obj_nbytes : pad + (i + 1) * obj_nbytes].view( + 1, num_layers, page_nbytes + ) + for i in range(num_blocks) + ] + block_ids = list(range(num_blocks)) + payload_bytes = num_blocks * obj_nbytes + + manager.experimental_map_shared_memory(pool_addr, pool_size) + try: + h2d_s = self._bench( + lambda: manager.experimental_h2d(block_ids, tensors, 0).wait() + ) + d2h_s = self._bench( + lambda: manager.experimental_d2h(block_ids, tensors, 0).wait() + ) + finally: + manager.experimental_unmap_shared_memory() + + mib = payload_bytes / (1024.0 * 1024.0) + print( + f"[object-tensor perf] layers={num_layers} blocks={num_blocks} " + f"page={page_nbytes / (1024.0 * 1024.0):.3f} MiB " + f"payload={mib:.1f} MiB iters={_PERF_ITERS.value}" + ) + rates = {} + for name, seconds in (("H2D", h2d_s), ("D2H", d2h_s)): + gbps = payload_bytes / seconds / 1e9 + gibps = payload_bytes / seconds / (1024.0**3) + rates[name] = gbps + print( + f"[object-tensor perf] {name} median {seconds * 1e3:.2f} ms -> " + f"{gbps:.1f} GB/s ({gibps:.1f} GiB/s)" + ) + + self.assertGreater(rates["H2D"], _MIN_GBPS.value) + self.assertGreater(rates["D2H"], _MIN_GBPS.value) + del backing + + +if __name__ == "__main__": + absltest.main()