Add DMA-mapped PyTorch KV cache offloader

[enkidu-214]

Summary

• Add a PyTorch KVCacheOffloader for bidirectional KV-cache page transfers.
• Register one caller-owned, page-locked host-memory pool through TpuClient::DmaMap, then reuse that registration for H2D and D2H transfers from smaller object-tensor views within the pool.
• Keep host-pool allocation and object-view construction outside tpu-sync, while retaining tensors and device buffers until asynchronous transfers complete.
• Serialize map/unmap transitions, reject duplicate lifecycle operations, drain in-flight copies before DmaUnmap, and surface the API through the native Torch extension and Python wrapper.
• Require libtpu>=0.0.47, which provides the DmaMap behavior used by this path.

Design

The caller owns the host-memory pool, its page locking, its process-local mapping lifetime, and the construction of mutually disjoint object-tensor views. map_shared_memory(address, size) registers the whole pool once with the same PJRT client that owns the KV-cache device buffers. Subsequent h2d and d2h calls operate on complete object-tensor views and issue raw page transfers without registering each object separately.

unmap_shared_memory() first prevents new submissions, waits for all previously submitted copies, and then releases the pool registration through TpuClient::DmaUnmap.

Validation

• bazel test //tpu_sync/frameworks/torch:kv_cache_offloader_test
• python -m unittest -v tpu_sync.api.torch.kv_cache_offloader_test
• TPU v7x functional DMA round trip on all 8 ranks with topology 2,2,1,2: H2D, clear host buffers, D2H, byte comparison, and unmap all completed successfully.

[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.