[Doc] Update CODEOWNERS with new document owners

.github/CODEOWNERS

@@ -8,6 +8,7 @@
 
 .github @stmatengss @ykwd @Ann-1024 @luketong777
 /docs @ShangmingCai @stmatengss @ykwd
+/docs/source/design/transfer-engine @alogfans @doujiang24 @chestnut-Q @amd-arozanov @ascend-direct-dev @whn09
 /mooncake-ep @UNIDY2002 @ympcMark @yuechen-sys
 /mooncake-integration/transfer_engine @ShangmingCai @alogfans
 /mooncake-integration/store @ykwd @stmatengss

.github/workflows/ci.yml

@@ -134,12 +134,10 @@ jobs:
         MASTER_PID=$!
         sleep 3
         cd mooncake-store/go
-        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 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 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-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
+        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"
         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
@@ -746,18 +744,6 @@ 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()
@@ -770,8 +756,6 @@ 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,6 +64,7 @@ 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"
         }
@@ -107,53 +108,22 @@ 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
-
-        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/
+        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 "Error: ../pybind11 does not exist. Cannot copy pybind11."
-              exit 1
+                echo "Detected that extern/pybind11 already exists, continuing execution...."
             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,7 +5,6 @@ build_ofed4
 old
 local_test
 go.sum
-!mooncake-common/etcd/go.sum
 *.so
 bin
 mod
@@ -205,4 +204,4 @@ CLAUDE.md
 _codeql_detected_source_root
 
 # CodeBuddy Memory
-.codebuddy/
+.codebuddy/

.pre-commit-config.yaml

@@ -23,16 +23,6 @@ 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:

CMakeLists.txt

@@ -45,21 +45,12 @@ 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 building targets that include ASIO headers.
+# Define ASIO macros before adding mooncake-asio subdirectory
 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,7 +41,6 @@ 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` |
@@ -54,8 +53,6 @@ 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.25.9
+GOVER=1.23.8
 
 # Function to print section headers
 print_section() {

docs/source/deployment/ssd-offload.md

@@ -88,11 +88,10 @@ 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_OFFLOAD_USE_URING` | `false` | Enable io_uring for async file I/O |
+| `MOONCAKE_USE_URING` | `false` | Enable io_uring for async file I/O |
 
 ### Bucket backend settings
 
@@ -102,8 +101,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_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` |
+| `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` |
 
 ---
 
@@ -128,11 +127,6 @@ 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`
@@ -149,7 +143,7 @@ Best for: high-concurrency scenarios with many small objects where restart durab
 
 ## Eviction (Bucket Backend Only)
 
-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.
+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.
 
 | Policy | Behavior |
 |--------|----------|
@@ -184,8 +178,8 @@ mooncake_master \
 ```bash
 export MOONCAKE_OFFLOAD_FILE_STORAGE_PATH=/nvme/mooncake_offload
 export MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR=bucket_storage_backend
-export MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE=$((200 * 1024 * 1024 * 1024))  # 200 GB
-export MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY=lru
+export MOONCAKE_BUCKET_MAX_TOTAL_SIZE=$((200 * 1024 * 1024 * 1024))  # 200 GB
+export MOONCAKE_BUCKET_EVICTION_POLICY=lru
 
 mooncake_client \
     --master_server_address="192.168.1.10:50051" \
@@ -256,7 +250,7 @@ mooncake_client \
 
 ### "Failed to register buffer with UringFile" warning in logs
 
-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.
+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.
 
 Check the current limit:
 

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_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`).
+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`).
 
 ### Policies
 
@@ -195,7 +195,7 @@ This ordering guarantees:
 
 ## io_uring File I/O
 
-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.
+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.
 
 ### Thread-local rings (`SharedUringRing`)
 

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

@@ -81,7 +81,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-common/libasio.so ../mooncake-wheel/mooncake/
+cp mooncake-asio/libasio.so ../mooncake-wheel/mooncake/
 
 # Install with pip
 pip install -e ../mooncake-wheel --no-build-isolation
@@ -224,30 +224,6 @@ Earlier CPU-to-CPU tuning results (before EFA striping optimization, when `MC_SL
 
 </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.
@@ -277,10 +253,8 @@ Tested on two p5en.48xlarge instances (Intel Xeon 8488C, 8× H200 141GB, 16 EFA
 
 | 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 |
@@ -299,27 +273,6 @@ Tested on two p5en.48xlarge instances (Intel Xeon 8488C, 8× H200 141GB, 16 EFA
 - 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

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-common/libasio.so ../mooncake-wheel/mooncake/
+cp mooncake-asio/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_OFFLOAD_USE_URING=1 \
+MOONCAKE_USE_URING=1 \
 mooncake_client \
     --host=127.0.0.1 \
     --global_segment_size=80GB \

mooncake-asio/CMakeLists.txt

@@ -0,0 +1,50 @@
+# 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,10 +2,6 @@ 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 go.mod go.sum build.sh
+    DEPENDS etcd_wrapper.go
 )
 
 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,8 +1,6 @@
 module github.com/kvcache-ai/Mooncake/mooncake-common/etcd
 
-go 1.25.0
-
-toolchain go1.25.9
+go 1.24.0
 
 require (
 	go.etcd.io/etcd/api/v3 v3.5.21

mooncake-common/etcd/go.sum

@@ -1,108 +0,0 @@
-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

@@ -1,20 +0,0 @@
-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

@@ -1,61 +0,0 @@
-// 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

@@ -1,60 +0,0 @@
-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

@@ -1,571 +0,0 @@
-//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

@@ -1,489 +0,0 @@
-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

@@ -1,293 +0,0 @@
-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,64 +1,13 @@
 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
@@ -67,5 +16,3 @@ 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,8 +1563,7 @@ 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,
-               const std::string &ssd_offload_path = "") {
+               bool enable_ssd_offload = false) {
                 auto real_client = self.init_real_client();
                 std::shared_ptr<mooncake::TransferEngine> transfer_engine =
                     nullptr;
@@ -1575,15 +1574,14 @@ 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,
-                    ssd_offload_path);
+                    master_server_addr, transfer_engine, "",
+                    enable_ssd_offload);
             },
             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("ssd_offload_path") = "")
+            py::arg("enable_ssd_offload") = false)
         .def(
             "setup",
             [](MooncakeStorePyWrapper &self, const py::dict &config_dict) {
@@ -1612,9 +1610,7 @@ 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).\n"
-            "  ssd_offload_path: SSD storage directory path (overrides env "
-            "var).")
+            "  enable_ssd_offload: Enable SSD offload (default false).")
         .def(
             "setup_dummy",
             [](MooncakeStorePyWrapper &self, size_t mem_pool_size,
@@ -2271,20 +2267,6 @@ 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-common"
+EXT_LDFLAGS+=" -L$BUILD_DIR/mooncake-asio"
 EXT_LDFLAGS+=" -ltransfer_engine -lbase -lasio -lstdc++ -lnuma -lglog -libverbs -ljsoncpp"
 
 if [ -d "/usr/local/cuda/lib64/stubs" ]; then

mooncake-pg/include/mooncake_worker.cuh

@@ -53,7 +53,7 @@ __global__ struct Task {
     size_t tensorSize;  // In bytes
     int64_t broadcastRoot;
     int bufferOffset;
-    uint64_t submitSequence = 0;
+    uint64_t submit_sequence = 0;
     BatchID batchID;
     void* transferGroupMeta;
 };

mooncake-pg/src/mooncake_worker.cu

@@ -3,7 +3,6 @@
 #include <memory>
 #include <thread>
 #include <mooncake_worker.cuh>
-#include <ATen/cuda/CUDAGraphsUtils.cuh>
 
 #include "pg_utils.h"
 
@@ -46,68 +45,36 @@ class MooncakeWorkCuda : public ::c10d::Work {
 
     bool wait(std::chrono::milliseconds timeout) override {
         // Wait until the task has been submitted to TransferEngine:
-        // This tries to ensure that the CUDA kernels required for the transfer
+        // This ensures 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).
+        // Although this is not required for CPU-only transports such as
+        // RdmaTransport, we keep this behavior to avoid invasive changes
+        // to TE/TENT.
         //
         // 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.
+        // assumption (cudaMemcpy(Async) may be called 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.
-        }
+        // This wait is primarily needed for two reasons:
+        //   1. PyTorch documents that `wait` should ensure the operation is
+        //      issued, though not necessarily completed, for CUDA work.
+        //      In practice, this means the transfer kernel must at least be
+        //      launched; otherwise, a hang may occur.
+        //   2. The current stream is 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
+        //      not block until all kernels needed for the transfer task have
+        //      been launched, in case TE/TENT use `current_stream` to
+        //      launch those kernels (though rare).
+        auto submitted =
+            worker_->waitUntilTasksSubmitted(submitted_tasks_, timeout);
         if (!submitted) return false;
 
-        // Once all tasks have been submitted, use the event to synchronize
+        // Once all tasks have been submitted, create an event to synchronize
         // the current stream and the enqueue stream, but do not wait on this
         // event.
         //
@@ -134,18 +101,6 @@ 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;
@@ -168,7 +123,7 @@ __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].submit_sequence = submitSequence;
     tasks[taskId].transferGroupMeta = meta;
 
     // Publish task metadata before notifying the host worker thread.
@@ -401,7 +356,7 @@ MooncakeWorker::MooncakeWorker(int cuda_device_index)
     }
     for (size_t i = 0; i < kNumTasks_; ++i) {
         tasks_[i].active = false;
-        tasks_[i].submitSequence = 0;
+        tasks_[i].submit_sequence = 0;
         submitted_task_sequence_[i].store(0, std::memory_order_relaxed);
     }
 }

mooncake-pg/src/mooncake_worker_thread.cpp

@@ -91,7 +91,7 @@ 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;
+                    const auto submit_sequence = task.submit_sequence;
                     if (skipTransfer) {
                         submitted_task_sequence_[i].store(
                             submit_sequence, 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-common"
+CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-asio"
 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,15 +1,10 @@
 #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>
@@ -51,56 +46,6 @@ 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",
@@ -116,8 +61,7 @@ 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),
-          start_time_(std::chrono::steady_clock::now()) {}
+                         "Put transfer latency (us)", kLatencyBucket, labels) {}
 
     ylt::metric::counter_t total_read_bytes;
     ylt::metric::counter_t total_write_bytes;
@@ -135,7 +79,7 @@ struct TransferMetric {
         put_latency_us.serialize(str);
     }
 
-    std::string summary_metrics(bool include_bandwidth = true) {
+    std::string summary_metrics() {
         std::stringstream ss;
         ss << "=== Transfer Metrics Summary ===\n";
 
@@ -144,14 +88,6 @@ 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";
@@ -166,14 +102,6 @@ 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 =
@@ -346,157 +274,6 @@ 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 = {
@@ -647,6 +424,7 @@ 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";
@@ -660,7 +438,6 @@ struct SsdMetric {
 struct ClientMetric {
     TransferMetric transfer_metric;
     MasterClientMetric master_client_metric;
-    TransferOperationMetric transfer_operation_metric;
     SsdMetric ssd_metric;
 
     /**
@@ -675,44 +452,25 @@ struct ClientMetric {
      *   (default: 0, 0 = collect but don't report)
      */
     static std::unique_ptr<ClientMetric> Create(
-        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);
-    }
+        const std::map<std::string, std::string>& labels = {});
 
     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 = {},
-                          bool bandwidth_reporting_enabled = true,
-                          bool master_rpc_metrics_enabled = true);
+    explicit ClientMetric(
+        uint64_t interval_seconds = 0,
+        const std::map<std::string, std::string>& labels = {});
     ~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,7 +24,6 @@
 #include "master_metric_manager.h"
 #include "count_min_sketch.h"
 #include "local_hot_cache.h"
-#include "pinned_buffer_pool.h"
 
 namespace mooncake {
 
@@ -62,8 +61,6 @@ 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)
@@ -447,15 +444,6 @@ 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) {
@@ -672,9 +660,6 @@ 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,7 +7,6 @@
 #include "pyclient.h"
 #include "real_client.h"
 #include "shm_helper.h"
-#include "client_metric.h"
 #include <memory>
 
 namespace mooncake {
@@ -26,8 +25,7 @@ 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,
-                   const std::string &ssd_offload_path = "") {
+                   bool enable_ssd_offload = false) {
         // Dummy client does not support real setup
         return -1;
     };
@@ -151,12 +149,6 @@ 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;
@@ -204,22 +196,6 @@ 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
@@ -278,11 +254,6 @@ 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,7 +3,6 @@
 #include "client_service.h"
 #include "client_buffer.hpp"
 #include "storage_backend.h"
-#include "pinned_buffer_pool.h"
 
 namespace mooncake {
 
@@ -106,8 +105,6 @@ 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

@@ -1,76 +0,0 @@
-#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,10 +81,6 @@ 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 {
@@ -144,8 +140,6 @@ 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
@@ -161,8 +155,6 @@ 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);
 
@@ -275,8 +267,6 @@ 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;
@@ -332,8 +322,6 @@ 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()) {
@@ -412,8 +400,6 @@ 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()) {
@@ -750,8 +736,6 @@ 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;
@@ -803,8 +787,6 @@ 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,14 +562,8 @@ 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);
@@ -897,10 +891,7 @@ class MasterService {
     }
 
     // Helper to clean up stale handles pointing to unmounted segments
-    // 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);
+    bool CleanupStaleHandles(ObjectMetadata& metadata);
 
     // Helper: allocate replicas, create ObjectMetadata, insert into shard,
     // and return descriptor list.  Shared by PutStart and UpsertStart.
@@ -971,21 +962,11 @@ 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 (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.
+            // Automatically clean up invalid handles
             if (it_ != shard_guard_->metadata.end()) {
-                // 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()) {
+                if (service_->CleanupStaleHandles(it_->second)) {
                     this->Erase();
+
                     if (processing_it_ != shard_guard_->processing_keys.end()) {
                         this->EraseFromProcessing();
                     }
@@ -1168,14 +1149,6 @@ 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

@@ -1,138 +0,0 @@
-#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,
-        const std::string &ssd_offload_path = "") = 0;
+        const std::string &ipc_socket_path,
+        bool enable_ssd_offload = false) = 0;
 
     virtual int setup_dummy(size_t mem_pool_size, size_t local_buffer_size,
                             const std::string &server_address,
@@ -330,10 +330,6 @@ 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,8 +79,7 @@ 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,
-        const std::string &ssd_offload_path = "");
+        bool enable_ssd_offload = false);
 
     int setup_dummy(size_t mem_pool_size, size_t local_buffer_size,
                     const std::string &server_address,
@@ -484,8 +483,7 @@ 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,
-        const std::string &ssd_offload_path = "");
+        bool enable_ssd_offload = false, bool start_offload_rpc_server = false);
 
     // Overload that accepts a configuration dictionary
     tl::expected<void, ErrorCode> setup_internal(const ConfigDict &config);
@@ -639,10 +637,6 @@ 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,12 +2,9 @@
 
 #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>
@@ -285,36 +282,6 @@ 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,12 +170,6 @@ 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,10 +174,7 @@ 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),
-          start_ts_(getCurrentTimeInMilli()) {}
+        : engine_(engine), batch_id_(batch_id), batch_size_(batch_size) {}
 
     ~TransferEngineOperationState() { engine_.freeBatchID(batch_id_); }
 
