kvcache-ai/Mooncake

Fork 0

mirror of https://github.com/kvcache-ai/Mooncake.git synced 2026-11-05 01:17:55 +00:00

[GH-ISSUE #2662] [RFC]: Add TPU Transport Support via PJRT #698

New issue

Open

opened 2026-10-29 21:18:21 +00:00 by bsakaguchi · 7 comments

bsakaguchi commented

2026-10-29 21:18:21 +00:00

Owner

Originally created by @Liwink on GitHub (Jun 29, 2026).
Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/2662

Originally assigned to: @staryxchen on GitHub.

Changes proposed

1. Summary

We propose adding TPU as a first-class accelerator target in Mooncake's Transfer Engine (TENT). The implementation introduces a TpuTransport class that enables KVCache disaggregation between TPU-based Prefill and Decode nodes, using a staged host-DRAM pipeline over existing TENT transports (RDMA / TCP) — without requiring any proprietary TPU SDK at build time.

The key architectural insight is that TPU KVCache transfer decomposes into two independent halves that map cleanly onto TENT's existing capabilities:

Stage	Mechanism	Owner
TPU HBM ↔ Host DRAM	PJRT async copy (`BufferToHostBuffer` / `CopyToHostAsync`)	New `TpuTransport`
Host DRAM ↔ Remote Host DRAM	Existing TENT RDMA/TCP transport	Existing TENT infra

This decomposition avoids flooding the TPU's Inter-Chip Interconnect (ICI) with asynchronous KVCache traffic, preserving ICI bandwidth for synchronous model-parallel collectives.

Estimated scope: ~800–1200 LOC (excluding tests).

2. Motivation

2.1 The Problem

Google Cloud TPUs (v4, v5e, v5p, v6e) are a major LLM serving platform — powering production workloads on JAX/MaxText, PyTorch/XLA, and vLLM-TPU. However, Mooncake currently has zero TPU support, despite supporting 8+ other accelerator platforms (NVIDIA CUDA, AMD ROCm, Huawei Ascend, Cambricon MLU, Moore Threads MUSA, MetaX MACA, Hygon DCU, Iluvatar CoreX).

This gap prevents the entire Google Cloud TPU ecosystem from benefiting from Mooncake's KVCache disaggregation, hierarchical storage, and P2P distribution capabilities.

2.2 Why Not GPUDirect RDMA?

Unlike NVIDIA GPUs, TPUs do not expose third-party NIC access to device HBM:

Capability	NVIDIA GPU	Google TPU
Direct NIC→VRAM DMA	✅ GPUDirect RDMA	❌ Not available
Memory registration API	`cuMemAlloc` + `ibv_reg_mr`	Opaque `PjRtBuffer`
Host driver model	Open CUDA toolkit	`libtpu.so` (closed-source runtime)

This means Mooncake's existing transport path (register device buffer → NIC DMA read/write) cannot work directly for TPU. A staged pipeline through host DRAM is required.

2.3 Why TENT?

TPU KVCache disaggregation needs the same infrastructure that TENT already provides:

Infrastructure	What TPU Needs	TENT Status
Multi-NIC topology discovery	"Which NIC is closest to this TPU host?"	`Topology::discover()` ✅
Peer discovery & handshake	Exchange buffer metadata, rkeys	`ControlService::bootstrap()` ✅
Host memory registration	Register staging DRAM for RDMA	`Transport::registerLocalMemory()` ✅
Bandwidth aggregation	Saturate multi-NIC links	Multi-path slicing ✅
Async transfer management	Batch pipeline stages	`submitTransfer` / `getTransferStatus` ✅

Building this outside TENT would duplicate significant effort.

3. Technical Design

3.1 Architecture Overview

flowchart TD
    subgraph Source["Source TPU VM"]
        direction TB
        HBM1["TPU HBM<br/>(KVCache)"]
        DRAM1["Host DRAM<br/>Staging Pool"]
        TE1["Mooncake Transfer Engine"]
        HBM1 -- "Layer 1: PJRT BufferToHostBuffer<br/>(PCIe DMA)" --> DRAM1
        DRAM1 -- "Layer 3: Zero-Copy IPC<br/>(shm_open / mmap)" --> TE1
    end

    subgraph Net["Data Center Network"]
        direction LR
        NIC1["Host NIC / gVNIC"]
        NIC2["Host NIC / gVNIC"]
        NIC1 == "RDMA / TCP" ==> NIC2
    end

    subgraph Dest["Destination TPU VM"]
        direction TB
        TE2["Mooncake Transfer Engine"]
        DRAM2["Host DRAM<br/>Staging Pool"]
        HBM2["TPU HBM<br/>(KVCache)"]
        TE2 -- "Layer 3: Zero-Copy IPC<br/>(shm_open / mmap)" --> DRAM2
        DRAM2 -- "Layer 1: PJRT BufferFromHostBuffer<br/>(PCIe DMA)" --> HBM2
    end

    TE1 --> NIC1
    NIC2 --> TE2

    style Source fill:#e8f5e9,stroke:#2e7d32,stroke-width:2px
    style Net fill:#e1f5fe,stroke:#0288d1,stroke-width:2px
    style Dest fill:#e8f5e9,stroke:#2e7d32,stroke-width:2px

3.2 Data Flow

The full KVCache transfer pipeline follows a staged 2-copy model:

\text{Source TPU HBM} \xrightarrow{\;\text{PJRT DMA}\;} \text{Source Host DRAM} \xrightarrow{\;\text{TENT RDMA/TCP}\;} \text{Dest Host DRAM} \xrightarrow{\;\text{PJRT DMA}\;} \text{Dest TPU HBM}

Why bypass the TPU ICI network?

Google's Inter-Chip Interconnect (ICI) is a high-bandwidth 2D/3D torus hardwired between TPUs within a Pod slice, optimized for synchronous model-parallel collectives (AllReduce, AllGather). Flooding ICI with asynchronous All-to-All KVCache traffic between disaggregated Prefill/Decode nodes causes contention that degrades matrix-multiply latency.

By staging KVCache through host DRAM and transmitting over multi-NIC datacenter networks (gVNIC or host RDMA), the TPU ICI links remain fully available for compute.

Division of labor

Pipeline Stage	Responsible Component	API Surface
HBM → Host DRAM	PJRT via `BufferToHostBuffer`	`TpuTransport::stageOut()`
Host DRAM → Remote Host DRAM	Existing TENT RDMA/TCP	`Transport::submitTransfer()`
Remote Host DRAM → HBM	PJRT via `BufferFromHostBuffer`	`TpuTransport::stageIn()`

3.3 Layer 1: C++ `TpuTransport`

3.3.1 Zero Build-Time Dependency via `dlopen`

Zero Build-Time Dependency: To ensure standard Linux build farms continue compiling smoothly, TpuTransport uses POSIX dlopen / dlsym at runtime to bind symbols from libtpu.so or libpjrt_c_api.so. This feature is gated behind option(USE_TPU "Enable TPU PJRT runtime support" OFF). No find_package or proprietary SDK is needed.

// mooncake-transfer-engine/src/transport/tpu_transport/tpu_transport.h

class TpuTransport : public Transport {
public:
    // Transport interface implementation
    int registerLocalMemory(void* addr, size_t length, 
                           const std::string& location, bool remote) override;
    int submitTransfer(BatchID batch_id, 
                      const TransferRequest& request) override;
    TransferStatus getTransferStatus(BatchID batch_id, 
                                     size_t task_id) override;
    
private:
    // PJRT symbols resolved at runtime via dlopen
    struct PjrtApi {
        void* handle;  // dlopen handle to libtpu.so / libpjrt_c_api.so
        
        // Function pointers resolved via dlsym
        decltype(&PJRT_Client_Create) client_create;
        decltype(&PJRT_Client_BufferFromHostBuffer) buffer_from_host;
        decltype(&PJRT_Buffer_ToHostBuffer) buffer_to_host;
        decltype(&PJRT_Buffer_OnDeviceSizeInBytes) buffer_size;
        decltype(&PJRT_Event_Await) event_await;
        // ...
    };
    
    PjrtApi pjrt_;
    std::unique_ptr<StagingBufferPool> staging_pool_;
};

3.3.2 Build System Integration

# Root CMakeLists.txt addition
option(USE_TPU "Enable TPU PJRT runtime support (no build-time TPU SDK required)" OFF)

if(USE_TPU)
    add_definitions(-DUSE_TPU)
    # No find_package needed — PJRT loaded at runtime via dlopen
    # Only need standard POSIX: dlopen, shm_open, mmap
endif()

The build guard follows the established pattern from USE_CUDA, USE_ROCM, USE_ASCEND, etc. in CMakeLists.txt and mooncake-common/common.cmake.

3.3.3 Staging Buffer Pool

// Shared memory staging pool for zero-copy between processes
class StagingBufferPool {
public:
    // Allocate staging buffer in POSIX shared memory
    StagingBuffer* allocate(size_t size);
    void release(StagingBuffer* buffer);
    
    // Get shared memory descriptor for IPC handoff
    ShmDescriptor getDescriptor(StagingBuffer* buffer);
    
private:
    int shm_fd_;              // POSIX shared memory file descriptor
    void* base_addr_;         // mmap'd base address
    size_t pool_size_;        // Total pool size
    FreeListAllocator alloc_; // Sub-allocator within the pool
};

struct StagingBuffer {
    void* host_ptr;       // Mapped host address
    size_t size;          // Buffer size
    int32_t offset;       // Offset within shared memory region
};

struct ShmDescriptor {
    char name[64];        // shm_open name
    int32_t offset;       // Offset within the region
    size_t length;        // Buffer length
};

3.4 Layer 2: Python `TPUAllocator`

Mirroring the existing NVLinkAllocator and BarexAllocator in mooncake-integration/allocator.py:

# mooncake-integration/allocator.py

class TPUAllocator(Allocator):
    """Allocator for TPU device buffers via PJRT/torch_xla."""
    
    def __init__(self, device_id: int = 0):
        super().__init__()
        self.device_id = device_id
        # Detect PJRT runtime availability
        try:
            import torch_xla.core.xla_model as xm
            self._xla = xm
            self._device = xm.xla_device(n=device_id)
        except ImportError:
            try:
                import jax
                self._jax = jax
                self._device = jax.devices('tpu')[device_id]
            except ImportError:
                raise RuntimeError(
                    "TPUAllocator requires torch_xla or jax. "
                    "Install via: pip install torch_xla or pip install jax[tpu]"
                )
    
    def allocate(self, size: int) -> MemoryHandle:
        """Allocate a KVCache staging buffer on the TPU device."""
        ...
    
    def stage_to_host(self, handle: MemoryHandle) -> HostBuffer:
        """Stage device buffer to host DRAM via PJRT async copy."""
        ...
    
    def stage_from_host(self, host_buf: HostBuffer, handle: MemoryHandle):
        """Copy host DRAM buffer back to TPU device via PJRT."""
        ...

3.5 Layer 3: Zero-Copy Shared Memory

When the serving engine (e.g., vLLM worker) and the Mooncake networking daemon run as separate processes on the same TPU VM host:

Allocator creates POSIX shared memory regions via shm_open / mmap
Engine process stages KVCache from TPU HBM into the shared DRAM region
Engine process passes a 16-byte ShmDescriptor (name + offset + length) to the TENT daemon via an IPC socket
TENT daemon accesses the same physical DRAM at zero CPU memcpy cost and transmits over the network

This avoids the memory bus saturation that would occur with inter-process memcpy of 500MB+ KVCache chunks.

3.6 Memory Location & Topology Integration

// mooncake-transfer-engine/src/memory_location.cpp — additions
case MemoryLocationType::TPU_HBM:
    return parseTPUDeviceAddress(addr);  // e.g., "tpu:0" → device 0 HBM

// mooncake-transfer-engine/src/topology.cpp — additions
// TPU topology: discover local TPU devices via PJRT device enumeration
// Map each TPU chip to its closest host NIC using PCIe topology

3.7 `patchelf` Build Consideration

Warning

PyTorch/XLA Symbol Collisions: When PyTorch/XLA C++ extensions are loaded alongside mooncake-ep/, XLA allocator symbol registration collisions can occur during Python module import. Our build configuration will incorporate a post-build patchelf step to inject a NEEDED dynamic entry linking against libpywrap_torch_tpu_common.so, ensuring a single unified global symbol table. This proven technique prevents fatal crashes during import.

4. Files Changed (Estimated Touchpoints)

Following the established contribution pattern from prior hardware PRs (#2118, #2019):

Category	Files	Description
Build	`CMakeLists.txt`, `mooncake-common/common.cmake`	`USE_TPU` option, conditional compilation guards
Transport	`mooncake-transfer-engine/include/transport/transport.h`	Register `TpuTransport` in transport enum/factory
New Transport	`mooncake-transfer-engine/src/transport/tpu_transport/tpu_transport.{h,cpp}`	Core `TpuTransport` implementation
New Transport	`mooncake-transfer-engine/src/transport/tpu_transport/pjrt_loader.{h,cpp}`	`dlopen`/`dlsym` PJRT C API loader
New Transport	`mooncake-transfer-engine/src/transport/tpu_transport/staging_buffer_pool.{h,cpp}`	POSIX shared memory staging pool
Memory	`mooncake-transfer-engine/src/memory_location.cpp`	TPU HBM address resolution
Topology	`mooncake-transfer-engine/src/topology.cpp`	TPU device ↔ NIC affinity discovery
Store	`mooncake-store/include/gpu_staging_utils.h`	TPU staging pointer integration
Store	`mooncake-store/include/pinned_buffer_pool.h`	TPU pinned buffer support
Python	`mooncake-integration/allocator.py`	`TPUAllocator` class
Tests	`mooncake-transfer-engine/tests/test_tpu_transport.cpp`	Unit tests (mock PJRT)
Tests	`mooncake-integration/tests/test_tpu_allocator.py`	Python allocator tests
Docs	`docs/source/getting_started/supported-protocols.md`	Document TPU transport

5. Implementation Plan

We propose splitting this into 3 incremental PRs to facilitate review:

PR 1: `[TransferEngine] Add TpuTransport with PJRT dynamic loading` (~500 LOC)

Build system: USE_TPU option in CMakeLists
pjrt_loader.{h,cpp}: Runtime dlopen / dlsym resolution
tpu_transport.{h,cpp}: Core Transport interface implementation
staging_buffer_pool.{h,cpp}: POSIX shared memory allocator
memory_location.cpp / topology.cpp: TPU address + topology
Unit tests with mocked PJRT symbols
Reviewable independently — pure C++, no Python dependencies

PR 2: `[Integration] Add TPUAllocator for vLLM/SGLang` (~200 LOC)

TPUAllocator class in allocator.py
PyTorch/XLA and JAX dual-backend support
patchelf build hook for symbol collision avoidance
Python integration tests

PR 3: `[Doc] Add TPU transport documentation and examples` (~100 LOC)

supported-protocols.md updates
TPU-specific getting-started guide
Architecture diagram

6. Testing Strategy

Test Level	Description	Environment
Unit (CI)	Mock PJRT C API symbols, test `TpuTransport` logic, staging pool alloc/free, `dlopen` error paths	Standard Linux (no TPU needed)
Integration	End-to-end KVCache transfer between 2 TPU VMs using RDMA/TCP	GCE TPU v4-8 (contributor-provisioned)
Benchmark	Measure staging latency (HBM→DRAM→RDMA→DRAM→HBM) vs. baseline ICI	GCE TPU v5e-16

CI compatibility: Since USE_TPU=OFF by default and all TPU dependencies are dlopen'd at runtime, enabling USE_TPU=ON in CI requires no TPU hardware — only the mock tests run. Real hardware tests are contributor-gated.

7. FAQ / Anticipated Questions

Q: Why not use ICI for KVCache transfer?

ICI is a fixed-topology, synchronous-optimized torus network designed for AllReduce / AllGather collectives within a Pod slice. Injecting asynchronous, unpredictable A2A KVCache traffic causes contention that degrades matrix-multiply throughput. Staging through host DRAM + datacenter NIC keeps ICI 100% available for compute.

Q: What's the performance overhead of the staged pipeline?

The staged pipeline adds ~2 PCIe DMA copies (HBM→DRAM, DRAM→HBM) relative to a hypothetical GPUDirect path. On TPU v5e with PCIe Gen5 x16:

PCIe bandwidth: ~64 GB/s per direction
Typical KVCache slice: ~256 MB
Per-copy latency: ~4 ms
Total staging overhead: ~8 ms (pipeline-able with network transfer)

The pipeline is double-buffered so staging and network transfer overlap, hiding most of the latency.

Q: Does this add any build-time dependencies?

No. The PJRT C API is loaded at runtime via dlopen("libtpu.so") or dlopen("libpjrt_c_api.so"). When USE_TPU=ON, the build only requires standard POSIX headers (dlfcn.h, sys/mman.h). The resulting binary runs on any Linux system and gracefully fails with a descriptive error if no PJRT runtime is found.

Before submitting a new issue...

Make sure you already searched for relevant issues and read the documentation
Checked that this aligns with the project Roadmap (#1883)

Before submitting a new issue...

Make sure you already searched for relevant issues and read the documentation

Originally created by @Liwink on GitHub (Jun 29, 2026). Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/2662 Originally assigned to: @staryxchen on GitHub. ### Changes proposed ## Changes proposed ## 1. Summary We propose adding **TPU** as a first-class accelerator target in Mooncake's Transfer Engine (TENT). The implementation introduces a `TpuTransport` class that enables KVCache disaggregation between TPU-based Prefill and Decode nodes, using a **staged host-DRAM pipeline** over existing TENT transports (RDMA / TCP) — without requiring any proprietary TPU SDK at build time. The key architectural insight is that TPU KVCache transfer decomposes into two independent halves that map cleanly onto TENT's existing capabilities: | Stage | Mechanism | Owner | |-------|-----------|-------| | TPU HBM ↔ Host DRAM | PJRT async copy (`BufferToHostBuffer` / `CopyToHostAsync`) | New `TpuTransport` | | Host DRAM ↔ Remote Host DRAM | Existing TENT RDMA/TCP transport | Existing TENT infra | This decomposition avoids flooding the TPU's Inter-Chip Interconnect (ICI) with asynchronous KVCache traffic, preserving ICI bandwidth for synchronous model-parallel collectives. **Estimated scope:** ~800–1200 LOC (excluding tests). --- ## 2. Motivation ### 2.1 The Problem Google Cloud TPUs (v4, v5e, v5p, v6e) are a major LLM serving platform — powering production workloads on JAX/MaxText, PyTorch/XLA, and vLLM-TPU. However, **Mooncake currently has zero TPU support**, despite supporting 8+ other accelerator platforms (NVIDIA CUDA, AMD ROCm, Huawei Ascend, Cambricon MLU, Moore Threads MUSA, MetaX MACA, Hygon DCU, Iluvatar CoreX). This gap prevents the entire Google Cloud TPU ecosystem from benefiting from Mooncake's KVCache disaggregation, hierarchical storage, and P2P distribution capabilities. ### 2.2 Why Not GPUDirect RDMA? Unlike NVIDIA GPUs, TPUs do not expose third-party NIC access to device HBM: | Capability | NVIDIA GPU | Google TPU | |-----------|-----------|-----------| | Direct NIC→VRAM DMA | ✅ GPUDirect RDMA | ❌ Not available | | Memory registration API | `cuMemAlloc` + `ibv_reg_mr` | Opaque `PjRtBuffer` | | Host driver model | Open CUDA toolkit | `libtpu.so` (closed-source runtime) | This means Mooncake's existing transport path (register device buffer → NIC DMA read/write) cannot work directly for TPU. A staged pipeline through host DRAM is required. ### 2.3 Why TENT? TPU KVCache disaggregation needs the same infrastructure that TENT already provides: | Infrastructure | What TPU Needs | TENT Status | |---------------|---------------|-------------| | Multi-NIC topology discovery | "Which NIC is closest to this TPU host?" | `Topology::discover()` ✅ | | Peer discovery & handshake | Exchange buffer metadata, rkeys | `ControlService::bootstrap()` ✅ | | Host memory registration | Register staging DRAM for RDMA | `Transport::registerLocalMemory()` ✅ | | Bandwidth aggregation | Saturate multi-NIC links | Multi-path slicing ✅ | | Async transfer management | Batch pipeline stages | `submitTransfer` / `getTransferStatus` ✅ | Building this outside TENT would duplicate significant effort. --- ## 3. Technical Design ### 3.1 Architecture Overview ```mermaid flowchart TD subgraph Source["Source TPU VM"] direction TB HBM1["TPU HBM (KVCache)"] DRAM1["Host DRAM Staging Pool"] TE1["Mooncake Transfer Engine"] HBM1 -- "Layer 1: PJRT BufferToHostBuffer (PCIe DMA)" --> DRAM1 DRAM1 -- "Layer 3: Zero-Copy IPC (shm_open / mmap)" --> TE1 end subgraph Net["Data Center Network"] direction LR NIC1["Host NIC / gVNIC"] NIC2["Host NIC / gVNIC"] NIC1 == "RDMA / TCP" ==> NIC2 end subgraph Dest["Destination TPU VM"] direction TB TE2["Mooncake Transfer Engine"] DRAM2["Host DRAM Staging Pool"] HBM2["TPU HBM (KVCache)"] TE2 -- "Layer 3: Zero-Copy IPC (shm_open / mmap)" --> DRAM2 DRAM2 -- "Layer 1: PJRT BufferFromHostBuffer (PCIe DMA)" --> HBM2 end TE1 --> NIC1 NIC2 --> TE2 style Source fill:#e8f5e9,stroke:#2e7d32,stroke-width:2px style Net fill:#e1f5fe,stroke:#0288d1,stroke-width:2px style Dest fill:#e8f5e9,stroke:#2e7d32,stroke-width:2px ``` ### 3.2 Data Flow The full KVCache transfer pipeline follows a **staged 2-copy model**: $$\text{Source TPU HBM} \xrightarrow{\;\text{PJRT DMA}\;} \text{Source Host DRAM} \xrightarrow{\;\text{TENT RDMA/TCP}\;} \text{Dest Host DRAM} \xrightarrow{\;\text{PJRT DMA}\;} \text{Dest TPU HBM}$$ #### Why bypass the TPU ICI network? Google's **Inter-Chip Interconnect (ICI)** is a high-bandwidth 2D/3D torus hardwired between TPUs within a Pod slice, optimized for synchronous model-parallel collectives (`AllReduce`, `AllGather`). Flooding ICI with asynchronous All-to-All KVCache traffic between disaggregated Prefill/Decode nodes causes contention that degrades matrix-multiply latency. By staging KVCache through host DRAM and transmitting over multi-NIC datacenter networks (gVNIC or host RDMA), the TPU ICI links remain fully available for compute. #### Division of labor | Pipeline Stage | Responsible Component | API Surface | |---------------|----------------------|-------------| | HBM → Host DRAM | PJRT via `BufferToHostBuffer` | `TpuTransport::stageOut()` | | Host DRAM → Remote Host DRAM | Existing TENT RDMA/TCP | `Transport::submitTransfer()` | | Remote Host DRAM → HBM | PJRT via `BufferFromHostBuffer` | `TpuTransport::stageIn()` | ### 3.3 Layer 1: C++ `TpuTransport` #### 3.3.1 Zero Build-Time Dependency via `dlopen` **Zero Build-Time Dependency:** To ensure standard Linux build farms continue compiling smoothly, `TpuTransport` uses POSIX `dlopen` / `dlsym` at runtime to bind symbols from `libtpu.so` or `libpjrt_c_api.so`. This feature is gated behind `option(USE_TPU "Enable TPU PJRT runtime support" OFF)`. No `find_package` or proprietary SDK is needed. ```cpp // mooncake-transfer-engine/src/transport/tpu_transport/tpu_transport.h class TpuTransport : public Transport { public: // Transport interface implementation int registerLocalMemory(void* addr, size_t length, const std::string& location, bool remote) override; int submitTransfer(BatchID batch_id, const TransferRequest& request) override; TransferStatus getTransferStatus(BatchID batch_id, size_t task_id) override; private: // PJRT symbols resolved at runtime via dlopen struct PjrtApi { void* handle; // dlopen handle to libtpu.so / libpjrt_c_api.so // Function pointers resolved via dlsym decltype(&PJRT_Client_Create) client_create; decltype(&PJRT_Client_BufferFromHostBuffer) buffer_from_host; decltype(&PJRT_Buffer_ToHostBuffer) buffer_to_host; decltype(&PJRT_Buffer_OnDeviceSizeInBytes) buffer_size; decltype(&PJRT_Event_Await) event_await; // ... }; PjrtApi pjrt_; std::unique_ptr<StagingBufferPool> staging_pool_; }; ``` #### 3.3.2 Build System Integration ```cmake # Root CMakeLists.txt addition option(USE_TPU "Enable TPU PJRT runtime support (no build-time TPU SDK required)" OFF) if(USE_TPU) add_definitions(-DUSE_TPU) # No find_package needed — PJRT loaded at runtime via dlopen # Only need standard POSIX: dlopen, shm_open, mmap endif() ``` The build guard follows the established pattern from `USE_CUDA`, `USE_ROCM`, `USE_ASCEND`, etc. in `CMakeLists.txt` and `mooncake-common/common.cmake`. #### 3.3.3 Staging Buffer Pool ```cpp // Shared memory staging pool for zero-copy between processes class StagingBufferPool { public: // Allocate staging buffer in POSIX shared memory StagingBuffer* allocate(size_t size); void release(StagingBuffer* buffer); // Get shared memory descriptor for IPC handoff ShmDescriptor getDescriptor(StagingBuffer* buffer); private: int shm_fd_; // POSIX shared memory file descriptor void* base_addr_; // mmap'd base address size_t pool_size_; // Total pool size FreeListAllocator alloc_; // Sub-allocator within the pool }; struct StagingBuffer { void* host_ptr; // Mapped host address size_t size; // Buffer size int32_t offset; // Offset within shared memory region }; struct ShmDescriptor { char name[64]; // shm_open name int32_t offset; // Offset within the region size_t length; // Buffer length }; ``` ### 3.4 Layer 2: Python `TPUAllocator` Mirroring the existing `NVLinkAllocator` and `BarexAllocator` in `mooncake-integration/allocator.py`: ```python # mooncake-integration/allocator.py class TPUAllocator(Allocator): """Allocator for TPU device buffers via PJRT/torch_xla.""" def __init__(self, device_id: int = 0): super().__init__() self.device_id = device_id # Detect PJRT runtime availability try: import torch_xla.core.xla_model as xm self._xla = xm self._device = xm.xla_device(n=device_id) except ImportError: try: import jax self._jax = jax self._device = jax.devices('tpu')[device_id] except ImportError: raise RuntimeError( "TPUAllocator requires torch_xla or jax. " "Install via: pip install torch_xla or pip install jax[tpu]" ) def allocate(self, size: int) -> MemoryHandle: """Allocate a KVCache staging buffer on the TPU device.""" ... def stage_to_host(self, handle: MemoryHandle) -> HostBuffer: """Stage device buffer to host DRAM via PJRT async copy.""" ... def stage_from_host(self, host_buf: HostBuffer, handle: MemoryHandle): """Copy host DRAM buffer back to TPU device via PJRT.""" ... ``` ### 3.5 Layer 3: Zero-Copy Shared Memory When the serving engine (e.g., vLLM worker) and the Mooncake networking daemon run as separate processes on the same TPU VM host: 1. **Allocator** creates POSIX shared memory regions via `shm_open` / `mmap` 2. **Engine process** stages KVCache from TPU HBM into the shared DRAM region 3. **Engine process** passes a 16-byte `ShmDescriptor` (name + offset + length) to the TENT daemon via an IPC socket 4. **TENT daemon** accesses the same physical DRAM at zero CPU `memcpy` cost and transmits over the network This avoids the memory bus saturation that would occur with inter-process `memcpy` of 500MB+ KVCache chunks. ### 3.6 Memory Location & Topology Integration ```cpp // mooncake-transfer-engine/src/memory_location.cpp — additions case MemoryLocationType::TPU_HBM: return parseTPUDeviceAddress(addr); // e.g., "tpu:0" → device 0 HBM // mooncake-transfer-engine/src/topology.cpp — additions // TPU topology: discover local TPU devices via PJRT device enumeration // Map each TPU chip to its closest host NIC using PCIe topology ``` ### 3.7 `patchelf` Build Consideration > [!WARNING] > **PyTorch/XLA Symbol Collisions:** When PyTorch/XLA C++ extensions are loaded alongside `mooncake-ep/`, XLA allocator symbol registration collisions can occur during Python module import. Our build configuration will incorporate a post-build `patchelf` step to inject a `NEEDED` dynamic entry linking against `libpywrap_torch_tpu_common.so`, ensuring a single unified global symbol table. This proven technique prevents fatal crashes during import. --- ## 4. Files Changed (Estimated Touchpoints) Following the established contribution pattern from prior hardware PRs (#2118, #2019): | Category | Files | Description | |----------|-------|-------------| | **Build** | `CMakeLists.txt`, `mooncake-common/common.cmake` | `USE_TPU` option, conditional compilation guards | | **Transport** | `mooncake-transfer-engine/include/transport/transport.h` | Register `TpuTransport` in transport enum/factory | | **New Transport** | `mooncake-transfer-engine/src/transport/tpu_transport/tpu_transport.{h,cpp}` | Core `TpuTransport` implementation | | **New Transport** | `mooncake-transfer-engine/src/transport/tpu_transport/pjrt_loader.{h,cpp}` | `dlopen`/`dlsym` PJRT C API loader | | **New Transport** | `mooncake-transfer-engine/src/transport/tpu_transport/staging_buffer_pool.{h,cpp}` | POSIX shared memory staging pool | | **Memory** | `mooncake-transfer-engine/src/memory_location.cpp` | TPU HBM address resolution | | **Topology** | `mooncake-transfer-engine/src/topology.cpp` | TPU device ↔ NIC affinity discovery | | **Store** | `mooncake-store/include/gpu_staging_utils.h` | TPU staging pointer integration | | **Store** | `mooncake-store/include/pinned_buffer_pool.h` | TPU pinned buffer support | | **Python** | `mooncake-integration/allocator.py` | `TPUAllocator` class | | **Tests** | `mooncake-transfer-engine/tests/test_tpu_transport.cpp` | Unit tests (mock PJRT) | | **Tests** | `mooncake-integration/tests/test_tpu_allocator.py` | Python allocator tests | | **Docs** | `docs/source/getting_started/supported-protocols.md` | Document TPU transport | --- ## 5. Implementation Plan We propose splitting this into 3 incremental PRs to facilitate review: ### PR 1: `[TransferEngine] Add TpuTransport with PJRT dynamic loading` (~500 LOC) - Build system: `USE_TPU` option in CMakeLists - `pjrt_loader.{h,cpp}`: Runtime `dlopen` / `dlsym` resolution - `tpu_transport.{h,cpp}`: Core Transport interface implementation - `staging_buffer_pool.{h,cpp}`: POSIX shared memory allocator - `memory_location.cpp` / `topology.cpp`: TPU address + topology - Unit tests with mocked PJRT symbols - **Reviewable independently** — pure C++, no Python dependencies ### PR 2: `[Integration] Add TPUAllocator for vLLM/SGLang` (~200 LOC) - `TPUAllocator` class in `allocator.py` - PyTorch/XLA and JAX dual-backend support - `patchelf` build hook for symbol collision avoidance - Python integration tests ### PR 3: `[Doc] Add TPU transport documentation and examples` (~100 LOC) - `supported-protocols.md` updates - TPU-specific getting-started guide - Architecture diagram --- ## 6. Testing Strategy | Test Level | Description | Environment | |-----------|-------------|-------------| | **Unit (CI)** | Mock PJRT C API symbols, test `TpuTransport` logic, staging pool alloc/free, `dlopen` error paths | Standard Linux (no TPU needed) | | **Integration** | End-to-end KVCache transfer between 2 TPU VMs using RDMA/TCP | GCE TPU v4-8 (contributor-provisioned) | | **Benchmark** | Measure staging latency (HBM→DRAM→RDMA→DRAM→HBM) vs. baseline ICI | GCE TPU v5e-16 | **CI compatibility:** Since `USE_TPU=OFF` by default and all TPU dependencies are `dlopen`'d at runtime, enabling `USE_TPU=ON` in CI requires no TPU hardware — only the mock tests run. Real hardware tests are contributor-gated. --- ## 7. FAQ / Anticipated Questions ### Q: Why not use ICI for KVCache transfer? ICI is a fixed-topology, synchronous-optimized torus network designed for `AllReduce` / `AllGather` collectives within a Pod slice. Injecting asynchronous, unpredictable A2A KVCache traffic causes contention that degrades matrix-multiply throughput. Staging through host DRAM + datacenter NIC keeps ICI 100% available for compute. ### Q: What's the performance overhead of the staged pipeline? The staged pipeline adds ~2 PCIe DMA copies (HBM→DRAM, DRAM→HBM) relative to a hypothetical GPUDirect path. On TPU v5e with PCIe Gen5 x16: - PCIe bandwidth: ~64 GB/s per direction - Typical KVCache slice: ~256 MB - Per-copy latency: ~4 ms - Total staging overhead: ~8 ms (pipeline-able with network transfer) The pipeline is **double-buffered** so staging and network transfer overlap, hiding most of the latency. ### Q: Does this add any build-time dependencies? **No.** The PJRT C API is loaded at runtime via `dlopen("libtpu.so")` or `dlopen("libpjrt_c_api.so")`. When `USE_TPU=ON`, the build only requires standard POSIX headers (`dlfcn.h`, `sys/mman.h`). The resulting binary runs on any Linux system and gracefully fails with a descriptive error if no PJRT runtime is found. --- ### Before submitting a new issue... - [x] Make sure you already searched for relevant issues and read the [documentation](https://kvcache-ai.github.io/Mooncake/) - [x] Checked that this aligns with the project [Roadmap (#1883)](https://github.com/kvcache-ai/Mooncake/issues/1883) ### Before submitting a new issue... - [x] Make sure you already searched for relevant issues and read the [documentation](https://kvcache-ai.github.io/Mooncake/)

bsakaguchi added the

RFC

label

2026-10-29 21:18:21 +00:00

bsakaguchi commented

2026-10-29 21:18:22 +00:00

Author

Owner

@github-actions[bot] commented on GitHub (Jun 29, 2026):

Thanks for opening this issue, @Liwink!

Field	Value
Issue	#2662
GitHub user ID	`6659863`
Reporter	@Liwink

A maintainer will triage this when possible. To help us respond faster, please include:

Mooncake version or commit SHA
Environment (OS, CUDA/driver, RDMA stack if relevant)
Steps to reproduce and expected vs. actual behavior

Useful links: Documentation · Contributing guide

This message was posted automatically by the issue bot.

@github-actions[bot] commented on GitHub (Jun 29, 2026): Thanks for opening this issue, @Liwink! | Field | Value | |-------|-------| | **Issue** | #2662 | | **GitHub user ID** | `6659863` | | **Reporter** | @Liwink | A maintainer will triage this when possible. To help us respond faster, please include: - Mooncake version or commit SHA - Environment (OS, CUDA/driver, RDMA stack if relevant) - Steps to reproduce and expected vs. actual behavior Useful links: [Documentation](https://kvcache-ai.github.io/Mooncake/) · [Contributing guide](https://github.com/kvcache-ai/Mooncake/blob/main/CONTRIBUTING.md) > This message was posted automatically by the issue bot.

bsakaguchi commented

2026-10-29 21:18:22 +00:00

Author

Owner

@staryxchen commented on GitHub (Jun 30, 2026):

Thanks for the detailed RFC! The motivation and staged-pipeline approach make sense for TPU's hardware constraints. Before diving into implementation details, I'd like to discuss two architectural questions:

1. Multi-transport collaboration: who owns the orchestration?

The proposed TpuTransport encapsulates a 3-stage pipeline:

TPU HBM → Host DRAM (PJRT)  →  Remote Host DRAM (RDMA/TCP)  →  TPU HBM (PJRT)

In the current architecture, each Transport subclass handles one hop — RdmaTransport does NIC-to-NIC, HipTransport does VRAM-to-VRAM via fabric memory, etc. They are composed as peers inside MultiTransport, not nested.

The middle hop (host DRAM → remote host DRAM) already has dedicated transports (RDMA/TCP). If TpuTransport internally invokes another transport instance, it breaks the flat composition model.

An alternative decomposition:

TpuTransport only handles HBM ↔ Host DRAM (the PJRT DMA stage)
Host-to-host transfer is delegated to the existing RDMA/TCP transport via the routing layer

This keeps each transport single-responsibility. But it raises questions about who orchestrates the full pipeline and how staging buffers are managed across the boundary.

Question: How do you envision TpuTransport interacting with RDMA/TCP for the middle hop? Should it call into another transport directly, or should there be a higher-level orchestrator (e.g., at the TransferEngine / TENT level) that chains the stages?

2. Target: Legacy `TransferEngineImpl` or TENT?

The RFC's code examples target the legacy interface (Transport::submitTransfer, registerLocalMemory). However, TENT (mooncake-transfer-engine/tent/) is actively developed with its own transport hierarchy (tent::Transport, tent::ShmTransport, etc.) and already has shared-memory infrastructure that overlaps with the proposed StagingBufferPool.

Given the direction of the project, should this land on TENT from the start? That would:

Avoid building on top of the legacy path and immediately accruing tech debt
Let you reuse TENT's existing ShmTransport for the staging buffer / IPC layer
Align with TENT's ControlService / Topology APIs directly

Or is there a reason to target the legacy path first (e.g., vLLM-TPU integration currently depends on legacy TE)?

Looking forward to your thoughts on these two points — they'll shape the PR structure significantly.

@staryxchen commented on GitHub (Jun 30, 2026): Thanks for the detailed RFC! The motivation and staged-pipeline approach make sense for TPU's hardware constraints. Before diving into implementation details, I'd like to discuss two architectural questions: ## 1. Multi-transport collaboration: who owns the orchestration? The proposed `TpuTransport` encapsulates a 3-stage pipeline: ``` TPU HBM → Host DRAM (PJRT) → Remote Host DRAM (RDMA/TCP) → TPU HBM (PJRT) ``` In the current architecture, each Transport subclass handles **one hop** — `RdmaTransport` does NIC-to-NIC, `HipTransport` does VRAM-to-VRAM via fabric memory, etc. They are composed as peers inside `MultiTransport`, not nested. The middle hop (host DRAM → remote host DRAM) already has dedicated transports (RDMA/TCP). If `TpuTransport` internally invokes another transport instance, it breaks the flat composition model. An alternative decomposition: - `TpuTransport` only handles **HBM ↔ Host DRAM** (the PJRT DMA stage) - Host-to-host transfer is delegated to the existing RDMA/TCP transport via the routing layer This keeps each transport single-responsibility. But it raises questions about who orchestrates the full pipeline and how staging buffers are managed across the boundary. **Question:** How do you envision `TpuTransport` interacting with RDMA/TCP for the middle hop? Should it call into another transport directly, or should there be a higher-level orchestrator (e.g., at the `TransferEngine` / TENT level) that chains the stages? ## 2. Target: Legacy `TransferEngineImpl` or TENT? The RFC's code examples target the legacy interface (`Transport::submitTransfer`, `registerLocalMemory`). However, TENT (`mooncake-transfer-engine/tent/`) is actively developed with its own transport hierarchy (`tent::Transport`, `tent::ShmTransport`, etc.) and already has shared-memory infrastructure that overlaps with the proposed `StagingBufferPool`. Given the direction of the project, should this land on TENT from the start? That would: - Avoid building on top of the legacy path and immediately accruing tech debt - Let you reuse TENT's existing `ShmTransport` for the staging buffer / IPC layer - Align with TENT's `ControlService` / `Topology` APIs directly Or is there a reason to target the legacy path first (e.g., vLLM-TPU integration currently depends on legacy TE)? --- Looking forward to your thoughts on these two points — they'll shape the PR structure significantly.

bsakaguchi commented

2026-10-29 21:18:22 +00:00

Author

Owner

@staryxchen commented on GitHub (Jun 30, 2026):

Thoughts on simplifying the implementation path

Thanks again for this RFC — the problem analysis and staged-pipeline design are solid, and TPU support is something we'd love to see in Mooncake.

After digging into the TENT internals, I wanted to share an observation that might actually make your life easier: TENT already has a staging mechanism (ProxyManager) that implements the same 3-stage pattern you described. This could significantly reduce the amount of code you need to write.

How ProxyManager works today

ProxyManager (tent/src/runtime/proxy_manager.cpp) handles the case where device memory can't be directly accessed by the NIC — exactly the TPU situation:

submitLocalStage()  → Device memory ↔ Host staging buffer (currently CUDA-based)
submitCrossStage()  → Host DRAM ↔ Remote Host DRAM (delegates to RDMA/TCP)
waitRemoteStage()   → Remote staging buffer ↔ Remote device memory

It already provides chunked double-buffering, async status tracking (TaskInfo::staging / staging_status), and stage buffer lifecycle management. The pipeline overlap optimization you described in §3.2 is essentially built-in.

A potentially lighter approach

Rather than implementing a full Transport subclass (with peer discovery, memory export, batch state machine, routing capabilities, etc.), you might be able to achieve TPU support by plugging PJRT into ProxyManager's existing staging path:

PjrtLoader — dlopen + symbol resolution (you already designed this well in §3.3.1)
MTYPE_TPU — extend the MemoryType enum so the routing layer knows when staging is needed
PJRT backend for the local staging step — where ProxyManager currently uses cudaMemcpy for GPU↔DRAM, add a PJRT-based path for TPU↔DRAM

This could bring the scope down to ~200–300 LOC in a single PR, while still delivering full end-to-end TPU KVCache transfer. The USE_TPU build gating and dlopen strategy from your RFC would carry over as-is.

Why this might be a better fit

The existing Transport implementations (NVLink, MNNVL, SHM, RDMA) all share a common trait: they are complete peer-to-peer data channels with handle exchange, remote memory mapping, and direct transfer capability. PJRT's BufferToHostBuffer / BufferFromHostBuffer is more of a local device I/O operation (similar in nature to cudaMemcpy D2H/H2D) — it doesn't have a remote peer or handle exchange. ProxyManager was designed precisely for wrapping this kind of local staging operation and chaining it with a real transport.

Suggested exploration

If this direction sounds reasonable, I'd suggest starting with:

Read through ProxyManager::transferEventLoop() to see how staging is orchestrated today
Check whether PJRT-allocated host buffers are compatible with ibv_reg_mr (for RDMA registration)
Prototype the PJRT staging path — the scope should be small enough for a single focused PR

Of course, if you've identified reasons why ProxyManager doesn't fit the TPU case (e.g., specific lifetime or threading constraints from the PJRT runtime), happy to discuss alternatives. Just wanted to flag this existing infrastructure since it seems like a shortcut.

@staryxchen commented on GitHub (Jun 30, 2026): ## Thoughts on simplifying the implementation path Thanks again for this RFC — the problem analysis and staged-pipeline design are solid, and TPU support is something we'd love to see in Mooncake. After digging into the TENT internals, I wanted to share an observation that might actually make your life easier: **TENT already has a staging mechanism (`ProxyManager`) that implements the same 3-stage pattern you described.** This could significantly reduce the amount of code you need to write. ### How ProxyManager works today `ProxyManager` (`tent/src/runtime/proxy_manager.cpp`) handles the case where device memory can't be directly accessed by the NIC — exactly the TPU situation: ``` submitLocalStage() → Device memory ↔ Host staging buffer (currently CUDA-based) submitCrossStage() → Host DRAM ↔ Remote Host DRAM (delegates to RDMA/TCP) waitRemoteStage() → Remote staging buffer ↔ Remote device memory ``` It already provides chunked double-buffering, async status tracking (`TaskInfo::staging` / `staging_status`), and stage buffer lifecycle management. The pipeline overlap optimization you described in §3.2 is essentially built-in. ### A potentially lighter approach Rather than implementing a full Transport subclass (with peer discovery, memory export, batch state machine, routing capabilities, etc.), you might be able to achieve TPU support by plugging PJRT into ProxyManager's existing staging path: 1. **`PjrtLoader`** — `dlopen` + symbol resolution (you already designed this well in §3.3.1) 2. **`MTYPE_TPU`** — extend the `MemoryType` enum so the routing layer knows when staging is needed 3. **PJRT backend for the local staging step** — where ProxyManager currently uses `cudaMemcpy` for GPU↔DRAM, add a PJRT-based path for TPU↔DRAM This could bring the scope down to ~200–300 LOC in a single PR, while still delivering full end-to-end TPU KVCache transfer. The `USE_TPU` build gating and dlopen strategy from your RFC would carry over as-is. ### Why this might be a better fit The existing Transport implementations (NVLink, MNNVL, SHM, RDMA) all share a common trait: they are **complete peer-to-peer data channels** with handle exchange, remote memory mapping, and direct transfer capability. PJRT's `BufferToHostBuffer` / `BufferFromHostBuffer` is more of a local device I/O operation (similar in nature to `cudaMemcpy D2H/H2D`) — it doesn't have a remote peer or handle exchange. ProxyManager was designed precisely for wrapping this kind of local staging operation and chaining it with a real transport. ### Suggested exploration If this direction sounds reasonable, I'd suggest starting with: 1. Read through `ProxyManager::transferEventLoop()` to see how staging is orchestrated today 2. Check whether PJRT-allocated host buffers are compatible with `ibv_reg_mr` (for RDMA registration) 3. Prototype the PJRT staging path — the scope should be small enough for a single focused PR Of course, if you've identified reasons why ProxyManager doesn't fit the TPU case (e.g., specific lifetime or threading constraints from the PJRT runtime), happy to discuss alternatives. Just wanted to flag this existing infrastructure since it seems like a shortcut.

bsakaguchi commented

2026-10-29 21:18:22 +00:00

Author

Owner

@Liwink commented on GitHub (Jun 30, 2026):

Hi Stary, Thanks a lot for detailed suggestion. Plugging PJRT into ProxyManager makes a lot sense!

I'm reading through ProxyManager::transferEventLoop() now to start prototyping the PJRT staging path. Thanks again for pointing out this shortcut which makes the architecture so much cleaner!

@Liwink commented on GitHub (Jun 30, 2026): Hi Stary, Thanks a lot for detailed suggestion. Plugging PJRT into `ProxyManager` makes a lot sense! I'm reading through ProxyManager::transferEventLoop() now to start prototyping the PJRT staging path. Thanks again for pointing out this shortcut which makes the architecture so much cleaner!

bsakaguchi commented

2026-10-29 21:18:22 +00:00

Author

Owner

@Liwink commented on GitHub (Jul 3, 2026):

Hi @staryxchen — following up on your ProxyManager suggestion. I've prototyped it and have first PR ready to submit; posting a summary here first to confirm alignment before opening it.

TL;DR

The ProxyManager direction works well and the staging pipeline is reused as-is — no new pipeline code, and no standalone networked TpuTransport like my original RFC. It lands entirely on TENT (not legacy TE), gated behind -DUSE_TPU (OFF by default), so CI and existing builds are unaffected.

One architectural finding worth sharing: the scope ended up a bit larger than "plug PJRT into the copy step," for a reason that only shows up when tracing the routing layer.

Finding: the copy seam is clean, but routing needs a capability provider

ProxyManager doesn't call cudaMemcpy directly. The local device↔host stage is submitted as a normal LOCAL_SEGMENT_ID sub-transfer (ProxyManager::submitLocalStage → submitStagingTransfer → submitTransfer), which is then routed like any other transfer — and routing only selects a transport that advertises the matching capability (gpu_to_dram / dram_to_gpu) via checkAvailability / TransportSelector::isTransportAvailable. The bytes actually move inside the selected transport's submitTransferTasks, which calls Platform::copy.

For CUDA, the transport answering "yes, I can do gpu_to_dram" is NVLinkTransport. On a TPU host, nothing answers it: ShmTransport only advertises dram_to_dram, and RdmaTransport only sets GPU caps when nvidia_peermem is present. So a PJRT-backed Platform::copy would exist but never be reached.

The fix is a thin local-stage executor transport — it advertises only gpu_to_dram/dram_to_gpu (never gpu_to_gpu, so cross-node TPU traffic is always staged through host DRAM) and its submitTransferTasks just runs Platform::copy for LOCAL_SEGMENT_ID requests. No network path, no peer discovery, no handle exchange. Most of its line count is the Transport interface surface, not logic.

PR #1 contents (~1.1k LOC incl. tests + docs; core is ~500)

MTYPE_TPU memory type + "tpu:N" location parsing, wired into both routing modes' device-type checks (isGpuType in transfer_engine_impl.cpp and the selector's is_gpu — they're separate and both need it).
TpuPlatform (derives CpuPlatform): host DRAM + NIC topology inherited; only device-aware paths overridden (copy, getMemoryType, getLocation, per-device MemEntry in probe).
Thin TpuTransport as above.
findStagingPolicy TPU case: local HBM↔host via the TPU transport, host↔host via whichever host transport is present. Note: typical cloud TPU deployments are TCP/multi-NIC, not RDMA, so the case is gated on RDMA || TCP and the cross stage is routed by dram_to_dram capability rather than hardcoded — TCP is selected when RDMA is absent, RDMA used if available.
Runtime-loaded PJRT adapter: all PJRT device I/O (HBM↔host DMA, pointer classification, device topology) sits behind a small C ABI (tpu_pjrt_abi.h) resolved via dlopen (TpuPjrtShim). TENT keeps zero build-time PJRT/XLA dependency, matching the RFC's original intent. Adapter path defaults to libmooncake_tpu_pjrt.so, overridable via MC_TPU_PJRT_LIB; if absent, the build still links and runs (TPU ops return non-OK Status).
Mock adapter + gtest exercising the shim ABI (pointer classification, D2H/H2D round-trip, topology) on plain Linux — no TPU hardware or PJRT runtime needed, per the testing approach in the RFC.
Docs: an experimental entry in supported-protocols.md, plus a README.md (adapter contract) and STAGING_ARCHITECTURE.md (the routing/capability walkthrough above, with code pointers) next to the platform code.

This also answers the ibv_reg_mr question you raised: staging buffers are allocated and RDMA-registered by TENT itself (ProxyManager::allocateStageBuffers → impl_->registerLocalMemory), and the adapter copies into that already-registered buffer — so there's no PJRT-allocated host memory on the RDMA path in PR #1.

Confirming three things before I open the PR

Shape — is a dedicated thin local-stage transport the right call, vs. teaching ShmTransport/the capability layer to handle device↔host generically? I chose the separate transport to avoid changing SHM semantics for all platforms.
Target — confirming TENT (not legacy TE) from the start, per your point about avoiding tech debt on the legacy path.
Scope split — first PR = the C++/TENT data path above. Follow-ups: serving-framework integration (JAX first; the PyTorch/XLA allocator-collision issue is framework-specific, so scoping the first backend to JAX avoids it), DMA-mapped pinned staging buffers for true async device DMA (plus the DMA-map ↔ RDMA-registration interaction on the same buffer), and on-hardware benchmarks.

If this looks right I'll open the PR — happy to break it into smaller pieces if you'd prefer.

@Liwink commented on GitHub (Jul 3, 2026): Hi @staryxchen — following up on your ProxyManager suggestion. I've prototyped it and have first PR ready to submit; posting a summary here first to confirm alignment before opening it. ## TL;DR The ProxyManager direction works well and the staging pipeline is reused as-is — no new pipeline code, and no standalone networked `TpuTransport` like my original RFC. It lands entirely on **TENT** (not legacy TE), gated behind `-DUSE_TPU` (OFF by default), so CI and existing builds are unaffected. One architectural finding worth sharing: the scope ended up a bit larger than "plug PJRT into the copy step," for a reason that only shows up when tracing the routing layer. ## Finding: the copy seam is clean, but routing needs a capability provider `ProxyManager` doesn't call `cudaMemcpy` directly. The local device↔host stage is submitted as a normal `LOCAL_SEGMENT_ID` sub-transfer (`ProxyManager::submitLocalStage` → `submitStagingTransfer` → `submitTransfer`), which is then **routed like any other transfer** — and routing only selects a transport that *advertises* the matching capability (`gpu_to_dram` / `dram_to_gpu`) via `checkAvailability` / `TransportSelector::isTransportAvailable`. The bytes actually move inside the selected transport's `submitTransferTasks`, which calls `Platform::copy`. For CUDA, the transport answering "yes, I can do `gpu_to_dram`" is `NVLinkTransport`. On a TPU host, nothing answers it: `ShmTransport` only advertises `dram_to_dram`, and `RdmaTransport` only sets GPU caps when `nvidia_peermem` is present. So a PJRT-backed `Platform::copy` would exist but never be reached. The fix is a **thin local-stage executor transport** — it advertises only `gpu_to_dram`/`dram_to_gpu` (never `gpu_to_gpu`, so cross-node TPU traffic is always staged through host DRAM) and its `submitTransferTasks` just runs `Platform::copy` for `LOCAL_SEGMENT_ID` requests. No network path, no peer discovery, no handle exchange. Most of its line count is the `Transport` interface surface, not logic. ## PR #1 contents (~1.1k LOC incl. tests + docs; core is ~500) 1. **`MTYPE_TPU`** memory type + `"tpu:N"` location parsing, wired into both routing modes' device-type checks (`isGpuType` in `transfer_engine_impl.cpp` and the selector's `is_gpu` — they're separate and both need it). 2. **`TpuPlatform`** (derives `CpuPlatform`): host DRAM + NIC topology inherited; only device-aware paths overridden (`copy`, `getMemoryType`, `getLocation`, per-device `MemEntry` in `probe`). 3. **Thin `TpuTransport`** as above. 4. **`findStagingPolicy` TPU case**: local HBM↔host via the TPU transport, host↔host via whichever host transport is present. Note: typical cloud TPU deployments are **TCP/multi-NIC, not RDMA**, so the case is gated on `RDMA || TCP` and the cross stage is routed by `dram_to_dram` capability rather than hardcoded — TCP is selected when RDMA is absent, RDMA used if available. 5. **Runtime-loaded PJRT adapter**: all PJRT device I/O (HBM↔host DMA, pointer classification, device topology) sits behind a small C ABI (`tpu_pjrt_abi.h`) resolved via `dlopen` (`TpuPjrtShim`). TENT keeps **zero build-time PJRT/XLA dependency**, matching the RFC's original intent. Adapter path defaults to `libmooncake_tpu_pjrt.so`, overridable via `MC_TPU_PJRT_LIB`; if absent, the build still links and runs (TPU ops return non-OK Status). 6. **Mock adapter + gtest** exercising the shim ABI (pointer classification, D2H/H2D round-trip, topology) on plain Linux — no TPU hardware or PJRT runtime needed, per the testing approach in the RFC. 7. **Docs**: an experimental entry in `supported-protocols.md`, plus a `README.md` (adapter contract) and `STAGING_ARCHITECTURE.md` (the routing/capability walkthrough above, with code pointers) next to the platform code. This also answers the `ibv_reg_mr` question you raised: staging buffers are allocated and RDMA-registered by TENT itself (`ProxyManager::allocateStageBuffers` → `impl_->registerLocalMemory`), and the adapter copies into that already-registered buffer — so there's no PJRT-allocated host memory on the RDMA path in PR #1. ## Confirming three things before I open the PR 1. **Shape** — is a dedicated thin local-stage transport the right call, vs. teaching `ShmTransport`/the capability layer to handle device↔host generically? I chose the separate transport to avoid changing SHM semantics for all platforms. 2. **Target** — confirming TENT (not legacy TE) from the start, per your point about avoiding tech debt on the legacy path. 3. **Scope split** — first PR = the C++/TENT data path above. Follow-ups: serving-framework integration (JAX first; the PyTorch/XLA allocator-collision issue is framework-specific, so scoping the first backend to JAX avoids it), DMA-mapped pinned staging buffers for true async device DMA (plus the DMA-map ↔ RDMA-registration interaction on the same buffer), and on-hardware benchmarks. If this looks right I'll open the PR — happy to break it into smaller pieces if you'd prefer.

bsakaguchi commented

2026-10-29 21:18:22 +00:00

Author

Owner

@staryxchen commented on GitHub (Jul 3, 2026):

Re: architecture and PR scope

Thanks for the thorough writeup — the analysis and code pointers are precise, and the capability-provider observation is a genuinely nice architectural insight. Platform::copy alone would never be reached without a transport advertising the capability; glad you caught this before writing the code.

On your three questions

1. Shape — thin dedicated transport vs. teaching ShmTransport / capability layer generically. Dedicated thin transport is the right call. Teaching ShmTransport to advertise gpu_to_dram/dram_to_gpu would change what SHM means to every existing deployment (a user configuring TCP+SHM would suddenly see SHM claim GPU staging). A generic "device-local staging transport" refactor would be premature — worth revisiting only when a second non-CUDA accelerator actually needs the same shape. Building the concrete TPU case first and generalizing later is the right order.

2. Target — TENT, not legacy TE. Confirmed, aligned with my earlier point about avoiding legacy tech debt.

3. Scope split — first PR = C++/TENT data path, follow-ups for JAX / DMA-pinned buffers / benchmarks. Well-scoped. The JAX-first choice neatly sidesteps the PyTorch/XLA symbol-collision problem the original RFC flagged (the patchelf workaround) — postponing that until it's actually needed keeps this PR focused. Splitting DMA-mapped pinned staging buffers into a follow-up also makes sense: the interaction between DMA-map and RDMA registration on the same buffer is a distinct optimization worth landing on its own with proper benchmarks.

Two small findings that may be worth double-checking

Both strengthen your approach rather than change it:

a) TransportSelector::isTransportAvailable has a hardcoded same-machine check at transport_selector.cpp:322:

if ((type == NVLINK || type == SHM) && !context.same_machine) {
    return false;  // NVLINK and SHM only work on same machine
}

Your thin TPU transport is also local-stage-only. Worth extending this check to include the TPU transport type so it can never be selected for a non-LOCAL_SEGMENT_ID request by mistake.

b) findStagingPolicy currently gates staging cases on NVLINK presence (e.g., if (transport_list_[RDMA] && transport_list_[NVLINK]) at line 1294 and similar for MNNVL at line 1316). Your TPU case will presumably use transport_list_[TPU] as the analogous gate — just flagging so the pattern is consistent with the existing cases.

Both are the kind of thing that would show up in review anyway, but easier to bake in from the start.

Go ahead with the PR

The direction is solid and the scope is well-controlled. Please open the PR — happy to review as soon as it's up. No need to break it into smaller pieces; the ~1.1k LOC as described (including tests + docs) is a coherent unit.

Also happy to dig into the DMA-mapped pinned buffer follow-up when you get to it — the RDMA-registration interaction is genuinely subtle.

@staryxchen commented on GitHub (Jul 3, 2026): ## Re: architecture and PR scope Thanks for the thorough writeup — the analysis and code pointers are precise, and the capability-provider observation is a genuinely nice architectural insight. `Platform::copy` alone would never be reached without a transport advertising the capability; glad you caught this before writing the code. ### On your three questions **1. Shape — thin dedicated transport vs. teaching ShmTransport / capability layer generically.** Dedicated thin transport is the right call. Teaching ShmTransport to advertise `gpu_to_dram`/`dram_to_gpu` would change what SHM means to every existing deployment (a user configuring TCP+SHM would suddenly see SHM claim GPU staging). A generic "device-local staging transport" refactor would be premature — worth revisiting only when a second non-CUDA accelerator actually needs the same shape. Building the concrete TPU case first and generalizing later is the right order. **2. Target — TENT, not legacy TE.** Confirmed, aligned with my earlier point about avoiding legacy tech debt. **3. Scope split — first PR = C++/TENT data path, follow-ups for JAX / DMA-pinned buffers / benchmarks.** Well-scoped. The JAX-first choice neatly sidesteps the PyTorch/XLA symbol-collision problem the original RFC flagged (the `patchelf` workaround) — postponing that until it's actually needed keeps this PR focused. Splitting DMA-mapped pinned staging buffers into a follow-up also makes sense: the interaction between DMA-map and RDMA registration on the same buffer is a distinct optimization worth landing on its own with proper benchmarks. ### Two small findings that may be worth double-checking Both strengthen your approach rather than change it: **a) `TransportSelector::isTransportAvailable` has a hardcoded same-machine check** at `transport_selector.cpp:322`: ```cpp if ((type == NVLINK || type == SHM) && !context.same_machine) { return false; // NVLINK and SHM only work on same machine } ``` Your thin TPU transport is also local-stage-only. Worth extending this check to include the TPU transport type so it can never be selected for a non-`LOCAL_SEGMENT_ID` request by mistake. **b) `findStagingPolicy` currently gates staging cases on `NVLINK` presence** (e.g., `if (transport_list_[RDMA] && transport_list_[NVLINK])` at line 1294 and similar for MNNVL at line 1316). Your TPU case will presumably use `transport_list_[TPU]` as the analogous gate — just flagging so the pattern is consistent with the existing cases. Both are the kind of thing that would show up in review anyway, but easier to bake in from the start. ### Go ahead with the PR The direction is solid and the scope is well-controlled. Please open the PR — happy to review as soon as it's up. No need to break it into smaller pieces; the ~1.1k LOC as described (including tests + docs) is a coherent unit. Also happy to dig into the DMA-mapped pinned buffer follow-up when you get to it — the RDMA-registration interaction is genuinely subtle.

bsakaguchi commented

2026-10-29 21:18:23 +00:00

Author

Owner

@Liwink commented on GitHub (Sep 1, 2026):

Quick update: working with TPU team on long-term support for mooncake - a new, more efficient TPU library is almost to use

@Liwink commented on GitHub (Sep 1, 2026): Quick update: working with TPU team on long-term support for mooncake - a new, more efficient TPU library is almost to use