[TPU Raiden] Backend Tier Extension (Part 3/4): Python Framework Bindings & Integration Tests (JAX & PyTorch)

[copybara-service]

[TPU Raiden] Backend Tier Extension (Part 3/4): Python Framework Bindings & Integration Tests (JAX & PyTorch)

Exposes secondary storage configuration to PyTorch and JAX serving runtimes, enabling end-to-end offload (store.save) and recall (store.load) between KVCacheStore (Coordinator) and KVCacheManager (Worker).

Key changes:

• Native Bindings: Binds BackendConfig, ParallelismConfig, and BlockStatus::SHARED_STORAGE in nanobind (_tpu_raiden_torch, _tpu_raiden_jax) and common.py.
• Worker & Coordinator APIs: Accepts optional backend_configs in PyTorch and JAX KVCacheManager and KVCacheStore, registering backends at initialization.
• JAX NumaAware Support: Forwards D2hWriteToBackend, H2dReadFromBackend, and GetKVBackend across sub-managers.
• E2E TPU Hardware Tests:
  • PyTorch (kv_cache_store_mpmd_secondary_storage_e2e_test_gf): 4-rank MPMD test where independent per-rank workers manage shards:
    • Writer step: Offloads KV tensors from HBM to storage and verifies per-rank shard directories (tp4_r0..tp4_r3) and .bin files are created.
    • Reader step: Discovers on-disk shards, verifies cold lookup returns SHARED_STORAGE, recalls blocks into HBM across 4 TPU chips verifying bit-for-bit correctness, and confirms status promotion to HOST_AND_HBM.
  • JAX (kv_cache_store_secondary_storage_e2e_test_gf): Single KVCacheManager managing shards across TPU devices executing the same workflow:
    • Writer step: Offloads KV tensors to disk and verifies shard .bin files are created.
    • Reader step: Clears cache, verifies cold lookup returns SHARED_STORAGE, recalls blocks from storage, and confirms bit-for-bit tensor restoration on TPU hardware.