@@ -202,7 +199,6 @@ class TransferEngineOperationState : public OperationState {
     TransferEngine& engine_;
     BatchID batch_id_;
     size_t batch_size_;
-    const int64_t start_ts_;
 };
 
 /**

mooncake-store/src/CMakeLists.txt

@@ -202,62 +202,15 @@ 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,27 +28,6 @@ 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) {
@@ -76,21 +55,12 @@ uint64_t parseMetricsInterval() {
 }  // anonymous namespace
 
 ClientMetric::ClientMetric(uint64_t interval_seconds,
-                           const std::map<std::string, std::string>& labels,
-                           bool bandwidth_reporting_enabled,
-                           bool master_rpc_metrics_enabled)
+                           const std::map<std::string, std::string>& labels)
     : transfer_metric(labels),
       master_client_metric(labels),
-      transfer_operation_metric(labels),
       ssd_metric(labels),
       should_stop_metrics_thread_(false),
-      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()};
+      metrics_interval_seconds_(interval_seconds) {
     if (metrics_interval_seconds_ > 0) {
         StartMetricsReportingThread();
     }
@@ -99,8 +69,7 @@ ClientMetric::ClientMetric(uint64_t interval_seconds,
 ClientMetric::~ClientMetric() { StopMetricsReportingThread(); }
 
 std::unique_ptr<ClientMetric> ClientMetric::Create(
-    const std::map<std::string, std::string>& labels,
-    bool master_rpc_metrics_enabled) {
+    const std::map<std::string, std::string>& labels) {
     if (!parseMetricsEnabled()) {
         LOG(INFO) << "Client metrics disabled (set MC_STORE_CLIENT_METRIC=0 to "
                      "disable)";
@@ -108,83 +77,29 @@ 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,
-                                          bandwidth_reporting_enabled,
-                                          master_rpc_metrics_enabled);
+    return std::make_unique<ClientMetric>(interval, labels);
 }
 
 void ClientMetric::serialize(std::string& str) {
     transfer_metric.serialize(str);
-    if (master_rpc_metrics_enabled_) {
-        master_client_metric.serialize(str);
-    }
-    transfer_operation_metric.serialize(str);
+    master_client_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(bandwidth_reporting_enabled_);
+    ss << transfer_metric.summary_metrics();
     ss << "\n";
-    if (master_rpc_metrics_enabled_) {
-        ss << master_client_metric.summary_metrics();
-        ss << "\n";
-    }
-    ss << transfer_operation_metric.summary_metrics();
+    ss << master_client_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_ =
@@ -209,12 +124,7 @@ void ClientMetric::StartMetricsReportingThread() {
 
                 // Print metrics summary
                 std::string summary = summary_metrics();
-                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 Report:\n" << summary;
             }
             LOG(INFO) << "Client metrics reporting thread stopped";
         });
@@ -230,4 +140,4 @@ void ClientMetric::StopMetricsReportingThread() {
     }
 }
 
-}  // namespace mooncake
+}  // namespace mooncake

mooncake-store/src/client_service.cpp

@@ -27,14 +27,9 @@
 #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) {
@@ -62,7 +57,6 @@ 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_;
@@ -2501,34 +2495,19 @@ 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) {
-        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);
-        }
+        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,7 +7,6 @@
 #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"
@@ -24,75 +23,6 @@
 #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;
@@ -199,10 +129,7 @@ std::vector<tl::expected<ResultType, ErrorCode>> DummyClient::invoke_batch_rpc(
         }());
 }
 
-DummyClient::DummyClient()
-    : client_id_(generate_uuid()),
-      metrics_(ClientMetric::Create(merge_labels({{"client_mode", "dummy"}}),
-                                    false)) {
+DummyClient::DummyClient() : client_id_(generate_uuid()) {
     // Initialize logging severity (leave as before)
     mooncake::init_ylt_log_level();
     // Initialize client pools
@@ -214,30 +141,6 @@ DummyClient::DummyClient()
 
 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);
@@ -621,40 +524,35 @@ 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 invoke_observed_void_rpc<&RealClient::put_dummy_helper>(
-        TransferOperationKind::kWrite, "put", value.size_bytes(), false, key,
-        value, config, client_id_);
+    return to_py_ret(invoke_rpc<&RealClient::put_dummy_helper, void>(
+        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 invoke_observed_void_rpc<&RealClient::put_batch_dummy_helper>(
-        TransferOperationKind::kWrite, "put_batch", sum_value_sizes(values),
-        true, keys, values, config, client_id_);
+    return to_py_ret(invoke_rpc<&RealClient::put_batch_dummy_helper, void>(
+        keys, values, config, client_id_));
 }
 
 int DummyClient::put_parts(const std::string& key,
                            std::vector<std::span<const char>> values,
                            const ReplicateConfig& config) {
-    return invoke_observed_void_rpc<&RealClient::put_parts_dummy_helper>(
-        TransferOperationKind::kWrite, "put_parts", sum_value_sizes(values),
-        false, key, values, config, client_id_);
+    return to_py_ret(invoke_rpc<&RealClient::put_parts_dummy_helper, void>(
+        key, values, config, client_id_));
 }
 
 int DummyClient::upsert(const std::string& key, std::span<const char> value,
                         const ReplicateConfig& config) {
-    return invoke_observed_void_rpc<&RealClient::upsert_dummy_helper>(
-        TransferOperationKind::kWrite, "upsert", value.size_bytes(), false, key,
-        value, config, client_id_);
+    return to_py_ret(invoke_rpc<&RealClient::upsert_dummy_helper, void>(
+        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 invoke_observed_void_rpc<&RealClient::upsert_from_dummy_helper>(
-        TransferOperationKind::kWrite, "upsert_from", size, false, key,
-        dummy_addr, size, config, client_id_);
+    return to_py_ret(invoke_rpc<&RealClient::upsert_from_dummy_helper, void>(
+        key, dummy_addr, size, config, client_id_));
 }
 
 std::vector<int> DummyClient::batch_upsert_from(
@@ -664,7 +562,6 @@ 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_);
@@ -673,29 +570,21 @@ 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 invoke_observed_void_rpc<&RealClient::upsert_parts_dummy_helper>(
-        TransferOperationKind::kWrite, "upsert_parts", sum_value_sizes(values),
-        false, key, values, config, client_id_);
+    return to_py_ret(invoke_rpc<&RealClient::upsert_parts_dummy_helper, void>(
+        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 invoke_observed_void_rpc<&RealClient::upsert_batch_dummy_helper>(
-        TransferOperationKind::kWrite, "upsert_batch", sum_value_sizes(values),
-        true, keys, values, config, client_id_);
+    return to_py_ret(invoke_rpc<&RealClient::upsert_batch_dummy_helper, void>(
+        keys, values, config, client_id_));
 }
 
 int DummyClient::remove(const std::string& key, bool force) {
@@ -761,7 +650,6 @@ 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,
@@ -780,8 +668,6 @@ 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));
         }
@@ -800,14 +686,11 @@ 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;
 
@@ -868,37 +751,19 @@ 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()));
     }
-    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;
+    return to_py_ret(*result);
 }
 
 std::vector<std::vector<std::vector<int64_t>>> DummyClient::get_into_ranges(
@@ -908,7 +773,6 @@ 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<
@@ -922,13 +786,8 @@ std::vector<std::vector<std::vector<int64_t>>> DummyClient::get_into_ranges(
                                                all_dst_offsets,
                                                internal_results.error());
     }
-    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;
+
+    return convert_ranged_read_results(internal_results.value());
 }
 
 std::string DummyClient::get_hostname() const {
@@ -940,7 +799,6 @@ 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_);
@@ -951,13 +809,6 @@ 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;
 }
 
@@ -971,7 +822,6 @@ 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_);
@@ -982,12 +832,6 @@ 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;
 }
 
@@ -1006,7 +850,6 @@ 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,
@@ -1016,13 +859,6 @@ 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;
 }
 
@@ -1033,7 +869,6 @@ 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,
@@ -1044,12 +879,6 @@ 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,17 +7,12 @@
 #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;
 
@@ -42,10 +37,9 @@ 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 = GetEnvOr<int64_t>(
-        "MOONCAKE_OFFLOAD_SCANMETA_ITERATOR_KEYS_LIMIT",
+    config.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);
@@ -64,9 +58,7 @@ 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_OFFLOAD_USE_URING",
-                       GetEnvStringOr("MOONCAKE_USE_URING", "false"));
+    auto use_uring_str = GetEnvStringOr("MOONCAKE_USE_URING", "false");
     config.use_uring = (use_uring_str == "true" || use_uring_str == "1");
 
     return config;
@@ -160,7 +152,6 @@ 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()) {
@@ -381,37 +372,6 @@ 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](
@@ -439,12 +399,7 @@ tl::expected<void, ErrorCode> FileStorage::OffloadObjects(
             return res;
         };
         auto offload_res = storage_backend_->BatchOffload(
-            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);
-        }
+            batch_object, bucket_complete_handler, eviction_handler);
         if (!offload_res) {
             LOG(ERROR) << "Failed to store objects with error: "
                        << offload_res.error();

mooncake-store/src/master.cpp

@@ -109,10 +109,6 @@ 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, "",
@@ -297,11 +293,6 @@ 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",
@@ -876,8 +867,6 @@ 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,13 +48,6 @@ 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,
@@ -166,19 +159,6 @@ 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";
@@ -371,23 +351,12 @@ 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, alive_clients)) {
+            if (CleanupStaleHandles(it->second)) {
                 // If the object is empty, we need to erase the iterator and
                 // also erase the key from processing_keys,
                 // replication_tasks, and offloading_tasks.
@@ -873,15 +842,13 @@ 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, alive_clients)) {
+    if (it != shard->metadata.end() && !CleanupStaleHandles(it->second)) {
         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
@@ -930,7 +897,7 @@ auto MasterService::PutEnd(const UUID& client_id, const std::string& key,
         },
         [](Replica& replica) { replica.mark_complete(); });
 
-    if (enable_offload_ && !offload_on_evict_) {
+    if (enable_offload_) {
         auto& shard = accessor.GetShard();
         metadata.VisitReplicas(
             &Replica::fn_is_completed, [this, &key, &shard](Replica& replica) {
@@ -1118,7 +1085,6 @@ 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));
 
@@ -1128,9 +1094,7 @@ 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.
-    // Also clean up local_disk replicas whose owner client has expired.
-    if (it != shard->metadata.end() &&
-        CleanupStaleHandles(it->second, alive_clients)) {
+    if (it != shard->metadata.end() && CleanupStaleHandles(it->second)) {
         shard->processing_keys.erase(key);
         shard->metadata.erase(it);
         it = shard->metadata.end();
@@ -1967,8 +1931,6 @@ 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);
@@ -1986,7 +1948,7 @@ auto MasterService::BatchRemove(const std::vector<std::string>& keys,
             }
 
             // Clean up stale replica handles (consistent with single Remove)
-            if (CleanupStaleHandles(it->second, alive_clients)) {
+            if (CleanupStaleHandles(it->second)) {
                 shard->processing_keys.erase(key);
                 shard->replication_tasks.erase(key);
                 shard->offloading_tasks.erase(key);
@@ -2034,14 +1996,10 @@ auto MasterService::BatchRemove(const std::vector<std::string>& keys,
     return results;
 }
 
-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);
+bool MasterService::CleanupStaleHandles(ObjectMetadata& metadata) {
+    // Remove those with invalid allocators
+    metadata.EraseReplicas([](const Replica& replica) {
+        return replica.has_invalid_mem_handle();
     });
 
     // Return true if no valid replicas remain after cleanup
@@ -3511,81 +3469,6 @@ 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;
@@ -3658,18 +3541,16 @@ void MasterService::BatchEvict(double evict_ratio_target,
                     continue;
                 }
                 if (it->second.lease_timeout <= target_timeout) {
-                    // Evict this object (or defer for offload)
-                    uint64_t freed =
-                        try_evict_or_offload(it->first, it->second, shard);
-                    total_freed_size += freed;
+                    // Evict this object
+                    total_freed_size +=
+                        it->second.size *
+                        evict_replicas(it->second);  // Erase memory replicas
                     if (it->second.IsValid() == false) {
                         it = shard->metadata.erase(it);
                     } else {
                         ++it;
                     }
-                    if (freed > 0) {
-                        shard_evicted_count++;
-                    }
+                    shard_evicted_count++;
                 } else {
                     // second pass candidates
                     no_pin_objects.push_back(it->second.lease_timeout);
@@ -3717,22 +3598,20 @@ 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 (or defer for offload)
-                        uint64_t freed =
-                            try_evict_or_offload(it->first, it->second, shard);
-                        total_freed_size += freed;
+                        // Evict this object
+                        total_freed_size +=
+                            it->second.size *
+                            evict_replicas(
+                                it->second);  // Erase memory replicas
                         if (it->second.IsValid() == false) {
                             it = shard->metadata.erase(it);
                         } else {
                             ++it;
                         }
-                        if (freed > 0) {
-                            evicted_count++;
-                        }
+                        evicted_count++;
                         target_evict_num--;
                     } else {
                         ++it;
@@ -3772,18 +3651,16 @@ 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) {
-                        // Evict this object (or defer for offload)
-                        uint64_t freed =
-                            try_evict_or_offload(it->first, it->second, shard);
-                        total_freed_size += freed;
+                        total_freed_size +=
+                            it->second.size *
+                            evict_replicas(
+                                it->second);  // Erase memory replicas
                         if (it->second.IsValid() == false) {
                             it = shard->metadata.erase(it);
                         } else {
                             ++it;
                         }
-                        if (freed > 0) {
-                            evicted_count++;
-                        }
+                        evicted_count++;
                         target_evict_num--;
                     } else {
                         ++it;
@@ -3804,11 +3681,7 @@ void MasterService::BatchEvict(double evict_ratio_target,
         }
     }
 
-    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.
+    if (evicted_count > 0 || released_discarded_cnt > 0) {
         need_eviction_ = false;
         MasterMetricManager::instance().inc_eviction_success(evicted_count,
                                                              total_freed_size);
@@ -3820,27 +3693,7 @@ 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() {
@@ -3917,15 +3770,9 @@ void MasterService::ClientMonitorFunc() {
             }  // Release the mutex before long-running ClearInvalidHandles and
                // avoid deadlocks
 
-            // 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();
+            if (!unmount_segments.empty()) {
+                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++) {
@@ -3935,9 +3782,6 @@ 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,87 +62,6 @@ 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,
@@ -535,8 +454,7 @@ 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,
-    const std::string &ssd_offload_path) {
+    bool enable_ssd_offload, bool start_offload_rpc_server) {
     this->protocol = protocol;
     this->ipc_socket_path_ = ipc_socket_path;
     const bool should_use_hugepage = use_hugepage_ &&
@@ -575,7 +493,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, {{"client_mode", "real"}});
+            master_server_addr, transfer_engine);
         if (!client_opt) {
             LOG(ERROR) << "Failed to create client";
             return tl::unexpected(ErrorCode::INVALID_PARAMS);
@@ -604,7 +522,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, {{"client_mode", "real"}});
+                master_server_addr, transfer_engine);
             if (client_opt) {
                 client_ = *client_opt;
                 success = true;
@@ -797,9 +715,6 @@ 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());
@@ -827,12 +742,11 @@ 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 &ssd_offload_path) {
+    const std::string &ipc_socket_path, bool enable_ssd_offload) {
     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, ssd_offload_path));
+        ipc_socket_path, 50052, enable_ssd_offload, true));
 }
 
 namespace {
@@ -923,8 +837,6 @@ 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(),
@@ -936,7 +848,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, ssd_offload_path);
+                          enable_ssd_offload, true);
 }
 
 tl::expected<void, ErrorCode> RealClient::initAll_internal(
@@ -1165,17 +1077,8 @@ tl::expected<void, ErrorCode> RealClient::put_dummy_helper(
 
 int RealClient::put(const std::string &key, std::span<const char> value,
                     const ReplicateConfig &config) {
-    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);
+    return to_py_ret(
+        put_internal(key, value, config, client_buffer_allocator_));
 }
 
 tl::expected<void, ErrorCode> RealClient::put_batch_internal(
@@ -1264,18 +1167,8 @@ 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) {
-    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);
+    return to_py_ret(
+        put_batch_internal(keys, values, config, client_buffer_allocator_));
 }
 
 tl::expected<void, ErrorCode> RealClient::put_parts_internal(
@@ -1356,18 +1249,8 @@ 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) {
-    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);
+    return to_py_ret(
+        put_parts_internal(key, values, config, client_buffer_allocator_));
 }
 
 tl::expected<void, ErrorCode> RealClient::remove_internal(
@@ -1964,14 +1847,7 @@ 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 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);
-        });
+    return get_buffer_internal(key, client_buffer_allocator_);
 }
 
 tl::expected<std::tuple<uint64_t, size_t>, ErrorCode>
@@ -2261,14 +2137,7 @@ 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 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);
-        });
+    return batch_get_buffer_internal(keys);
 }
 
 tl::expected<void, ErrorCode> RealClient::register_buffer_internal(
@@ -2423,17 +2292,7 @@ tl::expected<int64_t, ErrorCode> RealClient::get_into_range_internal(
 
 int64_t RealClient::get_into(const std::string &key, void *buffer,
                              size_t size) {
-    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);
+    return to_py_ret(get_into_range_internal(key, buffer, 0, 0, size, true));
 }
 
 std::vector<std::vector<std::vector<tl::expected<int64_t, ErrorCode>>>>
@@ -2553,20 +2412,8 @@ 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) {
-    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;
+    return convert_ranged_read_results(get_into_ranges_internal(
+        buffers, all_keys, all_dst_offsets, all_src_offsets, all_sizes));
 }
 
 std::string RealClient::get_hostname() const { return local_hostname; }
@@ -2575,21 +2422,7 @@ 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 =
-        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);
-            });
+        batch_put_from_internal(keys, buffers, sizes, config);
     std::vector<int> results;
     results.reserve(internal_results.size());
 
@@ -2732,14 +2565,7 @@ tl::expected<void, ErrorCode> RealClient::put_from_internal(
 
 int RealClient::put_from(const std::string &key, void *buffer, size_t size,
                          const ReplicateConfig &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);
+    return to_py_ret(put_from_internal(key, buffer, size, config));
 }
 
 // --- Upsert implementations ---
@@ -2780,18 +2606,8 @@ tl::expected<void, ErrorCode> RealClient::upsert_internal(
 
 int RealClient::upsert(const std::string &key, std::span<const char> value,
                        const ReplicateConfig &config) {
-    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);
+    return to_py_ret(
+        upsert_internal(key, value, config, client_buffer_allocator_));
 }
 
 tl::expected<void, ErrorCode> RealClient::upsert_dummy_helper(
@@ -2841,14 +2657,7 @@ tl::expected<void, ErrorCode> RealClient::upsert_from_internal(
 
 int RealClient::upsert_from(const std::string &key, void *buffer, size_t size,
                             const ReplicateConfig &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);
+    return to_py_ret(upsert_from_internal(key, buffer, size, config));
 }
 
 std::vector<tl::expected<void, ErrorCode>>
@@ -2894,21 +2703,7 @@ 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 =
-        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);
-            });
+        batch_upsert_from_internal(keys, buffers, sizes, config);
     std::vector<int> results;
     results.reserve(internal_results.size());
     for (const auto &result : internal_results) {
@@ -3032,18 +2827,8 @@ 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) {
-    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);
+    return to_py_ret(
+        upsert_parts_internal(key, values, config, client_buffer_allocator_));
 }
 
 tl::expected<void, ErrorCode> RealClient::upsert_parts_dummy_helper(
@@ -3146,18 +2931,8 @@ 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) {
-    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);
+    return to_py_ret(
+        upsert_batch_internal(keys, values, config, client_buffer_allocator_));
 }
 
 // --- End Upsert implementations ---
@@ -3165,20 +2940,7 @@ 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 =
-        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);
-            });
+    auto internal_results = batch_get_into_internal(keys, buffers, sizes);
     std::vector<int64_t> results;
     results.reserve(internal_results.size());
 
@@ -3582,7 +3344,6 @@ 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) {
@@ -3625,10 +3386,6 @@ 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;
 }
 
@@ -3637,24 +3394,10 @@ 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 =
-        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);
-            });
+        batch_put_from_multi_buffers_internal(keys, all_buffers, sizes, config);
     std::vector<int> results;
     results.reserve(internal_results.size());
 
@@ -3662,6 +3405,10 @@ 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;
 }
 
@@ -3707,30 +3454,19 @@ 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 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);
-            });
+    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);
     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;
 }
 
@@ -4156,21 +3892,6 @@ 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);
@@ -4266,18 +3987,6 @@ 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,10 +91,6 @@ 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.

mooncake-store/src/segment.cpp

@@ -273,15 +273,6 @@ 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,9 +49,8 @@ FilePerKeyConfig FilePerKeyConfig::FromEnvironment() {
 
     config.fsdir = GetEnvStringOr("MOONCAKE_OFFLOAD_FSDIR", config.fsdir);
 
-    config.enable_eviction = GetEnvOr<bool>(
-        "MOONCAKE_OFFLOAD_ENABLE_EVICTION",
-        GetEnvOr<bool>("ENABLE_EVICTION", config.enable_eviction));
+    config.enable_eviction =
+        GetEnvOr<bool>("ENABLE_EVICTION", config.enable_eviction);
 
     return config;
 }
@@ -65,14 +64,11 @@ 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_OFFLOAD_BUCKET_MAX_TOTAL_SIZE",
-                          GetEnvOr<int64_t>("MOONCAKE_BUCKET_MAX_TOTAL_SIZE",
-                                            config.max_total_size));
+    config.max_total_size = GetEnvOr<int64_t>("MOONCAKE_BUCKET_MAX_TOTAL_SIZE",
+                                              config.max_total_size);
 
-    const auto policy_str = GetEnvStringOr(
-        "MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY",
-        GetEnvStringOr("MOONCAKE_BUCKET_EVICTION_POLICY", "none"));
+    const auto policy_str =
+        GetEnvStringOr("MOONCAKE_BUCKET_EVICTION_POLICY", "none");
     if (policy_str == "fifo") {
         config.eviction_policy = BucketEvictionPolicy::FIFO;
     } else if (policy_str == "lru") {
@@ -1722,18 +1718,6 @@ 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. "
@@ -2203,7 +2187,8 @@ BucketStorageBackend::PendingEviction BucketStorageBackend::PrepareEviction(
     int64_t required_size) {
     PendingEviction result;
 
-    if (bucket_backend_config_.eviction_policy == BucketEvictionPolicy::NONE) {
+    if (bucket_backend_config_.eviction_policy == BucketEvictionPolicy::NONE ||
+        bucket_backend_config_.max_total_size <= 0) {
         return result;
     }
 

mooncake-store/src/transfer_task.cpp

@@ -4,8 +4,6 @@
 
 #include <algorithm>
 #include <cstdlib>
-#include <sstream>
-#include <vector>
 #include "transfer_engine.h"
 #include "transport/transport.h"
 
@@ -225,13 +223,9 @@ bool TransferEngineOperationState::is_completed() {
 }
 
 void TransferEngineOperationState::check_task_status() {
-    // 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;
+    // Check all transfers in the batch
+    bool all_completed = true;
+    bool has_failure = false;
 
     for (size_t i = 0; i < batch_size_; ++i) {
         TransferStatus status;
@@ -246,45 +240,38 @@ void TransferEngineOperationState::check_task_status() {
 
         switch (status.s) {
             case TransferStatusEnum::COMPLETED:
-                // This transfer is done successfully
+                // This transfer is done, continue checking others
                 break;
             case TransferStatusEnum::FAILED:
             case TransferStatusEnum::CANCELED:
             case TransferStatusEnum::INVALID:
 #ifndef USE_ASCEND_DIRECT
-                VLOG(1) << "Transfer failed for batch " << batch_id_ << " task "
-                        << i << " with status " << static_cast<int>(status.s);
+                LOG(ERROR) << "Transfer failed for batch " << batch_id_
+                           << " task " << i << " with status "
+                           << static_cast<int>(status.s);
 #endif
-                failed_task_ids.push_back(i);
+                has_failure = true;
                 break;
             default:
-                // Transfer is still in progress (WAITING, PENDING, etc.)
-                all_terminated = false;
+                // Transfer is still pending (PENDING, RUNNING, etc.)
+                all_completed = false;
                 break;
         }
     }
 
-    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.
+    if (has_failure) {
+        VLOG(1) << "Setting batch " << batch_id_
+                << " result to TRANSFER_FAIL due to task failures";
+        set_result_internal(ErrorCode::TRANSFER_FAIL);
         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;
+    if (all_completed) {
+        set_result_internal(ErrorCode::OK);
+        return;
     }
 
-    set_result_internal(ec);
+    return;
 }
 
 void TransferEngineOperationState::set_result_internal(ErrorCode error_code) {
@@ -307,8 +294,7 @@ void TransferEngineOperationState::wait_for_completion() {
         return;
     }
 
-    // 60 seconds
-    constexpr int64_t timeout_milliseconds = 60 * 1000;
+    constexpr int64_t timeout_seconds = 60;
 
 #ifdef USE_EVENT_DRIVEN_COMPLETION
     VLOG(1) << "Waiting for transfer engine completion for batch " << batch_id_;
@@ -328,18 +314,10 @@ 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);
-        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);
-                });
-        }
+        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);
+            });
     }  // Explicitly release completion_mutex before acquiring mutex_
 
     // Once completion is observed, read failure flag.
@@ -360,18 +338,20 @@ 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_milliseconds << " milliseconds for batch "
-                   << batch_id_;
+        LOG(ERROR) << "Failed to complete transfers after " << timeout_seconds
+                   << " seconds 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 (getCurrentTimeInMilli() - start_ts_ > timeout_milliseconds) {
+        if (getCurrentTimeInNano() - start_ts >
+            timeout_seconds * kOneSecondInNano) {
             LOG(ERROR) << "Failed to complete transfers after "
-                       << timeout_milliseconds << " milliseconds for batch "
-                       << batch_id_;
+                       << timeout_seconds << " seconds for batch " << batch_id_;
             set_result_internal(ErrorCode::TRANSFER_FAIL);
             return;
         }
@@ -429,17 +409,10 @@ 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.
-    // 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)
+    // Read MC_STORE_MEMCPY environment variable, default to false (disabled)
     const char* env_value = std::getenv("MC_STORE_MEMCPY");
     if (env_value == nullptr) {
-        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");
+        memcpy_enabled_ = false;  // Default: disabled
     } else {
         std::string env_str(env_value);
         // Convert to lowercase for case-insensitive comparison

mooncake-store/tests/CMakeLists.txt

@@ -39,7 +39,6 @@ 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,8 +45,6 @@ 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);
@@ -111,23 +109,15 @@ 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 transfer, RPC, and interface metrics
+    // Should contain both transfer and RPC 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;
 }
@@ -185,33 +175,6 @@ 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,8 +26,7 @@ class FileStorageTest : public ::testing::Test {
         FLAGS_logtostderr = true;
         UnsetEnv("MOONCAKE_OFFLOAD_FILE_STORAGE_PATH");
         UnsetEnv("MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES");
-        UnsetEnv("MOONCAKE_OFFLOAD_SCANMETA_ITERATOR_KEYS_LIMIT");
-        UnsetEnv("MOONCAKE_SCANMETA_ITERATOR_KEYS_LIMIT");
+        UnsetEnv("MOONCAKE_OFFLOAD_BUCKET_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

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

mooncake-transfer-engine/benchmark/CMakeLists.txt

@@ -17,4 +17,4 @@ 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-common")
+                     INSTALL_RPATH "$ORIGIN/../lib:$ORIGIN/../../mooncake-asio")

mooncake-transfer-engine/example/batch_register_bench.py

@@ -1,357 +0,0 @@
-#!/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/kvcache_prefix_bench.py

@@ -1,452 +0,0 @@
-#!/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,10 +113,6 @@ 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);
@@ -331,13 +327,9 @@ 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 && std::chrono::steady_clock::now() < deadline) {
+    while (nbytes) {
         ssize_t rc = read(fd, pos, nbytes);
         if (rc < 0 && (errno == EAGAIN || errno == EINTR))
             continue;
@@ -352,14 +344,7 @@ static inline ssize_t readFully(int fd, void *buf, size_t len) {
         pos += rc;
         nbytes -= rc;
     }
-    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;
+    return len;
 }
 
 static inline int writeString(int fd, const HandShakeRequestType type,

mooncake-transfer-engine/include/multi_transport.h

@@ -55,14 +55,6 @@ 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,14 +150,6 @@ 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,8 +105,6 @@ 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,
@@ -121,11 +119,6 @@ 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,8 +334,6 @@ 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,7 +31,6 @@
 #include <memory>
 #include <string>
 #include <thread>
-#include <map>
 #include <unordered_map>
 
 #include "common.h"
@@ -207,7 +206,7 @@ class EfaContext {
     std::atomic<int> next_cq_index_{0};
 
     RWSpinlock mr_lock_;
-    std::map<uint64_t, EfaMemoryRegionMeta> mr_map_;
+    std::unordered_map<uint64_t, EfaMemoryRegionMeta> mr_map_;
 };
 
 }  // namespace mooncake

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

@@ -52,96 +52,82 @@ 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;
-
-    // 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);
+        const std::vector<void *> &addr_list) override;
 
    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
@@ -153,16 +139,6 @@ 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,8 +38,6 @@ 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;
@@ -60,8 +58,6 @@ 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;
@@ -93,8 +89,6 @@ 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,9 +99,6 @@ 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

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

mooncake-transfer-engine/rust/build.rs

@@ -17,84 +17,16 @@ 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=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");
-    }
+    println!("cargo:rustc-link-lib=etcd-cpp-api");
 
     let bindings = bindgen::builder()
         .header("../include/transfer_engine_c.h")

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

@@ -91,27 +91,6 @@ 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);
@@ -263,20 +242,6 @@ 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/multi_transport.cpp

@@ -461,10 +461,6 @@ 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,8 +179,6 @@ 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);
 }
@@ -580,15 +578,6 @@ 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,9 +19,6 @@
 
 #include "transfer_engine.h"
 #include "transport/transport.h"
-#ifdef USE_EFA
-#include "transport/efa_transport/efa_transport.h"
-#endif
 
 using namespace mooncake;
 
@@ -39,11 +36,6 @@ 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();
@@ -85,21 +77,6 @@ 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,9 +29,7 @@
 
 #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

@@ -164,15 +164,6 @@ int EfaContext::construct(size_t num_cq_list, size_t num_comp_channels,
                           int max_endpoints) {
     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();
     if (!hints_) {
@@ -190,11 +181,8 @@ int EfaContext::construct(size_t num_cq_list, size_t num_comp_channels,
     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
-#if defined(USE_CUDA) || defined(USE_HIP)
-                                   | FI_MR_HMEM
-#endif
-        ;
+                                   FI_MR_ALLOCATED | FI_MR_PROV_KEY |
+                                   FI_MR_HMEM;
     hints_->domain_attr->threading = FI_THREAD_SAFE;
 
     // Get fabric info
@@ -441,35 +429,31 @@ int EfaContext::preTouchMemory(void* addr, size_t length) {
 
 uint64_t EfaContext::rkey(void* addr) {
     RWSpinlock::ReadGuard guard(mr_lock_);
-    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;
-        }
+    auto it = mr_map_.find((uint64_t)addr);
+    if (it != mr_map_.end() && it->second.mr) {
+        return it->second.key;
     }
     return 0;
 }
 
 uint64_t EfaContext::lkey(void* addr) {
     RWSpinlock::ReadGuard guard(mr_lock_);
-    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);
-        }
+    auto it = mr_map_.find((uint64_t)addr);
+    if (it != mr_map_.end() && it->second.mr) {
+        return fi_mr_key(it->second.mr);
     }
     return 0;
 }
 
 void* EfaContext::mrDesc(void* addr) {
     RWSpinlock::ReadGuard guard(mr_lock_);
-    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);
+    // 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);
+            }
         }
     }
     return nullptr;

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

@@ -17,13 +17,10 @@
 #include <glog/logging.h>
 #include <sys/mman.h>
 #include <sys/time.h>
-#include <unistd.h>
 
-#include <algorithm>
 #include <cassert>
 #include <chrono>
 #include <cstddef>
-#include <fstream>
 #include <future>
 #include <set>
 #include <thread>
@@ -39,57 +36,6 @@
 
 namespace mooncake {
 
-// Default PTE (page table entry) limit per EFA NIC.  EFA hardware supports
-// roughly 24 million PTEs per NIC, but we use 22M as a conservative default.
-// With 4KB pages: 22M × 4KB ≈ 88GB per NIC.
-// With 2MB hugepages: 22M × 2MB ≈ 44TB per NIC (effectively unlimited).
-// Override via MC_EFA_MAX_PTE_ENTRIES environment variable.
-static constexpr size_t kDefaultMaxPteEntries = 22ULL * 1024 * 1024;  // 22M
-
-// Detect the kernel page size backing the memory at `addr` by reading
-// /proc/self/smaps.  Falls back to sysconf(_SC_PAGESIZE) on any failure.
-static size_t detectBufferPageSize(void* addr) {
-    size_t fallback = static_cast<size_t>(sysconf(_SC_PAGESIZE));
-    std::ifstream smaps("/proc/self/smaps");
-    if (!smaps.is_open()) return fallback;
-
-    uintptr_t target = reinterpret_cast<uintptr_t>(addr);
-    std::string line;
-    bool in_range = false;
-
-    while (std::getline(smaps, line)) {
-        // VMA header: "start-end perms offset dev inode [pathname]"
-        if (!line.empty() && std::isxdigit(line[0])) {
-            unsigned long start = 0, end = 0;
-            if (sscanf(line.c_str(), "%lx-%lx", &start, &end) == 2) {
-                in_range = (target >= start && target < end);
-            }
-        } else if (in_range && line.compare(0, 15, "KernelPageSize:") == 0) {
-            unsigned long kb = 0;
-            if (sscanf(line.c_str(), "KernelPageSize: %lu kB", &kb) == 1 &&
-                kb > 0) {
-                return kb * 1024;
-            }
-        }
-    }
-    return fallback;
-}
-
-static size_t getMaxPteEntries() {
-    static size_t cached = []() {
-        const char* env = std::getenv("MC_EFA_MAX_PTE_ENTRIES");
-        if (env) {
-            size_t val = std::stoull(env);
-            if (val > 0) {
-                LOG(INFO) << "MC_EFA_MAX_PTE_ENTRIES override: " << val;
-                return val;
-            }
-        }
-        return kDefaultMaxPteEntries;
-    }();
-    return cached;
-}
-
 EfaTransport::EfaTransport() {
     LOG(INFO) << "[EFA] AWS Elastic Fabric Adapter transport initialized";
 }
@@ -211,6 +157,9 @@ int EfaTransport::preTouchMemory(void* addr, size_t length) {
 
     auto hwc = std::thread::hardware_concurrency();
     auto num_threads = hwc > 64 ? 16 : std::min(hwc, 8u);
+    if (length > (size_t)globalConfig().max_mr_size) {
+        length = (size_t)globalConfig().max_mr_size;
+    }
     size_t block_size = length / num_threads;
     if (block_size == 0) {
         return 0;
@@ -256,266 +205,105 @@ int EfaTransport::registerLocalMemoryInternal(void* addr, size_t length,
                                               bool update_metadata,
                                               bool force_sequential) {
     (void)remote_accessible;
+    BufferDesc buffer_desc;
     const int kBaseAccessRights = IBV_ACCESS_LOCAL_WRITE |
                                   IBV_ACCESS_REMOTE_WRITE |
                                   IBV_ACCESS_REMOTE_READ;
 
     int access_rights = kBaseAccessRights;
-    size_t max_mr = (size_t)globalConfig().max_mr_size;
 
-    // Compute chunk limit based on EFA PTE (page table entry) constraints.
-    // Each EFA NIC has a hardware limit on PTE entries (~24M).  The effective
-    // per-NIC MR size limit depends on the backing page size:
-    //   4KB pages  → 22M × 4KB  ≈ 88GB  (smaller than device max_mr_size)
-    //   2MB hugepg → 22M × 2MB  ≈ 44TB  (device max_mr_size is the limit)
-    // We detect the actual page size of the buffer and compute accordingly,
-    // so hugepage-backed memory avoids unnecessary splitting.
-    size_t page_size = detectBufferPageSize(addr);
-    size_t pte_limit = getMaxPteEntries() * page_size;
-    // When max_mr_size is not configured, fall back to pte_limit so that
-    // PTE-aware splitting still kicks in for large buffers on 4KB pages.
-    size_t chunk_limit = (max_mr > 0) ? std::min(max_mr, pte_limit) : pte_limit;
-    LOG(INFO) << "Auto-split params: page_size=" << page_size
-              << ", max_pte_entries=" << getMaxPteEntries()
-              << ", pte_limit=" << pte_limit << ", max_mr_size=" << max_mr
-              << ", chunk_limit=" << chunk_limit;
-
-    // Determine chunk boundaries
-    std::vector<std::pair<void*, size_t>> chunks;
-    if (length > chunk_limit) {
-        size_t offset = 0;
-        while (offset < length) {
-            size_t chunk_len = std::min(chunk_limit, length - offset);
-            chunks.emplace_back(static_cast<char*>(addr) + offset, chunk_len);
-            offset += chunk_len;
+    bool do_pre_touch = context_list_.size() > 0 &&
+                        std::thread::hardware_concurrency() >= 4 &&
+                        length >= (size_t)4 * 1024 * 1024 * 1024;
+    if (do_pre_touch) {
+        int ret = preTouchMemory(addr, length);
+        if (ret != 0) {
+            return ret;
         }
-        LOG(WARNING) << "Auto-splitting buffer " << addr << " (" << length
-                     << " bytes) into " << chunks.size()
-                     << " chunks of <= " << chunk_limit << " bytes each";
-    } else {
-        chunks.emplace_back(addr, length);
     }
 
-    // Resolve location name once (based on original buffer)
-    std::string resolved_name;
+    int use_parallel_reg = 0;
+    if (!force_sequential) {
+        use_parallel_reg = globalConfig().parallel_reg_mr;
+        if (use_parallel_reg == -1) {
+            use_parallel_reg = context_list_.size() > 1 && do_pre_touch;
+        }
+    }
+
+    auto reg_start = std::chrono::steady_clock::now();
+
+    if (use_parallel_reg) {
+        std::vector<std::thread> reg_threads;
+        reg_threads.reserve(context_list_.size());
+        std::vector<int> ret_codes(context_list_.size(), 0);
+        const int ar = access_rights;
+
+        for (size_t i = 0; i < context_list_.size(); ++i) {
+            reg_threads.emplace_back([this, &ret_codes, i, addr, length, ar]() {
+                ret_codes[i] =
+                    context_list_[i]->registerMemoryRegion(addr, length, ar);
+            });
+        }
+
+        for (auto& thread : reg_threads) {
+            thread.join();
+        }
+
+        for (size_t i = 0; i < ret_codes.size(); ++i) {
+            if (ret_codes[i] != 0) {
+                LOG(ERROR)
+                    << "Failed to register memory region with EFA context "
+                    << i;
+                return ret_codes[i];
+            }
+        }
+    } else {
+        for (size_t i = 0; i < context_list_.size(); ++i) {
+            int ret = context_list_[i]->registerMemoryRegion(addr, length,
+                                                             access_rights);
+            if (ret) {
+                LOG(ERROR)
+                    << "Failed to register memory region with EFA context "
+                    << i;
+                return ret;
+            }
+        }
+    }
+
+    auto reg_end = std::chrono::steady_clock::now();
+    auto reg_duration_ms =
+        std::chrono::duration_cast<std::chrono::milliseconds>(reg_end -
+                                                              reg_start)
+            .count();
+
+    if (globalConfig().trace) {
+        LOG(INFO) << "EFA registerMemoryRegion: addr=" << addr
+                  << ", length=" << length
+                  << ", contexts=" << context_list_.size()
+                  << ", parallel=" << (use_parallel_reg ? "true" : "false")
+                  << ", duration=" << reg_duration_ms << "ms";
+    }
+
+    // Collect keys from all contexts
+    for (auto& context : context_list_) {
+        buffer_desc.lkey.push_back(context->lkey(addr));
+        buffer_desc.rkey.push_back(context->rkey(addr));
+    }
+
     if (name == kWildcardLocation) {
         bool only_first_page = true;
         const std::vector<MemoryLocationEntry> entries =
             getMemoryLocation(addr, length, only_first_page);
         if (entries.empty()) return -1;
-        resolved_name = entries[0].location;
+        buffer_desc.name = entries[0].location;
     } else {
-        resolved_name = name;
-    }
-
-    // Pre-compute NIC assignments for each chunk.
-    // Strategy: if total PTE usage per NIC fits within the PTE budget,
-    // register ALL chunks on ALL NICs (full coverage → max throughput).
-    // Otherwise fall back to disjoint per-NIC partition.
-    size_t num_nics = context_list_.size();
-    size_t num_chunks = chunks.size();
-    size_t total_pages_per_nic = length / page_size;
-    bool use_full_coverage = (total_pages_per_nic <= getMaxPteEntries());
-
-    std::vector<std::vector<size_t>> nic_assignments(num_chunks);
-    if (chunks.size() <= 1) {
-        // Single chunk: all NICs
-        for (size_t n = 0; n < num_nics; ++n) {
-            nic_assignments[0].push_back(n);
-        }
-    } else if (use_full_coverage) {
-        // Multi-chunk, PTE budget OK: every chunk on every NIC
-        LOG(WARNING) << "Full NIC coverage: " << num_chunks << " chunks × "
-                     << num_nics
-                     << " NICs (total PTE/NIC=" << total_pages_per_nic
-                     << ", budget=" << getMaxPteEntries() << ")";
-        for (size_t ci = 0; ci < num_chunks; ++ci) {
-            for (size_t n = 0; n < num_nics; ++n) {
-                nic_assignments[ci].push_back(n);
-            }
-        }
-    } else if (num_chunks <= num_nics) {
-        // Multi-chunk, PTE exceeded, more NICs than chunks: disjoint partition
-        for (size_t ci = 0; ci < num_chunks; ++ci) {
-            size_t nics_per = num_nics / num_chunks;
-            size_t extra = num_nics % num_chunks;
-            size_t start = ci * nics_per + std::min(ci, extra);
-            size_t count = nics_per + (ci < extra ? 1 : 0);
-            for (size_t n = start; n < start + count; ++n) {
-                nic_assignments[ci].push_back(n);
-            }
-        }
-        LOG(WARNING) << "Disjoint NIC partition: PTE/NIC="
-                     << total_pages_per_nic
-                     << " exceeds budget=" << getMaxPteEntries();
-        for (size_t ci = 0; ci < num_chunks; ++ci) {
-            std::string nic_list;
-            for (size_t j = 0; j < nic_assignments[ci].size(); ++j) {
-                if (j > 0) nic_list += ",";
-                nic_list += std::to_string(nic_assignments[ci][j]);
-            }
-            LOG(WARNING) << "  chunk " << ci << " -> NICs [" << nic_list << "]";
-        }
-    } else {
-        // Multi-chunk, PTE exceeded, more chunks than NICs: round-robin
-        // Each NIC gets multiple chunks; verify per-NIC PTE stays in budget.
-        size_t pte_budget = getMaxPteEntries();
-        std::vector<size_t> pages_per_nic(num_nics, 0);
-        for (size_t ci = 0; ci < num_chunks; ++ci) {
-            size_t nic = ci % num_nics;
-            size_t chunk_pages = chunks[ci].second / page_size;
-            pages_per_nic[nic] += chunk_pages;
-            nic_assignments[ci].push_back(nic);
-        }
-        bool pte_ok = true;
-        for (size_t n = 0; n < num_nics; ++n) {
-            if (pages_per_nic[n] > pte_budget) {
-                pte_ok = false;
-                break;
-            }
-        }
-        if (!pte_ok) {
-            LOG(ERROR) << "Buffer requires " << num_chunks << " chunks ("
-                       << length << " bytes) but per-NIC PTE budget ("
-                       << pte_budget << " entries, page_size=" << page_size
-                       << ") is exceeded even with round-robin across "
-                       << num_nics << " NICs";
-            return ERR_INVALID_ARGUMENT;
-        }
-        LOG(WARNING) << "Round-robin NIC assignment: " << num_chunks
-                     << " chunks across " << num_nics << " NICs";
-        for (size_t ci = 0; ci < num_chunks; ++ci) {
-            LOG(WARNING) << "  chunk " << ci << " ("
-                         << chunks[ci].second / (1024 * 1024) << " MB) -> NIC "
-                         << nic_assignments[ci][0];
-        }
-    }
-
-    auto rollbackChunks = [&](size_t up_to_ci) {
-        for (size_t ri = 0; ri <= up_to_ci; ++ri) {
-            for (size_t nic_idx : nic_assignments[ri]) {
-                context_list_[nic_idx]->unregisterMemoryRegion(
-                    chunks[ri].first);
-            }
-        }
-    };
-
-    // Register each chunk on its assigned NICs
-    for (size_t ci = 0; ci < chunks.size(); ++ci) {
-        void* chunk_addr = chunks[ci].first;
-        size_t chunk_len = chunks[ci].second;
-        const auto& assigned_nics = nic_assignments[ci];
-
-        bool do_pre_touch = context_list_.size() > 0 &&
-                            std::thread::hardware_concurrency() >= 4 &&
-                            chunk_len >= (size_t)4 * 1024 * 1024 * 1024;
-        if (do_pre_touch) {
-            int ret = preTouchMemory(chunk_addr, chunk_len);
-            if (ret != 0) {
-                if (ci > 0) rollbackChunks(ci - 1);
-                return ret;
-            }
-        }
-
-        int use_parallel_reg = 0;
-        if (!force_sequential) {
-            use_parallel_reg = globalConfig().parallel_reg_mr;
-            if (use_parallel_reg == -1) {
-                use_parallel_reg = assigned_nics.size() > 1 && do_pre_touch;
-            }
-        }
-
-        auto reg_start = std::chrono::steady_clock::now();
-
-        if (use_parallel_reg) {
-            std::vector<std::thread> reg_threads;
-            reg_threads.reserve(assigned_nics.size());
-            std::vector<int> ret_codes(assigned_nics.size(), 0);
-            const int ar = access_rights;
-
-            for (size_t j = 0; j < assigned_nics.size(); ++j) {
-                size_t nic_idx = assigned_nics[j];
-                reg_threads.emplace_back([this, &ret_codes, j, nic_idx,
-                                          chunk_addr, chunk_len, ar]() {
-                    ret_codes[j] = context_list_[nic_idx]->registerMemoryRegion(
-                        chunk_addr, chunk_len, ar);
-                });
-            }
-
-            for (auto& thread : reg_threads) {
-                thread.join();
-            }
-
-            for (size_t j = 0; j < ret_codes.size(); ++j) {
-                if (ret_codes[j] != 0) {
-                    LOG(ERROR)
-                        << "Failed to register memory region chunk " << ci
-                        << " with EFA context " << assigned_nics[j];
-                    rollbackChunks(ci);
-                    return ret_codes[j];
-                }
-            }
-        } else {
-            for (size_t nic_idx : assigned_nics) {
-                int ret = context_list_[nic_idx]->registerMemoryRegion(
-                    chunk_addr, chunk_len, access_rights);
-                if (ret) {
-                    LOG(ERROR) << "Failed to register memory region chunk "
-                               << ci << " with EFA context " << nic_idx;
-                    rollbackChunks(ci);
-                    return ret;
-                }
-            }
-        }
-
-        auto reg_end = std::chrono::steady_clock::now();
-        auto reg_duration_ms =
-            std::chrono::duration_cast<std::chrono::milliseconds>(reg_end -
-                                                                  reg_start)
-                .count();
-
-        if (globalConfig().trace) {
-            LOG(INFO) << "EFA registerMemoryRegion: chunk " << ci
-                      << ", addr=" << chunk_addr << ", length=" << chunk_len
-                      << ", nics=" << assigned_nics.size() << "/"
-                      << context_list_.size()
-                      << ", parallel=" << (use_parallel_reg ? "true" : "false")
-                      << ", duration=" << reg_duration_ms << "ms";
-        }
-
-        LOG(WARNING) << "Chunk " << ci << "/" << chunks.size()
-                     << " registered on " << assigned_nics.size() << " NICs"
-                     << ", addr=" << chunk_addr << ", length=" << chunk_len
-                     << ", duration=" << reg_duration_ms << "ms";
-
-        // Collect keys: assigned NICs have valid keys, others get 0
-        BufferDesc buffer_desc;
-        for (auto& context : context_list_) {
-            buffer_desc.lkey.push_back(context->lkey(chunk_addr));
-            buffer_desc.rkey.push_back(context->rkey(chunk_addr));
-        }
-
-        buffer_desc.name = resolved_name;
-        buffer_desc.addr = (uint64_t)chunk_addr;
-        buffer_desc.length = chunk_len;
-        int rc = metadata_->addLocalMemoryBuffer(buffer_desc, update_metadata);
-        if (rc) {
-            rollbackChunks(ci);
-            return rc;
-        }
-    }
-
-    // Track chunks and NIC assignments for unregistration
-    if (chunks.size() > 1) {
-        std::lock_guard<std::mutex> lock(chunk_map_mutex_);
-        std::vector<ChunkRegistration> regs;
-        regs.reserve(chunks.size());
-        for (size_t ci = 0; ci < chunks.size(); ++ci) {
-            regs.push_back({(uint64_t)chunks[ci].first, nic_assignments[ci]});
-        }
-        chunk_map_[(uint64_t)addr] = std::move(regs);
+        buffer_desc.name = name;
     }
 
+    buffer_desc.addr = (uint64_t)addr;
+    buffer_desc.length = length;
+    int rc = metadata_->addLocalMemoryBuffer(buffer_desc, update_metadata);
+    if (rc) return rc;
     return 0;
 }
 
@@ -526,40 +314,6 @@ int EfaTransport::unregisterLocalMemory(void* addr, bool update_metadata) {
 int EfaTransport::unregisterLocalMemoryInternal(void* addr,
                                                 bool update_metadata,
                                                 bool force_sequential) {
-    // Check if this buffer was split into chunks (per-NIC partition)
-    std::vector<ChunkRegistration> chunk_regs;
-    {
-        std::lock_guard<std::mutex> lock(chunk_map_mutex_);
-        auto it = chunk_map_.find((uint64_t)addr);
-        if (it != chunk_map_.end()) {
-            chunk_regs = std::move(it->second);
-            chunk_map_.erase(it);
-        }
-    }
-
-    if (!chunk_regs.empty()) {
-        // Unregister each chunk from its assigned NICs only
-        for (auto& reg : chunk_regs) {
-            void* ca = (void*)reg.addr;
-            int rc = metadata_->removeLocalMemoryBuffer(ca, update_metadata);
-            if (rc) {
-                LOG(ERROR) << "Failed to remove chunk metadata at " << ca;
-                return rc;
-            }
-
-            for (size_t nic_idx : reg.nic_indices) {
-                int ret = context_list_[nic_idx]->unregisterMemoryRegion(ca);
-                if (ret) {
-                    LOG(ERROR) << "Failed to unregister chunk " << ca
-                               << " with EFA context " << nic_idx;
-                    return ret;
-                }
-            }
-        }
-        return 0;
-    }
-
-    // Non-chunked buffer: original path
     int rc = metadata_->removeLocalMemoryBuffer(addr, update_metadata);
     if (rc) return rc;
 
@@ -669,79 +423,6 @@ int EfaTransport::unregisterLocalMemoryBatch(
     return metadata_->updateLocalSegmentDesc();
 }
 
-int EfaTransport::warmupSegment(const std::string& segment_name) {
-    if (!metadata_) {
-        LOG(ERROR) << "EfaTransport::warmupSegment: metadata_ is null";
-        return ERR_INVALID_ARGUMENT;
-    }
-    if (segment_name.empty() || segment_name == local_server_name_) {
-        // Loopback / empty name — nothing to pre-connect.
-        return 0;
-    }
-
-    auto desc = metadata_->getSegmentDescByName(segment_name);
-    if (!desc) {
-        LOG(ERROR) << "EfaTransport::warmupSegment: segment '" << segment_name
-                   << "' not found in metadata (did you openSegment() first?)";
-        return ERR_INVALID_ARGUMENT;
-    }
-    if (desc->devices.empty()) {
-        LOG(WARNING) << "EfaTransport::warmupSegment: segment '" << segment_name
-                     << "' has no devices";
-        return 0;
-    }
-
-    // Build peer_nic_path list: "<segment_name>@<device_name>" for each NIC.
-    std::vector<std::string> peer_paths;
-    peer_paths.reserve(desc->devices.size());
-    for (const auto& dev : desc->devices) {
-        peer_paths.emplace_back(segment_name + "@" + dev.name);
-    }
-
-    // Warm up every (local_ctx, peer_nic) pair concurrently. Each
-    // EfaContext::endpoint() + setupConnectionsByActive() is idempotent and
-    // takes its own lock, so parallel calls across distinct peer_nic_paths
-    // (and distinct contexts) are safe. We use std::async to get roughly
-    // per-pair parallelism — the critical path is now max(handshake RTT) not
-    // sum(handshake RTT).
-    auto t0 = std::chrono::steady_clock::now();
-    size_t n_pairs = context_list_.size() * peer_paths.size();
-    std::vector<std::future<int>> futs;
-    futs.reserve(n_pairs);
-    for (auto& ctx : context_list_) {
-        for (const auto& path : peer_paths) {
-            futs.emplace_back(
-                std::async(std::launch::async, [ctx, path]() -> int {
-                    auto ep = ctx->endpoint(path);
-                    if (!ep) {
-                        LOG(WARNING) << "warmupSegment: endpoint() returned "
-                                        "null for "
-                                     << path;
-                        return -1;
-                    }
-                    if (ep->connected()) return 0;
-                    return ep->setupConnectionsByActive();
-                }));
-        }
-    }
-    int ok = 0, fail = 0;
-    for (auto& f : futs) {
-        int rc = f.get();
-        if (rc == 0)
-            ++ok;
-        else
-            ++fail;
-    }
-    auto elapsed =
-        std::chrono::duration<double>(std::chrono::steady_clock::now() - t0)
-            .count();
-    LOG(INFO) << "EfaTransport::warmupSegment('" << segment_name << "'): " << ok
-              << "/" << n_pairs << " endpoints connected (" << fail
-              << " failed) in " << elapsed << "s (" << context_list_.size()
-              << " local NICs × " << peer_paths.size() << " peer NICs)";
-    return fail == 0 ? 0 : ERR_ENDPOINT;
-}
-
 Status EfaTransport::submitTransfer(
     BatchID batch_id, const std::vector<TransferRequest>& entries) {
     auto& batch_desc = *((BatchDesc*)(batch_id));
@@ -806,9 +487,7 @@ Status EfaTransport::submitTransferTask(
                 if (static_cast<size_t>(request_buffer_id) <
                         local_segment_desc->buffers.size() &&
                     local_segment_desc->buffers[request_buffer_id].lkey.size() >
-                        static_cast<size_t>(d) &&
-                    local_segment_desc->buffers[request_buffer_id].lkey[d] !=
-                        0) {
+                        static_cast<size_t>(d)) {
                     active_devs.push_back({d, ctx});
                 }
             }
@@ -1109,39 +788,6 @@ int EfaTransport::initializeEfaResources() {
         LOG(ERROR) << "EfaTransport: No available EFA devices";
         return ERR_DEVICE_NOT_FOUND;
     }
-
-    // Query EFA device max_mr_size via ibverbs and clamp globalConfig.
-    // libfabric does not expose max_mr_size, so we go through the ibverbs
-    // layer.
-    {
-        int num_devices = 0;
-        struct ibv_device** dev_list = ibv_get_device_list(&num_devices);
-        if (dev_list) {
-            const std::string& first_efa = efa_devices[0];
-            for (int i = 0; i < num_devices; ++i) {
-                if (first_efa == ibv_get_device_name(dev_list[i])) {
-                    struct ibv_context* ctx = ibv_open_device(dev_list[i]);
-                    if (ctx) {
-                        struct ibv_device_attr attr;
-                        if (ibv_query_device(ctx, &attr) == 0) {
-                            auto& config = globalConfig();
-                            if (config.max_mr_size >
-                                (uint64_t)attr.max_mr_size) {
-                                config.max_mr_size = attr.max_mr_size;
-                                LOG(INFO) << "EfaTransport: Clamped "
-                                             "max_mr_size to device limit: "
-                                          << config.max_mr_size;
-                            }
-                        }
-                        ibv_close_device(ctx);
-                    }
-                    break;
-                }
-            }
-            ibv_free_device_list(dev_list);
-        }
-    }
-
     return 0;
 }
 
@@ -1167,31 +813,16 @@ int EfaTransport::selectDevice(SegmentDesc* desc, uint64_t offset,
             continue;
         }
 
-        // Try multiple attempts to find a device with valid MR registration.
-        // With per-NIC partition, not all devices have all buffers registered,
-        // so rkey[device_id] may be 0 for unassigned NICs.
-        int num_devices = static_cast<int>(desc->devices.size());
-        for (int attempt = 0; attempt < num_devices; ++attempt) {
-            int try_count = retry_count + attempt;
-            device_id =
-                hint.empty()
-                    ? desc->topology.selectDevice(buffer.name, try_count)
-                    : desc->topology.selectDevice(buffer.name, hint, try_count);
-            if (device_id >= 0 &&
-                static_cast<size_t>(device_id) < buffer.rkey.size() &&
-                buffer.rkey[device_id] != 0) {
-                return 0;
-            }
-            device_id = hint.empty() ? desc->topology.selectDevice(
-                                           kWildcardLocation, try_count)
-                                     : desc->topology.selectDevice(
-                                           kWildcardLocation, hint, try_count);
-            if (device_id >= 0 &&
-                static_cast<size_t>(device_id) < buffer.rkey.size() &&
-                buffer.rkey[device_id] != 0) {
-                return 0;
-            }
-        }
+        device_id =
+            hint.empty()
+                ? desc->topology.selectDevice(buffer.name, retry_count)
+                : desc->topology.selectDevice(buffer.name, hint, retry_count);
+        if (device_id >= 0) return 0;
+        device_id = hint.empty() ? desc->topology.selectDevice(
+                                       kWildcardLocation, retry_count)
+                                 : desc->topology.selectDevice(
+                                       kWildcardLocation, hint, retry_count);
+        if (device_id >= 0) return 0;
     }
     return ERR_ADDRESS_NOT_REGISTERED;
 }

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

@@ -35,17 +35,6 @@ 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_);
@@ -150,17 +139,6 @@ 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,11 +357,6 @@ 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,8 +161,9 @@ int WorkerPool::submitPostSend(
         slice_queue_lock_[shard_id].unlock();
     }
 
-    submitted_slice_count_.fetch_add(submitted_slice_count);
-    if (suspended_flag_.load()) {
+    submitted_slice_count_.fetch_add(submitted_slice_count,
+                                     std::memory_order_relaxed);
+    if (suspended_flag_.load(std::memory_order_relaxed)) {
         std::lock_guard<std::mutex> lock(cond_mutex_);
         cond_var_.notify_all();
     }
@@ -397,7 +398,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()) {
+                    submitted_slice_count_.load(std::memory_order_relaxed)) {
                     cond_var_.wait_for(lock, std::chrono::seconds(1));
                 }
                 suspended_flag_.fetch_sub(1);
@@ -414,62 +415,19 @@ 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_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);
-        }
+        auto endpoint = (RdmaEndPoint *)event.element.qp->qp_context;
+        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";
@@ -478,11 +436,7 @@ int WorkerPool::doProcessContextEvents() {
         LOG(INFO) << "Worker: Context " << context_.deviceName()
                   << " is now active";
     }
-
-    if (!event_acked) {
-        ibv_ack_async_event(&event);
-    }
-
+    ibv_ack_async_event(&event);
     return 0;
 }
 

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

@@ -1,213 +0,0 @@
-// 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

@@ -37,7 +37,6 @@ 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 {
@@ -84,7 +83,7 @@ class MnnvlTransport : public Transport {
     virtual Status freeLocalMemory(void *addr, size_t size);
 
    private:
-    void startTransfer(std::vector<MnnvlTask *> &tasks, MnnvlSubBatch *batch);
+    void startTransfer(MnnvlTask *task, MnnvlSubBatch *batch);
 
     void *createSharedMemory(const std::string &path, size_t size);
 

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

@@ -40,7 +40,6 @@ 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 {
@@ -82,7 +81,7 @@ class NVLinkTransport : public Transport {
     virtual const char *getName() const { return "nvlink"; }
 
    private:
-    void startTransfer(std::vector<NVLinkTask *> &tasks, NVLinkSubBatch *batch);
+    void startTransfer(NVLinkTask *task, NVLinkSubBatch *batch);
 
     void *createSharedMemory(const std::string &path, size_t size);
 

mooncake-transfer-engine/tent/src/metrics/CMakeLists.txt

@@ -7,7 +7,7 @@ add_library(tent_metrics STATIC ${TENT_METRICS_SOURCES})
 # ODR safety: yalantinglibs bundles ASIO headers but does NOT compile ASIO inline because
 # ASIO_SEPARATE_COMPILATION is set globally. All ASIO symbols live exclusively in asio_shared.so,
 # so there is no risk of duplicate symbols between yalantinglibs and the rest of TE.
-# (See mooncake-common/src/CMakeLists.txt: "Build asio as a shared library to avoid ODR violations")
+# (See mooncake-asio/CMakeLists.txt: "Build asio as a shared library to avoid ODR violations")
 target_link_libraries(tent_metrics PUBLIC tent_common tent_interface yalantinglibs::yalantinglibs glog pthread)
 
 # Pass compile definition based on option
@@ -17,4 +17,4 @@ if(TENT_METRICS_ENABLED)
 else()
     target_compile_definitions(tent_metrics PUBLIC TENT_METRICS_ENABLED=0)
     message(STATUS "TENT metrics: DISABLED (zero overhead)")
-endif()
+endif()

mooncake-transfer-engine/tent/src/transport/mnnvl/mnnvl_transport.cpp

@@ -168,7 +168,6 @@ Status MnnvlTransport::submitTransferTasks(
     if (request_list.size() + mnnvl_batch->task_list.size() >
         mnnvl_batch->max_size)
         return Status::TooManyRequests("Exceed batch capacity" LOC_MARK);
-    std::vector<MnnvlTask *> new_tasks;
     for (auto &request : request_list) {
         mnnvl_batch->task_list.push_back(MnnvlTask{});
         auto &task = mnnvl_batch->task_list[mnnvl_batch->task_list.size() - 1];
@@ -185,84 +184,72 @@ Status MnnvlTransport::submitTransferTasks(
         task.target_addr = target_addr;
         task.request = request;
         task.status_word = TransferStatusEnum::PENDING;
-        new_tasks.push_back(&task);
+        startTransfer(&task, mnnvl_batch);
     }
-    startTransfer(new_tasks, mnnvl_batch);
     return Status::OK();
 }
 
-void MnnvlTransport::startTransfer(std::vector<MnnvlTask *> &tasks,
-                                   MnnvlSubBatch *batch) {
-    if (tasks.empty()) return;
-
-    std::vector<void *> srcs;
-    std::vector<void *> dsts;
-    std::vector<size_t> sizes;
-    for (auto *task : tasks) {
-        void *src = nullptr, *dst = nullptr;
-        if (task->request.opcode == Request::READ) {
-            dst = task->request.source;       // read into source buffer
-            src = (void *)task->target_addr;  // from remote
-        } else {
-            src = task->request.source;       // write from source buffer
-            dst = (void *)task->target_addr;  // to remote
-        }
-        srcs.push_back(src);
-        dsts.push_back(dst);
-        sizes.push_back(task->request.length);
-    }
-
+void MnnvlTransport::startTransfer(MnnvlTask *task, MnnvlSubBatch *batch) {
     cudaError_t err;
+    void *src = nullptr, *dst = nullptr;
 
-#if CUDART_VERSION >= 13000
-    cudaMemcpyAttributes attr{};
-    attr.srcAccessOrder = cudaMemcpySrcAccessOrderStream;
-    size_t attrs_idx = 0;
-    err = cudaMemcpyBatchAsync(const_cast<const void **>(dsts.data()),
-                               const_cast<const void **>(srcs.data()),
-                               sizes.data(), srcs.size(), &attr, &attrs_idx, 1,
-                               batch->async_stream.get());
-#elif CUDART_VERSION >= 12080
-    cudaMemcpyAttributes attr{};
-    attr.srcAccessOrder = cudaMemcpySrcAccessOrderStream;
-    size_t attrs_idx = 0;
-    size_t fail_idx = tasks.size();
-    err = cudaMemcpyBatchAsync(dsts.data(), srcs.data(), sizes.data(),
-                               srcs.size(), &attr, &attrs_idx, 1, &fail_idx,
-                               batch->async_stream.get());
-    if (err != cudaSuccess && fail_idx < tasks.size()) {
-        LOG(ERROR) << "MnnvlTransport::startTransfer internal error: "
-                   << "cudaMemcpyBatchAsync failed at task index " << fail_idx
-                   << " (src=" << srcs[fail_idx] << ", dst=" << dsts[fail_idx]
-                   << ", size=" << sizes[fail_idx]
-                   << "): " << cudaGetErrorString(err);
-        tasks[fail_idx]->status_word = TransferStatusEnum::FAILED;
+    // Determine direction and addresses
+    if (task->request.opcode == Request::READ) {
+        dst = task->request.source;       // read into source buffer
+        src = (void *)task->target_addr;  // from remote
+    } else {
+        src = task->request.source;       // write from source buffer
+        dst = (void *)task->target_addr;  // to remote
     }
-#else
-    err = cudaSuccess;
-    for (size_t i = 0; i < tasks.size(); ++i) {
-        auto single_err =
-            cudaMemcpyAsync(dsts[i], srcs[i], sizes[i], cudaMemcpyDefault,
-                            batch->async_stream.get());
-        if (single_err != cudaSuccess) {
-            tasks[i]->status_word = TransferStatusEnum::FAILED;
-            err = single_err;
-        }
-    }
-#endif
 
-    if (err != cudaSuccess) {
-        for (auto *task : tasks) {
-            if (task->status_word == TransferStatusEnum::PENDING)
-                task->status_word = TransferStatusEnum::FAILED;
+    bool is_async = (task->request.length >= async_memcpy_threshold_);
+
+    cudaPointerAttributes src_attr_info, dst_attr_info;
+    cudaMemoryType src_type = cudaMemoryTypeHost, dst_type = cudaMemoryTypeHost;
+    if (cudaPointerGetAttributes(&src_attr_info, src) == cudaSuccess) {
+        src_type = src_attr_info.type;
+    }
+    if (cudaPointerGetAttributes(&dst_attr_info, dst) == cudaSuccess) {
+        dst_type = dst_attr_info.type;
+    }
+
+    cudaMemcpyKind kind = cudaMemcpyDefault;
+    if (src_type == cudaMemoryTypeDevice && dst_type == cudaMemoryTypeHost) {
+        kind = cudaMemcpyDeviceToHost;
+    } else if (src_type == cudaMemoryTypeHost &&
+               dst_type == cudaMemoryTypeDevice) {
+        kind = cudaMemcpyHostToDevice;
+    } else if (src_type == cudaMemoryTypeDevice &&
+               dst_type == cudaMemoryTypeDevice) {
+        kind = cudaMemcpyDeviceToDevice;
+    } else if (src_type == cudaMemoryTypeHost &&
+               dst_type == cudaMemoryTypeHost) {
+        kind = cudaMemcpyHostToHost;
+    }
+
+    if (!is_async) {
+        err = cudaMemcpyAsync(dst, src, task->request.length, kind,
+                              batch->sync_stream.get());
+        if (err != cudaSuccess) {
+            task->status_word = TransferStatusEnum::FAILED;
+            return;
         }
+
+        err = cudaStreamSynchronize(batch->sync_stream.get());
+        if (err != cudaSuccess) {
+            task->status_word = TransferStatusEnum::FAILED;
+            return;
+        }
+
+        task->transferred_bytes = task->request.length;
+        task->status_word = TransferStatusEnum::COMPLETED;
         return;
     }
 
-    cudaEvent_t event;
-    cudaEventCreateWithFlags(&event, cudaEventDisableTiming);
-    cudaEventRecord(event, batch->async_stream.get());
-    for (auto *task : tasks) task->completion_event = event;
+    err = cudaMemcpyAsync(dst, src, task->request.length, kind,
+                          batch->async_stream.get());
+
+    if (err != cudaSuccess) task->status_word = TransferStatusEnum::FAILED;
 }
 
 Status MnnvlTransport::getTransferStatus(SubBatchRef batch, int task_id,
@@ -274,8 +261,9 @@ Status MnnvlTransport::getTransferStatus(SubBatchRef batch, int task_id,
     auto &task = mnnvl_batch->task_list[task_id];
     status = TransferStatus{task.status_word, task.transferred_bytes};
     if (task.status_word == TransferStatusEnum::PENDING) {
-        auto err = cudaEventQuery(task.completion_event);
+        auto err = cudaStreamQuery(mnnvl_batch->async_stream.get());
         if (err == cudaSuccess) {
+            cudaStreamSynchronize(mnnvl_batch->async_stream.get());
             task.transferred_bytes = task.request.length;
             task.status_word = TransferStatusEnum::COMPLETED;
         } else if (err != cudaErrorNotReady) {

mooncake-transfer-engine/tent/src/transport/nvlink/nvlink_transport.cpp

@@ -106,7 +106,6 @@ Status NVLinkTransport::submitTransferTasks(
         return Status::InvalidArgument("Invalid NVLink sub-batch" LOC_MARK);
     if (request_list.size() + shm_batch->task_list.size() > shm_batch->max_size)
         return Status::TooManyRequests("Exceed batch capacity" LOC_MARK);
-    std::vector<NVLinkTask*> new_tasks;
     for (auto& request : request_list) {
         shm_batch->task_list.push_back(NVLinkTask{});
         auto& task = shm_batch->task_list[shm_batch->task_list.size() - 1];
@@ -119,83 +118,72 @@ Status NVLinkTransport::submitTransferTasks(
         task.target_addr = target_addr;
         task.request = request;
         task.status_word = TransferStatusEnum::PENDING;
-        new_tasks.push_back(&task);
+        startTransfer(&task, shm_batch);
     }
-    startTransfer(new_tasks, shm_batch);
     return Status::OK();
 }
 
-void NVLinkTransport::startTransfer(std::vector<NVLinkTask*>& tasks,
-                                    NVLinkSubBatch* batch) {
-    if (tasks.empty()) return;
-
-    std::vector<void*> srcs;
-    std::vector<void*> dsts;
-    std::vector<size_t> sizes;
-    for (auto* task : tasks) {
-        void *src = nullptr, *dst = nullptr;
-        if (task->request.opcode == Request::READ) {
-            dst = task->request.source;      // read into source buffer
-            src = (void*)task->target_addr;  // from remote
-        } else {
-            src = task->request.source;      // write from source buffer
-            dst = (void*)task->target_addr;  // to remote
-        }
-        srcs.push_back(src);
-        dsts.push_back(dst);
-        sizes.push_back(task->request.length);
-    }
-
+void NVLinkTransport::startTransfer(NVLinkTask* task, NVLinkSubBatch* batch) {
     cudaError_t err;
+    void *src = nullptr, *dst = nullptr;
 
-#if CUDART_VERSION >= 13000
-    cudaMemcpyAttributes attr{};
-    attr.srcAccessOrder = cudaMemcpySrcAccessOrderStream;
-    size_t attrs_idx = 0;
-    err = cudaMemcpyBatchAsync(const_cast<const void**>(dsts.data()),
-                               const_cast<const void**>(srcs.data()),
-                               sizes.data(), srcs.size(), &attr, &attrs_idx, 1,
-                               batch->async_stream.get());
-#elif CUDART_VERSION >= 12080
-    cudaMemcpyAttributes attr{};
-    attr.srcAccessOrder = cudaMemcpySrcAccessOrderStream;
-    size_t attrs_idx = 0;
-    size_t fail_idx = tasks.size();
-    err = cudaMemcpyBatchAsync(dsts.data(), srcs.data(), sizes.data(),
-                               srcs.size(), &attr, &attrs_idx, 1, &fail_idx,
-                               batch->async_stream.get());
-    if (err != cudaSuccess && fail_idx < tasks.size()) {
-        LOG(ERROR) << "NVLinkTransport::startTransfer internal error: "
-                   << "cudaMemcpyBatchAsync failed at task index " << fail_idx
-                   << " (src=" << srcs[fail_idx] << ", dst=" << dsts[fail_idx]
-                   << ", size=" << sizes[fail_idx]
-                   << "): " << cudaGetErrorString(err);
-        tasks[fail_idx]->status_word = TransferStatusEnum::FAILED;
+    // Determine direction and addresses
+    if (task->request.opcode == Request::READ) {
+        dst = task->request.source;      // read into source buffer
+        src = (void*)task->target_addr;  // from remote
+    } else {
+        src = task->request.source;      // write from source buffer
+        dst = (void*)task->target_addr;  // to remote
     }
-#else
-    err = cudaSuccess;
-    for (size_t i = 0; i < tasks.size(); ++i) {
-        auto single_err =
-            cudaMemcpyAsync(dsts[i], srcs[i], sizes[i], cudaMemcpyDefault,
-                            batch->async_stream.get());
-        if (single_err != cudaSuccess) {
-            tasks[i]->status_word = TransferStatusEnum::FAILED;
-            err = single_err;
+
+    bool is_async = (task->request.length >= async_memcpy_threshold_);
+
+    cudaPointerAttributes src_attr_info, dst_attr_info;
+    cudaMemoryType src_type = cudaMemoryTypeHost, dst_type = cudaMemoryTypeHost;
+    if (cudaPointerGetAttributes(&src_attr_info, src) == cudaSuccess) {
+        src_type = src_attr_info.type;
+    }
+    if (cudaPointerGetAttributes(&dst_attr_info, dst) == cudaSuccess) {
+        dst_type = dst_attr_info.type;
+    }
+
+    cudaMemcpyKind kind = cudaMemcpyDefault;
+    if (src_type == cudaMemoryTypeDevice && dst_type == cudaMemoryTypeHost) {
+        kind = cudaMemcpyDeviceToHost;
+    } else if (src_type == cudaMemoryTypeHost &&
+               dst_type == cudaMemoryTypeDevice) {
+        kind = cudaMemcpyHostToDevice;
+    } else if (src_type == cudaMemoryTypeDevice &&
+               dst_type == cudaMemoryTypeDevice) {
+        kind = cudaMemcpyDeviceToDevice;
+    } else if (src_type == cudaMemoryTypeHost &&
+               dst_type == cudaMemoryTypeHost) {
+        kind = cudaMemcpyHostToHost;
+    }
+
+    if (!is_async) {
+        err = cudaMemcpyAsync(dst, src, task->request.length, kind,
+                              batch->sync_stream.get());
+        if (err != cudaSuccess) {
+            task->status_word = TransferStatusEnum::FAILED;
+            return;
         }
-    }
-#endif
 
-    if (err != cudaSuccess) {
-        for (auto* task : tasks) {
-            if (task->status_word == TransferStatusEnum::PENDING)
-                task->status_word = TransferStatusEnum::FAILED;
+        err = cudaStreamSynchronize(batch->sync_stream.get());
+        if (err != cudaSuccess) {
+            task->status_word = TransferStatusEnum::FAILED;
+            return;
         }
+
+        task->transferred_bytes = task->request.length;
+        task->status_word = TransferStatusEnum::COMPLETED;
+        return;
     }
 
-    cudaEvent_t event;
-    cudaEventCreateWithFlags(&event, cudaEventDisableTiming);
-    cudaEventRecord(event, batch->async_stream.get());
-    for (auto* task : tasks) task->completion_event = event;
+    err = cudaMemcpyAsync(dst, src, task->request.length, kind,
+                          batch->async_stream.get());
+
+    if (err != cudaSuccess) task->status_word = TransferStatusEnum::FAILED;
 }
 
 Status NVLinkTransport::getTransferStatus(SubBatchRef batch, int task_id,
@@ -207,8 +195,9 @@ Status NVLinkTransport::getTransferStatus(SubBatchRef batch, int task_id,
     auto& task = shm_batch->task_list[task_id];
     status = TransferStatus{task.status_word, task.transferred_bytes};
     if (task.status_word == TransferStatusEnum::PENDING) {
-        auto err = cudaEventQuery(task.completion_event);
+        auto err = cudaStreamQuery(shm_batch->async_stream.get());
         if (err == cudaSuccess) {
+            cudaStreamSynchronize(shm_batch->async_stream.get());
             task.transferred_bytes = task.request.length;
             task.status_word = TransferStatusEnum::COMPLETED;
         } else if (err != cudaErrorNotReady) {

mooncake-transfer-engine/tent/tests/CMakeLists.txt

@@ -57,10 +57,3 @@ add_test(NAME tent_failover_test COMMAND tent_failover_test)
 add_executable(tent_endpoint_lifecycle_test endpoint_lifecycle_test.cpp)
 target_link_libraries(tent_endpoint_lifecycle_test PRIVATE gtest gtest_main)
 add_test(NAME tent_endpoint_lifecycle_test COMMAND tent_endpoint_lifecycle_test)
-
-add_executable(tent_fault_proxy_test fault_proxy_test.cpp)
-target_link_libraries(tent_fault_proxy_test PRIVATE gtest gtest_main
-                                                     tent_link_group)
-target_include_directories(tent_fault_proxy_test
-                           PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include)
-add_test(NAME tent_fault_proxy_test COMMAND tent_fault_proxy_test)

mooncake-transfer-engine/tent/tests/fault_proxy_test.cpp

@@ -1,365 +0,0 @@
-// 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.
-
-#include <gtest/gtest.h>
-
-#include <memory>
-#include <string>
-
-#include "tent/common/types.h"
-#include "tent/runtime/transfer_engine_impl.h"
-#include "tent/transport/fault_proxy/fault_proxy_transport.h"
-
-namespace mooncake {
-namespace tent {
-namespace {
-
-// ---------------------------------------------------------------------------
-// Minimal FakeTransport: always succeeds, tracks call counts.
-// ---------------------------------------------------------------------------
-
-class FakeSubBatch : public Transport::SubBatch {
-   public:
-    size_t size() const override { return task_count; }
-    size_t task_count = 0;
-    // Store status per task for getTransferStatus
-    std::vector<TransferStatus> statuses;
-};
-
-class FakeTransport : public Transport {
-   public:
-    int install_calls = 0;
-    int submit_calls = 0;
-    int status_calls = 0;
-
-    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 {
-        ++install_calls;
-        return Status::OK();
-    }
-
-    Status allocateSubBatch(SubBatchRef& batch, size_t /*max_size*/) override {
-        batch = new FakeSubBatch();
-        return Status::OK();
-    }
-
-    Status freeSubBatch(SubBatchRef& batch) override {
-        delete batch;
-        batch = nullptr;
-        return Status::OK();
-    }
-
-    Status submitTransferTasks(
-        SubBatchRef batch, const std::vector<Request>& request_list) override {
-        ++submit_calls;
-        auto* fb = static_cast<FakeSubBatch*>(batch);
-        for (size_t i = 0; i < request_list.size(); ++i) {
-            fb->statuses.push_back(
-                {TransferStatusEnum::COMPLETED, request_list[i].length});
-            fb->task_count++;
-        }
-        return Status::OK();
-    }
-
-    Status getTransferStatus(SubBatchRef batch, int task_id,
-                             TransferStatus& status) override {
-        ++status_calls;
-        auto* fb = static_cast<FakeSubBatch*>(batch);
-        if (task_id < 0 || task_id >= (int)fb->statuses.size()) {
-            return Status::InvalidArgument("bad task_id" LOC_MARK);
-        }
-        status = fb->statuses[task_id];
-        return Status::OK();
-    }
-
-    const char* getName() const override { return "<fake>"; }
-};
-
-// ---------------------------------------------------------------------------
-// Helper: create a Request for testing (no real memory needed)
-// ---------------------------------------------------------------------------
-
-static Request makeRequest(size_t length = 4096) {
-    Request req;
-    req.opcode = Request::WRITE;
-    req.source = nullptr;
-    req.target_id = 1;
-    req.target_offset = 0;
-    req.length = length;
-    return req;
-}
-
-// ===========================================================================
-// Group 1: Unit — FaultPolicy and FaultProxyTransport basics
-// ===========================================================================
-
-TEST(FaultPolicyTest, Defaults) {
-    FaultPolicy policy;
-    EXPECT_DOUBLE_EQ(policy.submit_fail_rate, 0.0);
-    EXPECT_DOUBLE_EQ(policy.status_corrupt_rate, 0.0);
-    EXPECT_EQ(policy.submit_delay_us, 0u);
-    EXPECT_EQ(policy.fail_after_n_submits, -1);
-    EXPECT_FALSE(policy.fail_install);
-}
-
-TEST(FaultProxyTest, DelegatesToReal) {
-    auto fake = std::make_shared<FakeTransport>();
-    FaultPolicy policy;  // all defaults — no faults
-    auto proxy = std::make_shared<FaultProxyTransport>(fake, policy);
-
-    // install
-    std::string seg_name = "test";
-    ASSERT_TRUE(proxy->install(seg_name, nullptr, nullptr).ok());
-    EXPECT_EQ(fake->install_calls, 1);
-
-    // allocateSubBatch + submitTransferTasks
-    Transport::SubBatchRef batch = nullptr;
-    ASSERT_TRUE(proxy->allocateSubBatch(batch, 16).ok());
-    ASSERT_NE(batch, nullptr);
-
-    auto req = makeRequest();
-    ASSERT_TRUE(proxy->submitTransferTasks(batch, {req}).ok());
-    EXPECT_EQ(fake->submit_calls, 1);
-    EXPECT_EQ(proxy->submitCount(), 1);
-
-    // getTransferStatus
-    TransferStatus ts{};
-    ASSERT_TRUE(proxy->getTransferStatus(batch, 0, ts).ok());
-    EXPECT_EQ(ts.s, TransferStatusEnum::COMPLETED);
-    EXPECT_EQ(ts.transferred_bytes, 4096u);
-    EXPECT_EQ(fake->status_calls, 1);
-
-    proxy->freeSubBatch(batch);
-}
-
-TEST(FaultProxyTest, FailsInstall) {
-    auto fake = std::make_shared<FakeTransport>();
-    FaultPolicy policy;
-    policy.fail_install = true;
-    auto proxy = std::make_shared<FaultProxyTransport>(fake, policy);
-
-    std::string seg_name = "test";
-    auto status = proxy->install(seg_name, nullptr, nullptr);
-    EXPECT_FALSE(status.ok());
-    EXPECT_EQ(fake->install_calls, 0);  // never reached the real transport
-}
-
-TEST(FaultProxyTest, FailsSubmitDeterministic) {
-    auto fake = std::make_shared<FakeTransport>();
-    FaultPolicy policy;
-    policy.fail_after_n_submits = 2;  // succeed twice, then fail
-    auto proxy = std::make_shared<FaultProxyTransport>(fake, policy);
-
-    Transport::SubBatchRef batch = nullptr;
-    ASSERT_TRUE(proxy->allocateSubBatch(batch, 16).ok());
-    auto req = makeRequest();
-
-    // First two submits succeed
-    EXPECT_TRUE(proxy->submitTransferTasks(batch, {req}).ok());
-    EXPECT_TRUE(proxy->submitTransferTasks(batch, {req}).ok());
-    EXPECT_EQ(fake->submit_calls, 2);
-
-    // Third submit fails
-    auto s = proxy->submitTransferTasks(batch, {req});
-    EXPECT_FALSE(s.ok());
-    EXPECT_EQ(fake->submit_calls, 2);  // real transport not called
-
-    // Fourth also fails
-    EXPECT_FALSE(proxy->submitTransferTasks(batch, {req}).ok());
-
-    proxy->freeSubBatch(batch);
-}
-
-TEST(FaultProxyTest, CorruptsStatus) {
-    auto fake = std::make_shared<FakeTransport>();
-    FaultPolicy policy;
-    policy.status_corrupt_rate = 1.0;  // always corrupt
-    auto proxy = std::make_shared<FaultProxyTransport>(fake, policy);
-
-    Transport::SubBatchRef batch = nullptr;
-    ASSERT_TRUE(proxy->allocateSubBatch(batch, 16).ok());
-    auto req = makeRequest();
-    ASSERT_TRUE(proxy->submitTransferTasks(batch, {req}).ok());
-
-    TransferStatus ts{};
-    ASSERT_TRUE(proxy->getTransferStatus(batch, 0, ts).ok());
-    // Real transport returned COMPLETED, but proxy flipped it to FAILED
-    EXPECT_EQ(ts.s, TransferStatusEnum::FAILED);
-
-    proxy->freeSubBatch(batch);
-}
-
-// ===========================================================================
-// Group 2: Integration — Failover state machine driven by proxy
-// ===========================================================================
-
-TEST(FaultProxyFailoverTest, FailoverFromProxyToReal) {
-    // Simulate: RDMA (proxied, always fails) → TCP (real, always succeeds)
-    auto fake_rdma = std::make_shared<FakeTransport>();
-    auto fake_tcp = std::make_shared<FakeTransport>();
-
-    FaultPolicy rdma_policy;
-    rdma_policy.submit_fail_rate = 1.0;  // RDMA always fails
-    auto proxied_rdma =
-        std::make_shared<FaultProxyTransport>(fake_rdma, rdma_policy);
-
-    // Allocate sub-batches for both transports
-    Transport::SubBatchRef rdma_batch = nullptr;
-    Transport::SubBatchRef tcp_batch = nullptr;
-    ASSERT_TRUE(proxied_rdma->allocateSubBatch(rdma_batch, 16).ok());
-    ASSERT_TRUE(fake_tcp->allocateSubBatch(tcp_batch, 16).ok());
-
-    // Simulate TransferEngineImpl::submitTransfer + resubmitTransferTask
-    constexpr int kMaxAttempts = 3;
-    TaskInfo task;
-    task.type = RDMA;
-    task.xport_priority = 0;
-    task.status = TransferStatusEnum::PENDING;
-    task.failover_count = 0;
-
-    auto req = makeRequest();
-
-    // Step 1: Submit on "RDMA" — should fail (proxy injects fault)
-    auto s = proxied_rdma->submitTransferTasks(rdma_batch, {req});
-    EXPECT_FALSE(s.ok());
-
-    // Step 2: Failover logic (mirrors resubmitTransferTask)
-    task.status = TransferStatusEnum::FAILED;
-    ++task.failover_count;
-    EXPECT_LE(task.failover_count, kMaxAttempts);
-    task.xport_priority++;
-    task.type = TCP;
-    task.status = TransferStatusEnum::PENDING;
-
-    // Step 3: Submit on "TCP" — should succeed
-    s = fake_tcp->submitTransferTasks(tcp_batch, {req});
-    EXPECT_TRUE(s.ok());
-
-    // Step 4: Verify completion
-    TransferStatus ts{};
-    ASSERT_TRUE(fake_tcp->getTransferStatus(tcp_batch, 0, ts).ok());
-    EXPECT_EQ(ts.s, TransferStatusEnum::COMPLETED);
-    task.status = TransferStatusEnum::COMPLETED;
-
-    // Assertions
-    EXPECT_EQ(task.type, TCP);
-    EXPECT_EQ(task.failover_count, 1);
-    EXPECT_EQ(task.status, TransferStatusEnum::COMPLETED);
-    EXPECT_EQ(fake_rdma->submit_calls, 0);  // proxy intercepted, never hit real
-    EXPECT_EQ(fake_tcp->submit_calls, 1);
-
-    proxied_rdma->freeSubBatch(rdma_batch);
-    fake_tcp->freeSubBatch(tcp_batch);
-}
-
-TEST(FaultProxyFailoverTest, ExhaustAllTransports) {
-    // Both RDMA and TCP are proxied with 100% failure
-    auto fake_rdma = std::make_shared<FakeTransport>();
-    auto fake_tcp = std::make_shared<FakeTransport>();
-
-    FaultPolicy always_fail;
-    always_fail.submit_fail_rate = 1.0;
-
-    auto proxy_rdma =
-        std::make_shared<FaultProxyTransport>(fake_rdma, always_fail);
-    auto proxy_tcp =
-        std::make_shared<FaultProxyTransport>(fake_tcp, always_fail);
-
-    Transport::SubBatchRef rdma_batch = nullptr;
-    Transport::SubBatchRef tcp_batch = nullptr;
-    ASSERT_TRUE(proxy_rdma->allocateSubBatch(rdma_batch, 16).ok());
-    ASSERT_TRUE(proxy_tcp->allocateSubBatch(tcp_batch, 16).ok());
-
-    constexpr int kMaxAttempts = 3;
-    TaskInfo task;
-    task.type = RDMA;
-    task.xport_priority = 0;
-    task.status = TransferStatusEnum::PENDING;
-    task.failover_count = 0;
-
-    auto req = makeRequest();
-
-    // Attempt on RDMA — fails
-    EXPECT_FALSE(proxy_rdma->submitTransferTasks(rdma_batch, {req}).ok());
-    task.status = TransferStatusEnum::FAILED;
-    ++task.failover_count;
-    task.xport_priority++;
-    task.type = TCP;
-    task.status = TransferStatusEnum::PENDING;
-
-    // Attempt on TCP — also fails
-    EXPECT_FALSE(proxy_tcp->submitTransferTasks(tcp_batch, {req}).ok());
-    task.status = TransferStatusEnum::FAILED;
-    ++task.failover_count;
-    task.xport_priority++;
-    task.status = TransferStatusEnum::PENDING;
-
-    // Third attempt — fails again
-    EXPECT_FALSE(proxy_rdma->submitTransferTasks(rdma_batch, {req}).ok());
-    task.status = TransferStatusEnum::FAILED;
-    ++task.failover_count;
-    EXPECT_LE(task.failover_count, kMaxAttempts);  // Still within limit
-
-    task.xport_priority++;
-    task.status = TransferStatusEnum::PENDING;
-
-    // Fourth attempt — exceeds limit
-    EXPECT_FALSE(proxy_tcp->submitTransferTasks(tcp_batch, {req}).ok());
-    task.status = TransferStatusEnum::FAILED;
-    ++task.failover_count;
-    EXPECT_GT(task.failover_count, kMaxAttempts);
-
-    // Task stays FAILED — no more failover
-    EXPECT_EQ(task.status, TransferStatusEnum::FAILED);
-    EXPECT_EQ(task.failover_count, kMaxAttempts + 1);
-
-    proxy_rdma->freeSubBatch(rdma_batch);
-    proxy_tcp->freeSubBatch(tcp_batch);
-}
-
-// ===========================================================================
-// Group 3: Policy mutation
-// ===========================================================================
-
-TEST(FaultProxyTest, ResetPolicyMidRun) {
-    auto fake = std::make_shared<FakeTransport>();
-    FaultPolicy fail_policy;
-    fail_policy.submit_fail_rate = 1.0;
-    auto proxy = std::make_shared<FaultProxyTransport>(fake, fail_policy);
-
-    Transport::SubBatchRef batch = nullptr;
-    ASSERT_TRUE(proxy->allocateSubBatch(batch, 16).ok());
-    auto req = makeRequest();
-
-    // Should fail
-    EXPECT_FALSE(proxy->submitTransferTasks(batch, {req}).ok());
-    EXPECT_EQ(fake->submit_calls, 0);
-
-    // Reset to clean policy
-    FaultPolicy clean_policy;
-    proxy->resetPolicy(clean_policy);
-
-    // Should succeed now
-    EXPECT_TRUE(proxy->submitTransferTasks(batch, {req}).ok());
-    EXPECT_EQ(fake->submit_calls, 1);
-    EXPECT_EQ(proxy->submitCount(), 1);  // counter was reset too
-
-    proxy->freeSubBatch(batch);
-}
-
-}  // namespace
-}  // namespace tent
-}  // namespace mooncake

