diff --git a/remote_read_to_hbm_plan/global_prefix_cache.md b/remote_read_to_hbm_plan/global_prefix_cache.md
new file mode 100644
index 00000000..ca228fd4
--- /dev/null
+++ b/remote_read_to_hbm_plan/global_prefix_cache.md
@@ -0,0 +1,398 @@
+# Global Prefix Caching in TPU Raiden
+
+This document details the architectural design, component layering, communication protocols, Python APIs, and E2E workflows for the **Global Prefix Caching** feature in TPU Raiden.
+
+---
+
+## 1. Overview & Functionality
+
+In distributed LLM serving, a request might be routed to a node that lacks the required Key-Value (KV) cache prefix in its local memory, even though another serving node has already computed and cached it.
+
+**Global Prefix Caching** solves this by allowing nodes to share KV cache blocks globally. When a local cache miss occurs, the node queries a centralized directory to locate the prefix on remote peers. If found, it fetches the blocks directly over the network, bypassing redundant prefill computation on the TPU.
+
+---
+
+## 2. Architecture & Component Layering
+
+The architecture cleanly decouples the logical directory management (control plane) from physical buffer allocation and network transfers (data plane). All control plane interactions are driven by gRPC, while data plane transfers utilize optimized local copies (PJRT) or raw TCP socket streams (PUSH mode).
+
+```mermaid
+graph TB
+ subgraph GLOBAL_ORCH ["GLOBAL ORCHESTRATOR LAYER"]
+ Registry["Global Registry
(gRPC GlobalRegistryService)"]:::control
+ Registry ~~~ Orchestrator
+ Orchestrator["RaidenOrchestrator
(gRPC OrchestratorService)"]:::control
+ end
+
+ subgraph ENGINE_0 ["ENGINE 0"]
+ direction TB
+ subgraph E0_CP ["ENGINE 0 CONTROL PLANE (inside vLLM connector scheduler)"]
+ subgraph E0_Ctrl ["vLLM connector scheduler"]
+ subgraph Store0_box ["KVCacheStore"]
+ Store0_LRU["Local LRU Directory"]:::control
+ Ctrl0["RaidenController
(gRPC RaidenControllerService)"]:::control
+ end
+ end
+ end
+ subgraph E0_DP ["ENGINE 0 DATA PLANE (vLLM connector worker)"]
+ subgraph E0_W0 ["vLLM connector worker 0"]
+ subgraph Daemon0_0_box ["WorkerServiceServer (gRPC)"]
+ Daemon0_0["gRPC Handlers"]:::control
+ Manager0_0["KVCacheManager
(DRAM Pool)"]:::control
+ end
+ end
+ subgraph E0_W1 ["vLLM connector worker 1"]
+ subgraph Daemon0_1_box ["WorkerServiceServer (gRPC)"]
+ Daemon0_1["gRPC Handlers"]:::control
+ Manager0_1["KVCacheManager
(DRAM Pool)"]:::control
+ end
+ end
+ end
+ E0_CP ===> E0_DP
+ end
+
+ subgraph ENGINE_1 ["ENGINE 1"]
+ direction TB
+ subgraph E1_CP ["ENGINE 1 CONTROL PLANE (inside vLLM connector scheduler)"]
+ subgraph E1_Ctrl ["vLLM connector scheduler"]
+ subgraph Store1_box ["KVCacheStore"]
+ Store1_LRU["Local LRU Directory"]:::control
+ Ctrl1["RaidenController
(gRPC RaidenControllerService)"]:::control
+ end
+ end
+ end
+ subgraph E1_DP ["ENGINE 1 DATA PLANE (vLLM connector worker)"]
+ subgraph E1_W0 ["vLLM connector worker 0"]
+ subgraph Daemon1_0_box ["WorkerServiceServer (gRPC)"]
+ Daemon1_0["gRPC Handlers"]:::control
+ Manager1_0["KVCacheManager
(DRAM Pool)"]:::control
+ end
+ end
+ subgraph E1_W1 ["vLLM connector worker 1"]
+ subgraph Daemon1_1_box ["WorkerServiceServer (gRPC)"]
+ Daemon1_1["gRPC Handlers"]:::control
+ Manager1_1["KVCacheManager
(DRAM Pool)"]:::control
+ end
+ end
+ end
+ E1_CP ===> E1_DP
+ end
+
+ %% Global Connections
+ Store0_LRU -.->|gRPC| Registry
+ Store1_LRU -.->|gRPC| Registry
+ Ctrl0 -.->|gRPC| Orchestrator
+ Ctrl1 -.->|gRPC| Orchestrator
+
+ %% Ownership: RaidenController is owned and managed by KVCacheStore
+
+ %% Local Control Plane to Data Plane gRPC Connections
+ Ctrl0 ===>|gRPC (WorkerServiceClient)| Daemon0_0
+ Ctrl0 ===>|gRPC (WorkerServiceClient)| Daemon0_1
+ Ctrl1 ===>|gRPC (WorkerServiceClient)| Daemon1_0
+ Ctrl1 ===>|gRPC (WorkerServiceClient)| Daemon1_1
+
+ %% Cross-Engine Data Connection
+ Manager0_1 <===>|Socket TCP (PUSH)| Manager1_0
+
+ %% Cross-Engine Control Plane Negotiation
+ Ctrl0 <===>|gRPC ReadRemote| Ctrl1
+
+ classDef global fill:#f5f5f5,stroke:#333,stroke-width:1px;
+ classDef control fill:#e1f5fe,stroke:#0288d1,stroke-width:1px;
+ classDef data fill:#efebe9,stroke:#5d4037,stroke-width:1px;
+
+ style E0_W0 fill:#fff,stroke:#7f8c8d,stroke-dasharray: 5 5;
+ style E0_W1 fill:#fff,stroke:#7f8c8d,stroke-dasharray: 5 5;
+ style E1_W0 fill:#fff,stroke:#7f8c8d,stroke-dasharray: 5 5;
+ style E1_W1 fill:#fff,stroke:#7f8c8d,stroke-dasharray: 5 5;
+
+ style Daemon0_0_box fill:#e1f5fe,stroke:#0288d1,stroke-width:3px;
+ style Daemon0_1_box fill:#e1f5fe,stroke:#0288d1,stroke-width:3px;
+ style Daemon1_0_box fill:#e1f5fe,stroke:#0288d1,stroke-width:3px;
+ style Daemon1_1_box fill:#e1f5fe,stroke:#0288d1,stroke-width:3px;
+
+ style Store0_box fill:#e1f5fe,stroke:#0288d1,stroke-width:3px;
+ style Store1_box fill:#e1f5fe,stroke:#0288d1,stroke-width:3px;
+```
+
+### Component Details
+
+1. **`KVCacheStore` (Logical Directory)**:
+ * **Layer**: Engine Control Plane (inside vLLM connector scheduler) (Engine 0 / Engine 1).
+ * **Role**: Manages the logical metadata of KV cache blocks. It uses an internal `LruCache` to map block hashes (strings) to `RaidenBlockID` descriptors. It tracks block status (e.g., `REMOTE`, `HOST`, `HBM`) and refcounted pins to protect active blocks from eviction.
+ * **Communication**: Invokes methods on `RaidenController` asynchronously, obtaining operation handles. It tracks active operations in internal maps and polls them for completion in a background `PollerLoop` thread. Updates the `GlobalRegistry` via an internal thread pool upon completed `Save` or `ReadRemote` operations.
+2. **`RaidenController` (Control Plane Coordinator)**:
+ * **Layer**: Engine Control Plane (inside vLLM connector scheduler) (Engine 0 / Engine 1).
+ * **Role**: As a part of `KVCacheStore`, it exposes the `RaidenControllerService` gRPC endpoint. Coordinates local worker daemons and negotiates remote reads with other engines.
+ * **Communication**:
+ * **Orchestrator**: Resolves logical `RaidenId`s to controller service addresses via `OrchestratorService` gRPC. The controller can also execute some proactive operations (e.g., prefetch remote KV) sent from Orchestrator, but this feature hasn't been done yet.
+ * **Remote Controllers**: Negotiates and triggers remote reads using the `RaidenControllerService::ReadRemote` gRPC.
+ * **Local Workers**: Dispatches buffer allocation and data transfer operations to worker daemons using `WorkerService` gRPC.
+3. **`WorkerServiceServer` (Data Plane Worker Daemon)**:
+ * **Layer**: Engine Data Plane (vLLM connector worker) (Engine 0 / Engine 1).
+ * **Role**: A gRPC daemon (`WorkerServiceServer`) running on each worker process (one per NUMA node in JAX). It listens for physical buffer commands from the local controller and drives the execution engine.
+ * **Communication**: Exposes `CreateBuffers`, `DeleteBuffers`, and `TransferBuffers` RPCs. It executes memory copies via PJRT or streams data directly to peer sockets.
+4. **`KVCacheManager` (Data Plane Execution)**:
+ * **Layer**: Engine Data Plane (vLLM connector worker) (Engine 0 / Engine 1).
+ * **Role**: `KVCacheManager` performs the actual data transfers: Host-to-Host over sockets, and Host-to-Device (H2D) / Device-to-Host (D2H) via PJRT.
+ * **Communication**: Streams data directly to remote peer workers over TCP sockets during network transfers (PUSH mode).
+5. **`GlobalRegistry`**:
+ * **Layer**: Global Orchestrator Layer.
+ * **Role**: A centralized gRPC service (`GlobalRegistryService`) that maps prefix hashes to `KVBlockMetadata` (which includes the owning `RaidenId` and its physical block ID on that host). It supports multiple owners and returns them in a round-robin fashion for load balancing. The inference request router can query this global registry for KV cache slot information to make request routing decisions.
+ * **Fault Tolerance**: Supports secondary (slave) registry nodes that replicate prefix mappings from the primary registry. If the primary registry fails, secondary nodes can serve lookups, avoiding a single point of failure (SPOF).
+6. **`RaidenOrchestrator`**:
+ * **Layer**: Global Orchestrator Layer.
+ * **Role**: A lightweight gRPC service (`OrchestratorService`) that maps logical `RaidenId`s to physical IP:port controller addresses. It acts as the routing table for peer discovery. Please also note, Orchestrator is capable of triggering proactive KV cache transfer operations (e.g., prefetch) without going through the control plane of inference engine.
+
+---
+
+## 3. Python APIs & Timeline Workflows [PENDING]
+
+### Python API Surface
+
+The bindings are exposed in `google3.third_party.tpu_raiden.tpu_raiden.api.jax.kv_cache_store` (Note: JAX version):
+
+#### `BlockStatus` (Enum)
+Represents the physical location and state of a cached block:
+* `INIT = 0`: Empty / Unallocated block.
+* `REMOTE = 1`: Discovered on a remote peer node (not local).
+* `HBM = 2`: Allocated and pinned solely in TPU HBM.
+* `HOST = 3`: Allocated in local Host DRAM, but not in device HBM (eligible for LRU eviction).
+* `HOST_AND_HBM = 4`: Synced in both Host DRAM and TPU HBM.
+
+#### `RaidenBlockID` (Class)
+Tracks physical coordinates and state for a cached block:
+* `property raiden_id: RaidenId` — The owning engine identifier.
+* `property host_block_id: int` — The host DRAM block offset index (local or remote peer).
+* `property device_block_id: int` — The local TPU HBM device block ID.
+* `property status: BlockStatus` — The active memory location status of the block.
+
+#### `KVCacheStore` (Class)
+
+* `__init__(capacity: int, global_registry_address: str = "", raiden_id: RaidenId = None, num_shards: int = 0, shard_size_bytes: int = 0, raiden_controller_port: int = 0, raiden_orchestrator_address: str = "")`
+ * **Description**: Initializes the `KVCacheStore` logical directory, starts the background poller thread, and initializes the local controller.
+ * **Inputs**:
+ * `capacity`: Total logical blocks this directory can manage.
+ * `global_registry_address`: GPRC endpoint for the centralized directory.
+ * `raiden_id`: The local engine's work unit identifier.
+ * `num_shards`: Number of local TPU chips.
+ * `shard_size_bytes`: Byte size of a single shard block.
+ * `raiden_controller_port`: Listening port for the background controller.
+ * `raiden_orchestrator_address`: Orchestrator gRPC address for resolving peer IPs.
+ * **Returns**: None.
+
+* `lookup(block_hashes: list[bytes], enable_global: bool = False) -> list[tuple[bytes, RaidenBlockID]]`
+ * **Description**: Checks the directory for block hashes. Performs a global registry lookup for any misses if `enable_global` is True.
+ * **Inputs**:
+ * `block_hashes`: List of binary prefix hashes to locate.
+ * `enable_global`: Fall back to querying the global registry on a local cache miss.
+ * **Returns**: List of `(hash, RaidenBlockID)` tuples. Halts and returns immediately upon the first complete miss (neither local nor global).
+
+* `insert(block_hashes: list[bytes], slices: list[RaidenBlockID], on_host: bool) -> tuple[bool, list[tuple[bytes, RaidenBlockID]]]`
+ * **Description**: Caches block metadata manually. Evicts older unpinned blocks if capacity is exceeded.
+ * **Inputs**:
+ * `block_hashes`: Hashes to insert.
+ * `slices`: Associated `RaidenBlockID` block descriptors.
+ * `on_host`: True if the backing buffers reside in Host DRAM.
+ * **Returns**: A tuple of `(all_inserted, evicted_entries)`, where `all_inserted` is True if all keys are new, and `evicted_entries` contains evicted `(hash, RaidenBlockID)` items.
+
+* `insert_and_lock(block_hashes: list[bytes], slices: list[RaidenBlockID], on_host: bool) -> bool`
+ * **Description**: Inserts block hashes and pins them in a single atomic transaction. Prevents eviction while in use by active attention prefill queries.
+ * **Inputs**:
+ * `block_hashes`: Hashes to insert and lock.
+ * `slices`: Associated `RaidenBlockID` descriptors.
+ * `on_host`: True if located in Host DRAM.
+ * **Returns**: True if the entire batch succeeded (all items locked, or inserted and locked).
+
+* `release_and_delete(block_hashes: list[bytes]) -> int`
+ * **Description**: Releases/unpins block hashes. If a block's status is `REMOTE` and its pin count drops to 0, it is deleted from the directory.
+ * **Inputs**:
+ * `block_hashes`: Hashes to unpin.
+ * **Returns**: The number of remote blocks deleted.
+
+* `pin(block_hashes: list[bytes]) -> bool`
+ * **Description**: Manually pins existing block hashes, protecting them from LRU eviction.
+ * **Inputs**:
+ * `block_hashes`: Hashes to pin.
+ * **Returns**: True if all hashes exist in the directory and were successfully pinned.
+
+* `release(block_hashes: list[bytes]) -> None`
+ * **Description**: Releases previously pinned block hashes, making them eligible for LRU eviction.
+ * **Inputs**:
+ * `block_hashes`: Hashes to release.
+ * **Returns**: None.
+
+* `save(block_hashes: list[bytes]) -> bool`
+ * **Description**: Launches an asynchronous Device-to-Host (D2H) copy from TPU HBM to local Host DRAM. Blocks must be locked before calling `save`.
+ * **Inputs**:
+ * `block_hashes`: Pinned HBM block hashes to save.
+ * **Returns**: True if the async copy was successfully dispatched.
+
+* `load(block_hashes: list[bytes], device_block_ids: list[int]) -> bool`
+ * **Description**: Launches an asynchronous Host-to-Device (H2D) copy from local Host DRAM to TPU HBM. Blocks must be pinned before calling `load`.
+ * **Inputs**:
+ * `block_hashes`: Pinned DRAM block hashes to load.
+ * `device_block_ids`: Destination TPU block IDs.
+ * **Returns**: True if the async copy was successfully dispatched.
+
+* `poll_save_status() -> tuple[list[bytes], list[bytes], list[bytes]]`
+ * **Description**: Polls active async `save` operations.
+ * **Returns**: A tuple of `(done_hashes, failed_hashes, pending_hashes)`. Completed items are transitioned to `HOST_AND_HBM` status in the LRU, and registered in the `GlobalRegistry`.
+
+* `poll_load_status() -> tuple[list[bytes], list[bytes], list[bytes]]`
+ * **Description**: Polls active async `load` operations.
+ * **Returns**: A tuple of `(done_hashes, failed_hashes, pending_hashes)`. Completed items are transitioned to `HOST_AND_HBM` status in the LRU.
+
+* `read_remote(block_hashes: list[bytes]) -> bool`
+ * **Description**: Launches an asynchronous network copy to read remote blocks from peer nodes into local host DRAM.
+ * **Inputs**:
+ * `block_hashes`: Remote block hashes (status `REMOTE`) to fetch.
+ * **Returns**: True if the async network copy was successfully dispatched.
+
+* `poll_remote_read_status() -> tuple[list[bytes], list[bytes], list[bytes]]`
+ * **Description**: Polls active async remote reads.
+ * **Returns**: A tuple of `(done_hashes, failed_hashes, pending_hashes)`. Completed items are transitioned to `HOST` status in the LRU (associated with their allocated local host block IDs) and registered in the `GlobalRegistry`.
+
+### Global Orchestration gRPC service
+
+> [!NOTE]
+> **Coming Soon**: This service is currently under development and design validation. Detailed gRPC schemas and service interfaces will be added here once finalized.
+
+---
+
+### Timeline Workflows
+
+#### A. Lookup (Global Fallback)
+
+Checks local cache and automatically queries the global registry on a local miss.
+
+```mermaid
+sequenceDiagram
+ participant Client as vLLM connector Scheduler
+ participant Store as KVCacheStore (Local)
+ participant Registry as Global Registry
+
+ Client->>Store: lookup(hashes, enable_global=True)
+ Store->>Store: Check local LRU Cache
+ Note over Store: Found H1 (HOST). H2 is a Miss.
+
+ rect rgb(240, 240, 240)
+ Note over Store: Query Global Registry for H2
+ Store->>Registry: Lookup([H2])
+ Registry-->>Store: Returns [Metadata(PeerRaidenId, peer_block_id)]
+ end
+
+ Store->>Store: Create remote RaidenBlockID for H2
(status=REMOTE, host_block_id=peer_block_id)
+ Store-->>Client: Returns [(H1, LocalBlockID), (H2, RemoteBlockID)]
+```
+
+#### B. ReadRemote (Network Copy)
+
+Pulls a remote block from a peer node into local DRAM staging memory.
+
+```mermaid
+sequenceDiagram
+ participant Store as KVCacheStore (Local)
+ participant CtrlD as Local Controller (Dest)
+ participant Orch as RaidenOrchestrator
+ participant CtrlS as Peer Controller (Source)
+ participant DaemonS as Peer Worker (Source)
+ participant DaemonD as Local Worker (Dest)
+ participant Registry as Global Registry
+
+ Store->>Store: 1. Verify status is REMOTE & pinned
+ Store->>Store: 2. Allocate local Host Block IDs (e.g. 5)
+ Store->>CtrlD: 3. ReadRemote(PeerRaidenId, src_host_ids, dst_host_ids)
+
+ CtrlD->>Orch: 4. Resolve address of PeerRaidenId
+ Orch-->>CtrlD: Peer Controller IP:port
+
+ CtrlD->>CtrlS: 5. gRPC ReadRemote(block_hashes, dst_host_block_ids, dest_worker_endpoints)
+
+ par Control Plane Dispatches to Workers
+ CtrlS->>CtrlS: 6a. Verify block_hashes with its LRU map
+ CtrlS->>DaemonS: 6b. TransferBuffer(DRAM -> remote DRAM)
+ end
+
+ DaemonS->>DaemonD: 7. Data Stream
+ Note over DaemonS, DaemonD: Direct Socket Transfer (PUSH)
+
+ DaemonS-->>CtrlS: 8. Transfer Completed (gRPC Response)
+ CtrlS-->>CtrlD: 9. ReadRemote gRPC Response (Success)
+
+ DaemonD->>DaemonD: 10. Local H2D (Optional)
+
+ Note over Store: Store Poller thread detects future completion
+ Store->>Store: 11. Update LRU (H2 status -> HOST or HOST_AND_HBM, host_block_id = 5)
+ Store->>Registry: 12. Async Write-Through Register(H2, host_block_id=5)
+```
+
+#### C. Load (H2D Copy)
+
+Copies blocks locally from host DRAM to TPU HBM before execution.
+
+```mermaid
+sequenceDiagram
+ participant Store as KVCacheStore
+ participant Ctrl as Local Controller
+ participant Daemon as WorkerServiceServer (Worker)
+ participant Manager as KVCacheManager (Worker)
+
+ Store->>Store: 1. Verify status is HOST & pinned
+ Store->>Ctrl: 2. TransferBuffers(DRAM -> HBM, src_host_block_ids, dst_hbm_block_ids)
+ Ctrl->>Daemon: 3. gRPC TransferBuffers(DRAM -> HBM, src_offsets, dst_offsets)
+ Daemon->>Manager: 4. H2d(src_offsets, dst_offsets)
+ Note over Manager: Async DMA Copy (DRAM -> HBM)
+ Manager-->>Daemon: 5. Copy Complete (Await)
+ Daemon-->>Ctrl: 6. gRPC TransferBuffers Response
+
+ Note over Store: Store Poller thread detects future completion
+ Store->>Store: 7. Update LRU (status -> HOST_AND_HBM, store hbm_block_id)
+```
+
+#### D. Save (D2H Copy & Registry Registration)
+
+Offloads blocks from TPU HBM to Host DRAM and registers them globally.
+
+```mermaid
+sequenceDiagram
+ participant Store as KVCacheStore
+ participant Ctrl as Local Controller
+ participant Daemon as WorkerServiceServer (Worker)
+ participant Manager as KVCacheManager (Worker)
+ participant Registry as Global Registry
+
+ Store->>Store: 1. Verify status is HBM & pinned
+ Store->>Store: 2. Allocate local Host Block IDs (e.g. 7)
+ Store->>Ctrl: 3. TransferBuffers(HBM -> DRAM, src_device_block_ids, dst_host_block_ids)
+ Ctrl->>Daemon: 4. gRPC TransferBuffers(HBM -> DRAM, src_offsets, dst_offsets)
+ Daemon->>Manager: 5. D2h(src_offsets, dst_offsets)
+ Note over Manager: Async DMA Copy (HBM -> DRAM)
+ Manager-->>Daemon: 6. Copy Complete (Await)
+ Daemon-->>Ctrl: 7. gRPC TransferBuffers Response
+
+ Note over Store: Store Poller thread detects future completion
+ Store->>Store: 8. Update LRU (status -> HOST_AND_HBM, store host_block_id)
+ Store->>Registry: 9. Async Write-Through Register(hashes, host_block_ids)
+```
+
+#### E. Evict (DRAM Cleanup & Registry Unregistration)
+
+Frees Host DRAM space by unlocking blocks and removing global directory entries.
+
+```mermaid
+sequenceDiagram
+ participant Store as KVCacheStore
+ participant Registry as Global Registry
+
+ Store->>Store: 1. Verify pin_count == 0 & status is HOST/HOST_AND_HBM
+ Store->>Store: 2. Immediate Status Transition:
HOST -> INIT, HOST_AND_HBM -> HBM
+ Store->>Registry: 3. Async Unregister(hashes)
+
+ Store->>Store: 4. DeallocateBlockIds(host_block_ids)
(Marks block offsets as free in LogicalBlockManager)
+ Store->>Store: 5. Final Directory Cleanup:
Erase INIT blocks, clear host_block_id to -1 for HBM blocks
+```
+
+
+
diff --git a/tpu_raiden/api/jax/kv_cache_store_e2e_test.py b/tpu_raiden/api/jax/kv_cache_store_e2e_test.py
index db7d99b8..f78fd027 100644
--- a/tpu_raiden/api/jax/kv_cache_store_e2e_test.py
+++ b/tpu_raiden/api/jax/kv_cache_store_e2e_test.py
@@ -180,6 +180,7 @@ def setUpClass(cls):
def setUp(self):
super().setUp()
+ os.environ["RAIDEN_LOCAL_IP"] = "127.0.0.1"
start_servers()
try:
self.devices = jax.devices("tpu")
@@ -239,6 +240,7 @@ def _run_e2e_test(self, enable_multi_numa: bool):
)
os.environ["ENABLE_MULTI_NUMA"] = "1" if enable_multi_numa else "0"
+ os.environ["RAIDEN_LOCAL_IP"] = "127.0.0.1"
tpu_sharding = self.setup_shardings()
num_blocks = 2
@@ -272,7 +274,7 @@ def _run_e2e_test(self, enable_multi_numa: bool):
# 4. Create KVCacheManager (Worker)
manager = kv_cache_manager.KVCacheManager(
kv_caches=[tpu_cache],
- local_control_port=0,
+ local_control_port=-1,
max_blocks=num_blocks,
num_slots=2,
unsafe_skip_buffer_lock=self.skip_lock,
@@ -404,6 +406,7 @@ def _run_remote_read_e2e_test(
)
os.environ["ENABLE_MULTI_NUMA"] = "1" if enable_multi_numa else "0"
+ os.environ["RAIDEN_LOCAL_IP"] = "127.0.0.1"
if len(self.devices) < 1:
self.skipTest(
@@ -451,7 +454,7 @@ def _run_remote_read_e2e_test(
)
manager_a = kv_cache_manager.KVCacheManager(
kv_caches=[tpu_cache_a],
- local_control_port=0,
+ local_control_port=-1,
max_blocks=num_blocks,
num_slots=2,
unsafe_skip_buffer_lock=self.skip_lock,
@@ -475,7 +478,7 @@ def _run_remote_read_e2e_test(
)
manager_b = kv_cache_manager.KVCacheManager(
kv_caches=[tpu_cache_b],
- local_control_port=0,
+ local_control_port=-1,
max_blocks=num_blocks,
num_slots=2,
unsafe_skip_buffer_lock=self.skip_lock,
@@ -669,7 +672,7 @@ def test_remote_read_e2e_source_missing_block_fails(self):
if len(self.devices) < 1:
self.skipTest("Requires at least 1 device")
os.environ["ENABLE_MULTI_NUMA"] = "0"
-
+ os.environ["RAIDEN_LOCAL_IP"] = "127.0.0.1"
sharding = self.setup_sharding_for_devices(self.devices)
num_blocks = 2
shape = (num_blocks, 128, 8, 8, 128)
@@ -701,7 +704,7 @@ def test_remote_read_e2e_source_missing_block_fails(self):
)
manager_a = kv_cache_manager.KVCacheManager(
kv_caches=[tpu_cache_a],
- local_control_port=0,
+ local_control_port=-1,
max_blocks=num_blocks,
num_slots=2,
unsafe_skip_buffer_lock=self.skip_lock,
@@ -723,7 +726,7 @@ def test_remote_read_e2e_source_missing_block_fails(self):
)
manager_b = kv_cache_manager.KVCacheManager(
kv_caches=[tpu_cache_b],
- local_control_port=0,
+ local_control_port=-1,
max_blocks=num_blocks,
num_slots=2,
unsafe_skip_buffer_lock=self.skip_lock,
@@ -768,7 +771,7 @@ def test_remote_read_e2e_source_wrong_status_fails(self):
if len(self.devices) < 1:
self.skipTest("Requires at least 1 device")
os.environ["ENABLE_MULTI_NUMA"] = "0"
-
+ os.environ["RAIDEN_LOCAL_IP"] = "127.0.0.1"
sharding = self.setup_sharding_for_devices(self.devices)
num_blocks = 2
shape = (num_blocks, 128, 8, 8, 128)
@@ -798,7 +801,7 @@ def test_remote_read_e2e_source_wrong_status_fails(self):
)
manager_a = kv_cache_manager.KVCacheManager(
kv_caches=[tpu_cache_a],
- local_control_port=0,
+ local_control_port=-1,
max_blocks=num_blocks,
num_slots=2,
unsafe_skip_buffer_lock=self.skip_lock,
@@ -835,7 +838,7 @@ def test_remote_read_e2e_source_wrong_status_fails(self):
)
manager_b = kv_cache_manager.KVCacheManager(
kv_caches=[tpu_cache_b],
- local_control_port=0,
+ local_control_port=-1,
max_blocks=num_blocks,
num_slots=2,
unsafe_skip_buffer_lock=self.skip_lock,
diff --git a/tpu_raiden/core/BUILD b/tpu_raiden/core/BUILD
index 597d8c79..6763459b 100644
--- a/tpu_raiden/core/BUILD
+++ b/tpu_raiden/core/BUILD
@@ -730,3 +730,15 @@ cc_test(
"@com_google_googletest//:gtest_main",
],
)
+
+cc_test(
+ name = "kv_manager_holder_test",
+ srcs = ["kv_manager_holder_test.cc"],
+ deps = [
+ ":kv_manager_holder",
+ ":raiden_transfer_endpoint",
+ ":raw_transfer_core",
+ "//tpu_raiden/core/controller:test_util",
+ "@com_google_googletest//:gtest_main",
+ ],
+)
diff --git a/tpu_raiden/core/controller/BUILD b/tpu_raiden/core/controller/BUILD
index 8893df51..026339f5 100644
--- a/tpu_raiden/core/controller/BUILD
+++ b/tpu_raiden/core/controller/BUILD
@@ -177,6 +177,7 @@ cc_library(
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/synchronization",
+ "@com_google_absl//absl/time",
"@com_google_absl//absl/types:span",
"@xla//xla/tsl/concurrency:future",
],
@@ -258,6 +259,7 @@ cc_test(
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
+ "@com_google_absl//absl/time",
"@com_google_absl//absl/types:span",
"@com_google_googletest//:gtest",
"@com_google_googletest//:gtest_main",
diff --git a/tpu_raiden/core/controller/controller_service.cc b/tpu_raiden/core/controller/controller_service.cc
index 8c68d6e2..4506b8fd 100644
--- a/tpu_raiden/core/controller/controller_service.cc
+++ b/tpu_raiden/core/controller/controller_service.cc
@@ -147,12 +147,7 @@ grpc::Status RaidenControllerServiceImpl::ReadRemote(
}
}
- tsl::Future<> future = (*cb)(src_buffers, dst_buffers);
- absl::Status status = future.Await();
- if (!status.ok()) {
- return grpc::Status(grpc::StatusCode::INTERNAL,
- std::string(status.message()));
- }
+ return grpc::Status::OK;
return grpc::Status::OK;
}
@@ -171,6 +166,138 @@ std::shared_ptr RaidenControllerServiceImpl::worker_registry()
return worker_registry_;
}
+RaidenControllerServiceImpl::~RaidenControllerServiceImpl() {
+ {
+ absl::MutexLock lock(&mutex_);
+ sweeper_running_ = false;
+ sweeper_cv_.SignalAll();
+ }
+ if (sweeper_thread_ && sweeper_thread_->joinable()) {
+ sweeper_thread_->join();
+ }
+}
+
+void RaidenControllerServiceImpl::StartSweeperIfNecessary() {
+ if (!sweeper_running_) {
+ sweeper_running_ = true;
+ sweeper_thread_ =
+ std::make_unique([this]() { SweeperLoop(); });
+ }
+}
+
+void RaidenControllerServiceImpl::SweeperLoop() {
+ absl::MutexLock lock(&mutex_);
+ while (sweeper_running_) {
+ if (active_pins_.empty()) {
+ sweeper_cv_.Wait(&mutex_);
+ continue;
+ }
+ absl::Time now = absl::Now();
+ absl::Time next_wakeup = absl::InfiniteFuture();
+ for (auto it = active_pins_.begin(); it != active_pins_.end();) {
+ if (it->expiration <= now) {
+ if (unpin_cb_) {
+ (*unpin_cb_)(it->block_hashes);
+ }
+ it = active_pins_.erase(it);
+ } else {
+ if (it->expiration < next_wakeup) {
+ next_wakeup = it->expiration;
+ }
+ ++it;
+ }
+ }
+ if (next_wakeup != absl::InfiniteFuture() && sweeper_running_) {
+ sweeper_cv_.WaitWithDeadline(&mutex_, next_wakeup);
+ }
+ }
+}
+
+grpc::Status RaidenControllerServiceImpl::PinRemoteBlocks(
+ grpc::ServerContext* context,
+ const ::tpu_raiden::proto::PinRemoteBlocksRequest* request,
+ ::tpu_raiden::proto::PinRemoteBlocksResponse* response) {
+ std::shared_ptr validate_cb;
+ std::shared_ptr registry;
+ {
+ absl::MutexLock lock(&mutex_);
+ validate_cb = validate_and_pin_cb_;
+ registry = worker_registry_;
+ }
+ if (!validate_cb) {
+ return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION,
+ "ReadRemote hooks not set.");
+ }
+ std::vector block_hashes(request->block_hashes().begin(),
+ request->block_hashes().end());
+ absl::StatusOr> ids_or = (*validate_cb)(block_hashes);
+ if (!ids_or.ok()) {
+ return grpc::Status(static_cast(ids_or.status().code()),
+ std::string(ids_or.status().message()));
+ }
+
+ if (registry) {
+ for (const auto& reg : registry->GetRegisteredWorkers()) {
+ auto* proto_endpoint_group = response->add_src_worker_endpoints();
+ proto_endpoint_group->set_node_id(reg.node_id);
+ proto_endpoint_group->set_worker_id(reg.worker_id);
+ for (const auto& ep : reg.raiden_transfer_endpoints) {
+ auto* ep_proto = proto_endpoint_group->add_endpoints();
+ ep_proto->set_endpoint(ep.endpoint);
+ for (int64_t shard : ep.shards) {
+ ep_proto->add_shards(shard);
+ }
+ }
+ }
+ }
+
+ {
+ absl::MutexLock lock(&mutex_);
+ PinLease lease;
+ lease.block_hashes = block_hashes;
+ lease.block_ids = *ids_or;
+ if (request->ttl_seconds() > 0) {
+ lease.expiration = absl::Now() + absl::Seconds(request->ttl_seconds());
+ } else {
+ lease.expiration = absl::Now() + absl::Seconds(60);
+ }
+ active_pins_.push_back(std::move(lease));
+ StartSweeperIfNecessary();
+ sweeper_cv_.Signal();
+ }
+
+ response->mutable_src_host_block_ids()->Assign(ids_or->begin(),
+ ids_or->end());
+ return grpc::Status::OK;
+}
+
+grpc::Status RaidenControllerServiceImpl::UnpinRemoteBlocks(
+ grpc::ServerContext* context,
+ const ::tpu_raiden::proto::UnpinRemoteBlocksRequest* request,
+ ::tpu_raiden::proto::UnpinRemoteBlocksResponse* response) {
+ std::shared_ptr cb;
+ std::vector hashes_to_unpin;
+ {
+ absl::MutexLock lock(&mutex_);
+ cb = unpin_cb_;
+ std::vector request_ids(request->src_host_block_ids().begin(),
+ request->src_host_block_ids().end());
+ for (auto it = active_pins_.begin(); it != active_pins_.end(); ++it) {
+ if (it->block_ids == request_ids) {
+ hashes_to_unpin = std::move(it->block_hashes);
+ active_pins_.erase(it);
+ break;
+ }
+ }
+ }
+
+ if (cb && !hashes_to_unpin.empty()) {
+ (*cb)(hashes_to_unpin);
+ }
+
+ return grpc::Status::OK;
+}
+
} // namespace controller
} // namespace core
} // namespace tpu_raiden
diff --git a/tpu_raiden/core/controller/controller_service.h b/tpu_raiden/core/controller/controller_service.h
index 447deb43..cdc6d15e 100644
--- a/tpu_raiden/core/controller/controller_service.h
+++ b/tpu_raiden/core/controller/controller_service.h
@@ -17,8 +17,10 @@
#include
+#include
#include
#include
+#include
#include
#include "absl/base/thread_annotations.h"
@@ -27,6 +29,7 @@
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/mutex.h"
+#include "absl/time/time.h"
#include "absl/types/span.h"
#include "grpcpp/server_context.h"
#include "grpcpp/support/status.h"
@@ -47,7 +50,7 @@ class RaidenControllerServiceImpl final
public:
explicit RaidenControllerServiceImpl(
std::shared_ptr worker_registry = nullptr);
- ~RaidenControllerServiceImpl() override = default;
+ ~RaidenControllerServiceImpl() override;
// Disallow copy and assign
RaidenControllerServiceImpl(const RaidenControllerServiceImpl&) = delete;
@@ -65,6 +68,16 @@ class RaidenControllerServiceImpl final
const ::tpu_raiden::proto::ReadRemoteRequest* request,
::tpu_raiden::proto::ReadRemoteResponse* response) override;
+ grpc::Status PinRemoteBlocks(
+ grpc::ServerContext* context,
+ const ::tpu_raiden::proto::PinRemoteBlocksRequest* request,
+ ::tpu_raiden::proto::PinRemoteBlocksResponse* response) override;
+
+ grpc::Status UnpinRemoteBlocks(
+ grpc::ServerContext* context,
+ const ::tpu_raiden::proto::UnpinRemoteBlocksRequest* request,
+ ::tpu_raiden::proto::UnpinRemoteBlocksResponse* response) override;
+
using TransferBuffersCallback = absl::AnyInvocable(
absl::Span src_buffers,
absl::Span dst_buffers) const>;
@@ -107,6 +120,19 @@ class RaidenControllerServiceImpl final
std::shared_ptr validate_and_pin_cb_
ABSL_GUARDED_BY(mutex_);
std::shared_ptr unpin_cb_ ABSL_GUARDED_BY(mutex_);
+
+ struct PinLease {
+ std::vector block_hashes;
+ absl::Time expiration;
+ std::vector block_ids;
+ };
+ std::vector active_pins_ ABSL_GUARDED_BY(mutex_);
+ bool sweeper_running_ ABSL_GUARDED_BY(mutex_) = false;
+ std::unique_ptr sweeper_thread_;
+ absl::CondVar sweeper_cv_;
+
+ void StartSweeperIfNecessary() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+ void SweeperLoop();
};
} // namespace controller
diff --git a/tpu_raiden/core/controller/raiden_controller.cc b/tpu_raiden/core/controller/raiden_controller.cc
index 671169a2..f554dc99 100644
--- a/tpu_raiden/core/controller/raiden_controller.cc
+++ b/tpu_raiden/core/controller/raiden_controller.cc
@@ -23,6 +23,7 @@
#include
#include
#include
+#include
#include
#include
@@ -377,6 +378,7 @@ absl::Status RaidenController::DeallocateBuffers(
absl::StatusOr
RaidenController::BuildTransferBuffersRequest(
absl::Span src_buffers, absl::Span dst_buffers,
+ absl::Span staging_host_buffers,
absl::Span copy_sizes) {
if (src_buffers.empty() || src_buffers.size() != dst_buffers.size()) {
return absl::InvalidArgumentError(
@@ -416,6 +418,15 @@ RaidenController::BuildTransferBuffersRequest(
for (int64_t size : copy_sizes) {
transfer->add_copy_sizes(size);
}
+ for (const auto& buf : staging_host_buffers) {
+ if (buf.index() < 0) {
+ return absl::InvalidArgumentError(absl::StrCat(
+ "Staging host buffer has invalid negative index: ", buf.index()));
+ }
+ auto* added_buf = transfer->add_staging_host_buffers();
+ *added_buf = buf.ToProto();
+ added_buf->set_index(buf.index());
+ }
return request;
}
@@ -423,9 +434,10 @@ RaidenController::BuildTransferBuffersRequest(
tsl::Future<> RaidenController::TransferBuffers(
absl::string_view worker_id, absl::Span src_buffers,
absl::Span dst_buffers,
+ absl::Span staging_host_buffers,
absl::Span copy_sizes) {
- auto request_or =
- BuildTransferBuffersRequest(src_buffers, dst_buffers, copy_sizes);
+ auto request_or = BuildTransferBuffersRequest(
+ src_buffers, dst_buffers, staging_host_buffers, copy_sizes);
if (!request_or.ok()) {
return tsl::Future<>(request_or.status());
}
@@ -447,6 +459,7 @@ tsl::Future<> RaidenController::TransferBuffers(
tsl::Future<> RaidenController::TransferBuffers(
absl::Span src_buffers, absl::Span dst_buffers,
+ absl::Span staging_host_buffers,
absl::Span copy_sizes) {
if (src_buffers.empty() || src_buffers.size() != dst_buffers.size()) {
return tsl::Future<>(absl::InvalidArgumentError(
@@ -476,6 +489,13 @@ tsl::Future<> RaidenController::TransferBuffers(
peer_node_id_to_endpoints.emplace(group.node_id, &group);
}
}
+ if (peer_node_id_to_endpoints.empty()) {
+ for (const auto& src_buf : src_buffers) {
+ for (const auto& group : src_buf.remote_worker_endpoints()) {
+ peer_node_id_to_endpoints.emplace(group.node_id, &group);
+ }
+ }
+ }
std::vector> worker_futures;
worker_futures.reserve(workers.size());
@@ -492,14 +512,32 @@ tsl::Future<> RaidenController::TransferBuffers(
// Every buffer/block is transferred by every worker (each worker owns a
// shard of every block), so there is no per-block partitioning here.
for (size_t i = 0; i < src_buffers.size(); ++i) {
- worker_src.push_back(src_buffers[i]);
+ Buffer src_buf = src_buffers[i];
+ if (!src_buf.remote_worker_endpoints().empty()) {
+ auto it = peer_node_id_to_endpoints.find(workers[w].node_id);
+ if (it == peer_node_id_to_endpoints.end()) {
+ LOG(ERROR) << "Could not match node_id: " << workers[w].node_id;
+ return tsl::Future<>(absl::FailedPreconditionError(absl::StrCat(
+ "ReadRemote: no source worker group with node_id ",
+ workers[w].node_id, " to match destination worker '",
+ workers[w].worker_id, "' (source provided ",
+ src_buf.remote_worker_endpoints().size(), " group(s))")));
+ }
+ LOG(INFO) << "Matched node_id " << workers[w].node_id
+ << " for src to endpoints[0]: "
+ << it->second->endpoints[0].endpoint;
+ src_buf.set_remote_descriptors(it->second->endpoints);
+ src_buf.set_remote_worker_endpoints({});
+ }
+ worker_src.push_back(std::move(src_buf));
Buffer dst_buf = dst_buffers[i];
if (!dst_buf.remote_worker_endpoints().empty()) {
auto it = peer_node_id_to_endpoints.find(workers[w].node_id);
if (it == peer_node_id_to_endpoints.end()) {
+ LOG(ERROR) << "Could not match node_id: " << workers[w].node_id;
return tsl::Future<>(absl::FailedPreconditionError(absl::StrCat(
- "ReadRemote: no destination worker group with node_id ",
+ "WriteRemote: no destination worker group with node_id ",
workers[w].node_id, " to match source worker '",
workers[w].worker_id, "' (destination provided ",
dst_buf.remote_worker_endpoints().size(), " group(s))")));
@@ -523,8 +561,10 @@ tsl::Future<> RaidenController::TransferBuffers(
if (worker_src.empty()) continue;
- auto req_or =
- BuildTransferBuffersRequest(worker_src, worker_dst, worker_copy_sizes);
+ // Every worker owns a shard of every block, so the (host) staging offsets
+ // are identical across workers.
+ auto req_or = BuildTransferBuffersRequest(
+ worker_src, worker_dst, staging_host_buffers, worker_copy_sizes);
if (!req_or.ok()) {
return tsl::Future<>(req_or.status());
}
@@ -593,24 +633,11 @@ tsl::Future<> RaidenController::ReadRemote(
rpc::MemoryType::MEMORY_TYPE_DRAM);
}
- // A KV block is sharded across every (destination) worker, so every block is
- // transferred by every worker. Attach the full list of destination workers'
- // endpoint groups to each dst buffer, each tagged with its node_id, so the
- // remote source controller can match each of its workers to the destination
- // peer worker with the same node_id.
- std::vector dest_worker_endpoints;
- dest_worker_endpoints.reserve(workers.size());
- for (const auto& reg : workers) {
- dest_worker_endpoints.push_back(
- {reg.node_id, reg.worker_id, reg.raiden_transfer_endpoints});
- }
-
std::vector dst_buffers;
dst_buffers.reserve(dest_host_block_ids.size());
for (size_t i = 0; i < dest_host_block_ids.size(); ++i) {
Buffer dst_buf(dest_host_block_ids[i], std::vector{},
std::nullopt, rpc::MemoryType::MEMORY_TYPE_DRAM);
- dst_buf.set_remote_worker_endpoints(dest_worker_endpoints);
dst_buffers.push_back(std::move(dst_buf));
}
@@ -655,44 +682,81 @@ tsl::Future<> RaidenController::ReadRemote(
}
}
- cproto::ReadRemoteRequest request;
- request.mutable_src_host_block_ids()->Reserve(src_host_block_ids.size());
- request.mutable_dest_host_block_ids()->Reserve(dest_host_block_ids.size());
- for (int32_t id : src_host_block_ids) {
- request.add_src_host_block_ids(id);
- }
- for (int32_t id : dest_host_block_ids) {
- request.add_dest_host_block_ids(id);
- }
+ cproto::PinRemoteBlocksRequest pin_req;
+ pin_req.set_ttl_seconds(60);
for (const auto& hash : block_hashes) {
- request.add_block_hashes(hash);
- }
- for (const auto& buf : src_buffers) {
- *request.add_src_buffers() = buf.ToProto();
- }
- for (const auto& buf : dst_buffers) {
- *request.add_dst_buffers() = buf.ToProto();
+ pin_req.add_block_hashes(hash);
}
auto [promise, future] = tsl::MakePromise<>();
auto context = std::make_shared();
- auto response = std::make_shared();
+ auto response = std::make_shared();
- stub->async()->ReadRemote(
- context.get(), &request, response.get(),
- [context, response, stub,
+ stub->async()->PinRemoteBlocks(
+ context.get(), &pin_req, response.get(),
+ [this, context, response, stub, dest_host_block_ids,
promise = std::move(promise).ToShared()](grpc::Status status) {
if (!status.ok()) {
- // Preserve the gRPC status code (grpc and absl codes share the same
- // canonical integers), so ReadRemote step-6a errors stay
- // distinguishable end-to-end: NOT_FOUND (BLOCK_HASH_NOT_FOUND) vs
- // FAILED_PRECONDITION (present but not host-resident) vs INTERNAL.
- promise->Set(absl::Status(
- static_cast(status.error_code()),
- absl::StrCat("ReadRemote RPC failed: ", status.error_message())));
- } else {
- promise->Set(absl::OkStatus());
+ promise->Set(
+ absl::Status(static_cast(status.error_code()),
+ absl::StrCat("PinRemoteBlocks RPC failed: ",
+ status.error_message())));
+ return;
}
+
+ std::vector src_endpoints;
+ for (const auto& ep : response->src_worker_endpoints()) {
+ std::vector endpoints;
+ for (const auto& r_ep : ep.endpoints()) {
+ endpoints.push_back(
+ {r_ep.endpoint(), std::vector(r_ep.shards().begin(),
+ r_ep.shards().end())});
+ }
+ src_endpoints.push_back({ep.node_id(), ep.worker_id(), endpoints});
+ }
+
+ std::vector src_buffers;
+ for (int32_t id : response->src_host_block_ids()) {
+ Buffer src_buf(id, {}, std::nullopt,
+ rpc::MemoryType::MEMORY_TYPE_DRAM);
+ src_buf.set_remote_worker_endpoints(src_endpoints);
+ src_buffers.push_back(std::move(src_buf));
+ }
+
+ std::vector staging_host_buffers;
+ std::vector dst_buffers;
+ for (int32_t id : dest_host_block_ids) {
+ staging_host_buffers.emplace_back(id, std::vector{},
+ std::nullopt,
+ rpc::MemoryType::MEMORY_TYPE_DRAM);
+ dst_buffers.emplace_back(id, std::vector{}, std::nullopt,
+ rpc::MemoryType::MEMORY_TYPE_HBM);
+ }
+
+ auto transfer_future = this->TransferBuffers(src_buffers, dst_buffers,
+ staging_host_buffers);
+ std::vector src_block_ids(
+ response->src_host_block_ids().begin(),
+ response->src_host_block_ids().end());
+
+ std::thread([transfer_future = std::move(transfer_future), promise,
+ stub, src_block_ids = std::move(src_block_ids)]() mutable {
+ absl::Status transfer_status = transfer_future.Await();
+
+ cproto::UnpinRemoteBlocksRequest unpin_req;
+ for (int32_t id : src_block_ids) {
+ unpin_req.add_src_host_block_ids(id);
+ }
+ auto unpin_ctx = std::make_shared();
+ auto unpin_resp =
+ std::make_shared();
+
+ stub->async()->UnpinRemoteBlocks(
+ unpin_ctx.get(), &unpin_req, unpin_resp.get(),
+ [unpin_ctx, unpin_resp](grpc::Status status) {});
+
+ promise->Set(transfer_status);
+ }).detach();
});
return future;
diff --git a/tpu_raiden/core/controller/raiden_controller.h b/tpu_raiden/core/controller/raiden_controller.h
index ab5a0d39..a011d927 100644
--- a/tpu_raiden/core/controller/raiden_controller.h
+++ b/tpu_raiden/core/controller/raiden_controller.h
@@ -118,16 +118,23 @@ class RaidenController {
// is performed.
absl::Status AllocateTargetBlockIds(absl::Span block_ids);
- // Targeted worker transfer
- tsl::Future<> TransferBuffers(absl::string_view worker_id,
- absl::Span src_buffers,
- absl::Span dst_buffers,
- absl::Span copy_sizes = {});
+ // Targeted worker transfer.
+ // staging_host_buffers: host DRAM staging (bridge) block offsets, required
+ // for 2-stage remote transfers (remote H2D read/write, remote D2H write);
+ // always the middle hop of the data flow. Unused for local transfers.
+ tsl::Future<> TransferBuffers(
+ absl::string_view worker_id, absl::Span src_buffers,
+ absl::Span dst_buffers,
+ absl::Span staging_host_buffers = {},
+ absl::Span copy_sizes = {});
- // Broadcast transfer to all registered workers
- tsl::Future<> TransferBuffers(absl::Span src_buffers,
- absl::Span dst_buffers,
- absl::Span copy_sizes = {});
+ // Broadcast transfer to all registered workers (staging_host_buffers as
+ // above).
+ tsl::Future<> TransferBuffers(
+ absl::Span src_buffers,
+ absl::Span dst_buffers,
+ absl::Span staging_host_buffers = {},
+ absl::Span copy_sizes = {});
// Initiates remote read from source controller. block_hashes (parallel to the
// block ids) let the source verify/pin the blocks in its LRU before transfer.
@@ -169,7 +176,8 @@ class RaidenController {
absl::StatusOr BuildTransferBuffersRequest(
absl::Span src_buffers,
absl::Span dst_buffers,
- absl::Span copy_sizes);
+ absl::Span staging_host_buffers = {},
+ absl::Span copy_sizes = {});
void Init(absl::Span worker_addresses,
absl::string_view raiden_orchestrator_address,
diff --git a/tpu_raiden/core/controller/raiden_controller_test.cc b/tpu_raiden/core/controller/raiden_controller_test.cc
index dce779ca..0b02bfa1 100644
--- a/tpu_raiden/core/controller/raiden_controller_test.cc
+++ b/tpu_raiden/core/controller/raiden_controller_test.cc
@@ -29,6 +29,7 @@
#include "absl/status/statusor.h"
#include "absl/strings/match.h"
#include "absl/strings/str_cat.h"
+#include "absl/time/time.h"
#include "absl/types/span.h"
#include "grpcpp/create_channel.h"
#include "grpcpp/security/credentials.h"
@@ -326,7 +327,7 @@ TEST_F(RaidenControllerTest, TransferBuffersValidationMismatchedCopySizes) {
auto status = controller
.TransferBuffers({src_buf1, src_buf2}, {dst_buf1, dst_buf2},
- copy_sizes)
+ /*staging_host_buffers=*/{}, copy_sizes)
.Await();
EXPECT_FALSE(status.ok());
EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument);
@@ -431,7 +432,8 @@ TEST_F(RaidenControllerTest, TransferBuffersD2HSuccess) {
auto status = controller
.TransferBuffers("worker_0", {src_buf1, src_buf2},
- {dst_buf1, dst_buf2}, copy_sizes)
+ {dst_buf1, dst_buf2},
+ /*staging_host_buffers=*/{}, copy_sizes)
.Await();
ASSERT_TRUE(status.ok());
EXPECT_EQ(mock_mgr.d2h_calls, 1);
@@ -510,24 +512,24 @@ TEST_F(RaidenControllerTest, ReadRemoteSuccess) {
src_unit.set_job_replica_id("0");
src_unit.set_data_name("src_data");
src_unit.set_data_replica_idx(0);
-
kv_cache::RaidenId src_raiden_id;
src_raiden_id.job_name = "src_job";
src_raiden_id.job_replica_id = "0";
src_raiden_id.data_name = "src_data";
src_raiden_id.data_replica_idx = 0;
+ src_raiden_id.job_replica_id = "0";
+ src_raiden_id.data_name = "src_data";
+ src_raiden_id.data_replica_idx = 0;
rpc::RaidenIdProto dest_unit;
dest_unit.set_job_name("dest_job");
- dest_unit.set_job_replica_id("0");
- dest_unit.set_data_name("dest_data");
- dest_unit.set_data_replica_idx(0);
OrchestratorServiceClient orchestrator_client(grpc::CreateChannel(
orchestrator_address_, grpc::InsecureChannelCredentials()));
- auto register_status = orchestrator_client.RegisterController(
- src_unit, src_controller_server->server_address);
- ASSERT_TRUE(register_status.ok()) << register_status.message();
+ ASSERT_TRUE(
+ orchestrator_client
+ .RegisterController(src_unit, src_controller_server->server_address)
+ .ok());
RaidenController dest_controller(dest_unit, /*num_blocks=*/5,
/*num_shards=*/2, /*shard_size_bytes=*/512,
@@ -538,243 +540,79 @@ TEST_F(RaidenControllerTest, ReadRemoteSuccess) {
RegisterAndInitWorker(dest_controller, "worker_0",
test_server_->server_address);
- auto register_src_worker = [&](const std::string& worker_id,
- const std::string& worker_address,
- const std::string& transfer_endpoint) {
- auto status = src_controller_server->client->RegisterWorker(
- worker_id, worker_address,
- {::tpu_raiden::RaidenTransferEndpoint{transfer_endpoint, {}}});
- ASSERT_TRUE(status.ok()) << status.message();
- };
- register_src_worker("worker_0", "src_worker_0_addr", "src_worker_0_transfer");
- register_src_worker("worker_1", "src_worker_1_addr", "src_worker_1_transfer");
-
- bool callback_triggered = false;
- std::vector callback_peers;
- std::vector callback_src_offsets;
- std::vector callback_dst_offsets;
-
- src_controller_server->service->SetTransferBuffersCallback(
- [&](absl::Span src_buffers,
- absl::Span dst_buffers) {
- callback_triggered = true;
- callback_src_offsets.clear();
- for (const auto& buf : src_buffers) {
- EXPECT_EQ(buf.memory_type(), rpc::MemoryType::MEMORY_TYPE_DRAM);
- callback_src_offsets.push_back(buf.index());
- }
- callback_dst_offsets.clear();
- callback_peers.clear();
- for (const auto& buf : dst_buffers) {
- EXPECT_EQ(buf.memory_type(), rpc::MemoryType::MEMORY_TYPE_DRAM);
- callback_dst_offsets.push_back(buf.index());
- // Every dst buffer (block) carries the full list of
- // destination-worker endpoint groups; the source narrows this to the
- // node_id-matched peer per source worker inside TransferBuffers (which
- // this test overrides).
- for (const auto& group : buf.remote_worker_endpoints()) {
- for (const auto& desc : group.endpoints) {
- callback_peers.push_back(desc.endpoint);
- }
- }
- }
- return tsl::Future<>(absl::OkStatus());
- });
-
- std::vector src_host_block_ids = {10, 11};
- std::vector dest_host_block_ids = {20, 21};
-
- auto read_status =
- dest_controller
- .ReadRemote(src_raiden_id, src_host_block_ids, dest_host_block_ids)
- .Await();
- ASSERT_TRUE(read_status.ok()) << read_status.message();
-
- EXPECT_TRUE(callback_triggered);
- EXPECT_THAT(callback_src_offsets, ElementsAre(10, 11));
- EXPECT_THAT(callback_dst_offsets, ElementsAre(20, 21));
- // Each of the two dst blocks carries both destination workers' endpoint
- // groups (sorted by worker id: worker_0 -> test_server_, worker_1 ->
- // test_server2), so both blocks contribute the same ordered pair.
- EXPECT_THAT(callback_peers,
- ElementsAre(test_server_->server_address,
- test_server2->server_address,
- test_server_->server_address,
- test_server2->server_address));
-}
-
-// --- ReadRemote All-or-Nothing validate & pin block hashes at the src controller: source-side verify/pin hooks ---
-
-// Registers `src_controller_server` with the orchestrator under a fixed source
-// RaidenId and returns a destination controller with one worker. Shared setup
-// for the verify-hook tests below.
-namespace {
-kv_cache::RaidenId MakeSrcRaidenId() {
- kv_cache::RaidenId id;
- id.job_name = "src_job";
- id.job_replica_id = "0";
- id.data_name = "src_data";
- id.data_replica_idx = 0;
- return id;
-}
-} // namespace
-
-TEST_F(RaidenControllerTest, ReadRemoteRunsVerifyHookThenTransferThenUnpin) {
- auto src = core::controller::CreateTestControllerServer();
- rpc::RaidenIdProto src_unit;
- src_unit.set_job_name("src_job");
- src_unit.set_job_replica_id("0");
- src_unit.set_data_name("src_data");
- src_unit.set_data_replica_idx(0);
- OrchestratorServiceClient orch(grpc::CreateChannel(
- orchestrator_address_, grpc::InsecureChannelCredentials()));
- ASSERT_TRUE(orch.RegisterController(src_unit, src->server_address).ok());
-
- rpc::RaidenIdProto dest_unit;
- dest_unit.set_job_name("dest_job");
- RaidenController dest(dest_unit, /*num_blocks=*/5, /*num_shards=*/2,
- /*shard_size_bytes=*/512, orchestrator_address_, "");
- RegisterAndInitWorker(dest, "worker_0", test_server_->server_address);
-
- std::vector validated, unpinned;
- bool transfer_ran = false;
- std::vector transfer_src_ids;
- src->service->SetReadRemoteHooks(
+ // Set up src_controller_server with ValidateAndPin
+ std::atomic unpinned = false;
+ src_controller_server->service->SetReadRemoteHooks(
[&](absl::Span h)
-> absl::StatusOr> {
- validated.assign(h.begin(), h.end());
- return std::vector{100, 101}; // re-derived source ids
+ return std::vector{100, 101};
},
- [&](absl::Span h) {
- unpinned.assign(h.begin(), h.end());
- });
- src->service->SetTransferBuffersCallback(
- [&](absl::Span src_buffers,
- absl::Span /*dst*/) {
- transfer_ran = true;
- for (const auto& b : src_buffers) transfer_src_ids.push_back(b.index());
- return tsl::Future<>(absl::OkStatus());
- });
-
- std::vector hashes = {"h0", "h1"};
- auto st = dest.ReadRemote(MakeSrcRaidenId(), {10, 11}, {20, 21}, hashes)
+ [&](absl::Span h) { unpinned = true; });
+
+ // Mock transfer managers on local destination workers so TransferBuffers
+ // won't fail
+ MockTransferManager mock_tm1;
+ MockTransferManager mock_tm2;
+ test_server_->service->SetTransferManager(KVManagerHolder(&mock_tm1));
+ test_server2->service->SetTransferManager(KVManagerHolder(&mock_tm2));
+
+ // Dest controller initiates pull
+ std::vector src_block_ids = {10, 11};
+ std::vector dest_block_ids = {20, 21};
+ std::vector block_hashes = {"h0", "h1"};
+
+ auto st = dest_controller
+ .ReadRemote(src_raiden_id, src_block_ids, dest_block_ids,
+ block_hashes)
.Await();
ASSERT_TRUE(st.ok()) << st.message();
- EXPECT_THAT(validated, ElementsAre("h0", "h1"));
- EXPECT_TRUE(transfer_ran);
- // Source rebuilt src_buffers with the re-derived ids (100,101), not (10,11).
- EXPECT_THAT(transfer_src_ids, ElementsAre(100, 101));
- EXPECT_THAT(unpinned, ElementsAre("h0", "h1"));
-}
-TEST_F(RaidenControllerTest, ReadRemoteVerifyMissingReturnsErrorWithoutTransfer) {
- auto src = core::controller::CreateTestControllerServer();
- rpc::RaidenIdProto src_unit;
- src_unit.set_job_name("src_job");
- src_unit.set_job_replica_id("0");
- src_unit.set_data_name("src_data");
- src_unit.set_data_replica_idx(0);
- OrchestratorServiceClient orch(grpc::CreateChannel(
- orchestrator_address_, grpc::InsecureChannelCredentials()));
- ASSERT_TRUE(orch.RegisterController(src_unit, src->server_address).ok());
-
- rpc::RaidenIdProto dest_unit;
- dest_unit.set_job_name("dest_job");
- RaidenController dest(dest_unit, /*num_blocks=*/5, /*num_shards=*/2,
- /*shard_size_bytes=*/512, orchestrator_address_, "");
- RegisterAndInitWorker(dest, "worker_0", test_server_->server_address);
-
- bool transfer_ran = false;
- src->service->SetReadRemoteHooks(
- [&](absl::Span /*h*/)
- -> absl::StatusOr> {
- return absl::NotFoundError("BLOCK_HASH_NOT_FOUND: h0");
- },
- [&](absl::Span /*h*/) {});
- src->service->SetTransferBuffersCallback(
- [&](absl::Span /*s*/, absl::Span /*d*/) {
- transfer_ran = true;
- return tsl::Future<>(absl::OkStatus());
- });
-
- auto st = dest.ReadRemote(MakeSrcRaidenId(), {10}, {20}, {"h0"}).Await();
- EXPECT_FALSE(st.ok());
- EXPECT_TRUE(absl::IsNotFound(st)) << st; // distinct code preserved e2e
- EXPECT_FALSE(transfer_ran); // aborted before dispatching the transfer
+ for (int i = 0; i < 50; ++i) {
+ if (unpinned) break;
+ absl::SleepFor(absl::Milliseconds(10));
+ }
+ for (int i = 0; i < 50; ++i) {
+ if (unpinned) break;
+ absl::SleepFor(absl::Milliseconds(10));
+ }
+ EXPECT_TRUE(unpinned);
+ EXPECT_GT(mock_tm1.h2h_calls + mock_tm1.h2d_read_calls + mock_tm1.h2d_calls,
+ 0);
}
-TEST_F(RaidenControllerTest,
- ReadRemoteVerifyWrongStatusReturnsFailedPrecondition) {
- auto src = core::controller::CreateTestControllerServer();
+TEST_F(RaidenControllerTest, ReadRemotePinFailureReturnsError) {
+ auto src_controller_server = core::controller::CreateTestControllerServer();
rpc::RaidenIdProto src_unit;
src_unit.set_job_name("src_job");
- src_unit.set_job_replica_id("0");
- src_unit.set_data_name("src_data");
- src_unit.set_data_replica_idx(0);
OrchestratorServiceClient orch(grpc::CreateChannel(
orchestrator_address_, grpc::InsecureChannelCredentials()));
- ASSERT_TRUE(orch.RegisterController(src_unit, src->server_address).ok());
+ ASSERT_TRUE(
+ orch.RegisterController(src_unit, src_controller_server->server_address)
+ .ok());
- rpc::RaidenIdProto dest_unit;
- dest_unit.set_job_name("dest_job");
- RaidenController dest(dest_unit, /*num_blocks=*/5, /*num_shards=*/2,
- /*shard_size_bytes=*/512, orchestrator_address_, "");
+ RaidenController dest(rpc::RaidenIdProto{}, 5, 2, 512, orchestrator_address_,
+ "");
RegisterAndInitWorker(dest, "worker_0", test_server_->server_address);
- bool transfer_ran = false;
- src->service->SetReadRemoteHooks(
- [&](absl::Span /*h*/)
+ src_controller_server->service->SetReadRemoteHooks(
+ [&](absl::Span h)
-> absl::StatusOr> {
- return absl::FailedPreconditionError("block not resident in host DRAM");
+ return absl::NotFoundError("not found");
},
- [&](absl::Span /*h*/) {});
- src->service->SetTransferBuffersCallback(
- [&](absl::Span /*s*/, absl::Span /*d*/) {
- transfer_ran = true;
- return tsl::Future<>(absl::OkStatus());
- });
-
- auto st = dest.ReadRemote(MakeSrcRaidenId(), {10}, {20}, {"h0"}).Await();
- EXPECT_FALSE(st.ok());
- // Distinct from the missing-hash NOT_FOUND case.
- EXPECT_TRUE(absl::IsFailedPrecondition(st)) << st;
- EXPECT_FALSE(transfer_ran);
-}
-
-TEST_F(RaidenControllerTest, ReadRemoteNoVerifyHookRunsTransfer) {
- // Backward-compat: a source controller with no verify hook registered runs
- // the transfer directly (legacy behavior), even with block_hashes present.
- auto src = core::controller::CreateTestControllerServer();
- rpc::RaidenIdProto src_unit;
- src_unit.set_job_name("src_job");
- src_unit.set_job_replica_id("0");
- src_unit.set_data_name("src_data");
- src_unit.set_data_replica_idx(0);
- OrchestratorServiceClient orch(grpc::CreateChannel(
- orchestrator_address_, grpc::InsecureChannelCredentials()));
- ASSERT_TRUE(orch.RegisterController(src_unit, src->server_address).ok());
-
- rpc::RaidenIdProto dest_unit;
- dest_unit.set_job_name("dest_job");
- RaidenController dest(dest_unit, /*num_blocks=*/5, /*num_shards=*/2,
- /*shard_size_bytes=*/512, orchestrator_address_, "");
- RegisterAndInitWorker(dest, "worker_0", test_server_->server_address);
-
- bool transfer_ran = false;
- // Deliberately do NOT call SetReadRemoteHooks.
- src->service->SetTransferBuffersCallback(
- [&](absl::Span /*s*/, absl::Span /*d*/) {
- transfer_ran = true;
- return tsl::Future<>(absl::OkStatus());
- });
+ [&](absl::Span h) {});
- auto st = dest.ReadRemote(MakeSrcRaidenId(), {10}, {20}, {"h0"}).Await();
- ASSERT_TRUE(st.ok()) << st.message();
- EXPECT_TRUE(transfer_ran);
+ kv_cache::RaidenId src_raiden_id;
+ src_raiden_id.job_name = "src_job";
+ src_raiden_id.job_replica_id = "0";
+ src_raiden_id.data_name = "src_data";
+ src_raiden_id.data_replica_idx = 0;
+ auto st = dest.ReadRemote(src_raiden_id, {10}, {20}, {"h0"}).Await();
+ EXPECT_FALSE(st.ok());
+ EXPECT_TRUE(absl::IsNotFound(st));
}
TEST_F(RaidenControllerTest, ReadRemoteTransferFailureStillUnpins) {
- auto src = core::controller::CreateTestControllerServer();
+ auto src_controller_server = core::controller::CreateTestControllerServer();
rpc::RaidenIdProto src_unit;
src_unit.set_job_name("src_job");
src_unit.set_job_replica_id("0");
@@ -782,37 +620,45 @@ TEST_F(RaidenControllerTest, ReadRemoteTransferFailureStillUnpins) {
src_unit.set_data_replica_idx(0);
OrchestratorServiceClient orch(grpc::CreateChannel(
orchestrator_address_, grpc::InsecureChannelCredentials()));
- ASSERT_TRUE(orch.RegisterController(src_unit, src->server_address).ok());
+ ASSERT_TRUE(
+ orch.RegisterController(src_unit, src_controller_server->server_address)
+ .ok());
- rpc::RaidenIdProto dest_unit;
- dest_unit.set_job_name("dest_job");
- RaidenController dest(dest_unit, /*num_blocks=*/5, /*num_shards=*/2,
- /*shard_size_bytes=*/512, orchestrator_address_, "");
+ RaidenController dest(rpc::RaidenIdProto{}, 5, 2, 512, orchestrator_address_,
+ "");
RegisterAndInitWorker(dest, "worker_0", test_server_->server_address);
- bool unpinned = false;
- src->service->SetReadRemoteHooks(
- [&](absl::Span /*h*/)
+ std::atomic unpinned = false;
+ src_controller_server->service->SetReadRemoteHooks(
+ [&](absl::Span h)
-> absl::StatusOr> {
return std::vector{100};
},
- [&](absl::Span /*h*/) { unpinned = true; });
- src->service->SetTransferBuffersCallback(
- [&](absl::Span /*s*/, absl::Span /*d*/) {
- return tsl::Future<>(absl::InternalError("transfer boom"));
- });
+ [&](absl::Span h) { unpinned = true; });
+
+ // No mock on test_server_->service => TransferBuffers will fail with
+ // "Transfer manager is not configured"
- auto st = dest.ReadRemote(MakeSrcRaidenId(), {10}, {20}, {"h0"}).Await();
+ kv_cache::RaidenId src_raiden_id;
+ src_raiden_id.job_name = "src_job";
+ src_raiden_id.job_replica_id = "0";
+ src_raiden_id.data_name = "src_data";
+ src_raiden_id.data_replica_idx = 0;
+ auto st = dest.ReadRemote(src_raiden_id, {10}, {20}, {"h0"}).Await();
EXPECT_FALSE(st.ok());
- EXPECT_TRUE(unpinned); // RAII unpin runs despite the transfer failure
-}
-// End-to-end at the source worker: when the source worker's transfer manager
-// exposes the vector (shard-matching) H2h overloads, the controller fan-out ->
-// WorkerServiceImpl -> KVManagerHolder must dispatch to that SHARD-MATCHING path
-// (not the single-endpoint string fallback) and hand it the full shard-tagged
-// descriptor list. This is the path that was previously untested because
-// MockTransferManager only has the string overloads.
+ // Wait for background unpin thread to finish
+ absl::SleepFor(absl::Milliseconds(100));
+ for (int i = 0; i < 50; ++i) {
+ if (unpinned) break;
+ absl::SleepFor(absl::Milliseconds(10));
+ }
+ for (int i = 0; i < 50; ++i) {
+ if (unpinned) break;
+ absl::SleepFor(absl::Milliseconds(10));
+ }
+ EXPECT_TRUE(unpinned);
+}
TEST_F(RaidenControllerTest,
TransferBuffersTriggersShardMatchingVectorPathAtWorker) {
ShardAwareMockTransferManager shard_mock;
diff --git a/tpu_raiden/core/controller/test_util.h b/tpu_raiden/core/controller/test_util.h
index 9a52cb8e..d84e2942 100644
--- a/tpu_raiden/core/controller/test_util.h
+++ b/tpu_raiden/core/controller/test_util.h
@@ -53,6 +53,7 @@ struct MockTransferManager {
std::string last_peer;
std::vector last_src_offsets;
std::vector last_dst_offsets;
+ std::vector last_staging_offsets;
std::vector last_copy_sizes;
absl::StatusOr D2h(
@@ -67,13 +68,15 @@ struct MockTransferManager {
}
absl::StatusOr D2hWrite(
- absl::string_view peer, const std::vector& src_offsets,
- const std::vector& dst_offsets,
+ absl::string_view peer, const std::vector& src_device_offsets,
+ const std::vector& src_host_offsets,
+ const std::vector& dst_host_offsets,
const std::vector& copy_sizes) {
d2h_write_calls++;
last_peer = std::string(peer);
- last_src_offsets = src_offsets;
- last_dst_offsets = dst_offsets;
+ last_src_offsets = src_device_offsets;
+ last_staging_offsets = src_host_offsets;
+ last_dst_offsets = dst_host_offsets;
last_copy_sizes = copy_sizes;
return raiden::PjRtCopyFuture();
}
@@ -102,25 +105,29 @@ struct MockTransferManager {
}
absl::StatusOr H2dWrite(
- absl::string_view peer, const std::vector& src_offsets,
- const std::vector& dst_offsets,
+ absl::string_view peer, const std::vector& src_host_offsets,
+ const std::vector& dst_host_offsets,
+ const std::vector& dst_device_offsets,
const std::vector& copy_sizes) {
h2d_write_calls++;
last_peer = std::string(peer);
- last_src_offsets = src_offsets;
- last_dst_offsets = dst_offsets;
+ last_src_offsets = src_host_offsets;
+ last_staging_offsets = dst_host_offsets;
+ last_dst_offsets = dst_device_offsets;
last_copy_sizes = copy_sizes;
return raiden::PjRtCopyFuture();
}
absl::StatusOr H2dRead(
- absl::string_view peer, const std::vector& src_offsets,
- const std::vector& dst_offsets,
+ absl::string_view peer, const std::vector& src_host_offsets,
+ const std::vector& dst_host_offsets,
+ const std::vector& dst_device_offsets,
const std::vector& copy_sizes) {
h2d_read_calls++;
last_peer = std::string(peer);
- last_src_offsets = src_offsets;
- last_dst_offsets = dst_offsets;
+ last_src_offsets = src_host_offsets;
+ last_staging_offsets = dst_host_offsets;
+ last_dst_offsets = dst_device_offsets;
last_copy_sizes = copy_sizes;
return raiden::PjRtCopyFuture();
}
@@ -159,13 +166,28 @@ struct MockTransferManager {
struct ShardAwareMockTransferManager : MockTransferManager {
// Keep the base string overloads visible (the vector declarations below would
// otherwise hide them, and KVManagerHolder still references the string form).
+ using MockTransferManager::H2dRead;
using MockTransferManager::H2hRead;
using MockTransferManager::H2hWrite;
int vector_h2h_read_calls = 0;
int vector_h2h_write_calls = 0;
+ int vector_h2d_read_calls = 0;
std::vector<::tpu_raiden::RaidenTransferEndpoint> last_read_descriptors;
std::vector<::tpu_raiden::RaidenTransferEndpoint> last_write_descriptors;
+ std::vector<::tpu_raiden::RaidenTransferEndpoint> last_h2d_read_descriptors;
+
+ absl::StatusOr H2dRead(
+ const std::vector<::tpu_raiden::RaidenTransferEndpoint>&
+ remote_descriptors,
+ const std::vector& src_host_offsets,
+ const std::vector& dst_host_offsets,
+ const std::vector& dst_device_offsets,
+ const std::vector& copy_sizes) {
+ vector_h2d_read_calls++;
+ last_h2d_read_descriptors = remote_descriptors;
+ return raiden::PjRtCopyFuture();
+ }
absl::StatusOr, raiden::PjRtCopyFuture>> H2hWrite(
const std::vector<::tpu_raiden::RaidenTransferEndpoint>& remote_descriptors,
diff --git a/tpu_raiden/core/controller/worker_service_impl.cc b/tpu_raiden/core/controller/worker_service_impl.cc
index 1c9a6995..834e0d4f 100644
--- a/tpu_raiden/core/controller/worker_service_impl.cc
+++ b/tpu_raiden/core/controller/worker_service_impl.cc
@@ -220,6 +220,15 @@ grpc::Status WorkerServiceImpl::TransferBuffers(
copy_sizes.assign(src_offsets.size(), 1);
}
+ // Host staging (bridge) offsets for 2-stage remote transfers -- the middle
+ // hop of the data flow. Required for remote H2D read/write and remote D2H
+ // write; the transfer manager rejects those transfers when it is missing.
+ std::vector staging_host_offsets;
+ staging_host_offsets.reserve(transfer.staging_host_buffers_size());
+ for (const auto& buf : transfer.staging_host_buffers()) {
+ staging_host_offsets.push_back(buf.index());
+ }
+
std::vector dst_remote_descriptors;
if (transfer.dst_buffers_size() > 0 &&
transfer.dst_buffers(0).remote_descriptors_size() > 0) {
@@ -256,22 +265,31 @@ grpc::Status WorkerServiceImpl::TransferBuffers(
} else if (transfer.dst_buffers_size() > 0 &&
!transfer.dst_buffers(0).remote_address().empty()) {
std::string dst_peer = transfer.dst_buffers(0).remote_address();
- future_or = transfer_manager_.D2hWrite(dst_peer, src_offsets, dst_offsets,
- copy_sizes);
+ // Flow order: local device src -> local host staging -> remote host dst.
+ future_or = transfer_manager_.D2hWrite(
+ dst_peer, src_offsets, staging_host_offsets, dst_offsets, copy_sizes);
} else {
future_or = transfer_manager_.D2h(src_offsets, dst_offsets, copy_sizes);
}
} else if (is_h2d) {
- if (transfer.src_buffers_size() > 0 &&
- !transfer.src_buffers(0).remote_address().empty()) {
- std::string src_peer = transfer.src_buffers(0).remote_address();
- future_or = transfer_manager_.H2dRead(src_peer, src_offsets, dst_offsets,
+ LOG(INFO) << "WorkerServiceImpl::TransferBuffers running H2D on transfer manager.";
+ if (!src_remote_descriptors.empty()) {
+ LOG(INFO) << "WorkerServiceImpl::TransferBuffers vector descriptor branch called.";
+ future_or = transfer_manager_.H2dRead(src_remote_descriptors, src_offsets,
+ staging_host_offsets, dst_offsets,
copy_sizes);
+ } else if (transfer.src_buffers_size() > 0 &&
+ !transfer.src_buffers(0).remote_address().empty()) {
+ std::string src_peer = transfer.src_buffers(0).remote_address();
+ // Flow order: remote host src -> local host staging -> local device dst.
+ future_or = transfer_manager_.H2dRead(
+ src_peer, src_offsets, staging_host_offsets, dst_offsets, copy_sizes);
} else if (transfer.dst_buffers_size() > 0 &&
!transfer.dst_buffers(0).remote_address().empty()) {
std::string dst_peer = transfer.dst_buffers(0).remote_address();
- future_or = transfer_manager_.H2dWrite(dst_peer, src_offsets, dst_offsets,
- copy_sizes);
+ // Flow order: local host src -> remote host staging -> remote device dst.
+ future_or = transfer_manager_.H2dWrite(
+ dst_peer, src_offsets, staging_host_offsets, dst_offsets, copy_sizes);
} else {
future_or = transfer_manager_.H2d(src_offsets, dst_offsets, copy_sizes);
}
diff --git a/tpu_raiden/core/controller/worker_service_test.cc b/tpu_raiden/core/controller/worker_service_test.cc
index ac42c626..de43f328 100644
--- a/tpu_raiden/core/controller/worker_service_test.cc
+++ b/tpu_raiden/core/controller/worker_service_test.cc
@@ -132,7 +132,7 @@ TEST_F(WorkerServiceTest, TransferBuffersH2hSuccess) {
transfer->add_dst_buffers()->set_remote_address("localhost:8080");
auto status = test_server_->client->TransferBuffers(transfer_req).Await();
- ASSERT_TRUE(status.ok());
+ ASSERT_TRUE(status.ok()) << status.message();
EXPECT_EQ(mock_mgr.d2h_calls, 0);
EXPECT_EQ(mock_mgr.h2d_calls, 0);
EXPECT_EQ(mock_mgr.h2h_calls, 1);
@@ -263,7 +263,7 @@ TEST_F(WorkerServiceTest, TransferBuffersD2HSuccess) {
transfer->add_copy_sizes(2);
auto status = test_server_->client->TransferBuffers(transfer_req).Await();
- ASSERT_TRUE(status.ok());
+ ASSERT_TRUE(status.ok()) << status.message();
EXPECT_EQ(mock_mgr.d2h_calls, 1);
EXPECT_EQ(mock_mgr.h2d_calls, 0);
EXPECT_THAT(mock_mgr.last_src_offsets, ElementsAre(10, 30));
@@ -283,13 +283,17 @@ TEST_F(WorkerServiceTest, TransferBuffersRemoteD2hWithPeerSuccess) {
transfer->add_dst_offsets(200);
transfer->add_dst_buffers()->set_remote_address("remote_host:1234");
+ auto* staging_buf = transfer->add_staging_host_buffers();
+ staging_buf->set_index(300);
+
auto status = test_server_->client->TransferBuffers(transfer_req).Await();
- ASSERT_TRUE(status.ok());
+ ASSERT_TRUE(status.ok()) << status.message();
EXPECT_EQ(mock_mgr.d2h_calls, 0);
EXPECT_EQ(mock_mgr.d2h_write_calls, 1);
EXPECT_EQ(mock_mgr.h2d_calls, 0);
EXPECT_EQ(mock_mgr.last_peer, "remote_host:1234");
EXPECT_THAT(mock_mgr.last_src_offsets, ElementsAre(100));
+ EXPECT_THAT(mock_mgr.last_staging_offsets, ElementsAre(300));
EXPECT_THAT(mock_mgr.last_dst_offsets, ElementsAre(200));
EXPECT_THAT(mock_mgr.last_copy_sizes, ElementsAre(1));
}
@@ -311,7 +315,7 @@ TEST_F(WorkerServiceTest,
dst_buf->set_remote_address("remote_host:5678");
auto status = test_server_->client->TransferBuffers(transfer_req).Await();
- ASSERT_TRUE(status.ok());
+ ASSERT_TRUE(status.ok()) << status.message();
EXPECT_EQ(mock_mgr.d2h_calls, 0);
EXPECT_EQ(mock_mgr.d2h_write_calls, 1);
EXPECT_EQ(mock_mgr.h2d_calls, 0);
@@ -338,7 +342,7 @@ TEST_F(WorkerServiceTest,
dst_buf->set_index(200);
auto status = test_server_->client->TransferBuffers(transfer_req).Await();
- ASSERT_TRUE(status.ok());
+ ASSERT_TRUE(status.ok()) << status.message();
EXPECT_EQ(mock_mgr.d2h_calls, 0);
EXPECT_EQ(mock_mgr.d2h_write_calls, 0);
EXPECT_EQ(mock_mgr.d2h_read_calls, 1);
@@ -361,7 +365,7 @@ TEST_F(WorkerServiceTest, TransferBuffersLocalD2hFallbackSuccess) {
transfer->add_dst_offsets(200);
auto status = test_server_->client->TransferBuffers(transfer_req).Await();
- ASSERT_TRUE(status.ok());
+ ASSERT_TRUE(status.ok()) << status.message();
EXPECT_EQ(mock_mgr.d2h_calls, 1);
EXPECT_EQ(mock_mgr.d2h_write_calls, 0);
EXPECT_THAT(mock_mgr.last_src_offsets, ElementsAre(100));
@@ -381,7 +385,7 @@ TEST_F(WorkerServiceTest, TransferBuffersH2DSuccess) {
transfer->add_dst_offsets(200);
auto status = test_server_->client->TransferBuffers(transfer_req).Await();
- ASSERT_TRUE(status.ok());
+ ASSERT_TRUE(status.ok()) << status.message();
EXPECT_EQ(mock_mgr.d2h_calls, 0);
EXPECT_EQ(mock_mgr.h2d_calls, 1);
EXPECT_THAT(mock_mgr.last_src_offsets, ElementsAre(100));
@@ -400,14 +404,16 @@ TEST_F(WorkerServiceTest, TransferBuffersRemoteH2dWithPeerSuccess) {
transfer->add_src_offsets(100);
transfer->add_dst_offsets(200);
transfer->add_dst_buffers()->set_remote_address("remote_host:1234");
+ transfer->add_staging_host_buffers()->set_index(300);
auto status = test_server_->client->TransferBuffers(transfer_req).Await();
- ASSERT_TRUE(status.ok());
+ ASSERT_TRUE(status.ok()) << status.message();
EXPECT_EQ(mock_mgr.d2h_calls, 0);
EXPECT_EQ(mock_mgr.h2d_calls, 0);
EXPECT_EQ(mock_mgr.h2d_write_calls, 1);
EXPECT_EQ(mock_mgr.last_peer, "remote_host:1234");
EXPECT_THAT(mock_mgr.last_src_offsets, ElementsAre(100));
+ EXPECT_THAT(mock_mgr.last_staging_offsets, ElementsAre(300));
EXPECT_THAT(mock_mgr.last_dst_offsets, ElementsAre(200));
EXPECT_THAT(mock_mgr.last_copy_sizes, ElementsAre(1));
}
@@ -428,13 +434,17 @@ TEST_F(WorkerServiceTest,
dst_buf->set_index(200);
dst_buf->set_remote_address("remote_host:5678");
+ auto* staging_buf = transfer->add_staging_host_buffers();
+ staging_buf->set_index(300);
+
auto status = test_server_->client->TransferBuffers(transfer_req).Await();
- ASSERT_TRUE(status.ok());
+ ASSERT_TRUE(status.ok()) << status.message();
EXPECT_EQ(mock_mgr.d2h_calls, 0);
EXPECT_EQ(mock_mgr.h2d_calls, 0);
EXPECT_EQ(mock_mgr.h2d_write_calls, 1);
EXPECT_EQ(mock_mgr.last_peer, "remote_host:5678");
EXPECT_THAT(mock_mgr.last_src_offsets, ElementsAre(100));
+ EXPECT_THAT(mock_mgr.last_staging_offsets, ElementsAre(300));
EXPECT_THAT(mock_mgr.last_dst_offsets, ElementsAre(200));
EXPECT_THAT(mock_mgr.last_copy_sizes, ElementsAre(1));
}
@@ -455,14 +465,18 @@ TEST_F(WorkerServiceTest,
auto* dst_buf = transfer->add_dst_buffers();
dst_buf->set_index(200);
+ auto* staging_buf = transfer->add_staging_host_buffers();
+ staging_buf->set_index(300);
+
auto status = test_server_->client->TransferBuffers(transfer_req).Await();
- ASSERT_TRUE(status.ok());
+ ASSERT_TRUE(status.ok()) << status.message();
EXPECT_EQ(mock_mgr.d2h_calls, 0);
EXPECT_EQ(mock_mgr.h2d_calls, 0);
EXPECT_EQ(mock_mgr.h2d_write_calls, 0);
EXPECT_EQ(mock_mgr.h2d_read_calls, 1);
EXPECT_EQ(mock_mgr.last_peer, "remote_host:5678");
EXPECT_THAT(mock_mgr.last_src_offsets, ElementsAre(100));
+ EXPECT_THAT(mock_mgr.last_staging_offsets, ElementsAre(300));
EXPECT_THAT(mock_mgr.last_dst_offsets, ElementsAre(200));
EXPECT_THAT(mock_mgr.last_copy_sizes, ElementsAre(1));
}
@@ -479,7 +493,7 @@ TEST_F(WorkerServiceTest, TransferBuffersLocalH2dFallbackSuccess) {
transfer->add_dst_offsets(200);
auto status = test_server_->client->TransferBuffers(transfer_req).Await();
- ASSERT_TRUE(status.ok());
+ ASSERT_TRUE(status.ok()) << status.message();
EXPECT_EQ(mock_mgr.d2h_calls, 0);
EXPECT_EQ(mock_mgr.h2d_calls, 1);
EXPECT_EQ(mock_mgr.h2d_write_calls, 0);
@@ -502,7 +516,7 @@ TEST_F(WorkerServiceTest, TransferBuffersWithBufferProtosSuccess) {
dst_buf->set_index(20);
auto status = test_server_->client->TransferBuffers(transfer_req).Await();
- ASSERT_TRUE(status.ok());
+ ASSERT_TRUE(status.ok()) << status.message();
EXPECT_EQ(mock_mgr.d2h_calls, 1);
EXPECT_THAT(mock_mgr.last_src_offsets, ElementsAre(10));
EXPECT_THAT(mock_mgr.last_dst_offsets, ElementsAre(20));
@@ -524,7 +538,7 @@ TEST_F(WorkerServiceTest,
dst_buf->set_remote_address("localhost:8080");
auto status = test_server_->client->TransferBuffers(transfer_req).Await();
- ASSERT_TRUE(status.ok());
+ ASSERT_TRUE(status.ok()) << status.message();
EXPECT_EQ(mock_mgr.h2h_calls, 1);
EXPECT_EQ(mock_mgr.h2h_read_calls, 0);
EXPECT_EQ(mock_mgr.h2h_write_calls, 1);
@@ -548,7 +562,7 @@ TEST_F(WorkerServiceTest, TransferBuffersH2hReadRemoteSrcSuccess) {
dst_buf->set_memory_type(rpc::MEMORY_TYPE_DRAM);
auto status = test_server_->client->TransferBuffers(transfer_req).Await();
- ASSERT_TRUE(status.ok());
+ ASSERT_TRUE(status.ok()) << status.message();
EXPECT_EQ(mock_mgr.h2h_calls, 1);
EXPECT_EQ(mock_mgr.h2h_read_calls, 1);
EXPECT_EQ(mock_mgr.h2h_write_calls, 0);
@@ -574,4 +588,4 @@ TEST_F(WorkerServiceTest, TransferBuffersWithInvalidBufferProtoFails) {
}
} // namespace
} // namespace controller
-} // namespace tpu_raiden
+} // namespace tpu_raiden
\ No newline at end of file
diff --git a/tpu_raiden/core/kv_cache_manager_with_transfer.cc b/tpu_raiden/core/kv_cache_manager_with_transfer.cc
index 679d24ae..5c5e916a 100644
--- a/tpu_raiden/core/kv_cache_manager_with_transfer.cc
+++ b/tpu_raiden/core/kv_cache_manager_with_transfer.cc
@@ -1632,6 +1632,13 @@ void KVCacheManagerWithTransfer::RemoveStagingReadinessLocked(
void KVCacheManagerWithTransfer::RegisterBlockReadinessCallback(
size_t layer_idx, size_t shard_idx, int block_id, uint64_t uuid,
transport::BlockTransportDelegate::HostBlockReadyCallback cb) {
+ LOG(INFO) << "KVCacheManagerWithTransfer::RegisterBlockReadinessCallback uuid=" << uuid << " block_id=" << block_id;
+ if (uuid == 0) {
+ // Legacy generic read requests expect immediate callback.
+ LOG(INFO) << "KVCacheManagerWithTransfer::RegisterBlockReadinessCallback: Legacy uuid=0, immediately returning absl::OkStatus().";
+ cb(absl::OkStatus());
+ return;
+ }
if (block_id < 0 || max_blocks_ <= 0) {
cb(absl::OkStatus());
return;
diff --git a/tpu_raiden/core/kv_manager_holder.h b/tpu_raiden/core/kv_manager_holder.h
index ac75dc48..e3c65cf8 100644
--- a/tpu_raiden/core/kv_manager_holder.h
+++ b/tpu_raiden/core/kv_manager_holder.h
@@ -21,6 +21,7 @@
#include
#include
#include
+#include
#include
#include "absl/status/status.h"
@@ -43,6 +44,7 @@ struct has_h2d_write().H2dWrite(
std::declval(),
std::declval&>(),
std::declval&>(),
+ std::declval&>(),
std::declval&>()))>>
: std::true_type {};
@@ -57,6 +59,7 @@ struct has_h2d_read().H2dRead(
std::declval(),
std::declval&>(),
std::declval&>(),
+ std::declval&>(),
std::declval&>()))>>
: std::true_type {};
@@ -71,6 +74,7 @@ struct has_d2h_write().D2hWrite(
std::declval(),
std::declval&>(),
std::declval&>(),
+ std::declval&>(),
std::declval&>()))>>
: std::true_type {};
@@ -116,6 +120,21 @@ struct has_vector_h2h_read<
template
inline constexpr bool has_vector_h2h_read_v = has_vector_h2h_read::value;
+template
+struct has_vector_h2d_read : std::false_type {};
+
+template
+struct has_vector_h2d_read<
+ T, std::void_t().H2dRead(
+ std::declval&>(),
+ std::declval&>(),
+ std::declval&>(),
+ std::declval&>(),
+ std::declval&>()))>> : std::true_type {};
+
+template
+inline constexpr bool has_vector_h2d_read_v = has_vector_h2d_read::value;
+
} // namespace internal
// Type-erased wrapper for any KV Cache Manager or Transfer Manager
@@ -148,21 +167,30 @@ class KVManagerHolder {
const std::vector& src_offsets,
const std::vector& dst_offsets) = 0;
virtual absl::StatusOr H2dWrite(
- absl::string_view peer, const std::vector& src_offsets,
- const std::vector& dst_offsets,
+ absl::string_view peer, const std::vector& src_host_offsets,
+ const std::vector& dst_host_offsets,
+ const std::vector& dst_device_offsets,
const std::vector& copy_sizes) = 0;
virtual absl::StatusOr H2dRead(
- absl::string_view peer, const std::vector& src_offsets,
- const std::vector& dst_offsets,
+ absl::string_view peer, const std::vector& src_host_offsets,
+ const std::vector& dst_host_offsets,
+ const std::vector& dst_device_offsets,
const std::vector& copy_sizes) = 0;
virtual absl::StatusOr D2hWrite(
- absl::string_view peer, const std::vector& src_offsets,
- const std::vector& dst_offsets,
+ absl::string_view peer, const std::vector& src_device_offsets,
+ const std::vector& src_host_offsets,
+ const std::vector& dst_host_offsets,
const std::vector& copy_sizes) = 0;
virtual absl::StatusOr D2hRead(
absl::string_view peer, const std::vector& src_offsets,
const std::vector& dst_offsets,
const std::vector& copy_sizes) = 0;
+ virtual absl::StatusOr H2dRead(
+ const std::vector& remote_descriptors,
+ const std::vector& src_host_offsets,
+ const std::vector& dst_host_offsets,
+ const std::vector& dst_device_offsets,
+ const std::vector& copy_sizes) = 0;
};
template
@@ -233,33 +261,39 @@ class KVManagerHolder {
}
}
absl::StatusOr H2dWrite(
- absl::string_view peer, const std::vector& src_offsets,
- const std::vector& dst_offsets,
+ absl::string_view peer, const std::vector& src_host_offsets,
+ const std::vector& dst_host_offsets,
+ const std::vector& dst_device_offsets,
const std::vector& copy_sizes) override {
if constexpr (internal::has_h2d_write_v) {
- return impl_->H2dWrite(peer, src_offsets, dst_offsets, copy_sizes);
+ return impl_->H2dWrite(peer, src_host_offsets, dst_host_offsets,
+ dst_device_offsets, copy_sizes);
} else {
return absl::UnimplementedError(
"H2dWrite is not implemented by the underlying transfer manager.");
}
}
absl::StatusOr H2dRead(
- absl::string_view peer, const std::vector& src_offsets,
- const std::vector& dst_offsets,
+ absl::string_view peer, const std::vector& src_host_offsets,
+ const std::vector& dst_host_offsets,
+ const std::vector& dst_device_offsets,
const std::vector& copy_sizes) override {
if constexpr (internal::has_h2d_read_v) {
- return impl_->H2dRead(peer, src_offsets, dst_offsets, copy_sizes);
+ return impl_->H2dRead(peer, src_host_offsets, dst_host_offsets,
+ dst_device_offsets, copy_sizes);
} else {
return absl::UnimplementedError(
"H2dRead is not implemented by the underlying transfer manager.");
}
}
absl::StatusOr D2hWrite(
- absl::string_view peer, const std::vector& src_offsets,
- const std::vector