[Store][K8s-Native][1/N] Go-library for k8s leader election (#1910)

---------

Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com>

.github/workflows/ci.yml

@@ -134,10 +134,12 @@ jobs:
         MASTER_PID=$!
         sleep 3
         cd mooncake-store/go
-        export LD_LIBRARY_PATH=$GITHUB_WORKSPACE/build/mooncake-asio:$GITHUB_WORKSPACE/build/mooncake-store/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base:$GITHUB_WORKSPACE/build/mooncake-common/etcd
+        export LD_LIBRARY_PATH=$GITHUB_WORKSPACE/build/mooncake-common:$GITHUB_WORKSPACE/build/mooncake-store/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base:$GITHUB_WORKSPACE/build/mooncake-common/etcd
         export CGO_ENABLED=1
         export CGO_CFLAGS="-I$GITHUB_WORKSPACE/mooncake-store/include -I$GITHUB_WORKSPACE/mooncake-transfer-engine/include"
-        export CGO_LDFLAGS="-L$GITHUB_WORKSPACE/build/mooncake-store/src -L$GITHUB_WORKSPACE/build/mooncake-store/src/cachelib_memory_allocator -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base -L$GITHUB_WORKSPACE/build/mooncake-asio -L$GITHUB_WORKSPACE/build/mooncake-common/etcd -lmooncake_store -lcachelib_memory_allocator -ltransfer_engine -lbase -lasio -letcd_wrapper -lstdc++ -lnuma -lglog -lgflags -libverbs -ljsoncpp -lzstd -lcurl -luring -lasan -lm -lgcov"
+        export CGO_LDFLAGS="-L$GITHUB_WORKSPACE/build/mooncake-store/src -L$GITHUB_WORKSPACE/build/mooncake-store/src/cachelib_memory_allocator -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base -L$GITHUB_WORKSPACE/build/mooncake-common -L$GITHUB_WORKSPACE/build/mooncake-common/etcd -lmooncake_store -lcachelib_memory_allocator -ltransfer_engine -lbase -lasio -letcd_wrapper -lstdc++ -lnuma -lglog -lgflags -libverbs -ljsoncpp -lzstd -lcurl -luring -lasan -lm -lgcov"
+        # Link cudart if CUDA is available (needed for D2H staging in mooncake_store)
+        if [ -d /usr/local/cuda/lib64 ]; then export CGO_LDFLAGS="$CGO_LDFLAGS -L/usr/local/cuda/lib64 -lcudart"; fi
         ASAN_OPTIONS=detect_leaks=0:verify_asan_link_order=0 MC_METADATA_SERVER=http://127.0.0.1:8080/metadata go test -v ./tests/...
         kill $MASTER_PID 2>/dev/null || true
       shell: bash
@@ -343,6 +345,14 @@ jobs:
 
     - name: Run tests with ssd
       run: |
+        # Reserve port 50052 (mooncake_client RPC port) so the kernel never
+        # auto-allocates it as ephemeral source port for other outbound
+        # connections in the test suite. Without this, a random Python test
+        # connection can pick src_port=50052, leave a TIME_WAIT on
+        # <eth0_ip>:50052 for 60s, and block mooncake_client's bind to
+        # 0.0.0.0:50052 even with SO_REUSEADDR (Linux only relaxes
+        # TIME_WAIT+bind conflict for same-IP or loopback).
+        sudo sysctl -w net.ipv4.ip_local_reserved_ports=50052
         source test_env/bin/activate
         MC_STORE_MEMCPY=false TEST_SSD_OFFLOAD_IN_EVICT=true ./scripts/run_tests.sh
         rm -rf /tmp/mooncake_test_ssd
@@ -736,6 +746,18 @@ jobs:
     uses: ./.github/workflows/ci_cu13.yml
     secrets: inherit
 
+  ascend-test:
+    needs: [build, check-paths]
+    if: needs.check-paths.outputs.should-run-downstream == 'true'
+    uses: ./.github/workflows/ci_ascend.yml
+    secrets: inherit
+
+  integration-test:
+    needs: [build, check-paths]
+    if: needs.check-paths.outputs.should-run-downstream == 'true'
+    uses: ./.github/workflows/integration-test.yml
+    secrets: inherit
+
   ci-gate:
     name: CI Gate
     if: always()
@@ -748,6 +770,8 @@ jobs:
       - build-docker
       - test-wheel-ubuntu
       - build-wheel-cu13
+      - ascend-test
+      - integration-test
     runs-on: ubuntu-latest
     steps:
       - name: Check required job results

.github/workflows/ci_ascend.yml

@@ -64,7 +64,6 @@ jobs:
           base="${base#${base%%[![:space:]]*}}"
           base="${base%${base##*[![:space:]]}}"
           [ -n "$base" ] || return 1
-          [ "$base" = "https://github.com" ] && base="https://github.com/"
           [ "$base" != "https://github.com/" ] && base="${base%/}/"
           printf '%s\n' "$base"
         }
@@ -108,22 +107,53 @@ jobs:
 
     - name: Configure CMake
       shell: bash
+      env:
+        ASCEND_GITHUB_MIRROR_URLS: ${{ vars.ASCEND_GITHUB_MIRROR_URLS }}
       run: |
         source /usr/local/Ascend/cann-9.0.0/set_env.sh
         pwd
-        if ! git submodule update --init --recursive; then
-            if [ ! -d "extern/pybind11" ] || [ -z "$(ls -A 'extern/pybind11' 2>/dev/null)" ]; then
-                echo "git submodule update failed, try to cp pybind11..."
-                if [ -d "../pybind11" ]; then
-                    cp -r ../pybind11 extern/
-                else
-                    echo "Error: ../pybind11 does not exist. Cannot copy pybind11."
-                    exit 1
-                fi
-            else
-                echo "Detected that extern/pybind11 already exists, continuing execution...."
+
+        submodule_updated=false
+        if git submodule update --init --recursive; then
+          submodule_updated=true
+        elif [ -n "${ASCEND_GITHUB_MIRROR_URLS:-}" ]; then
+          normalize_base() {
+            local base="$1"
+            base="${base#${base%%[![:space:]]*}}"
+            base="${base%${base##*[![:space:]]}}"
+            [ -n "$base" ] || return 1
+            [ "$base" != "https://github.com/" ] && base="${base%/}/"
+            printf '%s\n' "$base"
+          }
+
+          while IFS= read -r raw; do
+            base="$(normalize_base "$raw" || true)"
+            [ -n "$base" ] || continue
+            [ "$base" = "https://github.com/" ] && continue
+
+            echo "Retrying submodule update with ${base}"
+            if git -c url."${base}https://github.com/".insteadOf=https://github.com/ \
+              submodule update --init --recursive; then
+              submodule_updated=true
+              break
             fi
+          done < <(printf '%s\n' "$ASCEND_GITHUB_MIRROR_URLS" | tr ',;' '\n')
         fi
+
+        if [ "$submodule_updated" != true ]; then
+          if [ ! -d "extern/pybind11" ] || [ -z "$(ls -A 'extern/pybind11' 2>/dev/null)" ]; then
+            echo "git submodule update failed (mirrors also exhausted), trying to cp pybind11..."
+            if [ -d "../pybind11" ]; then
+              cp -r ../pybind11 extern/
+            else
+              echo "Error: ../pybind11 does not exist. Cannot copy pybind11."
+              exit 1
+            fi
+          else
+            echo "Detected that extern/pybind11 already exists, continuing execution...."
+          fi
+        fi
+
         bash scripts/ascend/dependencies_ascend_installation.sh
         echo "Configuring CMake..."
         rm -rf build

.gitignore

@@ -5,6 +5,7 @@ build_ofed4
 old
 local_test
 go.sum
+!mooncake-common/etcd/go.sum
 *.so
 bin
 mod
@@ -204,4 +205,4 @@ CLAUDE.md
 _codeql_detected_source_root
 
 # CodeBuddy Memory
-.codebuddy/
+.codebuddy/

.pre-commit-config.yaml

@@ -23,6 +23,16 @@ repos:
       - id: check-added-large-files
         args: ['--maxkb=1024']
 
+  - repo: local
+    hooks:
+      - id: mooncake-code-format
+        name: Run Mooncake code format script
+        entry: ./scripts/code_format.sh
+        language: system
+        pass_filenames: false
+        always_run: true
+        require_serial: true
+
   - repo: https://github.com/astral-sh/ruff-pre-commit
     rev: v0.6.9
     hooks:

.typos.toml

@@ -1,10 +1,11 @@
 [default]
-extend-ignore-words = ["CANN", "ASO", "fre"]
+extend-ignore-words = ["CANN", "ASO", "fre", "wqs"]
 
 [default.extend-words]
 CANN = "CANN"
 ASO = "ASO"
 fre = "fre"
+wqs = "wqs"
 
 [files]
 extend-exclude = [

CMakeLists.txt

@@ -45,12 +45,21 @@ option(STORE_USE_REDIS "build mooncake store with redis" OFF)
 if (STORE_USE_REDIS)
   add_compile_definitions(STORE_USE_REDIS)
 endif()
+option(STORE_USE_K8S_LEASE "build mooncake store with K8s Lease leader election" OFF)
+if (STORE_USE_K8S_LEASE)
+  if (STORE_USE_ETCD)
+    message(FATAL_ERROR "STORE_USE_K8S_LEASE and STORE_USE_ETCD cannot be enabled together because both build Go c-shared HA backends.")
+  endif()
+  if (USE_ETCD AND NOT USE_ETCD_LEGACY)
+    message(FATAL_ERROR "STORE_USE_K8S_LEASE cannot be enabled with non-legacy USE_ETCD because both build Go c-shared libraries in the same process.")
+  endif()
+  add_compile_definitions(STORE_USE_K8S_LEASE)
+endif()
 
 option(STORE_USE_JEMALLOC "Use jemalloc in mooncake store master" OFF)
 
-# Define ASIO macros before adding mooncake-asio subdirectory
+# Define ASIO macros before building targets that include ASIO headers.
 add_compile_definitions(ASIO_SEPARATE_COMPILATION ASIO_DYN_LINK)
-add_subdirectory(mooncake-asio)
 
 add_subdirectory(mooncake-common)
 include_directories(mooncake-common/etcd)

CONTRIBUTING.md

@@ -41,6 +41,7 @@ Mooncake uses [pre-commit](https://pre-commit.com/) to enforce consistent format
 | Type | Tool | Purpose |
 |------|------|---------|
 | Generic | trailing-whitespace / end-of-file-fixer | Basic hygiene |
+| Project | `./scripts/code_format.sh` | Enforce Mooncake C/C++ formatting script before commit |
 | Python | ruff / ruff-format | Lint + format (includes import sorting) |
 | Spelling | codespell | Catch common typos (ignores domain-specific words) |
 | C/C++ | clang-format | Apply style from the repository's `.clang-format` |
@@ -53,6 +54,8 @@ pip install -r requirements-dev.txt
 pre-commit install
 ```
 
+After installation, every commit will run `./scripts/code_format.sh` automatically. If it rewrites files, re-stage the changes and commit again.
+
 #### Usage
 Run on all files (first run will install hook environments):
 ```bash

dependencies.sh

@@ -23,7 +23,7 @@ NC="\033[0m" # No Color
 # Configuration
 REPO_ROOT=`pwd`
 GITHUB_PROXY=${GITHUB_PROXY:-"https://github.com"}
-GOVER=1.23.8
+GOVER=1.25.9
 
 # Function to print section headers
 print_section() {

docs/source/deployment/ssd-offload.md

@@ -88,10 +88,11 @@ store.setup(
 | `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` | `/data/file_storage` | Absolute path to the SSD storage directory |
 | `MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR` | `bucket_storage_backend` | Storage backend type (see below) |
 | `MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES` | `1342177280` (1.25 GB) | Client-side staging buffer size |
+| `MOONCAKE_OFFLOAD_SCANMETA_ITERATOR_KEYS_LIMIT` | `20000` | Max keys processed per iteration when scanning existing SSD metadata on startup |
 | `MOONCAKE_OFFLOAD_TOTAL_SIZE_LIMIT_BYTES` | `2199023255552` (2 TB) | Maximum disk usage |
 | `MOONCAKE_OFFLOAD_TOTAL_KEYS_LIMIT` | `10000000` | Maximum number of objects on disk |
 | `MOONCAKE_OFFLOAD_HEARTBEAT_INTERVAL_SECONDS` | `10` | Interval for offload heartbeat to master (seconds) |
-| `MOONCAKE_USE_URING` | `false` | Enable io_uring for async file I/O |
+| `MOONCAKE_OFFLOAD_USE_URING` | `false` | Enable io_uring for async file I/O |
 
 ### Bucket backend settings
 
@@ -101,8 +102,8 @@ Applies when `MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR=bucket_storage_backend
 |---|---|---|
 | `MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES` | `268435456` (256 MB) | Max size per bucket |
 | `MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT` | `500` | Max keys per bucket |
-| `MOONCAKE_BUCKET_MAX_TOTAL_SIZE` | `0` | Eviction threshold in bytes. When set to `0`, the backend uses **90% of the physical disk capacity** as the quota — it does not mean unlimited. Set an explicit value to control disk usage precisely. |
-| `MOONCAKE_BUCKET_EVICTION_POLICY` | `none` | Eviction policy: `none` / `fifo` / `lru` |
+| `MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE` | `0` | Eviction threshold in bytes. When set to `0`, the backend uses **90% of the physical disk capacity** as the quota — it does not mean unlimited. Set an explicit value to control disk usage precisely. |
+| `MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY` | `none` | Eviction policy: `none` / `fifo` / `lru` |
 
 ---
 
@@ -127,6 +128,11 @@ Best for: general-purpose use, large-scale deployments.
 
 Stores each object in an individual file. Simple and easy to inspect, but generates many small files at scale.
 
+| Environment Variable | Default | Description |
+|---|---|---|
+| `MOONCAKE_OFFLOAD_FSDIR` | `file_per_key_dir` | Subdirectory name under `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` where objects are stored |
+| `MOONCAKE_OFFLOAD_ENABLE_EVICTION` | `true` | Enable disk eviction when the total size exceeds the quota |
+
 Best for: debugging or small-scale deployments.
 
 ### `offset_allocator_storage_backend`
@@ -143,7 +149,7 @@ Best for: high-concurrency scenarios with many small objects where restart durab
 
 ## Eviction (Bucket Backend Only)
 
-When `MOONCAKE_BUCKET_MAX_TOTAL_SIZE` is set, the backend automatically evicts buckets before writing new ones if total disk usage would exceed the limit.
+When `MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE` is set, the backend automatically evicts buckets before writing new ones if total disk usage would exceed the limit.
 
 | Policy | Behavior |
 |--------|----------|
@@ -178,8 +184,8 @@ mooncake_master \
 ```bash
 export MOONCAKE_OFFLOAD_FILE_STORAGE_PATH=/nvme/mooncake_offload
 export MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR=bucket_storage_backend
-export MOONCAKE_BUCKET_MAX_TOTAL_SIZE=$((200 * 1024 * 1024 * 1024))  # 200 GB
-export MOONCAKE_BUCKET_EVICTION_POLICY=lru
+export MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE=$((200 * 1024 * 1024 * 1024))  # 200 GB
+export MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY=lru
 
 mooncake_client \
     --master_server_address="192.168.1.10:50051" \
@@ -250,7 +256,7 @@ mooncake_client \
 
 ### "Failed to register buffer with UringFile" warning in logs
 
-This warning appears when `MOONCAKE_USE_URING=true` and the io_uring fixed-buffer registration fails. The most common cause is that `MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES` exceeds the process's locked-memory limit (`RLIMIT_MEMLOCK`). io_uring requires the registered buffer to be pinned in physical memory, which counts against this limit.
+This warning appears when `MOONCAKE_OFFLOAD_USE_URING=true` and the io_uring fixed-buffer registration fails. The most common cause is that `MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES` exceeds the process's locked-memory limit (`RLIMIT_MEMLOCK`). io_uring requires the registered buffer to be pinned in physical memory, which counts against this limit.
 
 Check the current limit:
 

docs/source/design/mooncake-store.md

@@ -36,13 +36,14 @@ It is possible to configure a `Client` instance to act in only one of its two ro
 * If `global_segment_size` is set to zero, the instance functions as a **pure client**, issuing requests but not contributing memory to the system.
 * If `local_buffer_size` is set to zero, it acts as a **pure server**, providing memory for storage. In this case, request operations such as `Get` or `Put` are not permitted from this instance.
 
-The `Client` can be used in two modes:
-1. **Embedded mode**: Runs in the same process as the LLM inference program (e.g., a vLLM instance), by being imported as a shared library.
-2. **Standalone mode**: Runs as an independent process. In this mode, the `Client` is separated into two parts: a **dummy** `Client` and a **real** `Client`: The **real** `Client` is a full-featured implementation that runs as a standalone process and directly communicates with other Mooncake Store components. It handles all RPC communications, memory management, and data transfer operations. The **real** `Client` is typically deployed on nodes that contribute memory to the distributed cache pool; The **dummy** `Client` is a lightweight wrapper that forwards all operations to a local **real** `Client` via RPC calls, which is designed for scenarios where the client needs to be embedded in the same process as the application (such as vLLM), but the actual Mooncake Store operations should be handled by a standalone process. The **dummy** `Client` and the **real** `Client` communicate via RPC calls and shared memory to make sure that Zero-copy transfers are still possible.
+The `Client` can be used in three ways:
+1. **Embedded mode**: Runs in the same process as the LLM inference program (e.g., a vLLM instance), by being imported as a shared library. Embedded clients issue requests directly, and when configured with `global_segment_size > 0` they also contribute memory resources to the cluster.
+2. **Embedded mode with dummy-real clients**: Each LLM inference **rank** holds an embedded **dummy** client (which holds no resources). Each LLM inference **instance** has one resource-owning **real** client (for example, with TP=8 there can be 8 dummy clients and 1 real client). All dummy clients of the same inference instance forward requests to that one real client. The real client owns the global segment (optionally) and is responsible for RPC handling, memory management, and data transfer. Dummy and real clients communicate via RPC, and use shared memory/zero-copy mechanisms for data transfer, so that the data path remains efficient.
+3. **Standalone store service**: A standalone store service (e.g., `python -m mooncake.mooncake_store_service`) wraps a client and provides the global memory/SSD resource pool. With this service, embedded clients can be configured with `global_segment_size = 0` so they contribute network/NIC resources only, while the standalone store service owns memory and storage management. This service can be deployed on the same server as the inference engine or on separate servers.
 
 Mooncake store supports two deployment methods to accommodate different availability requirements:
 1. **Default mode**: In this mode, the master service consists of a single master node, which simplifies deployment but introduces a single point of failure. If the master crashes or becomes unreachable, the system cannot continue to serve requests until it is restored.
-2. **High availability mode (unstable)**: This mode enhances fault tolerance by running the master service as a cluster of multiple master nodes coordinated through an etcd cluster. The master nodes use etcd to elect a leader, which is responsible for handling client requests.
+2. **High availability mode**: This mode enhances fault tolerance by running the master service as a cluster of multiple master nodes coordinated through an etcd cluster. The master nodes use etcd to elect a leader, which is responsible for handling client requests.
 If the current leader fails or becomes partitioned from the network, the remaining master nodes automatically perform a new leader election, ensuring continuous availability.
 
 In both modes, the leader monitors the health of all client nodes through periodic heartbeats. If a client crashes or becomes unreachable, the leader quickly detects the failure and takes appropriate action. When a client node recovers or reconnects, it can automatically rejoin the cluster without manual intervention.

docs/source/design/ssd-offload.md

@@ -154,7 +154,7 @@ A single pre-allocated file (`kv_cache.data`) is shared by all objects. Space wi
 
 ## Eviction (BucketStorageBackend)
 
-When `MOONCAKE_BUCKET_MAX_TOTAL_SIZE` is set, the backend evicts existing buckets to make room before writing a new one. Eviction is disabled by default (`BucketEvictionPolicy::NONE`).
+When `MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE` is set, the backend evicts existing buckets to make room before writing a new one. Eviction is disabled by default (`BucketEvictionPolicy::NONE`).
 
 ### Policies
 
@@ -195,7 +195,7 @@ This ordering guarantees:
 
 ## io_uring File I/O
 
-When `MOONCAKE_USE_URING=true`, the storage backends replace POSIX `pread`/`pwrite` calls with an io_uring-based implementation (`UringFile`). The design prioritizes eliminating inter-thread lock contention, which was the dominant latency source in the previous global-ring approach.
+When `MOONCAKE_OFFLOAD_USE_URING=true`, the storage backends replace POSIX `pread`/`pwrite` calls with an io_uring-based implementation (`UringFile`). The design prioritizes eliminating inter-thread lock contention, which was the dominant latency source in the previous global-ring approach.
 
 ### Thread-local rings (`SharedUringRing`)
 

docs/source/design/transfer-engine/efa_transport.md

@@ -6,7 +6,7 @@ This document describes how to build and use Mooncake with AWS Elastic Fabric Ad
 
 ### 1. AWS EFA Driver and libfabric
 
-EFA driver and libfabric should be pre-installed on AWS instances with EFA support (e.g., p6-b200.48xlarge, p5e.48xlarge, p4d.24xlarge).
+EFA driver and libfabric should be pre-installed on AWS instances with EFA support (e.g., p6-b300.48xlarge, p6-b200.48xlarge, p5en.48xlarge, p5e.48xlarge, p5.48xlarge).
 
 Verify installation:
 ```bash
@@ -22,48 +22,30 @@ If not installed, follow [AWS EFA documentation](https://docs.aws.amazon.com/AWS
 
 ### 2. Build Dependencies
 
-```bash
-# Ubuntu/Debian
-sudo apt-get update
-sudo apt-get install -y \
-    build-essential \
-    cmake \
-    git \
-    libgflags-dev \
-    libgoogle-glog-dev \
-    libjsoncpp-dev \
-    libnuma-dev \
-    libibverbs-dev \
-    libboost-all-dev \
-    libcurl4-openssl-dev \
-    libgtest-dev \
-    libmsgpack-dev \
-    libxxhash-dev \
-    libyaml-cpp-dev \
-    pybind11-dev \
-    python3-dev
-
-# Install yalantinglibs (required)
-cd /tmp
-git clone https://github.com/alibaba/yalantinglibs.git
-cd yalantinglibs
-mkdir build && cd build
-cmake .. -DCMAKE_INSTALL_PREFIX=/usr/local
-make -j$(nproc)
-sudo make install
-```
-
-## Building Mooncake with EFA Support
-
-### 1. Clone the Repository
+Clone the repository and install all dependencies:
 
 ```bash
 git clone https://github.com/kvcache-ai/Mooncake.git
 cd Mooncake
-git submodule update --init --recursive
+sudo ./dependencies.sh -y
 ```
 
-### 2. Build with EFA Enabled
+This installs all system packages, git submodules (including pybind11 and yalantinglibs), and Go.
+
+**Additional EFA-specific dependencies** (not covered by `dependencies.sh`):
+
+```bash
+# gflags is needed by transfer_engine_bench and EFA unit tests
+sudo apt-get install -y libgflags-dev
+```
+
+> **Note:** The EFA driver and libfabric are **not** installed by `dependencies.sh`. They must be pre-installed on the instance (see section 1 above).
+
+## Building Mooncake with EFA Support
+
+### 1. Build with EFA Enabled
+
+**GPU memory transfers (e.g., KV cache in vLLM):**
 
 ```bash
 mkdir build && cd build
@@ -78,13 +60,28 @@ make -j$(nproc)
 
 > **Note:** `-DUSE_CUDA=ON` is required when transferring GPU memory (e.g., KV cache in vLLM). Without it, the TCP transport (used as fallback when `mooncake_protocol` is set to `"tcp"`) cannot detect GPU memory and will fail with "Bad address" (EFAULT) errors.
 
-### 3. Install Python Package
+**CPU memory transfers only (no GPU dependency):**
+
+```bash
+mkdir build && cd build
+
+cmake .. \
+    -DUSE_EFA=ON \
+    -DUSE_CUDA=OFF \
+    -DCMAKE_BUILD_TYPE=RelWithDebInfo
+
+make -j$(nproc)
+```
+
+> **Note:** With `-DUSE_CUDA=OFF`, the benchmark tool uses DRAM buffers allocated via `numa_alloc_onnode`. This is useful for measuring EFA transport throughput independently of GPU hardware.
+
+### 2. Install Python Package
 
 ```bash
 # Copy built modules to wheel directory
 cp mooncake-integration/engine.cpython-*.so ../mooncake-wheel/mooncake/
 cp mooncake-integration/store.cpython-*.so ../mooncake-wheel/mooncake/
-cp mooncake-asio/libasio.so ../mooncake-wheel/mooncake/
+cp mooncake-common/libasio.so ../mooncake-wheel/mooncake/
 
 # Install with pip
 pip install -e ../mooncake-wheel --no-build-isolation
@@ -105,27 +102,6 @@ print(f'Initialize result: {result}')  # Should be 0
 # EFA device (libfabric): rdmap79s0, domain: rdmap79s0-rdm, provider: efa
 ```
 
-## Usage with vLLM
-
-### Prefill Instance
-
-```bash
-VLLM_MOONCAKE_BOOTSTRAP_PORT=8998 \
-vllm serve <model_path> -tp 8 \
-   --port 8010 \
-   --trust-remote-code \
-   --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_producer","kv_connector_extra_config":{"mooncake_protocol":"efa"}}'
-```
-
-### Decode Instance
-
-```bash
-vllm serve <model_path> -tp 8 \
-   --port 8020 \
-   --trust-remote-code \
-   --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_consumer","kv_connector_extra_config":{"mooncake_protocol":"efa"}}'
-```
-
 ## Unit Tests
 
 Run the EFA transport unit tests (requires EFA hardware):
@@ -187,6 +163,8 @@ Use `transfer_engine_bench` to measure EFA transport throughput between two node
     --report_unit=GB
 ```
 
+> **Tip:** For CPU-to-CPU benchmarks, prepend `CUDA_VISIBLE_DEVICES=""` to prevent the CUDA runtime from being initialized. Without it, `nvidia-smi` may show GPU memory usage (due to CUDA context initialization) even though the benchmark only uses DRAM.
+
 Replace `<target_hostname>:<target_port>` with the target node's address shown in the target's startup log (e.g., `ip-172-31-29-226:12345`).
 
 ### Key Parameters
@@ -196,68 +174,219 @@ Replace `<target_hostname>:<target_port>` with the target node's address shown i
 | `--block_size` | 65536 | Bytes per transfer request |
 | `--batch_size` | 128 | Requests per batch |
 | `--threads` | 12 | Concurrent submission threads |
-| `--buffer_size` | 1 GB | Total buffer size |
+| `--buffer_size` | 1 GB | Total buffer size (per GPU when `--gpu_id=-1`) |
 | `--duration` | 10 | Test duration in seconds |
-| `--operation` | read | `read` or `write` |
+| `--operation` | write | `read` or `write` |
 | `--report_unit` | GB | `GB\|GiB\|Gb\|MB\|MiB\|Mb` |
+| `--gpu_id` | 0 | GPU device ID; `-1` to use all GPUs (requires `-DUSE_CUDA=ON`) |
+
+| Environment Variable | Default | Description |
+|---------------------|---------|-------------|
+| `MC_SLICE_SIZE` | 65536 | Slice size for RDMA transport. **Not used by EFA transport** (see note below). |
+| `MC_EFA_STRIPING_THRESHOLD` | 2097152 | Transfers larger than this (bytes) are striped across all NICs |
+
+> **Note on EFA slicing:** Unlike RDMA transport which splits every transfer into fixed `MC_SLICE_SIZE` chunks, EFA transport uses a different strategy: transfers ≤ `MC_EFA_STRIPING_THRESHOLD` (default 2MB) are sent as a **single `fi_write`/`fi_read`** whose size equals `block_size`; transfers larger than the threshold are striped across all NICs (one chunk per NIC). This means **`block_size` directly determines per-operation size** and is the key tuning parameter for EFA, while `MC_SLICE_SIZE` has no effect.
+
+> **Note:** `buffer_size` must be >= `block_size * batch_size * threads`. The benchmark auto-adjusts if too small.
 
 ### Benchmark Results
 
-Tested on two p6-b200.48xlarge instances (8 EFA devices each, 8×400 Gbps) in the same AWS placement group.
+#### p6-b200.48xlarge (B200, 8 EFA × 400 Gbps)
 
-#### Optimized Results
+Tested on two p6-b200.48xlarge instances in the same AWS placement group.
 
-With tuned parameters (`MC_SLICE_SIZE=262144`):
+**GPU-to-GPU** (build with `-DUSE_CUDA=ON`, `--gpu_id=-1` for all 8 GPUs):
 
-| Operation | Throughput | Configuration |
-|-----------|-----------|---------------|
-| **Write** | **167.63 GB/s** | threads=48, block_size=128KB, batch_size=128, MC_SLICE_SIZE=256KB |
-| **Read**  | **171.89 GB/s** | threads=48, block_size=128KB, batch_size=128, MC_SLICE_SIZE=256KB |
+| Configuration | Write | Read |
+|---------------|-------|------|
+| block=1MB, threads=32, batch=64, buf=2GB/GPU | 285-296 GB/s | 312 GB/s |
+| **block=1MB, threads=16, batch=128, buf=2GB/GPU** | **302 GB/s** | **313 GB/s** |
 
-#### Parameter Tuning Results
+**CPU-to-CPU** (build with `-DUSE_CUDA=OFF`):
 
-The following table shows how different parameters affect write throughput:
+| Configuration | Write | Read |
+|---------------|-------|------|
+| block=1MB, threads=32, batch=128, buf=4GB | **222 GB/s** (stable over 6 runs) | **226 GB/s** |
+
+<details>
+<summary>CPU Parameter Tuning History (p6-b200)</summary>
+
+Earlier CPU-to-CPU tuning results (before EFA striping optimization, when `MC_SLICE_SIZE` was still used by EFA):
 
 | block_size | threads | batch_size | MC_SLICE_SIZE | Throughput |
 |-----------|---------|------------|---------------|-----------|
 | 64KB  | 8  | 128 | default (64KB) | 69.47 GB/s |
-| 256KB | 8  | 128 | default        | 70.09 GB/s |
-| 64KB  | 16 | 128 | default        | 78.80 GB/s |
-| 64KB  | 32 | 256 | default        | 87.65 GB/s |
-| 64KB  | 64 | 256 | default        | 85.72 GB/s |
 | 128KB | 32 | 128 | default        | 92.33 GB/s |
-| 128KB | 32 | 128 | 128KB          | 152.26 GB/s |
 | 128KB | 32 | 128 | 256KB          | 156.18 GB/s |
-| 128KB | 48 | 128 | 256KB          | **160.34 GB/s** |
-| 128KB | 64 | 128 | 256KB          | 158.82 GB/s |
+| 128KB | 48 | 128 | 256KB          | 160.34 GB/s |
 
-Key findings:
-- **MC_SLICE_SIZE** is the most impactful tuning parameter — increasing from default 64KB to 256KB nearly **doubles** throughput (92→160 GB/s)
-- **block_size=128KB** outperforms 64KB by ~10-15%
-- **threads=48** is optimal for 8 EFA devices; 64 threads shows slight diminishing returns
-- **batch_size=128** is sufficient; increasing to 256+ causes "Cannot select device" errors at higher thread counts
+> **Note:** These results predate the EFA striping optimization. With the current code, `MC_SLICE_SIZE` no longer affects EFA performance. Use `--block_size=1048576` (1MB) instead, which achieves 222 GB/s.
+
+</details>
+
+#### p6-b300.48xlarge (B300, 16 EFA × 400 Gbps)
+
+Tested on two p6-b300.48xlarge instances (Intel Xeon Platinum 8559C, 8× B300, 16 EFA devices) in the same AWS placement group.
+
+**GPU-to-GPU** (build with `-DUSE_CUDA=ON`, `--gpu_id=-1` for all 8 GPUs, `--buffer_size=2147483648`):
+
+| Configuration | Write | Read |
+|---------------|-------|------|
+| block=1MB, threads=16, batch=128 | 701 GB/s | **697 GB/s** |
+| **block=1MB, threads=32, batch=64** | **752 GB/s** | 713 GB/s |
+| block=1MB, threads=32, batch=32 | 751 GB/s | - |
+| block=1MB, threads=64, batch=32 | 728 GB/s | - |
+
+> **Peak: 752 GB/s write**, reaching ~94% of the 800 GB/s theoretical line rate (16×400 Gbps). GPUDirect RDMA bypasses DRAM entirely (HBM3e → PCIe switch → NIC), so performance is not bottlenecked by CPU memory bandwidth.
+
+**CPU-to-CPU** (build with `-DUSE_CUDA=OFF`):
+
+| Configuration | Write | Read |
+|---------------|-------|------|
+| **block=1MB, threads=32, batch=128, buf=4GB** | **230 GB/s** | 180 GB/s |
+| block=16MB, threads=32, batch=8, buf=8GB (striping off) | 233 GB/s | - |
+
+> CPU-to-CPU is bounded by DRAM bandwidth (~250 GB/s/socket on Xeon 8559C). Per-NIC sampling shows NUMA-0 NICs at 90 Gbps and NUMA-1 NICs at 53 Gbps, confirming DRAM controller saturation rather than NIC limit.
+
+#### p5en.48xlarge (H200, 16 EFA × 200 Gbps)
+
+Tested on two p5en.48xlarge instances (Intel Xeon 8488C, 8× H200 141GB, 16 EFA devices) in the same AWS placement group.
+
+**GPU-to-GPU** (build with `-DUSE_CUDA=ON`, `--gpu_id=-1` for all 8 GPUs):
+
+| Configuration | Write | Read |
+|---------------|-------|------|
+| block=1MB, threads=8, batch=128, buf=1GB/GPU | 236 GB/s | 271 GB/s |
+| block=1MB, threads=16, batch=128, buf=2GB/GPU | 271 GB/s | **297-308 GB/s** |
+| **block=1MB, threads=32, batch=64, buf=2GB/GPU** | **337-347 GB/s** | 274 GB/s |
+
+> GPU HBM bandwidth (>3 TB/s) eliminates the memory bottleneck, allowing full EFA utilization. Write and read have different optimal thread counts: write peaks at 32 threads, read peaks at 16 threads.
+
+> **Note:** EFA memory region registration (fi_mr_reg) for GPU memory segfaults at 4GB+ per GPU. Use `--buffer_size=2147483648` (2GB) as the maximum per-GPU buffer.
+
+**CPU-to-CPU** (build with `-DUSE_CUDA=OFF`):
+
+| Configuration | Write | Read |
+|---------------|-------|------|
+| Single instance (block=1MB, threads=32, batch=128, buf=4GB) | 179 GB/s | 185 GB/s |
+| NUMA-split (block=1MB, 2 instances, 8 NICs each, threads=16, buf=2GB) | **192 GB/s** | **182 GB/s** |
+
+> CPU-to-CPU throughput is bottlenecked by DRAM bandwidth (~155 GB/s per NUMA node, measured with STREAM Copy).
 
 #### Cross-Transport Comparison
 
-| Transport | Throughput | Per-NIC Bandwidth | Notes |
-|-----------|-----------|-------------------|-------|
-| **EFA (tuned)** | **168-172 GB/s** | ~207-214 Gbps × 8 NICs | MC_SLICE_SIZE=256KB, threads=48 |
-| **EFA (default)** | **69.47 GB/s** | ~86 Gbps × 8 NICs | Default parameters |
-| TCP (iperf3 baseline) | 9.5 GB/s | 76 Gbps total | Kernel TCP stack, 8 parallel streams |
-| TCP (Mooncake) | 0.11 GB/s | — | Mooncake TCP transport, unoptimized for throughput |
+| Transport | Throughput | Notes |
+|-----------|-----------|-------|
+| **EFA GPU-to-GPU (B300)** | **752 GB/s** | p6-b300.48xlarge, 16×400G, block=1MB, ~94% line rate |
+| **EFA GPU-to-GPU (H200)** | **347 GB/s** | p5en.48xlarge, 16×200G, block=1MB |
+| **EFA GPU-to-GPU (B200)** | **313 GB/s** | p6-b200.48xlarge, 8×400G, block=1MB |
+| **EFA CPU-to-CPU (B300)** | **230 GB/s** | p6-b300.48xlarge, 16×400G, block=1MB, DRAM-limited |
+| **EFA CPU-to-CPU (B200)** | **222 GB/s** | p6-b200.48xlarge, 8×400G, block=1MB, DRAM-limited |
+| **EFA CPU-to-CPU (H200)** | **192 GB/s** | p5en.48xlarge, block=1MB, NUMA-split, DRAM-limited |
+| EFA (default params) | 69.47 GB/s | Default block=64KB |
+| TCP (iperf3 baseline) | 9.5 GB/s | Kernel TCP stack, 8 parallel streams |
 
-**EFA (tuned) vs TCP**: EFA delivers **17.7x** the raw TCP bandwidth by bypassing the kernel network stack.
-
-**EFA vs RoCE RDMA**: On comparable 8×400 Gbps RoCE networks, Mooncake's RDMA transport achieves ~190 GB/s. Tuned EFA reaches **~88%** of RoCE performance, demonstrating that proper parameter tuning can largely close the gap between SRD-based EFA and hardware-offloaded RDMA.
+**EFA vs RoCE RDMA**: On comparable 8×400 Gbps RoCE networks, Mooncake's RDMA transport achieves ~190 GB/s. Tuned EFA **exceeds** RoCE performance with GPU memory (313-347 GB/s) and on CPU-to-CPU (222 GB/s).
 
 ### Tuning Tips
 
-- **Set `MC_SLICE_SIZE=262144` (256KB)** — this is the single most important tuning knob, nearly doubling throughput from defaults
-- Increase `--threads` to 32-48 to saturate multiple EFA devices (6 threads per device is a good starting point)
-- Use `--block_size=131072` (128KB) for optimal per-request efficiency
-- Keep `--batch_size=128`; higher values may cause device selection failures with many threads
-- Allocate buffers on both NUMA nodes for balanced NIC utilization (the bench tool does this by default)
-- Avoid `--block_size=256KB` or larger with many threads — this can trigger "Cannot select device" errors due to buffer boundary alignment across 8 EFA devices
+- **Use `--block_size=1048576` (1MB)** — this is the most important tuning parameter for EFA. Each `block_size`-sized transfer becomes a single `fi_write`/`fi_read` call, so larger blocks amortize per-operation overhead. 1MB gives ~2× throughput over the 64KB default.
+- `MC_SLICE_SIZE` has **no effect** on EFA transport (it only applies to RDMA transport). Use `block_size` instead.
+- Increase `--threads` to 32-48 to saturate multiple EFA devices (2-4 threads per device is a good starting point)
+- For **CPU-to-CPU**: use `--block_size=1048576` (1MB) with NUMA-split (separate instances per NUMA node) for best results
+- For **GPU-to-GPU**: use `--block_size=1048576` (1MB), `--gpu_id=-1` (all GPUs), and `--buffer_size=2147483648` (2GB max per GPU). Write peaks at threads=32, read at threads=16
+- Keep `--batch_size` such that `block_size * batch_size * threads <= buffer_size`
+- Allocate buffers on both NUMA nodes for balanced NIC utilization (the bench tool does this by default for CPU mode)
+- On 16-NIC instances (p5en), writes are NUMA-sensitive: 8 local-NUMA NICs reach 90 Gbps each, while 8 cross-NUMA NICs only reach ~20 Gbps without NUMA-split
+
+### Eager endpoint warmup (first-request latency)
+
+libfabric `FI_EP_RDM` endpoints resolve peer addresses lazily: `fi_av_insert()` and the metadata handshake fire on the first send to each `(local_ctx, peer_nic)` pair. On 16-NIC instances that gives `16 × N_peer_NICs` serial handshakes inside the first `submitTransfer`, which shows up as a single-digit-second first-batch stall (measured ~4 s on p6-B300 for a 100 × 0.5 MB batch; the first batch runs at <0.1 GB/s while the CQ drains, steady-state afterwards is unaffected).
+
+Mooncake exposes an explicit eager-warmup API to eliminate the stall:
+
+- C++: `EfaTransport::warmupSegment(const std::string& segment_name)`
+- C: `int warmupEfaSegment(transfer_engine_t engine, const char *segment_name)`
+- Rust: `TransferEngine::warmup_efa_segment(name: &str)`
+
+Call it once per peer segment, right after `openSegment` (or after any metadata change that adds a new peer). Every `(local_ctx, peer_nic)` endpoint is connected concurrently via `std::async`; the critical path becomes `max(handshake RTT)` instead of `sum(handshake RTT)`. The call is idempotent — safe to re-run.
+
+Measured on p6-B300 (16 local NICs × 16 peer NICs, dual-NUMA initiator, 100 × 0.5 MB batch):
+
+| | first-batch latency | steady-state |
+|---|---:|---:|
+| No warmup | 4,043 ms | 141 GB/s |
+| `warmup_efa_segment` (256 endpoints connected in 4.1 s) | **13.5 ms** (~300×) | 230 GB/s |
+
+The warmup call itself takes roughly the same wall time as the stall it replaces — the win is that it's a one-time setup cost decoupled from the critical path of the first real transfer, not paid inside your latency budget.
+
+## Usage with vLLM
+
+### Prefill Instance
+
+```bash
+VLLM_MOONCAKE_BOOTSTRAP_PORT=8998 \
+vllm serve <model_path> -tp 8 \
+   --port 8010 \
+   --trust-remote-code \
+   --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_producer","kv_connector_extra_config":{"mooncake_protocol":"efa"}}'
+```
+
+### Decode Instance
+
+```bash
+vllm serve <model_path> -tp 8 \
+   --port 8020 \
+   --trust-remote-code \
+   --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_consumer","kv_connector_extra_config":{"mooncake_protocol":"efa"}}'
+```
+
+## Usage with SGLang
+
+SGLang's Mooncake integration currently hardcodes the `"rdma"` protocol. To use EFA transport, apply the provided patch and set environment variables.
+
+### 1. Apply EFA Patch
+
+SGLang's transfer engine initialization needs to be patched to read the protocol from an environment variable instead of using hardcoded `"rdma"`. Use the [patch script](https://github.com/whn09/kimi-k2-sglang):
+
+```bash
+bash patch_sglang_efa.sh
+```
+
+This is idempotent and safe to rerun.
+
+### 2. Environment Variables
+
+```bash
+export MOONCAKE_PROTOCOL=efa
+export FI_PROVIDER=efa
+export FI_EFA_USE_DEVICE_RDMA=1
+export GLOO_SOCKET_IFNAME=enp71s0  # adjust to your instance's primary interface
+```
+
+For multi-node expert parallelism (EP) deployments, also set:
+
+```bash
+export NVSHMEM_REMOTE_TRANSPORT=libfabric
+export NVSHMEM_LIBFABRIC_PROVIDER=efa
+```
+
+> **Warning:** Do **not** set NVSHMEM variables on single-node deployments — doing so causes segmentation faults.
+
+### 3. Docker Launch Example
+
+```bash
+docker run -d --name sglang \
+  --runtime=nvidia --gpus all --network host \
+  --privileged --shm-size=600g \
+  --device=/dev/infiniband \
+  -e MOONCAKE_PROTOCOL=efa \
+  -e FI_PROVIDER=efa \
+  -e FI_EFA_USE_DEVICE_RDMA=1 \
+  <image> bash start.sh
+```
+
+> **Note:** Ensure the Docker image's libfabric version matches the host's EFA driver. If not, mount the host's EFA libraries into the container (see [Troubleshooting](#libfabric-version-mismatch-in-docker)).
 
 ## Technical Details
 
@@ -290,11 +419,11 @@ AWS EFA exposes RDMA-like devices through the ibverbs interface, but does not su
 
 ### Thread Safety
 
-The EFA transport requests `FI_THREAD_SAFE` from the libfabric provider and adds per-endpoint spinlocks to serialize `fi_write` calls. This is necessary because:
+The EFA transport requests `FI_THREAD_SAFE` from the libfabric provider and adds per-endpoint spinlocks to serialize `fi_write`/`fi_read` calls. This is necessary because:
 
 - Multiple submission threads may route slices to the same endpoint concurrently
 - libfabric RDM endpoints default to `FI_THREAD_UNSPEC` (no thread safety guarantees)
-- Concurrent `fi_write` without serialization corrupts provider internals, causing completions to silently vanish
+- Concurrent `fi_write`/`fi_read` without serialization corrupts provider internals, causing completions to silently vanish
 
 CQ completion queues are polled by dedicated worker threads (one per EFA device) that run independently of submission threads.
 
@@ -306,14 +435,18 @@ CQ completion queues are polled by dedicated worker threads (one per EFA device)
 | Endpoint type | `FI_EP_RDM` (message-based) | Queue Pairs (true RDMA) |
 | Write operation | Software-emulated via messages + ACKs | Hardware-offloaded one-sided RDMA |
 | CPU overhead | Moderate (provider processes ACKs) | Minimal (NIC handles everything) |
-| Throughput (8×400G) | ~170 GB/s (tuned) | ~190 GB/s |
+| Throughput CPU-to-CPU (8×400G) | 222 GB/s (tuned) | ~190 GB/s |
+| Throughput GPU-to-GPU (16×200G) | 347 GB/s (tuned) | N/A |
+| Throughput GPU-to-GPU (8×400G) | 313 GB/s (tuned) | N/A |
 | AWS availability | All EFA-enabled instances | Not available on AWS |
 
 ### Supported AWS Instance Types
 
-- p6-b200.48xlarge (8 EFA devices, `rdmap*` naming)
-- p5e.48xlarge (16 EFA devices, `rdmap*` naming)
-- p4d.24xlarge (4 EFA devices)
+- p6-b300.48xlarge (16 EFA devices × 400 Gbps = 6,400 Gbps, `rdmap*` naming)
+- p6-b200.48xlarge (8 EFA devices × 400 Gbps = 3,200 Gbps, `rdmap*` naming)
+- p5en.48xlarge (16 EFA devices × 200 Gbps = 3,200 Gbps, `rdmap*` naming)
+- p5e.48xlarge (32 EFA devices × 100 Gbps = 3,200 Gbps, `rdmap*` naming)
+- p5.48xlarge (32 EFA devices × 100 Gbps = 3,200 Gbps, `rdmap*` naming)
 - Other EFA-enabled instances
 
 Use `fi_info -p efa` to list available EFA devices on your instance.
@@ -354,3 +487,59 @@ If `transfer_engine_bench` hangs with some workers never completing:
 1. **Ensure both nodes are running the same build** — the CQ backpressure and thread-safety fixes must be present on both sides
 2. **Reduce concurrency** to verify basic connectivity: `--threads=1 --batch_size=16`
 3. **Check CQ poller threads**: logs should show "Started N CQ polling worker threads" where N matches the number of EFA devices
+
+### Building on AWS Deep Learning AMI
+
+On AWS Deep Learning AMI (e.g., Ubuntu 24.04), the system Python and CUDA toolkit are bundled inside the `/opt/pytorch` virtual environment. You must activate it and set CUDA paths before building:
+
+```bash
+# Activate the PyTorch environment (provides Python 3.13 + CUDA toolkit)
+source /opt/pytorch/bin/activate
+
+# Set CUDA paths (nvcc, headers and libs are inside the pip-installed nvidia packages)
+export CUDA_HOME=/opt/pytorch/lib/python3.13/site-packages/nvidia/cu13
+export PATH=$CUDA_HOME/bin:$PATH
+export CPLUS_INCLUDE_PATH=$CUDA_HOME/include:$CPLUS_INCLUDE_PATH
+export LD_LIBRARY_PATH=$CUDA_HOME/lib:$LD_LIBRARY_PATH
+export LIBRARY_PATH=$CUDA_HOME/lib:$LIBRARY_PATH
+
+# Build with CUDA support
+cd ~/Mooncake
+mkdir -p build && cd build
+cmake .. -DUSE_EFA=ON -DUSE_CUDA=ON -DCMAKE_BUILD_TYPE=RelWithDebInfo
+make -j$(nproc)
+```
+
+Without activating the environment, you may encounter:
+- `Could not find nvcc, please set CUDAToolkit_ROOT` — nvcc is not in PATH
+- `fatal error: cuda.h: No such file or directory` — CUDA headers not in include path, set `CPLUS_INCLUDE_PATH`
+- `cannot find -lcudart: No such file or directory` — CUDA libs not in library path, set `LIBRARY_PATH` and `LD_LIBRARY_PATH`
+- `ModuleNotFoundError: No module named 'mooncake.engine'` — `.so` built against wrong Python version (e.g., 3.12 vs 3.13)
+
+### libfabric version mismatch in Docker
+
+```
+fi_ep_bind (av) failed: Function not implemented
+```
+
+or:
+
+```
+undefined reference to `efadv_query_qp_wqs@EFA_1.4'
+```
+
+This happens when the Docker container's libfabric version is older than the host's EFA driver. Check with `fi_info --version` on both host and container.
+
+Solution: Mount the host's EFA libraries into the container:
+
+```bash
+docker run --gpus all --device=/dev/infiniband --net=host --privileged \
+  -v /opt/amazon/efa:/opt/amazon/efa \
+  -v /lib/x86_64-linux-gnu/libefa.so.1:/lib/x86_64-linux-gnu/libefa.so.1 \
+  -v /lib/x86_64-linux-gnu/libefa.so:/lib/x86_64-linux-gnu/libefa.so \
+  -v /lib/x86_64-linux-gnu/libibverbs.so.1:/lib/x86_64-linux-gnu/libibverbs.so.1 \
+  -e LD_LIBRARY_PATH=/opt/amazon/efa/lib:$LD_LIBRARY_PATH \
+  -it <image>
+```
+
+Then rebuild Mooncake inside the container to link against the host's libfabric.

docs/source/design/transfer-engine/kunpeng_ub_transport.md

@@ -99,7 +99,7 @@ make -j$(nproc)
 # Copy built modules to wheel directory
 cp mooncake-integration/engine.cpython-*.so ../mooncake-wheel/mooncake/
 cp mooncake-integration/store.cpython-*.so ../mooncake-wheel/mooncake/
-cp mooncake-asio/libasio.so ../mooncake-wheel/mooncake/
+cp mooncake-common/libasio.so ../mooncake-wheel/mooncake/
 
 # Install with pip
 pip install -e ../mooncake-wheel --no-build-isolation

docs/source/performance/ssd-offload-benchmark-results.md

@@ -141,7 +141,7 @@ mooncake_master \
 # SSD Offload Buffer: 20GB
 MOONCAKE_OFFLOAD_FILE_STORAGE_PATH="/mnt/data/file_storage" \
 MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES=21474836480 \
-MOONCAKE_USE_URING=1 \
+MOONCAKE_OFFLOAD_USE_URING=1 \
 mooncake_client \
     --host=127.0.0.1 \
     --global_segment_size=80GB \

mooncake-asio/CMakeLists.txt

@@ -1,50 +0,0 @@
-# Build asio as a shared library to avoid ODR violations
-# when multiple shared libraries use asio
-
-# Try to find ASIO using find_package first
-find_package(asio QUIET)
-
-if(asio_FOUND)
-    message(STATUS "Found ASIO via find_package")
-    set(ASIO_INCLUDE_DIR ${asio_INCLUDE_DIR})
-else()
-    # Fallback to find_path if find_package fails
-    find_path(ASIO_INCLUDE_DIR
-        NAMES asio.hpp
-        PATHS
-            /usr/local/include
-            /usr/include
-            ${CMAKE_INSTALL_PREFIX}/include
-        DOC "Path to ASIO headers"
-    )
-
-    if(NOT ASIO_INCLUDE_DIR)
-        message(FATAL_ERROR "ASIO not found. Please install ASIO or set ASIO_INCLUDE_DIR manually.")
-    endif()
-
-    message(STATUS "Found ASIO at: ${ASIO_INCLUDE_DIR}")
-endif()
-
-add_library(asio_shared SHARED asio_impl.cpp)
-
-target_compile_definitions(asio_shared
-    PUBLIC
-        ASIO_SEPARATE_COMPILATION
-        ASIO_DYN_LINK
-)
-
-target_include_directories(asio_shared
-    PUBLIC
-        ${ASIO_INCLUDE_DIR}
-)
-
-set_target_properties(asio_shared PROPERTIES
-    POSITION_INDEPENDENT_CODE ON
-    INSTALL_RPATH "$ORIGIN"
-    BUILD_WITH_INSTALL_RPATH TRUE
-    OUTPUT_NAME "asio"
-)
-
-target_link_libraries(asio_shared PUBLIC pthread)
-
-install(TARGETS asio_shared DESTINATION lib)

mooncake-common/CMakeLists.txt

@@ -2,6 +2,10 @@ if ((USE_ETCD AND NOT USE_ETCD_LEGACY) OR STORE_USE_ETCD)
     add_subdirectory(etcd)
 endif()
 
+if (STORE_USE_K8S_LEASE)
+    add_subdirectory(k8s-lease)
+endif()
+
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include)
 add_subdirectory(src)
 

mooncake-common/etcd/CMakeLists.txt

@@ -3,7 +3,7 @@ add_custom_command(
     COMMAND bash -c "go mod tidy" && bash -c "go build -buildmode=c-shared -o ${CMAKE_CURRENT_BINARY_DIR}/libetcd_wrapper.so etcd_wrapper.go" && cp ${CMAKE_CURRENT_BINARY_DIR}/libetcd_wrapper.h ${CMAKE_CURRENT_SOURCE_DIR}
     WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
     COMMENT "Building Go shared library"
-    DEPENDS etcd_wrapper.go
+    DEPENDS etcd_wrapper.go go.mod go.sum build.sh
 )
 
 set(ETCD_WRAPPER_INCLUDE ${CMAKE_CURRENT_BINARY_DIR}/libetcd_wrapper.h)
@@ -17,4 +17,4 @@ add_custom_target(
 install(
     FILES ${ETCD_WRAPPER_LIB}
     DESTINATION lib
-)
+)

mooncake-common/etcd/go.mod

@@ -1,6 +1,8 @@
 module github.com/kvcache-ai/Mooncake/mooncake-common/etcd
 
-go 1.24.0
+go 1.25.0
+
+toolchain go1.25.9
 
 require (
 	go.etcd.io/etcd/api/v3 v3.5.21

mooncake-common/etcd/go.sum

@@ -0,0 +1,108 @@
+github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
+github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
+github.com/coreos/go-semver v0.3.0 h1:wkHLiw0WNATZnSG7epLsujiMCgPAc9xhjJ4tgnAxmfM=
+github.com/coreos/go-semver v0.3.0/go.mod h1:nnelYz7RCh+5ahJtPPxZlU+153eP4D4r3EedlOD2RNk=
+github.com/coreos/go-systemd/v22 v22.3.2 h1:D9/bQk5vlXQFZ6Kwuu6zaiXJ9oTPe68++AzAJc1DzSI=
+github.com/coreos/go-systemd/v22 v22.3.2/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI=
+github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
+github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
+github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
+github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
+github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
+github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
+github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
+github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
+github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
+github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
+github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
+github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=
+github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
+github.com/pkg/errors v0.8.1 h1:iURUrRGxPUNPdy5/HRSm+Yj6okJ6UtLINN0Q9M4+h3I=
+github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
+github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
+github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
+go.etcd.io/etcd/api/v3 v3.5.21 h1:A6O2/JDb3tvHhiIz3xf9nJ7REHvtEFJJ3veW3FbCnS8=
+go.etcd.io/etcd/api/v3 v3.5.21/go.mod h1:c3aH5wcvXv/9dqIw2Y810LDXJfhSYdHQ0vxmP3CCHVY=
+go.etcd.io/etcd/client/pkg/v3 v3.5.21 h1:lPBu71Y7osQmzlflM9OfeIV2JlmpBjqBNlLtcoBqUTc=
+go.etcd.io/etcd/client/pkg/v3 v3.5.21/go.mod h1:BgqT/IXPjK9NkeSDjbzwsHySX3yIle2+ndz28nVsjUs=
+go.etcd.io/etcd/client/v3 v3.5.21 h1:T6b1Ow6fNjOLOtM0xSoKNQt1ASPCLWrF9XMHcH9pEyY=
+go.etcd.io/etcd/client/v3 v3.5.21/go.mod h1:mFYy67IOqmbRf/kRUvsHixzo3iG+1OF2W2+jVIQRAnU=
+go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64=
+go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y=
+go.opentelemetry.io/otel v1.39.0 h1:8yPrr/S0ND9QEfTfdP9V+SiwT4E0G7Y5MO7p85nis48=
+go.opentelemetry.io/otel v1.39.0/go.mod h1:kLlFTywNWrFyEdH0oj2xK0bFYZtHRYUdv1NklR/tgc8=
+go.opentelemetry.io/otel/metric v1.39.0 h1:d1UzonvEZriVfpNKEVmHXbdf909uGTOQjA0HF0Ls5Q0=
+go.opentelemetry.io/otel/metric v1.39.0/go.mod h1:jrZSWL33sD7bBxg1xjrqyDjnuzTUB0x1nBERXd7Ftcs=
+go.opentelemetry.io/otel/sdk v1.39.0 h1:nMLYcjVsvdui1B/4FRkwjzoRVsMK8uL/cj0OyhKzt18=
+go.opentelemetry.io/otel/sdk v1.39.0/go.mod h1:vDojkC4/jsTJsE+kh+LXYQlbL8CgrEcwmt1ENZszdJE=
+go.opentelemetry.io/otel/sdk/metric v1.39.0 h1:cXMVVFVgsIf2YL6QkRF4Urbr/aMInf+2WKg+sEJTtB8=
+go.opentelemetry.io/otel/sdk/metric v1.39.0/go.mod h1:xq9HEVH7qeX69/JnwEfp6fVq5wosJsY1mt4lLfYdVew=
+go.opentelemetry.io/otel/trace v1.39.0 h1:2d2vfpEDmCJ5zVYz7ijaJdOF59xLomrvj7bjt6/qCJI=
+go.opentelemetry.io/otel/trace v1.39.0/go.mod h1:88w4/PnZSazkGzz/w84VHpQafiU4EtqqlVdxWy+rNOA=
+go.uber.org/atomic v1.7.0 h1:ADUqmZGgLDDfbSL9ZmPxKTybcoEYHgpYfELNoN+7hsw=
+go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
+go.uber.org/multierr v1.6.0 h1:y6IPFStTAIT5Ytl7/XYmHvzXQ7S3g/IeZW9hyZ5thw4=
+go.uber.org/multierr v1.6.0/go.mod h1:cdWPpRnG4AhwMwsgIHip0KRBQjJy5kYEpYjJxpXp9iU=
+go.uber.org/zap v1.17.0 h1:MTjgFu6ZLKvY6Pvaqk97GlxNBuMpV4Hy/3P6tRGlI2U=
+go.uber.org/zap v1.17.0/go.mod h1:MXVU+bhUf/A7Xi2HNOnopQOrmycQ5Ih87HtOu4q5SSo=
+golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
+golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
+golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
+golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
+golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
+golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
+golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
+golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
+golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
+golang.org/x/net v0.48.0 h1:zyQRTTrjc33Lhh0fBgT/H3oZq9WuvRR5gPC70xpDiQU=
+golang.org/x/net v0.48.0/go.mod h1:+ndRgGjkh8FGtu1w1FGbEC31if4VrNVMuKTgcAAnQRY=
+golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
+golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.39.0 h1:CvCKL8MeisomCi6qNZ+wbb0DN9E5AATixKsvNtMoMFk=
+golang.org/x/sys v0.39.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
+golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
+golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
+golang.org/x/text v0.32.0 h1:ZD01bjUt1FQ9WJ0ClOL5vxgxOI/sVCNgX1YtKwcY0mU=
+golang.org/x/text v0.32.0/go.mod h1:o/rUWzghvpD5TXrTIBuJU77MTaN0ljMWE47kxGJQ7jY=
+golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
+golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
+golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
+golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
+golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk=
+gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E=
+google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217 h1:fCvbg86sFXwdrl5LgVcTEvNC+2txB5mgROGmRL5mrls=
+google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217/go.mod h1:+rXWjjaukWZun3mLfjmVnQi18E1AsFbDN9QdJ5YXLto=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217 h1:gRkg/vSppuSQoDjxyiGfN4Upv/h/DQmIR10ZU8dh4Ww=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217/go.mod h1:7i2o+ce6H/6BluujYR+kqX3GKH+dChPTQU19wjRPiGk=
+google.golang.org/grpc v1.79.3 h1:sybAEdRIEtvcD68Gx7dmnwjZKlyfuc61Dyo9pGXXkKE=
+google.golang.org/grpc v1.79.3/go.mod h1:KmT0Kjez+0dde/v2j9vzwoAScgEPx/Bw1CYChhHLrHQ=
+google.golang.org/protobuf v1.36.10 h1:AYd7cD/uASjIL6Q9LiTjz8JLcrh/88q5UObnmY3aOOE=
+google.golang.org/protobuf v1.36.10/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
+gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

mooncake-common/k8s-lease/CMakeLists.txt

@@ -0,0 +1,20 @@
+add_custom_command(
+    OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/libk8s_lease_wrapper.so
+    COMMAND bash -c "go mod tidy" && bash -c "go build -buildmode=c-shared -o ${CMAKE_CURRENT_BINARY_DIR}/libk8s_lease_wrapper.so k8s_lease_wrapper.go" && cp ${CMAKE_CURRENT_BINARY_DIR}/libk8s_lease_wrapper.h ${CMAKE_CURRENT_SOURCE_DIR}
+    WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+    COMMENT "Building K8s Lease Go shared library"
+    DEPENDS k8s_lease_wrapper.go
+)
+
+set(K8S_LEASE_WRAPPER_INCLUDE ${CMAKE_CURRENT_BINARY_DIR}/libk8s_lease_wrapper.h)
+set(K8S_LEASE_WRAPPER_LIB ${CMAKE_CURRENT_BINARY_DIR}/libk8s_lease_wrapper.so)
+
+add_custom_target(
+    build_k8s_lease_wrapper
+    DEPENDS ${K8S_LEASE_WRAPPER_LIB}
+)
+
+install(
+    FILES ${K8S_LEASE_WRAPPER_LIB}
+    DESTINATION lib
+)

mooncake-common/k8s-lease/cmd/envtest-server/main.go

@@ -0,0 +1,61 @@
+// envtest-server starts a real kube-apiserver + etcd via envtest, writes the
+// KUBECONFIG path to stdout, and blocks until SIGTERM or SIGINT.  This lets
+// C++ tests launch it as a subprocess and talk to a real K8s API without a
+// full cluster.
+package main
+
+import (
+	"fmt"
+	"os"
+	"os/signal"
+	"path/filepath"
+	"syscall"
+
+	"k8s.io/client-go/tools/clientcmd"
+	clientcmdapi "k8s.io/client-go/tools/clientcmd/api"
+	"sigs.k8s.io/controller-runtime/pkg/envtest"
+)
+
+func main() {
+	env := &envtest.Environment{}
+
+	cfg, err := env.Start()
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "envtest start failed: %v\n", err)
+		os.Exit(1)
+	}
+
+	// Write a KUBECONFIG file that points at the envtest kube-apiserver.
+	kubeconfigPath := filepath.Join(os.TempDir(), fmt.Sprintf("envtest-kubeconfig-%d", os.Getpid()))
+	kubeconfig := clientcmdapi.NewConfig()
+	kubeconfig.Clusters["envtest"] = &clientcmdapi.Cluster{
+		Server:                   cfg.Host,
+		CertificateAuthorityData: cfg.CAData,
+	}
+	kubeconfig.AuthInfos["envtest"] = &clientcmdapi.AuthInfo{
+		ClientCertificateData: cfg.CertData,
+		ClientKeyData:         cfg.KeyData,
+	}
+	kubeconfig.Contexts["envtest"] = &clientcmdapi.Context{
+		Cluster:  "envtest",
+		AuthInfo: "envtest",
+	}
+	kubeconfig.CurrentContext = "envtest"
+
+	if err := clientcmd.WriteToFile(*kubeconfig, kubeconfigPath); err != nil {
+		fmt.Fprintf(os.Stderr, "failed to write kubeconfig: %v\n", err)
+		env.Stop()
+		os.Exit(1)
+	}
+
+	// Print the kubeconfig path — the parent process reads this from stdout.
+	fmt.Println(kubeconfigPath)
+
+	// Block until SIGTERM or SIGINT.
+	sigCh := make(chan os.Signal, 1)
+	signal.Notify(sigCh, syscall.SIGTERM, syscall.SIGINT)
+	<-sigCh
+
+	os.Remove(kubeconfigPath)
+	env.Stop()
+}

mooncake-common/k8s-lease/go.mod

@@ -0,0 +1,60 @@
+module github.com/kvcache-ai/Mooncake/mooncake-common/k8s-lease
+
+go 1.24.0
+
+require (
+	k8s.io/api v0.34.3
+	k8s.io/apimachinery v0.34.3
+	k8s.io/client-go v0.34.3
+	k8s.io/utils v0.0.0-20251002143259-bc988d571ff4
+	sigs.k8s.io/controller-runtime v0.22.5
+)
+
+require (
+	github.com/beorn7/perks v1.0.1 // indirect
+	github.com/cespare/xxhash/v2 v2.3.0 // indirect
+	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
+	github.com/emicklei/go-restful/v3 v3.12.2 // indirect
+	github.com/evanphx/json-patch/v5 v5.9.11 // indirect
+	github.com/fxamacker/cbor/v2 v2.9.0 // indirect
+	github.com/go-logr/logr v1.4.3 // indirect
+	github.com/go-openapi/jsonpointer v0.21.0 // indirect
+	github.com/go-openapi/jsonreference v0.20.2 // indirect
+	github.com/go-openapi/swag v0.23.0 // indirect
+	github.com/gogo/protobuf v1.3.2 // indirect
+	github.com/google/gnostic-models v0.7.0 // indirect
+	github.com/google/go-cmp v0.7.0 // indirect
+	github.com/google/uuid v1.6.0 // indirect
+	github.com/josharian/intern v1.0.0 // indirect
+	github.com/json-iterator/go v1.1.12 // indirect
+	github.com/mailru/easyjson v0.7.7 // indirect
+	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
+	github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee // indirect
+	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/prometheus/client_golang v1.23.2 // indirect
+	github.com/prometheus/client_model v0.6.2 // indirect
+	github.com/prometheus/common v0.66.1 // indirect
+	github.com/prometheus/procfs v0.16.1 // indirect
+	github.com/spf13/pflag v1.0.9 // indirect
+	github.com/x448/float16 v0.8.4 // indirect
+	go.yaml.in/yaml/v2 v2.4.3 // indirect
+	go.yaml.in/yaml/v3 v3.0.4 // indirect
+	golang.org/x/net v0.47.0 // indirect
+	golang.org/x/oauth2 v0.30.0 // indirect
+	golang.org/x/sys v0.38.0 // indirect
+	golang.org/x/term v0.37.0 // indirect
+	golang.org/x/text v0.31.0 // indirect
+	golang.org/x/time v0.9.0 // indirect
+	google.golang.org/protobuf v1.36.8 // indirect
+	gopkg.in/evanphx/json-patch.v4 v4.13.0 // indirect
+	gopkg.in/inf.v0 v0.9.1 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+	k8s.io/apiextensions-apiserver v0.34.3 // indirect
+	k8s.io/klog/v2 v2.130.1 // indirect
+	k8s.io/kube-openapi v0.0.0-20250910181357-589584f1c912 // indirect
+	sigs.k8s.io/json v0.0.0-20250730193827-2d320260d730 // indirect
+	sigs.k8s.io/randfill v1.0.0 // indirect
+	sigs.k8s.io/structured-merge-diff/v6 v6.3.2-0.20260122202528-d9cc6641c482 // indirect
+	sigs.k8s.io/yaml v1.6.0 // indirect
+)

mooncake-common/k8s-lease/k8s_lease_integration_test.go

@@ -0,0 +1,571 @@
+//go:build integration
+
+package main
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"sync"
+	"testing"
+	"time"
+
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/client-go/kubernetes"
+	"k8s.io/client-go/rest"
+	"k8s.io/client-go/tools/leaderelection"
+	"k8s.io/client-go/tools/leaderelection/resourcelock"
+
+	"sigs.k8s.io/controller-runtime/pkg/envtest"
+)
+
+var (
+	testEnv    *envtest.Environment
+	testConfig *rest.Config
+)
+
+type electionStateNoRelease struct {
+	cancel  context.CancelFunc
+	elected chan struct{}
+	lost    chan struct{}
+}
+
+func TestMain(m *testing.M) {
+	testEnv = &envtest.Environment{}
+
+	var err error
+	testConfig, err = testEnv.Start()
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "failed to start envtest: %v\n", err)
+		os.Exit(1)
+	}
+
+	// Set up global client for the wrapper
+	client, err := kubernetes.NewForConfig(testConfig)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "failed to create clientset: %v\n", err)
+		testEnv.Stop()
+		os.Exit(1)
+	}
+	clientMutex.Lock()
+	globalClient = client
+	clientMutex.Unlock()
+
+	code := m.Run()
+
+	testEnv.Stop()
+	os.Exit(code)
+}
+
+func runElectionWithoutRelease(namespace, leaseName, identity string,
+	leaseDurationSec, renewDeadlineSec, retryPeriodSec int) (*electionStateNoRelease, error) {
+	if err := ensureClientInitialized(); err != nil {
+		return nil, err
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	state := &electionStateNoRelease{
+		cancel:  cancel,
+		elected: make(chan struct{}),
+		lost:    make(chan struct{}),
+	}
+
+	lock := &resourcelock.LeaseLock{
+		LeaseMeta: metav1.ObjectMeta{
+			Name:      leaseName,
+			Namespace: namespace,
+		},
+		Client: globalClient.CoordinationV1(),
+		LockConfig: resourcelock.ResourceLockConfig{
+			Identity: identity,
+		},
+	}
+
+	le, err := leaderelection.NewLeaderElector(leaderelection.LeaderElectionConfig{
+		Lock:            lock,
+		LeaseDuration:   time.Duration(leaseDurationSec) * time.Second,
+		RenewDeadline:   time.Duration(renewDeadlineSec) * time.Second,
+		RetryPeriod:     time.Duration(retryPeriodSec) * time.Second,
+		ReleaseOnCancel: false,
+		Callbacks: leaderelection.LeaderCallbacks{
+			OnStartedLeading: func(ctx context.Context) {
+				close(state.elected)
+				<-ctx.Done()
+			},
+			OnStoppedLeading: func() {
+				close(state.lost)
+			},
+		},
+	})
+	if err != nil {
+		cancel()
+		return nil, fmt.Errorf("failed to create leader elector: %w", err)
+	}
+
+	go le.Run(ctx)
+	return state, nil
+}
+
+// TestSingleLeaderElection verifies a single candidate becomes leader.
+func TestSingleLeaderElection(t *testing.T) {
+	ns := "default"
+	lease := "single-election-test"
+	identity := "node-1:8080"
+
+	err := runElection(ns, lease, identity, 5, 4, 1)
+	if err != nil {
+		t.Fatalf("runElection failed: %v", err)
+	}
+
+	// Wait for elected
+	key := electionKey(ns, lease)
+	electionMutex.Lock()
+	state := elections[key]
+	electionMutex.Unlock()
+
+	select {
+	case <-state.elected:
+		// success
+	case <-time.After(15 * time.Second):
+		t.Fatal("timed out waiting for election")
+	}
+
+	// Verify holder via getHolder
+	holder, transitions, err := getHolder(ns, lease)
+	if err != nil {
+		t.Fatalf("getHolder failed: %v", err)
+	}
+	if holder != identity {
+		t.Errorf("expected holder %q, got %q", identity, holder)
+	}
+	// First election — transitions should be 0 or 1
+	if transitions < 0 {
+		t.Errorf("expected non-negative transitions, got %d", transitions)
+	}
+
+	// Cancel the election
+	electionMutex.Lock()
+	state = elections[key]
+	electionMutex.Unlock()
+	state.cancel()
+
+	select {
+	case <-state.lost:
+		// success
+	case <-time.After(10 * time.Second):
+		t.Fatal("timed out waiting for election loss after cancel")
+	}
+}
+
+// TestLeaderEpoch verifies leaseTransitions increments across elections.
+func TestLeaderEpoch(t *testing.T) {
+	ns := "default"
+	lease := "epoch-test"
+
+	// First election
+	err := runElection(ns, lease, "node-epoch-1:8080", 5, 4, 1)
+	if err != nil {
+		t.Fatalf("first runElection failed: %v", err)
+	}
+
+	key := electionKey(ns, lease)
+	electionMutex.Lock()
+	state1 := elections[key]
+	electionMutex.Unlock()
+
+	select {
+	case <-state1.elected:
+	case <-time.After(15 * time.Second):
+		t.Fatal("timed out on first election")
+	}
+
+	_, trans1, _ := getHolder(ns, lease)
+
+	// Cancel first election and wait for loss
+	state1.cancel()
+	select {
+	case <-state1.lost:
+	case <-time.After(10 * time.Second):
+		t.Fatal("timed out waiting for first election loss")
+	}
+
+	// Wait for lease to expire / be released
+	time.Sleep(2 * time.Second)
+
+	// Second election
+	err = runElection(ns, lease, "node-epoch-2:8080", 5, 4, 1)
+	if err != nil {
+		t.Fatalf("second runElection failed: %v", err)
+	}
+
+	electionMutex.Lock()
+	state2 := elections[key]
+	electionMutex.Unlock()
+
+	select {
+	case <-state2.elected:
+	case <-time.After(15 * time.Second):
+		t.Fatal("timed out on second election")
+	}
+
+	_, trans2, _ := getHolder(ns, lease)
+	if trans2 <= trans1 {
+		t.Errorf("expected transitions to increment: first=%d, second=%d", trans1, trans2)
+	}
+
+	state2.cancel()
+	select {
+	case <-state2.lost:
+	case <-time.After(10 * time.Second):
+		t.Fatal("timed out waiting for second election loss")
+	}
+}
+
+// TestSequentialLeadershipHandoff tests that a second candidate can acquire
+// leadership after the first one releases it.
+func TestSequentialLeadershipHandoff(t *testing.T) {
+	ns := "default"
+	lease := "two-candidate-test"
+
+	err1 := runElection(ns, lease, "candidate-a:8080", 5, 4, 1)
+	if err1 != nil {
+		t.Fatalf("first runElection failed: %v", err1)
+	}
+
+	key := electionKey(ns, lease)
+	electionMutex.Lock()
+	stateA := elections[key]
+	electionMutex.Unlock()
+
+	// Wait for first candidate to win
+	select {
+	case <-stateA.elected:
+	case <-time.After(15 * time.Second):
+		t.Fatal("timed out waiting for first candidate")
+	}
+
+	// Verify holder is candidate-a
+	holder, _, err := getHolder(ns, lease)
+	if err != nil {
+		t.Fatalf("getHolder failed: %v", err)
+	}
+	if holder != "candidate-a:8080" {
+		t.Errorf("expected candidate-a, got %q", holder)
+	}
+
+	// Cancel candidate-a
+	stateA.cancel()
+	select {
+	case <-stateA.lost:
+	case <-time.After(10 * time.Second):
+		t.Fatal("timed out waiting for candidate-a loss")
+	}
+
+	// Wait for lease to expire
+	time.Sleep(2 * time.Second)
+
+	// Start candidate-b
+	err2 := runElection(ns, lease, "candidate-b:8080", 5, 4, 1)
+	if err2 != nil {
+		t.Fatalf("second runElection failed: %v", err2)
+	}
+
+	electionMutex.Lock()
+	stateB := elections[key]
+	electionMutex.Unlock()
+
+	select {
+	case <-stateB.elected:
+	case <-time.After(15 * time.Second):
+		t.Fatal("timed out waiting for candidate-b")
+	}
+
+	holder, _, err = getHolder(ns, lease)
+	if err != nil {
+		t.Fatalf("getHolder after takeover failed: %v", err)
+	}
+	if holder != "candidate-b:8080" {
+		t.Errorf("expected candidate-b, got %q", holder)
+	}
+
+	stateB.cancel()
+	select {
+	case <-stateB.lost:
+	case <-time.After(10 * time.Second):
+		t.Fatal("timed out waiting for candidate-b loss")
+	}
+}
+
+// TestConcurrentCandidateElection starts two candidates simultaneously and
+// verifies that exactly one wins leadership.
+func TestConcurrentCandidateElection(t *testing.T) {
+	ns := "default"
+	lease := "concurrent-election-test"
+
+	type result struct {
+		identity string
+		elected  bool
+	}
+
+	candidates := []string{"candidate-a:8080", "candidate-b:8080"}
+	results := make(chan result, len(candidates))
+
+	lock := func(identity string) *resourcelock.LeaseLock {
+		return &resourcelock.LeaseLock{
+			LeaseMeta: metav1.ObjectMeta{
+				Name:      lease,
+				Namespace: ns,
+			},
+			Client: globalClient.CoordinationV1(),
+			LockConfig: resourcelock.ResourceLockConfig{
+				Identity: identity,
+			},
+		}
+	}
+
+	var wg sync.WaitGroup
+	for _, id := range candidates {
+		wg.Add(1)
+		go func(identity string) {
+			defer wg.Done()
+
+			// Short timeout: enough for one to acquire, but the loser
+			// times out before the winner's lease could expire.
+			ctx, cancel := context.WithTimeout(context.Background(), 8*time.Second)
+			defer cancel()
+
+			elected := make(chan struct{})
+			le, err := leaderelection.NewLeaderElector(leaderelection.LeaderElectionConfig{
+				Lock:            lock(identity),
+				LeaseDuration:   5 * time.Second,
+				RenewDeadline:   3 * time.Second,
+				RetryPeriod:     1 * time.Second,
+				ReleaseOnCancel: true,
+				Callbacks: leaderelection.LeaderCallbacks{
+					OnStartedLeading: func(ctx context.Context) {
+						close(elected)
+						<-ctx.Done()
+					},
+					OnStoppedLeading: func() {},
+				},
+			})
+			if err != nil {
+				t.Errorf("NewLeaderElector(%s): %v", identity, err)
+				return
+			}
+
+			go le.Run(ctx)
+
+			select {
+			case <-elected:
+				results <- result{identity, true}
+				// Keep holding until context expires (8s total).
+				// Winner does NOT release early, so loser cannot
+				// re-acquire within its own 8s window.
+				<-ctx.Done()
+			case <-ctx.Done():
+				results <- result{identity, false}
+			}
+		}(id)
+	}
+
+	wg.Wait()
+	close(results)
+
+	winners := 0
+	for r := range results {
+		if r.elected {
+			winners++
+			t.Logf("winner: %s", r.identity)
+		}
+	}
+
+	if winners != 1 {
+		t.Fatalf("expected exactly 1 winner, got %d", winners)
+	}
+}
+
+// TestCancelElection tests that cancelling an election makes WaitLost return.
+func TestCancelElection(t *testing.T) {
+	ns := "default"
+	lease := "cancel-test"
+
+	err := runElection(ns, lease, "cancel-node:8080", 5, 4, 1)
+	if err != nil {
+		t.Fatalf("runElection failed: %v", err)
+	}
+
+	key := electionKey(ns, lease)
+	electionMutex.Lock()
+	state := elections[key]
+	electionMutex.Unlock()
+
+	// Wait for elected
+	select {
+	case <-state.elected:
+	case <-time.After(15 * time.Second):
+		t.Fatal("timed out waiting for election")
+	}
+
+	// Cancel
+	state.cancel()
+
+	// WaitLost should return promptly
+	select {
+	case <-state.lost:
+		// success
+	case <-time.After(10 * time.Second):
+		t.Fatal("WaitLost did not return after cancel")
+	}
+}
+
+// TestGetHolderDuringElection verifies getHolder works while election is active.
+func TestGetHolderDuringElection(t *testing.T) {
+	ns := "default"
+	lease := "active-get-holder-test"
+	identity := "active-node:8080"
+
+	err := runElection(ns, lease, identity, 5, 4, 1)
+	if err != nil {
+		t.Fatalf("runElection failed: %v", err)
+	}
+
+	key := electionKey(ns, lease)
+	electionMutex.Lock()
+	state := elections[key]
+	electionMutex.Unlock()
+
+	select {
+	case <-state.elected:
+	case <-time.After(15 * time.Second):
+		t.Fatal("timed out waiting for election")
+	}
+
+	// Concurrent getHolder calls during active election
+	var wg sync.WaitGroup
+	for i := 0; i < 5; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			holder, _, err := getHolder(ns, lease)
+			if err != nil {
+				t.Errorf("getHolder during election failed: %v", err)
+				return
+			}
+			if holder != identity {
+				t.Errorf("expected %q, got %q", identity, holder)
+			}
+		}()
+	}
+	wg.Wait()
+
+	state.cancel()
+	<-state.lost
+}
+
+// TestGetHolderReturnsEmptyAfterLeaderDeath verifies that after a leader stops
+// renewing its lease without releasing it, getHolder returns an empty holder
+// once the lease expires. This is the integration-level counterpart to the
+// unit test TestGetHolderReturnsEmptyForExpiredLease.
+func TestGetHolderReturnsEmptyAfterLeaderDeath(t *testing.T) {
+	ns := "default"
+	lease := "expired-leader-test"
+	identity := "doomed-leader:8080"
+
+	// Acquire leadership without ReleaseOnCancel so canceling simulates a dead
+	// leader that stops renewing and leaves the old holder until expiry.
+	state, err := runElectionWithoutRelease(ns, lease, identity, 5, 4, 1)
+	if err != nil {
+		t.Fatalf("runElection failed: %v", err)
+	}
+
+	select {
+	case <-state.elected:
+	case <-time.After(15 * time.Second):
+		t.Fatal("timed out waiting for election")
+	}
+
+	// Verify holder while active.
+	holder, _, err := getHolder(ns, lease)
+	if err != nil {
+		t.Fatalf("getHolder (active) failed: %v", err)
+	}
+	if holder != identity {
+		t.Fatalf("expected active holder %q, got %q", identity, holder)
+	}
+
+	// Simulate leader death: stop renewing without explicitly releasing.
+	state.cancel()
+	select {
+	case <-state.lost:
+	case <-time.After(10 * time.Second):
+		t.Fatal("timed out waiting for loss")
+	}
+
+	// Wait for the lease to expire (leaseDuration=5s, add margin).
+	time.Sleep(7 * time.Second)
+
+	// After expiry, getHolder must return empty holder so that the
+	// supervisor will attempt acquisition.
+	holder, _, err = getHolder(ns, lease)
+	if err != nil {
+		t.Fatalf("getHolder (expired) failed: %v", err)
+	}
+	if holder != "" {
+		t.Errorf("expected empty holder after lease expiry, got %q", holder)
+	}
+}
+
+// TestFailoverAfterLeaderDeath verifies that a new candidate can acquire
+// leadership after the previous leader dies and its lease expires.
+func TestFailoverAfterLeaderDeath(t *testing.T) {
+	ns := "default"
+	lease := "failover-test"
+
+	// First leader acquires without ReleaseOnCancel so canceling leaves the
+	// old holder in place until the lease naturally expires.
+	state1, err := runElectionWithoutRelease(ns, lease, "leader-1:8080", 5, 4, 1)
+	if err != nil {
+		t.Fatalf("first runElection failed: %v", err)
+	}
+
+	select {
+	case <-state1.elected:
+	case <-time.After(15 * time.Second):
+		t.Fatal("timed out waiting for first election")
+	}
+
+	// Simulate crash: cancel without release, wait for expiry.
+	state1.cancel()
+	<-state1.lost
+	time.Sleep(7 * time.Second)
+
+	// Second candidate should be able to acquire.
+	err = runElection(ns, lease, "leader-2:8080", 5, 4, 1)
+	if err != nil {
+		t.Fatalf("second runElection failed: %v", err)
+	}
+
+	key := electionKey(ns, lease)
+	electionMutex.Lock()
+	state2 := elections[key]
+	electionMutex.Unlock()
+
+	select {
+	case <-state2.elected:
+		// success — failover worked
+	case <-time.After(15 * time.Second):
+		t.Fatal("second candidate failed to acquire after leader death")
+	}
+
+	holder, _, err := getHolder(ns, lease)
+	if err != nil {
+		t.Fatalf("getHolder after failover failed: %v", err)
+	}
+	if holder != "leader-2:8080" {
+		t.Errorf("expected new leader %q, got %q", "leader-2:8080", holder)
+	}
+
+	state2.cancel()
+	<-state2.lost
+}

mooncake-common/k8s-lease/k8s_lease_wrapper.go

@@ -0,0 +1,489 @@
+package main
+
+/*
+#include <stdint.h>
+#include <stdlib.h>
+#include <string.h>
+
+// Trampoline to invoke C/C++ callback safely from Go via cgo.
+typedef void (*holder_change_cb_t)(void* ctx,
+                                    const char* holder, size_t holderSize,
+                                    int64_t leaseTransitions);
+
+static inline void call_holder_change_cb(holder_change_cb_t func, void* ctx,
+                                          const char* holder, size_t holderSize,
+                                          int64_t leaseTransitions) {
+  func(ctx, holder, holderSize, leaseTransitions);
+}
+*/
+import "C"
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"sync"
+	"time"
+	"unsafe"
+
+	coordinationv1 "k8s.io/api/coordination/v1"
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/watch"
+	"k8s.io/client-go/kubernetes"
+	"k8s.io/client-go/rest"
+	"k8s.io/client-go/tools/clientcmd"
+	"k8s.io/client-go/tools/leaderelection"
+	"k8s.io/client-go/tools/leaderelection/resourcelock"
+)
+
+// electionState holds the runtime state for a single leader election.
+type electionState struct {
+	cancel      context.CancelFunc
+	elected     chan struct{} // closed when OnStartedLeading fires
+	lost        chan struct{} // closed when OnStoppedLeading fires
+	err         error         // set before lost is closed, if any
+	transitions int64         // set before elected is closed
+}
+
+// watchState holds the runtime state for a single Lease watch.
+type watchState struct {
+	cancel context.CancelFunc
+}
+
+var (
+	globalClient kubernetes.Interface
+	clientMutex  sync.Mutex
+	initClientFn = initClient
+
+	elections     = make(map[string]*electionState)
+	electionMutex sync.Mutex
+
+	watches    = make(map[string]*watchState)
+	watchMutex sync.Mutex
+)
+
+func electionKey(namespace, leaseName string) string {
+	return namespace + "/" + leaseName
+}
+
+func ensureClientInitialized() error {
+	clientMutex.Lock()
+	initialized := globalClient != nil
+	clientMutex.Unlock()
+	if initialized {
+		return nil
+	}
+	return initClientFn()
+}
+
+// initClient creates the K8s clientset from in-cluster config or KUBECONFIG.
+func initClient() error {
+	clientMutex.Lock()
+	defer clientMutex.Unlock()
+	if globalClient != nil {
+		return nil
+	}
+
+	config, err := rest.InClusterConfig()
+	if err != nil {
+		// Fall back to KUBECONFIG
+		kubeconfig := os.Getenv("KUBECONFIG")
+		if kubeconfig == "" {
+			home := os.Getenv("HOME")
+			if home != "" {
+				kubeconfig = home + "/.kube/config"
+			}
+		}
+		config, err = clientcmd.BuildConfigFromFlags("", kubeconfig)
+		if err != nil {
+			return fmt.Errorf("failed to build k8s config: %w", err)
+		}
+	}
+
+	client, err := kubernetes.NewForConfig(config)
+	if err != nil {
+		return fmt.Errorf("failed to create k8s clientset: %w", err)
+	}
+	globalClient = client
+	return nil
+}
+
+// runElection starts a leader election goroutine for the given namespace/leaseName.
+func runElection(namespace, leaseName, identity string,
+	leaseDurationSec, renewDeadlineSec, retryPeriodSec int) error {
+	if err := ensureClientInitialized(); err != nil {
+		return err
+	}
+
+	key := electionKey(namespace, leaseName)
+
+	electionMutex.Lock()
+	if _, exists := elections[key]; exists {
+		electionMutex.Unlock()
+		return fmt.Errorf("election already running for %s", key)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	state := &electionState{
+		cancel:  cancel,
+		elected: make(chan struct{}),
+		lost:    make(chan struct{}),
+	}
+	elections[key] = state
+	electionMutex.Unlock()
+
+	lock := &resourcelock.LeaseLock{
+		LeaseMeta: metav1.ObjectMeta{
+			Name:      leaseName,
+			Namespace: namespace,
+		},
+		Client: globalClient.CoordinationV1(),
+		LockConfig: resourcelock.ResourceLockConfig{
+			Identity: identity,
+		},
+	}
+
+	le, err := leaderelection.NewLeaderElector(leaderelection.LeaderElectionConfig{
+		Lock:            lock,
+		LeaseDuration:   time.Duration(leaseDurationSec) * time.Second,
+		RenewDeadline:   time.Duration(renewDeadlineSec) * time.Second,
+		RetryPeriod:     time.Duration(retryPeriodSec) * time.Second,
+		ReleaseOnCancel: true,
+		Callbacks: leaderelection.LeaderCallbacks{
+			OnStartedLeading: func(ctx context.Context) {
+				_, transitions, err := getHolder(namespace, leaseName)
+				if err == nil {
+					state.transitions = transitions
+				}
+				close(state.elected)
+				// Block until context is cancelled (leadership lost or explicit cancel)
+				<-ctx.Done()
+			},
+			OnStoppedLeading: func() {
+				close(state.lost)
+				// Auto-cleanup: remove from map so the same key can be reused.
+				electionMutex.Lock()
+				if elections[key] == state {
+					delete(elections, key)
+				}
+				electionMutex.Unlock()
+			},
+		},
+	})
+	if err != nil {
+		electionMutex.Lock()
+		delete(elections, key)
+		electionMutex.Unlock()
+		cancel()
+		return fmt.Errorf("failed to create leader elector: %w", err)
+	}
+
+	go le.Run(ctx)
+	return nil
+}
+
+// getHolder reads the current Lease holder identity and transitions.
+func getHolder(namespace, leaseName string) (string, int64, error) {
+	if err := ensureClientInitialized(); err != nil {
+		return "", 0, err
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	lease, err := globalClient.CoordinationV1().Leases(namespace).Get(ctx, leaseName, metav1.GetOptions{})
+	if err != nil {
+		return "", 0, fmt.Errorf("failed to get lease: %w", err)
+	}
+
+	holder := ""
+	if lease.Spec.HolderIdentity != nil {
+		holder = *lease.Spec.HolderIdentity
+	}
+	transitions := int64(0)
+	if lease.Spec.LeaseTransitions != nil {
+		transitions = int64(*lease.Spec.LeaseTransitions)
+	}
+
+	// Treat expired leases as having no holder so that the C++ supervisor
+	// will attempt acquisition instead of going to standby.
+	if holder != "" && lease.Spec.RenewTime != nil && lease.Spec.LeaseDurationSeconds != nil {
+		expiry := lease.Spec.RenewTime.Time.Add(time.Duration(*lease.Spec.LeaseDurationSeconds) * time.Second)
+		if time.Now().After(expiry) {
+			holder = ""
+		}
+	}
+	return holder, transitions, nil
+}
+
+//export K8sLeaseInit
+func K8sLeaseInit(errMsg **C.char) C.int {
+	if err := ensureClientInitialized(); err != nil {
+		*errMsg = C.CString(err.Error())
+		return -1
+	}
+	return 0
+}
+
+//export K8sLeaseRunElection
+func K8sLeaseRunElection(
+	ns, leaseName, identity *C.char,
+	leaseDurationSec, renewDeadlineSec, retryPeriodSec C.int,
+	errMsg **C.char,
+) C.int {
+	nsStr := C.GoString(ns)
+	ln := C.GoString(leaseName)
+	id := C.GoString(identity)
+
+	err := runElection(nsStr, ln, id,
+		int(leaseDurationSec), int(renewDeadlineSec), int(retryPeriodSec))
+	if err != nil {
+		*errMsg = C.CString(err.Error())
+		return -1
+	}
+	return 0
+}
+
+//export K8sLeaseWaitElected
+func K8sLeaseWaitElected(
+	ns, leaseName *C.char,
+	timeoutSec C.int,
+	leaseTransitions *C.longlong,
+	errMsg **C.char,
+) C.int {
+	key := electionKey(C.GoString(ns), C.GoString(leaseName))
+
+	electionMutex.Lock()
+	state, exists := elections[key]
+	electionMutex.Unlock()
+
+	if !exists {
+		*errMsg = C.CString("no election running for " + key)
+		return -1
+	}
+
+	timeout := time.Duration(timeoutSec) * time.Second
+
+	// Wait for elected, lost, or timeout
+	select {
+	case <-state.elected:
+		*leaseTransitions = C.longlong(state.transitions)
+		return 0
+	case <-state.lost:
+		*errMsg = C.CString("election lost before becoming leader")
+		return -1
+	case <-time.After(timeout):
+		state.cancel()
+		<-state.lost
+		*errMsg = C.CString("election timed out after " + fmt.Sprintf("%d", int(timeoutSec)) + "s")
+		return -1
+	}
+}
+
+//export K8sLeaseWaitLost
+func K8sLeaseWaitLost(
+	ns, leaseName *C.char,
+	errMsg **C.char,
+) C.int {
+	key := electionKey(C.GoString(ns), C.GoString(leaseName))
+
+	electionMutex.Lock()
+	state, exists := elections[key]
+	electionMutex.Unlock()
+
+	if !exists {
+		// Already cleaned up by OnStoppedLeading — election is over.
+		return 0
+	}
+
+	<-state.lost
+
+	if state.err != nil {
+		*errMsg = C.CString(state.err.Error())
+		return -1
+	}
+	return 0
+}
+
+//export K8sLeaseCancelElection
+func K8sLeaseCancelElection(
+	ns, leaseName *C.char,
+	errMsg **C.char,
+) C.int {
+	key := electionKey(C.GoString(ns), C.GoString(leaseName))
+
+	electionMutex.Lock()
+	state, exists := elections[key]
+	electionMutex.Unlock()
+
+	if !exists {
+		// Idempotent — no error if no election
+		return 0
+	}
+
+	state.cancel()
+	return 0
+}
+
+//export K8sLeaseGetHolder
+func K8sLeaseGetHolder(
+	ns, leaseName *C.char,
+	holderIdentity **C.char,
+	leaseTransitions *C.longlong,
+	errMsg **C.char,
+) C.int {
+	nsStr := C.GoString(ns)
+	ln := C.GoString(leaseName)
+
+	holder, transitions, err := getHolder(nsStr, ln)
+	if err != nil {
+		if apierrors.IsNotFound(err) {
+			*holderIdentity = nil
+			*leaseTransitions = 0
+			return 1
+		}
+		errStr := err.Error()
+		*errMsg = C.CString(errStr)
+		return -1
+	}
+
+	if holder == "" {
+		*holderIdentity = nil
+	} else {
+		*holderIdentity = C.CString(holder)
+	}
+	*leaseTransitions = C.longlong(transitions)
+	return 0
+}
+
+//export K8sLeaseWatchHolder
+func K8sLeaseWatchHolder(
+	ns, leaseName *C.char,
+	callbackCtx unsafe.Pointer,
+	callbackFunc C.holder_change_cb_t,
+	errMsg **C.char,
+) C.int {
+	nsStr := C.GoString(ns)
+	ln := C.GoString(leaseName)
+	key := electionKey(nsStr, ln)
+
+	if callbackFunc == nil {
+		*errMsg = C.CString("callback function is nil")
+		return -1
+	}
+	if err := ensureClientInitialized(); err != nil {
+		*errMsg = C.CString(err.Error())
+		return -1
+	}
+
+	watchMutex.Lock()
+	if _, exists := watches[key]; exists {
+		watchMutex.Unlock()
+		*errMsg = C.CString("watch already running for " + key)
+		return -1
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	watches[key] = &watchState{cancel: cancel}
+	watchMutex.Unlock()
+
+	go func() {
+		defer func() {
+			watchMutex.Lock()
+			delete(watches, key)
+			watchMutex.Unlock()
+		}()
+
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			default:
+			}
+
+			watcher, err := globalClient.CoordinationV1().Leases(nsStr).Watch(ctx, metav1.ListOptions{
+				FieldSelector: "metadata.name=" + ln,
+			})
+			if err != nil {
+				select {
+				case <-ctx.Done():
+					return
+				default:
+					time.Sleep(time.Second)
+					continue
+				}
+			}
+
+			for event := range watcher.ResultChan() {
+				select {
+				case <-ctx.Done():
+					watcher.Stop()
+					return
+				default:
+				}
+
+				if event.Type == watch.Modified || event.Type == watch.Added {
+					lease, ok := event.Object.(*coordinationv1.Lease)
+					if !ok {
+						continue
+					}
+					holder := ""
+					if lease.Spec.HolderIdentity != nil {
+						holder = *lease.Spec.HolderIdentity
+					}
+					transitions := int64(0)
+					if lease.Spec.LeaseTransitions != nil {
+						transitions = int64(*lease.Spec.LeaseTransitions)
+					}
+
+					var holderPtr *C.char
+					var holderSize C.size_t
+					if holder != "" {
+						holderPtr = C.CString(holder)
+						holderSize = C.size_t(len(holder))
+					}
+
+					C.call_holder_change_cb(callbackFunc, callbackCtx,
+						holderPtr, holderSize, C.int64_t(transitions))
+
+					if holderPtr != nil {
+						C.free(unsafe.Pointer(holderPtr))
+					}
+				}
+			}
+
+			// Watch channel closed — retry unless cancelled
+			select {
+			case <-ctx.Done():
+				return
+			default:
+				time.Sleep(time.Second)
+			}
+		}
+	}()
+
+	return 0
+}
+
+//export K8sLeaseCancelWatch
+func K8sLeaseCancelWatch(
+	ns, leaseName *C.char,
+	errMsg **C.char,
+) C.int {
+	key := electionKey(C.GoString(ns), C.GoString(leaseName))
+
+	watchMutex.Lock()
+	state, exists := watches[key]
+	watchMutex.Unlock()
+
+	if !exists {
+		// Idempotent
+		return 0
+	}
+
+	state.cancel()
+	return 0
+}
+
+func main() {}

mooncake-common/k8s-lease/k8s_lease_wrapper_test.go

@@ -0,0 +1,293 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	coordinationv1 "k8s.io/api/coordination/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/client-go/kubernetes"
+	"k8s.io/client-go/kubernetes/fake"
+	"k8s.io/utils/ptr"
+)
+
+// swapClient replaces globalClient and returns the old one.
+func swapClient(newClient kubernetes.Interface) kubernetes.Interface {
+	clientMutex.Lock()
+	defer clientMutex.Unlock()
+	old := globalClient
+	globalClient = newClient
+	return old
+}
+
+// TestGetHolderWithFakeClient tests getHolder using a fake K8s clientset.
+func TestGetHolderWithFakeClient(t *testing.T) {
+	holderID := "node-1:8080"
+	transitions := int32(3)
+	lease := &coordinationv1.Lease{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "test-lease",
+			Namespace: "default",
+		},
+		Spec: coordinationv1.LeaseSpec{
+			HolderIdentity:   &holderID,
+			LeaseTransitions: &transitions,
+		},
+	}
+
+	fakeClient := fake.NewSimpleClientset(lease)
+	old := swapClient(fakeClient)
+	defer swapClient(old)
+
+	holder, trans, err := getHolder("default", "test-lease")
+	if err != nil {
+		t.Fatalf("getHolder failed: %v", err)
+	}
+	if holder != holderID {
+		t.Errorf("expected holder %q, got %q", holderID, holder)
+	}
+	if trans != int64(transitions) {
+		t.Errorf("expected transitions %d, got %d", transitions, trans)
+	}
+}
+
+// TestGetHolderNotFound tests getHolder when the Lease does not exist.
+func TestGetHolderNotFound(t *testing.T) {
+	fakeClient := fake.NewSimpleClientset()
+	old := swapClient(fakeClient)
+	defer swapClient(old)
+
+	_, _, err := getHolder("default", "nonexistent")
+	if err == nil {
+		t.Fatal("expected error for nonexistent lease, got nil")
+	}
+}
+
+// TestGetHolderEmptyIdentity tests getHolder when holder is nil.
+func TestGetHolderEmptyIdentity(t *testing.T) {
+	lease := &coordinationv1.Lease{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "empty-lease",
+			Namespace: "default",
+		},
+		Spec: coordinationv1.LeaseSpec{},
+	}
+	fakeClient := fake.NewSimpleClientset(lease)
+	old := swapClient(fakeClient)
+	defer swapClient(old)
+
+	holder, trans, err := getHolder("default", "empty-lease")
+	if err != nil {
+		t.Fatalf("getHolder failed: %v", err)
+	}
+	if holder != "" {
+		t.Errorf("expected empty holder, got %q", holder)
+	}
+	if trans != 0 {
+		t.Errorf("expected 0 transitions, got %d", trans)
+	}
+}
+
+// TestGetHolderReturnsEmptyForExpiredLease verifies that getHolder treats a
+// lease whose renewTime + leaseDuration is in the past as having no holder.
+// This is critical for failover: when a leader pod dies without releasing the
+// lease, standbys must see an empty holder so the supervisor attempts
+// acquisition instead of looping in standby.
+func TestGetHolderReturnsEmptyForExpiredLease(t *testing.T) {
+	holderID := "dead-leader:8080"
+	leaseDuration := int32(5)
+	transitions := int32(2)
+	expiredRenewTime := metav1.NewMicroTime(time.Now().Add(-10 * time.Second))
+
+	lease := &coordinationv1.Lease{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "expired-lease",
+			Namespace: "default",
+		},
+		Spec: coordinationv1.LeaseSpec{
+			HolderIdentity:       &holderID,
+			LeaseDurationSeconds: &leaseDuration,
+			LeaseTransitions:     &transitions,
+			RenewTime:            &expiredRenewTime,
+		},
+	}
+
+	fakeClient := fake.NewSimpleClientset(lease)
+	old := swapClient(fakeClient)
+	defer swapClient(old)
+
+	holder, trans, err := getHolder("default", "expired-lease")
+	if err != nil {
+		t.Fatalf("getHolder failed: %v", err)
+	}
+	if holder != "" {
+		t.Errorf("expected empty holder for expired lease, got %q", holder)
+	}
+	// Transitions should still be reported even for expired leases.
+	if trans != int64(transitions) {
+		t.Errorf("expected transitions %d, got %d", transitions, trans)
+	}
+}
+
+// TestGetHolderReturnsHolderForActiveLease verifies that getHolder returns the
+// holder identity when the lease is still active (renewTime + leaseDuration is
+// in the future).
+func TestGetHolderReturnsHolderForActiveLease(t *testing.T) {
+	holderID := "active-leader:8080"
+	leaseDuration := int32(15)
+	transitions := int32(1)
+	recentRenewTime := metav1.NewMicroTime(time.Now())
+
+	lease := &coordinationv1.Lease{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "active-lease",
+			Namespace: "default",
+		},
+		Spec: coordinationv1.LeaseSpec{
+			HolderIdentity:       &holderID,
+			LeaseDurationSeconds: &leaseDuration,
+			LeaseTransitions:     &transitions,
+			RenewTime:            &recentRenewTime,
+		},
+	}
+
+	fakeClient := fake.NewSimpleClientset(lease)
+	old := swapClient(fakeClient)
+	defer swapClient(old)
+
+	holder, trans, err := getHolder("default", "active-lease")
+	if err != nil {
+		t.Fatalf("getHolder failed: %v", err)
+	}
+	if holder != holderID {
+		t.Errorf("expected holder %q, got %q", holderID, holder)
+	}
+	if trans != int64(transitions) {
+		t.Errorf("expected transitions %d, got %d", transitions, trans)
+	}
+}
+
+// TestElectionKeyFormat tests the election key construction.
+func TestElectionKeyFormat(t *testing.T) {
+	tests := []struct {
+		ns, name, want string
+	}{
+		{"default", "leader", "default/leader"},
+		{"kube-system", "my-lock", "kube-system/my-lock"},
+		{"", "bare", "/bare"},
+	}
+	for _, tc := range tests {
+		got := electionKey(tc.ns, tc.name)
+		if got != tc.want {
+			t.Errorf("electionKey(%q, %q) = %q, want %q", tc.ns, tc.name, got, tc.want)
+		}
+	}
+}
+
+// TestLeaseCRUDWithFakeClient tests basic Lease CRUD via the K8s API.
+func TestLeaseCRUDWithFakeClient(t *testing.T) {
+	fakeClient := fake.NewSimpleClientset()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	holderID := "node-a:9090"
+	transitions := int32(0)
+	lease := &coordinationv1.Lease{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "crud-test",
+			Namespace: "default",
+		},
+		Spec: coordinationv1.LeaseSpec{
+			HolderIdentity:   &holderID,
+			LeaseTransitions: &transitions,
+		},
+	}
+	created, err := fakeClient.CoordinationV1().Leases("default").Create(ctx, lease, metav1.CreateOptions{})
+	if err != nil {
+		t.Fatalf("create lease failed: %v", err)
+	}
+	if *created.Spec.HolderIdentity != holderID {
+		t.Errorf("created holder = %q, want %q", *created.Spec.HolderIdentity, holderID)
+	}
+
+	newHolder := "node-b:9090"
+	newTransitions := int32(1)
+	created.Spec.HolderIdentity = &newHolder
+	created.Spec.LeaseTransitions = &newTransitions
+	updated, err := fakeClient.CoordinationV1().Leases("default").Update(ctx, created, metav1.UpdateOptions{})
+	if err != nil {
+		t.Fatalf("update lease failed: %v", err)
+	}
+	if *updated.Spec.HolderIdentity != newHolder {
+		t.Errorf("updated holder = %q, want %q", *updated.Spec.HolderIdentity, newHolder)
+	}
+	if *updated.Spec.LeaseTransitions != newTransitions {
+		t.Errorf("updated transitions = %d, want %d", *updated.Spec.LeaseTransitions, newTransitions)
+	}
+
+	got, err := fakeClient.CoordinationV1().Leases("default").Get(ctx, "crud-test", metav1.GetOptions{})
+	if err != nil {
+		t.Fatalf("get lease failed: %v", err)
+	}
+	if *got.Spec.HolderIdentity != newHolder {
+		t.Errorf("got holder = %q, want %q", *got.Spec.HolderIdentity, newHolder)
+	}
+
+	err = fakeClient.CoordinationV1().Leases("default").Delete(ctx, "crud-test", metav1.DeleteOptions{})
+	if err != nil {
+		t.Fatalf("delete lease failed: %v", err)
+	}
+
+	_, err = fakeClient.CoordinationV1().Leases("default").Get(ctx, "crud-test", metav1.GetOptions{})
+	if err == nil {
+		t.Fatal("expected error after delete, got nil")
+	}
+}
+
+// TestConcurrentGetHolder tests concurrent calls to getHolder.
+func TestConcurrentGetHolder(t *testing.T) {
+	holderID := "concurrent-node:8080"
+	lease := &coordinationv1.Lease{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "concurrent-lease",
+			Namespace: "default",
+		},
+		Spec: coordinationv1.LeaseSpec{
+			HolderIdentity:   &holderID,
+			LeaseTransitions: ptr.To(int32(5)),
+		},
+	}
+	fakeClient := fake.NewSimpleClientset(lease)
+	old := swapClient(fakeClient)
+	defer swapClient(old)
+
+	const n = 10
+	errCh := make(chan error, n)
+	for i := 0; i < n; i++ {
+		go func() {
+			holder, trans, err := getHolder("default", "concurrent-lease")
+			if err != nil {
+				errCh <- err
+				return
+			}
+			if holder != holderID {
+				errCh <- fmt.Errorf("expected holder %q, got %q", holderID, holder)
+				return
+			}
+			if trans != 5 {
+				errCh <- fmt.Errorf("expected transitions 5, got %d", trans)
+				return
+			}
+			errCh <- nil
+		}()
+	}
+
+	for i := 0; i < n; i++ {
+		if err := <-errCh; err != nil {
+			t.Fatalf("concurrent getHolder failed: %v", err)
+		}
+	}
+}

mooncake-common/src/CMakeLists.txt

@@ -1,13 +1,64 @@
 find_package(yaml-cpp REQUIRED)
 
+find_package(asio QUIET)
+
+if(asio_FOUND)
+    message(STATUS "Found ASIO via find_package")
+    set(ASIO_INCLUDE_DIR ${asio_INCLUDE_DIR})
+else()
+    find_path(ASIO_INCLUDE_DIR
+        NAMES asio.hpp
+        PATHS
+            /usr/local/include
+            /usr/include
+            ${CMAKE_INSTALL_PREFIX}/include
+        DOC "Path to ASIO headers"
+    )
+
+    if(NOT ASIO_INCLUDE_DIR)
+        message(FATAL_ERROR "ASIO not found. Please install ASIO or set ASIO_INCLUDE_DIR manually.")
+    endif()
+
+    message(STATUS "Found ASIO at: ${ASIO_INCLUDE_DIR}")
+endif()
+
 set(MOONCAKE_COMMON_SOURCES
     default_config.cpp
     environ.cpp
 )
+
+add_library(asio_shared SHARED asio_impl.cpp)
+
+target_compile_definitions(asio_shared
+    PUBLIC
+        ASIO_SEPARATE_COMPILATION
+        ASIO_DYN_LINK
+)
+
+target_include_directories(asio_shared
+    PUBLIC
+        ${ASIO_INCLUDE_DIR}
+)
+
+set_target_properties(asio_shared PROPERTIES
+    POSITION_INDEPENDENT_CODE ON
+    INSTALL_RPATH "$ORIGIN"
+    BUILD_WITH_INSTALL_RPATH TRUE
+    OUTPUT_NAME "asio"
+    LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/mooncake-common"
+)
+
+target_link_libraries(asio_shared PUBLIC pthread)
+
 add_library(mooncake_common
     ${MOONCAKE_COMMON_SOURCES}
 )
 
+target_include_directories(mooncake_common PUBLIC
+    $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/../include>
+    $<INSTALL_INTERFACE:include>
+)
+
 target_link_libraries(mooncake_common PUBLIC
     yaml-cpp
     jsoncpp
@@ -16,3 +67,5 @@ target_link_libraries(mooncake_common PUBLIC
 if (BUILD_SHARED_LIBS)
   install(TARGETS mooncake_common DESTINATION lib)
 endif()
+
+install(TARGETS asio_shared DESTINATION lib)

mooncake-integration/store/store_py.cpp

@@ -1563,7 +1563,8 @@ PYBIND11_MODULE(store, m) {
                const std::string &rdma_devices = "",
                const std::string &master_server_addr = "127.0.0.1:50051",
                const py::object &engine = py::none(),
-               bool enable_ssd_offload = false) {
+               bool enable_ssd_offload = false,
+               const std::string &ssd_offload_path = "") {
                 auto real_client = self.init_real_client();
                 std::shared_ptr<mooncake::TransferEngine> transfer_engine =
                     nullptr;
@@ -1574,14 +1575,15 @@ PYBIND11_MODULE(store, m) {
                 return real_client->setup_real(
                     local_hostname, metadata_server, global_segment_size,
                     local_buffer_size, protocol, rdma_devices,
-                    master_server_addr, transfer_engine, "",
-                    enable_ssd_offload);
+                    master_server_addr, transfer_engine, "", enable_ssd_offload,
+                    ssd_offload_path);
             },
             py::arg("local_hostname"), py::arg("metadata_server"),
             py::arg("global_segment_size"), py::arg("local_buffer_size"),
             py::arg("protocol"), py::arg("rdma_devices"),
             py::arg("master_server_addr"), py::arg("engine") = py::none(),
-            py::arg("enable_ssd_offload") = false)
+            py::arg("enable_ssd_offload") = false,
+            py::arg("ssd_offload_path") = "")
         .def(
             "setup",
             [](MooncakeStorePyWrapper &self, const py::dict &config_dict) {
@@ -1610,7 +1612,9 @@ PYBIND11_MODULE(store, m) {
             "  rdma_devices: RDMA device list.\n"
             "  master_server_addr: Master server address.\n"
             "  ipc_socket_path: IPC socket path.\n"
-            "  enable_ssd_offload: Enable SSD offload (default false).")
+            "  enable_ssd_offload: Enable SSD offload (default false).\n"
+            "  ssd_offload_path: SSD storage directory path (overrides env "
+            "var).")
         .def(
             "setup_dummy",
             [](MooncakeStorePyWrapper &self, size_t mem_pool_size,
@@ -2267,6 +2271,20 @@ PYBIND11_MODULE(store, m) {
                 return self.store_->batch_get_replica_desc(keys);
             },
             py::arg("keys"))
+        .def(
+            "batch_replica_clear",
+            [](MooncakeStorePyWrapper &self,
+               const std::vector<std::string> &keys,
+               const std::string &segment_name) {
+                if (!self.is_client_initialized()) {
+                    LOG(ERROR) << "Client is not initialized";
+                    return std::vector<std::string>{};
+                }
+                py::gil_scoped_release release;
+                return self.store_->batch_replica_clear(keys, segment_name);
+            },
+            py::arg("keys"), py::arg("segment_name") = "",
+            "Clear replicas for the given keys. Requires lease to be expired.")
         .def(
             "create_copy_task",
             [](MooncakeStorePyWrapper &self, const std::string &key,

mooncake-p2p-store/build.sh

@@ -33,7 +33,7 @@ fi
 
 EXT_LDFLAGS="-L$BUILD_DIR/mooncake-transfer-engine/src"
 EXT_LDFLAGS+=" -L$BUILD_DIR/mooncake-transfer-engine/src/common/base"
-EXT_LDFLAGS+=" -L$BUILD_DIR/mooncake-asio"
+EXT_LDFLAGS+=" -L$BUILD_DIR/mooncake-common"
 EXT_LDFLAGS+=" -ltransfer_engine -lbase -lasio -lstdc++ -lnuma -lglog -libverbs -ljsoncpp"
 
 if [ -d "/usr/local/cuda/lib64/stubs" ]; then

mooncake-pg/include/mooncake_backend.h

@@ -125,17 +125,33 @@ class MooncakeBackend final : public ::c10d::Backend {
     }
 
     std::string getPreferredHca(std::string location) {
-        auto matrix = engine_->getLocalTopology()->getMatrix();
+        static std::once_flag topo_once;
+        static std::shared_ptr<Topology> topology;
+        static TopologyMatrix matrix;
+        std::call_once(topo_once, [this] {
+            // FIXME: getLocalTopology is deprecated in TENT
+            topology = engine_->getLocalTopology();
+            if (topology) {
+                matrix = topology->getMatrix();
+            }
+            if (!topology || matrix.empty()) {
+                topology = std::make_shared<Topology>();
+                topology->discover();
+                matrix = topology->getMatrix();
+            }
+        });
+
         auto it = matrix.find(location);
         if (it == matrix.end()) {
-            LOG(INFO) << "Topology is "
-                      << engine_->getLocalTopology()->toJson();
+            LOG(INFO) << "Topology is " << topology->toJson();
             LOG(ERROR) << "Topology entry not found for location: " << location;
-        } else if (it->second.preferred_hca.empty()) {
-            LOG(INFO) << "Topology is "
-                      << engine_->getLocalTopology()->toJson();
+            return "";
+        }
+        if (it->second.preferred_hca.empty()) {
+            LOG(INFO) << "Topology is " << topology->toJson();
             LOG(ERROR) << "Preferred HCA list is empty for location: "
                        << location;
+            return "";
         }
         return it->second.preferred_hca[0];
     }

mooncake-pg/include/mooncake_worker.cuh

@@ -11,9 +11,11 @@
 #include <transfer_engine.h>
 
 #include <memory>
+#include <atomic>
 #include <mutex>
 #include <thread>
 #include <unordered_map>
+#include <vector>
 
 namespace mooncake {
 
@@ -51,6 +53,7 @@ __global__ struct Task {
     size_t tensorSize;  // In bytes
     int64_t broadcastRoot;
     int bufferOffset;
+    uint64_t submitSequence = 0;
     BatchID batchID;
     void* transferGroupMeta;
 };
@@ -64,6 +67,12 @@ void launchReduceCpu(at::Tensor dst, size_t pos, size_t realSize, void* src,
 void preloadReduceKernels();
 
 class ConnectionContext;
+
+struct CudaTaskSubmissionToken {
+    size_t task_id;
+    uint64_t sequence;
+};
+
 class MooncakeWorker {
    public:
     explicit MooncakeWorker(int cuda_device_index = -1);
@@ -82,11 +91,11 @@ class MooncakeWorker {
         c10d::OpType opType, size_t tensorSize, int64_t broadcastRoot,
         const std::shared_ptr<TransferGroupMeta>& meta,
         const std::shared_ptr<ConnectionContext>& connection_ctx,
-        const at::cuda::CUDAStream& stream,
-        const std::function<void(void* dst, size_t pos, size_t realSize)>&
-            tensorToBuffer,
-        const std::function<void(void* src, size_t pos, size_t realSize)>&
-            bufferToTensor);
+        const at::cuda::CUDAStream& issue_stream,
+        const std::function<void(void* dst, size_t pos, size_t realSize,
+                                 const at::cuda::CUDAStream&)>& tensorToBuffer,
+        const std::function<void(void* src, size_t pos, size_t realSize,
+                                 const at::cuda::CUDAStream&)>& bufferToTensor);
 
     void Start();
 
@@ -104,6 +113,10 @@ class MooncakeWorker {
      */
     bool drainTasks(const TransferGroupMeta* meta) const;
 
+    bool waitUntilTasksSubmitted(
+        const std::vector<CudaTaskSubmissionToken>& tasks,
+        std::chrono::milliseconds timeout) const;
+
    private:
     void startWorker();
 
@@ -122,6 +135,8 @@ class MooncakeWorker {
 
     int cpuTaskCount = 0;
     int cudaTaskCount = 0;
+    std::atomic<uint64_t> next_cuda_task_sequence_{1};
+    std::atomic<uint64_t> submitted_task_sequence_[kNumTasks_]{};
 
     std::thread worker_thread_;
 };

mooncake-pg/src/mooncake_backend.cpp

@@ -103,17 +103,18 @@ MooncakeBackend::MooncakeBackend(
     const int size = distBackendOpts.group_size;
     const auto& globalRanks = distBackendOpts.global_ranks_in_group;
 
-    // Get device data
-    std::string location;
-    int deviceCount = 0;
-    cudaError_t err = cudaGetDeviceCount(&deviceCount);
-    if (err != cudaSuccess || deviceCount == 0) {
-        location = kWildcardLocation;
-    } else {
-        int deviceId_;
-        err = cudaGetDevice(&deviceId_);
-        TORCH_CHECK(!err, c10::str("Failed to get device id"));
-        location = GPU_PREFIX + std::to_string(deviceId_);
+    // Memory location for device specific buffers
+    // always kWildcardLocation for cpu backend
+    std::string location = kWildcardLocation;
+    if (!isCpu) {
+        int deviceCount = 0;
+        cudaError_t err = cudaGetDeviceCount(&deviceCount);
+        if (err == cudaSuccess && deviceCount != 0) {
+            int deviceId_;
+            err = cudaGetDevice(&deviceId_);
+            TORCH_CHECK(!err, c10::str("Failed to get device id"));
+            location = GPU_PREFIX + std::to_string(deviceId_);
+        }
     }
 
     // Initialize transfer engine
@@ -390,15 +391,18 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::broadcast(
         return worker_->putTaskCuda(
             c10d::OpType::BROADCAST, tensorSize, root, meta_, connection_ctx_,
             stream,
-            [=](void* dst, size_t pos, size_t realSize) {
+            [=](void* dst, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 if (isRoot) {
                     cudaMemcpyAsync(dst, (char*)tensor.data_ptr() + pos,
-                                    realSize, cudaMemcpyDeviceToDevice, stream);
+                                    realSize, cudaMemcpyDeviceToDevice,
+                                    enq_stream);
                 }
             },
-            [=](void* src, size_t pos, size_t realSize) {
+            [=](void* src, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 cudaMemcpyAsync((char*)tensor.data_ptr() + pos, src, realSize,
-                                cudaMemcpyDeviceToDevice, stream);
+                                cudaMemcpyDeviceToDevice, enq_stream);
             });
     }
 }
@@ -426,16 +430,18 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::allreduce(
         return worker_->putTaskCuda(
             c10d::OpType::ALLREDUCE, tensorSize, 0, meta_, connection_ctx_,
             stream,
-            [=](void* dst, size_t pos, size_t realSize) {
+            [=](void* dst, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 cudaMemcpyAsync(dst, (char*)tensor.data_ptr() + pos, realSize,
-                                cudaMemcpyDeviceToDevice, stream);
+                                cudaMemcpyDeviceToDevice, enq_stream);
             },
-            [=, this](void* src, size_t pos, size_t realSize) {
+            [=, this](void* src, size_t pos, size_t realSize,
+                      const at::cuda::CUDAStream& enq_stream) {
                 cudaMemsetAsync((char*)tensor.data_ptr() + pos, 0, realSize,
-                                stream);
+                                enq_stream);
                 launchReduceKernel(tensor, pos, realSize, src, meta_->size,
                                    opts.reduceOp, meta_->activeRanksDevice,
-                                   stream);
+                                   enq_stream);
             });
     }
 }
@@ -466,15 +472,17 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::allgather(
         return worker_->putTaskCuda(
             c10d::OpType::ALLGATHER, tensorSize, 0, meta_, connection_ctx_,
             stream,
-            [=](void* dst, size_t pos, size_t realSize) {
+            [=](void* dst, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 cudaMemcpyAsync(dst, (char*)inputTensor.data_ptr() + pos,
-                                realSize, cudaMemcpyDeviceToDevice, stream);
+                                realSize, cudaMemcpyDeviceToDevice, enq_stream);
             },
-            [=](void* src, size_t pos, size_t realSize) {
+            [=](void* src, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 for (const auto j : c10::irange(outputTensors_.size())) {
                     cudaMemcpyAsync((char*)outputTensors_[j].data_ptr() + pos,
                                     (char*)src + j * realSize, realSize,
-                                    cudaMemcpyDeviceToDevice, stream);
+                                    cudaMemcpyDeviceToDevice, enq_stream);
                 }
             });
     }
@@ -505,16 +513,18 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::_allgather_base(
         return worker_->putTaskCuda(
             c10d::OpType::_ALLGATHER_BASE, tensorSize, 0, meta_,
             connection_ctx_, stream,
-            [=](void* dst, size_t pos, size_t realSize) {
+            [=](void* dst, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 cudaMemcpyAsync(dst, (char*)inputBuffer.data_ptr() + pos,
-                                realSize, cudaMemcpyDeviceToDevice, stream);
+                                realSize, cudaMemcpyDeviceToDevice, enq_stream);
             },
-            [=, this](void* src, size_t pos, size_t realSize) {
+            [=, this](void* src, size_t pos, size_t realSize,
+                      const at::cuda::CUDAStream& enq_stream) {
                 for (const auto j : c10::irange(meta_->size)) {
                     cudaMemcpyAsync(
                         (char*)outputBuffer.data_ptr() + j * tensorSize + pos,
                         (char*)src + j * realSize, realSize,
-                        cudaMemcpyDeviceToDevice, stream);
+                        cudaMemcpyDeviceToDevice, enq_stream);
                 }
             });
     }
@@ -547,20 +557,22 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::_reduce_scatter_base(
         return worker_->putTaskCuda(
             c10d::OpType::_REDUCE_SCATTER_BASE, tensorSize, 0, meta_,
             connection_ctx_, stream,
-            [=, this](void* dst, size_t pos, size_t realSize) {
+            [=, this](void* dst, size_t pos, size_t realSize,
+                      const at::cuda::CUDAStream& enq_stream) {
                 for (const auto j : c10::irange(meta_->size)) {
                     cudaMemcpyAsync(
                         (char*)dst + j * realSize,
                         (char*)inputBuffer.data_ptr() + j * tensorSize + pos,
-                        realSize, cudaMemcpyDeviceToDevice, stream);
+                        realSize, cudaMemcpyDeviceToDevice, enq_stream);
                 }
             },
-            [=, this](void* src, size_t pos, size_t realSize) {
+            [=, this](void* src, size_t pos, size_t realSize,
+                      const at::cuda::CUDAStream& enq_stream) {
                 cudaMemsetAsync((char*)outputBuffer.data_ptr() + pos, 0,
-                                realSize, stream);
+                                realSize, enq_stream);
                 launchReduceKernel(outputBuffer, pos, realSize, src,
                                    meta_->size, opts.reduceOp,
-                                   meta_->activeRanksDevice, stream);
+                                   meta_->activeRanksDevice, enq_stream);
             });
     }
 }
@@ -591,18 +603,21 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::alltoall(
         return worker_->putTaskCuda(
             c10d::OpType::ALLTOALL, tensorSize, 0, meta_, connection_ctx_,
             stream,
-            [=](void* dst, size_t pos, size_t realSize) {
+            [=](void* dst, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 for (const auto j : c10::irange(inputTensors.size())) {
                     cudaMemcpyAsync((char*)dst + j * realSize,
                                     (char*)inputTensors[j].data_ptr() + pos,
-                                    realSize, cudaMemcpyDeviceToDevice, stream);
+                                    realSize, cudaMemcpyDeviceToDevice,
+                                    enq_stream);
                 }
             },
-            [=](void* src, size_t pos, size_t realSize) {
+            [=](void* src, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 for (const auto j : c10::irange(outputTensors.size())) {
                     cudaMemcpyAsync((char*)outputTensors[j].data_ptr() + pos,
                                     (char*)src + j * realSize, realSize,
-                                    cudaMemcpyDeviceToDevice, stream);
+                                    cudaMemcpyDeviceToDevice, enq_stream);
                 }
             });
     }
@@ -621,8 +636,9 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::barrier(
         auto stream = at::cuda::getCurrentCUDAStream(device_index);
         return worker_->putTaskCuda(
             c10d::OpType::BARRIER, kBarrierDummyTensorSize, 0, meta_,
-            connection_ctx_, stream, [=](void*, size_t, size_t) {},
-            [=](void*, size_t, size_t) {});
+            connection_ctx_, stream,
+            [=](void*, size_t, size_t, const at::cuda::CUDAStream&) {},
+            [=](void*, size_t, size_t, const at::cuda::CUDAStream&) {});
     }
 }
 
@@ -652,17 +668,19 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::reduce(
         return worker_->putTaskCuda(
             c10d::OpType::REDUCE, tensorSize, root, meta_, connection_ctx_,
             stream,
-            [=](void* dst, size_t pos, size_t realSize) {
+            [=](void* dst, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 cudaMemcpyAsync(dst, (char*)tensor.data_ptr() + pos, realSize,
-                                cudaMemcpyDeviceToDevice, stream);
+                                cudaMemcpyDeviceToDevice, enq_stream);
             },
-            [=, this](void* src, size_t pos, size_t realSize) {
+            [=, this](void* src, size_t pos, size_t realSize,
+                      const at::cuda::CUDAStream& enq_stream) {
                 if (isRoot) {
                     cudaMemsetAsync((char*)tensor.data_ptr() + pos, 0, realSize,
-                                    stream);
+                                    enq_stream);
                     launchReduceKernel(tensor, pos, realSize, src, meta_->size,
                                        opts.reduceOp, meta_->activeRanksDevice,
-                                       stream);
+                                       enq_stream);
                 }
             });
     }
@@ -700,18 +718,20 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::gather(
         return worker_->putTaskCuda(
             c10d::OpType::GATHER, tensorSize, root, meta_, connection_ctx_,
             stream,
-            [=](void* dst, size_t pos, size_t realSize) {
+            [=](void* dst, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 cudaMemcpyAsync(dst, (char*)inputTensor.data_ptr() + pos,
-                                realSize, cudaMemcpyDeviceToDevice, stream);
+                                realSize, cudaMemcpyDeviceToDevice, enq_stream);
             },
-            [=](void* src, size_t pos, size_t realSize) {
+            [=](void* src, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 if (isRoot) {
                     auto outputTensors_ = outputTensors.back();
                     for (const auto j : c10::irange(outputTensors_.size())) {
                         cudaMemcpyAsync(
                             (char*)outputTensors_[j].data_ptr() + pos,
                             (char*)src + j * realSize, realSize,
-                            cudaMemcpyDeviceToDevice, stream);
+                            cudaMemcpyDeviceToDevice, enq_stream);
                     }
                 }
             });
@@ -752,20 +772,22 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::scatter(
         return worker_->putTaskCuda(
             c10d::OpType::SCATTER, tensorSize, root, meta_, connection_ctx_,
             stream,
-            [=](void* dst, size_t pos, size_t realSize) {
+            [=](void* dst, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 if (isRoot) {
                     auto inputTensors_ = inputTensors.back();
                     for (const auto j : c10::irange(inputTensors_.size())) {
                         cudaMemcpyAsync(
                             (char*)dst + j * realSize,
                             (char*)inputTensors_[j].data_ptr() + pos, realSize,
-                            cudaMemcpyDeviceToDevice, stream);
+                            cudaMemcpyDeviceToDevice, enq_stream);
                     }
                 }
             },
-            [=](void* src, size_t pos, size_t realSize) {
+            [=](void* src, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 cudaMemcpyAsync((char*)outputTensor.data_ptr() + pos, src,
-                                realSize, cudaMemcpyDeviceToDevice, stream);
+                                realSize, cudaMemcpyDeviceToDevice, enq_stream);
             });
     }
 }

mooncake-pg/src/mooncake_worker.cu

@@ -3,6 +3,9 @@
 #include <memory>
 #include <thread>
 #include <mooncake_worker.cuh>
+#include <ATen/cuda/CUDAGraphsUtils.cuh>
+
+#include "pg_utils.h"
 
 namespace mooncake {
 
@@ -30,18 +33,100 @@ class MooncakeWorkCpu : public ::c10d::Work {
 class MooncakeWorkCuda : public ::c10d::Work {
    public:
     MooncakeWorkCuda(c10d::OpType opType, std::shared_ptr<torch::Event> event,
-                     std::shared_ptr<TransferGroupMeta> meta)
-        : Work(-1, opType), event_(std::move(event)), meta_(std::move(meta)) {}
+                     std::shared_ptr<TransferGroupMeta> meta,
+                     const MooncakeWorker* worker,
+                     std::vector<CudaTaskSubmissionToken> submitted_tasks)
+        : Work(-1, opType),
+          event_(std::move(event)),
+          meta_(std::move(meta)),
+          worker_(worker),
+          submitted_tasks_(std::move(submitted_tasks)) {}
 
     bool isCompleted() override { return event_->query(); }
 
     bool wait(std::chrono::milliseconds timeout) override {
-        return true;  // This should be a no-op
+        // Wait until the task has been submitted to TransferEngine:
+        // This tries to ensure that the CUDA kernels required for the transfer
+        // have been launched by the time `waitUntilTasksSubmitted` returns.
+        //
+        // Why is this needed? PyTorch documentation implies that collective
+        // operations should be enqueued when `wait()` returns. In practice, we
+        // found that violating this causes hangs.
+        //
+        // Our current hypothesis for the hang is: PyTorch assumes the kernels
+        // needed for the transfer are already launched when `wait` returns
+        // true. It may then launch subsequent operations after the collective
+        // (e.g., `.cpu()`). Such operations may acquire a process-wide lock in
+        // the CUDA runtime. Also, they may rely on the data produced by the
+        // collective, thus causing a synchronization on enq_stream. However,
+        // holding that runtime lock prevents cudaMemcpy(Async) in TE/TENT from
+        // launching. This means the transfer can't finish, and enq_stream won't
+        // complete. Thus, a deadlock occurs.
+        // (In practice, we found that replacing all cudaMemcpyAsync in TENT
+        // with cuMemcpyAsync actually alleviates this, which further suggests a
+        // deadlock in the CUDA runtime. However, that change is too invasive
+        // for TE/TENT, so we do not adopt it here.)
+        //
+        // Strictly speaking, the wait is needed for another reason: The current
+        // stream will be blocked on the event below. Any subsequent work on
+        // `current_stream` will wait on that event, which effectively waits for
+        // the task to be done. Therefore, we must ensure all kernels needed for
+        // the transfer task are launched BEFORE blocking the current stream, in
+        // case TE/TENT use `current_stream` to launch those kernels (though it
+        // is rare).
+        //
+        // Please note that this logic relies on the assumption that TE/TENT
+        // will launch all CUDA operations in `submitTransfer`.
+        // Unfortunately, TcpTransport in TE and TENT currently violates this
+        // assumption (cudaMemcpy(Async) may be called later from a callback),
+        // which can cause hangs in PG when a CUDA operation such as
+        // `x.cpu().item()` follows the collective. For TE's TcpTransport, the
+        // use of cudaMemcpy on the default stream may also contribute to the
+        // hang.
+        //
+        // Besides, for CPU-only transports (like RdmaTransport),
+        // waitUntilTasksSubmitted is totally unnecessary, but we keep it for
+        // uniform behavior to avoid invasive changes to TE/TENT.
+        bool submitted = true;
+        if (at::cuda::currentStreamCaptureStatus() ==
+            c10::cuda::CaptureStatus::None) {
+            // Normal execution: block until tasks are submitted.
+            submitted =
+                worker_->waitUntilTasksSubmitted(submitted_tasks_, timeout);
+        } else {
+            // During CUDA graph capture, kernels are recorded but not actually
+            // executed. The enqueueTaskKernel would never run, so
+            // waitUntilTasksSubmitted would hang because the CPU worker thread
+            // never sees task.active == true.
+            //
+            // Note that this also means NvlinkTransport (and TcpTransport too,
+            // of course) won't work with CUDA Graphs: Kernels launched inside
+            // TE/TENT can't be captured by the graph, and during replay they
+            // are not ordered with the graph execution. This may trigger the
+            // same deadlock described above.
+        }
+        if (!submitted) return false;
+
+        // Once all tasks have been submitted, use the event to synchronize
+        // the current stream and the enqueue stream, but do not wait on this
+        // event.
+        //
+        // See PyTorch docs for more details:
+        // https://docs.pytorch.org/docs/stable/distributed.html#synchronous-and-asynchronous-collective-operations
+        //   "wait() - in the case of CPU collectives, will block the process
+        //    until the operation is completed. In the case of CUDA collectives,
+        //    will block the currently active CUDA stream until the operation
+        //    is completed (but will not block the CPU)."
+        auto current_stream = at::cuda::getCurrentCUDAStream();
+        event_->block(current_stream);
+        return true;
     }
 
    protected:
     std::shared_ptr<torch::Event> event_;
     std::shared_ptr<TransferGroupMeta> meta_;
+    const MooncakeWorker* worker_;
+    std::vector<CudaTaskSubmissionToken> submitted_tasks_;
 };
 
 class MooncakeBarrierWorkCuda : public MooncakeWorkCuda {
@@ -49,29 +134,33 @@ class MooncakeBarrierWorkCuda : public MooncakeWorkCuda {
     using MooncakeWorkCuda::MooncakeWorkCuda;
 
     bool wait(std::chrono::milliseconds timeout) override {
+        // Skip host-side synchronization during CUDA graph capture.
+        // cudaEventSynchronize is not permitted while a stream is capturing.
+        if (at::cuda::currentStreamCaptureStatus() !=
+            c10::cuda::CaptureStatus::None) {
+            // We still need stream-level synchronization so that subsequent
+            // operations on the capture stream are ordered after the barrier
+            // task on the enqueue stream.
+            auto current_stream = at::cuda::getCurrentCUDAStream();
+            event_->block(current_stream);
+            return true;
+        }
+
         if (timeout == kNoTimeout) {
             event_->synchronize();
             return true;
         }
 
-        auto start = std::chrono::steady_clock::now();
-        while (!event_->query()) {
-            auto now = std::chrono::steady_clock::now();
-            auto elapsed =
-                std::chrono::duration_cast<std::chrono::milliseconds>(now -
-                                                                      start);
-            if (elapsed >= timeout) {
-                return false;
-            }
-            std::this_thread::sleep_for(std::chrono::microseconds(10));
-        }
-        return true;
+        BackoffWaiter waiter(
+            BackoffWaiterConfig::constantSleep(std::chrono::microseconds(10)));
+        return waiter.wait_for(timeout, [this] { return event_->query(); });
     }
 };
 
 __global__ void enqueueTaskKernel(c10d::OpType opType, size_t tensorSize,
                                   int64_t broadcastRoot, int bufferOffset,
-                                  void* meta, Task* tasks, int numRanks,
+                                  uint64_t submitSequence, void* meta,
+                                  Task* tasks, int numRanks,
                                   const bool* activeRanks,
                                   int* activeRanksTensor, size_t taskId) {
     // Copy task into slot
@@ -79,15 +168,16 @@ __global__ void enqueueTaskKernel(c10d::OpType opType, size_t tensorSize,
     tasks[taskId].tensorSize = tensorSize;
     tasks[taskId].broadcastRoot = broadcastRoot;
     tasks[taskId].bufferOffset = bufferOffset;
+    tasks[taskId].submitSequence = submitSequence;
     tasks[taskId].transferGroupMeta = meta;
 
-    // Mark active
-    __threadfence();  // Ensure writes visible to host
+    // Publish task metadata before notifying the host worker thread.
+    __threadfence_system();
     tasks[taskId].active = true;
 
     // Spin-wait until CPU proxy sets DONE
     while (tasks[taskId].active) {
-        __threadfence();
+        __threadfence_system();
     }
     for (int i = 0; i < numRanks; ++i) {
         activeRanksTensor[i] = activeRanks[i] ? 1 : 0;
@@ -311,6 +401,8 @@ MooncakeWorker::MooncakeWorker(int cuda_device_index)
     }
     for (size_t i = 0; i < kNumTasks_; ++i) {
         tasks_[i].active = false;
+        tasks_[i].submitSequence = 0;
+        submitted_task_sequence_[i].store(0, std::memory_order_relaxed);
     }
 }
 
@@ -400,45 +492,65 @@ c10::intrusive_ptr<c10d::Work> MooncakeWorker::putTaskCuda(
     c10d::OpType opType, size_t tensorSize, int64_t broadcastRoot,
     const std::shared_ptr<TransferGroupMeta>& meta,
     const std::shared_ptr<ConnectionContext>& connection_ctx,
-    const at::cuda::CUDAStream& stream,
-    const std::function<void(void* dst, size_t pos, size_t realSize)>&
-        tensorToBuffer,
-    const std::function<void(void* src, size_t pos, size_t realSize)>&
-        bufferToTensor) {
+    const at::cuda::CUDAStream& issue_stream,
+    const std::function<void(void* dst, size_t pos, size_t realSize,
+                             const at::cuda::CUDAStream&)>& tensorToBuffer,
+    const std::function<void(void* src, size_t pos, size_t realSize,
+                             const at::cuda::CUDAStream&)>& bufferToTensor) {
     connection_ctx->waitUntilNewRanksConnected();
 
     // TORCH_CHECK(tensorSize * meta->size < kBufferSize, "Too large!");
     //  Alternately use even-odd items to maintain tasks
     size_t chunkSize = ((kBufferSize - 1) / meta->size) & ~(size_t)7;
 
+    // Get a non-blocking stream for enqueue:
+    //   The incoming `issue_stream` may be the Null Stream, which enforces
+    //   implicit synchronization semantics. Launching a spin-wait kernel
+    //   (enqueueTaskKernel) on such a stream can introduce potential deadlock.
+    at::cuda::CUDAStream enq_stream =
+        at::cuda::getStreamFromPool(false, issue_stream.device_index());
+
+    // Synchronize: enq_stream waits for issue_stream
+    auto event_start = std::make_shared<torch::Event>(torch::kCUDA);
+    event_start->record(issue_stream);
+    event_start->block(enq_stream);
+
+    std::vector<CudaTaskSubmissionToken> submitted_tasks;
+    submitted_tasks.reserve((tensorSize + chunkSize - 1) / chunkSize);
     for (size_t pos = 0; pos < tensorSize; pos += chunkSize) {
         size_t realSize = min(tensorSize, pos + chunkSize) - pos;
         int taskId = cudaTaskCount % 2 + 2;
         int bufferOffset = meta->taskCount % 2;
+        const uint64_t taskSequence =
+            next_cuda_task_sequence_.fetch_add(1, std::memory_order_relaxed);
+        submitted_tasks.push_back(
+            {.task_id = static_cast<size_t>(taskId), .sequence = taskSequence});
         tensorToBuffer(
             (void*)meta->segmentInfos[meta->rank].send_buffer[bufferOffset],
-            pos, realSize);
+            pos, realSize, enq_stream);
 
         hasCallback_[taskId] = false;
-        enqueueTaskKernel<<<1, 1, 0, stream>>>(
-            opType, realSize, broadcastRoot, bufferOffset, meta.get(),
-            tasks_device_, meta->size, meta->activeRanksDevice,
+        enqueueTaskKernel<<<1, 1, 0, enq_stream>>>(
+            opType, realSize, broadcastRoot, bufferOffset, taskSequence,
+            meta.get(), tasks_device_, meta->size, meta->activeRanksDevice,
             meta->activeRanksTensor.data_ptr<int>(), taskId);
         bufferToTensor(
             (void*)meta->segmentInfos[meta->rank].recv_buffer[bufferOffset],
-            pos, realSize);
+            pos, realSize, enq_stream);
 
         ++cudaTaskCount;
         ++meta->taskCount;
     }
 
-    auto event = std::make_shared<torch::Event>(torch::kCUDA);
-    event->record(stream);
+    auto event_end = std::make_shared<torch::Event>(torch::kCUDA);
+    event_end->record(enq_stream);
+
     if (opType == c10d::OpType::BARRIER) {
-        return c10::make_intrusive<MooncakeBarrierWorkCuda>(opType, event,
-                                                            meta);
+        return c10::make_intrusive<MooncakeBarrierWorkCuda>(
+            opType, event_end, meta, this, std::move(submitted_tasks));
     }
-    return c10::make_intrusive<MooncakeWorkCuda>(opType, event, meta);
+    return c10::make_intrusive<MooncakeWorkCuda>(opType, event_end, meta, this,
+                                                 std::move(submitted_tasks));
 }
 
 }  // namespace mooncake

mooncake-pg/src/mooncake_worker_thread.cpp

@@ -35,6 +35,36 @@ bool MooncakeWorker::drainTasks(const TransferGroupMeta* meta) const {
         });
 }
 
+bool MooncakeWorker::waitUntilTasksSubmitted(
+    const std::vector<CudaTaskSubmissionToken>& tasks,
+    std::chrono::milliseconds timeout) const {
+    if (tasks.empty()) {
+        return true;
+    }
+
+    auto submitted = [this, &tasks] {
+        for (const auto& task : tasks) {
+            if (task.task_id >= kNumTasks_) {
+                LOG(ERROR) << "Invalid task id.";
+                return true;
+            }
+            if (submitted_task_sequence_[task.task_id].load(
+                    std::memory_order_acquire) < task.sequence) {
+                return false;
+            }
+        }
+        return true;
+    };
+
+    BackoffWaiter waiter(
+        BackoffWaiterConfig::constantSleep(std::chrono::microseconds(10)));
+    if (timeout == kNoTimeout) {
+        waiter.wait(submitted);
+        return true;
+    }
+    return waiter.wait_for(timeout, submitted);
+}
+
 void MooncakeWorker::startWorker() {
     running_ = true;
     worker_thread_ = std::thread([this] {
@@ -61,7 +91,10 @@ void MooncakeWorker::startWorker() {
                                      group->rank != task.broadcastRoot) ||
                                     task.opType == c10d::OpType::BARRIER;
                 if (task_status[i].load(std::memory_order_acquire) == IDLE) {
+                    const auto submit_sequence = task.submitSequence;
                     if (skipTransfer) {
+                        submitted_task_sequence_[i].store(
+                            submit_sequence, std::memory_order_release);
                         task_status[i].store(TRANSFERRED_1,
                                              std::memory_order_release);
                         continue;
@@ -133,6 +166,8 @@ void MooncakeWorker::startWorker() {
                     task.batchID =
                         group->engine->allocateBatchID(entries.size());
                     group->engine->submitTransfer(task.batchID, entries);
+                    submitted_task_sequence_[i].store(
+                        submit_sequence, std::memory_order_release);
                     activeTime[i] = clock::now();
                     task_status[i].store(TRANSFERRED_1,
                                          std::memory_order_release);

mooncake-store/go/build.sh

@@ -39,7 +39,7 @@ CGO_LDFLAGS="-L${BUILD_DIR}/mooncake-store/src"
 CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-store/src/cachelib_memory_allocator"
 CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-transfer-engine/src"
 CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-transfer-engine/src/common/base"
-CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-asio"
+CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-common"
 CGO_LDFLAGS+=" -lmooncake_store -lcachelib_memory_allocator -ltransfer_engine -lbase -lasio"
 CGO_LDFLAGS+=" -lstdc++ -lnuma -lglog -lgflags -libverbs -ljsoncpp -lzstd -lcurl"
 

mooncake-store/include/client_metric.h

@@ -1,10 +1,15 @@
 #pragma once
 
+#include <algorithm>
+#include <array>
 #include <atomic>
 #include <chrono>
 #include <iomanip>
+#include <mutex>
+#include <optional>
 #include <sstream>
 #include <thread>
+#include <unordered_set>
 #include <vector>
 #include <ylt/metric/counter.hpp>
 #include <ylt/metric/histogram.hpp>
@@ -46,6 +51,56 @@ const inline std::map<std::string, std::string> merge_labels(
     return merged_labels;
 }
 
+inline std::string format_metric_rate(double value, const char* suffix) {
+    const double KB = 1024.0;
+    const double MB = KB * 1024.0;
+    const double GB = MB * 1024.0;
+    const double TB = GB * 1024.0;
+
+    std::ostringstream oss;
+    oss << std::fixed << std::setprecision(2);
+    if (value >= TB) {
+        oss << value / TB << " T" << suffix;
+    } else if (value >= GB) {
+        oss << value / GB << " G" << suffix;
+    } else if (value >= MB) {
+        oss << value / MB << " M" << suffix;
+    } else if (value >= KB) {
+        oss << value / KB << " K" << suffix;
+    } else {
+        oss << value << " " << suffix;
+    }
+    return oss.str();
+}
+
+inline std::string format_metric_bandwidth(uint64_t total_bytes,
+                                           double elapsed_seconds) {
+    return format_metric_rate(total_bytes / elapsed_seconds, "B/s");
+}
+
+inline uint64_t elapsed_us_since(
+    std::chrono::steady_clock::time_point start_time) {
+    return static_cast<uint64_t>(
+        std::chrono::duration_cast<std::chrono::microseconds>(
+            std::chrono::steady_clock::now() - start_time)
+            .count());
+}
+
+template <typename Result, typename Operation, typename SuccessFn,
+          typename ObserveFn>
+Result execute_timed_operation(Operation&& operation, SuccessFn&& success_fn,
+                               ObserveFn&& observe_fn) {
+    const auto start_time = std::chrono::steady_clock::now();
+    Result result = std::forward<Operation>(operation)();
+    if (std::forward<SuccessFn>(success_fn)(result)) {
+        std::forward<ObserveFn>(observe_fn)(elapsed_us_since(start_time),
+                                            result);
+    }
+    return result;
+}
+
+enum class TransferOperationKind { kRead, kWrite };
+
 struct TransferMetric {
     TransferMetric(std::map<std::string, std::string> labels = {})
         : total_read_bytes("mooncake_transfer_read_bytes", "Total bytes read",
@@ -61,7 +116,8 @@ struct TransferMetric {
           get_latency_us("mooncake_transfer_get_latency",
                          "Get transfer latency (us)", kLatencyBucket, labels),
           put_latency_us("mooncake_transfer_put_latency",
-                         "Put transfer latency (us)", kLatencyBucket, labels) {}
+                         "Put transfer latency (us)", kLatencyBucket, labels),
+          start_time_(std::chrono::steady_clock::now()) {}
 
     ylt::metric::counter_t total_read_bytes;
     ylt::metric::counter_t total_write_bytes;
@@ -79,7 +135,7 @@ struct TransferMetric {
         put_latency_us.serialize(str);
     }
 
-    std::string summary_metrics() {
+    std::string summary_metrics(bool include_bandwidth = true) {
         std::stringstream ss;
         ss << "=== Transfer Metrics Summary ===\n";
 
@@ -88,6 +144,14 @@ struct TransferMetric {
         auto write_bytes = total_write_bytes.value();
         ss << "Total Read: " << byte_size_to_string(read_bytes) << "\n";
         ss << "Total Write: " << byte_size_to_string(write_bytes) << "\n";
+        if (include_bandwidth) {
+            ss << "Average Read Throughput: "
+               << format_metric_bandwidth(read_bytes, elapsed_seconds())
+               << "\n";
+            ss << "Average Write Throughput: "
+               << format_metric_bandwidth(write_bytes, elapsed_seconds())
+               << "\n";
+        }
 
         // Latency summaries
         ss << "\n=== Latency Summary (microseconds) ===\n";
@@ -102,6 +166,14 @@ struct TransferMetric {
     }
 
    private:
+    std::chrono::steady_clock::time_point start_time_;
+
+    double elapsed_seconds() const {
+        const auto elapsed = std::chrono::duration<double>(
+            std::chrono::steady_clock::now() - start_time_);
+        return std::max(elapsed.count(), 1e-9);
+    }
+
     std::string format_latency_summary(ylt::metric::histogram_t& hist) {
         // Access the internal sum and bucket counts
         auto sum_ptr =
@@ -274,6 +346,157 @@ struct MasterClientMetric {
     }
 };
 
+struct TransferOperationMetric {
+    std::array<std::string, 1> op_names = {"op_name"};
+
+    explicit TransferOperationMetric(
+        std::map<std::string, std::string> labels = {})
+        : read_op_count("mooncake_transfer_read_operation_count",
+                        "Total read operations by interface type", labels,
+                        op_names),
+          read_op_bytes("mooncake_transfer_read_operation_bytes",
+                        "Total read bytes by interface type", labels, op_names),
+          read_op_latency_us("mooncake_transfer_read_operation_latency",
+                             "Read operation latency by interface type (us)",
+                             kLatencyBucket, labels, op_names),
+          write_op_count("mooncake_transfer_write_operation_count",
+                         "Total write operations by interface type", labels,
+                         op_names),
+          write_op_bytes("mooncake_transfer_write_operation_bytes",
+                         "Total write bytes by interface type", labels,
+                         op_names),
+          write_op_latency_us("mooncake_transfer_write_operation_latency",
+                              "Write operation latency by interface type (us)",
+                              kLatencyBucket, labels, op_names) {}
+
+    ylt::metric::hybrid_counter_1t read_op_count;
+    ylt::metric::hybrid_counter_1t read_op_bytes;
+    ylt::metric::hybrid_histogram_1t read_op_latency_us;
+    ylt::metric::hybrid_counter_1t write_op_count;
+    ylt::metric::hybrid_counter_1t write_op_bytes;
+    ylt::metric::hybrid_histogram_1t write_op_latency_us;
+
+    void Observe(TransferOperationKind kind, const std::string& op_name,
+                 uint64_t bytes, uint64_t latency_us) {
+        const std::array<std::string, 1> label = {op_name};
+        {
+            std::lock_guard<std::mutex> lock(observed_ops_mutex_);
+            if (kind == TransferOperationKind::kRead) {
+                observed_read_ops_.insert(op_name);
+            } else {
+                observed_write_ops_.insert(op_name);
+            }
+        }
+
+        if (kind == TransferOperationKind::kRead) {
+            read_op_count.inc(label);
+            read_op_bytes.inc(label, bytes);
+            read_op_latency_us.observe(label, latency_us);
+        } else {
+            write_op_count.inc(label);
+            write_op_bytes.inc(label, bytes);
+            write_op_latency_us.observe(label, latency_us);
+        }
+    }
+
+    void serialize(std::string& str) {
+        read_op_count.serialize(str);
+        read_op_bytes.serialize(str);
+        read_op_latency_us.serialize(str);
+        write_op_count.serialize(str);
+        write_op_bytes.serialize(str);
+        write_op_latency_us.serialize(str);
+    }
+
+    std::string summary_metrics() {
+        std::stringstream ss;
+        ss << "=== Interface Operation Metrics Summary ===\n";
+        ss << format_operation_group_summary(
+                  "Read Interfaces", snapshot_operations(observed_read_ops_),
+                  read_op_count, read_op_bytes, read_op_latency_us)
+           << "\n";
+        ss << format_operation_group_summary(
+            "Write Interfaces", snapshot_operations(observed_write_ops_),
+            write_op_count, write_op_bytes, write_op_latency_us);
+        return ss.str();
+    }
+
+   private:
+    std::mutex observed_ops_mutex_;
+    std::unordered_set<std::string> observed_read_ops_;
+    std::unordered_set<std::string> observed_write_ops_;
+
+    std::vector<std::string> snapshot_operations(
+        const std::unordered_set<std::string>& source) {
+        std::lock_guard<std::mutex> lock(observed_ops_mutex_);
+        std::vector<std::string> ops(source.begin(), source.end());
+        std::sort(ops.begin(), ops.end());
+        return ops;
+    }
+
+    std::string format_operation_group_summary(
+        const std::string& group_name, const std::vector<std::string>& ops,
+        ylt::metric::hybrid_counter_1t& op_count,
+        ylt::metric::hybrid_counter_1t& op_bytes,
+        ylt::metric::hybrid_histogram_1t& op_latency_us) {
+        std::stringstream ss;
+        ss << group_name << ":\n";
+        if (ops.empty()) {
+            ss << "No data";
+            return ss.str();
+        }
+
+        auto bucket_counts = op_latency_us.get_bucket_counts();
+        bool found_any = false;
+        for (const auto& op_name : ops) {
+            const std::array<std::string, 1> label = {op_name};
+            const int64_t total_count = op_count.value(label);
+            if (total_count == 0) {
+                continue;
+            }
+
+            found_any = true;
+            ss << op_name << ": count=" << total_count << ", bytes="
+               << byte_size_to_string(
+                      static_cast<uint64_t>(op_bytes.value(label)));
+
+            int64_t p95_target = (total_count * 95) / 100;
+            int64_t cumulative = 0;
+            double p95_bucket = 0;
+            for (size_t i = 0;
+                 i < bucket_counts.size() && i < kLatencyBucket.size(); ++i) {
+                cumulative += bucket_counts[i]->value(label);
+                if (cumulative >= p95_target && p95_bucket == 0) {
+                    p95_bucket = kLatencyBucket[i];
+                    break;
+                }
+            }
+            if (p95_bucket > 0) {
+                ss << ", p95<" << p95_bucket << "μs";
+            }
+
+            double max_bucket = 0;
+            for (size_t i = bucket_counts.size(); i > 0; --i) {
+                const size_t idx = i - 1;
+                if (idx < kLatencyBucket.size() &&
+                    bucket_counts[idx]->value(label) > 0) {
+                    max_bucket = kLatencyBucket[idx];
+                    break;
+                }
+            }
+            if (max_bucket > 0) {
+                ss << ", max<" << max_bucket << "μs";
+            }
+            ss << "\n";
+        }
+
+        if (!found_any) {
+            ss << "No data";
+        }
+        return ss.str();
+    }
+};
+
 // SSD latency bucket: microseconds, tuned for SSD/network storage
 // Range: 50us (high-end NVMe) to 30s (3fs/nfs large object batch writes)
 inline const std::vector<double> kSsdLatencyBucket = {
@@ -424,7 +647,6 @@ struct SsdMetric {
         std::stringstream ss;
         ss << std::fixed << std::setprecision(1);
         ss << "count=" << count;
-        // rates order matches quantiles {0.5, 0.9, 0.99}
         if (rates.size() >= 1) ss << ", p50=" << rates[0] << "us";
         if (rates.size() >= 2) ss << ", p90=" << rates[1] << "us";
         if (rates.size() >= 3) ss << ", p99=" << rates[2] << "us";
@@ -438,6 +660,7 @@ struct SsdMetric {
 struct ClientMetric {
     TransferMetric transfer_metric;
     MasterClientMetric master_client_metric;
+    TransferOperationMetric transfer_operation_metric;
     SsdMetric ssd_metric;
 
     /**
@@ -452,25 +675,44 @@ struct ClientMetric {
      *   (default: 0, 0 = collect but don't report)
      */
     static std::unique_ptr<ClientMetric> Create(
-        const std::map<std::string, std::string>& labels = {});
+        const std::map<std::string, std::string>& labels = {},
+        bool master_rpc_metrics_enabled = true);
+
+    void ObserveTransferOperation(TransferOperationKind kind,
+                                  const std::string& op_name, uint64_t bytes,
+                                  uint64_t latency_us) {
+        transfer_operation_metric.Observe(kind, op_name, bytes, latency_us);
+    }
 
     void serialize(std::string& str);
     std::string summary_metrics();
 
     uint64_t GetReportingInterval() const { return metrics_interval_seconds_; }
 
-    explicit ClientMetric(
-        uint64_t interval_seconds = 0,
-        const std::map<std::string, std::string>& labels = {});
+    explicit ClientMetric(uint64_t interval_seconds = 0,
+                          const std::map<std::string, std::string>& labels = {},
+                          bool bandwidth_reporting_enabled = true,
+                          bool master_rpc_metrics_enabled = true);
     ~ClientMetric();
 
    private:
+    struct TransferSnapshot {
+        uint64_t read_bytes;
+        uint64_t write_bytes;
+        std::chrono::steady_clock::time_point timestamp;
+    };
+
     // Metrics reporting thread management
     std::jthread metrics_reporting_thread_;
     std::atomic<bool> should_stop_metrics_thread_{false};
     uint64_t metrics_interval_seconds_{0};
+    bool bandwidth_reporting_enabled_{true};
+    bool master_rpc_metrics_enabled_{true};
+    std::mutex snapshot_mutex_;
+    std::optional<TransferSnapshot> last_report_snapshot_;
 
     void StartMetricsReportingThread();
     void StopMetricsReportingThread();
+    std::string BuildBandwidthReport();
 };
-};  // namespace mooncake
+};  // namespace mooncake

mooncake-store/include/client_service.h

@@ -24,6 +24,7 @@
 #include "master_metric_manager.h"
 #include "count_min_sketch.h"
 #include "local_hot_cache.h"
+#include "pinned_buffer_pool.h"
 
 namespace mooncake {
 
@@ -61,6 +62,8 @@ class Client {
    public:
     ~Client();
 
+    const UUID& getClientId() const { return client_id_; }
+
     /**
      * @brief Creates and initializes a new Client instance
      * @param local_hostname Local host address (IP:Port)
@@ -444,6 +447,15 @@ class Client {
         return master_client_.CalcCacheStats();
     }
 
+    void ObserveTransferOperation(TransferOperationKind kind,
+                                  const std::string& op_name, uint64_t bytes,
+                                  uint64_t latency_us) {
+        if (metrics_ != nullptr) {
+            metrics_->ObserveTransferOperation(kind, op_name, bytes,
+                                               latency_us);
+        }
+    }
+
     // For Prometheus-style metrics
     tl::expected<std::string, ErrorCode> SerializeMetrics() {
         if (metrics_ == nullptr) {
@@ -660,6 +672,9 @@ class Client {
     const std::string protocol_;
 
     // Client persistent thread pool for async operations
+    // Pinned host memory pool for GPU D2H staging (must outlive
+    // write_thread_pool_)
+    std::unique_ptr<PinnedBufferPool> pinned_buffer_pool_;
     ThreadPool write_thread_pool_;
     std::shared_ptr<StorageBackend> storage_backend_;
 

mooncake-store/include/dummy_client.h

@@ -7,6 +7,7 @@
 #include "pyclient.h"
 #include "real_client.h"
 #include "shm_helper.h"
+#include "client_metric.h"
 #include <memory>
 
 namespace mooncake {
@@ -25,7 +26,8 @@ class DummyClient : public PyClient {
                    const std::string &master_server_addr,
                    const std::shared_ptr<TransferEngine> &transfer_engine,
                    const std::string &ipc_socket_path,
-                   bool enable_ssd_offload = false) {
+                   bool enable_ssd_offload = false,
+                   const std::string &ssd_offload_path = "") {
         // Dummy client does not support real setup
         return -1;
     };
@@ -149,6 +151,12 @@ class DummyClient : public PyClient {
     batch_get_replica_desc(const std::vector<std::string> &keys);
     std::vector<Replica::Descriptor> get_replica_desc(const std::string &key);
 
+    std::vector<std::string> batch_replica_clear(
+        const std::vector<std::string> &keys,
+        const std::string &segment_name = "") override {
+        return {};
+    }
+
     int tearDownAll();
 
     int health_check() override;
@@ -196,6 +204,22 @@ class DummyClient : public PyClient {
     [[nodiscard]] std::vector<tl::expected<ResultType, ErrorCode>>
     invoke_batch_rpc(size_t input_size, Args &&...args);
 
+    template <auto ServiceMethod, typename... Args>
+    int invoke_observed_void_rpc(TransferOperationKind kind,
+                                 const char *op_name, size_t bytes, bool batch,
+                                 Args &&...args) {
+        auto result = execute_timed_operation<tl::expected<void, ErrorCode>>(
+            [&]() {
+                return invoke_rpc<ServiceMethod, void>(
+                    std::forward<Args>(args)...);
+            },
+            [](const auto &ret) { return ret.has_value(); },
+            [&](uint64_t latency_us, const auto &) {
+                ObserveTransferMetric(kind, op_name, bytes, latency_us, batch);
+            });
+        return to_py_ret(result);
+    }
+
     /**
      * @brief Accessor for the coro_rpc_client pool. Since coro_rpc_client
      * pool cannot reconnect to a different address, a new coro_rpc_client
@@ -254,6 +278,11 @@ class DummyClient : public PyClient {
 
     // Ascend physical device id for dummy-real RPC to real, set in setup_dummy
     int32_t device_id_ = 0;
+
+    std::unique_ptr<ClientMetric> metrics_;
+
+    void ObserveTransferMetric(TransferOperationKind kind, const char *op_name,
+                               size_t bytes, uint64_t latency_us, bool batch);
 };
 
 }  // namespace mooncake

mooncake-store/include/file_storage.h

@@ -3,6 +3,7 @@
 #include "client_service.h"
 #include "client_buffer.hpp"
 #include "storage_backend.h"
+#include "pinned_buffer_pool.h"
 
 namespace mooncake {
 
@@ -105,6 +106,8 @@ class FileStorage {
     std::shared_ptr<Client> client_;
     SsdMetric* ssd_metric_{nullptr};
     std::string local_rpc_addr_;
+    // Pinned host memory pool for GPU D2H staging in OffloadObjects
+    std::unique_ptr<PinnedBufferPool> pinned_buffer_pool_;
     std::shared_ptr<StorageBackendInterface> storage_backend_;
     std::shared_ptr<ClientBufferAllocator> client_buffer_allocator_;
     mutable Mutex client_buffer_mutex_;

mooncake-store/include/gpu_staging_utils.h

@@ -0,0 +1,76 @@
+#pragma once
+
+#include "cuda_alike.h"
+
+#if defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM)
+#include <acl/acl_rt.h>
+#endif
+
+#include <cstddef>
+#include <glog/logging.h>
+
+namespace mooncake {
+namespace gpu_staging {
+
+// Detect whether ptr resides in accelerator device memory.
+// If so, writes the device ID to *out_device_id for subsequent SetDevice.
+inline bool IsDevicePointer(const void* ptr, int* out_device_id) {
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)
+    cudaPointerAttributes attr{};
+    if (cudaPointerGetAttributes(&attr, ptr) == cudaSuccess &&
+        attr.type == cudaMemoryTypeDevice) {
+        if (out_device_id) *out_device_id = attr.device;
+        return true;
+    }
+#elif defined(USE_HIP)
+    hipPointerAttribute_t attr{};
+    if (hipPointerGetAttributes(&attr, ptr) == hipSuccess &&
+        attr.type == hipMemoryTypeDevice) {
+        if (out_device_id) *out_device_id = attr.device;
+        return true;
+    }
+#elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM)
+    aclrtPtrAttributes attr{};
+    if (aclrtPointerGetAttributes(const_cast<void*>(ptr), &attr) ==
+            ACL_SUCCESS &&
+        attr.location.type == ACL_MEM_LOCATION_TYPE_DEVICE) {
+        if (out_device_id) *out_device_id = static_cast<int>(attr.location.id);
+        return true;
+    }
+#endif
+    (void)ptr;
+    (void)out_device_id;
+    return false;
+}
+
+// Copy device memory to host. Caller must have called SetDevice first.
+inline bool CopyDeviceToHost(void* dst, const void* src, size_t size) {
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)
+    return cudaMemcpy(dst, src, size, cudaMemcpyDeviceToHost) == cudaSuccess;
+#elif defined(USE_HIP)
+    return hipMemcpy(dst, src, size, hipMemcpyDeviceToHost) == hipSuccess;
+#elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM)
+    return aclrtMemcpy(dst, size, src, size, ACL_MEMCPY_DEVICE_TO_HOST) ==
+           ACL_SUCCESS;
+#else
+    (void)dst;
+    (void)src;
+    (void)size;
+    return false;
+#endif
+}
+
+// Bind the calling thread to the given device context.
+inline void SetDevice(int device_id) {
+    if (device_id < 0) return;
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)
+    cudaSetDevice(device_id);
+#elif defined(USE_HIP)
+    hipSetDevice(device_id);
+#elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM)
+    aclrtSetDevice(device_id);
+#endif
+}
+
+}  // namespace gpu_staging
+}  // namespace mooncake

mooncake-store/include/master_config.h

@@ -81,6 +81,10 @@ struct MasterConfig {
     std::string cxl_path;
     size_t cxl_size;
     bool enable_cxl = false;
+
+    // Offload-on-evict: defer LOCAL_DISK offload to eviction time
+    bool offload_on_evict = false;
+    bool offload_force_evict = false;
 };
 
 class MasterServiceSupervisorConfig {
@@ -140,6 +144,8 @@ class MasterServiceSupervisorConfig {
     std::string cxl_path = DEFAULT_CXL_PATH;
     size_t cxl_size = DEFAULT_CXL_SIZE;
     bool enable_cxl = false;
+    bool offload_on_evict = false;
+    bool offload_force_evict = false;
     MasterServiceSupervisorConfig() = default;
 
     // From MasterConfig
@@ -155,6 +161,8 @@ class MasterServiceSupervisorConfig {
         eviction_high_watermark_ratio = config.eviction_high_watermark_ratio;
         client_live_ttl_sec = config.client_live_ttl_sec;
         enable_offload = config.enable_offload;
+        offload_on_evict = config.offload_on_evict;
+        offload_force_evict = config.offload_force_evict;
         rpc_port = static_cast<int>(config.rpc_port);
         rpc_thread_num = static_cast<size_t>(config.rpc_thread_num);
 
@@ -267,6 +275,8 @@ class WrappedMasterServiceConfig {
     int64_t client_live_ttl_sec = DEFAULT_CLIENT_LIVE_TTL_SEC;
     bool enable_ha = false;
     bool enable_offload = false;
+    bool offload_on_evict = false;
+    bool offload_force_evict = false;
     std::string ha_backend_type = "etcd";
     std::string ha_backend_connstring;
     std::string cluster_id = DEFAULT_CLUSTER_ID;
@@ -322,6 +332,8 @@ class WrappedMasterServiceConfig {
         client_live_ttl_sec = config.client_live_ttl_sec;
         enable_ha = config.enable_ha;
         enable_offload = config.enable_offload;
+        offload_on_evict = config.offload_on_evict;
+        offload_force_evict = config.offload_force_evict;
         ha_backend_type = config.ha_backend_type;
         ha_backend_connstring = config.ha_backend_connstring;
         if (ha_backend_connstring.empty()) {
@@ -400,6 +412,8 @@ class WrappedMasterServiceConfig {
         enable_ha =
             true;  // This is used in HA mode, so enable_ha should be true
         enable_offload = config.enable_offload;
+        offload_on_evict = config.offload_on_evict;
+        offload_force_evict = config.offload_force_evict;
         ha_backend_type = config.ha_backend_type;
         ha_backend_connstring = config.ha_backend_connstring;
         if (ha_backend_connstring.empty()) {
@@ -736,6 +750,8 @@ class MasterServiceConfig {
     int64_t client_live_ttl_sec = DEFAULT_CLIENT_LIVE_TTL_SEC;
     bool enable_ha = false;
     bool enable_offload = false;
+    bool offload_on_evict = false;
+    bool offload_force_evict = false;
     std::string ha_backend_type = "etcd";
     std::string ha_backend_connstring;
     std::string cluster_id = DEFAULT_CLUSTER_ID;
@@ -787,6 +803,8 @@ class MasterServiceConfig {
         client_live_ttl_sec = config.client_live_ttl_sec;
         enable_ha = config.enable_ha;
         enable_offload = config.enable_offload;
+        offload_on_evict = config.offload_on_evict;
+        offload_force_evict = config.offload_force_evict;
         ha_backend_type = config.ha_backend_type;
         ha_backend_connstring = config.ha_backend_connstring;
         cluster_id = config.cluster_id;

mooncake-store/include/master_service.h

@@ -562,8 +562,14 @@ class MasterService {
     // fulfill evict ratio lowerbound.
     void BatchEvict(double evict_ratio_target, double evict_ratio_lowerbound);
 
+    // Helper to get a snapshot of alive clients (under client_mutex_ shared
+    // lock)
+    std::unordered_set<UUID, boost::hash<UUID>> getAliveClientsSnapshot() const;
+
     // Clear invalid handles in all shards
     void ClearInvalidHandles();
+    void ClearInvalidHandles(
+        const std::unordered_set<UUID, boost::hash<UUID>>& alive_clients);
 
     std::string FormatTimestamp(
         const std::chrono::system_clock::time_point& tp);
@@ -891,7 +897,10 @@ class MasterService {
     }
 
     // Helper to clean up stale handles pointing to unmounted segments
-    bool CleanupStaleHandles(ObjectMetadata& metadata);
+    // or local_disk replicas whose owner client is no longer alive.
+    bool CleanupStaleHandles(
+        ObjectMetadata& metadata,
+        const std::unordered_set<UUID, boost::hash<UUID>>& alive_clients);
 
     // Helper: allocate replicas, create ObjectMetadata, insert into shard,
     // and return descriptor list.  Shared by PutStart and UpsertStart.
@@ -962,11 +971,21 @@ class MasterService {
               it_(shard_guard_->metadata.find(key)),
               processing_it_(shard_guard_->processing_keys.find(key)),
               replication_task_it_(shard_guard_->replication_tasks.find(key)) {
-            // Automatically clean up invalid handles
+            // Automatically clean up invalid handles (memory replicas only).
+            // Note: We only check memory replicas here to avoid lock order
+            // violation (client_mutex_ must be acquired before metadata shard).
+            // local_disk replicas are cleaned up by ClearInvalidHandles() in
+            // ClientMonitorFunc.
             if (it_ != shard_guard_->metadata.end()) {
-                if (service_->CleanupStaleHandles(it_->second)) {
+                // Erase invalid memory replicas (those with unmounted
+                // segments). No client_mutex_ needed since we only check memory
+                // replicas.
+                it_->second.EraseReplicas([](const Replica& replica) {
+                    return replica.has_invalid_mem_handle();
+                });
+                // If no valid replicas remain, delete the whole object.
+                if (!it_->second.IsValid()) {
                     this->Erase();
-
                     if (processing_it_ != shard_guard_->processing_keys.end()) {
                         this->EraseFromProcessing();
                     }
@@ -1149,6 +1168,14 @@ class MasterService {
 
     const bool enable_offload_;
 
+    // Offload-on-evict: defer disk offload to eviction time
+    // (config: offload_on_evict)
+    bool offload_on_evict_{false};
+    // Force-evict: allow evicting MEMORY replicas without disk offload when cap
+    // exceeded (config: offload_force_evict, only effective when
+    // offload_on_evict_=true)
+    bool offload_force_evict_{false};
+
     const std::string ha_backend_type_;
 
     const std::string ha_backend_connstring_;

mooncake-store/include/pinned_buffer_pool.h

@@ -0,0 +1,138 @@
+#pragma once
+
+#include <mutex>
+#include <vector>
+#include <cstdlib>
+#include "cuda_alike.h"
+
+// Ascend CANN is not covered by cuda_alike.h
+#if defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM)
+#include <acl/acl_rt.h>
+#endif
+
+namespace mooncake {
+
+/**
+ * PinnedBufferPool: Thread-safe pool of reusable pinned host memory buffers.
+ *
+ * Platform pinned alloc APIs:
+ *   CUDA / MUSA / MACA : cudaMallocHost   (mapped via cuda_alike.h)
+ *   HIP                : hipHostMalloc     (not mapped in hip.h, native API)
+ *   Ascend             : aclrtMallocHost
+ *   Other              : new char[]        (pageable fallback)
+ *
+ * Pinned memory provides 10x~100x higher D2H bandwidth than pageable memory.
+ * Falls back to new char[] if pinned allocation fails.
+ *
+ * The pool enforces a maximum number of cached buffers (kDefaultMaxPoolSize).
+ * When the pool is full, Release() frees the buffer immediately instead of
+ * caching it, preventing unbounded pinned memory growth.
+ */
+class PinnedBufferPool {
+   public:
+    static constexpr size_t kDefaultMaxPoolSize = 32;
+
+    struct Buffer {
+        char* data = nullptr;
+        size_t capacity = 0;
+        bool is_pinned = false;  // Selects correct free API in FreeBuffer
+    };
+
+    explicit PinnedBufferPool(size_t max_pool_size = kDefaultMaxPoolSize)
+        : max_pool_size_(max_pool_size) {}
+
+    ~PinnedBufferPool() { Clear(); }
+
+    Buffer Acquire(size_t size) {
+        {
+            std::lock_guard<std::mutex> lk(mutex_);
+            for (size_t i = 0; i < pool_.size(); ++i) {
+                if (pool_[i].capacity >= size) {
+                    Buffer buf = pool_[i];
+                    // O(1) erase: swap with back then pop
+                    pool_[i] = pool_.back();
+                    pool_.pop_back();
+                    return buf;
+                }
+            }
+        }
+        return AllocNew(size);
+    }
+
+    void Release(Buffer buf) {
+        std::lock_guard<std::mutex> lk(mutex_);
+        if (pool_.size() < max_pool_size_) {
+            pool_.push_back(buf);
+        } else {
+            // Pool full — free immediately to bound pinned memory usage
+            FreeBuffer(buf);
+        }
+    }
+
+    void Clear() {
+        std::lock_guard<std::mutex> lk(mutex_);
+        for (auto& buf : pool_) {
+            FreeBuffer(buf);
+        }
+        pool_.clear();
+    }
+
+   private:
+    static Buffer AllocNew(size_t size) {
+        Buffer buf;
+        buf.capacity = size;
+        buf.is_pinned = false;
+
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)
+        if (cudaMallocHost(reinterpret_cast<void**>(&buf.data), size) ==
+            cudaSuccess) {
+            buf.is_pinned = true;
+        } else {
+            buf.data = new char[size];
+        }
+
+#elif defined(USE_HIP)
+        if (hipHostMalloc(reinterpret_cast<void**>(&buf.data), size, 0) ==
+            hipSuccess) {
+            buf.is_pinned = true;
+        } else {
+            buf.data = new char[size];
+        }
+
+#elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM)
+        if (aclrtMallocHost(reinterpret_cast<void**>(&buf.data), size) ==
+            ACL_SUCCESS) {
+            buf.is_pinned = true;
+        } else {
+            buf.data = new char[size];
+        }
+
+#else
+        buf.data = new char[size];
+#endif
+        return buf;
+    }
+
+    static void FreeBuffer(Buffer& buf) {
+        if (!buf.data) return;
+        if (!buf.is_pinned) {
+            delete[] buf.data;
+            return;
+        }
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)
+        cudaFreeHost(buf.data);
+#elif defined(USE_HIP)
+        hipHostFree(buf.data);
+#elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM)
+        aclrtFreeHost(buf.data);
+#else
+        delete[] buf.data;
+#endif
+    }
+
+    const size_t max_pool_size_;
+    std::mutex mutex_;
+    std::vector<Buffer> pool_;
+};
+
+}  // namespace mooncake

mooncake-store/include/pyclient.h

@@ -211,8 +211,8 @@ class PyClient {
         const std::string &protocol, const std::string &rdma_devices,
         const std::string &master_server_addr,
         const std::shared_ptr<TransferEngine> &transfer_engine,
-        const std::string &ipc_socket_path,
-        bool enable_ssd_offload = false) = 0;
+        const std::string &ipc_socket_path, bool enable_ssd_offload = false,
+        const std::string &ssd_offload_path = "") = 0;
 
     virtual int setup_dummy(size_t mem_pool_size, size_t local_buffer_size,
                             const std::string &server_address,
@@ -330,6 +330,10 @@ class PyClient {
     virtual std::vector<Replica::Descriptor> get_replica_desc(
         const std::string &key) = 0;
 
+    virtual std::vector<std::string> batch_replica_clear(
+        const std::vector<std::string> &keys,
+        const std::string &segment_name = "") = 0;
+
     virtual int tearDownAll() = 0;
 
     virtual int health_check() = 0;

mooncake-store/include/real_client.h

@@ -79,7 +79,8 @@ class RealClient : public PyClient {
         const std::string &master_server_addr = "127.0.0.1:50051",
         const std::shared_ptr<TransferEngine> &transfer_engine = nullptr,
         const std::string &ipc_socket_path = "",
-        bool enable_ssd_offload = false);
+        bool enable_ssd_offload = false,
+        const std::string &ssd_offload_path = "");
 
     int setup_dummy(size_t mem_pool_size, size_t local_buffer_size,
                     const std::string &server_address,
@@ -483,7 +484,8 @@ class RealClient : public PyClient {
         const std::string &master_server_addr = "127.0.0.1:50051",
         const std::shared_ptr<TransferEngine> &transfer_engine = nullptr,
         const std::string &ipc_socket_path = "", int local_rpc_port = 50052,
-        bool enable_ssd_offload = false, bool start_offload_rpc_server = false);
+        bool enable_ssd_offload = false, bool start_offload_rpc_server = false,
+        const std::string &ssd_offload_path = "");
 
     // Overload that accepts a configuration dictionary
     tl::expected<void, ErrorCode> setup_internal(const ConfigDict &config);
@@ -637,6 +639,10 @@ class RealClient : public PyClient {
     batch_get_replica_desc(const std::vector<std::string> &keys);
     std::vector<Replica::Descriptor> get_replica_desc(const std::string &key);
 
+    std::vector<std::string> batch_replica_clear(
+        const std::vector<std::string> &keys,
+        const std::string &segment_name = "") override;
+
     tl::expected<PingResponse, ErrorCode> ping(const UUID &client_id);
 
     tl::expected<BatchGetOffloadObjectResponse, ErrorCode>

mooncake-store/include/replica.h

@@ -2,9 +2,12 @@
 
 #include <glog/logging.h>
 
+#include <boost/functional/hash.hpp>
+
 #include <cstdint>
 #include <memory>
 #include <string>
+#include <unordered_set>
 #include <variant>
 #include <vector>
 #include <unordered_map>
@@ -282,6 +285,36 @@ class Replica {
         return false;  // DiskReplicaData does not have handles
     }
 
+    /**
+     * @brief Check if a local_disk replica's owner client is still alive.
+     * Used by CleanupStaleHandles to remove replicas belonging to expired
+     * clients. For non-local_disk replicas, always returns false.
+     * @param alive_clients Set of currently alive client IDs.
+     * @return true if this is a local_disk replica whose client is not alive.
+     */
+    [[nodiscard]] bool has_stale_local_disk_client(
+        const std::unordered_set<UUID, boost::hash<UUID>>& alive_clients)
+        const {
+        auto client_id = get_local_disk_client_id();
+        if (client_id.has_value()) {
+            return alive_clients.find(client_id.value()) == alive_clients.end();
+        }
+        return false;
+    }
+
+    /**
+     * @brief Get the client_id for local_disk replicas.
+     * @return The client_id if this is a local_disk replica, std::nullopt
+     * otherwise.
+     */
+    [[nodiscard]] std::optional<UUID> get_local_disk_client_id() const {
+        if (is_local_disk_replica()) {
+            const auto& disk_data = std::get<LocalDiskReplicaData>(data_);
+            return disk_data.client_id;
+        }
+        return std::nullopt;
+    }
+
     [[nodiscard]] size_t get_memory_buffer_size() const {
         if (is_memory_replica()) {
             const auto& mem_data = std::get<MemoryReplicaData>(data_);

mooncake-store/include/segment.h

@@ -170,6 +170,12 @@ class ScopedSegmentAccess {
     ErrorCode SetSegmentStatusByName(const std::string& segment_name,
                                      SegmentStatus status);
 
+    /**
+     * @brief Remove the local disk segment entry for a client.
+     * Called when a client expires to clean up its local disk segment.
+     */
+    void UnmountLocalDiskSegment(const UUID& client_id);
+
    private:
     SegmentManager* segment_manager_;
     std::unique_lock<std::shared_mutex> lock_;

mooncake-store/include/transfer_task.h

@@ -174,7 +174,10 @@ class TransferEngineOperationState : public OperationState {
    public:
     TransferEngineOperationState(TransferEngine& engine, BatchID batch_id,
                                  size_t batch_size)
-        : engine_(engine), batch_id_(batch_id), batch_size_(batch_size) {}
+        : engine_(engine),
+          batch_id_(batch_id),
+          batch_size_(batch_size),
+          start_ts_(getCurrentTimeInMilli()) {}
 
     ~TransferEngineOperationState() { engine_.freeBatchID(batch_id_); }
 
@@ -199,6 +202,7 @@ class TransferEngineOperationState : public OperationState {
     TransferEngine& engine_;
     BatchID batch_id_;
     size_t batch_size_;
+    const int64_t start_ts_;
 };
 
 /**

mooncake-store/src/CMakeLists.txt

@@ -202,15 +202,62 @@ if(STORE_USE_ETCD)
   add_dependencies(mooncake_master build_etcd_wrapper)
 endif()
 
-target_compile_options(mooncake_master PRIVATE -Os)
-target_link_options(mooncake_master PRIVATE -Os -s)
-
 # Client server binary
 add_executable(mooncake_client real_client_main.cpp)
 # Client needs transfer_engine for data transfer operations
 target_link_libraries(mooncake_client PRIVATE mooncake_store transfer_engine
                                               asio_shared)
-target_compile_options(mooncake_client PRIVATE -Os)
-target_link_options(mooncake_client PRIVATE -Os -s)
 
+# Optimize binary sizes only in Release mode
+string(TOUPPER "${CMAKE_BUILD_TYPE}" CMAKE_BUILD_TYPE_UPPER)
+if (CMAKE_BUILD_TYPE_UPPER STREQUAL "RELEASE")
+  target_compile_options(mooncake_master PRIVATE -Os)
+  target_link_options(mooncake_master PRIVATE -Os -s)
+  target_compile_options(mooncake_client PRIVATE -Os)
+  target_link_options(mooncake_client PRIVATE -Os -s)
+endif()
+
+
+# GPU runtime library for D2H staging in PutToLocalFile / OffloadObjects.
+# transfer_engine is PRIVATE-linked, so its CUDA/HIP/Ascend dependencies
+# are not propagated; we must detect and link them independently.
+#
+# Auto-detect each toolkit regardless of global USE_CUDA/USE_HIP flags,
+# because USE_CUDA may be OFF even when GPU pointers are present
+# (e.g. WITH_NVIDIA_PEERMEM=ON uses nvidia-peermem for RDMA without cudart).
+# Each detected toolkit gets both link libraries AND compile definitions,
+# so that gpu_staging_utils.h / pinned_buffer_pool.h pick the correct backend.
+#
+# NOTE: mooncake_store is a static library (.a). External consumers (Go CGo,
+# Python pybind) that link it must also link the GPU runtime (e.g. -lcudart).
+# Go's build.sh already does this; CI workflows must do the same.
+
+find_package(CUDAToolkit QUIET)
+if(CUDAToolkit_FOUND)
+    message(STATUS "mooncake_store: CUDAToolkit detected, enabling D2H staging")
+    target_compile_definitions(mooncake_store PRIVATE USE_CUDA)
+    target_compile_definitions(mooncake_client PRIVATE USE_CUDA)
+    target_include_directories(mooncake_store PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
+    target_include_directories(mooncake_client PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
+    target_link_libraries(mooncake_store PRIVATE CUDA::cudart)
+    target_link_libraries(mooncake_client PRIVATE CUDA::cudart)
+endif()
+
+if(NOT CUDAToolkit_FOUND)
+    find_package(hip QUIET)
+    if(hip_FOUND)
+        message(STATUS "mooncake_store: HIP detected, enabling D2H staging")
+        target_compile_definitions(mooncake_store PRIVATE USE_HIP)
+        target_compile_definitions(mooncake_client PRIVATE USE_HIP)
+        target_link_libraries(mooncake_store PRIVATE hip::host)
+        target_link_libraries(mooncake_client PRIVATE hip::host)
+    endif()
+endif()
+
+if(USE_ASCEND OR USE_ASCEND_DIRECT OR USE_UBSHMEM)
+    target_include_directories(mooncake_store PRIVATE $ENV{ASCEND_PATH}/include)
+    target_link_libraries(mooncake_store PRIVATE ascendcl)
+    target_include_directories(mooncake_client PRIVATE $ENV{ASCEND_PATH}/include)
+    target_link_libraries(mooncake_client PRIVATE ascendcl)
+endif()
 install(TARGETS mooncake_master mooncake_client DESTINATION bin)

mooncake-store/src/client_metric.cpp

@@ -28,6 +28,27 @@ bool parseMetricsEnabled() {
             value == "on" || value == "enable");
 }
 
+bool parseBoolEnv(const char* env_name, bool default_value) {
+    const char* env_value = std::getenv(env_name);
+    if (!env_value) {
+        return default_value;
+    }
+
+    std::string value = toLower(env_value);
+    if (value == "1" || value == "true" || value == "yes" || value == "on" ||
+        value == "enable") {
+        return true;
+    }
+    if (value == "0" || value == "false" || value == "no" || value == "off" ||
+        value == "disable") {
+        return false;
+    }
+
+    LOG(WARNING) << "Failed to parse " << env_name << ": " << env_value
+                 << ", fallback to default=" << default_value;
+    return default_value;
+}
+
 uint64_t parseMetricsInterval() {
     const char* interval_env = std::getenv("MC_STORE_CLIENT_METRIC_INTERVAL");
     if (!interval_env) {
@@ -55,12 +76,21 @@ uint64_t parseMetricsInterval() {
 }  // anonymous namespace
 
 ClientMetric::ClientMetric(uint64_t interval_seconds,
-                           const std::map<std::string, std::string>& labels)
+                           const std::map<std::string, std::string>& labels,
+                           bool bandwidth_reporting_enabled,
+                           bool master_rpc_metrics_enabled)
     : transfer_metric(labels),
       master_client_metric(labels),
+      transfer_operation_metric(labels),
       ssd_metric(labels),
       should_stop_metrics_thread_(false),
-      metrics_interval_seconds_(interval_seconds) {
+      metrics_interval_seconds_(interval_seconds),
+      bandwidth_reporting_enabled_(bandwidth_reporting_enabled),
+      master_rpc_metrics_enabled_(master_rpc_metrics_enabled) {
+    last_report_snapshot_ = TransferSnapshot{
+        static_cast<uint64_t>(transfer_metric.total_read_bytes.value()),
+        static_cast<uint64_t>(transfer_metric.total_write_bytes.value()),
+        std::chrono::steady_clock::now()};
     if (metrics_interval_seconds_ > 0) {
         StartMetricsReportingThread();
     }
@@ -69,7 +99,8 @@ ClientMetric::ClientMetric(uint64_t interval_seconds,
 ClientMetric::~ClientMetric() { StopMetricsReportingThread(); }
 
 std::unique_ptr<ClientMetric> ClientMetric::Create(
-    const std::map<std::string, std::string>& labels) {
+    const std::map<std::string, std::string>& labels,
+    bool master_rpc_metrics_enabled) {
     if (!parseMetricsEnabled()) {
         LOG(INFO) << "Client metrics disabled (set MC_STORE_CLIENT_METRIC=0 to "
                      "disable)";
@@ -77,29 +108,83 @@ std::unique_ptr<ClientMetric> ClientMetric::Create(
     }
 
     uint64_t interval = parseMetricsInterval();
+    bool bandwidth_reporting_enabled =
+        parseBoolEnv("MC_STORE_CLIENT_METRIC_BANDWIDTH", true);
 
     LOG(INFO) << "Client metrics enabled (default enabled)";
+    LOG(INFO) << "Client bandwidth summary "
+              << (bandwidth_reporting_enabled ? "enabled" : "disabled")
+              << " via MC_STORE_CLIENT_METRIC_BANDWIDTH";
 
-    return std::make_unique<ClientMetric>(interval, labels);
+    return std::make_unique<ClientMetric>(interval, labels,
+                                          bandwidth_reporting_enabled,
+                                          master_rpc_metrics_enabled);
 }
 
 void ClientMetric::serialize(std::string& str) {
     transfer_metric.serialize(str);
-    master_client_metric.serialize(str);
+    if (master_rpc_metrics_enabled_) {
+        master_client_metric.serialize(str);
+    }
+    transfer_operation_metric.serialize(str);
     ssd_metric.serialize(str);
 }
 
 std::string ClientMetric::summary_metrics() {
     std::stringstream ss;
     ss << "Client Metrics Summary\n";
-    ss << transfer_metric.summary_metrics();
+    ss << transfer_metric.summary_metrics(bandwidth_reporting_enabled_);
     ss << "\n";
-    ss << master_client_metric.summary_metrics();
+    if (master_rpc_metrics_enabled_) {
+        ss << master_client_metric.summary_metrics();
+        ss << "\n";
+    }
+    ss << transfer_operation_metric.summary_metrics();
     ss << "\n";
     ss << ssd_metric.summary_metrics();
     return ss.str();
 }
 
+std::string ClientMetric::BuildBandwidthReport() {
+    if (!bandwidth_reporting_enabled_) {
+        return "";
+    }
+
+    const auto now = std::chrono::steady_clock::now();
+    const uint64_t read_bytes = transfer_metric.total_read_bytes.value();
+    const uint64_t write_bytes = transfer_metric.total_write_bytes.value();
+
+    std::lock_guard<std::mutex> lock(snapshot_mutex_);
+    if (!last_report_snapshot_.has_value()) {
+        last_report_snapshot_ = TransferSnapshot{read_bytes, write_bytes, now};
+        return "";
+    }
+
+    const auto previous = *last_report_snapshot_;
+    last_report_snapshot_ = TransferSnapshot{read_bytes, write_bytes, now};
+
+    const double elapsed_seconds = std::max(
+        std::chrono::duration<double>(now - previous.timestamp).count(), 1e-9);
+    const uint64_t read_delta = read_bytes >= previous.read_bytes
+                                    ? read_bytes - previous.read_bytes
+                                    : 0;
+    const uint64_t write_delta = write_bytes >= previous.write_bytes
+                                     ? write_bytes - previous.write_bytes
+                                     : 0;
+
+    std::stringstream ss;
+    ss << "=== Interval Throughput Summary ===\n";
+    ss << "Read Throughput: "
+       << format_metric_rate(read_delta / elapsed_seconds, "B/s") << " ("
+       << byte_size_to_string(read_delta) << " over " << std::fixed
+       << std::setprecision(2) << elapsed_seconds << "s)\n";
+    ss << "Write Throughput: "
+       << format_metric_rate(write_delta / elapsed_seconds, "B/s") << " ("
+       << byte_size_to_string(write_delta) << " over " << std::fixed
+       << std::setprecision(2) << elapsed_seconds << "s)";
+    return ss.str();
+}
+
 void ClientMetric::StartMetricsReportingThread() {
     should_stop_metrics_thread_ = false;
     metrics_reporting_thread_ =
@@ -124,7 +209,12 @@ void ClientMetric::StartMetricsReportingThread() {
 
                 // Print metrics summary
                 std::string summary = summary_metrics();
-                LOG(INFO) << "Client Metrics Report:\n" << summary;
+                std::string bandwidth_report = BuildBandwidthReport();
+                std::string report = "Client Metrics Report:\n" + summary;
+                if (!bandwidth_report.empty()) {
+                    report += "\n" + bandwidth_report;
+                }
+                LOG(INFO) << report;
             }
             LOG(INFO) << "Client metrics reporting thread stopped";
         });
@@ -140,4 +230,4 @@ void ClientMetric::StopMetricsReportingThread() {
     }
 }
 
-}  // namespace mooncake
+}  // namespace mooncake

mooncake-store/src/client_service.cpp

@@ -27,9 +27,14 @@
 #include "utils.h"
 #include "rpc_types.h"
 #include "local_hot_cache.h"
+#include "gpu_staging_utils.h"
 
 namespace mooncake {
 
+using gpu_staging::CopyDeviceToHost;
+using gpu_staging::IsDevicePointer;
+using gpu_staging::SetDevice;
+
 [[nodiscard]] size_t CalculateSliceSize(const std::vector<Slice>& slices) {
     size_t slice_size = 0;
     for (const auto& slice : slices) {
@@ -57,6 +62,7 @@ Client::Client(const std::string& local_hostname,
       local_hostname_(local_hostname),
       metadata_connstring_(metadata_connstring),
       protocol_(protocol),
+      pinned_buffer_pool_(std::make_unique<PinnedBufferPool>()),
       write_thread_pool_(2),
       task_thread_pool_(4) {
     LOG(INFO) << "client_id=" << client_id_;
@@ -2495,19 +2501,34 @@ void Client::PutToLocalFile(const std::string& key,
     }
 
     std::string path = disk_descriptor.file_path;
-    // Currently, persistence is achieved through asynchronous writes, but
-    // before asynchronous writing in 3FS, significant performance degradation
-    // may occur due to data copying. Profiling reveals that the number of page
-    // faults triggered in this scenario is nearly double the normal count.
-    // Future plans include introducing a reuse buffer list to address this
-    // performance degradation issue.
 
+    // Synchronous D2H staging + copy into std::string.
+    // Done on the calling thread to guarantee GPU buffers are still valid
+    // (BatchPut has not yet returned to Python, so blocks are not reused).
     std::string value;
     value.reserve(total_size);
+
     for (const auto& slice : slices) {
-        value.append(static_cast<char*>(slice.ptr), slice.size);
+        int device_id = -1;
+        if (IsDevicePointer(slice.ptr, &device_id)) {
+            SetDevice(device_id);
+            auto buf = pinned_buffer_pool_->Acquire(slice.size);
+            if (!CopyDeviceToHost(buf.data, slice.ptr, slice.size)) {
+                LOG(ERROR) << "D2H copy failed for key: " << key
+                           << ", triggering PutRevoke for disk replica";
+                pinned_buffer_pool_->Release(buf);
+                // Must revoke to avoid phantom replica in master
+                master_client_.PutRevoke(key, ReplicaType::DISK);
+                return;
+            }
+            value.append(buf.data, slice.size);
+            pinned_buffer_pool_->Release(buf);
+        } else {
+            value.append(static_cast<char*>(slice.ptr), slice.size);
+        }
     }
 
+    // Async StoreObject + PutEnd (unchanged from original)
     write_thread_pool_.enqueue([this, backend = storage_backend_, key,
                                 value = std::move(value), path] {
         // Store the object

mooncake-store/src/dummy_client.cpp

@@ -7,6 +7,7 @@
 #include <sys/stat.h>  // For S_IRUSR, S_IWUSR
 #include <fcntl.h>     // For O_CREAT, O_RDWR
 #include <unistd.h>    // For ftruncate, close, shm_unlink
+#include <chrono>
 #include <cstdlib>
 
 #include "real_client.h"
@@ -23,6 +24,75 @@
 #endif
 
 namespace {
+size_t sum_value_sizes(const std::vector<std::span<const char>>& values) {
+    size_t total = 0;
+    for (const auto& value : values) {
+        total += value.size_bytes();
+    }
+    return total;
+}
+
+size_t sum_sizes(const std::vector<size_t>& sizes) {
+    size_t total = 0;
+    for (size_t size : sizes) {
+        total += size;
+    }
+    return total;
+}
+
+size_t sum_successful_sizes(const std::vector<int>& results,
+                            const std::vector<size_t>& sizes) {
+    size_t total = 0;
+    for (size_t i = 0; i < results.size() && i < sizes.size(); ++i) {
+        if (results[i] == 0) {
+            total += sizes[i];
+        }
+    }
+    return total;
+}
+
+size_t sum_successful_nested_sizes(
+    const std::vector<int>& results,
+    const std::vector<std::vector<size_t>>& nested_sizes) {
+    size_t total = 0;
+    for (size_t i = 0; i < results.size() && i < nested_sizes.size(); ++i) {
+        if (results[i] == 0) {
+            total += sum_sizes(nested_sizes[i]);
+        }
+    }
+    return total;
+}
+
+size_t sum_positive_results(const std::vector<int64_t>& results) {
+    size_t total = 0;
+    for (int64_t result : results) {
+        if (result > 0) {
+            total += static_cast<size_t>(result);
+        }
+    }
+    return total;
+}
+
+size_t sum_positive_results(const std::vector<int>& results) {
+    size_t total = 0;
+    for (int result : results) {
+        if (result > 0) {
+            total += static_cast<size_t>(result);
+        }
+    }
+    return total;
+}
+
+size_t sum_positive_ranges(
+    const std::vector<std::vector<std::vector<int64_t>>>& results) {
+    size_t total = 0;
+    for (const auto& key_rows : results) {
+        for (const auto& row : key_rows) {
+            total += sum_positive_results(row);
+        }
+    }
+    return total;
+}
 
 std::vector<uint64_t> void_ptrs_to_u64(const std::vector<void*>& ptrs) {
     std::vector<uint64_t> out;
@@ -129,7 +199,10 @@ std::vector<tl::expected<ResultType, ErrorCode>> DummyClient::invoke_batch_rpc(
         }());
 }
 
-DummyClient::DummyClient() : client_id_(generate_uuid()) {
+DummyClient::DummyClient()
+    : client_id_(generate_uuid()),
+      metrics_(ClientMetric::Create(merge_labels({{"client_mode", "dummy"}}),
+                                    false)) {
     // Initialize logging severity (leave as before)
     mooncake::init_ylt_log_level();
     // Initialize client pools
@@ -141,6 +214,30 @@ DummyClient::DummyClient() : client_id_(generate_uuid()) {
 
 DummyClient::~DummyClient() { tearDownAll(); }
 
+void DummyClient::ObserveTransferMetric(TransferOperationKind kind,
+                                        const char* op_name, size_t bytes,
+                                        uint64_t latency_us, bool batch) {
+    if (!metrics_) {
+        return;
+    }
+    metrics_->ObserveTransferOperation(kind, op_name, bytes, latency_us);
+    if (kind == TransferOperationKind::kRead) {
+        metrics_->transfer_metric.total_read_bytes.inc(bytes);
+        if (batch) {
+            metrics_->transfer_metric.batch_get_latency_us.observe(latency_us);
+        } else {
+            metrics_->transfer_metric.get_latency_us.observe(latency_us);
+        }
+    } else {
+        metrics_->transfer_metric.total_write_bytes.inc(bytes);
+        if (batch) {
+            metrics_->transfer_metric.batch_put_latency_us.observe(latency_us);
+        } else {
+            metrics_->transfer_metric.put_latency_us.observe(latency_us);
+        }
+    }
+}
+
 ErrorCode DummyClient::connect(const std::string& server_address) {
     ScopedVLogTimer timer(1, "DummyClient::Connect");
     timer.LogRequest("real_client_addr=", server_address);
@@ -524,35 +621,40 @@ uint64_t DummyClient::alloc_from_mem_pool(size_t size) {
 
 int DummyClient::put(const std::string& key, std::span<const char> value,
                      const ReplicateConfig& config) {
-    return to_py_ret(invoke_rpc<&RealClient::put_dummy_helper, void>(
-        key, value, config, client_id_));
+    return invoke_observed_void_rpc<&RealClient::put_dummy_helper>(
+        TransferOperationKind::kWrite, "put", value.size_bytes(), false, key,
+        value, config, client_id_);
 }
 
 int DummyClient::put_batch(const std::vector<std::string>& keys,
                            const std::vector<std::span<const char>>& values,
                            const ReplicateConfig& config) {
-    return to_py_ret(invoke_rpc<&RealClient::put_batch_dummy_helper, void>(
-        keys, values, config, client_id_));
+    return invoke_observed_void_rpc<&RealClient::put_batch_dummy_helper>(
+        TransferOperationKind::kWrite, "put_batch", sum_value_sizes(values),
+        true, keys, values, config, client_id_);
 }
 
 int DummyClient::put_parts(const std::string& key,
                            std::vector<std::span<const char>> values,
                            const ReplicateConfig& config) {
-    return to_py_ret(invoke_rpc<&RealClient::put_parts_dummy_helper, void>(
-        key, values, config, client_id_));
+    return invoke_observed_void_rpc<&RealClient::put_parts_dummy_helper>(
+        TransferOperationKind::kWrite, "put_parts", sum_value_sizes(values),
+        false, key, values, config, client_id_);
 }
 
 int DummyClient::upsert(const std::string& key, std::span<const char> value,
                         const ReplicateConfig& config) {
-    return to_py_ret(invoke_rpc<&RealClient::upsert_dummy_helper, void>(
-        key, value, config, client_id_));
+    return invoke_observed_void_rpc<&RealClient::upsert_dummy_helper>(
+        TransferOperationKind::kWrite, "upsert", value.size_bytes(), false, key,
+        value, config, client_id_);
 }
 
 int DummyClient::upsert_from(const std::string& key, void* buffer, size_t size,
                              const ReplicateConfig& config) {
     uint64_t dummy_addr = reinterpret_cast<uint64_t>(buffer);
-    return to_py_ret(invoke_rpc<&RealClient::upsert_from_dummy_helper, void>(
-        key, dummy_addr, size, config, client_id_));
+    return invoke_observed_void_rpc<&RealClient::upsert_from_dummy_helper>(
+        TransferOperationKind::kWrite, "upsert_from", size, false, key,
+        dummy_addr, size, config, client_id_);
 }
 
 std::vector<int> DummyClient::batch_upsert_from(
@@ -562,6 +664,7 @@ std::vector<int> DummyClient::batch_upsert_from(
     for (auto ptr : buffer_ptrs) {
         buffers.push_back(reinterpret_cast<uint64_t>(ptr));
     }
+    const auto start_time = std::chrono::steady_clock::now();
     auto internal_results =
         invoke_batch_rpc<&RealClient::batch_upsert_from_dummy_helper, void>(
             keys.size(), keys, buffers, sizes, config, client_id_);
@@ -570,21 +673,29 @@ std::vector<int> DummyClient::batch_upsert_from(
     for (const auto& result : internal_results) {
         results.push_back(to_py_ret(result));
     }
+    const size_t successful_bytes = sum_successful_sizes(results, sizes);
+    if (successful_bytes > 0) {
+        ObserveTransferMetric(TransferOperationKind::kWrite,
+                              "batch_upsert_from", successful_bytes,
+                              elapsed_us_since(start_time), true);
+    }
     return results;
 }
 
 int DummyClient::upsert_parts(const std::string& key,
                               std::vector<std::span<const char>> values,
                               const ReplicateConfig& config) {
-    return to_py_ret(invoke_rpc<&RealClient::upsert_parts_dummy_helper, void>(
-        key, values, config, client_id_));
+    return invoke_observed_void_rpc<&RealClient::upsert_parts_dummy_helper>(
+        TransferOperationKind::kWrite, "upsert_parts", sum_value_sizes(values),
+        false, key, values, config, client_id_);
 }
 
 int DummyClient::upsert_batch(const std::vector<std::string>& keys,
                               const std::vector<std::span<const char>>& values,
                               const ReplicateConfig& config) {
-    return to_py_ret(invoke_rpc<&RealClient::upsert_batch_dummy_helper, void>(
-        keys, values, config, client_id_));
+    return invoke_observed_void_rpc<&RealClient::upsert_batch_dummy_helper>(
+        TransferOperationKind::kWrite, "upsert_batch", sum_value_sizes(values),
+        true, keys, values, config, client_id_);
 }
 
 int DummyClient::remove(const std::string& key, bool force) {
@@ -650,6 +761,7 @@ int64_t DummyClient::getSize(const std::string& key) {
 }
 
 std::shared_ptr<BufferHandle> DummyClient::get_buffer(const std::string& key) {
+    const auto start_time = std::chrono::steady_clock::now();
     // Try hot cache path if shm is mapped
     if (hot_cache_base_) {
         auto result = invoke_rpc<&RealClient::acquire_hot_cache,
@@ -668,6 +780,8 @@ std::shared_ptr<BufferHandle> DummyClient::get_buffer(const std::string& key) {
                 (void)invoke_rpc<&RealClient::release_hot_cache, void>(
                     key_copy);
             };
+            ObserveTransferMetric(TransferOperationKind::kRead, "get_buffer",
+                                  size, elapsed_us_since(start_time), false);
             return std::make_shared<BufferHandle>(local_ptr, size,
                                                   std::move(release));
         }
@@ -686,11 +800,14 @@ std::shared_ptr<BufferHandle> DummyClient::get_buffer(const std::string& key) {
         (void)invoke_rpc<&RealClient::release_buffer_dummy, void>(dummy_addr,
                                                                   client_id_);
     };
+    ObserveTransferMetric(TransferOperationKind::kRead, "get_buffer", size,
+                          elapsed_us_since(start_time), false);
     return std::make_shared<BufferHandle>(local_ptr, size, std::move(release));
 }
 
 std::vector<std::shared_ptr<BufferHandle>> DummyClient::batch_get_buffer(
     const std::vector<std::string>& keys) {
+    const auto start_time = std::chrono::steady_clock::now();
     std::vector<std::shared_ptr<BufferHandle>> results(keys.size(), nullptr);
     if (keys.empty()) return results;
 
@@ -751,19 +868,37 @@ std::vector<std::shared_ptr<BufferHandle>> DummyClient::batch_get_buffer(
             std::make_shared<BufferHandle>(ptr, size, std::move(release));
     }
 
+    size_t total_bytes = 0;
+    for (const auto& result : results) {
+        if (result != nullptr) {
+            total_bytes += result->size();
+        }
+    }
+    if (total_bytes > 0) {
+        ObserveTransferMetric(TransferOperationKind::kRead, "batch_get_buffer",
+                              total_bytes, elapsed_us_since(start_time), true);
+    }
+
     return results;
 }
 
 int64_t DummyClient::get_into(const std::string& key, void* buffer,
                               size_t size) {
     uint64_t buf_addr = reinterpret_cast<uint64_t>(buffer);
+    const auto start_time = std::chrono::steady_clock::now();
     auto result = invoke_rpc<&RealClient::get_into_range_shm_helper,
                              tl::expected<int64_t, ErrorCode>>(
         key, buf_addr, 0, 0, size, client_id_);
     if (!result) {
         return static_cast<int64_t>(toInt(result.error()));
     }
-    return to_py_ret(*result);
+    const int64_t bytes_read = to_py_ret(*result);
+    if (bytes_read >= 0) {
+        ObserveTransferMetric(TransferOperationKind::kRead, "get_into",
+                              static_cast<size_t>(bytes_read),
+                              elapsed_us_since(start_time), false);
+    }
+    return bytes_read;
 }
 
 std::vector<std::vector<std::vector<int64_t>>> DummyClient::get_into_ranges(
@@ -773,6 +908,7 @@ std::vector<std::vector<std::vector<int64_t>>> DummyClient::get_into_ranges(
     const std::vector<std::vector<std::vector<size_t>>>& all_src_offsets,
     const std::vector<std::vector<std::vector<size_t>>>& all_sizes) {
     std::vector<uint64_t> dummy_buffers = void_ptrs_to_u64(buffers);
+    const auto start_time = std::chrono::steady_clock::now();
     auto internal_results =
         invoke_rpc<&RealClient::get_into_ranges_shm_helper,
                    std::vector<std::vector<
@@ -786,8 +922,13 @@ std::vector<std::vector<std::vector<int64_t>>> DummyClient::get_into_ranges(
                                                all_dst_offsets,
                                                internal_results.error());
     }
-
-    return convert_ranged_read_results(internal_results.value());
+    auto results = convert_ranged_read_results(internal_results.value());
+    const size_t total_bytes = sum_positive_ranges(results);
+    if (total_bytes > 0) {
+        ObserveTransferMetric(TransferOperationKind::kRead, "get_into_ranges",
+                              total_bytes, elapsed_us_since(start_time), true);
+    }
+    return results;
 }
 
 std::string DummyClient::get_hostname() const {
@@ -799,6 +940,7 @@ std::vector<int> DummyClient::batch_put_from(
     const std::vector<std::string>& keys, const std::vector<void*>& buffer_ptrs,
     const std::vector<size_t>& sizes, const ReplicateConfig& config) {
     std::vector<uint64_t> buffers = void_ptrs_to_u64(buffer_ptrs);
+    const auto start_time = std::chrono::steady_clock::now();
     auto internal_results =
         invoke_batch_rpc<&RealClient::batch_put_from_dummy_helper, void>(
             keys.size(), keys, buffers, sizes, config, device_id_, client_id_);
@@ -809,6 +951,13 @@ std::vector<int> DummyClient::batch_put_from(
         results.push_back(to_py_ret(result));
     }
 
+    const size_t successful_bytes = sum_successful_sizes(results, sizes);
+    if (successful_bytes > 0) {
+        ObserveTransferMetric(TransferOperationKind::kWrite, "batch_put_from",
+                              successful_bytes, elapsed_us_since(start_time),
+                              true);
+    }
+
     return results;
 }
 
@@ -822,6 +971,7 @@ std::vector<int64_t> DummyClient::batch_get_into(
     const std::vector<std::string>& keys, const std::vector<void*>& buffer_ptrs,
     const std::vector<size_t>& sizes) {
     std::vector<uint64_t> buffers = void_ptrs_to_u64(buffer_ptrs);
+    const auto start_time = std::chrono::steady_clock::now();
     auto internal_results =
         invoke_batch_rpc<&RealClient::batch_get_into_dummy_helper, int64_t>(
             keys.size(), keys, buffers, sizes, device_id_, client_id_);
@@ -832,6 +982,12 @@ std::vector<int64_t> DummyClient::batch_get_into(
         results.push_back(to_py_ret(result));
     }
 
+    const size_t total_bytes = sum_positive_results(results);
+    if (total_bytes > 0) {
+        ObserveTransferMetric(TransferOperationKind::kRead, "batch_get_into",
+                              total_bytes, elapsed_us_since(start_time), true);
+    }
+
     return results;
 }
 
@@ -850,6 +1006,7 @@ std::vector<int> DummyClient::batch_put_from_multi_buffers(
     const ReplicateConfig& config) {
     std::vector<std::vector<uint64_t>> dummy_nested =
         void_ptr_rows_to_u64_nested(all_buffer_ptrs);
+    const auto start_time = std::chrono::steady_clock::now();
     auto internal_results =
         invoke_batch_rpc<&RealClient::batch_put_from_multi_buffers_dummy_helper,
                          void>(keys.size(), keys, dummy_nested, all_sizes,
@@ -859,6 +1016,13 @@ std::vector<int> DummyClient::batch_put_from_multi_buffers(
     for (const auto& result : internal_results) {
         results.push_back(to_py_ret(result));
     }
+    const size_t successful_bytes =
+        sum_successful_nested_sizes(results, all_sizes);
+    if (successful_bytes > 0) {
+        ObserveTransferMetric(TransferOperationKind::kWrite,
+                              "batch_put_from_multi_buffers", successful_bytes,
+                              elapsed_us_since(start_time), true);
+    }
     return results;
 }
 
@@ -869,6 +1033,7 @@ std::vector<int> DummyClient::batch_get_into_multi_buffers(
     bool prefer_alloc_in_same_node) {
     std::vector<std::vector<uint64_t>> dummy_nested =
         void_ptr_rows_to_u64_nested(all_buffer_ptrs);
+    const auto start_time = std::chrono::steady_clock::now();
     auto internal_results =
         invoke_batch_rpc<&RealClient::batch_get_into_multi_buffers_dummy_helper,
                          int64_t>(keys.size(), keys, dummy_nested, all_sizes,
@@ -879,6 +1044,12 @@ std::vector<int> DummyClient::batch_get_into_multi_buffers(
     for (const auto& result : internal_results) {
         results.push_back(to_py_ret(result));
     }
+    const size_t total_bytes = sum_positive_results(results);
+    if (total_bytes > 0) {
+        ObserveTransferMetric(TransferOperationKind::kRead,
+                              "batch_get_into_multi_buffers", total_bytes,
+                              elapsed_us_since(start_time), true);
+    }
     return results;
 }
 

mooncake-store/src/file_storage.cpp

@@ -7,12 +7,17 @@
 #include "storage_backend.h"
 #include "client_metric.h"
 #include "utils.h"
+#include "gpu_staging_utils.h"
 #ifdef USE_URING
 #include "file_interface.h"
 #endif
 
 namespace mooncake {
 
+using gpu_staging::CopyDeviceToHost;
+using gpu_staging::IsDevicePointer;
+using gpu_staging::SetDevice;
+
 FileStorageConfig FileStorageConfig::FromEnvironment() {
     FileStorageConfig config;
 
@@ -37,9 +42,10 @@ FileStorageConfig FileStorageConfig::FromEnvironment() {
     config.local_buffer_size = GetEnvOr<int64_t>(
         "MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES", config.local_buffer_size);
 
-    config.scanmeta_iterator_keys_limit =
+    config.scanmeta_iterator_keys_limit = GetEnvOr<int64_t>(
+        "MOONCAKE_OFFLOAD_SCANMETA_ITERATOR_KEYS_LIMIT",
         GetEnvOr<int64_t>("MOONCAKE_SCANMETA_ITERATOR_KEYS_LIMIT",
-                          config.scanmeta_iterator_keys_limit);
+                          config.scanmeta_iterator_keys_limit));
 
     config.total_keys_limit = GetEnvOr<int64_t>(
         "MOONCAKE_OFFLOAD_TOTAL_KEYS_LIMIT", config.total_keys_limit);
@@ -58,7 +64,9 @@ FileStorageConfig FileStorageConfig::FromEnvironment() {
         GetEnvOr<uint64_t>("MOONCAKE_OFFLOAD_CLIENT_BUFFER_GC_TTL_MS",
                            config.client_buffer_gc_ttl_ms);
 
-    auto use_uring_str = GetEnvStringOr("MOONCAKE_USE_URING", "false");
+    auto use_uring_str =
+        GetEnvStringOr("MOONCAKE_OFFLOAD_USE_URING",
+                       GetEnvStringOr("MOONCAKE_USE_URING", "false"));
     config.use_uring = (use_uring_str == "true" || use_uring_str == "1");
 
     return config;
@@ -152,6 +160,7 @@ FileStorage::FileStorage(const FileStorageConfig& config,
       client_(client),
       ssd_metric_(ssd_metric),
       local_rpc_addr_(local_rpc_addr),
+      pinned_buffer_pool_(std::make_unique<PinnedBufferPool>()),
       client_buffer_allocator_(
           AlignedClientBufferAllocator::create(config.local_buffer_size, "")) {
     if (!config.Validate()) {
@@ -372,6 +381,37 @@ tl::expected<void, ErrorCode> FileStorage::OffloadObjects(
             }
         };
 
+        // D2H staging: replace device slices with host memory slices
+        // so that storage_backend (ConcatSlicesToString / BuildBucket /
+        // WriteBucket) always receives host pointers.
+        std::unordered_map<std::string, std::vector<Slice>> host_batch_object;
+        std::vector<PinnedBufferPool::Buffer> staging_bufs;
+
+        for (auto& [obj_key, slices] : batch_object) {
+            std::vector<Slice> host_slices;
+            bool obj_success = true;
+            for (const auto& slice : slices) {
+                int device_id = -1;
+                if (IsDevicePointer(slice.ptr, &device_id)) {
+                    SetDevice(device_id);
+                    auto buf = pinned_buffer_pool_->Acquire(slice.size);
+                    if (!CopyDeviceToHost(buf.data, slice.ptr, slice.size)) {
+                        LOG(ERROR) << "D2H staging failed for key: " << obj_key;
+                        pinned_buffer_pool_->Release(buf);
+                        obj_success = false;
+                        break;
+                    }
+                    host_slices.emplace_back(Slice{buf.data, slice.size});
+                    staging_bufs.push_back(buf);
+                } else {
+                    host_slices.push_back(slice);
+                }
+            }
+            if (obj_success) {
+                host_batch_object[obj_key] = std::move(host_slices);
+            }
+        }
+
         auto offload_start = std::chrono::steady_clock::now();
         auto bucket_complete_handler =
             [this, offload_start, complete_handler](
@@ -399,7 +439,12 @@ tl::expected<void, ErrorCode> FileStorage::OffloadObjects(
             return res;
         };
         auto offload_res = storage_backend_->BatchOffload(
-            batch_object, bucket_complete_handler, eviction_handler);
+            host_batch_object, bucket_complete_handler, eviction_handler);
+
+        // Release staging buffers back to pool (Buffer is POD, no destructor)
+        for (auto& buf : staging_bufs) {
+            pinned_buffer_pool_->Release(buf);
+        }
         if (!offload_res) {
             LOG(ERROR) << "Failed to store objects with error: "
                        << offload_res.error();

mooncake-store/src/master.cpp

@@ -109,6 +109,10 @@ DEFINE_validator(eviction_ratio, [](const char* flagname, double value) {
 DEFINE_bool(enable_ha, false,
             "Enable high availability, which depends on etcd");
 DEFINE_bool(enable_offload, false, "Enable offload availability");
+DEFINE_bool(offload_on_evict, false,
+            "Defer LOCAL_DISK offload to eviction time instead of PutEnd");
+DEFINE_bool(offload_force_evict, false,
+            "Force-evict objects exceeding offload cap without disk offload");
 DEFINE_string(ha_backend_type, "etcd",
               "HA backend type, e.g. etcd | redis | k8s");
 DEFINE_string(ha_backend_connstring, "",
@@ -293,6 +297,11 @@ void InitMasterConf(const mooncake::DefaultConfig& default_config,
                            FLAGS_enable_ha);
     default_config.GetBool("enable_offload", &master_config.enable_offload,
                            FLAGS_enable_offload);
+    default_config.GetBool("offload_on_evict", &master_config.offload_on_evict,
+                           FLAGS_offload_on_evict);
+    default_config.GetBool("offload_force_evict",
+                           &master_config.offload_force_evict,
+                           FLAGS_offload_force_evict);
     default_config.GetString("ha_backend_type", &master_config.ha_backend_type,
                              FLAGS_ha_backend_type);
     default_config.GetString("ha_backend_connstring",
@@ -867,6 +876,8 @@ int main(int argc, char* argv[]) {
         << master_config.eviction_high_watermark_ratio
         << ", enable_ha=" << master_config.enable_ha
         << ", enable_offload=" << master_config.enable_offload
+        << ", offload_on_evict=" << master_config.offload_on_evict
+        << ", offload_force_evict=" << master_config.offload_force_evict
         << ", ha_backend_type=" << master_config.ha_backend_type
         << ", ha_backend_connstring=" << ha_backend_connstring
         << ", etcd_endpoints=" << master_config.etcd_endpoints

mooncake-store/src/master_service.cpp

@@ -48,6 +48,13 @@ namespace {
 
 constexpr size_t kUnlimitedSnapshotList = 0;
 
+// Per-cycle offload cap as a fraction of `offloading_queue_limit_`. Used only
+// when offload-on-evict mode is active. Defers memory eviction for at most
+// this fraction of the queue limit per BatchEvict cycle; beyond that, eviction
+// falls back according to `offload_force_evict_`. A future change may expose
+// this as a configurable parameter if workloads demand tuning.
+constexpr double kOffloadCapRatio = 0.5;
+
 enum class SnapshotCatalogBackendKind {
     kEmbedded,
     kRedis,
@@ -159,6 +166,19 @@ MasterService::MasterService(const MasterServiceConfig& config)
             "put_start_discard_timeout_sec");
     }
 
+    // Offload-on-evict: defer LOCAL_DISK offload to eviction time
+    offload_on_evict_ = enable_offload_ && config.offload_on_evict;
+    if (offload_on_evict_) {
+        LOG(INFO) << "Offload-on-evict mode enabled: DRAM offload to "
+                     "LOCAL_DISK will occur at eviction time instead of "
+                     "PutEnd";
+        offload_force_evict_ = config.offload_force_evict;
+        if (offload_force_evict_) {
+            LOG(INFO) << "Force-evict enabled: objects exceeding offload "
+                         "cap will be evicted without disk offload";
+        }
+    }
+
     eviction_running_ = true;
     eviction_thread_ = std::thread(&MasterService::EvictionThreadFunc, this);
     VLOG(1) << "action=start_eviction_thread";
@@ -351,12 +371,23 @@ auto MasterService::ReMountSegment(const std::vector<Segment>& segments,
     return {};
 }
 
+std::unordered_set<UUID, boost::hash<UUID>>
+MasterService::getAliveClientsSnapshot() const {
+    std::shared_lock<std::shared_mutex> lock(client_mutex_);
+    return ok_client_;
+}
+
 void MasterService::ClearInvalidHandles() {
+    ClearInvalidHandles(getAliveClientsSnapshot());
+}
+
+void MasterService::ClearInvalidHandles(
+    const std::unordered_set<UUID, boost::hash<UUID>>& alive_clients) {
     for (size_t i = 0; i < kNumShards; i++) {
         MetadataShardAccessorRW shard(this, i);
         auto it = shard->metadata.begin();
         while (it != shard->metadata.end()) {
-            if (CleanupStaleHandles(it->second)) {
+            if (CleanupStaleHandles(it->second, alive_clients)) {
                 // If the object is empty, we need to erase the iterator and
                 // also erase the key from processing_keys,
                 // replication_tasks, and offloading_tasks.
@@ -842,13 +873,15 @@ auto MasterService::PutStart(const UUID& client_id, const std::string& key,
     VLOG(1) << "key=" << key << ", value_length=" << slice_length
             << ", config=" << config << ", action=put_start_begin";
 
+    auto alive_clients = getAliveClientsSnapshot();
     std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
     // Lock the shard and check if object already exists
     MetadataShardAccessorRW shard(this, getShardIndex(key));
 
     const auto now = std::chrono::system_clock::now();
     auto it = shard->metadata.find(key);
-    if (it != shard->metadata.end() && !CleanupStaleHandles(it->second)) {
+    if (it != shard->metadata.end() &&
+        !CleanupStaleHandles(it->second, alive_clients)) {
         auto& metadata = it->second;
         // If the object's PutStart expired and has not completed any
         // replicas, we can discard it and allow the new PutStart to
@@ -897,7 +930,7 @@ auto MasterService::PutEnd(const UUID& client_id, const std::string& key,
         },
         [](Replica& replica) { replica.mark_complete(); });
 
-    if (enable_offload_) {
+    if (enable_offload_ && !offload_on_evict_) {
         auto& shard = accessor.GetShard();
         metadata.VisitReplicas(
             &Replica::fn_is_completed, [this, &key, &shard](Replica& replica) {
@@ -1085,6 +1118,7 @@ auto MasterService::UpsertStart(const UUID& client_id, const std::string& key,
     //   during full metadata snapshots.
     // shard lock (exclusive via MetadataShardAccessorRW): serializes all
     //   operations on keys that hash to the same shard.
+    auto alive_clients = getAliveClientsSnapshot();
     std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
     MetadataShardAccessorRW shard(this, getShardIndex(key));
 
@@ -1094,7 +1128,9 @@ auto MasterService::UpsertStart(const UUID& client_id, const std::string& key,
     // --- Step 0: stale handle cleanup ---
     // If all memory replicas point to unmounted segments (node crashed and
     // restarted), the metadata is useless — erase it and treat as new key.
-    if (it != shard->metadata.end() && CleanupStaleHandles(it->second)) {
+    // Also clean up local_disk replicas whose owner client has expired.
+    if (it != shard->metadata.end() &&
+        CleanupStaleHandles(it->second, alive_clients)) {
         shard->processing_keys.erase(key);
         shard->metadata.erase(it);
         it = shard->metadata.end();
@@ -1931,6 +1967,8 @@ auto MasterService::BatchRemove(const std::vector<std::string>& keys,
 
     std::shared_lock<std::shared_mutex> snapshot_lock(snapshot_mutex_);
 
+    auto alive_clients = getAliveClientsSnapshot();
+
     // Process each shard once, acquiring lock per shard
     for (auto& [shard_idx, key_group] : keys_by_shard) {
         MetadataShardAccessorRW shard(this, shard_idx);
@@ -1948,7 +1986,7 @@ auto MasterService::BatchRemove(const std::vector<std::string>& keys,
             }
 
             // Clean up stale replica handles (consistent with single Remove)
-            if (CleanupStaleHandles(it->second)) {
+            if (CleanupStaleHandles(it->second, alive_clients)) {
                 shard->processing_keys.erase(key);
                 shard->replication_tasks.erase(key);
                 shard->offloading_tasks.erase(key);
@@ -1996,10 +2034,14 @@ auto MasterService::BatchRemove(const std::vector<std::string>& keys,
     return results;
 }
 
-bool MasterService::CleanupStaleHandles(ObjectMetadata& metadata) {
-    // Remove those with invalid allocators
-    metadata.EraseReplicas([](const Replica& replica) {
-        return replica.has_invalid_mem_handle();
+bool MasterService::CleanupStaleHandles(
+    ObjectMetadata& metadata,
+    const std::unordered_set<UUID, boost::hash<UUID>>& alive_clients) {
+    // Remove those with invalid allocators (memory replicas on unmounted
+    // segments) and local_disk replicas whose owner client is no longer alive.
+    metadata.EraseReplicas([&alive_clients](const Replica& replica) {
+        return replica.has_invalid_mem_handle() ||
+               replica.has_stale_local_disk_client(alive_clients);
     });
 
     // Return true if no valid replicas remain after cleanup
@@ -3469,6 +3511,81 @@ void MasterService::BatchEvict(double evict_ratio_target,
         });
     };
 
+    // --- Offload-on-evict support ---
+    long offload_queued_this_cycle = 0;
+    long offload_deferred_count = 0;
+    long offload_cap_forced_count = 0;    // #keys force-evicted due to cap
+    long offload_push_failed_forced = 0;  // #keys force-evicted on push fail
+    const long offload_cap =
+        offload_on_evict_
+            ? static_cast<long>(offloading_queue_limit_ * kOffloadCapRatio)
+            : 0;
+
+    auto has_local_disk_replica = [](const ObjectMetadata& metadata) {
+        return metadata.HasReplica(&Replica::fn_is_local_disk_replica);
+    };
+
+    // Returns freed bytes. Returns 0 if offload-queued and no additional
+    // replicas were evicted (all MEMORY replicas of the key are now pinned).
+    auto try_evict_or_offload =
+        [&, this](const std::string& key, ObjectMetadata& metadata,
+                  MetadataShardAccessorRW& shard) -> uint64_t {
+        if (!offload_on_evict_) {
+            // Original behavior
+            return metadata.size * evict_replicas(metadata);
+        }
+
+        // LOCAL_DISK replica already exists — safe to delete MEMORY immediately
+        if (has_local_disk_replica(metadata)) {
+            return metadata.size * evict_replicas(metadata);
+        }
+
+        // Force-evict cap: if force_evict enabled and cap reached, force
+        // delete. Warning is aggregated at the end of the cycle to avoid log
+        // flooding.
+        if (offload_force_evict_ && offload_queued_this_cycle >= offload_cap) {
+            offload_cap_forced_count++;
+            return metadata.size * evict_replicas(metadata);
+        }
+
+        // Queue one MEMORY replica for offload; others will be evicted below.
+        bool queued = false;
+        metadata.VisitReplicas(
+            [](const Replica& r) {
+                return r.is_memory_replica() && r.is_completed() &&
+                       r.get_refcnt() == 0;
+            },
+            [this, &key, &shard, &queued, &now](Replica& replica) {
+                if (queued) return;  // only need to pin one replica for offload
+                auto result = PushOffloadingQueue(key, replica);
+                if (result) {
+                    replica.inc_refcnt();
+                    shard->offloading_tasks.emplace(
+                        key, OffloadingTask{replica.id(), now});
+                    queued = true;
+                }
+            });
+
+        if (queued) {
+            offload_queued_this_cycle++;
+            offload_deferred_count++;
+            // Any remaining MEMORY replicas with refcnt==0 are redundant copies
+            // (data survives via the pinned replica → disk). Evict them now to
+            // reclaim memory immediately rather than waiting another cycle.
+            return metadata.size * evict_replicas(metadata);
+        }
+
+        // PushOffloadingQueue failed. Default (data-preserving) behavior is to
+        // skip this cycle — the outer eviction loop will retry after the
+        // offload queue drains. Only force-evict when explicitly opted in, to
+        // prevent silent data loss when the queue is unavailable.
+        if (offload_force_evict_) {
+            offload_push_failed_forced++;
+            return metadata.size * evict_replicas(metadata);
+        }
+        return 0;
+    };
+
     // Randomly select a starting shard to avoid imbalance eviction between
     // shards. No need to use expensive random_device here.
     size_t start_idx = rand() % kNumShards;
@@ -3541,16 +3658,18 @@ void MasterService::BatchEvict(double evict_ratio_target,
                     continue;
                 }
                 if (it->second.lease_timeout <= target_timeout) {
-                    // Evict this object
-                    total_freed_size +=
-                        it->second.size *
-                        evict_replicas(it->second);  // Erase memory replicas
+                    // Evict this object (or defer for offload)
+                    uint64_t freed =
+                        try_evict_or_offload(it->first, it->second, shard);
+                    total_freed_size += freed;
                     if (it->second.IsValid() == false) {
                         it = shard->metadata.erase(it);
                     } else {
                         ++it;
                     }
-                    shard_evicted_count++;
+                    if (freed > 0) {
+                        shard_evicted_count++;
+                    }
                 } else {
                     // second pass candidates
                     no_pin_objects.push_back(it->second.lease_timeout);
@@ -3598,20 +3717,22 @@ void MasterService::BatchEvict(double evict_ratio_target,
                 auto it = shard->metadata.begin();
                 while (it != shard->metadata.end() && target_evict_num > 0) {
                     if (!it->second.IsHardPinned() &&
+                        it->second.IsLeaseExpired(now) &&
                         it->second.lease_timeout <= target_timeout &&
                         !it->second.IsSoftPinned(now) &&
                         can_evict_replicas(it->second)) {
-                        // Evict this object
-                        total_freed_size +=
-                            it->second.size *
-                            evict_replicas(
-                                it->second);  // Erase memory replicas
+                        // Evict this object (or defer for offload)
+                        uint64_t freed =
+                            try_evict_or_offload(it->first, it->second, shard);
+                        total_freed_size += freed;
                         if (it->second.IsValid() == false) {
                             it = shard->metadata.erase(it);
                         } else {
                             ++it;
                         }
-                        evicted_count++;
+                        if (freed > 0) {
+                            evicted_count++;
+                        }
                         target_evict_num--;
                     } else {
                         ++it;
@@ -3651,16 +3772,18 @@ void MasterService::BatchEvict(double evict_ratio_target,
                     // and lease timeout less than or equal to target.
                     if (!it->second.IsSoftPinned(now) ||
                         it->second.lease_timeout <= soft_target_timeout) {
-                        total_freed_size +=
-                            it->second.size *
-                            evict_replicas(
-                                it->second);  // Erase memory replicas
+                        // Evict this object (or defer for offload)
+                        uint64_t freed =
+                            try_evict_or_offload(it->first, it->second, shard);
+                        total_freed_size += freed;
                         if (it->second.IsValid() == false) {
                             it = shard->metadata.erase(it);
                         } else {
                             ++it;
                         }
-                        evicted_count++;
+                        if (freed > 0) {
+                            evicted_count++;
+                        }
                         target_evict_num--;
                     } else {
                         ++it;
@@ -3681,7 +3804,11 @@ void MasterService::BatchEvict(double evict_ratio_target,
         }
     }
 
-    if (evicted_count > 0 || released_discarded_cnt > 0) {
+    if (evicted_count > 0 || released_discarded_cnt > 0 ||
+        offload_deferred_count > 0) {
+        // Offload-deferred counts as partial success: work was done (objects
+        // queued for disk offload), so suppress re-triggering until the next
+        // watermark breach or explicit need_eviction_ signal.
         need_eviction_ = false;
         MasterMetricManager::instance().inc_eviction_success(evicted_count,
                                                              total_freed_size);
@@ -3693,7 +3820,27 @@ void MasterService::BatchEvict(double evict_ratio_target,
         MasterMetricManager::instance().inc_eviction_fail();
     }
     VLOG(1) << "action=evict_objects" << ", evicted_count=" << evicted_count
+            << ", offload_deferred=" << offload_deferred_count
+            << ", offload_cap_forced=" << offload_cap_forced_count
+            << ", offload_push_failed_forced=" << offload_push_failed_forced
             << ", total_freed_size=" << total_freed_size;
+    if (offload_on_evict_ && evicted_count == 0 && offload_deferred_count > 0) {
+        LOG(WARNING) << "[EVICT] No memory freed this cycle; "
+                     << offload_deferred_count
+                     << " objects deferred for disk offload. "
+                        "Consider lowering eviction_high_watermark_ratio.";
+    }
+    if (offload_cap_forced_count > 0) {
+        LOG(WARNING) << "[EVICT] Offload cap (" << offload_cap
+                     << ") reached; force-evicted " << offload_cap_forced_count
+                     << " object(s) without disk offload this cycle.";
+    }
+    if (offload_push_failed_forced > 0) {
+        LOG(WARNING) << "[EVICT] PushOffloadingQueue failed for "
+                     << offload_push_failed_forced
+                     << " object(s); force-evicted without disk offload "
+                        "(offload_force_evict=true).";
+    }
 }
 
 void MasterService::ClientMonitorFunc() {
@@ -3770,9 +3917,15 @@ void MasterService::ClientMonitorFunc() {
             }  // Release the mutex before long-running ClearInvalidHandles and
                // avoid deadlocks
 
-            if (!unmount_segments.empty()) {
-                ClearInvalidHandles();
+            // Always clean up invalid handles when there are expired clients,
+            // even if no memory segments were unmounted. This is necessary
+            // to clean up local_disk replicas whose owner client has expired.
+            ClearInvalidHandles();
 
+            // Commit unmount of memory segments and clean up local_disk
+            // segments for expired clients. Both require the exclusive
+            // segment lock.
+            {
                 ScopedSegmentAccess segment_access =
                     segment_manager_.getSegmentAccess();
                 for (size_t i = 0; i < unmount_segments.size(); i++) {
@@ -3782,6 +3935,9 @@ void MasterService::ClientMonitorFunc() {
                               << ", segment_name=" << segment_names[i]
                               << ", action=unmount_expired_segment";
                 }
+                for (auto& client_id : expired_clients) {
+                    segment_access.UnmountLocalDiskSegment(client_id);
+                }
             }
         }
 

mooncake-store/src/real_client.cpp

@@ -62,6 +62,87 @@ struct PreparedRangedReadRequest {
     bool has_any_valid_fragment = false;
 };
 
+size_t sum_value_sizes(const std::vector<std::span<const char>> &values) {
+    size_t total = 0;
+    for (const auto &value : values) {
+        total += value.size_bytes();
+    }
+    return total;
+}
+
+size_t sum_sizes(const std::vector<size_t> &sizes) {
+    size_t total = 0;
+    for (size_t size : sizes) {
+        total += size;
+    }
+    return total;
+}
+
+size_t sum_successful_sizes(const std::vector<int> &results,
+                            const std::vector<size_t> &sizes) {
+    size_t total = 0;
+    for (size_t i = 0; i < results.size() && i < sizes.size(); ++i) {
+        if (results[i] == 0) {
+            total += sizes[i];
+        }
+    }
+    return total;
+}
+
+size_t sum_successful_nested_sizes(
+    const std::vector<int> &results,
+    const std::vector<std::vector<size_t>> &nested_sizes) {
+    size_t total = 0;
+    for (size_t i = 0; i < results.size() && i < nested_sizes.size(); ++i) {
+        if (results[i] == 0) {
+            total += sum_sizes(nested_sizes[i]);
+        }
+    }
+    return total;
+}
+
+size_t sum_positive_results(const std::vector<int64_t> &results) {
+    size_t total = 0;
+    for (int64_t result : results) {
+        if (result > 0) {
+            total += static_cast<size_t>(result);
+        }
+    }
+    return total;
+}
+
+size_t sum_positive_results(const std::vector<int> &results) {
+    size_t total = 0;
+    for (int result : results) {
+        if (result > 0) {
+            total += static_cast<size_t>(result);
+        }
+    }
+    return total;
+}
+
+size_t sum_positive_ranges(
+    const std::vector<std::vector<std::vector<int64_t>>> &results) {
+    size_t total = 0;
+    for (const auto &key_rows : results) {
+        for (const auto &row : key_rows) {
+            total += sum_positive_results(row);
+        }
+    }
+    return total;
+}
+
+size_t sum_buffer_handle_sizes(
+    const std::vector<std::shared_ptr<BufferHandle>> &buffers) {
+    size_t total = 0;
+    for (const auto &buffer : buffers) {
+        if (buffer != nullptr) {
+            total += buffer->size();
+        }
+    }
+    return total;
+}
+
 PreparedRangedReadRequest prepare_ranged_read_request(
     size_t buffer_count, const std::vector<std::vector<std::string>> &all_keys,
     const std::vector<std::vector<std::vector<size_t>>> &all_dst_offsets,
@@ -454,7 +535,8 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
     const std::string &master_server_addr,
     const std::shared_ptr<TransferEngine> &transfer_engine,
     const std::string &ipc_socket_path, int local_rpc_port,
-    bool enable_ssd_offload, bool start_offload_rpc_server) {
+    bool enable_ssd_offload, bool start_offload_rpc_server,
+    const std::string &ssd_offload_path) {
     this->protocol = protocol;
     this->ipc_socket_path_ = ipc_socket_path;
     const bool should_use_hugepage = use_hugepage_ &&
@@ -493,7 +575,7 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
             hostname.substr(0, colon_pos + 1) + std::to_string(local_rpc_port);
         auto client_opt = mooncake::Client::Create(
             this->local_hostname, metadata_server, protocol, device_name,
-            master_server_addr, transfer_engine);
+            master_server_addr, transfer_engine, {{"client_mode", "real"}});
         if (!client_opt) {
             LOG(ERROR) << "Failed to create client";
             return tl::unexpected(ErrorCode::INVALID_PARAMS);
@@ -522,7 +604,7 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
                 hostname + ":" + std::to_string(local_rpc_port);
             auto client_opt = mooncake::Client::Create(
                 this->local_hostname, metadata_server, protocol, device_name,
-                master_server_addr, transfer_engine);
+                master_server_addr, transfer_engine, {{"client_mode", "real"}});
             if (client_opt) {
                 client_ = *client_opt;
                 success = true;
@@ -715,6 +797,9 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
     }
     if (enable_ssd_offload) {
         auto file_storage_config = FileStorageConfig::FromEnvironment();
+        if (!ssd_offload_path.empty()) {
+            file_storage_config.storage_filepath = ssd_offload_path;
+        }
         file_storage_ = std::make_shared<FileStorage>(
             file_storage_config, client_, this->local_rpc_addr,
             client_->GetSsdMetricPtr());
@@ -742,11 +827,12 @@ int RealClient::setup_real(
     const std::string &protocol, const std::string &rdma_devices,
     const std::string &master_server_addr,
     const std::shared_ptr<TransferEngine> &transfer_engine,
-    const std::string &ipc_socket_path, bool enable_ssd_offload) {
+    const std::string &ipc_socket_path, bool enable_ssd_offload,
+    const std::string &ssd_offload_path) {
     return to_py_ret(setup_internal(
         local_hostname, metadata_server, global_segment_size, local_buffer_size,
         protocol, rdma_devices, master_server_addr, transfer_engine,
-        ipc_socket_path, 50052, enable_ssd_offload, true));
+        ipc_socket_path, 50052, enable_ssd_offload, true, ssd_offload_path));
 }
 
 namespace {
@@ -837,6 +923,8 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
         return tl::unexpected(ErrorCode::INVALID_PARAMS);
     }
 
+    std::string ssd_offload_path = get_config(config, "ssd_offload_path");
+
     std::string enable_ssd_offload_str =
         get_config(config, "enable_ssd_offload", "false");
     std::transform(enable_ssd_offload_str.begin(), enable_ssd_offload_str.end(),
@@ -848,7 +936,7 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
     return setup_internal(local_hostname, metadata_server, global_segment_size,
                           local_buffer_size, protocol, rdma_devices,
                           master_server_addr, nullptr, ipc_socket_path, 50052,
-                          enable_ssd_offload, true);
+                          enable_ssd_offload, true, ssd_offload_path);
 }
 
 tl::expected<void, ErrorCode> RealClient::initAll_internal(
@@ -1077,8 +1165,17 @@ tl::expected<void, ErrorCode> RealClient::put_dummy_helper(
 
 int RealClient::put(const std::string &key, std::span<const char> value,
                     const ReplicateConfig &config) {
-    return to_py_ret(
-        put_internal(key, value, config, client_buffer_allocator_));
+    auto result = execute_timed_operation<tl::expected<void, ErrorCode>>(
+        [&]() {
+            return put_internal(key, value, config, client_buffer_allocator_);
+        },
+        [](const auto &ret) { return ret.has_value(); },
+        [&](uint64_t latency_us, const auto &) {
+            client_->ObserveTransferOperation(TransferOperationKind::kWrite,
+                                              "put", value.size_bytes(),
+                                              latency_us);
+        });
+    return to_py_ret(result);
 }
 
 tl::expected<void, ErrorCode> RealClient::put_batch_internal(
@@ -1167,8 +1264,18 @@ tl::expected<void, ErrorCode> RealClient::put_batch_dummy_helper(
 int RealClient::put_batch(const std::vector<std::string> &keys,
                           const std::vector<std::span<const char>> &values,
                           const ReplicateConfig &config) {
-    return to_py_ret(
-        put_batch_internal(keys, values, config, client_buffer_allocator_));
+    auto result = execute_timed_operation<tl::expected<void, ErrorCode>>(
+        [&]() {
+            return put_batch_internal(keys, values, config,
+                                      client_buffer_allocator_);
+        },
+        [](const auto &ret) { return ret.has_value(); },
+        [&](uint64_t latency_us, const auto &) {
+            client_->ObserveTransferOperation(
+                TransferOperationKind::kWrite, "put_batch",
+                sum_value_sizes(values), latency_us);
+        });
+    return to_py_ret(result);
 }
 
 tl::expected<void, ErrorCode> RealClient::put_parts_internal(
@@ -1249,8 +1356,18 @@ tl::expected<void, ErrorCode> RealClient::put_parts_dummy_helper(
 int RealClient::put_parts(const std::string &key,
                           std::vector<std::span<const char>> values,
                           const ReplicateConfig &config) {
-    return to_py_ret(
-        put_parts_internal(key, values, config, client_buffer_allocator_));
+    auto result = execute_timed_operation<tl::expected<void, ErrorCode>>(
+        [&]() {
+            return put_parts_internal(key, values, config,
+                                      client_buffer_allocator_);
+        },
+        [](const auto &ret) { return ret.has_value(); },
+        [&](uint64_t latency_us, const auto &) {
+            client_->ObserveTransferOperation(
+                TransferOperationKind::kWrite, "put_parts",
+                sum_value_sizes(values), latency_us);
+        });
+    return to_py_ret(result);
 }
 
 tl::expected<void, ErrorCode> RealClient::remove_internal(
@@ -1847,7 +1964,14 @@ std::shared_ptr<BufferHandle> RealClient::get_buffer_internal(
 
 // Implementation of get_buffer method
 std::shared_ptr<BufferHandle> RealClient::get_buffer(const std::string &key) {
-    return get_buffer_internal(key, client_buffer_allocator_);
+    return execute_timed_operation<std::shared_ptr<BufferHandle>>(
+        [&]() { return get_buffer_internal(key, client_buffer_allocator_); },
+        [](const auto &buffer) { return buffer != nullptr; },
+        [&](uint64_t latency_us, const auto &buffer) {
+            client_->ObserveTransferOperation(TransferOperationKind::kRead,
+                                              "get_buffer", buffer->size(),
+                                              latency_us);
+        });
 }
 
 tl::expected<std::tuple<uint64_t, size_t>, ErrorCode>
@@ -2137,7 +2261,14 @@ RealClient::batch_get_buffer_internal(
 // Implementation of batch_get_buffer method
 std::vector<std::shared_ptr<BufferHandle>> RealClient::batch_get_buffer(
     const std::vector<std::string> &keys) {
-    return batch_get_buffer_internal(keys);
+    return execute_timed_operation<std::vector<std::shared_ptr<BufferHandle>>>(
+        [&]() { return batch_get_buffer_internal(keys); },
+        [](const auto &) { return true; },
+        [&](uint64_t latency_us, const auto &buffers) {
+            client_->ObserveTransferOperation(
+                TransferOperationKind::kRead, "batch_get_buffer",
+                sum_buffer_handle_sizes(buffers), latency_us);
+        });
 }
 
 tl::expected<void, ErrorCode> RealClient::register_buffer_internal(
@@ -2292,7 +2423,17 @@ tl::expected<int64_t, ErrorCode> RealClient::get_into_range_internal(
 
 int64_t RealClient::get_into(const std::string &key, void *buffer,
                              size_t size) {
-    return to_py_ret(get_into_range_internal(key, buffer, 0, 0, size, true));
+    auto result = execute_timed_operation<tl::expected<int64_t, ErrorCode>>(
+        [&]() {
+            return get_into_range_internal(key, buffer, 0, 0, size, true);
+        },
+        [](const auto &ret) { return ret.has_value(); },
+        [&](uint64_t latency_us, const auto &ret) {
+            client_->ObserveTransferOperation(
+                TransferOperationKind::kRead, "get_into",
+                static_cast<uint64_t>(ret.value()), latency_us);
+        });
+    return to_py_ret(result);
 }
 
 std::vector<std::vector<std::vector<tl::expected<int64_t, ErrorCode>>>>
@@ -2412,8 +2553,20 @@ std::vector<std::vector<std::vector<int64_t>>> RealClient::get_into_ranges(
     const std::vector<std::vector<std::vector<size_t>>> &all_dst_offsets,
     const std::vector<std::vector<std::vector<size_t>>> &all_src_offsets,
     const std::vector<std::vector<std::vector<size_t>>> &all_sizes) {
-    return convert_ranged_read_results(get_into_ranges_internal(
-        buffers, all_keys, all_dst_offsets, all_src_offsets, all_sizes));
+    auto results =
+        execute_timed_operation<std::vector<std::vector<std::vector<int64_t>>>>(
+            [&]() {
+                return convert_ranged_read_results(
+                    get_into_ranges_internal(buffers, all_keys, all_dst_offsets,
+                                             all_src_offsets, all_sizes));
+            },
+            [](const auto &) { return true; },
+            [&](uint64_t latency_us, const auto &ret) {
+                client_->ObserveTransferOperation(
+                    TransferOperationKind::kRead, "get_into_ranges",
+                    sum_positive_ranges(ret), latency_us);
+            });
+    return results;
 }
 
 std::string RealClient::get_hostname() const { return local_hostname; }
@@ -2422,7 +2575,21 @@ std::vector<int> RealClient::batch_put_from(
     const std::vector<std::string> &keys, const std::vector<void *> &buffers,
     const std::vector<size_t> &sizes, const ReplicateConfig &config) {
     auto internal_results =
-        batch_put_from_internal(keys, buffers, sizes, config);
+        execute_timed_operation<std::vector<tl::expected<void, ErrorCode>>>(
+            [&]() {
+                return batch_put_from_internal(keys, buffers, sizes, config);
+            },
+            [](const auto &) { return true; },
+            [&](uint64_t latency_us, const auto &ret) {
+                std::vector<int> py_results;
+                py_results.reserve(ret.size());
+                for (const auto &item : ret) {
+                    py_results.push_back(to_py_ret(item));
+                }
+                client_->ObserveTransferOperation(
+                    TransferOperationKind::kWrite, "batch_put_from",
+                    sum_successful_sizes(py_results, sizes), latency_us);
+            });
     std::vector<int> results;
     results.reserve(internal_results.size());
 
@@ -2565,7 +2732,14 @@ tl::expected<void, ErrorCode> RealClient::put_from_internal(
 
 int RealClient::put_from(const std::string &key, void *buffer, size_t size,
                          const ReplicateConfig &config) {
-    return to_py_ret(put_from_internal(key, buffer, size, config));
+    auto result = execute_timed_operation<tl::expected<void, ErrorCode>>(
+        [&]() { return put_from_internal(key, buffer, size, config); },
+        [](const auto &ret) { return ret.has_value(); },
+        [&](uint64_t latency_us, const auto &) {
+            client_->ObserveTransferOperation(TransferOperationKind::kWrite,
+                                              "put_from", size, latency_us);
+        });
+    return to_py_ret(result);
 }
 
 // --- Upsert implementations ---
@@ -2606,8 +2780,18 @@ tl::expected<void, ErrorCode> RealClient::upsert_internal(
 
 int RealClient::upsert(const std::string &key, std::span<const char> value,
                        const ReplicateConfig &config) {
-    return to_py_ret(
-        upsert_internal(key, value, config, client_buffer_allocator_));
+    auto result = execute_timed_operation<tl::expected<void, ErrorCode>>(
+        [&]() {
+            return upsert_internal(key, value, config,
+                                   client_buffer_allocator_);
+        },
+        [](const auto &ret) { return ret.has_value(); },
+        [&](uint64_t latency_us, const auto &) {
+            client_->ObserveTransferOperation(TransferOperationKind::kWrite,
+                                              "upsert", value.size_bytes(),
+                                              latency_us);
+        });
+    return to_py_ret(result);
 }
 
 tl::expected<void, ErrorCode> RealClient::upsert_dummy_helper(
@@ -2657,7 +2841,14 @@ tl::expected<void, ErrorCode> RealClient::upsert_from_internal(
 
 int RealClient::upsert_from(const std::string &key, void *buffer, size_t size,
                             const ReplicateConfig &config) {
-    return to_py_ret(upsert_from_internal(key, buffer, size, config));
+    auto result = execute_timed_operation<tl::expected<void, ErrorCode>>(
+        [&]() { return upsert_from_internal(key, buffer, size, config); },
+        [](const auto &ret) { return ret.has_value(); },
+        [&](uint64_t latency_us, const auto &) {
+            client_->ObserveTransferOperation(TransferOperationKind::kWrite,
+                                              "upsert_from", size, latency_us);
+        });
+    return to_py_ret(result);
 }
 
 std::vector<tl::expected<void, ErrorCode>>
@@ -2703,7 +2894,21 @@ std::vector<int> RealClient::batch_upsert_from(
     const std::vector<std::string> &keys, const std::vector<void *> &buffers,
     const std::vector<size_t> &sizes, const ReplicateConfig &config) {
     auto internal_results =
-        batch_upsert_from_internal(keys, buffers, sizes, config);
+        execute_timed_operation<std::vector<tl::expected<void, ErrorCode>>>(
+            [&]() {
+                return batch_upsert_from_internal(keys, buffers, sizes, config);
+            },
+            [](const auto &) { return true; },
+            [&](uint64_t latency_us, const auto &ret) {
+                std::vector<int> py_results;
+                py_results.reserve(ret.size());
+                for (const auto &item : ret) {
+                    py_results.push_back(to_py_ret(item));
+                }
+                client_->ObserveTransferOperation(
+                    TransferOperationKind::kWrite, "batch_upsert_from",
+                    sum_successful_sizes(py_results, sizes), latency_us);
+            });
     std::vector<int> results;
     results.reserve(internal_results.size());
     for (const auto &result : internal_results) {
@@ -2827,8 +3032,18 @@ tl::expected<void, ErrorCode> RealClient::upsert_parts_internal(
 int RealClient::upsert_parts(const std::string &key,
                              std::vector<std::span<const char>> values,
                              const ReplicateConfig &config) {
-    return to_py_ret(
-        upsert_parts_internal(key, values, config, client_buffer_allocator_));
+    auto result = execute_timed_operation<tl::expected<void, ErrorCode>>(
+        [&]() {
+            return upsert_parts_internal(key, values, config,
+                                         client_buffer_allocator_);
+        },
+        [](const auto &ret) { return ret.has_value(); },
+        [&](uint64_t latency_us, const auto &) {
+            client_->ObserveTransferOperation(
+                TransferOperationKind::kWrite, "upsert_parts",
+                sum_value_sizes(values), latency_us);
+        });
+    return to_py_ret(result);
 }
 
 tl::expected<void, ErrorCode> RealClient::upsert_parts_dummy_helper(
@@ -2931,8 +3146,18 @@ tl::expected<void, ErrorCode> RealClient::upsert_batch_dummy_helper(
 int RealClient::upsert_batch(const std::vector<std::string> &keys,
                              const std::vector<std::span<const char>> &values,
                              const ReplicateConfig &config) {
-    return to_py_ret(
-        upsert_batch_internal(keys, values, config, client_buffer_allocator_));
+    auto result = execute_timed_operation<tl::expected<void, ErrorCode>>(
+        [&]() {
+            return upsert_batch_internal(keys, values, config,
+                                         client_buffer_allocator_);
+        },
+        [](const auto &ret) { return ret.has_value(); },
+        [&](uint64_t latency_us, const auto &) {
+            client_->ObserveTransferOperation(
+                TransferOperationKind::kWrite, "upsert_batch",
+                sum_value_sizes(values), latency_us);
+        });
+    return to_py_ret(result);
 }
 
 // --- End Upsert implementations ---
@@ -2940,7 +3165,20 @@ int RealClient::upsert_batch(const std::vector<std::string> &keys,
 std::vector<int64_t> RealClient::batch_get_into(
     const std::vector<std::string> &keys, const std::vector<void *> &buffers,
     const std::vector<size_t> &sizes) {
-    auto internal_results = batch_get_into_internal(keys, buffers, sizes);
+    auto internal_results =
+        execute_timed_operation<std::vector<tl::expected<int64_t, ErrorCode>>>(
+            [&]() { return batch_get_into_internal(keys, buffers, sizes); },
+            [](const auto &) { return true; },
+            [&](uint64_t latency_us, const auto &ret) {
+                std::vector<int64_t> py_results;
+                py_results.reserve(ret.size());
+                for (const auto &item : ret) {
+                    py_results.push_back(to_py_ret(item));
+                }
+                client_->ObserveTransferOperation(
+                    TransferOperationKind::kRead, "batch_get_into",
+                    sum_positive_results(py_results), latency_us);
+            });
     std::vector<int64_t> results;
     results.reserve(internal_results.size());
 
@@ -3344,6 +3582,7 @@ int RealClient::put_from_with_metadata(const std::string &key, void *buffer,
                                        void *metadata_buffer, size_t size,
                                        size_t metadata_size,
                                        const ReplicateConfig &config) {
+    const auto start_time = std::chrono::steady_clock::now();
     // NOTE: The buffer address must be previously registered with
     // register_buffer() for zero-copy RDMA operations to work correctly
     if (config.prefer_alloc_in_same_node) {
@@ -3386,6 +3625,10 @@ int RealClient::put_from_with_metadata(const std::string &key, void *buffer,
                    << toString(put_result.error());
         return -toInt(put_result.error());
     }
+
+    client_->ObserveTransferOperation(
+        TransferOperationKind::kWrite, "put_from_with_metadata",
+        size + metadata_size, elapsed_us_since(start_time));
     return 0;
 }
 
@@ -3394,10 +3637,24 @@ std::vector<int> RealClient::batch_put_from_multi_buffers(
     const std::vector<std::vector<void *>> &all_buffers,
     const std::vector<std::vector<size_t>> &sizes,
     const ReplicateConfig &config) {
-    auto start = std::chrono::steady_clock::now();
-
     auto internal_results =
-        batch_put_from_multi_buffers_internal(keys, all_buffers, sizes, config);
+        execute_timed_operation<std::vector<tl::expected<void, ErrorCode>>>(
+            [&]() {
+                return batch_put_from_multi_buffers_internal(keys, all_buffers,
+                                                             sizes, config);
+            },
+            [](const auto &) { return true; },
+            [&](uint64_t latency_us, const auto &ret) {
+                std::vector<int> py_results;
+                py_results.reserve(ret.size());
+                for (const auto &item : ret) {
+                    py_results.push_back(to_py_ret(item));
+                }
+                client_->ObserveTransferOperation(
+                    TransferOperationKind::kWrite,
+                    "batch_put_from_multi_buffers",
+                    sum_successful_nested_sizes(py_results, sizes), latency_us);
+            });
     std::vector<int> results;
     results.reserve(internal_results.size());
 
@@ -3405,10 +3662,6 @@ std::vector<int> RealClient::batch_put_from_multi_buffers(
         results.push_back(to_py_ret(result));
     }
 
-    auto duration_call = std::chrono::duration_cast<std::chrono::microseconds>(
-        std::chrono::steady_clock::now() - start);
-    VLOG(1) << "batch_put_from_multi_buffers: " << duration_call.count()
-            << " us";
     return results;
 }
 
@@ -3454,19 +3707,30 @@ std::vector<int> RealClient::batch_get_into_multi_buffers(
     const std::vector<std::vector<void *>> &all_buffers,
     const std::vector<std::vector<size_t>> &all_sizes,
     bool prefer_alloc_in_same_node) {
-    auto start = std::chrono::steady_clock::now();
-    auto internal_results = batch_get_into_multi_buffers_internal(
-        keys, all_buffers, all_sizes, prefer_alloc_in_same_node);
+    auto internal_results =
+        execute_timed_operation<std::vector<tl::expected<int64_t, ErrorCode>>>(
+            [&]() {
+                return batch_get_into_multi_buffers_internal(
+                    keys, all_buffers, all_sizes, prefer_alloc_in_same_node);
+            },
+            [](const auto &) { return true; },
+            [&](uint64_t latency_us, const auto &ret) {
+                std::vector<int> py_results;
+                py_results.reserve(ret.size());
+                for (const auto &item : ret) {
+                    py_results.push_back(to_py_ret(item));
+                }
+                client_->ObserveTransferOperation(
+                    TransferOperationKind::kRead,
+                    "batch_get_into_multi_buffers",
+                    sum_positive_results(py_results), latency_us);
+            });
     std::vector<int> results;
     results.reserve(internal_results.size());
 
     for (const auto &result : internal_results) {
         results.push_back(to_py_ret(result));
     }
-    auto duration_call = std::chrono::duration_cast<std::chrono::microseconds>(
-        std::chrono::steady_clock::now() - start);
-    VLOG(1) << "batch_get_into_multi_buffers: " << duration_call.count()
-            << " us";
     return results;
 }
 
@@ -3892,6 +4156,21 @@ RealClient::batch_get_replica_desc(const std::vector<std::string> &keys) {
     return replica_map;
 }
 
+std::vector<std::string> RealClient::batch_replica_clear(
+    const std::vector<std::string> &keys, const std::string &segment_name) {
+    if (!client_) {
+        LOG(ERROR) << "batch_replica_clear: client not initialized";
+        return {};
+    }
+    auto result =
+        client_->BatchReplicaClear(keys, client_->getClientId(), segment_name);
+    if (result) {
+        return result.value();
+    }
+    LOG(ERROR) << "batch_replica_clear failed: " << toString(result.error());
+    return {};
+}
+
 tl::expected<UUID, ErrorCode> RealClient::create_copy_task(
     const std::string &key, const std::vector<std::string> &targets) {
     return client_->CreateCopyTask(key, targets);
@@ -3987,6 +4266,18 @@ ClientRequester::ClientRequester() {
         pool_conf.client_config.socket_config =
             coro_io::ib_socket_t::config_t{};
     }
+    // Configure reasonable retry limits for SSD offload RPC connections.
+    // - connect_retry_count: Maximum connection retry attempts (default: 3)
+    // - reconnect_wait_time: Wait time between retries (default: 1000ms)
+    // - host_alive_detect_duration: Duration for background alive detection.
+    //   Set to 0 to disable infinite background reconnection attempts when
+    //   a Store node goes down. This prevents continuous "Connection refused"
+    //   logs. When Master cleans up stale local_disk replicas (via
+    //   CleanupStaleHandles), new requests won't route to dead nodes anyway.
+    pool_conf.connect_retry_count = 3;
+    pool_conf.reconnect_wait_time = std::chrono::milliseconds{1000};
+    pool_conf.host_alive_detect_duration = std::chrono::milliseconds{0};
+
     client_pools_ =
         std::make_shared<coro_io::client_pools<coro_rpc::coro_rpc_client>>(
             pool_conf);

mooncake-store/src/real_client_main.cpp

@@ -91,6 +91,10 @@ int main(int argc, char *argv[]) {
     mooncake::ResourceTracker::getInstance();
 
     gflags::ParseCommandLineFlags(&argc, &argv, true);
+    if (!FLAGS_log_dir.empty()) {
+        google::InitGoogleLogging(argv[0]);
+    }
+
     size_t global_segment_size = string_to_byte_size(FLAGS_global_segment_size);
 #ifdef USE_ASCEND_DIRECT
     // just set to true, does not affect GPU process.
@@ -113,10 +117,11 @@ int main(int argc, char *argv[]) {
         return -1;
     }
 
-    coro_rpc::coro_rpc_server server(FLAGS_threads, FLAGS_port, "127.0.0.1");
+    coro_rpc::coro_rpc_server server(FLAGS_threads, FLAGS_port, FLAGS_host);
     RegisterClientRpcService(server, *client_inst);
 
-    LOG(INFO) << "Starting real client service on 127.0.0.1:" << FLAGS_port;
+    LOG(INFO) << "Starting real client service on " << FLAGS_host << ":"
+              << FLAGS_port;
 
     return server.start();
 }

mooncake-store/src/segment.cpp

@@ -273,6 +273,15 @@ ErrorCode ScopedSegmentAccess::GetClientSegments(
     return ErrorCode::OK;
 }
 
+void ScopedSegmentAccess::UnmountLocalDiskSegment(const UUID& client_id) {
+    auto it = segment_manager_->client_local_disk_segment_.find(client_id);
+    if (it != segment_manager_->client_local_disk_segment_.end()) {
+        segment_manager_->client_local_disk_segment_.erase(it);
+        LOG(INFO) << "client_id=" << client_id
+                  << ", action=unmount_local_disk_segment";
+    }
+}
+
 ErrorCode ScopedSegmentAccess::GetAllSegments(
     std::vector<std::string>& all_segments) {
     all_segments.clear();

mooncake-store/src/storage_backend.cpp

@@ -49,8 +49,9 @@ FilePerKeyConfig FilePerKeyConfig::FromEnvironment() {
 
     config.fsdir = GetEnvStringOr("MOONCAKE_OFFLOAD_FSDIR", config.fsdir);
 
-    config.enable_eviction =
-        GetEnvOr<bool>("ENABLE_EVICTION", config.enable_eviction);
+    config.enable_eviction = GetEnvOr<bool>(
+        "MOONCAKE_OFFLOAD_ENABLE_EVICTION",
+        GetEnvOr<bool>("ENABLE_EVICTION", config.enable_eviction));
 
     return config;
 }
@@ -64,11 +65,14 @@ BucketBackendConfig BucketBackendConfig::FromEnvironment() {
     config.bucket_size_limit = GetEnvOr<int64_t>(
         "MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES", config.bucket_size_limit);
 
-    config.max_total_size = GetEnvOr<int64_t>("MOONCAKE_BUCKET_MAX_TOTAL_SIZE",
-                                              config.max_total_size);
+    config.max_total_size =
+        GetEnvOr<int64_t>("MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE",
+                          GetEnvOr<int64_t>("MOONCAKE_BUCKET_MAX_TOTAL_SIZE",
+                                            config.max_total_size));
 
-    const auto policy_str =
-        GetEnvStringOr("MOONCAKE_BUCKET_EVICTION_POLICY", "none");
+    const auto policy_str = GetEnvStringOr(
+        "MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY",
+        GetEnvStringOr("MOONCAKE_BUCKET_EVICTION_POLICY", "none"));
     if (policy_str == "fifo") {
         config.eviction_policy = BucketEvictionPolicy::FIFO;
     } else if (policy_str == "lru") {
@@ -1718,6 +1722,18 @@ tl::expected<void, ErrorCode> BucketStorageBackend::Init() {
                       << orphaned_space_freed << " bytes";
         }
 
+        // When max_total_size is not explicitly set (<= 0), default to 90% of
+        // the physical disk capacity to match FilePerKey backend behavior.
+        if (bucket_backend_config_.max_total_size <= 0) {
+            constexpr double kDefaultQuotaPercentage = 0.9;
+            const auto space_info = fs::space(storage_path_);
+            bucket_backend_config_.max_total_size = static_cast<int64_t>(
+                space_info.capacity * kDefaultQuotaPercentage);
+            LOG(INFO) << "Bucket backend max_total_size not set; using "
+                      << kDefaultQuotaPercentage * 100 << "% of disk capacity: "
+                      << bucket_backend_config_.max_total_size << " bytes";
+        }
+
         bucket_id_generator_.emplace(max_bucket_id);
         if (max_bucket_id == BucketIdGenerator::INIT_NEW_START_ID) {
             LOG(INFO) << "Initialized BucketIdGenerator with fresh start. "
@@ -2187,8 +2203,7 @@ BucketStorageBackend::PendingEviction BucketStorageBackend::PrepareEviction(
     int64_t required_size) {
     PendingEviction result;
 
-    if (bucket_backend_config_.eviction_policy == BucketEvictionPolicy::NONE ||
-        bucket_backend_config_.max_total_size <= 0) {
+    if (bucket_backend_config_.eviction_policy == BucketEvictionPolicy::NONE) {
         return result;
     }
 

mooncake-store/src/transfer_task.cpp

@@ -4,6 +4,8 @@
 
 #include <algorithm>
 #include <cstdlib>
+#include <sstream>
+#include <vector>
 #include "transfer_engine.h"
 #include "transport/transport.h"
 
@@ -223,9 +225,13 @@ bool TransferEngineOperationState::is_completed() {
 }
 
 void TransferEngineOperationState::check_task_status() {
-    // Check all transfers in the batch
-    bool all_completed = true;
-    bool has_failure = false;
+    // Check all transfers in the batch.
+    // Wait for ALL tasks to reach a terminal state before setting the result,
+    // even if some have already failed. This prevents the caller from seeing
+    // "completed" while background transfers are still in progress, which
+    // could cause issues when freeBatchID is called in the destructor.
+    bool all_terminated = true;
+    std::vector<size_t> failed_task_ids;
 
     for (size_t i = 0; i < batch_size_; ++i) {
         TransferStatus status;
@@ -240,38 +246,45 @@ void TransferEngineOperationState::check_task_status() {
 
         switch (status.s) {
             case TransferStatusEnum::COMPLETED:
-                // This transfer is done, continue checking others
+                // This transfer is done successfully
                 break;
             case TransferStatusEnum::FAILED:
             case TransferStatusEnum::CANCELED:
             case TransferStatusEnum::INVALID:
 #ifndef USE_ASCEND_DIRECT
-                LOG(ERROR) << "Transfer failed for batch " << batch_id_
-                           << " task " << i << " with status "
-                           << static_cast<int>(status.s);
+                VLOG(1) << "Transfer failed for batch " << batch_id_ << " task "
+                        << i << " with status " << static_cast<int>(status.s);
 #endif
-                has_failure = true;
+                failed_task_ids.push_back(i);
                 break;
             default:
-                // Transfer is still pending (PENDING, RUNNING, etc.)
-                all_completed = false;
+                // Transfer is still in progress (WAITING, PENDING, etc.)
+                all_terminated = false;
                 break;
         }
     }
 
-    if (has_failure) {
-        VLOG(1) << "Setting batch " << batch_id_
-                << " result to TRANSFER_FAIL due to task failures";
-        set_result_internal(ErrorCode::TRANSFER_FAIL);
+    if (!all_terminated) {
+        // Some tasks are still in progress; wait for next poll iteration.
+        // Do NOT set result yet, even if some tasks have already failed.
         return;
     }
 
-    if (all_completed) {
-        set_result_internal(ErrorCode::OK);
-        return;
+    // All tasks have reached a terminal state.
+    ErrorCode ec = ErrorCode::OK;
+    if (!failed_task_ids.empty()) {
+        std::ostringstream oss;
+        for (size_t j = 0; j < failed_task_ids.size(); ++j) {
+            if (j > 0) oss << ", ";
+            oss << failed_task_ids[j];
+        }
+        LOG(ERROR) << "Batch " << batch_id_
+                   << " completed with task failures: task_ids=[" << oss.str()
+                   << "]";
+        ec = ErrorCode::TRANSFER_FAIL;
     }
 
-    return;
+    set_result_internal(ec);
 }
 
 void TransferEngineOperationState::set_result_internal(ErrorCode error_code) {
@@ -294,7 +307,8 @@ void TransferEngineOperationState::wait_for_completion() {
         return;
     }
 
-    constexpr int64_t timeout_seconds = 60;
+    // 60 seconds
+    constexpr int64_t timeout_milliseconds = 60 * 1000;
 
 #ifdef USE_EVENT_DRIVEN_COMPLETION
     VLOG(1) << "Waiting for transfer engine completion for batch " << batch_id_;
@@ -314,10 +328,18 @@ void TransferEngineOperationState::wait_for_completion() {
         // lock. Under the mutex, relaxed is sufficient; the mutex acquire
         // orders prior writes.
         std::unique_lock<std::mutex> lock(batch_desc.completion_mutex);
-        completed = batch_desc.completion_cv.wait_for(
-            lock, std::chrono::seconds(timeout_seconds), [&batch_desc] {
-                return batch_desc.is_finished.load(std::memory_order_relaxed);
-            });
+        const int64_t elapsed_milliseconds =
+            getCurrentTimeInMilli() - start_ts_;
+        if (elapsed_milliseconds < timeout_milliseconds) {
+            completed = batch_desc.completion_cv.wait_for(
+                lock,
+                std::chrono::milliseconds(timeout_milliseconds -
+                                          elapsed_milliseconds),
+                [&batch_desc] {
+                    return batch_desc.is_finished.load(
+                        std::memory_order_relaxed);
+                });
+        }
     }  // Explicitly release completion_mutex before acquiring mutex_
 
     // Once completion is observed, read failure flag.
@@ -338,20 +360,18 @@ void TransferEngineOperationState::wait_for_completion() {
         VLOG(1) << "Transfer engine operation completed for batch " << batch_id_
                 << " with result: " << static_cast<int>(error_code);
     } else {
-        LOG(ERROR) << "Failed to complete transfers after " << timeout_seconds
-                   << " seconds for batch " << batch_id_;
+        LOG(ERROR) << "Failed to complete transfers after "
+                   << timeout_milliseconds << " milliseconds for batch "
+                   << batch_id_;
     }
 #else
     VLOG(1) << "Starting transfer engine polling for batch " << batch_id_;
 
-    constexpr int64_t kOneSecondInNano = 1000 * 1000 * 1000;
-    const int64_t start_ts = getCurrentTimeInNano();
-
     while (true) {
-        if (getCurrentTimeInNano() - start_ts >
-            timeout_seconds * kOneSecondInNano) {
+        if (getCurrentTimeInMilli() - start_ts_ > timeout_milliseconds) {
             LOG(ERROR) << "Failed to complete transfers after "
-                       << timeout_seconds << " seconds for batch " << batch_id_;
+                       << timeout_milliseconds << " milliseconds for batch "
+                       << batch_id_;
             set_result_internal(ErrorCode::TRANSFER_FAIL);
             return;
         }
@@ -409,10 +429,17 @@ TransferSubmitter::TransferSubmitter(TransferEngine& engine,
       memcpy_pool_(std::make_unique<MemcpyWorkerPool>()),
       fileread_pool_(std::make_unique<FilereadWorkerPool>(backend)),
       transfer_metric_(transfer_metric) {
-    // Read MC_STORE_MEMCPY environment variable, default to false (disabled)
+    // Read MC_STORE_MEMCPY environment variable.
+    // When not set, auto-detect based on transport type:
+    //   - TCP-only environment: enable memcpy (avoids TCP loopback overhead)
+    //   - RDMA/other transports: disable memcpy (RDMA is more efficient)
     const char* env_value = std::getenv("MC_STORE_MEMCPY");
     if (env_value == nullptr) {
-        memcpy_enabled_ = false;  // Default: disabled
+        memcpy_enabled_ = engine_.isTcpOnly();
+        LOG(INFO) << "MC_STORE_MEMCPY not set, auto-detected: "
+                  << (memcpy_enabled_ ? "TCP-only environment, memcpy enabled"
+                                      : "non-TCP transport available, memcpy "
+                                        "disabled");
     } else {
         std::string env_str(env_value);
         // Convert to lowercase for case-insensitive comparison

mooncake-store/tests/CMakeLists.txt

@@ -39,6 +39,7 @@ add_store_test(eviction_strategy_test eviction_strategy_test.cpp)
 add_store_test(master_service_test master_service_test.cpp)
 add_store_test(batch_remove_test batch_remove_test.cpp)
 add_store_test(master_service_ssd_test master_service_ssd_test.cpp)
+add_store_test(offload_on_evict_test offload_on_evict_test.cpp)
 add_store_test(master_service_ssd_test_for_snapshot
                ha/snapshot/master_service_ssd_test_for_snapshot.cpp)
 add_store_test(client_integration_test client_integration_test.cpp)

mooncake-store/tests/client_metrics_test.cpp

@@ -45,6 +45,8 @@ TEST_F(ClientMetricsTest, TransferMetricsSummaryTest) {
     // Check byte formatting
     EXPECT_TRUE(summary.find("Total Read: 1.00 KB") != std::string::npos);
     EXPECT_TRUE(summary.find("Total Write: 2.00 MB") != std::string::npos);
+    EXPECT_TRUE(summary.find("Average Read Throughput:") != std::string::npos);
+    EXPECT_TRUE(summary.find("Average Write Throughput:") != std::string::npos);
 
     // Check latency summaries
     EXPECT_TRUE(summary.find("Get: count=3") != std::string::npos);
@@ -109,15 +111,23 @@ TEST_F(ClientMetricsTest, ClientMetricsSummaryTest) {
     std::array<std::string, 1> exist_key_label = {"ExistKey"};
     metrics.master_client_metric.rpc_count.inc(exist_key_label);
     metrics.master_client_metric.rpc_latency.observe(exist_key_label, 180);
+    metrics.ObserveTransferOperation(TransferOperationKind::kRead, "get_buffer",
+                                     2 * 1024, 220);
+    metrics.ObserveTransferOperation(TransferOperationKind::kWrite, "put_batch",
+                                     4 * 1024, 420);
 
     std::string summary = metrics.summary_metrics();
 
-    // Should contain both transfer and RPC metrics
+    // Should contain transfer, RPC, and interface metrics
     EXPECT_TRUE(summary.find("Transfer Metrics Summary") != std::string::npos);
     EXPECT_TRUE(summary.find("RPC Metrics Summary") != std::string::npos);
+    EXPECT_TRUE(summary.find("Interface Operation Metrics Summary") !=
+                std::string::npos);
     EXPECT_TRUE(summary.find("Total Read: 5.00 MB") != std::string::npos);
     EXPECT_TRUE(summary.find("Total Write: 10.00 MB") != std::string::npos);
     EXPECT_TRUE(summary.find("ExistKey: count=1") != std::string::npos);
+    EXPECT_TRUE(summary.find("get_buffer: count=1") != std::string::npos);
+    EXPECT_TRUE(summary.find("put_batch: count=1") != std::string::npos);
 
     std::cout << "Full Client Metrics Summary:\n" << summary << std::endl;
 }
@@ -175,6 +185,33 @@ TEST_F(ClientMetricsTest, CompareWithSerializedMetrics) {
                 summary.find("No data") != std::string::npos);
 }
 
+TEST_F(ClientMetricsTest, BandwidthSummaryRespectsEnvFlag) {
+    setenv("MC_STORE_CLIENT_METRIC_BANDWIDTH", "0", 1);
+    auto metrics = ClientMetric::Create();
+    ASSERT_NE(metrics, nullptr);
+
+    metrics->transfer_metric.total_read_bytes.inc(1024);
+    std::string summary = metrics->summary_metrics();
+    EXPECT_TRUE(summary.find("Average Read Throughput:") == std::string::npos);
+
+    unsetenv("MC_STORE_CLIENT_METRIC_BANDWIDTH");
+}
+
+TEST_F(ClientMetricsTest, SummaryCanOmitMasterRpcMetrics) {
+    auto metrics = ClientMetric::Create({}, false);
+    ASSERT_NE(metrics, nullptr);
+
+    metrics->ObserveTransferOperation(TransferOperationKind::kRead,
+                                      "get_buffer", 1024, 200);
+    std::string summary = metrics->summary_metrics();
+    std::string serialized;
+    metrics->serialize(serialized);
+
+    EXPECT_TRUE(summary.find("RPC Metrics Summary") == std::string::npos);
+    EXPECT_TRUE(serialized.find("mooncake_client_rpc_count") ==
+                std::string::npos);
+}
+
 TEST_F(ClientMetricsTest, SerializeWithDynamicLabels) {
     auto verify = [](const std::string& str) {
         EXPECT_TRUE(str.find("instance_id=\"12345\"") != std::string::npos);

mooncake-store/tests/file_storage_test.cpp

@@ -26,7 +26,8 @@ class FileStorageTest : public ::testing::Test {
         FLAGS_logtostderr = true;
         UnsetEnv("MOONCAKE_OFFLOAD_FILE_STORAGE_PATH");
         UnsetEnv("MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES");
-        UnsetEnv("MOONCAKE_OFFLOAD_BUCKET_ITERATOR_KEYS_LIMIT");
+        UnsetEnv("MOONCAKE_OFFLOAD_SCANMETA_ITERATOR_KEYS_LIMIT");
+        UnsetEnv("MOONCAKE_SCANMETA_ITERATOR_KEYS_LIMIT");
         UnsetEnv("MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT");
         UnsetEnv("MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES");
         UnsetEnv("MOONCAKE_OFFLOAD_TOTAL_KEYS_LIMIT");

mooncake-store/tests/offload_on_evict_test.cpp

@@ -0,0 +1,388 @@
+#include "master_service.h"
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+
+#include <atomic>
+#include <chrono>
+#include <memory>
+#include <string>
+#include <thread>
+#include <vector>
+
+#include "types.h"
+
+namespace mooncake::test {
+
+class OffloadOnEvictTest : public ::testing::Test {
+   protected:
+    void SetUp() override {
+        google::InitGoogleLogging("OffloadOnEvictTest");
+        FLAGS_logtostderr = true;
+    }
+
+    void TearDown() override { google::ShutdownGoogleLogging(); }
+
+    static constexpr size_t kDefaultSegmentBase = 0x300000000;
+
+    Segment MakeSegment(std::string name, size_t base, size_t size) const {
+        Segment segment;
+        segment.id = generate_uuid();
+        segment.name = std::move(name);
+        segment.base = base;
+        segment.size = size;
+        segment.te_endpoint = segment.name;
+        return segment;
+    }
+
+    struct MountedSegmentContext {
+        UUID segment_id;
+        UUID client_id;
+    };
+
+    MountedSegmentContext PrepareSegment(MasterService& service,
+                                         std::string name, size_t base,
+                                         size_t size) const {
+        Segment segment = MakeSegment(std::move(name), base, size);
+        UUID client_id = generate_uuid();
+        auto mount_result = service.MountSegment(segment, client_id);
+        EXPECT_TRUE(mount_result.has_value());
+        return {.segment_id = segment.id, .client_id = client_id};
+    }
+
+    // Put an object and complete it.
+    void PutObject(MasterService& service, const UUID& client_id,
+                   const std::string& key, size_t size = 1024) {
+        ReplicateConfig config;
+        config.replica_num = 1;
+        auto put_start = service.PutStart(client_id, key, size, config);
+        ASSERT_TRUE(put_start.has_value()) << "PutStart failed for key=" << key;
+        auto put_end = service.PutEnd(client_id, key, ReplicaType::MEMORY);
+        ASSERT_TRUE(put_end.has_value()) << "PutEnd failed for key=" << key;
+    }
+
+    // Drain the offload queue via OffloadObjectHeartbeat.
+    std::unordered_map<std::string, int64_t> DrainOffloadQueue(
+        MasterService& service, const UUID& client_id) {
+        auto res = service.OffloadObjectHeartbeat(client_id, true);
+        if (!res) {
+            return {};
+        }
+        return std::move(res.value());
+    }
+
+    template <typename Predicate>
+    void WaitUntil(
+        Predicate&& predicate,
+        std::chrono::milliseconds timeout = std::chrono::milliseconds(4000),
+        std::chrono::milliseconds interval =
+            std::chrono::milliseconds(50)) const {
+        const auto deadline = std::chrono::steady_clock::now() + timeout;
+        while (std::chrono::steady_clock::now() < deadline) {
+            if (predicate()) {
+                return;
+            }
+            std::this_thread::sleep_for(interval);
+        }
+        EXPECT_TRUE(predicate());
+    }
+
+    // Fill a segment until PutStart fails, triggering eviction.
+    // Returns the number of successful puts.
+    int FillSegmentUntilEviction(MasterService& service, const UUID& client_id,
+                                 const std::string& key_prefix,
+                                 size_t object_size, int max_puts) {
+        int success_puts = 0;
+        for (int i = 0; i < max_puts; ++i) {
+            std::string key = key_prefix + std::to_string(i);
+            ReplicateConfig config;
+            config.replica_num = 1;
+            auto result = service.PutStart(client_id, key, object_size, config);
+            if (result.has_value()) {
+                auto end = service.PutEnd(client_id, key, ReplicaType::MEMORY);
+                EXPECT_TRUE(end.has_value());
+                success_puts++;
+            } else {
+                // Wait for eviction to process
+                std::this_thread::sleep_for(std::chrono::milliseconds(50));
+            }
+        }
+        return success_puts;
+    }
+};
+
+// =============================================================================
+// Combo A: Default config (offload at PutEnd)
+// =============================================================================
+
+TEST_F(OffloadOnEvictTest, ComboA_OffloadAtPutEnd) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.default_kv_lease_ttl = 2000;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto ctx =
+        PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
+
+    // Mount local disk segment with offloading ENABLED
+    auto mount_ld = service->MountLocalDiskSegment(ctx.client_id, true);
+    ASSERT_TRUE(mount_ld.has_value());
+
+    // Put objects
+    PutObject(*service, ctx.client_id, "key_a1");
+    PutObject(*service, ctx.client_id, "key_a2");
+    PutObject(*service, ctx.client_id, "key_a3");
+
+    // Default mode: PutEnd pushes to offload queue immediately
+    auto queued = DrainOffloadQueue(*service, ctx.client_id);
+    EXPECT_EQ(queued.size(), 3u)
+        << "Default: all 3 objects should be in offload queue after PutEnd";
+    EXPECT_TRUE(queued.count("key_a1"));
+    EXPECT_TRUE(queued.count("key_a2"));
+    EXPECT_TRUE(queued.count("key_a3"));
+
+    service->RemoveAll();
+}
+
+TEST_F(OffloadOnEvictTest, ComboA_EvictionWorks) {
+    // Regression: eviction still works in default mode
+    const uint64_t kv_lease_ttl = 2000;
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.default_kv_lease_ttl = kv_lease_ttl;
+    auto service = std::make_unique<MasterService>(config);
+
+    // Large segment: can hold ~16K objects of 15KB
+    constexpr size_t seg_size = 1024 * 1024 * 16 * 15;
+    constexpr size_t object_size = 1024 * 15;
+    auto ctx =
+        PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
+
+    // Put more objects than the segment can hold
+    int success_puts = FillSegmentUntilEviction(
+        *service, ctx.client_id, "evict_a_", object_size, 1024 * 16 + 50);
+    EXPECT_GT(success_puts, 1024 * 16)
+        << "Default: eviction should allow more puts than capacity";
+
+    std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
+    service->RemoveAll();
+}
+
+// =============================================================================
+// Combo B: offload_on_evict=true (offload on evict, no force-evict)
+// =============================================================================
+
+TEST_F(OffloadOnEvictTest, ComboB_PutEndSkipsOffloadQueue) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.offload_on_evict = true;
+    config.default_kv_lease_ttl = 2000;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto ctx =
+        PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
+    auto mount_ld = service->MountLocalDiskSegment(ctx.client_id, true);
+    ASSERT_TRUE(mount_ld.has_value());
+
+    PutObject(*service, ctx.client_id, "key_b1");
+    PutObject(*service, ctx.client_id, "key_b2");
+    PutObject(*service, ctx.client_id, "key_b3");
+
+    // Offload-on-evict: PutEnd should NOT push to offload queue
+    auto queued = DrainOffloadQueue(*service, ctx.client_id);
+    EXPECT_EQ(queued.size(), 0u)
+        << "Offload-on-evict: queue should be empty after PutEnd";
+
+    service->RemoveAll();
+}
+
+TEST_F(OffloadOnEvictTest, ComboB_EvictionTriggersOffload) {
+    const uint64_t kv_lease_ttl = 2000;
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.offload_on_evict = true;
+    config.default_kv_lease_ttl = kv_lease_ttl;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16 * 15;
+    constexpr size_t object_size = 1024 * 15;
+    auto ctx =
+        PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
+    auto mount_ld = service->MountLocalDiskSegment(ctx.client_id, true);
+    ASSERT_TRUE(mount_ld.has_value());
+
+    // Fill segment to trigger eviction
+    bool eviction_triggered = false;
+    int success_puts = 0;
+    for (int i = 0; i < 1024 * 16 + 50; ++i) {
+        std::string key = "evict_b_" + std::to_string(i);
+        ReplicateConfig config;
+        config.replica_num = 1;
+        auto result =
+            service->PutStart(ctx.client_id, key, object_size, config);
+        if (result.has_value()) {
+            auto end = service->PutEnd(ctx.client_id, key, ReplicaType::MEMORY);
+            ASSERT_TRUE(end.has_value());
+            success_puts++;
+        } else {
+            eviction_triggered = true;
+            std::this_thread::sleep_for(std::chrono::milliseconds(50));
+        }
+    }
+
+    EXPECT_TRUE(eviction_triggered)
+        << "Eviction should trigger when segment fills up";
+
+    // Offload-on-evict: eviction should push objects to offload queue
+    auto queued = DrainOffloadQueue(*service, ctx.client_id);
+    EXPECT_GT(queued.size(), 0u)
+        << "Offload-on-evict: eviction should push to offload queue";
+
+    std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
+    service->RemoveAll();
+}
+
+TEST_F(OffloadOnEvictTest, ComboB_NoFallbackWithoutForceEvict) {
+    // Without force_evict AND without a LocalDiskSegment, offload queue push
+    // fails and eviction does NOT force-delete MEMORY (data-preserving).
+    // The segment fills and subsequent puts fail — this is the safe default.
+    const uint64_t kv_lease_ttl = 2000;
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.offload_on_evict = true;
+    config.default_kv_lease_ttl = kv_lease_ttl;
+    auto service = std::make_unique<MasterService>(config);
+
+    // NO local disk segment mounted — PushOffloadingQueue will fail
+    constexpr size_t seg_size = 1024 * 1024 * 16 * 15;
+    constexpr size_t object_size = 1024 * 15;
+    auto ctx =
+        PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
+
+    // Without force_evict, push failures mean DRAM cannot be freed,
+    // so we can only put up to segment capacity (no overflow).
+    int success_puts = FillSegmentUntilEviction(
+        *service, ctx.client_id, "evict_b2_", object_size, 1024 * 16 + 50);
+    EXPECT_LE(success_puts, 1024 * 16)
+        << "Without force_evict, segment should fill and stay full";
+
+    std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
+    service->RemoveAll();
+}
+
+// =============================================================================
+// Combo C: offload_on_evict=true + offload_force_evict=true
+// =============================================================================
+
+TEST_F(OffloadOnEvictTest, ComboC_PutEndSkipsOffloadQueue) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.offload_on_evict = true;
+    config.offload_force_evict = true;
+    config.default_kv_lease_ttl = 2000;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto ctx =
+        PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
+    auto mount_ld = service->MountLocalDiskSegment(ctx.client_id, true);
+    ASSERT_TRUE(mount_ld.has_value());
+
+    PutObject(*service, ctx.client_id, "key_c1");
+    PutObject(*service, ctx.client_id, "key_c2");
+
+    // Same as Combo B: PutEnd should skip offload queue
+    auto queued = DrainOffloadQueue(*service, ctx.client_id);
+    EXPECT_EQ(queued.size(), 0u)
+        << "Combo C: offload queue should be empty after PutEnd";
+
+    service->RemoveAll();
+}
+
+TEST_F(OffloadOnEvictTest, ComboC_EvictionWithForceEvict) {
+    const uint64_t kv_lease_ttl = 2000;
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.offload_on_evict = true;
+    config.offload_force_evict = true;
+    config.default_kv_lease_ttl = kv_lease_ttl;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16 * 15;
+    constexpr size_t object_size = 1024 * 15;
+    auto ctx =
+        PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
+    auto mount_ld = service->MountLocalDiskSegment(ctx.client_id, true);
+    ASSERT_TRUE(mount_ld.has_value());
+
+    // With force-evict, eviction should work effectively.
+    // Note: without a real FileStorage heartbeat, offloaded objects' refcnt
+    // never decreases, so DRAM isn't fully freed beyond what direct eviction
+    // allows. We verify eviction doesn't deadlock (can fill to capacity).
+    int success_puts = FillSegmentUntilEviction(
+        *service, ctx.client_id, "evict_c_", object_size, 1024 * 16 + 50);
+    EXPECT_GE(success_puts, 1024 * 16)
+        << "Combo C: eviction should work with force-evict enabled";
+
+    std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
+    service->RemoveAll();
+}
+
+// =============================================================================
+// Combo D: offload_force_evict=true only (should be no-op without on_evict)
+// =============================================================================
+
+TEST_F(OffloadOnEvictTest, ComboD_ForceEvictAloneIsIgnored) {
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.offload_force_evict = true;  // on_evict is false → force is ignored
+    config.default_kv_lease_ttl = 2000;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16;
+    auto ctx =
+        PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
+    auto mount_ld = service->MountLocalDiskSegment(ctx.client_id, true);
+    ASSERT_TRUE(mount_ld.has_value());
+
+    // Should behave like Combo A (default: offload at PutEnd)
+    PutObject(*service, ctx.client_id, "key_d1");
+    PutObject(*service, ctx.client_id, "key_d2");
+
+    auto queued = DrainOffloadQueue(*service, ctx.client_id);
+    EXPECT_EQ(queued.size(), 2u)
+        << "Combo D: FORCE_EVICT alone should not change default behavior";
+
+    service->RemoveAll();
+}
+
+TEST_F(OffloadOnEvictTest, ComboD_EvictionWorks) {
+    const uint64_t kv_lease_ttl = 2000;
+    MasterServiceConfig config;
+    config.enable_offload = true;
+    config.offload_force_evict = true;  // on_evict is false → force is ignored
+    config.default_kv_lease_ttl = kv_lease_ttl;
+    auto service = std::make_unique<MasterService>(config);
+
+    constexpr size_t seg_size = 1024 * 1024 * 16 * 15;
+    constexpr size_t object_size = 1024 * 15;
+    auto ctx =
+        PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
+
+    int success_puts = FillSegmentUntilEviction(
+        *service, ctx.client_id, "evict_d_", object_size, 1024 * 16 + 50);
+    EXPECT_GT(success_puts, 1024 * 16)
+        << "Combo D: eviction should work normally";
+
+    std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
+    service->RemoveAll();
+}
+
+}  // namespace mooncake::test
+
+int main(int argc, char** argv) {
+    ::testing::InitGoogleTest(&argc, argv);
+    return RUN_ALL_TESTS();
+}

mooncake-transfer-engine/CMakeLists.txt

@@ -9,10 +9,11 @@ if (NOT GLOBAL_CONFIG)
   include(../mooncake-common/FindGLOG.cmake)
   include(../mooncake-common/common.cmake)
 
-  # Build asio_shared library for standalone build
-  # Set compile definitions before add_subdirectory to ensure proper compilation
+  # Build asio_shared in mooncake-common for standalone builds.
+  # Set compile definitions before adding targets that include ASIO headers.
   add_compile_definitions(ASIO_SEPARATE_COMPILATION ASIO_DYN_LINK)
-  add_subdirectory(../mooncake-asio ${CMAKE_CURRENT_BINARY_DIR}/mooncake-asio)
+  add_subdirectory(../mooncake-common/src
+                   ${CMAKE_CURRENT_BINARY_DIR}/mooncake-common-src)
 endif() # GLOBAL_CONFIG
 
 if (USE_ASCEND)

mooncake-transfer-engine/benchmark/CMakeLists.txt

@@ -9,11 +9,12 @@ endif()
 
 file(GLOB TEBENCH_SOURCES "*.cpp")
 add_executable(tebench ${TEBENCH_SOURCES})
-target_link_libraries(tebench PUBLIC transfer_engine tent)
+target_link_libraries(tebench PUBLIC transfer_engine tent_link_group)
 if(USE_CUDA)
   target_link_libraries(tebench PUBLIC CUDA::cudart)
 endif()
 
 # Set RPATH for finding libasio.so at runtime
-set_target_properties(tebench PROPERTIES BUILD_WITH_INSTALL_RPATH TRUE
-                                         INSTALL_RPATH "$ORIGIN/../lib:$ORIGIN/../../mooncake-asio")
+set_target_properties(
+  tebench PROPERTIES BUILD_WITH_INSTALL_RPATH TRUE
+                     INSTALL_RPATH "$ORIGIN/../lib:$ORIGIN/../../mooncake-common")

mooncake-transfer-engine/example/CMakeLists.txt

@@ -1,47 +1,57 @@
 set(WORKSPACE "${CMAKE_CURRENT_SOURCE_DIR}")
 
-if (USE_HIP)
+if(USE_HIP)
   file(GLOB EXAMPLE_SOURCES "*.cpp")
   hipify_files(EXAMPLE_SOURCES)
 
-  file(RELATIVE_PATH EXAMPLE_REL_PATH "${CMAKE_SOURCE_DIR}" "${CMAKE_CURRENT_SOURCE_DIR}")
+  file(RELATIVE_PATH EXAMPLE_REL_PATH "${CMAKE_SOURCE_DIR}"
+       "${CMAKE_CURRENT_SOURCE_DIR}")
   set(WORKSPACE "${CMAKE_BINARY_DIR}/${EXAMPLE_REL_PATH}")
 endif()
 
 add_executable(transfer_engine_bench ${WORKSPACE}/transfer_engine_bench.cpp)
 target_link_libraries(transfer_engine_bench PUBLIC transfer_engine)
-if (USE_TENT)
-    target_link_libraries(transfer_engine_bench PUBLIC tent)
-    target_compile_definitions(transfer_engine_bench PRIVATE USE_TENT)
+if(USE_TENT)
+  target_link_libraries(transfer_engine_bench PUBLIC tent_link_group)
+  target_compile_definitions(transfer_engine_bench PRIVATE USE_TENT)
 endif()
 
-add_executable(transfer_engine_validator ${WORKSPACE}/transfer_engine_validator.cpp)
+add_executable(transfer_engine_validator
+               ${WORKSPACE}/transfer_engine_validator.cpp)
 target_link_libraries(transfer_engine_validator PUBLIC transfer_engine)
 
-add_executable(transfer_engine_bench_with_notify ${WORKSPACE}/transfer_engine_bench_with_notify.cpp)
+add_executable(transfer_engine_bench_with_notify
+               ${WORKSPACE}/transfer_engine_bench_with_notify.cpp)
 target_link_libraries(transfer_engine_bench_with_notify PUBLIC transfer_engine)
 
 add_executable(memory_pool ${WORKSPACE}/memory_pool.cpp)
 target_link_libraries(memory_pool PUBLIC transfer_engine)
 
-if (USE_ASCEND)
-    add_executable(transfer_engine_ascend_one_sided ${WORKSPACE}/transfer_engine_ascend_one_sided.cpp)
-    target_link_libraries(transfer_engine_ascend_one_sided PUBLIC transfer_engine)
+if(USE_ASCEND)
+  add_executable(transfer_engine_ascend_one_sided
+                 ${WORKSPACE}/transfer_engine_ascend_one_sided.cpp)
+  target_link_libraries(transfer_engine_ascend_one_sided PUBLIC transfer_engine)
 
-    add_executable(transfer_engine_ascend_perf ${WORKSPACE}/transfer_engine_ascend_perf.cpp)
-    target_link_libraries(transfer_engine_ascend_perf PUBLIC transfer_engine)
+  add_executable(transfer_engine_ascend_perf
+                 ${WORKSPACE}/transfer_engine_ascend_perf.cpp)
+  target_link_libraries(transfer_engine_ascend_perf PUBLIC transfer_engine)
 endif()
 
-if (USE_ASCEND_DIRECT)
-    add_executable(transfer_engine_ascend_direct_perf ${WORKSPACE}/transfer_engine_ascend_direct_perf.cpp)
-    target_link_libraries(transfer_engine_ascend_direct_perf PUBLIC ascendcl transfer_engine)
+if(USE_ASCEND_DIRECT)
+  add_executable(transfer_engine_ascend_direct_perf
+                 ${WORKSPACE}/transfer_engine_ascend_direct_perf.cpp)
+  target_link_libraries(transfer_engine_ascend_direct_perf
+                        PUBLIC ascendcl transfer_engine)
 endif()
 
-if (USE_ASCEND_HETEROGENEOUS)
-    add_executable(transfer_engine_heterogeneous_ascend_perf_initiator ${WORKSPACE}/transfer_engine_heterogeneous_ascend_perf_initiator.cpp)
-    target_link_libraries(transfer_engine_heterogeneous_ascend_perf_initiator PUBLIC transfer_engine)
+if(USE_ASCEND_HETEROGENEOUS)
+  add_executable(
+    transfer_engine_heterogeneous_ascend_perf_initiator
+    ${WORKSPACE}/transfer_engine_heterogeneous_ascend_perf_initiator.cpp)
+  target_link_libraries(transfer_engine_heterogeneous_ascend_perf_initiator
+                        PUBLIC transfer_engine)
 endif()
 
-if (USE_UBSHMEM)
-    target_link_libraries(transfer_engine_bench PUBLIC transfer_engine)
-endif()
+if(USE_UBSHMEM)
+  target_link_libraries(transfer_engine_bench PUBLIC transfer_engine)
+endif()

mooncake-transfer-engine/example/batch_register_bench.py

@@ -0,0 +1,357 @@
+#!/usr/bin/env python3
+"""
+Batch Memory Registration Benchmark for Mooncake EFA Transport.
+
+Tests registering multiple independent memory blocks (simulating
+multi-tenant or sharded KV cache pools) and measures registration
+time and transfer throughput across blocks.
+
+Usage:
+  # Target node (registers memory blocks):
+  python batch_register_bench.py --mode target \
+      --local_server_name <target_ip>:12345 \
+      --num_blocks 40 --block_size_gb 4 --protocol efa
+
+  # Initiator node (pulls data from random blocks):
+  python batch_register_bench.py --mode initiator \
+      --local_server_name <initiator_ip>:12346 \
+      --target_server_name <target_ip>:12345 \
+      --num_blocks 40 --block_size_gb 4 --protocol efa
+
+Requires: mooncake Python package (pip install -e mooncake-wheel)
+"""
+
+import argparse
+import ctypes
+import ctypes.util
+import json
+import os
+import random
+import signal
+import statistics
+import sys
+import time
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Batch Memory Registration Benchmark"
+    )
+    parser.add_argument(
+        "--mode",
+        choices=["target", "initiator"],
+        required=True,
+        help="Run as target (memory holder) or initiator (puller)",
+    )
+    parser.add_argument(
+        "--local_server_name",
+        required=True,
+        help="Local address, e.g. 172.31.6.162:12345",
+    )
+    parser.add_argument(
+        "--target_server_name",
+        default="",
+        help="Target address (initiator mode only)",
+    )
+    parser.add_argument(
+        "--metadata_server",
+        default="P2PHANDSHAKE",
+        help="Metadata server address (default: P2PHANDSHAKE)",
+    )
+    parser.add_argument(
+        "--protocol", default="efa", help="Transport protocol (default: efa)"
+    )
+    parser.add_argument(
+        "--num_blocks",
+        type=int,
+        default=40,
+        help="Number of memory blocks to register (default: 40)",
+    )
+    parser.add_argument(
+        "--block_size_gb",
+        type=float,
+        default=4.0,
+        help="Size of each memory block in GB (default: 4)",
+    )
+    parser.add_argument(
+        "--transfer_size_mb",
+        type=float,
+        default=368.0,
+        help="Transfer size in MB per iteration (default: 368, ~4K tokens GLM-5.1)",
+    )
+    parser.add_argument(
+        "--iterations",
+        type=int,
+        default=20,
+        help="Number of transfer iterations (default: 20)",
+    )
+    parser.add_argument(
+        "--warmup",
+        type=int,
+        default=5,
+        help="Number of warmup iterations (default: 5)",
+    )
+    parser.add_argument(
+        "--use_batch_api",
+        action="store_true",
+        help="Use batch_register_memory API instead of per-block register_memory",
+    )
+    return parser.parse_args()
+
+
+def allocate_block(size_bytes):
+    """Allocate a single page-aligned memory block using mmap."""
+    libc_name = ctypes.util.find_library("c")
+    libc = ctypes.CDLL(libc_name, use_errno=True)
+
+    PROT_READ = 0x1
+    PROT_WRITE = 0x2
+    MAP_PRIVATE = 0x02
+    MAP_ANONYMOUS = 0x20
+    MAP_HUGETLB = 0x40000
+    MAP_FAILED = ctypes.c_void_p(-1).value
+
+    libc.mmap.restype = ctypes.c_void_p
+    libc.mmap.argtypes = [
+        ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int,
+        ctypes.c_int, ctypes.c_int, ctypes.c_long,
+    ]
+
+    flags = MAP_PRIVATE | MAP_ANONYMOUS
+
+    # Try hugepages first
+    ptr = libc.mmap(None, size_bytes, PROT_READ | PROT_WRITE,
+                    flags | MAP_HUGETLB, -1, 0)
+    if ptr and ptr != MAP_FAILED:
+        return ptr, True
+
+    # Fall back to regular pages
+    ptr = libc.mmap(None, size_bytes, PROT_READ | PROT_WRITE, flags, -1, 0)
+    if not ptr or ptr == MAP_FAILED:
+        raise RuntimeError(
+            f"mmap failed for {size_bytes} bytes: errno={ctypes.get_errno()}"
+        )
+    return ptr, False
+
+
+def run_target(args):
+    """Run as target: allocate blocks, register, and wait."""
+    from mooncake.engine import TransferEngine
+
+    block_bytes = int(args.block_size_gb * 1024 * 1024 * 1024)
+    total_gb = args.num_blocks * args.block_size_gb
+
+    print(f"=== Target Node ===")
+    print(f"Blocks: {args.num_blocks} x {args.block_size_gb} GB = {total_gb} GB total")
+    print(f"Protocol: {args.protocol}")
+    print(f"Registration API: {'batch' if args.use_batch_api else 'per-block'}")
+
+    engine = TransferEngine()
+    ret = engine.initialize(
+        args.local_server_name, args.metadata_server, args.protocol, ""
+    )
+    if ret != 0:
+        raise RuntimeError(f"Engine initialization failed: {ret}")
+
+    if args.metadata_server == "P2PHANDSHAKE":
+        host = args.local_server_name.rpartition(":")[0]
+        rpc_port = engine.get_rpc_port()
+        actual_name = f"{host}:{rpc_port}"
+        print(f"Actual server name: {actual_name}")
+
+    # Allocate blocks
+    print(f"\nAllocating {args.num_blocks} blocks of {args.block_size_gb} GB...")
+    blocks = []  # (addr, size, is_hugepage)
+    for i in range(args.num_blocks):
+        try:
+            ptr, hp = allocate_block(block_bytes)
+            blocks.append((ptr, block_bytes, hp))
+            if (i + 1) % 10 == 0 or i == 0:
+                page_type = "hugepage" if hp else "4KB"
+                print(f"  Block {i+1}/{args.num_blocks}: 0x{ptr:x} ({page_type})")
+        except RuntimeError as e:
+            print(f"  Block {i+1} allocation FAILED: {e}")
+            break
+
+    if not blocks:
+        raise RuntimeError("No blocks allocated")
+
+    hp_count = sum(1 for _, _, hp in blocks if hp)
+    reg_count = sum(1 for _, _, hp in blocks if not hp)
+    print(f"  Allocated {len(blocks)} blocks: {hp_count} hugepage, {reg_count} regular")
+
+    # Register
+    addrs = [b[0] for b in blocks]
+    sizes = [b[1] for b in blocks]
+
+    print(f"\nRegistering {len(blocks)} blocks...")
+    t0 = time.time()
+
+    if args.use_batch_api:
+        ret = engine.batch_register_memory(addrs, sizes)
+        if ret != 0:
+            print(f"  batch_register_memory FAILED: {ret}")
+            print(f"  Registration took {time.time() - t0:.1f}s before failure")
+            return
+    else:
+        for i, (addr, size, _) in enumerate(blocks):
+            ret = engine.register_memory(addr, size)
+            if ret != 0:
+                print(f"  register_memory FAILED at block {i}: {ret}")
+                print(f"  Registration took {time.time() - t0:.1f}s before failure")
+                return
+
+    reg_time = time.time() - t0
+    print(f"  Registration OK: {reg_time:.1f}s for {len(blocks)} blocks")
+    print(f"  Per-block: {reg_time / len(blocks) * 1000:.0f}ms")
+
+    # Publish block info
+    info = {
+        "blocks": [{"addr": b[0], "size": b[1]} for b in blocks],
+        "num_blocks": len(blocks),
+    }
+    print(f"\nTarget ready with {len(blocks)} blocks.")
+    print(f"TARGET_INFO:{json.dumps(info)}")
+    print("Waiting for initiator (Ctrl+C to stop)...")
+
+    try:
+        signal.pause()
+    except KeyboardInterrupt:
+        print("\nShutting down target.")
+
+
+def run_initiator(args):
+    """Run as initiator: pull data from random blocks."""
+    from mooncake.engine import TransferEngine
+
+    block_bytes = int(args.block_size_gb * 1024 * 1024 * 1024)
+    transfer_bytes = int(args.transfer_size_mb * 1024 * 1024)
+
+    print(f"=== Initiator Node ===")
+    print(f"Target: {args.target_server_name}")
+    print(f"Blocks: {args.num_blocks} x {args.block_size_gb} GB")
+    print(f"Transfer size: {args.transfer_size_mb} MB")
+    print(f"Protocol: {args.protocol}")
+
+    if not args.target_server_name:
+        raise RuntimeError("--target_server_name required in initiator mode")
+
+    engine = TransferEngine()
+    ret = engine.initialize(
+        args.local_server_name, args.metadata_server, args.protocol, ""
+    )
+    if ret != 0:
+        raise RuntimeError(f"Engine initialization failed: {ret}")
+
+    if args.metadata_server == "P2PHANDSHAKE":
+        host = args.local_server_name.rpartition(":")[0]
+        rpc_port = engine.get_rpc_port()
+        actual_name = f"{host}:{rpc_port}"
+        print(f"Actual server name: {actual_name}")
+
+    # Allocate local receive buffer
+    recv_bytes = transfer_bytes
+    print(f"\nAllocating {recv_bytes / 1e6:.0f} MB receive buffer...")
+    recv_addr, hp = allocate_block(recv_bytes)
+    print(f"  0x{recv_addr:x} ({'hugepage' if hp else '4KB'})")
+    ret = engine.register_memory(recv_addr, recv_bytes)
+    if ret != 0:
+        raise RuntimeError(f"Local memory registration failed: {ret}")
+
+    # Get target's first buffer address
+    print(f"Connecting to target {args.target_server_name}...")
+    remote_base = engine.get_first_buffer_address(args.target_server_name)
+    if remote_base == 0:
+        raise RuntimeError("Cannot get target buffer address")
+    print(f"  Remote first buffer at 0x{remote_base:x}")
+
+    # Connection warmup
+    print("Warming up connection...")
+    for _ in range(3):
+        engine.transfer_sync_read(
+            args.target_server_name, recv_addr, remote_base, transfer_bytes
+        )
+    print("  Connection ready.")
+
+    # Note: blocks are independently mmap'd on the target, so they are NOT
+    # contiguous. We can only access the first block via get_first_buffer_address.
+    # The primary goal is to validate that registration of N separate blocks works
+    # and that data can be transferred from a registered block.
+    print(f"\nBenchmarking transfers from first registered block...")
+    print(f"  Each transfer: {args.transfer_size_mb} MB")
+
+    # Warmup
+    print(f"  Warming up ({args.warmup} iterations)...")
+    for w in range(args.warmup):
+        ret = engine.transfer_sync_read(
+            args.target_server_name, recv_addr, remote_base, transfer_bytes
+        )
+        if ret != 0:
+            print(f"  WARNING: warmup failed iter {w}: {ret}")
+
+    # Benchmark
+    latencies = []
+    errors = 0
+    for i in range(args.iterations):
+        t0 = time.perf_counter()
+        ret = engine.transfer_sync_read(
+            args.target_server_name, recv_addr, remote_base, transfer_bytes
+        )
+        elapsed = time.perf_counter() - t0
+
+        if ret != 0:
+            errors += 1
+            if errors <= 3:
+                print(f"  ERROR: transfer failed iter {i}: {ret}")
+            continue
+
+        latencies.append(elapsed * 1000)
+
+    if not latencies:
+        print(f"  ALL FAILED ({errors} errors)")
+        return
+
+    latencies.sort()
+    avg_ms = statistics.mean(latencies)
+    p50_ms = latencies[len(latencies) // 2]
+    p99_ms = latencies[int(len(latencies) * 0.99)]
+    throughput_gbs = (transfer_bytes / 1e9) / (p50_ms / 1000)
+
+    print(f"\n{'='*60}")
+    print(f"  Registered blocks: {args.num_blocks} x {args.block_size_gb}GB")
+    print(f"  Transfer: {args.transfer_size_mb} MB")
+    print(f"  Iterations: {len(latencies)} (errors: {errors})")
+    print(f"  Avg latency:  {avg_ms:.2f} ms")
+    print(f"  p50 latency:  {p50_ms:.2f} ms")
+    print(f"  p99 latency:  {p99_ms:.2f} ms")
+    print(f"  Throughput:   {throughput_gbs:.2f} GB/s")
+    print(f"{'='*60}")
+
+    results = {
+        "num_blocks": args.num_blocks,
+        "block_size_gb": args.block_size_gb,
+        "transfer_size_mb": args.transfer_size_mb,
+        "iterations": len(latencies),
+        "errors": errors,
+        "avg_latency_ms": round(avg_ms, 3),
+        "p50_latency_ms": round(p50_ms, 3),
+        "p99_latency_ms": round(p99_ms, 3),
+        "throughput_gbs": round(throughput_gbs, 3),
+    }
+
+    out_file = f"batch_bench_{args.num_blocks}x{args.block_size_gb}gb.json"
+    with open(out_file, "w") as f:
+        json.dump(results, f, indent=2)
+    print(f"Results saved to {out_file}")
+
+
+def main():
+    args = parse_args()
+    if args.mode == "target":
+        run_target(args)
+    else:
+        run_initiator(args)
+
+
+if __name__ == "__main__":
+    main()

mooncake-transfer-engine/example/efa_per_transfer_latency_bench.py

@@ -0,0 +1,241 @@
+#!/usr/bin/env python3
+"""
+EFA Per-Transfer Latency Benchmark
+
+Measures single-transfer latency by running transfer_engine_bench with
+threads=1, batch_size=1 across multiple block sizes. This isolates the
+per-request overhead (slice creation, spinlock, atomic ops, MR lookup)
+that the P0 NIC-striping optimization aims to eliminate.
+
+Usage:
+    python3 efa_per_transfer_latency_bench.py \
+        --target_host=HOST_A --initiator_host=HOST_B \
+        --build_dir=/path/to/build \
+        --ssh_opts="-i /path/to/key.pem"
+"""
+
+import argparse
+import os
+import re
+import subprocess
+import sys
+import time
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="EFA Per-Transfer Latency Benchmark"
+    )
+    parser.add_argument("--target_host", required=True)
+    parser.add_argument("--initiator_host", required=True)
+    parser.add_argument(
+        "--build_dir",
+        default="/opt/dlami/nvme/Mooncake/build",
+    )
+    parser.add_argument("--duration", type=int, default=10)
+    parser.add_argument("--operation", default="write", choices=["read", "write"])
+    parser.add_argument("--ssh_user", default="ubuntu")
+    parser.add_argument(
+        "--ssh_opts",
+        default="-o StrictHostKeyChecking=no -o ConnectTimeout=10",
+    )
+    parser.add_argument(
+        "--block_sizes",
+        default="65536,131072,262144,524288,1048576,2097152,4194304,8388608,16777216",
+        help="Comma-separated block sizes in bytes",
+    )
+    parser.add_argument("--threads", type=int, default=1, help="Number of threads")
+    parser.add_argument("--batch_size", type=int, default=1, help="Batch size")
+    parser.add_argument(
+        "--env", action="append", default=[],
+        help="Environment variables to pass to remote bench (e.g. --env MC_EFA_STRIPING_THRESHOLD=67108864)",
+    )
+    parser.add_argument("--output", default=None, help="Output file for results")
+    return parser.parse_args()
+
+
+def run_ssh(host, command, user, ssh_opts, timeout=None):
+    ssh_args = ["ssh", *ssh_opts.split(), f"{user}@{host}", command]
+    try:
+        result = subprocess.run(ssh_args, capture_output=True, text=True, timeout=timeout)
+        return result.returncode, result.stdout, result.stderr
+    except subprocess.TimeoutExpired:
+        return -1, "", "Command timed out"
+
+
+def kill_bench(host, user, ssh_opts):
+    cmd = (
+        "ps aux | grep '[t]ransfer_engine_bench' | awk '{print $2}' "
+        "| xargs -r kill 2>/dev/null; sleep 1; echo done"
+    )
+    run_ssh(host, cmd, user, ssh_opts, timeout=15)
+    time.sleep(2)
+
+
+def start_target(host, build_dir, user, ssh_opts):
+    bench_bin = os.path.join(
+        build_dir, "mooncake-transfer-engine/example/transfer_engine_bench"
+    )
+    log_file = "/tmp/efa_latency_target.log"
+    target_cmd = (
+        f"cd {build_dir} && "
+        f"{bench_bin} "
+        f"--mode=target --protocol=efa --metadata_server=P2PHANDSHAKE "
+        f"> {log_file} 2>&1"
+    )
+    ssh_args = ["ssh", "-n", *ssh_opts.split(), f"{user}@{host}", target_cmd]
+    subprocess.Popen(
+        ssh_args,
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+        stdin=subprocess.DEVNULL,
+    )
+    for _ in range(20):
+        time.sleep(1)
+        rc, stdout, _ = run_ssh(
+            host, f"grep 'listening on' {log_file} 2>/dev/null",
+            user, ssh_opts, timeout=10,
+        )
+        if rc == 0 and "listening on" in stdout:
+            match = re.search(r"listening on (\S+:\d+)", stdout)
+            if match:
+                return match.group(1)
+    return None
+
+
+def run_single_bench(host, build_dir, target_addr, block_size,
+                     duration, operation, user, ssh_opts,
+                     threads=1, batch_size=1, env_vars=None):
+    """Run bench with configurable threads and batch_size."""
+    bench_bin = os.path.join(
+        build_dir, "mooncake-transfer-engine/example/transfer_engine_bench"
+    )
+    env_prefix = " ".join(env_vars) + " " if env_vars else ""
+    bench_cmd = (
+        f"cd {build_dir} && "
+        f"{env_prefix}"
+        f"{bench_bin} "
+        f"--mode=initiator --protocol=efa --metadata_server=P2PHANDSHAKE "
+        f"--segment_id={target_addr} "
+        f"--operation={operation} "
+        f"--duration={duration} "
+        f"--threads={threads} "
+        f"--block_size={block_size} "
+        f"--batch_size={batch_size} "
+        f"2>&1"
+    )
+    timeout = duration + 60
+    rc, stdout, stderr = run_ssh(host, bench_cmd, user, ssh_opts, timeout=timeout)
+    combined = stdout + "\n" + stderr
+
+    # Parse throughput
+    match = re.search(r"throughput\s+([\d.]+)\s+GB/s", combined)
+    if match:
+        return float(match.group(1))
+
+    # Try MB/s
+    match = re.search(r"throughput\s+([\d.]+)\s+MB/s", combined)
+    if match:
+        return float(match.group(1)) / 1024.0
+
+    print(f"    WARNING: Could not parse throughput", file=sys.stderr)
+    for line in combined.strip().split("\n")[-3:]:
+        print(f"      {line}", file=sys.stderr)
+    return None
+
+
+def format_size(size_bytes):
+    if size_bytes >= 1048576:
+        return f"{size_bytes / 1048576:.0f}MB"
+    return f"{size_bytes / 1024:.0f}KB"
+
+
+def main():
+    args = parse_args()
+    block_sizes = [int(x) for x in args.block_sizes.split(",")]
+
+    print("=" * 70)
+    print("EFA Per-Transfer Latency Benchmark")
+    print("=" * 70)
+    print(f"  Target:     {args.target_host}")
+    print(f"  Initiator:  {args.initiator_host}")
+    print(f"  Build dir:  {args.build_dir}")
+    print(f"  Duration:   {args.duration}s per point")
+    print(f"  Operation:  {args.operation}")
+    print(f"  Mode:       threads={args.threads}, batch_size={args.batch_size}")
+    print(f"  Block sizes: {[format_size(b) for b in block_sizes]}")
+    print()
+
+    # Start target
+    kill_bench(args.target_host, args.ssh_user, args.ssh_opts)
+    print("Starting target...", end="", flush=True)
+    target_addr = start_target(
+        args.target_host, args.build_dir, args.ssh_user, args.ssh_opts
+    )
+    if not target_addr:
+        print(" FAILED")
+        sys.exit(1)
+    print(f" ready ({target_addr})")
+    print()
+
+    results = []
+    for i, block_size in enumerate(block_sizes):
+        tag = format_size(block_size)
+        print(f"  [{i+1}/{len(block_sizes)}] {tag:>6} ...", end="", flush=True)
+
+        tp = run_single_bench(
+            args.initiator_host, args.build_dir, target_addr,
+            block_size, args.duration, args.operation,
+            args.ssh_user, args.ssh_opts,
+            threads=args.threads, batch_size=args.batch_size,
+            env_vars=args.env,
+        )
+
+        if tp is None or tp == 0:
+            print(" FAILED")
+            results.append((block_size, None, None))
+            continue
+
+        # latency = block_size / throughput
+        tp_bytes = tp * 1e9  # GB/s -> bytes/s
+        latency_us = (block_size / tp_bytes) * 1e6  # microseconds
+        results.append((block_size, tp, latency_us))
+        print(f" {tp:7.2f} GB/s  latency={latency_us:8.1f} us")
+
+    # Cleanup
+    kill_bench(args.target_host, args.ssh_user, args.ssh_opts)
+
+    # Summary
+    print()
+    print("=" * 70)
+    print("Results Summary")
+    print("=" * 70)
+    print(f"{'Block Size':>12}  {'Throughput':>12}  {'Latency (us)':>14}")
+    print("-" * 42)
+    for block_size, tp, lat in results:
+        tag = format_size(block_size)
+        if tp is not None:
+            print(f"{tag:>12}  {tp:>9.2f} GB/s  {lat:>11.1f} us")
+        else:
+            print(f"{tag:>12}  {'N/A':>12}  {'N/A':>14}")
+
+    # Write output file
+    if args.output:
+        with open(args.output, "w") as f:
+            f.write(f"# EFA Per-Transfer Latency Benchmark\n")
+            f.write(f"# Operation: {args.operation}\n")
+            f.write(f"# Duration: {args.duration}s per point\n")
+            f.write(f"# Mode: threads={args.threads}, batch_size={args.batch_size}\n")
+            f.write(f"#\n")
+            f.write(f"{'block_bytes':>12}  {'block_size':>10}  {'gbps':>10}  {'latency_us':>12}\n")
+            for block_size, tp, lat in results:
+                tag = format_size(block_size)
+                if tp is not None:
+                    f.write(f"{block_size:>12}  {tag:>10}  {tp:>10.2f}  {lat:>12.1f}\n")
+                else:
+                    f.write(f"{block_size:>12}  {tag:>10}  {'N/A':>10}  {'N/A':>12}\n")
+        print(f"\nResults saved to: {args.output}")
+
+
+if __name__ == "__main__":
+    main()

mooncake-transfer-engine/example/kvcache_prefix_bench.py

@@ -0,0 +1,452 @@
+#!/usr/bin/env python3
+"""
+KV Cache Prefix Transfer Benchmark for Mooncake EFA Transport.
+
+Simulates cross-node KV cache transfer for prefix cache hits in LLM inference.
+Tests how registered memory pool size affects transfer latency/throughput
+under the per-NIC partition auto-split strategy.
+
+Usage:
+  # Target node (holds KV cache pool):
+  python kvcache_prefix_bench.py --mode target \
+      --local_server_name <target_ip>:12345 \
+      --pool_size_gb 100 --protocol efa
+
+  # Initiator node (pulls prefix KV cache):
+  python kvcache_prefix_bench.py --mode initiator \
+      --local_server_name <initiator_ip>:12346 \
+      --target_server_name <target_ip>:12345 \
+      --pool_size_gb 100 --protocol efa \
+      --prefix_tokens 4096,8192,16384,32768
+
+Requires: mooncake Python package (pip install -e mooncake-wheel)
+Branch: feat/efa-auto-split-mr (per-NIC partition auto-split)
+"""
+
+import argparse
+import ctypes
+import ctypes.util
+import json
+import os
+import signal
+import statistics
+import sys
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="KV Cache Prefix Transfer Benchmark"
+    )
+    parser.add_argument(
+        "--mode",
+        choices=["target", "initiator"],
+        required=True,
+        help="Run as target (KV cache holder) or initiator (puller)",
+    )
+    parser.add_argument(
+        "--local_server_name",
+        required=True,
+        help="Local address, e.g. 172.31.6.162:12345",
+    )
+    parser.add_argument(
+        "--target_server_name",
+        default="",
+        help="Target address (initiator mode only)",
+    )
+    parser.add_argument(
+        "--metadata_server",
+        default="P2PHANDSHAKE",
+        help="Metadata server address (default: P2PHANDSHAKE)",
+    )
+    parser.add_argument(
+        "--protocol", default="efa", help="Transport protocol (default: efa)"
+    )
+    parser.add_argument(
+        "--pool_size_gb",
+        type=float,
+        default=10.0,
+        help="KV cache pool size in GB to register (default: 10)",
+    )
+    parser.add_argument(
+        "--prefix_tokens",
+        default="4096,8192,16384,32768",
+        help="Comma-separated list of prefix token counts to test",
+    )
+    parser.add_argument(
+        "--kv_bytes_per_token",
+        type=int,
+        default=89856,
+        help="KV cache bytes per token. "
+        "Default: 89856 for GLM-5.1 (754B MoE, MLA attention): "
+        "(kv_lora_rank=512 + qk_rope_head_dim=64) * 2 bytes * 78 layers. "
+        "For GLM-4-9B (standard MHA): use 40960",
+    )
+    parser.add_argument(
+        "--iterations",
+        type=int,
+        default=50,
+        help="Number of transfer iterations per test (default: 50)",
+    )
+    parser.add_argument(
+        "--warmup",
+        type=int,
+        default=5,
+        help="Number of warmup iterations (default: 5)",
+    )
+    parser.add_argument(
+        "--use_gpu",
+        action="store_true",
+        help="Use GPU memory instead of CPU memory",
+    )
+    parser.add_argument(
+        "--gpu_id", type=int, default=0, help="GPU device ID (default: 0)"
+    )
+    parser.add_argument(
+        "--threads",
+        type=int,
+        default=1,
+        help="Number of threads for concurrent transfer (default: 1). "
+        "Each thread transfers a chunk of the prefix in parallel.",
+    )
+    return parser.parse_args()
+
+
+def allocate_cpu_memory(size_bytes):
+    """Allocate page-aligned CPU memory using mmap."""
+    libc_name = ctypes.util.find_library("c")
+    libc = ctypes.CDLL(libc_name, use_errno=True)
+
+    # mmap(NULL, size, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0)
+    PROT_READ = 0x1
+    PROT_WRITE = 0x2
+    MAP_PRIVATE = 0x02
+    MAP_ANONYMOUS = 0x20
+    MAP_HUGETLB = 0x40000
+    MAP_FAILED = ctypes.c_void_p(-1).value
+
+    libc.mmap.restype = ctypes.c_void_p
+    libc.mmap.argtypes = [
+        ctypes.c_void_p,
+        ctypes.c_size_t,
+        ctypes.c_int,
+        ctypes.c_int,
+        ctypes.c_int,
+        ctypes.c_long,
+    ]
+
+    flags = MAP_PRIVATE | MAP_ANONYMOUS
+
+    # Try hugepages first (recommended for large EFA registrations)
+    ptr = libc.mmap(None, size_bytes, PROT_READ | PROT_WRITE,
+                    flags | MAP_HUGETLB, -1, 0)
+    if ptr and ptr != MAP_FAILED:
+        print(f"  Allocated {size_bytes / 1e9:.1f} GB with 2MB hugepages")
+        return ptr, size_bytes, True
+
+    # Fall back to regular pages
+    ptr = libc.mmap(None, size_bytes, PROT_READ | PROT_WRITE, flags, -1, 0)
+    if not ptr or ptr == MAP_FAILED:
+        raise RuntimeError(
+            f"mmap failed for {size_bytes} bytes: "
+            f"errno={ctypes.get_errno()}"
+        )
+    print(
+        f"  Allocated {size_bytes / 1e9:.1f} GB with 4KB pages "
+        f"(hugepages unavailable — configure vm.nr_hugepages for better EFA performance)"
+    )
+    return ptr, size_bytes, False
+
+
+def allocate_gpu_memory(size_bytes, gpu_id):
+    """Allocate GPU memory using PyTorch."""
+    import torch
+
+    torch.cuda.set_device(gpu_id)
+    # Allocate as uint8 tensor
+    tensor = torch.zeros(size_bytes, dtype=torch.uint8, device=f"cuda:{gpu_id}")
+    ptr = tensor.data_ptr()
+    print(f"  Allocated {size_bytes / 1e9:.1f} GB on GPU {gpu_id}")
+    return ptr, tensor  # keep tensor alive
+
+
+def run_target(args):
+    """Run as target node: allocate KV cache pool and wait."""
+    from mooncake.engine import TransferEngine
+
+    print(f"=== Target Node ===")
+    print(f"Pool size: {args.pool_size_gb} GB")
+    print(f"Protocol: {args.protocol}")
+
+    engine = TransferEngine()
+    ret = engine.initialize(
+        args.local_server_name, args.metadata_server, args.protocol, ""
+    )
+    if ret != 0:
+        raise RuntimeError(f"Engine initialization failed: {ret}")
+
+    if args.metadata_server == "P2PHANDSHAKE":
+        host = args.local_server_name.rpartition(":")[0]
+        rpc_port = engine.get_rpc_port()
+        actual_name = f"{host}:{rpc_port}"
+        print(f"Actual server name: {actual_name}")
+
+    pool_bytes = int(args.pool_size_gb * 1024 * 1024 * 1024)
+
+    print(f"Allocating {args.pool_size_gb} GB KV cache pool...")
+    if args.use_gpu:
+        pool_addr, _tensor = allocate_gpu_memory(pool_bytes, args.gpu_id)
+    else:
+        pool_addr, _, _ = allocate_cpu_memory(pool_bytes)
+
+    print(f"Registering memory with transfer engine...")
+    t0 = time.time()
+    ret = engine.register_memory(pool_addr, pool_bytes)
+    reg_time = time.time() - t0
+    if ret != 0:
+        raise RuntimeError(f"Memory registration failed: {ret}")
+    print(f"  Registration took {reg_time:.1f}s")
+
+    print(f"\nTarget ready. Pool addr: 0x{pool_addr:x}, size: {pool_bytes}")
+    print("Waiting for initiator (Ctrl+C to stop)...")
+
+    # Write pool info for initiator to read
+    info = {"pool_addr": pool_addr, "pool_bytes": pool_bytes}
+    print(f"TARGET_INFO:{json.dumps(info)}")
+
+    try:
+        signal.pause()
+    except KeyboardInterrupt:
+        print("\nShutting down target.")
+
+
+def run_initiator(args):
+    """Run as initiator node: pull prefix KV cache and measure performance."""
+    from mooncake.engine import TransferEngine
+
+    print(f"=== Initiator Node ===")
+    print(f"Target: {args.target_server_name}")
+    print(f"Pool size: {args.pool_size_gb} GB")
+    print(f"Protocol: {args.protocol}")
+    print(f"KV bytes/token: {args.kv_bytes_per_token}")
+    print(f"Threads: {args.threads}")
+
+    if not args.target_server_name:
+        raise RuntimeError("--target_server_name required in initiator mode")
+
+    prefix_tokens_list = [int(x) for x in args.prefix_tokens.split(",")]
+    transfer_sizes = [
+        tokens * args.kv_bytes_per_token for tokens in prefix_tokens_list
+    ]
+
+    print(f"\nTest matrix:")
+    for tokens, size in zip(prefix_tokens_list, transfer_sizes):
+        print(f"  {tokens:>6} tokens -> {size / 1e6:.1f} MB transfer")
+
+    # Initialize engine
+    engine = TransferEngine()
+    ret = engine.initialize(
+        args.local_server_name, args.metadata_server, args.protocol, ""
+    )
+    if ret != 0:
+        raise RuntimeError(f"Engine initialization failed: {ret}")
+
+    if args.metadata_server == "P2PHANDSHAKE":
+        host = args.local_server_name.rpartition(":")[0]
+        rpc_port = engine.get_rpc_port()
+        actual_name = f"{host}:{rpc_port}"
+        print(f"Actual server name: {actual_name}")
+
+    # Allocate local receive buffer (large enough for the biggest transfer)
+    max_transfer = max(transfer_sizes)
+    recv_bytes = max_transfer
+    print(f"\nAllocating {recv_bytes / 1e6:.1f} MB local receive buffer...")
+    if args.use_gpu:
+        recv_addr, _tensor = allocate_gpu_memory(recv_bytes, args.gpu_id)
+    else:
+        recv_addr, _, _ = allocate_cpu_memory(recv_bytes)
+
+    ret = engine.register_memory(recv_addr, recv_bytes)
+    if ret != 0:
+        raise RuntimeError(f"Local memory registration failed: {ret}")
+
+    # Get target's buffer address
+    print(f"Connecting to target {args.target_server_name}...")
+    remote_addr = engine.get_first_buffer_address(args.target_server_name)
+    if remote_addr == 0:
+        raise RuntimeError(
+            "Cannot get target buffer address. "
+            "Is the target running and registered?"
+        )
+    print(f"  Remote buffer at 0x{remote_addr:x}")
+
+    # Connection warmup: a small transfer to establish the EFA connection
+    # (openSegment, endpoint creation, etc.) so it doesn't skew the first
+    # prefix size's measurements.
+    print("Warming up connection...")
+    warmup_size = min(transfer_sizes[0], recv_bytes)
+    for _ in range(3):
+        engine.transfer_sync_read(
+            args.target_server_name, recv_addr, remote_addr, warmup_size
+        )
+    print("  Connection ready.")
+
+    num_threads = args.threads
+
+    def do_transfer(local_addr, remote_addr_with_offset, size):
+        """Single transfer_sync_read call, suitable for thread pool."""
+        return engine.transfer_sync_read(
+            args.target_server_name, local_addr,
+            remote_addr_with_offset, size
+        )
+
+    def threaded_transfer(local_base, remote_base, total_size, pool):
+        """Split transfer across threads; return max of 0 (ok) or error code."""
+        if num_threads <= 1:
+            return engine.transfer_sync_read(
+                args.target_server_name, local_base, remote_base, total_size
+            )
+        chunk = total_size // num_threads
+        # Align chunk to 4KB
+        chunk = chunk & ~0xFFF
+        futures = []
+        for t in range(num_threads):
+            off = t * chunk
+            sz = chunk if t < num_threads - 1 else (total_size - off)
+            futures.append(
+                pool.submit(do_transfer, local_base + off,
+                            remote_base + off, sz)
+            )
+        return max(f.result() for f in futures)
+
+    # Run benchmarks
+    print(f"\n{'='*72}")
+    print(
+        f"{'Prefix':>8} {'Size':>10} {'Latency(ms)':>12} "
+        f"{'p50(ms)':>10} {'p99(ms)':>10} {'Tput(GB/s)':>12}"
+    )
+    print(f"{'='*72}")
+
+    results = []
+    with ThreadPoolExecutor(max_workers=num_threads) as pool:
+        for tokens, transfer_size in zip(prefix_tokens_list, transfer_sizes):
+            if transfer_size > recv_bytes:
+                print(f"  SKIP {tokens} tokens: transfer {transfer_size} > recv buffer")
+                continue
+
+            # Warmup: use the same offset pattern as the benchmark to pre-warm
+            # remote memory pages and DMA paths at each offset.
+            pool_bytes = int(args.pool_size_gb * 1024 * 1024 * 1024)
+            max_offset = pool_bytes - transfer_size
+            for w in range(args.warmup):
+                for i in range(args.iterations):
+                    if max_offset > 0:
+                        offset = ((i * transfer_size) % max_offset) & ~0xFFF
+                    else:
+                        offset = 0
+                    ret = threaded_transfer(
+                        recv_addr, remote_addr + offset, transfer_size, pool
+                    )
+                    if ret != 0:
+                        print(f"  WARNING: warmup transfer failed: {ret}")
+                        break
+
+            # Benchmark
+            latencies = []
+            errors = 0
+            for i in range(args.iterations):
+                # Use different offset within remote pool for each iteration
+                # to simulate accessing different prefix locations
+                if max_offset > 0:
+                    # Align to 4KB boundary
+                    offset = ((i * transfer_size) % max_offset) & ~0xFFF
+                else:
+                    offset = 0
+
+                t0 = time.perf_counter()
+                ret = threaded_transfer(
+                    recv_addr, remote_addr + offset, transfer_size, pool
+                )
+                elapsed = time.perf_counter() - t0
+
+                if ret != 0:
+                    errors += 1
+                    if errors <= 3:
+                        print(f"  ERROR: transfer failed at iter {i}: {ret}")
+                    continue
+
+                latencies.append(elapsed * 1000)  # ms
+
+            if not latencies:
+                print(f"  {tokens:>6}k  ALL FAILED ({errors} errors)")
+                continue
+
+            latencies.sort()
+            avg_ms = statistics.mean(latencies)
+            p50_ms = latencies[len(latencies) // 2]
+            p99_ms = latencies[int(len(latencies) * 0.99)]
+            throughput_gbs = (transfer_size / 1e9) / (p50_ms / 1000)
+
+            print(
+                f"  {tokens:>6} {transfer_size/1e6:>8.1f}MB "
+                f"{avg_ms:>11.2f} {p50_ms:>10.2f} {p99_ms:>10.2f} "
+                f"{throughput_gbs:>11.2f}"
+            )
+
+            results.append(
+                {
+                    "prefix_tokens": tokens,
+                    "transfer_bytes": transfer_size,
+                    "transfer_mb": transfer_size / 1e6,
+                    "pool_size_gb": args.pool_size_gb,
+                    "iterations": len(latencies),
+                    "errors": errors,
+                    "avg_latency_ms": round(avg_ms, 3),
+                    "p50_latency_ms": round(p50_ms, 3),
+                    "p99_latency_ms": round(p99_ms, 3),
+                    "throughput_gbs": round(throughput_gbs, 3),
+                    "threads": num_threads,
+                }
+            )
+
+    print(f"{'='*72}")
+
+    # Summary
+    if results:
+        print(f"\n=== Summary (pool_size={args.pool_size_gb}GB, threads={num_threads}) ===")
+        print(json.dumps(results, indent=2))
+
+        # Save results
+        thread_tag = f"_t{num_threads}" if num_threads > 1 else ""
+        out_file = (
+            f"kvcache_bench_pool{args.pool_size_gb}gb"
+            f"{'_gpu' if args.use_gpu else '_cpu'}{thread_tag}.json"
+        )
+        with open(out_file, "w") as f:
+            json.dump(
+                {
+                    "pool_size_gb": args.pool_size_gb,
+                    "protocol": args.protocol,
+                    "use_gpu": args.use_gpu,
+                    "kv_bytes_per_token": args.kv_bytes_per_token,
+                    "threads": num_threads,
+                    "iterations": args.iterations,
+                    "results": results,
+                },
+                f,
+                indent=2,
+            )
+        print(f"Results saved to {out_file}")
+
+
+def main():
+    args = parse_args()
+    if args.mode == "target":
+        run_target(args)
+    else:
+        run_initiator(args)
+
+
+if __name__ == "__main__":
+    main()

mooncake-transfer-engine/include/common.h

@@ -113,6 +113,10 @@ static inline int64_t getCurrentTimeInNano() {
     return (int64_t{ts.tv_sec} * kNanosPerSecond + int64_t{ts.tv_nsec});
 }
 
+static inline int64_t getCurrentTimeInMilli() {
+    return getCurrentTimeInNano() / 1000 / 1000;
+}
+
 static inline std::string getCurrentDateTime() {
     auto now = std::chrono::system_clock::now();
     auto time_t_now = std::chrono::system_clock::to_time_t(now);
@@ -327,9 +331,13 @@ static inline ssize_t writeFully(int fd, const void *buf, size_t len) {
 }
 
 static inline ssize_t readFully(int fd, void *buf, size_t len) {
+    // Set a timeout for read to avoid hanging forever.
+    constexpr std::chrono::seconds kReadTimeout = std::chrono::seconds(300);
+    const std::chrono::steady_clock::time_point deadline =
+        std::chrono::steady_clock::now() + kReadTimeout;
     char *pos = (char *)buf;
     size_t nbytes = len;
-    while (nbytes) {
+    while (nbytes && std::chrono::steady_clock::now() < deadline) {
         ssize_t rc = read(fd, pos, nbytes);
         if (rc < 0 && (errno == EAGAIN || errno == EINTR))
             continue;
@@ -344,7 +352,14 @@ static inline ssize_t readFully(int fd, void *buf, size_t len) {
         pos += rc;
         nbytes -= rc;
     }
-    return len;
+    if (nbytes != 0) {
+        LOG(WARNING) << "Socket read timed out, timeout: "
+                     << kReadTimeout.count()
+                     << ", deadline: " << deadline.time_since_epoch().count()
+                     << ", read " << len - nbytes << " out of " << len
+                     << " bytes";
+    }
+    return len - nbytes;
 }
 
 static inline int writeString(int fd, const HandShakeRequestType type,
@@ -457,6 +472,21 @@ static inline const std::string MakeNicPath(const std::string &server_name,
     return server_name + NIC_PATH_DELIM + nic_name;
 }
 
+// Strip the port from a nic_path to get a stable key for endpoint reuse.
+// "ip-172-31-45-191:15365@rdmap135s0" → "ip-172-31-45-191@rdmap135s0"
+// This allows the same physical peer to reuse endpoints across reconnections
+// (each run picks a random P2P handshake port).
+static inline std::string normalizeNicPath(const std::string &nic_path) {
+    std::string server_name = getServerNameFromNicPath(nic_path);
+    std::string nic_name = getNicNameFromNicPath(nic_path);
+    if (server_name.empty() || nic_name.empty()) return nic_path;
+    size_t colon = server_name.rfind(':');
+    if (colon != std::string::npos) {
+        server_name = server_name.substr(0, colon);
+    }
+    return server_name + NIC_PATH_DELIM + nic_name;
+}
+
 static inline bool overlap(const void *a, size_t a_len, const void *b,
                            size_t b_len) {
     return (a >= b && a < (char *)b + b_len) ||

mooncake-transfer-engine/include/config.h

@@ -65,6 +65,7 @@ struct GlobalConfig {
     int ib_pci_relaxed_ordering_mode = 0;
     bool ascend_use_fabric_mem = false;
     bool ascend_agent_mode = false;
+    size_t efa_striping_threshold = 2 * 1024 * 1024;  // 2MB default
     // ub config parameters
     size_t num_jfc_per_ctx = 2;
     size_t num_jfce_per_ctx = 2;
@@ -81,13 +82,13 @@ struct RpcCommunicatorConfig {
     size_t pool_size = 10;
 };
 
-void loadGlobalConfig(GlobalConfig &config);
+void loadGlobalConfig(GlobalConfig& config);
 
 void dumpGlobalConfig();
 
-void updateGlobalConfig(ibv_device_attr &device_attr);
+void updateGlobalConfig(ibv_device_attr& device_attr);
 
-GlobalConfig &globalConfig();
+GlobalConfig& globalConfig();
 
 uint16_t getDefaultHandshakePort();
 

mooncake-transfer-engine/include/multi_transport.h

@@ -55,6 +55,14 @@ class MultiTransport {
 
     Transport *getTransport(const std::string &proto);
 
+    /**
+     * @brief Check if TCP is the only installed transport.
+     *
+     * When only TCP transport is available (no RDMA, NVLink, etc.),
+     * local memcpy is preferred over TCP loopback for same-host transfers.
+     */
+    bool isTcpOnly() const;
+
     std::vector<Transport *> listTransports();
 
     void *getBaseAddr();

mooncake-transfer-engine/include/transfer_engine.h

@@ -150,6 +150,14 @@ class TransferEngine {
 
     Transport* getTransport(const std::string& proto);
 
+    /**
+     * @brief Check if TCP is the only installed transport.
+     *
+     * When only TCP transport is available (no RDMA, NVLink, etc.),
+     * local memcpy is preferred over TCP loopback for same-host transfers.
+     */
+    bool isTcpOnly() const;
+
     int syncSegmentCache(const std::string& segment_name = "");
 
     std::shared_ptr<TransferMetadata> getMetadata();

mooncake-transfer-engine/include/transfer_engine_c.h

@@ -105,6 +105,8 @@ transfer_engine_t createTransferEngine(const char *metadata_conn_string,
                                        const char *ip_or_host_name,
                                        uint64_t rpc_port, int auto_discover);
 
+int discoverTopology(transfer_engine_t engine);
+
 int getLocalIpAndPort(transfer_engine_t engine, char *buf_out, size_t buf_len);
 
 transport_t installTransport(transfer_engine_t engine, const char *proto,
@@ -119,6 +121,11 @@ segment_id_t openSegmentNoCache(transfer_engine_t engine,
 
 int closeSegment(transfer_engine_t engine, segment_id_t segment_id);
 
+// Eagerly pre-connect all EFA endpoints to `segment_name`. Eliminates the
+// first-batch fi_av_insert stall (observed ~6 s for 16 local NICs × N peer
+// NICs). No-op on non-EFA installs. Idempotent. Returns 0 on success.
+int warmupEfaSegment(transfer_engine_t engine, const char *segment_name);
+
 int removeLocalSegment(transfer_engine_t engine, const char *segment_name);
 
 void destroyTransferEngine(transfer_engine_t engine);

mooncake-transfer-engine/include/transfer_engine_impl.h

@@ -334,6 +334,8 @@ class TransferEngineImpl {
         return multi_transports_->getTransport(proto);
     }
 
+    bool isTcpOnly() const { return multi_transports_->isTcpOnly(); }
+
     int syncSegmentCache(const std::string& segment_name = "") {
         return metadata_->syncSegmentCache(segment_name);
     }

mooncake-transfer-engine/include/transport/efa_transport/efa_context.h

@@ -31,6 +31,7 @@
 #include <memory>
 #include <string>
 #include <thread>
+#include <map>
 #include <unordered_map>
 
 #include "common.h"
@@ -44,63 +45,87 @@ class EfaTransport;
 
 struct EfaCq {
     EfaCq() : cq(nullptr), outstanding(0) {}
-    struct fid_cq *cq;
+    struct fid_cq* cq;
     volatile int outstanding;
 };
 
 struct EfaMemoryRegionMeta {
-    void *addr;
+    void* addr;
     size_t length;
-    struct fid_mr *mr;
+    struct fid_mr* mr;
     uint64_t key;
 };
 
-// Simple endpoint store for EFA
+// Endpoint store for EFA with LRU eviction support.
+// When the store approaches the AV capacity (max_endpoints), stale endpoints
+// that have been inactive beyond `inactive_timeout_sec` are evicted to free
+// AV slots.  This prevents AV exhaustion in long-running services that
+// communicate with many transient peers.
 class EfaEndpointStore {
    public:
-    std::shared_ptr<EfaEndPoint> get(const std::string &peer_nic_path);
+    static constexpr double kDefaultInactiveTimeoutSec = 5.0;  // 5 seconds
+
+    explicit EfaEndpointStore(
+        size_t max_endpoints = 65536,
+        double inactive_timeout_sec = kDefaultInactiveTimeoutSec);
+
+    std::shared_ptr<EfaEndPoint> get(const std::string& peer_nic_path);
     // Atomically get-or-insert: returns existing endpoint or inserts new_ep.
     // Prevents duplicate endpoint creation from concurrent callers.
+    // Triggers eviction when the store is at or above max_endpoints.
     std::shared_ptr<EfaEndPoint> getOrInsert(
-        const std::string &peer_nic_path, std::shared_ptr<EfaEndPoint> new_ep);
-    void add(const std::string &peer_nic_path,
+        const std::string& peer_nic_path, std::shared_ptr<EfaEndPoint> new_ep);
+    void add(const std::string& peer_nic_path,
              std::shared_ptr<EfaEndPoint> endpoint);
-    void remove(const std::string &peer_nic_path);
+    void remove(const std::string& peer_nic_path);
     int disconnectAll();
     size_t size() const;
 
+    // Evict endpoints that have been inactive longer than the configured
+    // timeout.  Returns the number of evicted endpoints.
+    size_t evictStale();
+
+    // Remove endpoints whose connections are broken (not connected and not
+    // initializing).  Returns the number of removed endpoints.
+    size_t removeDisconnected();
+
    private:
+    // Must be called with write lock held.
+    size_t evictStaleLocked();
+
     mutable RWSpinlock lock_;
     std::unordered_map<std::string, std::shared_ptr<EfaEndPoint>> endpoints_;
+    size_t max_endpoints_;
+    double inactive_timeout_sec_;
 };
 
 // EfaContext represents the set of resources controlled by each local EFA
 // device, including Memory Region, CQ, EndPoint, etc. using libfabric
 class EfaContext {
    public:
-    EfaContext(EfaTransport &engine, const std::string &device_name);
+    EfaContext(EfaTransport& engine, const std::string& device_name);
 
     ~EfaContext();
 
     int construct(size_t num_cq_list = 1, size_t num_comp_channels = 1,
                   uint8_t port = 1, int gid_index = -1, size_t max_cqe = 4096,
-                  int max_endpoints = 256);
+                  int max_endpoints = 65536);
 
    private:
     int deconstruct();
 
    public:
     // Memory Region Management
-    int registerMemoryRegion(void *addr, size_t length, int access);
-    int unregisterMemoryRegion(void *addr);
-    int preTouchMemory(void *addr, size_t length);
-    uint64_t rkey(void *addr);
-    uint64_t lkey(void *addr);
-    void *mrDesc(void *addr);  // Get MR descriptor for fi_write local_desc
+    int registerMemoryRegion(void* addr, size_t length, int access);
+    int unregisterMemoryRegion(void* addr);
+    int preTouchMemory(void* addr, size_t length);
+    uint64_t rkey(void* addr);
+    uint64_t lkey(void* addr);
+    void* mrDesc(void* addr);  // Get MR descriptor for fi_write local_desc
 
    private:
-    int registerMemoryRegionInternal(void *addr, size_t length, int access,
-                                     EfaMemoryRegionMeta &mrMeta);
+    int registerMemoryRegionInternal(void* addr, size_t length, int access,
+                                     EfaMemoryRegionMeta& mrMeta);
 
    public:
     bool active() const { return active_; }
@@ -108,27 +133,39 @@ class EfaContext {
 
    public:
     // EndPoint Management
-    std::shared_ptr<EfaEndPoint> endpoint(const std::string &peer_nic_path);
-    int deleteEndpoint(const std::string &peer_nic_path);
+    std::shared_ptr<EfaEndPoint> endpoint(const std::string& peer_nic_path);
+    int deleteEndpoint(const std::string& peer_nic_path);
     int disconnectAllEndpoints();
     size_t getTotalQPNumber() const;
 
    public:
     // Access to engine for endpoint handshake
-    EfaTransport &engine() { return engine_; }
-    const EfaTransport &engine() const { return engine_; }
+    EfaTransport& engine() { return engine_; }
+    const EfaTransport& engine() const { return engine_; }
 
     // Submit slices for transfer
-    int submitPostSend(const std::vector<Transport::Slice *> &slice_list);
+    int submitPostSend(const std::vector<Transport::Slice*>& slice_list);
 
     // Poll completion queue for completed operations
     int pollCq(int max_entries, int cq_index = 0);
 
+    // Evict stale endpoints to free EFA resources
+    size_t evictStaleEndpoints() {
+        return endpoint_store_ ? endpoint_store_->evictStale() : 0;
+    }
+
     // Get CQ count
     size_t cqCount() const { return cq_list_.size(); }
 
+    // Round-robin CQ assignment for new endpoints (mirrors RDMA transport)
+    std::shared_ptr<EfaCq> nextCq() {
+        if (cq_list_.empty()) return nullptr;
+        int index = (next_cq_index_++) % cq_list_.size();
+        return cq_list_[index];
+    }
+
     // Get CQ outstanding count pointer
-    volatile int *cqOutstandingCount(int cq_index) {
+    volatile int* cqOutstandingCount(int cq_index) {
         if (cq_index < 0 || (size_t)cq_index >= cq_list_.size()) return nullptr;
         return &cq_list_[cq_index]->outstanding;
     }
@@ -142,10 +179,10 @@ class EfaContext {
 
    public:
     // Libfabric accessors
-    struct fid_fabric *fabric() const { return fabric_; }
-    struct fid_domain *domain() const { return domain_; }
-    struct fid_av *av() const { return av_; }
-    struct fi_info *info() const { return fi_info_; }
+    struct fid_fabric* fabric() const { return fabric_; }
+    struct fid_domain* domain() const { return domain_; }
+    struct fid_av* av() const { return av_; }
+    struct fi_info* info() const { return fi_info_; }
     std::string localAddr() const;
 
     // Compatibility methods (libfabric doesn't use lid/gid like ibverbs)
@@ -153,23 +190,24 @@ class EfaContext {
     std::string gid() const { return localAddr(); }
 
    private:
-    EfaTransport &engine_;
+    EfaTransport& engine_;
     std::string device_name_;
 
     // Libfabric objects
-    struct fi_info *fi_info_;
-    struct fi_info *hints_;
-    struct fid_fabric *fabric_;
-    struct fid_domain *domain_;
-    struct fid_av *av_;  // Address vector for peer addressing
+    struct fi_info* fi_info_;
+    struct fi_info* hints_;
+    struct fid_fabric* fabric_;
+    struct fid_domain* domain_;
+    struct fid_av* av_;  // Address vector for peer addressing
 
     bool active_;
 
     std::shared_ptr<EfaEndpointStore> endpoint_store_;
     std::vector<std::shared_ptr<EfaCq>> cq_list_;
+    std::atomic<int> next_cq_index_{0};
 
     RWSpinlock mr_lock_;
-    std::unordered_map<uint64_t, EfaMemoryRegionMeta> mr_map_;
+    std::map<uint64_t, EfaMemoryRegionMeta> mr_map_;
 };
 
 }  // namespace mooncake

mooncake-transfer-engine/include/transport/efa_transport/efa_endpoint.h

@@ -41,8 +41,8 @@ class EfaContext;
 // handling This struct MUST have fi_context as its first member
 struct EfaOpContext {
     struct fi_context fi_ctx;  // Must be first member
-    Transport::Slice *slice;   // Slice pointer for completion handling
-    volatile int *wr_depth;    // Pointer to endpoint's wr_depth_ for CQ
+    Transport::Slice* slice;   // Slice pointer for completion handling
+    volatile int* wr_depth;    // Pointer to endpoint's wr_depth_ for CQ
                                // completion decrement
 };
 
@@ -55,28 +55,29 @@ class EfaEndPoint {
 
     enum Status { INITIALIZING, UNCONNECTED, CONNECTED };
 
-    EfaEndPoint(EfaContext &context);
+    EfaEndPoint(EfaContext& context);
     ~EfaEndPoint();
 
     // Construct endpoint with specified completion queue
-    int construct(struct fid_cq *cq, size_t num_qp_list = 1, size_t max_sge = 4,
+    int construct(struct fid_cq* cq, volatile int* cq_outstanding,
+                  size_t num_qp_list = 1, size_t max_sge = 4,
                   size_t max_wr = 256, size_t max_inline = 64);
 
    private:
     int deconstruct();
 
    public:
-    void setPeerNicPath(const std::string &peer_nic_path);
+    void setPeerNicPath(const std::string& peer_nic_path);
 
     int setupConnectionsByActive();
 
-    int setupConnectionsByActive(const std::string &peer_nic_path) {
+    int setupConnectionsByActive(const std::string& peer_nic_path) {
         setPeerNicPath(peer_nic_path);
         return setupConnectionsByActive();
     }
 
-    int setupConnectionsByPassive(const HandShakeDesc &peer_desc,
-                                  HandShakeDesc &local_desc);
+    int setupConnectionsByPassive(const HandShakeDesc& peer_desc,
+                                  HandShakeDesc& local_desc);
 
     bool hasOutstandingSlice() const;
 
@@ -93,6 +94,12 @@ class EfaEndPoint {
         return (getCurrentTimeInNano() - inactive_time_) / 1000000000.0;
     }
 
+    void touchLastUsed() { last_used_time_ = getCurrentTimeInNano(); }
+
+    double lastUsedAge() const {
+        return (getCurrentTimeInNano() - last_used_time_) / 1000000000.0;
+    }
+
    public:
     bool connected() const {
         return status_.load(std::memory_order_relaxed) == CONNECTED;
@@ -108,8 +115,8 @@ class EfaEndPoint {
     const std::string toString() const;
 
     // Submit RDMA write/read operations via libfabric
-    int submitPostSend(std::vector<Transport::Slice *> &slice_list,
-                       std::vector<Transport::Slice *> &failed_slice_list);
+    int submitPostSend(std::vector<Transport::Slice*>& slice_list,
+                       std::vector<Transport::Slice*>& failed_slice_list);
 
     // Get the number of endpoints (always 1 for EFA RDM)
     size_t getQPNumber() const { return 1; }
@@ -120,27 +127,27 @@ class EfaEndPoint {
     // Get peer's fi_addr
     fi_addr_t getPeerFiAddr() const { return peer_fi_addr_; }
 
-    EfaContext &context() { return context_; }
+    EfaContext& context() { return context_; }
 
    private:
     // Setup connection using peer's address from handshake
-    int doSetupConnection(const std::string &peer_addr,
-                          std::string *reply_msg = nullptr);
+    int doSetupConnection(const std::string& peer_addr,
+                          std::string* reply_msg = nullptr);
 
     // Insert peer address into address vector
-    int insertPeerAddr(const std::string &peer_addr);
+    int insertPeerAddr(const std::string& peer_addr);
 
    private:
-    EfaContext &context_;
+    EfaContext& context_;
     std::atomic<Status> status_;
 
     RWSpinlock lock_;
     std::string peer_nic_path_;
 
     // Libfabric endpoint
-    struct fid_ep *ep_;
-    struct fid_cq *tx_cq_;
-    struct fid_cq *rx_cq_;
+    struct fid_ep* ep_;
+    struct fid_cq* tx_cq_;
+    struct fid_cq* rx_cq_;
     fi_addr_t peer_fi_addr_;  // Peer's address in the AV
 
     // Local endpoint address (for handshake)
@@ -149,14 +156,15 @@ class EfaEndPoint {
 
     volatile int wr_depth_;
     int max_wr_depth_;
-    volatile int *cq_outstanding_;
+    volatile int* cq_outstanding_;
 
-    // Spinlock to serialize fi_write calls on this endpoint.
+    // Spinlock to serialize fi_write/fi_read calls on this endpoint.
     // libfabric RDM endpoints are not thread-safe by default.
     std::atomic_flag post_lock_ = ATOMIC_FLAG_INIT;
 
     volatile bool active_;
     volatile uint64_t inactive_time_;
+    volatile uint64_t last_used_time_;  // Updated on connection and I/O
 };
 
 }  // namespace mooncake

mooncake-transfer-engine/include/transport/efa_transport/efa_transport.h

@@ -52,82 +52,96 @@ class EfaTransport : public Transport {
 
     ~EfaTransport();
 
-    int install(std::string &local_server_name,
+    int install(std::string& local_server_name,
                 std::shared_ptr<TransferMetadata> meta,
                 std::shared_ptr<Topology> topo) override;
 
-    const char *getName() const override { return "efa"; }
+    const char* getName() const override { return "efa"; }
 
-    int registerLocalMemory(void *addr, size_t length,
-                            const std::string &location, bool remote_accessible,
+    int registerLocalMemory(void* addr, size_t length,
+                            const std::string& location, bool remote_accessible,
                             bool update_metadata) override;
 
-    int unregisterLocalMemory(void *addr, bool update_metadata = true) override;
+    int unregisterLocalMemory(void* addr, bool update_metadata = true) override;
 
-    int registerLocalMemoryBatch(const std::vector<BufferEntry> &buffer_list,
-                                 const std::string &location) override;
+    int registerLocalMemoryBatch(const std::vector<BufferEntry>& buffer_list,
+                                 const std::string& location) override;
 
     int unregisterLocalMemoryBatch(
-        const std::vector<void *> &addr_list) override;
+        const std::vector<void*>& addr_list) override;
+
+    // Eagerly establish EFA endpoints to every NIC of `segment_name`.
+    //
+    // Rationale: libfabric FI_EP_RDM endpoints resolve peer addresses lazily
+    // via fi_av_insert() on first send. With 16 local NICs × N peer NICs,
+    // the first submitTransfer() of a batch serializes ~N*16
+    // handshake+fi_av_insert round-trips, producing a single-digit-second
+    // stall (observed ~6 s on B300 for the first 100 × 0.5 MB batch). After
+    // this call returns, every (local_ctx, peer_nic) endpoint is CONNECTED
+    // and the first real submitTransfer() goes straight to fi_write/fi_read.
+    //
+    // Safe to call multiple times (idempotent: endpoint() + setup are both
+    // idempotent). Re-run after any openSegment() on a new peer.
+    int warmupSegment(const std::string& segment_name);
 
    private:
     // Internal version with force_sequential option to avoid nested parallelism
-    int registerLocalMemoryInternal(void *addr, size_t length,
-                                    const std::string &location,
+    int registerLocalMemoryInternal(void* addr, size_t length,
+                                    const std::string& location,
                                     bool remote_accessible,
                                     bool update_metadata,
                                     bool force_sequential);
 
-    int unregisterLocalMemoryInternal(void *addr, bool update_metadata,
+    int unregisterLocalMemoryInternal(void* addr, bool update_metadata,
                                       bool force_sequential);
 
     // TRANSFER
 
     Status submitTransfer(BatchID batch_id,
-                          const std::vector<TransferRequest> &entries) override;
+                          const std::vector<TransferRequest>& entries) override;
 
     Status submitTransferTask(
-        const std::vector<TransferTask *> &task_list) override;
+        const std::vector<TransferTask*>& task_list) override;
 
     Status getTransferStatus(BatchID batch_id,
-                             std::vector<TransferStatus> &status);
+                             std::vector<TransferStatus>& status);
 
     Status getTransferStatus(BatchID batch_id, size_t task_id,
-                             TransferStatus &status) override;
+                             TransferStatus& status) override;
 
-    SegmentID getSegmentID(const std::string &segment_name);
+    SegmentID getSegmentID(const std::string& segment_name);
 
    private:
     int allocateLocalSegmentID();
 
-    int preTouchMemory(void *addr, size_t length);
+    int preTouchMemory(void* addr, size_t length);
 
    public:
-    int onSetupEfaConnections(const HandShakeDesc &peer_desc,
-                              HandShakeDesc &local_desc);
+    int onSetupEfaConnections(const HandShakeDesc& peer_desc,
+                              HandShakeDesc& local_desc);
 
-    int sendHandshake(const std::string &peer_server_name,
-                      const HandShakeDesc &local_desc,
-                      HandShakeDesc &peer_desc) {
+    int sendHandshake(const std::string& peer_server_name,
+                      const HandShakeDesc& local_desc,
+                      HandShakeDesc& peer_desc) {
         return metadata_->sendHandshake(peer_server_name, local_desc,
                                         peer_desc);
     }
 
-    const std::string &local_server_name() const { return local_server_name_; }
+    const std::string& local_server_name() const { return local_server_name_; }
 
     std::shared_ptr<TransferMetadata> meta() { return metadata_; }
 
    private:
     int initializeEfaResources();
 
-    int startHandshakeDaemon(std::string &local_server_name);
+    int startHandshakeDaemon(std::string& local_server_name);
 
    public:
-    static int selectDevice(SegmentDesc *desc, uint64_t offset, size_t length,
-                            int &buffer_id, int &device_id, int retry_cnt = 0);
-    static int selectDevice(SegmentDesc *desc, uint64_t offset, size_t length,
-                            std::string_view hint, int &buffer_id,
-                            int &device_id, int retry_cnt = 0);
+    static int selectDevice(SegmentDesc* desc, uint64_t offset, size_t length,
+                            int& buffer_id, int& device_id, int retry_cnt = 0);
+    static int selectDevice(SegmentDesc* desc, uint64_t offset, size_t length,
+                            std::string_view hint, int& buffer_id,
+                            int& device_id, int retry_cnt = 0);
 
    private:
     // Start/stop CQ polling worker threads
@@ -139,6 +153,16 @@ class EfaTransport : public Transport {
     std::vector<std::shared_ptr<EfaContext>> context_list_;
     std::shared_ptr<Topology> local_topology_;
 
+    // Track chunked MR registrations for per-NIC partitioned buffers.
+    // When a buffer exceeds max_mr_size, it is split into chunks, each
+    // registered on a disjoint subset of NICs (per-NIC partition).
+    struct ChunkRegistration {
+        uint64_t addr;
+        std::vector<size_t> nic_indices;
+    };
+    std::mutex chunk_map_mutex_;
+    std::unordered_map<uint64_t, std::vector<ChunkRegistration>> chunk_map_;
+
     // CQ polling worker threads
     std::atomic<bool> worker_running_{false};
     std::vector<std::thread> worker_threads_;

mooncake-transfer-engine/include/transport/rdma_transport/endpoint_store.h

@@ -38,6 +38,8 @@ class EndpointStore {
    public:
     virtual std::shared_ptr<RdmaEndPoint> getEndpoint(
         const std::string &peer_nic_path) = 0;
+    virtual std::shared_ptr<RdmaEndPoint> getEndpointByPtr(
+        const RdmaEndPoint *endpoint_ptr) = 0;
     virtual std::shared_ptr<RdmaEndPoint> insertEndpoint(
         const std::string &peer_nic_path, RdmaContext *context) = 0;
     virtual int deleteEndpoint(const std::string &peer_nic_path) = 0;
@@ -58,6 +60,8 @@ class FIFOEndpointStore : public EndpointStore {
     FIFOEndpointStore(size_t max_size) : max_size_(max_size) {}
     std::shared_ptr<RdmaEndPoint> getEndpoint(
         const std::string &peer_nic_path) override;
+    std::shared_ptr<RdmaEndPoint> getEndpointByPtr(
+        const RdmaEndPoint *endpoint_ptr) override;
     std::shared_ptr<RdmaEndPoint> insertEndpoint(
         const std::string &peer_nic_path, RdmaContext *context) override;
     int deleteEndpoint(const std::string &peer_nic_path) override;
@@ -89,6 +93,8 @@ class SIEVEEndpointStore : public EndpointStore {
         : waiting_list_len_(0), max_size_(max_size) {}
     std::shared_ptr<RdmaEndPoint> getEndpoint(
         const std::string &peer_nic_path) override;
+    std::shared_ptr<RdmaEndPoint> getEndpointByPtr(
+        const RdmaEndPoint *endpoint_ptr) override;
     std::shared_ptr<RdmaEndPoint> insertEndpoint(
         const std::string &peer_nic_path, RdmaContext *context) override;
     int deleteEndpoint(const std::string &peer_nic_path) override;

mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h

@@ -99,6 +99,9 @@ class RdmaContext {
     // EndPoint Management
     std::shared_ptr<RdmaEndPoint> endpoint(const std::string &peer_nic_path);
 
+    std::shared_ptr<RdmaEndPoint> getEndpointByPtr(
+        const RdmaEndPoint *endpoint_ptr);
+
     int deleteEndpoint(const std::string &peer_nic_path);
 
     int disconnectAllEndpoints();

mooncake-transfer-engine/rust/.gitignore

@@ -0,0 +1 @@
+src/customer_pattern.rs

mooncake-transfer-engine/rust/build.rs

@@ -17,16 +17,84 @@ use std::path::PathBuf;
 
 fn main() {
     println!("cargo:rustc-link-search=native=../build/src");
+    println!("cargo:rustc-link-search=native=../../build/mooncake-transfer-engine/src");
     println!("cargo:rustc-link-lib=static=transfer_engine");
 
+    // libbase.a holds mooncake::Status, which libtransfer_engine.a references.
+    println!("cargo:rustc-link-search=native=../build/src/common/base");
+    println!("cargo:rustc-link-search=native=../../build/mooncake-transfer-engine/src/common/base");
+    println!("cargo:rustc-link-lib=static=base");
+
+    // The transfer_engine build uses ASIO_SEPARATE_COMPILATION + ASIO_DYN_LINK,
+    // so the asio symbols live in mooncake-asio/libasio.so.  Link it whenever
+    // we can find it (standalone cmake build places it alongside src/).
+    println!("cargo:rustc-link-search=native=../build/mooncake-asio");
+    println!("cargo:rustc-link-search=native=../../build/mooncake-asio");
+    println!("cargo:rustc-link-lib=asio");
+
+    // EFA on AWS installs libfabric under /opt/amazon/efa/lib.
+    if std::path::Path::new("/opt/amazon/efa/lib").exists() {
+        println!("cargo:rustc-link-search=native=/opt/amazon/efa/lib");
+    }
+
     println!("cargo:rustc-link-lib=stdc++");
     println!("cargo:rustc-link-lib=ibverbs");
+    // libfabric (fi_*): only needed for EFA transport, but harmless when
+    // the system has it installed; required on AWS EFA hosts.  Opt-out by
+    // setting MOONCAKE_WITHOUT_LIBFABRIC=1 if building on a box without it.
+    if env::var("MOONCAKE_WITHOUT_LIBFABRIC")
+        .map(|v| v == "1" || v.eq_ignore_ascii_case("on") || v.eq_ignore_ascii_case("true"))
+        .unwrap_or(false)
+    {
+        // skip
+    } else {
+        println!("cargo:rustc-link-lib=fabric");
+    }
     println!("cargo:rustc-link-lib=glog");
     println!("cargo:rustc-link-lib=gflags");
     println!("cargo:rustc-link-lib=pthread");
     println!("cargo:rustc-link-lib=jsoncpp");
     println!("cargo:rustc-link-lib=numa");
-    println!("cargo:rustc-link-lib=etcd-cpp-api");
+    println!("cargo:rustc-link-lib=curl");
+
+    let flag_on = |name: &str| {
+        env::var(name)
+            .map(|v| v == "1" || v.eq_ignore_ascii_case("on") || v.eq_ignore_ascii_case("true"))
+            .unwrap_or(false)
+    };
+
+    // etcd-cpp-api: only needed when transfer_engine was built with
+    // USE_ETCD=ON.  Opt-in via MOONCAKE_WITH_ETCD=1 to keep non-etcd builds
+    // (e.g. EFA-only on AWS) linkable.
+    if flag_on("MOONCAKE_WITH_ETCD") {
+        println!("cargo:rustc-link-lib=etcd-cpp-api");
+    }
+
+    // CUDA runtime: libtransfer_engine.a built with USE_CUDA=ON pulls in
+    // cudaMemcpy/cudaStream* symbols.  The Rust demos themselves don't call
+    // CUDA — this is purely a transitive archive dep.  Enable with
+    // MOONCAKE_WITH_CUDA=1 and optional CUDA_HOME override for lib path.
+    if flag_on("MOONCAKE_WITH_CUDA") {
+        // Accept either a CUDA_HOME (append lib64/lib) or an explicit
+        // CUDART_LIB_DIR that already points at the directory containing
+        // libcudart.so.  This covers both /usr/local/cuda installs and
+        // pip-wheel layouts like .../nvidia/cu13/lib.
+        if let Ok(dir) = env::var("CUDART_LIB_DIR") {
+            println!("cargo:rustc-link-search=native={}", dir);
+        } else if let Ok(cuda_home) = env::var("CUDA_HOME") {
+            let lib64 = PathBuf::from(&cuda_home).join("lib64");
+            let lib = PathBuf::from(&cuda_home).join("lib");
+            if lib64.exists() {
+                println!("cargo:rustc-link-search=native={}", lib64.display());
+            }
+            if lib.exists() {
+                println!("cargo:rustc-link-search=native={}", lib.display());
+            }
+        } else {
+            println!("cargo:rustc-link-search=native=/usr/local/cuda/lib64");
+        }
+        println!("cargo:rustc-link-lib=cudart");
+    }
 
     let bindings = bindgen::builder()
         .header("../include/transfer_engine_c.h")

mooncake-transfer-engine/rust/src/transfer_engine.rs

@@ -91,6 +91,27 @@ impl TransferEngine {
         Ok(Self { engine })
     }
 
+    pub fn discover_topology(&self) -> Result<()> {
+        let ret = unsafe { bindings::discoverTopology(self.engine) };
+        if ret != 0 {
+            bail!("Failed to discover topology")
+        } else {
+            Ok(())
+        }
+    }
+
+    pub fn install_transport(&self, proto: &str) -> Result<()> {
+        let proto_c = CString::new(proto).map_err(|_| anyhow!("CString::new failed"))?;
+        let ret = unsafe {
+            bindings::installTransport(self.engine, proto_c.as_ptr(), std::ptr::null_mut())
+        };
+        if ret.is_null() {
+            bail!("Failed to install transport '{}'", proto)
+        } else {
+            Ok(())
+        }
+    }
+
     pub fn close(&mut self) -> Result<()> {
         unsafe {
             bindings::destroyTransferEngine(self.engine);
@@ -242,6 +263,20 @@ impl TransferEngine {
         }
     }
 
+    /// Eagerly establish EFA endpoints to `segment_name` so the first
+    /// `submit_transfer` doesn't pay the serial fi_av_insert cost. No-op on
+    /// non-EFA transports. Call after `open_segment` (and after the metadata
+    /// has the peer's NIC list published).
+    pub fn warmup_efa_segment(&self, name: &str) -> Result<()> {
+        let name_c = CString::new(name).map_err(|_| anyhow!("CString::new failed"))?;
+        let ret = unsafe { bindings::warmupEfaSegment(self.engine, name_c.as_ptr()) };
+        if ret < 0 {
+            bail!("warmupEfaSegment failed for {}: {}", name, ret)
+        } else {
+            Ok(())
+        }
+    }
+
     pub fn sync_segment_cache(&self) -> Result<()> {
         let ret = unsafe { bindings::syncSegmentCache(self.engine) };
         if ret < 0 {

mooncake-transfer-engine/src/CMakeLists.txt

@@ -91,20 +91,22 @@ if(USE_MLU)
   target_link_libraries(transfer_engine PUBLIC cnrt cndrv)
 endif()
 
-if (USE_ASCEND)
-  target_link_libraries(transfer_engine PUBLIC ascendcl hccl ascend_transport MPI::MPI)
+if(USE_ASCEND)
+  target_link_libraries(transfer_engine PUBLIC ascendcl hccl ascend_transport
+                                               MPI::MPI)
 endif()
 
-if (USE_ASCEND_DIRECT)
+if(USE_ASCEND_DIRECT)
   target_link_libraries(transfer_engine PUBLIC ascend_transport)
 endif()
 
-if (USE_UBSHMEM)
+if(USE_UBSHMEM)
   target_link_libraries(transfer_engine PUBLIC ascend_transport)
 endif()
 
-if (USE_ASCEND_HETEROGENEOUS)
-  file(GLOB ASCEND_TOOLKIT_ROOT "/usr/local/Ascend/ascend-toolkit/latest/*-linux")
+if(USE_ASCEND_HETEROGENEOUS)
+  file(GLOB ASCEND_TOOLKIT_ROOT
+       "/usr/local/Ascend/ascend-toolkit/latest/*-linux")
   set(ASCEND_LIB_DIR "${ASCEND_TOOLKIT_ROOT}/lib64")
   link_directories(${ASCEND_LIB_DIR})
   target_link_libraries(transfer_engine PUBLIC ascendcl ascend_transport)
@@ -112,7 +114,7 @@ endif()
 
 if(USE_TENT)
   add_compile_definitions(transfer_engine PUBLIC USE_TENT)
-  target_link_libraries(transfer_engine PUBLIC tent)
+  target_link_libraries(transfer_engine PUBLIC tent_link_group)
 endif()
 
 if(USE_INTRA_NVLINK)

mooncake-transfer-engine/src/config.cpp

@@ -21,8 +21,8 @@
 #include <unistd.h>
 
 namespace mooncake {
-void loadGlobalConfig(GlobalConfig &config) {
-    const char *num_cq_per_ctx_env = std::getenv("MC_NUM_CQ_PER_CTX");
+void loadGlobalConfig(GlobalConfig& config) {
+    const char* num_cq_per_ctx_env = std::getenv("MC_NUM_CQ_PER_CTX");
     if (num_cq_per_ctx_env) {
         int val = atoi(num_cq_per_ctx_env);
         if (val > 0 && val < 256) {
@@ -35,7 +35,7 @@ void loadGlobalConfig(GlobalConfig &config) {
                 << "Ignore value from environment variable MC_NUM_CQ_PER_CTX";
     }
 
-    const char *num_comp_channels_per_ctx_env =
+    const char* num_comp_channels_per_ctx_env =
         std::getenv("MC_NUM_COMP_CHANNELS_PER_CTX");
     if (num_comp_channels_per_ctx_env) {
         int val = atoi(num_comp_channels_per_ctx_env);
@@ -46,7 +46,7 @@ void loadGlobalConfig(GlobalConfig &config) {
                             "MC_NUM_COMP_CHANNELS_PER_CTX";
     }
 
-    const char *port_env = std::getenv("MC_IB_PORT");
+    const char* port_env = std::getenv("MC_IB_PORT");
     if (port_env) {
         int val = atoi(port_env);
         if (val >= 0 && val < 256)
@@ -55,7 +55,7 @@ void loadGlobalConfig(GlobalConfig &config) {
             LOG(WARNING) << "Ignore value from environment variable MC_IB_PORT";
     }
 
-    const char *gid_index_env = std::getenv("MC_GID_INDEX");
+    const char* gid_index_env = std::getenv("MC_GID_INDEX");
     if (!gid_index_env) gid_index_env = std::getenv("NCCL_IB_GID_INDEX");
 
     if (gid_index_env) {
@@ -67,7 +67,7 @@ void loadGlobalConfig(GlobalConfig &config) {
                 << "Ignore value from environment variable MC_GID_INDEX";
     }
 
-    const char *max_cqe_per_ctx_env = std::getenv("MC_MAX_CQE_PER_CTX");
+    const char* max_cqe_per_ctx_env = std::getenv("MC_MAX_CQE_PER_CTX");
     if (max_cqe_per_ctx_env) {
         size_t val = atoi(max_cqe_per_ctx_env);
         if (val > 0 && val <= UINT16_MAX)
@@ -77,7 +77,7 @@ void loadGlobalConfig(GlobalConfig &config) {
                 << "Ignore value from environment variable MC_MAX_CQE_PER_CTX";
     }
 
-    const char *max_ep_per_ctx_env = std::getenv("MC_MAX_EP_PER_CTX");
+    const char* max_ep_per_ctx_env = std::getenv("MC_MAX_EP_PER_CTX");
     if (max_ep_per_ctx_env) {
         size_t val = atoi(max_ep_per_ctx_env);
         if (val > 0 && val <= UINT16_MAX)
@@ -87,7 +87,7 @@ void loadGlobalConfig(GlobalConfig &config) {
                 << "Ignore value from environment variable MC_MAX_EP_PER_CTX";
     }
 
-    const char *num_qp_per_ep_env = std::getenv("MC_NUM_QP_PER_EP");
+    const char* num_qp_per_ep_env = std::getenv("MC_NUM_QP_PER_EP");
     if (num_qp_per_ep_env) {
         int val = atoi(num_qp_per_ep_env);
         if (val > 0 && val < 256)
@@ -97,7 +97,7 @@ void loadGlobalConfig(GlobalConfig &config) {
                 << "Ignore value from environment variable MC_NUM_QP_PER_EP";
     }
 
-    const char *max_sge_env = std::getenv("MC_MAX_SGE");
+    const char* max_sge_env = std::getenv("MC_MAX_SGE");
     if (max_sge_env) {
         size_t val = atoi(max_sge_env);
         if (val > 0 && val <= UINT16_MAX)
@@ -106,7 +106,7 @@ void loadGlobalConfig(GlobalConfig &config) {
             LOG(WARNING) << "Ignore value from environment variable MC_MAX_SGE";
     }
 
-    const char *max_wr_env = std::getenv("MC_MAX_WR");
+    const char* max_wr_env = std::getenv("MC_MAX_WR");
     if (max_wr_env) {
         size_t val = atoi(max_wr_env);
         if (val > 0 && val <= UINT16_MAX)
@@ -115,7 +115,7 @@ void loadGlobalConfig(GlobalConfig &config) {
             LOG(WARNING) << "Ignore value from environment variable MC_MAX_WR";
     }
 
-    const char *max_inline_env = std::getenv("MC_MAX_INLINE");
+    const char* max_inline_env = std::getenv("MC_MAX_INLINE");
     if (max_inline_env) {
         size_t val = atoi(max_inline_env);
         if (val <= UINT16_MAX)
@@ -125,7 +125,7 @@ void loadGlobalConfig(GlobalConfig &config) {
                 << "Ignore value from environment variable MC_MAX_INLINE";
     }
 
-    const char *mtu_length_env = std::getenv("MC_MTU");
+    const char* mtu_length_env = std::getenv("MC_MTU");
     if (mtu_length_env) {
         size_t val = atoi(mtu_length_env);
         if (val == 512)
@@ -143,7 +143,7 @@ void loadGlobalConfig(GlobalConfig &config) {
         }
     }
 
-    const char *handshake_port_env = std::getenv("MC_HANDSHAKE_PORT");
+    const char* handshake_port_env = std::getenv("MC_HANDSHAKE_PORT");
     if (handshake_port_env) {
         int val = atoi(handshake_port_env);
         if (val > 0 && val < 65536)
@@ -153,7 +153,7 @@ void loadGlobalConfig(GlobalConfig &config) {
                 << "Ignore value from environment variable MC_HANDSHAKE_PORT";
     }
 
-    const char *workers_per_ctx_env = std::getenv("MC_WORKERS_PER_CTX");
+    const char* workers_per_ctx_env = std::getenv("MC_WORKERS_PER_CTX");
     if (workers_per_ctx_env) {
         size_t val = atoi(workers_per_ctx_env);
         if (val > 0 && val <= 8)
@@ -163,7 +163,7 @@ void loadGlobalConfig(GlobalConfig &config) {
                 << "Ignore value from environment variable MC_WORKERS_PER_CTX";
     }
 
-    const char *slice_size_env = std::getenv("MC_SLICE_SIZE");
+    const char* slice_size_env = std::getenv("MC_SLICE_SIZE");
     if (slice_size_env) {
         size_t val = atoi(slice_size_env);
         if (val > 0)
@@ -173,7 +173,7 @@ void loadGlobalConfig(GlobalConfig &config) {
                 << "Ignore value from environment variable MC_SLICE_SIZE";
     }
 
-    const char *min_reg_size_env = std::getenv("MC_MIN_REG_SIZE");
+    const char* min_reg_size_env = std::getenv("MC_MIN_REG_SIZE");
     if (min_reg_size_env) {
         size_t val = atoll(min_reg_size_env);
         if (val > 0) {
@@ -184,7 +184,7 @@ void loadGlobalConfig(GlobalConfig &config) {
                 << "Ignore value from environment variable MC_MIN_REG_SIZE";
     }
 
-    const char *max_mr_size_env = std::getenv("MC_MAX_MR_SIZE");
+    const char* max_mr_size_env = std::getenv("MC_MAX_MR_SIZE");
     if (max_mr_size_env) {
         uint64_t val = atoll(max_mr_size_env);
         if (val > 0) {
@@ -195,7 +195,17 @@ void loadGlobalConfig(GlobalConfig &config) {
         }
     }
 
-    const char *retry_cnt_env = std::getenv("MC_RETRY_CNT");
+    const char* efa_striping_env = std::getenv("MC_EFA_STRIPING_THRESHOLD");
+    if (efa_striping_env) {
+        size_t val = strtoull(efa_striping_env, nullptr, 10);
+        if (val > 0)
+            config.efa_striping_threshold = val;
+        else
+            LOG(WARNING) << "Ignore value from environment variable "
+                            "MC_EFA_STRIPING_THRESHOLD";
+    }
+
+    const char* retry_cnt_env = std::getenv("MC_RETRY_CNT");
     if (retry_cnt_env) {
         size_t val = atoi(retry_cnt_env);
         if (val > 0 && val < 128)
@@ -205,12 +215,12 @@ void loadGlobalConfig(GlobalConfig &config) {
                 << "Ignore value from environment variable MC_RETRY_CNT";
     }
 
-    const char *disable_metacache = std::getenv("MC_DISABLE_METACACHE");
+    const char* disable_metacache = std::getenv("MC_DISABLE_METACACHE");
     if (disable_metacache) {
         config.metacache = false;
     }
 
-    const char *handshake_listen_backlog =
+    const char* handshake_listen_backlog =
         std::getenv("MC_HANDSHAKE_LISTEN_BACKLOG");
     if (handshake_listen_backlog) {
         int val = std::stoi(handshake_listen_backlog);
@@ -222,7 +232,7 @@ void loadGlobalConfig(GlobalConfig &config) {
         }
     }
 
-    const char *log_level = std::getenv("MC_LOG_LEVEL");
+    const char* log_level = std::getenv("MC_LOG_LEVEL");
     config.trace = false;
     if (log_level) {
         if (strcmp(log_level, "TRACE") == 0) {
@@ -238,7 +248,7 @@ void loadGlobalConfig(GlobalConfig &config) {
     }
     FLAGS_minloglevel = config.log_level;
 
-    const char *slice_timeout_env = std::getenv("MC_SLICE_TIMEOUT");
+    const char* slice_timeout_env = std::getenv("MC_SLICE_TIMEOUT");
     if (slice_timeout_env) {
         int val = atoi(slice_timeout_env);
         if (val > 0 && val < 65536)
@@ -248,7 +258,7 @@ void loadGlobalConfig(GlobalConfig &config) {
                 << "Ignore value from environment variable MC_SLICE_TIMEOUT";
     }
 
-    const char *log_dir_path = std::getenv("MC_LOG_DIR");
+    const char* log_dir_path = std::getenv("MC_LOG_DIR");
     if (log_dir_path) {
         google::InitGoogleLogging("mooncake-transfer-engine");
         if (opendir(log_dir_path) == NULL) {
@@ -267,7 +277,7 @@ void loadGlobalConfig(GlobalConfig &config) {
         }
     }
 
-    const char *min_port_env = std::getenv("MC_MIN_PRC_PORT");
+    const char* min_port_env = std::getenv("MC_MIN_PRC_PORT");
     if (min_port_env) {
         int val = atoi(min_port_env);
         if (val > 0 && val < 65536)
@@ -277,7 +287,7 @@ void loadGlobalConfig(GlobalConfig &config) {
                 << "Ignore value from environment variable MC_PRC_MIN_PORT";
     }
 
-    const char *max_port_env = std::getenv("MC_MAX_PRC_PORT");
+    const char* max_port_env = std::getenv("MC_MAX_PRC_PORT");
     if (max_port_env) {
         int val = atoi(max_port_env);
         if (val > 0 && val < 65536)
@@ -291,7 +301,7 @@ void loadGlobalConfig(GlobalConfig &config) {
         config.use_ipv6 = true;
     }
 
-    const char *fragment_ratio = std::getenv("MC_FRAGMENT_RATIO");
+    const char* fragment_ratio = std::getenv("MC_FRAGMENT_RATIO");
     if (fragment_ratio) {
         size_t val = atoi(fragment_ratio);
         if (val > 0 && val < config.slice_size)
@@ -307,7 +317,7 @@ void loadGlobalConfig(GlobalConfig &config) {
         config.enable_dest_device_affinity = true;
     }
 
-    const char *enable_parallel_reg_mr =
+    const char* enable_parallel_reg_mr =
         std::getenv("MC_ENABLE_PARALLEL_REG_MR");
     if (enable_parallel_reg_mr) {
         int val = atoi(enable_parallel_reg_mr);
@@ -319,7 +329,7 @@ void loadGlobalConfig(GlobalConfig &config) {
         }
     }
 
-    const char *endpoint_store_type_env = std::getenv("MC_ENDPOINT_STORE_TYPE");
+    const char* endpoint_store_type_env = std::getenv("MC_ENDPOINT_STORE_TYPE");
     if (endpoint_store_type_env) {
         if (strcmp(endpoint_store_type_env, "FIFO") == 0) {
             config.endpoint_store_type = EndpointStoreType::FIFO;
@@ -331,7 +341,7 @@ void loadGlobalConfig(GlobalConfig &config) {
         }
     }
 
-    const char *traffic_class_env = std::getenv("MC_IB_TC");
+    const char* traffic_class_env = std::getenv("MC_IB_TC");
     if (traffic_class_env) {
         try {
             int val = std::stoi(traffic_class_env);
@@ -343,13 +353,13 @@ void loadGlobalConfig(GlobalConfig &config) {
                     << "value " << traffic_class_env
                     << " out of range (should be 0-255)";
             }
-        } catch (const std::exception &e) {
+        } catch (const std::exception& e) {
             LOG(WARNING) << "Invalid MC_IB_TC environment value: "
                          << traffic_class_env << ". Error: " << e.what();
         }
     }
 
-    const char *ib_relaxed_ordering_env =
+    const char* ib_relaxed_ordering_env =
         std::getenv("MC_IB_PCI_RELAXED_ORDERING");
     if (ib_relaxed_ordering_env) {
         int val = atoi(ib_relaxed_ordering_env);
@@ -374,8 +384,8 @@ std::string mtuLengthToString(ibv_mtu mtu) {
         return "UNKNOWN";
 }
 
-void updateGlobalConfig(ibv_device_attr &device_attr) {
-    auto &config = globalConfig();
+void updateGlobalConfig(ibv_device_attr& device_attr) {
+    auto& config = globalConfig();
     if (config.max_ep_per_ctx * config.num_qp_per_ep >
         (size_t)device_attr.max_qp)
         config.max_ep_per_ctx = device_attr.max_qp / config.num_qp_per_ep;
@@ -392,7 +402,7 @@ void updateGlobalConfig(ibv_device_attr &device_attr) {
 }
 
 void dumpGlobalConfig() {
-    auto &config = globalConfig();
+    auto& config = globalConfig();
     LOG(INFO) << "=== GlobalConfig ===";
     LOG(INFO) << "num_cq_per_ctx = " << config.num_cq_per_ctx;
     LOG(INFO) << "num_comp_channels_per_ctx = "
@@ -411,7 +421,7 @@ void dumpGlobalConfig() {
     LOG(INFO) << "ib_traffic_class = " << config.ib_traffic_class;
 }
 
-GlobalConfig &globalConfig() {
+GlobalConfig& globalConfig() {
     static GlobalConfig config;
     static std::once_flag g_once_flag;
     std::call_once(g_once_flag, []() { loadGlobalConfig(config); });

mooncake-transfer-engine/src/multi_transport.cpp

@@ -461,6 +461,10 @@ Transport* MultiTransport::getTransport(const std::string& proto) {
     return transport_map_[proto].get();
 }
 
+bool MultiTransport::isTcpOnly() const {
+    return transport_map_.size() == 1 && transport_map_.count("tcp") == 1;
+}
+
 std::vector<Transport*> MultiTransport::listTransports() {
     std::vector<Transport*> transport_list;
     for (auto& entry : transport_map_)

mooncake-transfer-engine/src/transfer_engine.cpp

@@ -179,6 +179,8 @@ Transport* TransferEngine::getTransport(const std::string& proto) {
     return impl_->getTransport(proto);
 }
 
+bool TransferEngine::isTcpOnly() const { return impl_->isTcpOnly(); }
+
 int TransferEngine::syncSegmentCache(const std::string& segment_name) {
     return impl_->syncSegmentCache(segment_name);
 }
@@ -578,6 +580,15 @@ Transport* TransferEngine::getTransport(const std::string& proto) {
         return impl_->getTransport(proto);
 }
 
+bool TransferEngine::isTcpOnly() const {
+    if (use_tent_)
+        // TENT already rejects TCP loopback transfers when MC_STORE_MEMCPY
+        // is disabled, so auto-enabling memcpy is unnecessary in TENT mode.
+        return false;
+    else
+        return impl_->isTcpOnly();
+}
+
 int TransferEngine::syncSegmentCache(const std::string& segment_name) {
     if (use_tent_)
         return 0;

mooncake-transfer-engine/src/transfer_engine_c.cpp

@@ -19,6 +19,9 @@
 
 #include "transfer_engine.h"
 #include "transport/transport.h"
+#ifdef USE_EFA
+#include "transport/efa_transport/efa_transport.h"
+#endif
 
 using namespace mooncake;
 
@@ -36,6 +39,11 @@ transfer_engine_t createTransferEngine(const char *metadata_conn_string,
     return (transfer_engine_t)native;
 }
 
+int discoverTopology(transfer_engine_t engine) {
+    TransferEngine *native = (TransferEngine *)engine;
+    return native->getLocalTopology()->discover({});
+}
+
 int getLocalIpAndPort(transfer_engine_t engine, char *buf_out, size_t buf_len) {
     TransferEngine *native = (TransferEngine *)engine;
     auto str = native->getLocalIpAndPort();
@@ -77,6 +85,21 @@ int closeSegment(transfer_engine_t engine, segment_id_t segment_id) {
     return native->closeSegment(segment_id);
 }
 
+int warmupEfaSegment(transfer_engine_t engine, const char *segment_name) {
+#ifdef USE_EFA
+    TransferEngine *native = (TransferEngine *)engine;
+    auto *t = native->getTransport("efa");
+    if (!t) return 0;  // Non-EFA build or EFA not installed; nothing to do.
+    auto *efa = dynamic_cast<EfaTransport *>(t);
+    if (!efa) return 0;
+    return efa->warmupSegment(segment_name ? segment_name : "");
+#else
+    (void)engine;
+    (void)segment_name;
+    return 0;
+#endif
+}
+
 int removeLocalSegment(transfer_engine_t engine, const char *segment_name) {
     TransferEngine *native = (TransferEngine *)engine;
     return native->removeLocalSegment(segment_name);

mooncake-transfer-engine/src/transfer_engine_impl.cpp

@@ -29,7 +29,9 @@
 
 #include "transfer_metadata_plugin.h"
 #include "transport/transport.h"
+#ifdef USE_BAREX
 #include "transport/barex_transport/barex_transport.h"
+#endif
 
 namespace mooncake {
 

mooncake-transfer-engine/src/transport/efa_transport/efa_context.cpp

@@ -27,6 +27,7 @@
 #include <sstream>
 
 #include "config.h"
+#include "cuda_alike.h"
 #include "transport/efa_transport/efa_endpoint.h"
 #include "transport/efa_transport/efa_transport.h"
 #include "transport/transport.h"
@@ -34,41 +35,58 @@
 namespace mooncake {
 
 // EfaEndpointStore implementation
+
+EfaEndpointStore::EfaEndpointStore(size_t max_endpoints,
+                                   double inactive_timeout_sec)
+    : max_endpoints_(max_endpoints),
+      inactive_timeout_sec_(inactive_timeout_sec) {}
+
 std::shared_ptr<EfaEndPoint> EfaEndpointStore::get(
-    const std::string &peer_nic_path) {
+    const std::string& peer_nic_path) {
     RWSpinlock::ReadGuard guard(lock_);
     auto it = endpoints_.find(peer_nic_path);
     if (it != endpoints_.end()) {
+        it->second->set_active(true);
         return it->second;
     }
     return nullptr;
 }
 
 std::shared_ptr<EfaEndPoint> EfaEndpointStore::getOrInsert(
-    const std::string &peer_nic_path, std::shared_ptr<EfaEndPoint> new_ep) {
+    const std::string& peer_nic_path, std::shared_ptr<EfaEndPoint> new_ep) {
     RWSpinlock::WriteGuard guard(lock_);
     auto it = endpoints_.find(peer_nic_path);
     if (it != endpoints_.end()) {
+        it->second->set_active(true);
         return it->second;  // Another thread already created it
     }
+    // Evict stale endpoints if at capacity
+    if (endpoints_.size() >= max_endpoints_) {
+        size_t evicted = evictStaleLocked();
+        if (evicted == 0 && endpoints_.size() >= max_endpoints_) {
+            LOG(WARNING) << "EfaEndpointStore at capacity (" << max_endpoints_
+                         << ") with no stale endpoints to evict";
+        }
+    }
     endpoints_[peer_nic_path] = new_ep;
+    new_ep->set_active(true);
     return new_ep;
 }
 
-void EfaEndpointStore::add(const std::string &peer_nic_path,
+void EfaEndpointStore::add(const std::string& peer_nic_path,
                            std::shared_ptr<EfaEndPoint> endpoint) {
     RWSpinlock::WriteGuard guard(lock_);
     endpoints_[peer_nic_path] = endpoint;
 }
 
-void EfaEndpointStore::remove(const std::string &peer_nic_path) {
+void EfaEndpointStore::remove(const std::string& peer_nic_path) {
     RWSpinlock::WriteGuard guard(lock_);
     endpoints_.erase(peer_nic_path);
 }
 
 int EfaEndpointStore::disconnectAll() {
     RWSpinlock::WriteGuard guard(lock_);
-    for (auto &entry : endpoints_) {
+    for (auto& entry : endpoints_) {
         if (entry.second) {
             entry.second->disconnect();
         }
@@ -81,8 +99,53 @@ size_t EfaEndpointStore::size() const {
     return endpoints_.size();
 }
 
+size_t EfaEndpointStore::evictStale() {
+    RWSpinlock::WriteGuard guard(lock_);
+    return evictStaleLocked();
+}
+
+size_t EfaEndpointStore::evictStaleLocked() {
+    size_t evicted = 0;
+    for (auto it = endpoints_.begin(); it != endpoints_.end();) {
+        auto& ep = it->second;
+        double age = ep ? ep->lastUsedAge() : 0;
+        bool outstanding = ep ? ep->hasOutstandingSlice() : false;
+        if (ep && !outstanding && age > inactive_timeout_sec_) {
+            LOG(INFO) << "Evicting stale EFA endpoint: " << it->first
+                      << " (idle " << age
+                      << "s, timeout=" << inactive_timeout_sec_ << "s)";
+            ep->disconnect();
+            it = endpoints_.erase(it);
+            ++evicted;
+        } else {
+            ++it;
+        }
+    }
+    if (endpoints_.size() > 0) {
+        VLOG(1) << "evictStale: " << endpoints_.size() << " endpoints, "
+                << evicted << " evicted";
+    }
+    return evicted;
+}
+
+size_t EfaEndpointStore::removeDisconnected() {
+    RWSpinlock::WriteGuard guard(lock_);
+    size_t removed = 0;
+    for (auto it = endpoints_.begin(); it != endpoints_.end();) {
+        auto& ep = it->second;
+        if (ep && !ep->connected() && !ep->hasOutstandingSlice()) {
+            LOG(INFO) << "Removing disconnected EFA endpoint: " << it->first;
+            it = endpoints_.erase(it);
+            ++removed;
+        } else {
+            ++it;
+        }
+    }
+    return removed;
+}
+
 // EfaContext implementation
-EfaContext::EfaContext(EfaTransport &engine, const std::string &device_name)
+EfaContext::EfaContext(EfaTransport& engine, const std::string& device_name)
     : engine_(engine),
       device_name_(device_name),
       fi_info_(nullptr),
@@ -99,7 +162,16 @@ EfaContext::~EfaContext() {
 int EfaContext::construct(size_t num_cq_list, size_t num_comp_channels,
                           uint8_t port, int gid_index, size_t max_cqe,
                           int max_endpoints) {
-    endpoint_store_ = std::make_shared<EfaEndpointStore>();
+    endpoint_store_ = std::make_shared<EfaEndpointStore>(max_endpoints);
+
+#if !defined(USE_CUDA) && !defined(USE_HIP)
+    // When built without GPU support, prevent libfabric's EFA provider from
+    // dlopen-ing libcudart/libcuda at fi_getinfo/fi_domain time. That
+    // initialization creates a CUDA primary context on GPU 0 and leaks
+    // ~616 MiB of device memory even when no GPU memory is ever registered.
+    // Only set if the user hasn't explicitly configured FI_HMEM.
+    setenv("FI_HMEM", "system", 0);
+#endif
 
     // Setup hints for EFA provider
     hints_ = fi_allocinfo();
@@ -117,8 +189,12 @@ int EfaContext::construct(size_t num_cq_list, size_t num_comp_channels,
     // Specify the domain (device) name - append "-rdm" for RDM endpoint
     std::string domain_name = device_name_ + "-rdm";
     hints_->domain_attr->name = strdup(domain_name.c_str());
-    hints_->domain_attr->mr_mode =
-        FI_MR_LOCAL | FI_MR_VIRT_ADDR | FI_MR_ALLOCATED | FI_MR_PROV_KEY;
+    hints_->domain_attr->mr_mode = FI_MR_LOCAL | FI_MR_VIRT_ADDR |
+                                   FI_MR_ALLOCATED | FI_MR_PROV_KEY
+#if defined(USE_CUDA) || defined(USE_HIP)
+                                   | FI_MR_HMEM
+#endif
+        ;
     hints_->domain_attr->threading = FI_THREAD_SAFE;
 
     // Get fabric info
@@ -208,7 +284,7 @@ int EfaContext::deconstruct() {
 
     {
         RWSpinlock::WriteGuard guard(mr_lock_);
-        for (auto &entry : mr_map_) {
+        for (auto& entry : mr_map_) {
             if (entry.second.mr) {
                 fi_close(&entry.second.mr->fid);
             }
@@ -216,7 +292,7 @@ int EfaContext::deconstruct() {
         mr_map_.clear();
     }
 
-    for (auto &cq : cq_list_) {
+    for (auto& cq : cq_list_) {
         if (cq && cq->cq) {
             fi_close(&cq->cq->fid);
             cq->cq = nullptr;
@@ -252,34 +328,74 @@ int EfaContext::deconstruct() {
     return 0;
 }
 
-int EfaContext::registerMemoryRegionInternal(void *addr, size_t length,
+int EfaContext::registerMemoryRegionInternal(void* addr, size_t length,
                                              int access,
-                                             EfaMemoryRegionMeta &mrMeta) {
+                                             EfaMemoryRegionMeta& mrMeta) {
     if (length > (size_t)globalConfig().max_mr_size) {
-        PLOG(WARNING) << "The buffer length exceeds device max_mr_size, "
-                      << "shrink it to " << globalConfig().max_mr_size;
-        length = (size_t)globalConfig().max_mr_size;
+        LOG(ERROR) << "Buffer length " << length
+                   << " exceeds device max_mr_size "
+                   << globalConfig().max_mr_size
+                   << ". Use EfaTransport::registerLocalMemory() which "
+                      "auto-splits large buffers.";
+        return ERR_CONTEXT;
     }
 
     mrMeta.addr = addr;
     mrMeta.length = length;
 
-    // Convert access flags to libfabric flags
-    uint64_t fi_access = 0;
-    if (access & FI_READ) fi_access |= FI_READ;
-    if (access & FI_WRITE) fi_access |= FI_WRITE;
-    if (access & FI_REMOTE_READ) fi_access |= FI_REMOTE_READ;
-    if (access & FI_REMOTE_WRITE) fi_access |= FI_REMOTE_WRITE;
-
     // For EFA, we need local read/write and remote read/write
-    fi_access = FI_READ | FI_WRITE | FI_REMOTE_READ | FI_REMOTE_WRITE;
+    uint64_t fi_access = FI_READ | FI_WRITE | FI_REMOTE_READ | FI_REMOTE_WRITE;
 
-    int ret = fi_mr_reg(domain_, addr, length, fi_access, 0, 0, 0, &mrMeta.mr,
+    // Detect memory type and use fi_mr_regattr() for GPU memory.
+    // The EFA provider's fi_mr_reg() hardcodes iface=FI_HMEM_SYSTEM,
+    // so GPU memory must go through fi_mr_regattr() with explicit
+    // iface/device fields (per libfabric spec and EFA provider impl).
+    enum fi_hmem_iface iface = FI_HMEM_SYSTEM;
+    int device_ordinal = 0;
+
+#if defined(USE_CUDA)
+    cudaPointerAttributes attributes;
+    cudaError_t cuda_ret = cudaPointerGetAttributes(&attributes, addr);
+    if (cuda_ret == cudaSuccess && attributes.type == cudaMemoryTypeDevice) {
+        iface = FI_HMEM_CUDA;
+        device_ordinal = attributes.device;
+    }
+#elif defined(USE_HIP)
+    hipPointerAttribute_t attributes;
+    hipError_t hip_ret = hipPointerGetAttributes(&attributes, addr);
+    if (hip_ret == hipSuccess && attributes.type == hipMemoryTypeDevice) {
+        iface = FI_HMEM_ROCR;
+        device_ordinal = attributes.device;
+    }
+#endif
+
+    int ret;
+    if (iface != FI_HMEM_SYSTEM) {
+        // GPU memory: use fi_mr_regattr with explicit iface and device
+        struct iovec iov = {.iov_base = addr, .iov_len = length};
+        struct fi_mr_attr attr = {};
+        attr.mr_iov = &iov;
+        attr.iov_count = 1;
+        attr.access = fi_access;
+        attr.iface = iface;
+        attr.device.cuda = device_ordinal;
+
+        ret = fi_mr_regattr(domain_, &attr, 0, &mrMeta.mr);
+        if (ret) {
+            LOG(ERROR) << "fi_mr_regattr failed for GPU memory " << addr
+                       << " (device " << device_ordinal
+                       << "): " << fi_strerror(-ret);
+            return ERR_CONTEXT;
+        }
+    } else {
+        // CPU memory: fi_mr_reg is sufficient
+        ret = fi_mr_reg(domain_, addr, length, fi_access, 0, 0, 0, &mrMeta.mr,
                         nullptr);
-    if (ret) {
-        LOG(ERROR) << "fi_mr_reg failed for " << addr << ": "
-                   << fi_strerror(-ret);
-        return ERR_CONTEXT;
+        if (ret) {
+            LOG(ERROR) << "fi_mr_reg failed for " << addr << ": "
+                       << fi_strerror(-ret);
+            return ERR_CONTEXT;
+        }
     }
 
     mrMeta.key = fi_mr_key(mrMeta.mr);
@@ -287,7 +403,7 @@ int EfaContext::registerMemoryRegionInternal(void *addr, size_t length,
     return 0;
 }
 
-int EfaContext::registerMemoryRegion(void *addr, size_t length, int access) {
+int EfaContext::registerMemoryRegion(void* addr, size_t length, int access) {
     EfaMemoryRegionMeta mrMeta;
     int ret = registerMemoryRegionInternal(addr, length, access, mrMeta);
     if (ret != 0) {
@@ -298,7 +414,7 @@ int EfaContext::registerMemoryRegion(void *addr, size_t length, int access) {
     return 0;
 }
 
-int EfaContext::unregisterMemoryRegion(void *addr) {
+int EfaContext::unregisterMemoryRegion(void* addr) {
     RWSpinlock::WriteGuard guard(mr_lock_);
     auto it = mr_map_.find((uint64_t)addr);
     if (it != mr_map_.end()) {
@@ -315,72 +431,89 @@ int EfaContext::unregisterMemoryRegion(void *addr) {
     return 0;
 }
 
-int EfaContext::preTouchMemory(void *addr, size_t length) {
-    volatile char *ptr = (volatile char *)addr;
+int EfaContext::preTouchMemory(void* addr, size_t length) {
+    volatile char* ptr = (volatile char*)addr;
     for (size_t i = 0; i < length; i += 4096) {
         ptr[i] = ptr[i];
     }
     return 0;
 }
 
-uint64_t EfaContext::rkey(void *addr) {
+uint64_t EfaContext::rkey(void* addr) {
     RWSpinlock::ReadGuard guard(mr_lock_);
-    auto it = mr_map_.find((uint64_t)addr);
-    if (it != mr_map_.end() && it->second.mr) {
-        return it->second.key;
+    auto it = mr_map_.upper_bound((uint64_t)addr);
+    if (it != mr_map_.begin()) {
+        --it;
+        if ((uint64_t)addr < it->first + it->second.length && it->second.mr) {
+            return it->second.key;
+        }
     }
     return 0;
 }
 
-uint64_t EfaContext::lkey(void *addr) {
+uint64_t EfaContext::lkey(void* addr) {
     RWSpinlock::ReadGuard guard(mr_lock_);
-    auto it = mr_map_.find((uint64_t)addr);
-    if (it != mr_map_.end() && it->second.mr) {
-        return fi_mr_key(it->second.mr);
+    auto it = mr_map_.upper_bound((uint64_t)addr);
+    if (it != mr_map_.begin()) {
+        --it;
+        if ((uint64_t)addr < it->first + it->second.length && it->second.mr) {
+            return fi_mr_key(it->second.mr);
+        }
     }
     return 0;
 }
 
-void *EfaContext::mrDesc(void *addr) {
+void* EfaContext::mrDesc(void* addr) {
     RWSpinlock::ReadGuard guard(mr_lock_);
-    // Find the MR that contains this address
-    for (auto &entry : mr_map_) {
-        if ((uint64_t)addr >= entry.first &&
-            (uint64_t)addr < entry.first + entry.second.length) {
-            if (entry.second.mr) {
-                return fi_mr_desc(entry.second.mr);
-            }
+    auto it = mr_map_.upper_bound((uint64_t)addr);
+    if (it != mr_map_.begin()) {
+        --it;
+        if ((uint64_t)addr < it->first + it->second.length && it->second.mr) {
+            return fi_mr_desc(it->second.mr);
         }
     }
     return nullptr;
 }
 
 std::shared_ptr<EfaEndPoint> EfaContext::endpoint(
-    const std::string &peer_nic_path) {
+    const std::string& peer_nic_path) {
     if (!endpoint_store_) return nullptr;
 
-    // Fast path: endpoint already exists
-    auto ep = endpoint_store_->get(peer_nic_path);
-    if (ep) return ep;
+    // Use normalized key (strip port) so the same physical peer reuses its
+    // endpoint across reconnections.  Each P2PHANDSHAKE run picks a random
+    // port, producing a different peer_nic_path for the same peer host+NIC.
+    std::string key = normalizeNicPath(peer_nic_path);
+
+    // Fast path: endpoint already exists for this physical peer.
+    // Update peer_nic_path in case the port changed (new initiator run).
+    // setPeerNicPath disconnects old connection (fi_av_remove) if needed.
+    auto ep = endpoint_store_->get(key);
+    if (ep) {
+        ep->setPeerNicPath(peer_nic_path);
+        return ep;
+    }
 
     // Slow path: create new endpoint, then atomically insert (or get existing
     // if another thread raced us).  getOrInsert prevents duplicate endpoints
     // and duplicate AV entries for the same peer.
     auto new_endpoint = std::make_shared<EfaEndPoint>(*this);
-    if (!cq_list_.empty() && cq_list_[0]) {
-        int ret = new_endpoint->construct(cq_list_[0]->cq);
+    auto cq = nextCq();
+    if (cq) {
+        int ret = new_endpoint->construct(cq->cq, &cq->outstanding, 1, 4,
+                                          globalConfig().max_wr, 64);
         if (ret != 0) {
             LOG(ERROR) << "Failed to construct EFA endpoint";
             return nullptr;
         }
     }
+    // Still set the full peer_nic_path (with port) for handshake routing
     new_endpoint->setPeerNicPath(peer_nic_path);
-    ep = endpoint_store_->getOrInsert(peer_nic_path, new_endpoint);
+    ep = endpoint_store_->getOrInsert(key, new_endpoint);
     // If another thread won the race, new_endpoint is discarded (RAII cleanup)
     return ep;
 }
 
-int EfaContext::deleteEndpoint(const std::string &peer_nic_path) {
+int EfaContext::deleteEndpoint(const std::string& peer_nic_path) {
     if (endpoint_store_) {
         endpoint_store_->remove(peer_nic_path);
     }
@@ -409,7 +542,7 @@ std::string EfaContext::localAddr() const {
     }
 
     std::ostringstream oss;
-    const uint8_t *addr = static_cast<const uint8_t *>(fi_info_->src_addr);
+    const uint8_t* addr = static_cast<const uint8_t*>(fi_info_->src_addr);
     for (size_t i = 0; i < fi_info_->src_addrlen; ++i) {
         oss << std::hex << std::setw(2) << std::setfill('0') << (int)addr[i];
     }
@@ -417,17 +550,26 @@ std::string EfaContext::localAddr() const {
 }
 
 int EfaContext::submitPostSend(
-    const std::vector<Transport::Slice *> &slice_list) {
+    const std::vector<Transport::Slice*>& slice_list) {
     // Route slices to appropriate endpoints for sending
     // Group slices by peer NIC path
-    std::unordered_map<std::string, std::vector<Transport::Slice *>>
+    std::unordered_map<std::string, std::vector<Transport::Slice*>>
         slices_by_peer;
-    std::vector<Transport::Slice *> failed_slices;
+    std::vector<Transport::Slice*> failed_slices;
 
-    for (auto *slice : slice_list) {
+    for (auto* slice : slice_list) {
         if (!slice) continue;
 
-        // Get peer segment descriptor to find dest_rkey and peer device info
+        // Fast path: peer info already resolved by submitTransferTask's
+        // striping path (dest_rkey and peer_nic_path pre-set on slice).
+        // This eliminates per-slice metadata lookup, selectDevice(), and
+        // string construction — the main bottleneck for multi-NIC striping.
+        if (!slice->peer_nic_path.empty()) {
+            slices_by_peer[slice->peer_nic_path].push_back(slice);
+            continue;
+        }
+
+        // Slow path: resolve peer info per-slice (non-striped transfers)
         auto peer_segment_desc =
             engine_.meta()->getSegmentDescByID(slice->target_id);
         if (!peer_segment_desc) {
@@ -443,7 +585,7 @@ int EfaContext::submitPostSend(
                                        slice->rdma.dest_addr, slice->length,
                                        buffer_id, device_id)) {
             LOG(ERROR) << "Cannot select device for dest_addr "
-                       << (void *)slice->rdma.dest_addr;
+                       << (void*)slice->rdma.dest_addr;
             slice->markFailed();
             continue;
         }
@@ -461,26 +603,26 @@ int EfaContext::submitPostSend(
     }
 
     // Now send to each peer endpoint
-    for (auto &entry : slices_by_peer) {
-        const std::string &peer_nic_path = entry.first;
-        auto &peer_slices = entry.second;
+    for (auto& entry : slices_by_peer) {
+        const std::string& peer_nic_path = entry.first;
+        auto& peer_slices = entry.second;
 
         // Get or create endpoint for this peer
         auto ep = endpoint(peer_nic_path);
         if (!ep) {
             LOG(ERROR) << "Cannot create endpoint for peer " << peer_nic_path;
-            for (auto *slice : peer_slices) {
+            for (auto* slice : peer_slices) {
                 slice->markFailed();
             }
             continue;
         }
 
         // Submit to endpoint
-        std::vector<Transport::Slice *> failed_slice_list;
+        std::vector<Transport::Slice*> failed_slice_list;
         ep->submitPostSend(peer_slices, failed_slice_list);
 
         // Handle any slices that failed to post
-        for (auto *slice : failed_slice_list) {
+        for (auto* slice : failed_slice_list) {
             slice->markFailed();
         }
     }
@@ -493,7 +635,7 @@ int EfaContext::pollCq(int max_entries, int cq_index) {
         return 0;
     }
 
-    struct fid_cq *cq = cq_list_[cq_index]->cq;
+    struct fid_cq* cq = cq_list_[cq_index]->cq;
     if (!cq) return 0;
 
     // Use fi_cq_data format for completions
@@ -504,10 +646,10 @@ int EfaContext::pollCq(int max_entries, int cq_index) {
 
     if (ret > 0) {
         // Process completions outside the lock (markSuccess / delete are safe)
-        std::unordered_map<volatile int *, int> wr_depth_set;
+        std::unordered_map<volatile int*, int> wr_depth_set;
         for (ssize_t i = 0; i < ret; i++) {
-            EfaOpContext *op_ctx =
-                reinterpret_cast<EfaOpContext *>(entries[i].op_context);
+            EfaOpContext* op_ctx =
+                reinterpret_cast<EfaOpContext*>(entries[i].op_context);
             if (op_ctx && op_ctx->slice) {
                 op_ctx->slice->markSuccess();
                 if (op_ctx->wr_depth) {
@@ -516,7 +658,7 @@ int EfaContext::pollCq(int max_entries, int cq_index) {
                 delete op_ctx;
             }
         }
-        for (auto &entry : wr_depth_set) {
+        for (auto& entry : wr_depth_set) {
             __sync_fetch_and_sub(entry.first, entry.second);
         }
         __sync_fetch_and_sub(&cq_list_[cq_index]->outstanding,
@@ -528,11 +670,11 @@ int EfaContext::pollCq(int max_entries, int cq_index) {
         // CQ error - drain all queued error entries under the domain lock
         int err_count = 0;
         struct fi_cq_err_entry err_entry;
-        std::unordered_map<volatile int *, int> wr_depth_set;
+        std::unordered_map<volatile int*, int> wr_depth_set;
 
         while ((ret = fi_cq_readerr(cq, &err_entry, 0)) > 0) {
-            EfaOpContext *op_ctx =
-                reinterpret_cast<EfaOpContext *>(err_entry.op_context);
+            EfaOpContext* op_ctx =
+                reinterpret_cast<EfaOpContext*>(err_entry.op_context);
             if (op_ctx && op_ctx->slice) {
                 LOG(ERROR) << "EFA CQ error: "
                            << fi_cq_strerror(cq, err_entry.prov_errno,
@@ -547,7 +689,7 @@ int EfaContext::pollCq(int max_entries, int cq_index) {
             err_count++;
         }
 
-        for (auto &entry : wr_depth_set) {
+        for (auto& entry : wr_depth_set) {
             __sync_fetch_and_sub(entry.first, entry.second);
         }
         if (err_count > 0) {

mooncake-transfer-engine/src/transport/efa_transport/efa_endpoint.cpp

@@ -27,7 +27,7 @@
 
 namespace mooncake {
 
-EfaEndPoint::EfaEndPoint(EfaContext &context)
+EfaEndPoint::EfaEndPoint(EfaContext& context)
     : context_(context),
       status_(INITIALIZING),
       ep_(nullptr),
@@ -39,14 +39,16 @@ EfaEndPoint::EfaEndPoint(EfaContext &context)
       max_wr_depth_(0),
       cq_outstanding_(nullptr),
       active_(true),
-      inactive_time_(0) {}
+      inactive_time_(0),
+      last_used_time_(getCurrentTimeInNano()) {}
 
 EfaEndPoint::~EfaEndPoint() {
     if (ep_) deconstruct();
 }
 
-int EfaEndPoint::construct(struct fid_cq *cq, size_t num_qp_list,
-                           size_t max_sge, size_t max_wr, size_t max_inline) {
+int EfaEndPoint::construct(struct fid_cq* cq, volatile int* cq_outstanding,
+                           size_t num_qp_list, size_t max_sge, size_t max_wr,
+                           size_t max_inline) {
     if (status_.load(std::memory_order_relaxed) != INITIALIZING) {
         LOG(ERROR) << "EFA Endpoint has already been constructed";
         return ERR_ENDPOINT;
@@ -55,7 +57,7 @@ int EfaEndPoint::construct(struct fid_cq *cq, size_t num_qp_list,
     tx_cq_ = cq;
     rx_cq_ = cq;  // Use same CQ for TX and RX
     max_wr_depth_ = max_wr;
-    cq_outstanding_ = context_.cqOutstandingCount(0);
+    cq_outstanding_ = cq_outstanding;
 
     // Create endpoint
     int ret = fi_endpoint(context_.domain(), context_.info(), &ep_, nullptr);
@@ -126,10 +128,12 @@ int EfaEndPoint::deconstruct() {
 
 int EfaEndPoint::destroyQP() { return deconstruct(); }
 
-void EfaEndPoint::setPeerNicPath(const std::string &peer_nic_path) {
+void EfaEndPoint::setPeerNicPath(const std::string& peer_nic_path) {
     RWSpinlock::WriteGuard guard(lock_);
+    if (peer_nic_path_ == peer_nic_path) return;  // No change
     if (connected()) {
-        LOG(WARNING) << "Previous EFA connection will be discarded";
+        LOG(INFO) << "Peer reconnected with new address, re-establishing: "
+                  << peer_nic_path_ << " -> " << peer_nic_path;
         disconnectUnlocked();
     }
     peer_nic_path_ = peer_nic_path;
@@ -144,7 +148,7 @@ std::string EfaEndPoint::getLocalAddr() const {
     return oss.str();
 }
 
-int EfaEndPoint::insertPeerAddr(const std::string &peer_addr) {
+int EfaEndPoint::insertPeerAddr(const std::string& peer_addr) {
     // Convert hex string to binary address
     std::vector<uint8_t> addr_bin;
     addr_bin.reserve(peer_addr.size() / 2);
@@ -215,13 +219,14 @@ int EfaEndPoint::setupConnectionsByActive() {
     }
 
     status_.store(CONNECTED, std::memory_order_release);
+    touchLastUsed();
     VLOG(1) << "EFA connection established: " << toString()
             << " peer_fi_addr=" << peer_fi_addr_;
     return 0;
 }
 
-int EfaEndPoint::setupConnectionsByPassive(const HandShakeDesc &peer_desc,
-                                           HandShakeDesc &local_desc) {
+int EfaEndPoint::setupConnectionsByPassive(const HandShakeDesc& peer_desc,
+                                           HandShakeDesc& local_desc) {
     RWSpinlock::WriteGuard guard(lock_);
     if (connected()) {
         LOG(WARNING) << "Re-establish EFA connection: " << toString();
@@ -259,6 +264,7 @@ int EfaEndPoint::setupConnectionsByPassive(const HandShakeDesc &peer_desc,
     // reply_msg should be empty on success
 
     status_.store(CONNECTED, std::memory_order_release);
+    touchLastUsed();
     VLOG(1) << "EFA connection established (passive): " << toString();
     return 0;
 }
@@ -269,8 +275,16 @@ void EfaEndPoint::disconnect() {
 }
 
 void EfaEndPoint::disconnectUnlocked() {
-    // For EFA RDM endpoints, we don't need to reset QP state
-    // Just remove peer from AV if needed and mark as disconnected
+    // Remove peer address from AV to free the slot for reuse.
+    // Without this, AV entries leak on every initiator disconnect,
+    // eventually exhausting the AV and causing fi_av_insert/fi_endpoint
+    // failures ("No data available") after 3-5 initiator runs.
+    if (peer_fi_addr_ != FI_ADDR_UNSPEC) {
+        int ret = fi_av_remove(context_.av(), &peer_fi_addr_, 1, 0);
+        if (ret) {
+            LOG(WARNING) << "fi_av_remove failed: " << fi_strerror(-ret);
+        }
+    }
     peer_fi_addr_ = FI_ADDR_UNSPEC;
     status_.store(UNCONNECTED, std::memory_order_release);
 }
@@ -281,8 +295,8 @@ const std::string EfaEndPoint::toString() const {
 
 bool EfaEndPoint::hasOutstandingSlice() const { return wr_depth_ > 0; }
 
-int EfaEndPoint::doSetupConnection(const std::string &peer_addr,
-                                   std::string *reply_msg) {
+int EfaEndPoint::doSetupConnection(const std::string& peer_addr,
+                                   std::string* reply_msg) {
     int ret = insertPeerAddr(peer_addr);
     if (ret != 0) {
         if (reply_msg) *reply_msg = "Failed to insert peer address into AV";
@@ -294,14 +308,14 @@ int EfaEndPoint::doSetupConnection(const std::string &peer_addr,
 }
 
 int EfaEndPoint::submitPostSend(
-    std::vector<Transport::Slice *> &slice_list,
-    std::vector<Transport::Slice *> &failed_slice_list) {
+    std::vector<Transport::Slice*>& slice_list,
+    std::vector<Transport::Slice*>& failed_slice_list) {
     if (!connected()) {
         // Try to establish connection first
         int ret = setupConnectionsByActive();
         if (ret != 0) {
             // Move all slices to failed list
-            for (auto *slice : slice_list) {
+            for (auto* slice : slice_list) {
                 failed_slice_list.push_back(slice);
             }
             slice_list.clear();
@@ -309,130 +323,174 @@ int EfaEndPoint::submitPostSend(
         }
     }
 
-    // Process slices - using fi_write for RDMA write operations.
-    // Use atomic reserve-before-post to prevent CQ overflow when multiple
-    // threads post to endpoints sharing the same CQ.  The CQ has a fixed
-    // capacity (max_cqe); if more completions arrive than it can hold, the
-    // provider silently drops them and those slices never complete (hang).
+    // Batched submission — mirrors RDMA's ibv_post_send approach:
+    // 1. Reserve N WR+CQ slots in bulk (single CAS each)
+    // 2. Prepare MR descriptors and op contexts outside the lock
+    // 3. Hold post_lock once for the entire batch of fi_write calls
+    // This reduces per-slice overhead from 2 CAS + 1 spinlock to amortized
+    // 2/N CAS + 1/N spinlock per slice.
     const int kMaxBackoffYields = 100000;
     const int cq_limit = static_cast<int>(globalConfig().max_cqe);
 
-    for (auto it = slice_list.begin(); it != slice_list.end();) {
-        // --- Atomically reserve CQ and WR capacity before posting ---
-        // This eliminates the TOCTOU race where multiple threads pass the
-        // capacity check simultaneously and collectively overflow the CQ.
+    struct BatchEntry {
+        Transport::Slice* slice;
+        void* local_desc;
+        EfaOpContext* op_ctx;
+    };
+
+    while (!slice_list.empty()) {
+        // --- Step 1: Atomically reserve a batch of WR+CQ slots ---
+        int batch_count = 0;
         int backoff = 0;
-        bool reserved = false;
-        while (!reserved) {
-            // Try to reserve one WR slot
+        bool timed_out = false;
+        while (batch_count == 0) {
             int cur_wr = wr_depth_;
-            if (cur_wr >= max_wr_depth_) {
-                if (++backoff > kMaxBackoffYields) goto timeout;
+            int wr_avail = max_wr_depth_ - cur_wr;
+            if (wr_avail <= 0) {
+                if (++backoff > kMaxBackoffYields) {
+                    timed_out = true;
+                    break;
+                }
                 std::this_thread::yield();
                 continue;
             }
-            if (!__sync_bool_compare_and_swap(&wr_depth_, cur_wr, cur_wr + 1)) {
-                continue;  // CAS failed, retry immediately
-            }
-            // WR slot reserved. Now try to reserve CQ slot.
+            int want = std::min(wr_avail, (int)slice_list.size());
             if (cq_outstanding_) {
                 int cur_cq = *cq_outstanding_;
-                while (cur_cq < cq_limit) {
-                    if (__sync_bool_compare_and_swap(cq_outstanding_, cur_cq,
-                                                     cur_cq + 1)) {
-                        reserved = true;
+                int cq_avail = cq_limit - cur_cq;
+                if (cq_avail <= 0) {
+                    if (++backoff > kMaxBackoffYields) {
+                        timed_out = true;
                         break;
                     }
-                    cur_cq = *cq_outstanding_;
-                }
-                if (!reserved) {
-                    // CQ full - release WR reservation and back off
-                    __sync_fetch_and_sub(&wr_depth_, 1);
-                    if (++backoff > kMaxBackoffYields) goto timeout;
                     std::this_thread::yield();
                     continue;
                 }
+                want = std::min(want, cq_avail);
+                if (!__sync_bool_compare_and_swap(&wr_depth_, cur_wr,
+                                                  cur_wr + want)) {
+                    continue;
+                }
+                cur_cq = *cq_outstanding_;
+                cq_avail = cq_limit - cur_cq;
+                if (cq_avail < want) {
+                    __sync_fetch_and_sub(&wr_depth_, want);
+                    continue;
+                }
+                if (!__sync_bool_compare_and_swap(cq_outstanding_, cur_cq,
+                                                  cur_cq + want)) {
+                    __sync_fetch_and_sub(&wr_depth_, want);
+                    continue;
+                }
             } else {
-                reserved = true;
+                if (!__sync_bool_compare_and_swap(&wr_depth_, cur_wr,
+                                                  cur_wr + want)) {
+                    continue;
+                }
             }
+            batch_count = want;
         }
 
-        {
-            Transport::Slice *slice = *it;
+        if (timed_out) {
+            LOG(WARNING) << "EFA submitPostSend: timed out waiting for CQ drain"
+                         << " (wr_depth=" << wr_depth_
+                         << ", max=" << max_wr_depth_ << ", cq_outstanding="
+                         << (cq_outstanding_ ? *cq_outstanding_ : -1)
+                         << ", max_cqe=" << cq_limit << ")";
+            for (auto* slice : slice_list) {
+                failed_slice_list.push_back(slice);
+            }
+            slice_list.clear();
+            return 0;
+        }
 
-            // Get memory region descriptor for the local buffer
-            void *local_desc = context_.mrDesc(slice->source_addr);
+        // --- Step 2: Prepare MR descriptors outside the lock ---
+        std::vector<BatchEntry> batch(batch_count);
+        int valid_count = 0;
+
+        for (int i = 0; i < batch_count; i++) {
+            Transport::Slice* slice = slice_list[i];
+            void* local_desc = context_.mrDesc(slice->source_addr);
             if (!local_desc) {
                 LOG(ERROR) << "No MR descriptor found for address "
                            << slice->source_addr;
-                // Release reservations
-                __sync_fetch_and_sub(&wr_depth_, 1);
-                if (cq_outstanding_) __sync_fetch_and_sub(cq_outstanding_, 1);
                 failed_slice_list.push_back(slice);
-                it = slice_list.erase(it);
                 continue;
             }
-
-            // Allocate operation context to track the slice for completion
-            // Note: This memory is freed after CQ completion in pollCq
-            EfaOpContext *op_ctx = new EfaOpContext();
+            EfaOpContext* op_ctx = new EfaOpContext();
             memset(op_ctx, 0, sizeof(EfaOpContext));
             op_ctx->slice = slice;
             op_ctx->wr_depth = &wr_depth_;
 
-            // Serialize fi_write per-endpoint: concurrent fi_write on the
-            // same RDM endpoint corrupts provider state.  Cross-endpoint
-            // safety is handled by the FI_THREAD_SAFE hint.
+            batch[valid_count++] = {slice, local_desc, op_ctx};
+        }
+
+        int mr_failures = batch_count - valid_count;
+        if (mr_failures > 0) {
+            __sync_fetch_and_sub(&wr_depth_, mr_failures);
+            if (cq_outstanding_)
+                __sync_fetch_and_sub(cq_outstanding_, mr_failures);
+        }
+
+        // --- Step 3: Post entire batch under a single lock hold ---
+        std::vector<Transport::Slice*> retry_slices;
+        if (valid_count > 0) {
             while (post_lock_.test_and_set(std::memory_order_acquire)) {
             }
-            ssize_t ret = fi_write(ep_,
-                                   (void *)slice->source_addr,  // local buffer
-                                   slice->length, local_desc, peer_fi_addr_,
-                                   slice->rdma.dest_addr,  // remote address
-                                   slice->rdma.dest_rkey,  // remote key
-                                   &op_ctx->fi_ctx);  // context for completion
-            post_lock_.clear(std::memory_order_release);
-
-            if (ret == 0) {
-                // Successfully posted - do NOT mark success here!
-                // Success is marked only after CQ completion in pollCq.
-                // WR and CQ reservations are already accounted for.
-                slice->status = Transport::Slice::PENDING;
-                it = slice_list.erase(it);
-            } else if (ret == -FI_EAGAIN) {
-                // Provider queue full - release reservations and retry
-                delete op_ctx;
-                __sync_fetch_and_sub(&wr_depth_, 1);
-                if (cq_outstanding_) __sync_fetch_and_sub(cq_outstanding_, 1);
-                std::this_thread::yield();
-                // Don't advance iterator - retry the same slice
-            } else {
-                // Hard error - release reservations
-                LOG(ERROR) << "fi_write failed: " << fi_strerror(-ret)
-                           << " (source=" << slice->source_addr
-                           << ", len=" << slice->length
-                           << ", dest=" << (void *)slice->rdma.dest_addr
-                           << ", rkey=" << slice->rdma.dest_rkey << ")";
-                delete op_ctx;
-                __sync_fetch_and_sub(&wr_depth_, 1);
-                if (cq_outstanding_) __sync_fetch_and_sub(cq_outstanding_, 1);
-                failed_slice_list.push_back(slice);
-                it = slice_list.erase(it);
+            for (int i = 0; i < valid_count; i++) {
+                auto& entry = batch[i];
+                ssize_t ret;
+                if (entry.slice->opcode == Transport::TransferRequest::READ) {
+                    ret = fi_read(ep_, (void*)entry.slice->source_addr,
+                                  entry.slice->length, entry.local_desc,
+                                  peer_fi_addr_, entry.slice->rdma.dest_addr,
+                                  entry.slice->rdma.dest_rkey,
+                                  &entry.op_ctx->fi_ctx);
+                } else {
+                    ret = fi_write(ep_, (void*)entry.slice->source_addr,
+                                   entry.slice->length, entry.local_desc,
+                                   peer_fi_addr_, entry.slice->rdma.dest_addr,
+                                   entry.slice->rdma.dest_rkey,
+                                   &entry.op_ctx->fi_ctx);
+                }
+                if (ret == 0) {
+                    entry.slice->status = Transport::Slice::PENDING;
+                } else if (ret == -FI_EAGAIN) {
+                    delete entry.op_ctx;
+                    int not_posted = valid_count - i;
+                    __sync_fetch_and_sub(&wr_depth_, not_posted);
+                    if (cq_outstanding_)
+                        __sync_fetch_and_sub(cq_outstanding_, not_posted);
+                    for (int j = i; j < valid_count; j++) {
+                        if (j > i) delete batch[j].op_ctx;
+                        retry_slices.push_back(batch[j].slice);
+                    }
+                    break;
+                } else {
+                    LOG(ERROR)
+                        << "fi_read/fi_write failed: " << fi_strerror(-ret)
+                        << " (source=" << entry.slice->source_addr
+                        << ", len=" << entry.slice->length
+                        << ", dest=" << (void*)entry.slice->rdma.dest_addr
+                        << ", rkey=" << entry.slice->rdma.dest_rkey << ")";
+                    delete entry.op_ctx;
+                    __sync_fetch_and_sub(&wr_depth_, 1);
+                    if (cq_outstanding_)
+                        __sync_fetch_and_sub(cq_outstanding_, 1);
+                    failed_slice_list.push_back(entry.slice);
+                }
             }
+            post_lock_.clear(std::memory_order_release);
+            touchLastUsed();
         }
-        continue;
 
-    timeout:
-        LOG(WARNING) << "EFA submitPostSend: timed out waiting for CQ drain"
-                     << " (wr_depth=" << wr_depth_ << ", max=" << max_wr_depth_
-                     << ", cq_outstanding="
-                     << (cq_outstanding_ ? *cq_outstanding_ : -1)
-                     << ", max_cqe=" << cq_limit << ")";
-        for (; it != slice_list.end(); ++it) {
-            failed_slice_list.push_back(*it);
+        // Remove consumed batch from slice_list, prepend any retry slices
+        slice_list.erase(slice_list.begin(), slice_list.begin() + batch_count);
+        if (!retry_slices.empty()) {
+            slice_list.insert(slice_list.begin(), retry_slices.begin(),
+                              retry_slices.end());
+            std::this_thread::yield();
         }
-        slice_list.clear();
-        return 0;
     }
 
     return 0;

mooncake-transfer-engine/src/transport/rdma_transport/endpoint_store.cpp

@@ -35,6 +35,17 @@ std::shared_ptr<RdmaEndPoint> FIFOEndpointStore::getEndpoint(
     return nullptr;
 }
 
+std::shared_ptr<RdmaEndPoint> FIFOEndpointStore::getEndpointByPtr(
+    const RdmaEndPoint *endpoint_ptr) {
+    RWSpinlock::ReadGuard guard(endpoint_map_lock_);
+    for (auto &kv : endpoint_map_) {
+        if (kv.second.get() == endpoint_ptr) return kv.second;
+    }
+    for (auto &endpoint : waiting_list_)
+        if (endpoint.get() == endpoint_ptr) return endpoint;
+    return nullptr;
+}
+
 std::shared_ptr<RdmaEndPoint> FIFOEndpointStore::insertEndpoint(
     const std::string &peer_nic_path, RdmaContext *context) {
     RWSpinlock::WriteGuard guard(endpoint_map_lock_);
@@ -139,6 +150,17 @@ std::shared_ptr<RdmaEndPoint> SIEVEEndpointStore::getEndpoint(
     return nullptr;
 }
 
+std::shared_ptr<RdmaEndPoint> SIEVEEndpointStore::getEndpointByPtr(
+    const RdmaEndPoint *endpoint_ptr) {
+    RWSpinlock::ReadGuard guard(endpoint_map_lock_);
+    for (auto &kv : endpoint_map_) {
+        if (kv.second.first.get() == endpoint_ptr) return kv.second.first;
+    }
+    for (auto &endpoint : waiting_list_)
+        if (endpoint.get() == endpoint_ptr) return endpoint;
+    return nullptr;
+}
+
 std::shared_ptr<RdmaEndPoint> SIEVEEndpointStore::insertEndpoint(
     const std::string &peer_nic_path, RdmaContext *context) {
     RWSpinlock::WriteGuard guard(endpoint_map_lock_);

mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp

@@ -357,6 +357,11 @@ std::shared_ptr<RdmaEndPoint> RdmaContext::endpoint(
     return endpoint;
 }
 
+std::shared_ptr<RdmaEndPoint> RdmaContext::getEndpointByPtr(
+    const RdmaEndPoint *endpoint_ptr) {
+    return endpoint_store_->getEndpointByPtr(endpoint_ptr);
+}
+
 int RdmaContext::disconnectAllEndpoints() {
     return endpoint_store_->disconnectQPs();
 }

mooncake-transfer-engine/src/transport/rdma_transport/worker_pool.cpp

@@ -161,9 +161,8 @@ int WorkerPool::submitPostSend(
         slice_queue_lock_[shard_id].unlock();
     }
 
-    submitted_slice_count_.fetch_add(submitted_slice_count,
-                                     std::memory_order_relaxed);
-    if (suspended_flag_.load(std::memory_order_relaxed)) {
+    submitted_slice_count_.fetch_add(submitted_slice_count);
+    if (suspended_flag_.load()) {
         std::lock_guard<std::mutex> lock(cond_mutex_);
         cond_var_.notify_all();
     }
@@ -398,7 +397,7 @@ void WorkerPool::transferWorker(int thread_id) {
                 // Double-check condition after acquiring lock to avoid lost
                 // wakeup
                 if (processed_slice_count_.load(std::memory_order_relaxed) ==
-                    submitted_slice_count_.load(std::memory_order_relaxed)) {
+                    submitted_slice_count_.load()) {
                     cond_var_.wait_for(lock, std::chrono::seconds(1));
                 }
                 suspended_flag_.fetch_sub(1);
@@ -415,19 +414,62 @@ void WorkerPool::transferWorker(int thread_id) {
 
 int WorkerPool::doProcessContextEvents() {
     ibv_async_event event;
+    bool event_acked = false;
     if (ibv_get_async_event(context_.context(), &event) < 0) return ERR_CONTEXT;
     LOG(WARNING) << "Worker: Received context async event "
                  << ibv_event_type_str(event.event_type) << " for context "
                  << context_.deviceName();
     if (event.event_type == IBV_EVENT_QP_FATAL) {
-        auto endpoint = (RdmaEndPoint *)event.element.qp->qp_context;
-        endpoint->set_active(false);
+        auto endpoint_ptr = (RdmaEndPoint *)event.element.qp->qp_context;
+
+        /**
+         * There might be a deadlock if we call endpoint->set_active(false)
+         * before ack the event:
+         *
+         * Thread A:
+         *     Holding endpoint->lock_ and calling ibv_destroy_qp (if using
+         * eRDMA), ibv_destroy_qp will block until the event is acked.
+         *
+         * Thread B (this thread):
+         *     Calling endpoint->set_active(false), which blocks as
+         * endpoint->lock_ is held by Thread A.
+         */
+        ibv_ack_async_event(&event);
+        event_acked = true;
+
+        /**
+         * After ack the event, the endpoint might be destroyed if it happened
+         * to be destroying event.element.qp. Therefore, we cannot just
+         * dereference endpoint_ptr. Instead, we need to get the shared_ptr of
+         * the endpoint from context_ and use that shared_ptr to access the
+         * endpoint.
+         */
+        auto endpoint = context_.getEndpointByPtr(endpoint_ptr);
+        if (endpoint) {
+            endpoint->set_active(false);
+        }
     } else if (event.event_type == IBV_EVENT_DEVICE_FATAL ||
                event.event_type == IBV_EVENT_CQ_ERR ||
                event.event_type == IBV_EVENT_WQ_FATAL ||
                event.event_type == IBV_EVENT_PORT_ERR ||
                event.event_type == IBV_EVENT_LID_CHANGE) {
         context_.set_active(false);
+
+        /**
+         * Similar deadlock might happen if we call
+         * context_.disconnectAllEndpoints() before ack the event:
+         *
+         * Thread A:
+         *     Holding endpoint->lock_ and calling ibv_destroy_qp (if using
+         * eRDMA), ibv_destroy_qp will block until the event is acked.
+         *
+         * Thread B (this thread):
+         *     Calling endpoint->disconnect(), which blocks as endpoint->lock_
+         * is held by Thread A.
+         */
+        ibv_ack_async_event(&event);
+        event_acked = true;
+
         context_.disconnectAllEndpoints();
         LOG(INFO) << "Worker: Context " << context_.deviceName()
                   << " is now inactive";
@@ -436,7 +478,11 @@ int WorkerPool::doProcessContextEvents() {
         LOG(INFO) << "Worker: Context " << context_.deviceName()
                   << " is now active";
     }
-    ibv_ack_async_event(&event);
+
+    if (!event_acked) {
+        ibv_ack_async_event(&event);
+    }
+
     return 0;
 }
 

mooncake-transfer-engine/tent/include/tent/platform/cuda.h

@@ -17,34 +17,117 @@
 
 #include "tent/runtime/platform.h"
 #include "tent/common/config.h"
+#include "tent/common/concurrent/rw_spinlock.h"
+
+#include <cuda_runtime.h>
 
 namespace mooncake {
 namespace tent {
 
+class CUDAStreamPool;
+
+// RAII Wrapper for cudaStream_t.
+// When this handle destructs, the stream is automatically returned to the pool.
+class CUDAStreamHandle {
+   public:
+    CUDAStreamHandle() = default;
+
+    CUDAStreamHandle(cudaStream_t stream, int deviceId, CUDAStreamPool* pool)
+        : stream_(stream), deviceId_(deviceId), pool_(pool) {}
+
+    // exclusive ownership
+    CUDAStreamHandle(const CUDAStreamHandle&) = delete;
+    CUDAStreamHandle& operator=(const CUDAStreamHandle&) = delete;
+
+    CUDAStreamHandle(CUDAStreamHandle&& other) noexcept
+        : stream_(other.stream_),
+          deviceId_(other.deviceId_),
+          pool_(other.pool_) {
+        other.stream_ = nullptr;
+        other.pool_ = nullptr;
+    }
+
+    CUDAStreamHandle& operator=(CUDAStreamHandle&& other) noexcept;
+    ~CUDAStreamHandle();
+
+    // Get the underlying CUDA stream
+    [[nodiscard]] cudaStream_t get() const { return stream_; }
+
+   private:
+    void releaseToPool();
+
+    cudaStream_t stream_ = nullptr;
+    int deviceId_ = -1;
+    CUDAStreamPool* pool_ = nullptr;
+};
+
+// CUDA Stream Pool managing all devices.
+class CUDAStreamPool {
+    friend class CUDAStreamHandle;
+
+   public:
+    CUDAStreamPool() = default;
+    ~CUDAStreamPool() = default;
+
+    // Non-copyable
+    CUDAStreamPool(const CUDAStreamPool&) = delete;
+    CUDAStreamPool& operator=(const CUDAStreamPool&) = delete;
+
+    // Acquires a stream for the specified device.
+    static constexpr int kCurrentDevice = -1;
+    Status acquire(CUDAStreamHandle& outHandle, int deviceId = kCurrentDevice);
+
+   private:
+    class DevicePool {
+       public:
+        explicit DevicePool(int deviceId);
+        ~DevicePool();
+
+        Status acquire(cudaStream_t& outStream);
+        void release(cudaStream_t stream);
+
+       private:
+        int deviceId_;
+        RWSpinlock dev_lock_;
+        std::vector<cudaStream_t> availableStreams_;
+    };
+
+    // called by CUDAStreamHandle::releaseToPool
+    void release(int deviceId, cudaStream_t stream);
+    DevicePool* getDevicePool(int deviceId);
+
+    RWSpinlock pools_lock_;
+    std::vector<std::unique_ptr<DevicePool>> devicePools_;
+};
+
 class CudaPlatform : public Platform {
    public:
     CudaPlatform(std::shared_ptr<Config> config) : conf(std::move(config)) {}
 
     virtual ~CudaPlatform() {}
 
-    virtual Status probe(std::vector<Topology::NicEntry> &nic_list,
-                         std::vector<Topology::MemEntry> &mem_list);
+    virtual Status probe(std::vector<Topology::NicEntry>& nic_list,
+                         std::vector<Topology::MemEntry>& mem_list);
 
-    virtual Status allocate(void **pptr, size_t size, MemoryOptions &options);
+    virtual Status allocate(void** pptr, size_t size, MemoryOptions& options);
 
-    virtual Status free(void *ptr, size_t size);
+    virtual Status free(void* ptr, size_t size);
 
-    virtual Status copy(void *dst, void *src, size_t length);
+    virtual Status copy(void* dst, void* src, size_t length);
 
-    virtual MemoryType getMemoryType(void *addr);
+    virtual MemoryType getMemoryType(void* addr);
 
     virtual const std::vector<RangeLocation> getLocation(
-        void *start, size_t len, bool skip_prefault = false);
+        void* start, size_t len, bool skip_prefault = false);
 
     virtual const std::string type() const { return "cuda"; }
 
+    Status getStreamFromPool(CUDAStreamHandle& outHandle,
+                             int deviceId = CUDAStreamPool::kCurrentDevice);
+
    private:
     std::shared_ptr<Config> conf;
+    CUDAStreamPool stream_pool;
 };
 
 }  // namespace tent

mooncake-transfer-engine/tent/include/tent/transport/fault_proxy/fault_proxy_transport.h

@@ -0,0 +1,213 @@
+// Copyright 2026 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef TENT_FAULT_PROXY_TRANSPORT_H
+#define TENT_FAULT_PROXY_TRANSPORT_H
+
+#include <atomic>
+#include <cassert>
+#include <chrono>
+#include <memory>
+#include <random>
+#include <thread>
+
+#include "tent/common/status.h"
+#include "tent/runtime/transport.h"
+
+namespace mooncake {
+namespace tent {
+
+// Configurable fault injection policy.  All fields default to "no fault".
+struct FaultPolicy {
+    // Probability [0.0, 1.0] that submitTransferTasks() returns an error.
+    double submit_fail_rate = 0.0;
+
+    // Probability [0.0, 1.0] that getTransferStatus() flips COMPLETED→FAILED.
+    double status_corrupt_rate = 0.0;
+
+    // Artificial latency (microseconds) added before each submit.
+    uint64_t submit_delay_us = 0;
+
+    // Deterministic mode: succeed for the first N submits, then always fail.
+    // -1 disables this mode (rate-based mode used instead).
+    int fail_after_n_submits = -1;
+
+    // If true, install() returns an error immediately.
+    bool fail_install = false;
+};
+
+// Decorator that wraps any Transport and injects faults according to a
+// FaultPolicy.  The engine sees a normal Transport; failures trigger the
+// real failover / retry machinery without requiring hardware.
+class FaultProxyTransport : public Transport {
+   public:
+    FaultProxyTransport(std::shared_ptr<Transport> real, FaultPolicy policy)
+        : real_(std::move(real)), policy_(policy), submit_count_(0) {
+        assert(real_ && "FaultProxyTransport: real transport must not be null");
+    }
+
+    // -- Lifecycle -----------------------------------------------------------
+
+    Status install(std::string& local_segment_name,
+                   std::shared_ptr<ControlService> metadata,
+                   std::shared_ptr<Topology> local_topology,
+                   std::shared_ptr<Config> conf = nullptr) override {
+        if (policy_.fail_install) {
+            return Status::InternalError(
+                "fault injected: install failure" LOC_MARK);
+        }
+        return real_->install(local_segment_name, metadata, local_topology,
+                              conf);
+    }
+
+    Status uninstall() override { return real_->uninstall(); }
+
+    // -- Capabilities --------------------------------------------------------
+
+    const Capabilities capabilities() const override {
+        return real_->capabilities();
+    }
+
+    // -- Batch management ----------------------------------------------------
+
+    Status allocateSubBatch(SubBatchRef& batch, size_t max_size) override {
+        return real_->allocateSubBatch(batch, max_size);
+    }
+
+    Status freeSubBatch(SubBatchRef& batch) override {
+        return real_->freeSubBatch(batch);
+    }
+
+    // -- Transfer (primary injection points) ---------------------------------
+
+    Status submitTransferTasks(
+        SubBatchRef batch, const std::vector<Request>& request_list) override {
+        // Artificial delay
+        if (policy_.submit_delay_us > 0) {
+            std::this_thread::sleep_for(
+                std::chrono::microseconds(policy_.submit_delay_us));
+        }
+
+        int count = submit_count_.fetch_add(1, std::memory_order_relaxed);
+
+        // Deterministic count-based failure (takes precedence over rate-based)
+        if (policy_.fail_after_n_submits >= 0 &&
+            count >= policy_.fail_after_n_submits) {
+            return Status::InternalError(
+                "fault injected: submit failure (count)" LOC_MARK);
+        }
+
+        // Rate-based probabilistic failure
+        if (policy_.submit_fail_rate > 0.0 &&
+            randomDouble() < policy_.submit_fail_rate) {
+            return Status::InternalError(
+                "fault injected: submit failure (rate)" LOC_MARK);
+        }
+
+        return real_->submitTransferTasks(batch, request_list);
+    }
+
+    Status getTransferStatus(SubBatchRef batch, int task_id,
+                             TransferStatus& status) override {
+        auto s = real_->getTransferStatus(batch, task_id, status);
+        if (!s.ok()) return s;
+
+        // Status corruption: flip COMPLETED → FAILED
+        if (status.s == TransferStatusEnum::COMPLETED &&
+            policy_.status_corrupt_rate > 0.0 &&
+            randomDouble() < policy_.status_corrupt_rate) {
+            status.s = TransferStatusEnum::FAILED;
+        }
+        return s;
+    }
+
+    // -- Memory management (pass-through) ------------------------------------
+
+    Status addMemoryBuffer(BufferDesc& desc,
+                           const MemoryOptions& options) override {
+        return real_->addMemoryBuffer(desc, options);
+    }
+
+    Status addMemoryBuffer(std::vector<BufferDesc>& desc_list,
+                           const MemoryOptions& options) override {
+        return real_->addMemoryBuffer(desc_list, options);
+    }
+
+    Status removeMemoryBuffer(BufferDesc& desc) override {
+        return real_->removeMemoryBuffer(desc);
+    }
+
+    Status allocateLocalMemory(void** addr, size_t size,
+                               MemoryOptions& options) override {
+        return real_->allocateLocalMemory(addr, size, options);
+    }
+
+    Status freeLocalMemory(void* addr, size_t size) override {
+        return real_->freeLocalMemory(addr, size);
+    }
+
+    bool warmupMemory(void* addr, size_t length) override {
+        return real_->warmupMemory(addr, length);
+    }
+
+    // -- Notifications (pass-through) ----------------------------------------
+
+    bool supportNotification() const override {
+        return real_->supportNotification();
+    }
+
+    Status sendNotification(SegmentID target_id,
+                            const Notification& notify) override {
+        return real_->sendNotification(target_id, notify);
+    }
+
+    Status receiveNotification(
+        std::vector<Notification>& notify_list) override {
+        return real_->receiveNotification(notify_list);
+    }
+
+    // -- Identity ------------------------------------------------------------
+
+    const char* getName() const override { return "<fault-proxy>"; }
+
+    // -- Test helpers --------------------------------------------------------
+    // Note: resetPolicy() is intended for single-threaded test scenarios.
+    // It is NOT safe to call concurrently with submitTransferTasks().
+
+    void resetPolicy(FaultPolicy new_policy) {
+        policy_ = new_policy;
+        submit_count_.store(0, std::memory_order_relaxed);
+    }
+
+    int submitCount() const {
+        return submit_count_.load(std::memory_order_relaxed);
+    }
+
+   private:
+    // Thread-safe random double in [0.0, 1.0).
+    static double randomDouble() {
+        thread_local std::mt19937 rng(std::random_device{}());
+        thread_local std::uniform_real_distribution<double> dist(0.0, 1.0);
+        return dist(rng);
+    }
+
+    std::shared_ptr<Transport> real_;
+    FaultPolicy policy_;
+    std::atomic<int> submit_count_;
+};
+
+}  // namespace tent
+}  // namespace mooncake
+
+#endif  // TENT_FAULT_PROXY_TRANSPORT_H

mooncake-transfer-engine/tent/include/tent/transport/mnnvl/mnnvl_transport.h

@@ -26,6 +26,7 @@
 
 #include "tent/runtime/control_plane.h"
 #include "tent/runtime/transport.h"
+#include "tent/platform/cuda.h"
 
 namespace mooncake {
 namespace tent {
@@ -36,12 +37,14 @@ struct MnnvlTask {
     volatile size_t transferred_bytes;
     uint64_t target_addr = 0;
     int cuda_id = 0;
+    cudaEvent_t completion_event = nullptr;
 };
 
 struct MnnvlSubBatch : public Transport::SubBatch {
     std::vector<MnnvlTask> task_list;
     size_t max_size;
-    cudaStream_t stream;
+    CUDAStreamHandle sync_stream;
+    CUDAStreamHandle async_stream;
     virtual size_t size() const { return task_list.size(); }
 };
 
@@ -81,7 +84,7 @@ class MnnvlTransport : public Transport {
     virtual Status freeLocalMemory(void *addr, size_t size);
 
    private:
-    void startTransfer(MnnvlTask *task, MnnvlSubBatch *batch);
+    void startTransfer(std::vector<MnnvlTask *> &tasks, MnnvlSubBatch *batch);
 
     void *createSharedMemory(const std::string &path, size_t size);
 
@@ -95,6 +98,7 @@ class MnnvlTransport : public Transport {
     std::string local_segment_name_;
     std::shared_ptr<Topology> local_topology_;
     std::shared_ptr<ControlService> metadata_;
+    CudaPlatform *platform_;
 
     struct OpenedMnnvlEntry {
         void *mnnvl_addr;

mooncake-transfer-engine/tent/include/tent/transport/nvlink/nvlink_transport.h

@@ -28,6 +28,7 @@
 #include "tent/common/concurrent/ticket_lock.h"
 #include "tent/runtime/control_plane.h"
 #include "tent/runtime/transport.h"
+#include "tent/platform/cuda.h"
 
 namespace mooncake {
 namespace tent {
@@ -39,12 +40,14 @@ struct NVLinkTask {
     uint64_t target_addr = 0;
     bool is_cuda_ipc;
     int cuda_id = 0;
+    cudaEvent_t completion_event = nullptr;
 };
 
 struct NVLinkSubBatch : public Transport::SubBatch {
     std::vector<NVLinkTask> task_list;
     size_t max_size;
-    cudaStream_t stream;
+    CUDAStreamHandle sync_stream;
+    CUDAStreamHandle async_stream;
     virtual size_t size() const { return task_list.size(); }
 };
 
@@ -79,7 +82,7 @@ class NVLinkTransport : public Transport {
     virtual const char *getName() const { return "nvlink"; }
 
    private:
-    void startTransfer(NVLinkTask *task, NVLinkSubBatch *batch);
+    void startTransfer(std::vector<NVLinkTask *> &tasks, NVLinkSubBatch *batch);
 
     void *createSharedMemory(const std::string &path, size_t size);
 
@@ -93,6 +96,7 @@ class NVLinkTransport : public Transport {
     std::string local_segment_name_;
     std::shared_ptr<Topology> local_topology_;
     std::shared_ptr<ControlService> metadata_;
+    CudaPlatform *platform_;
 
     struct OpenedShmEntry {
         void *shm_addr;