mooncake-transfer-engine/tests/CMakeLists.txt

@@ -63,16 +63,6 @@ if (USE_EFA)
     add_executable(efa_transport_test ${WORKSPACE}/efa_transport_test.cpp)
     target_link_libraries(efa_transport_test PUBLIC transfer_engine gtest gtest_main)
     add_test(NAME efa_transport_test COMMAND efa_transport_test)
-
-    add_executable(efa_c_api_test ${WORKSPACE}/efa_c_api_test.cpp)
-    target_link_libraries(efa_c_api_test PUBLIC transfer_engine gtest gtest_main)
-    add_test(NAME efa_c_api_test COMMAND efa_c_api_test)
-
-    add_executable(efa_single_nic_large_mr_test ${WORKSPACE}/efa_single_nic_large_mr_test.cpp)
-    target_link_libraries(efa_single_nic_large_mr_test PUBLIC transfer_engine gflags::gflags glog::glog)
-
-    add_executable(efa_transfer_test ${WORKSPACE}/efa_transfer_test.cpp)
-    target_link_libraries(efa_transfer_test PUBLIC transfer_engine gflags::gflags glog::glog)
 endif()
 
 # UB transport test with URMA endpoint and mock support

mooncake-transfer-engine/tests/efa_c_api_test.cpp

@@ -1,104 +0,0 @@
-// Copyright 2024 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.
-
-#include <gtest/gtest.h>
-#include <numa.h>
-
-#include <cstring>
-
-#include "transfer_engine_c.h"
-
-// Test: discoverTopology + installTransport("efa") via pure C API
-TEST(EfaCApiTest, DiscoverTopologyAndInstallEfa) {
-    transfer_engine_t engine = createTransferEngine(
-        "P2PHANDSHAKE", "127.0.0.1:12345", "127.0.0.1", 12345, 0);
-    ASSERT_NE(engine, nullptr) << "createTransferEngine failed";
-
-    int ret = discoverTopology(engine);
-    ASSERT_EQ(ret, 0) << "discoverTopology failed with code " << ret;
-
-    transport_t xport = installTransport(engine, "efa", NULL);
-    ASSERT_NE(xport, nullptr) << "installTransport(\"efa\") failed";
-
-    destroyTransferEngine(engine);
-}
-
-// Test: registerLocalMemory works after C API EFA setup
-TEST(EfaCApiTest, RegisterMemoryAfterCApiSetup) {
-    transfer_engine_t engine = createTransferEngine(
-        "P2PHANDSHAKE", "127.0.0.1:22345", "127.0.0.1", 22345, 0);
-    ASSERT_NE(engine, nullptr);
-
-    ASSERT_EQ(discoverTopology(engine), 0);
-    ASSERT_NE(installTransport(engine, "efa", NULL), nullptr);
-
-    size_t buf_size = 1 << 20;  // 1 MB
-    void* buf = numa_alloc_onnode(buf_size, 0);
-    ASSERT_NE(buf, nullptr);
-    memset(buf, 0xAB, buf_size);
-
-    int ret = registerLocalMemory(engine, buf, buf_size, "cpu:0", 1);
-    EXPECT_EQ(ret, 0) << "registerLocalMemory failed with code " << ret;
-
-    unregisterLocalMemory(engine, buf);
-    numa_free(buf, buf_size);
-    destroyTransferEngine(engine);
-}
-
-// Test: registerLocalMemoryBatch works after C API EFA setup
-TEST(EfaCApiTest, RegisterMemoryBatchAfterCApiSetup) {
-    transfer_engine_t engine = createTransferEngine(
-        "P2PHANDSHAKE", "127.0.0.1:32345", "127.0.0.1", 32345, 0);
-    ASSERT_NE(engine, nullptr);
-
-    ASSERT_EQ(discoverTopology(engine), 0);
-    ASSERT_NE(installTransport(engine, "efa", NULL), nullptr);
-
-    const int num_bufs = 4;
-    size_t buf_size = 1 << 20;
-    buffer_entry_t entries[4];
-    for (int i = 0; i < num_bufs; i++) {
-        entries[i].addr = numa_alloc_onnode(buf_size, 0);
-        ASSERT_NE(entries[i].addr, nullptr);
-        entries[i].length = buf_size;
-        memset(entries[i].addr, 0xCD, buf_size);
-    }
-
-    int ret = registerLocalMemoryBatch(engine, entries, num_bufs, "cpu:0");
-    EXPECT_EQ(ret, 0) << "registerLocalMemoryBatch failed with code " << ret;
-
-    void* addrs[4];
-    for (int i = 0; i < num_bufs; i++) addrs[i] = entries[i].addr;
-    unregisterLocalMemoryBatch(engine, addrs, num_bufs);
-    for (int i = 0; i < num_bufs; i++) numa_free(entries[i].addr, buf_size);
-    destroyTransferEngine(engine);
-}
-
-// Negative test: without discoverTopology, installTransport should fail
-TEST(EfaCApiTest, InstallEfaWithoutDiscoverFails) {
-    transfer_engine_t engine = createTransferEngine(
-        "P2PHANDSHAKE", "127.0.0.1:42345", "127.0.0.1", 42345, 0);
-    ASSERT_NE(engine, nullptr);
-
-    transport_t xport = installTransport(engine, "efa", NULL);
-    EXPECT_EQ(xport, nullptr)
-        << "installTransport should fail without discover";
-
-    destroyTransferEngine(engine);
-}
-
-int main(int argc, char** argv) {
-    ::testing::InitGoogleTest(&argc, argv);
-    return RUN_ALL_TESTS();
-}

mooncake-transfer-engine/tests/efa_single_nic_large_mr_test.cpp

@@ -1,178 +0,0 @@
-// Copyright 2024 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.
-
-// Reproduce: single-NIC + 200GB 4KB-page buffer registration.
-// Expected: auto-split into chunks and register all on the one NIC.
-// Current bug: "Buffer requires N chunks but only 1 NICs available."
-//
-// Usage:
-//   ./efa_single_nic_large_mr_test [--nic rdmap85s0] [--size_gb 200]
-//
-// The buffer is allocated via mmap (no MAP_HUGETLB) to get 4KB pages,
-// matching the customer's /dev/shm mmap pattern.
-
-#include <gflags/gflags.h>
-#include <glog/logging.h>
-#include <sys/mman.h>
-
-#include <chrono>
-#include <cstdint>
-#include <cstdio>
-#include <cstring>
-#include <string>
-#include <vector>
-
-#include "transfer_engine_c.h"
-
-DEFINE_string(nic, "",
-              "EFA device name (e.g. rdmap85s0). "
-              "If empty, auto-detects the first rdmap* device.");
-DEFINE_double(size_gb, 200.0, "Total memory in GB to register");
-DEFINE_double(chunk_gb, 0, "Per-buffer chunk size in GB (0 = single buffer)");
-DEFINE_string(server, "127.0.0.1:12345", "Local server name");
-
-static std::string detectFirstEfaDevice() {
-    FILE* fp = popen(
-        "ls /sys/class/infiniband/ 2>/dev/null | grep rdmap | head -1", "r");
-    if (!fp) return "";
-    char buf[256] = {};
-    if (fgets(buf, sizeof(buf), fp)) {
-        size_t len = strlen(buf);
-        if (len > 0 && buf[len - 1] == '\n') buf[len - 1] = '\0';
-    }
-    pclose(fp);
-    return std::string(buf);
-}
-
-int main(int argc, char** argv) {
-    google::InitGoogleLogging(argv[0]);
-    gflags::ParseCommandLineFlags(&argc, &argv, true);
-    FLAGS_logtostderr = 1;
-
-    std::string nic = FLAGS_nic;
-
-    size_t size_bytes = static_cast<size_t>(FLAGS_size_gb * 1024 * 1024 * 1024);
-    LOG(INFO) << "Buffer size: " << FLAGS_size_gb << " GB (" << size_bytes
-              << " bytes)";
-
-    // Create engine
-    transfer_engine_t engine = createTransferEngine(
-        "P2PHANDSHAKE", FLAGS_server.c_str(), "127.0.0.1", 12345, 0);
-    if (!engine) {
-        LOG(ERROR) << "createTransferEngine failed";
-        return 1;
-    }
-
-    // discoverTopology to populate device list
-    int ret = discoverTopology(engine);
-    if (ret != 0) {
-        LOG(ERROR) << "discoverTopology failed: " << ret;
-        destroyTransferEngine(engine);
-        return 1;
-    }
-
-    // Install EFA transport — optionally restrict to single NIC
-    transport_t xport = nullptr;
-    if (!nic.empty()) {
-        std::string matrix = "{\"cpu:0\": [[\"" + nic + "\"], []]}";
-        LOG(INFO) << "nic_priority_matrix: " << matrix;
-        char* matrix_cstr = const_cast<char*>(matrix.c_str());
-        void* args[] = {matrix_cstr};
-        xport = installTransport(engine, "efa", args);
-    } else {
-        LOG(INFO) << "Using all available NICs (no nic restriction)";
-        xport = installTransport(engine, "efa", nullptr);
-    }
-    if (!xport) {
-        LOG(ERROR) << "installTransport(efa) failed";
-        destroyTransferEngine(engine);
-        return 1;
-    }
-
-    // Determine buffer count and per-buffer size
-    size_t chunk_bytes = 0;
-    int num_bufs = 1;
-    if (FLAGS_chunk_gb > 0) {
-        chunk_bytes = static_cast<size_t>(FLAGS_chunk_gb * 1024 * 1024 * 1024);
-        num_bufs = (size_bytes + chunk_bytes - 1) / chunk_bytes;
-    } else {
-        chunk_bytes = size_bytes;
-    }
-    LOG(INFO) << "Plan: " << num_bufs << " buffers × "
-              << chunk_bytes / (1024 * 1024)
-              << " MB = " << (num_bufs * chunk_bytes) / (1024ULL * 1024 * 1024)
-              << " GB";
-
-    // Allocate all buffers with hugepages
-    std::vector<void*> bufs;
-    bufs.reserve(num_bufs);
-    LOG(INFO) << "Allocating " << num_bufs << " buffers...";
-    for (int i = 0; i < num_bufs; ++i) {
-        void* buf = mmap(nullptr, chunk_bytes, PROT_READ | PROT_WRITE,
-                         MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB, -1, 0);
-        if (buf == MAP_FAILED) {
-            LOG(ERROR) << "mmap failed at buffer " << i << "/" << num_bufs
-                       << ": " << strerror(errno);
-            for (auto* p : bufs) munmap(p, chunk_bytes);
-            destroyTransferEngine(engine);
-            return 1;
-        }
-        bufs.push_back(buf);
-        if ((i + 1) % 50 == 0 || i == num_bufs - 1) {
-            LOG(INFO) << "  allocated " << (i + 1) << "/" << num_bufs;
-        }
-    }
-    LOG(INFO) << "All buffers allocated";
-
-    // Register each buffer
-    LOG(INFO) << "Registering " << num_bufs << " × "
-              << chunk_bytes / (1024 * 1024) << " MB on single NIC...";
-    auto t0 = std::chrono::steady_clock::now();
-    int failures = 0;
-    for (int i = 0; i < num_bufs; ++i) {
-        ret = registerLocalMemory(engine, bufs[i], chunk_bytes, "cpu:0", 1);
-        if (ret != 0) {
-            LOG(ERROR) << "FAILED at buffer " << i << "/" << num_bufs
-                       << " (addr=" << bufs[i] << "): ret=" << ret;
-            failures++;
-            break;
-        }
-        if ((i + 1) % 50 == 0 || i == num_bufs - 1) {
-            auto now = std::chrono::steady_clock::now();
-            double elapsed = std::chrono::duration<double>(now - t0).count();
-            LOG(INFO) << "  registered " << (i + 1) << "/" << num_bufs << " ("
-                      << elapsed << "s)";
-        }
-    }
-    auto t1 = std::chrono::steady_clock::now();
-    double total_time = std::chrono::duration<double>(t1 - t0).count();
-
-    if (failures == 0) {
-        LOG(INFO) << "SUCCESS: registered " << num_bufs << " × "
-                  << chunk_bytes / (1024 * 1024) << " MB ("
-                  << (num_bufs * chunk_bytes) / (1024ULL * 1024 * 1024)
-                  << " GB) in " << total_time << "s";
-    } else {
-        LOG(ERROR) << "FAILED after " << total_time << "s";
-    }
-
-    // Cleanup
-    for (auto* p : bufs) {
-        unregisterLocalMemory(engine, p);
-        munmap(p, chunk_bytes);
-    }
-    destroyTransferEngine(engine);
-
-    return failures == 0 ? 0 : 1;
-}

mooncake-transfer-engine/tests/efa_transfer_test.cpp

@@ -1,446 +0,0 @@
-// Copyright 2024 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.
-
-// EFA multi-NIC transfer test: register 200x2GB on all NICs, then transfer.
-//
-// Target node: allocate 200x2GB buffers, register on all NICs, wait.
-// Initiator node: allocate receive buffer, register, pull from target.
-//
-// Usage:
-//   # Target (holds KV cache):
-//   ./efa_transfer_test --mode target --server <target_ip>:12345 \
-//       --num_bufs 200 --buf_size_gb 2
-//
-//   # Initiator (pulls data):
-//   ./efa_transfer_test --mode initiator --server <initiator_ip>:12346 \
-//       --target <target_ip>:12345 --num_bufs 200 --buf_size_gb 2 \
-//       --transfer_mb 368
-
-#include <gflags/gflags.h>
-#include <glog/logging.h>
-#include <signal.h>
-#include <sys/mman.h>
-
-#include <algorithm>
-#include <atomic>
-#include <chrono>
-#include <cstdint>
-#include <cstring>
-#include <numeric>
-#include <string>
-#include <thread>
-#include <vector>
-
-#include "transfer_engine.h"
-
-using namespace mooncake;
-
-DEFINE_string(mode, "target", "Running mode: target or initiator");
-DEFINE_string(server, "", "Local server name, e.g. 172.31.6.162:12345");
-DEFINE_string(target, "", "Target server name (initiator mode only)");
-DEFINE_string(metadata, "P2PHANDSHAKE", "Metadata server");
-DEFINE_int32(num_bufs, 200, "Number of buffers to allocate and register");
-DEFINE_double(buf_size_gb, 2.0, "Size of each buffer in GB");
-DEFINE_double(transfer_mb, 368.0,
-              "Transfer size per iteration in MB (initiator)");
-DEFINE_int32(iterations, 50, "Number of benchmark iterations");
-DEFINE_int32(warmup, 5, "Number of warmup iterations");
-DEFINE_int32(batch_size, 1, "Batch size for each transfer submission");
-DEFINE_int32(threads, 1, "Number of initiator worker threads");
-DEFINE_uint64(block_size, 65536, "Block size for transfer requests (64KB)");
-
-static std::atomic<bool> g_running(true);
-
-static void signalHandler(int) {
-    g_running.store(false, std::memory_order_relaxed);
-}
-
-static void setupSignalHandler() {
-    struct sigaction sa;
-    sa.sa_handler = signalHandler;
-    sigemptyset(&sa.sa_mask);
-    sa.sa_flags = 0;
-    sigaction(SIGINT, &sa, nullptr);
-    sigaction(SIGTERM, &sa, nullptr);
-}
-
-static void* allocateHugepage(size_t size) {
-    void* buf = mmap(nullptr, size, PROT_READ | PROT_WRITE,
-                     MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB, -1, 0);
-    if (buf == MAP_FAILED) {
-        LOG(WARNING) << "Hugepage mmap failed (" << strerror(errno)
-                     << "), falling back to regular pages";
-        buf = mmap(nullptr, size, PROT_READ | PROT_WRITE,
-                   MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
-        if (buf == MAP_FAILED) return nullptr;
-    }
-    return buf;
-}
-
-static int runTarget(TransferEngine* engine) {
-    size_t buf_bytes =
-        static_cast<size_t>(FLAGS_buf_size_gb * 1024 * 1024 * 1024);
-    int num_bufs = FLAGS_num_bufs;
-
-    LOG(INFO) << "=== Target Node ===";
-    LOG(INFO) << "Registering " << num_bufs << " x " << FLAGS_buf_size_gb
-              << " GB = " << num_bufs * FLAGS_buf_size_gb << " GB";
-
-    // Allocate buffers
-    std::vector<void*> bufs;
-    bufs.reserve(num_bufs);
-    LOG(INFO) << "Allocating " << num_bufs << " buffers...";
-    for (int i = 0; i < num_bufs; ++i) {
-        void* buf = allocateHugepage(buf_bytes);
-        if (!buf) {
-            LOG(ERROR) << "Allocation failed at buffer " << i;
-            for (auto* p : bufs) munmap(p, buf_bytes);
-            return 1;
-        }
-        bufs.push_back(buf);
-        if ((i + 1) % 50 == 0 || i == num_bufs - 1)
-            LOG(INFO) << "  allocated " << (i + 1) << "/" << num_bufs;
-    }
-
-    // Register all buffers
-    LOG(INFO) << "Registering " << num_bufs << " buffers on all NICs...";
-    auto t0 = std::chrono::steady_clock::now();
-    for (int i = 0; i < num_bufs; ++i) {
-        int ret = engine->registerLocalMemory(bufs[i], buf_bytes, "*", true);
-        if (ret != 0) {
-            LOG(ERROR) << "registerLocalMemory failed at buffer " << i
-                       << ": ret=" << ret;
-            for (int j = 0; j < i; ++j) engine->unregisterLocalMemory(bufs[j]);
-            for (auto* p : bufs) munmap(p, buf_bytes);
-            return 1;
-        }
-        if ((i + 1) % 50 == 0 || i == num_bufs - 1) {
-            auto now = std::chrono::steady_clock::now();
-            double elapsed = std::chrono::duration<double>(now - t0).count();
-            LOG(INFO) << "  registered " << (i + 1) << "/" << num_bufs << " ("
-                      << elapsed << "s)";
-        }
-    }
-    auto t1 = std::chrono::steady_clock::now();
-    LOG(INFO) << "Registration complete: "
-              << std::chrono::duration<double>(t1 - t0).count() << "s";
-
-    LOG(INFO) << "Target ready. First buffer at " << bufs[0]
-              << ". Waiting for initiator (Ctrl+C to stop)...";
-
-    while (g_running) sleep(1);
-
-    LOG(INFO) << "Shutting down target...";
-    for (auto* p : bufs) {
-        engine->unregisterLocalMemory(p);
-        munmap(p, buf_bytes);
-    }
-    return 0;
-}
-
-struct LatencyStats {
-    double avg_ms;
-    double p50_ms;
-    double p99_ms;
-    double throughput_gbs;
-};
-
-static LatencyStats computeStats(std::vector<double>& latencies_ms,
-                                 size_t transfer_bytes) {
-    std::sort(latencies_ms.begin(), latencies_ms.end());
-    double sum = std::accumulate(latencies_ms.begin(), latencies_ms.end(), 0.0);
-    size_t n = latencies_ms.size();
-    LatencyStats stats;
-    stats.avg_ms = sum / n;
-    stats.p50_ms = latencies_ms[n / 2];
-    stats.p99_ms = latencies_ms[std::min(n - 1, (size_t)(n * 0.99))];
-    stats.throughput_gbs = (transfer_bytes / 1e9) / (stats.p50_ms / 1000.0);
-    return stats;
-}
-
-static int runInitiator(TransferEngine* engine) {
-    size_t transfer_bytes =
-        static_cast<size_t>(FLAGS_transfer_mb * 1024 * 1024);
-
-    LOG(INFO) << "=== Initiator Node ===";
-    LOG(INFO) << "Target: " << FLAGS_target;
-    LOG(INFO) << "Transfer size: " << FLAGS_transfer_mb << " MB";
-    LOG(INFO) << "Threads: " << FLAGS_threads;
-
-    if (FLAGS_target.empty()) {
-        LOG(ERROR) << "--target required in initiator mode";
-        return 1;
-    }
-
-    // Allocate local receive buffer (one per thread)
-    size_t recv_bytes = transfer_bytes;
-    std::vector<void*> recv_bufs(FLAGS_threads);
-    for (int t = 0; t < FLAGS_threads; ++t) {
-        recv_bufs[t] = allocateHugepage(recv_bytes);
-        if (!recv_bufs[t]) {
-            LOG(ERROR) << "Failed to allocate receive buffer for thread " << t;
-            return 1;
-        }
-        int ret = engine->registerLocalMemory(recv_bufs[t], recv_bytes, "cpu:0",
-                                              true);
-        if (ret != 0) {
-            LOG(ERROR) << "Failed to register receive buffer: " << ret;
-            return 1;
-        }
-    }
-    LOG(INFO) << "Allocated and registered " << FLAGS_threads
-              << " receive buffers of " << recv_bytes / 1e6 << " MB each";
-
-    // Open remote segment
-    auto segment_id = engine->openSegment(FLAGS_target);
-    if (segment_id < 0) {
-        LOG(ERROR) << "openSegment failed for " << FLAGS_target;
-        return 1;
-    }
-
-    // Get remote buffer info
-    auto segment_desc = engine->getMetadata()->getSegmentDescByID(segment_id);
-    if (!segment_desc || segment_desc->buffers.empty()) {
-        LOG(ERROR) << "No remote buffers found";
-        return 1;
-    }
-    size_t num_remote_bufs = segment_desc->buffers.size();
-    LOG(INFO) << "Remote has " << num_remote_bufs << " buffers";
-    LOG(INFO) << "First buffer: addr=0x" << std::hex
-              << segment_desc->buffers[0].addr << std::dec
-              << " size=" << segment_desc->buffers[0].length;
-
-    // Connection warmup: small transfer to establish endpoints
-    LOG(INFO) << "Warming up connection...";
-    {
-        size_t warmup_size = std::min(transfer_bytes, (size_t)(64 * 1024));
-        auto batch_id = engine->allocateBatchID(1);
-        TransferRequest req;
-        req.opcode = TransferRequest::READ;
-        req.source = (uint8_t*)recv_bufs[0];
-        req.target_id = segment_id;
-        req.target_offset = segment_desc->buffers[0].addr;
-        req.length = warmup_size;
-        auto s = engine->submitTransfer(batch_id, {req});
-        if (!s.ok()) {
-            LOG(ERROR) << "Warmup transfer failed: " << s.ToString();
-            return 1;
-        }
-        while (true) {
-            TransferStatus status;
-            engine->getTransferStatus(batch_id, 0, status);
-            if (status.s == TransferStatusEnum::COMPLETED) break;
-            if (status.s == TransferStatusEnum::FAILED) {
-                LOG(ERROR) << "Warmup transfer FAILED";
-                return 1;
-            }
-        }
-        engine->freeBatchID(batch_id);
-    }
-    LOG(INFO) << "Connection ready.";
-
-    // Refresh segment desc after connection warmup (endpoints are now up)
-    engine->syncSegmentCache(FLAGS_target);
-    segment_desc = engine->getMetadata()->getSegmentDescByID(segment_id);
-
-    // Worker function: each thread runs its own transfer loop
-    struct ThreadResult {
-        std::vector<double> latencies;
-        int errors = 0;
-    };
-
-    auto workerFn = [&](int tid, int warmup_iters, int bench_iters,
-                        ThreadResult* result) {
-        void* my_recv = recv_bufs[tid];
-        for (int w = 0; w < warmup_iters; ++w) {
-            size_t buf_idx = (tid + w * FLAGS_threads) % num_remote_bufs;
-            uint64_t raddr = segment_desc->buffers[buf_idx].addr;
-            size_t rlen = segment_desc->buffers[buf_idx].length;
-            size_t xfer = std::min(transfer_bytes, rlen);
-
-            auto bid = engine->allocateBatchID(1);
-            TransferRequest req;
-            req.opcode = TransferRequest::READ;
-            req.source = (uint8_t*)my_recv;
-            req.target_id = segment_id;
-            req.target_offset = raddr;
-            req.length = xfer;
-            engine->submitTransfer(bid, {req});
-            while (true) {
-                TransferStatus st;
-                engine->getTransferStatus(bid, 0, st);
-                if (st.s == TransferStatusEnum::COMPLETED ||
-                    st.s == TransferStatusEnum::FAILED)
-                    break;
-            }
-            engine->freeBatchID(bid);
-        }
-
-        for (int i = 0; i < bench_iters; ++i) {
-            size_t buf_idx = (tid + i * FLAGS_threads) % num_remote_bufs;
-            uint64_t raddr = segment_desc->buffers[buf_idx].addr;
-            size_t rlen = segment_desc->buffers[buf_idx].length;
-            size_t xfer = std::min(transfer_bytes, rlen);
-
-            auto t0 = std::chrono::steady_clock::now();
-            auto bid = engine->allocateBatchID(1);
-            TransferRequest req;
-            req.opcode = TransferRequest::READ;
-            req.source = (uint8_t*)my_recv;
-            req.target_id = segment_id;
-            req.target_offset = raddr;
-            req.length = xfer;
-            auto s = engine->submitTransfer(bid, {req});
-            if (!s.ok()) {
-                result->errors++;
-                engine->freeBatchID(bid);
-                continue;
-            }
-            bool ok = false;
-            while (true) {
-                TransferStatus st;
-                engine->getTransferStatus(bid, 0, st);
-                if (st.s == TransferStatusEnum::COMPLETED) {
-                    ok = true;
-                    break;
-                }
-                if (st.s == TransferStatusEnum::FAILED) {
-                    result->errors++;
-                    break;
-                }
-            }
-            engine->freeBatchID(bid);
-            if (ok) {
-                auto t1 = std::chrono::steady_clock::now();
-                result->latencies.push_back(
-                    std::chrono::duration<double, std::milli>(t1 - t0).count());
-            }
-        }
-    };
-
-    // Run warmup + benchmark with threads
-    int num_threads = FLAGS_threads;
-    LOG(INFO) << "Running with " << num_threads << " threads, " << FLAGS_warmup
-              << " warmup + " << FLAGS_iterations
-              << " bench iterations per thread...";
-
-    std::vector<ThreadResult> results(num_threads);
-    std::vector<std::thread> threads;
-
-    auto wall_t0 = std::chrono::steady_clock::now();
-    for (int t = 0; t < num_threads; ++t) {
-        threads.emplace_back(workerFn, t, FLAGS_warmup * FLAGS_iterations,
-                             FLAGS_iterations, &results[t]);
-    }
-    for (auto& th : threads) th.join();
-    auto wall_t1 = std::chrono::steady_clock::now();
-    double wall_ms =
-        std::chrono::duration<double, std::milli>(wall_t1 - wall_t0).count();
-
-    // Aggregate results
-    std::vector<double> all_latencies;
-    int total_errors = 0;
-    for (auto& r : results) {
-        all_latencies.insert(all_latencies.end(), r.latencies.begin(),
-                             r.latencies.end());
-        total_errors += r.errors;
-    }
-
-    if (all_latencies.empty()) {
-        LOG(ERROR) << "All transfers failed";
-        return 1;
-    }
-
-    auto stats = computeStats(all_latencies, transfer_bytes);
-    size_t total_xfers = all_latencies.size();
-    double total_bytes = (double)total_xfers * transfer_bytes;
-    double agg_throughput = total_bytes / 1e9 / (wall_ms / 1000.0);
-
-    LOG(INFO) << "=== Results (" << num_threads << " threads) ===";
-    LOG(INFO) << "Transfer: " << FLAGS_transfer_mb << " MB x " << total_xfers
-              << " transfers";
-    LOG(INFO) << "Wall time: " << wall_ms << " ms";
-    LOG(INFO) << "Per-transfer p50: " << stats.p50_ms << " ms"
-              << "  p99: " << stats.p99_ms << " ms";
-    LOG(INFO) << "Per-transfer throughput: " << stats.throughput_gbs << " GB/s";
-    LOG(INFO) << "Aggregate throughput: " << agg_throughput << " GB/s";
-    LOG(INFO) << "Errors: " << total_errors;
-
-    // Per-thread stats
-    for (int t = 0; t < num_threads; ++t) {
-        if (results[t].latencies.empty()) continue;
-        auto ts = computeStats(results[t].latencies, transfer_bytes);
-        LOG(INFO) << "  Thread " << t << ": p50=" << ts.p50_ms
-                  << "ms  tput=" << ts.throughput_gbs << " GB/s"
-                  << "  iters=" << results[t].latencies.size()
-                  << "  errors=" << results[t].errors;
-    }
-
-    // Cleanup
-    for (int t = 0; t < FLAGS_threads; ++t) {
-        engine->unregisterLocalMemory(recv_bufs[t]);
-        munmap(recv_bufs[t], recv_bytes);
-    }
-    return total_errors > 0 ? 1 : 0;
-}
-
-int main(int argc, char** argv) {
-    google::InitGoogleLogging(argv[0]);
-    gflags::ParseCommandLineFlags(&argc, &argv, true);
-    FLAGS_logtostderr = 1;
-
-    setupSignalHandler();
-
-    if (FLAGS_server.empty()) {
-        LOG(ERROR) << "--server is required";
-        return 1;
-    }
-
-    // Parse host:port
-    auto colon = FLAGS_server.rfind(':');
-    std::string host = FLAGS_server.substr(0, colon);
-    uint64_t port = 12345;
-    if (colon != std::string::npos)
-        port = std::stoull(FLAGS_server.substr(colon + 1));
-
-    auto engine = std::make_unique<TransferEngine>(false);
-    int ret = engine->init(FLAGS_metadata, FLAGS_server, host, port);
-    if (ret != 0) {
-        LOG(ERROR) << "Engine init failed: " << ret;
-        return 1;
-    }
-
-    // Discover topology and install EFA transport (all NICs)
-    engine->getLocalTopology()->discover({});
-    auto* xport = engine->installTransport("efa", nullptr);
-    if (!xport) {
-        LOG(ERROR) << "installTransport(efa) failed";
-        return 1;
-    }
-
-    std::string actual_server = engine->getLocalIpAndPort();
-    LOG(INFO) << "Actual server name (use this for --target): "
-              << actual_server;
-
-    if (FLAGS_mode == "target") {
-        ret = runTarget(engine.get());
-    } else if (FLAGS_mode == "initiator") {
-        ret = runInitiator(engine.get());
-    } else {
-        LOG(ERROR) << "Unknown mode: " << FLAGS_mode;
-        ret = 1;
-    }
-
-    return ret;
-}

mooncake-wheel/mooncake/mooncake_connector_v1.py

@@ -28,19 +28,6 @@ from vllm.attention.selector import get_attn_backend
 from vllm.config import VllmConfig
 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
     KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole)
-
-# SupportsHMA was introduced in vllm-project/vllm PR #25712 and is enforced
-# for KV connectors by PR #27592. It is the marker the Hybrid Memory
-# Allocator uses to allow PD-disaggregation with hybrid (e.g. attention +
-# Mamba2) models. We import it conditionally so this connector keeps
-# working on older vLLM releases that pre-date the interface; on those
-# releases the marker is a no-op object base and the
-# `request_finished_all_groups` shim below is dead code.
-try:
-    from vllm.distributed.kv_transfer.kv_connector.v1.base import (
-        SupportsHMA)
-except ImportError:  # pragma: no cover - older vLLM
-    SupportsHMA = object  # type: ignore[assignment, misc]
 from vllm.distributed.parallel_state import (get_tensor_model_parallel_rank,
                                              get_tp_group)
 from vllm.forward_context import ForwardContext
@@ -117,12 +104,7 @@ class MooncakeConnectorMetadata(KVConnectorMetadata):
             self.reqs_to_send[request_id] = local_block_ids
 
 
-class MooncakeConnector(KVConnectorBase_V1, SupportsHMA):
-    # Subclassing SupportsHMA gates this connector through vLLM's Hybrid
-    # Memory Allocator path, which is required to PD-disaggregate hybrid
-    # attention + Mamba2 models (e.g. nvidia/NVIDIA-Nemotron-Nano-9B-v2).
-    # On vLLM versions that pre-date PR #25712 the marker resolves to
-    # `object` and this is identical to the previous single-base class.
+class MooncakeConnector(KVConnectorBase_V1):
 
     def __init__(self, vllm_config: VllmConfig, role: KVConnectorRole):
         assert vllm_config.kv_transfer_config is not None
@@ -172,28 +154,6 @@ class MooncakeConnector(KVConnectorBase_V1, SupportsHMA):
         assert self.connector_scheduler is not None
         return self.connector_scheduler.request_finished(request, block_ids)
 
-    def request_finished_all_groups(
-        self,
-        request: "Request",
-        block_ids: tuple[list[int], ...],
-    ) -> tuple[bool, Optional[dict[str, Any]]]:
-        """SupportsHMA hook for hybrid (multi-group) KV cache layouts.
-
-        Hybrid models (e.g. attention + Mamba2) expose one block-id list
-        per KV cache group instead of a single flat list. The Mooncake
-        transport itself does not yet distinguish groups on the wire, so
-        we flatten the per-group lists and delegate to the existing
-        single-group `request_finished`. This is the minimum-viable shim
-        that satisfies the `SupportsHMA` contract and unblocks the
-        Hybrid Memory Allocator gate so the engine can start up; cross-
-        node fidelity for non-attention SSM/Mamba state is not asserted
-        by this method and remains a follow-up (the Mamba2 backend in
-        vLLM still raises NotImplementedError from
-        `get_kv_cache_shape()`).
-        """
-        flat: list[int] = [b for group in block_ids for b in group]
-        return self.request_finished(request, flat)
-
     ############################################################
     # Worker Side Methods
     ############################################################

mooncake-wheel/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "mooncake-transfer-engine"
-version = "0.3.10.post2"
+version = "0.3.10.post1"
 description = "Python binding of a Mooncake library using pybind11"
 authors = [
     { name = "Mooncake Authors" }