docs: implement documentation plan for missing Mooncake features

Agent-Logs-Url: https://github.com/kvcache-ai/Mooncake/sessions/3432c219-559e-4d6a-b137-ce729304fc49

Co-authored-by: stmatengss <11641725+stmatengss@users.noreply.github.com>

.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
@@ -345,14 +343,6 @@ jobs:
 
     - name: Run tests with ssd
       run: |
-        # Reserve port 50052 (mooncake_client RPC port) so the kernel never
-        # auto-allocates it as ephemeral source port for other outbound
-        # connections in the test suite. Without this, a random Python test
-        # connection can pick src_port=50052, leave a TIME_WAIT on
-        # <eth0_ip>:50052 for 60s, and block mooncake_client's bind to
-        # 0.0.0.0:50052 even with SO_REUSEADDR (Linux only relaxes
-        # TIME_WAIT+bind conflict for same-IP or loopback).
-        sudo sysctl -w net.ipv4.ip_local_reserved_ports=50052
         source test_env/bin/activate
         MC_STORE_MEMCPY=false TEST_SSD_OFFLOAD_IN_EVICT=true ./scripts/run_tests.sh
         rm -rf /tmp/mooncake_test_ssd
@@ -746,18 +736,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 +748,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

@@ -15,8 +15,8 @@ jobs:
 
     container:
       image: localhost:5000/mooncake-hixl-ci:v5
-      options: --privileged --user 0:0 --device /dev/davinci0 --device /dev/davinci1 --device /dev/davinci2 --device /dev/davinci3
-                    --device /dev/davinci4 --device /dev/davinci5 --device /dev/davinci6 --device /dev/davinci7
+      options: --privileged --user 0:0 --device /dev/davinci0 --device /dev/davinci1 --device /dev/davinci2 --device /dev/davinci3 
+                    --device /dev/davinci4 --device /dev/davinci5 --device /dev/davinci6 --device /dev/davinci7 
                     --device /dev/davinci_manager --device /dev/devmm_svm --device /dev/hisi_hdc --ulimit nproc=65535:65535
       env:
         GITHUB_ACTIONS: "true"
@@ -28,132 +28,31 @@ jobs:
         - /etc/hccn.conf:/etc/hccn.conf
 
     steps:
-    - name: Configure GitHub fetch defaults
-      shell: bash
-      run: |
-        git config --global protocol.version 2
-        git config --global http.version HTTP/1.1
-        git config --global http.lowSpeedLimit 1024
-        git config --global http.lowSpeedTime 30
-
     - name: Checkout code
-      id: checkout_code
-      continue-on-error: true
       uses: actions/checkout@v4
       with:
         ref: ${{ inputs.checkout_ref || github.sha }}
         fetch-depth: 1
         persist-credentials: false
 
-    - name: Retry checkout via GitHub mirror
-      if: steps.checkout_code.outcome == 'failure'
-      shell: bash
-      env:
-        ASCEND_GITHUB_MIRROR_URLS: ${{ vars.ASCEND_GITHUB_MIRROR_URLS }}
-        CHECKOUT_REF: ${{ inputs.checkout_ref || github.sha }}
-      run: |
-        set -euo pipefail
-
-        if [ -z "${ASCEND_GITHUB_MIRROR_URLS:-}" ]; then
-          echo "Checkout from GitHub failed and ASCEND_GITHUB_MIRROR_URLS is not set"
-          exit 1
-        fi
-
-        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"
-        }
-
-        candidates=()
-        while IFS= read -r raw; do
-          base="$(normalize_base "$raw" || true)"
-          [ -n "$base" ] || continue
-          [ "$base" = "https://github.com/" ] && continue
-          candidates+=("$base")
-        done < <(printf '%s\n' "$ASCEND_GITHUB_MIRROR_URLS" | tr ',;' '\n')
-
-        if [ ${#candidates[@]} -eq 0 ]; then
-          echo "Checkout from GitHub failed and no valid mirror candidates were configured"
-          exit 1
-        fi
-
-        workdir="${GITHUB_WORKSPACE}"
-        git config --global --add safe.directory "$workdir"
-
-        for base in "${candidates[@]}"; do
-          mirror_url="${base}https://github.com/${GITHUB_REPOSITORY}.git"
-          echo "Retrying checkout with ${mirror_url}"
-
-          find "$workdir" -mindepth 1 -maxdepth 1 -exec rm -rf {} +
-          git init "$workdir"
-          git -C "$workdir" remote add origin "$mirror_url"
-
-          if git -C "$workdir" fetch --depth=1 origin "$CHECKOUT_REF" && \
-             git -C "$workdir" checkout --force --detach FETCH_HEAD; then
-            echo "Mirror checkout succeeded via ${base}"
-            exit 0
-          fi
-
-          echo "Mirror checkout failed via ${base}"
-          rm -rf "$workdir/.git"
-        done
-
-        echo "Direct GitHub checkout failed and all mirror retries failed"
-        exit 1
-
     - 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
@@ -221,7 +120,7 @@ jobs:
           # Check if master is running
           if ! kill -0 $MASTER_PID 2>/dev/null; then
               echo "Error: Mooncake Master failed to start"
-              cat /tmp/mooncake_master.log
+              cat /tmp/mooncake_master.log 
               exit 1
           fi
 
@@ -345,7 +244,7 @@ jobs:
 
           echo ""
           echo "All Hixl Mooncake Store tests completed successfully!"
-
+          
 
     - name: Test Summary
       if: always()

.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:
@@ -47,7 +37,7 @@ repos:
     hooks:
       - id: codespell
         exclude: '^(extern/|FAST25-release/)'
-        args: ['--ignore-words-list=te,mooncake,KVCache,cann']
+        args: ['--ignore-words-list=te,mooncake,KVCache']
 
   - repo: https://github.com/pre-commit/mirrors-clang-format
     rev: v20.1.8

.typos.toml

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

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/ha-hot-standby.md

@@ -0,0 +1,146 @@
+# Mooncake Store HA Hot Standby
+
+Mooncake Store supports a **Hot Standby** mode that keeps a passive replica of the master metadata in sync with the primary. When the primary fails, the standby can be promoted to take over with minimal downtime and no data loss.
+
+## Architecture
+
+```
+  ┌─────────────────────┐          OpLog stream           ┌─────────────────────┐
+  │   Primary Master    │ ──────────────────────────────▶ │   Standby Master    │
+  │  (MasterService)    │                                  │ (HotStandbyService) │
+  │                     │  ①  Snapshot bootstrap (once)   │                     │
+  │  Oplog  ─────────── │ ──────────────────────────────▶ │  ─── apply oplogs   │
+  │  (OpLogManager)     │  ②  Oplog replication (steady)  │  (OpLogApplier)     │
+  └─────────────────────┘                                  └─────────────────────┘
+           ▲                                                        │
+           │           Leader Election (etcd / Redis)              │  promotion
+           └────────────────────────────────────────────────────────┘
+```
+
+### Phase 1 — Snapshot Bootstrap
+
+On startup, the standby optionally downloads the latest snapshot produced by the primary (see [Snapshot / Restore flags](mooncake-store-deployment-guide.md)). This baseline allows the standby to catch up quickly without replaying the entire oplog history.
+
+Enable snapshot bootstrap via `HotStandbyConfig.enable_snapshot_bootstrap = true`.
+
+### Phase 2 — Oplog Replication
+
+After the snapshot is applied, the standby enters **steady-state replication**: it continuously polls the primary's `OpLogReplicator` for new oplog entries and applies them locally through `OpLogApplier`. The lag is bounded by `max_replication_lag_entries` (default 1000 entries).
+
+### Leader Election
+
+Mooncake uses either **etcd** or **Redis** as the distributed coordination backend for leader election:
+
+- **etcd** (`STORE_USE_ETCD`): Build with `-DSTORE_USE_ETCD=ON`. Set `--etcd_endpoints` on the master.
+- **Redis** (`STORE_USE_REDIS`): Build with `-DSTORE_USE_REDIS=ON`. Set connection details via environment variables (see [Redis HA Backend](mooncake-store-deployment-guide.md#redis-ha-backend)).
+
+When the primary is unresponsive for longer than `--client_ttl` seconds, the standby acquires the leader lease and promotes itself.
+
+### Promotion
+
+Promotion is handled by `StandbyStateMachine`. On promotion the standby:
+1. Stops polling the (now-dead) primary.
+2. Registers itself in the metadata service under the primary's endpoint address.
+3. Begins accepting client RPCs.
+
+Clients that retry their connections will automatically reconnect to the promoted standby without any application-level change.
+
+## Configuration Reference
+
+### `HotStandbyConfig` Fields
+
+| Field | Default | Description |
+|-------|---------|-------------|
+| `standby_id` | — | Unique identifier for this standby instance |
+| `primary_address` | — | `ip:port` of the primary master's RPC endpoint |
+| `replication_port` | 0 (auto) | Port for the oplog replication channel |
+| `verification_interval_sec` | 30 | How often to verify sync status |
+| `max_replication_lag_entries` | 1000 | Alert threshold for oplog lag |
+| `enable_verification` | true | Enable periodic lag verification |
+| `enable_snapshot_bootstrap` | false | Download latest snapshot before replicating |
+| `enable_oplog_following` | true | Enable steady-state oplog following |
+| `oplog_store_type` | default | Where to persist the local oplog copy |
+| `oplog_store_root_dir` | default | Root directory for local oplog storage |
+| `oplog_poll_interval_ms` | default | Polling interval for new oplog entries |
+
+### Master Startup Flags (HA-related)
+
+The following flags from [Mooncake Store Deployment Guide](mooncake-store-deployment-guide.md) apply to HA setups:
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--enable_ha` | `false` | Enable HA mode (requires etcd or Redis) |
+| `--etcd_endpoints` | — | Semicolon-separated etcd endpoints |
+| `--client_ttl` | `10` s | How long a client (or standby) has to re-ping before considered dead |
+| `--cluster_id` | `mooncake_cluster` | Cluster ID used for persistence keys in HA mode |
+
+## Deployment Example
+
+### Step 1: Start the Primary Master (with HA enabled)
+
+```bash
+mooncake_master \
+    --rpc_port=50051 \
+    --enable_ha=true \
+    --etcd_endpoints="http://etcd-0:2379;http://etcd-1:2379;http://etcd-2:2379" \
+    --cluster_id=prod-cluster \
+    --client_ttl=15 \
+    --enable_snapshot=true \
+    --snapshot_backend_type=local \
+    --snapshot_interval_seconds=300
+```
+
+### Step 2: Start the Standby Master
+
+The standby is a separate `mooncake_master` process with identical flags **plus** the standby-specific ones:
+
+```bash
+MOONCAKE_SNAPSHOT_LOCAL_PATH=/data/mooncake_snapshots \
+mooncake_master \
+    --rpc_port=50052 \
+    --enable_ha=true \
+    --etcd_endpoints="http://etcd-0:2379;http://etcd-1:2379;http://etcd-2:2379" \
+    --cluster_id=prod-cluster \
+    --client_ttl=15 \
+    --enable_snapshot_restore=true \
+    --snapshot_backend_type=local
+```
+
+The standby will:
+1. Detect that another master holds the leader lease via etcd.
+2. Start the `HotStandbyService` with snapshot bootstrap enabled.
+3. Download the latest snapshot from the primary and apply it.
+4. Begin following the oplog in steady state.
+
+### Step 3: Verify Sync Status
+
+The standby logs its sync status periodically. Look for lines like:
+
+```
+[HotStandbyService] applied_seq_id=12345 lag=0 entries
+[HotStandbyService] verification OK: primary_seq_id=12345 standby_seq_id=12345
+```
+
+A non-zero lag that is growing indicates the standby cannot keep up with the primary write rate — consider reducing the primary write load or increasing `oplog_poll_interval_ms`.
+
+### Step 4: Simulate a Failover
+
+Kill the primary master process. Within `--client_ttl` seconds, the standby will:
+
+1. Detect the primary is down (etcd lease expiry).
+2. Acquire the leader lease.
+3. Promote itself and start serving client RPCs on its own `--rpc_port`.
+
+Point clients to the standby's address or use a DNS/load-balancer alias.
+
+## Tuning Tips
+
+- **`max_replication_lag_entries`**: Lower values trigger earlier alerts but may cause false positives during write bursts. The default of 1000 is conservative for most workloads.
+- **`verification_interval_sec`**: Increase to reduce overhead if the standby's oplog applier is a bottleneck.
+- **`enable_snapshot_bootstrap`**: Always enable in production to reduce time-to-sync after a standby restart.
+- **`client_ttl`**: Set this equal to or slightly higher than your deployment's network heartbeat interval. A value too low causes spurious failovers; too high delays recovery.
+
+## See Also
+
+- [Mooncake Store Deployment Guide](mooncake-store-deployment-guide) — snapshot flags, S3 backend, Redis HA
+- [Mooncake Store Design](../design/mooncake-store)

docs/source/deployment/monitoring.md

@@ -0,0 +1,130 @@
+# Monitoring Mooncake with Prometheus and Grafana
+
+Mooncake Master exposes a Prometheus-compatible `/metrics` endpoint. The `monitoring/` directory in the repository contains a ready-to-use Docker Compose stack that wires Prometheus and Grafana together with a pre-built dashboard for `mooncake_master`.
+
+## Quick Start
+
+### Prerequisites
+
+- Docker and Docker Compose installed on the monitoring host.
+- `mooncake_master` accessible from the monitoring host.
+
+### Step 1: Start the Monitoring Stack
+
+```bash
+cd monitoring
+docker-compose up -d
+```
+
+This starts two containers:
+- **Prometheus** — scrapes metrics from `mooncake_master` every 15 s.
+- **Grafana** — pre-configured with a Prometheus data source and a sample dashboard.
+
+### Step 2: Open the UIs
+
+| UI | URL | Credentials |
+|----|-----|-------------|
+| Prometheus | <http://localhost:9090> | — |
+| Grafana | <http://localhost:3000> | `admin` / `admin` |
+
+Navigate to **Grafana → Dashboards** to find the pre-built `mooncake_master` dashboard.
+
+### Step 3: Start `mooncake_master` with Metrics Enabled
+
+```bash
+./build/mooncake_master \
+    --metrics_port=9003 \
+    --enable_metric_reporting=true \
+    --rpc_port=50051
+```
+
+Verify the metrics endpoint is live:
+
+```bash
+curl -s http://localhost:9003/metrics | head -20
+```
+
+You should see Prometheus-format lines such as:
+
+```
+# HELP mooncake_master_kv_object_count Total number of KV objects in the store
+# TYPE mooncake_master_kv_object_count gauge
+mooncake_master_kv_object_count 4096
+```
+
+Check **Prometheus → Status → Targets** — the `mooncake-master` job should show `UP`.
+
+## Configuration Files
+
+| File | Description |
+|------|-------------|
+| `monitoring/docker-compose.yml` | Service definitions for Prometheus and Grafana |
+| `monitoring/prometheus/prometheus.yml` | Prometheus scrape configuration |
+| `monitoring/grafana/` | Grafana provisioning (data source + dashboard JSON) |
+
+### Prometheus Scrape Target
+
+By default, `prometheus.yml` scrapes `host.docker.internal:9003`. On Linux, `host.docker.internal` may not be available — add the following to the `prometheus` service in `docker-compose.yml`:
+
+```yaml
+extra_hosts:
+  - "host.docker.internal:host-gateway"
+```
+
+To scrape a remote `mooncake_master`, change the target in `prometheus/prometheus.yml`:
+
+```yaml
+scrape_configs:
+  - job_name: mooncake-master
+    static_configs:
+      - targets:
+          - "10.0.0.1:9003"   # replace with actual master host
+```
+
+## Available Metrics
+
+The `/metrics/summary` endpoint (human-readable) and `/metrics` endpoint (Prometheus format) expose the following categories:
+
+| Category | Example Metric | Description |
+|----------|---------------|-------------|
+| KV objects | `mooncake_master_kv_object_count` | Total objects in the store |
+| Memory | `mooncake_master_segment_free_bytes` | Free bytes per segment |
+| Eviction | `mooncake_master_eviction_total` | Eviction events |
+| Tasks | `mooncake_master_pending_tasks` | Pending transfer tasks |
+| RPC | `mooncake_master_rpc_requests_total` | Total RPC requests served |
+
+Browse all available metrics via:
+
+```bash
+curl -s http://localhost:9003/metrics/summary
+```
+
+## Grafana Alerts (Optional)
+
+To set up alerts in Grafana:
+
+1. Open the dashboard and click the panel you want to alert on.
+2. Choose **Edit → Alert → Create alert rule**.
+3. Example: alert when `mooncake_master_segment_free_bytes` drops below 10 % of total.
+
+## Multi-Master Monitoring
+
+To monitor multiple masters in one Prometheus instance, add multiple targets:
+
+```yaml
+scrape_configs:
+  - job_name: mooncake-masters
+    static_configs:
+      - targets:
+          - "master-0:9003"
+          - "master-1:9003"
+          - "master-2:9003"
+    relabel_configs:
+      - source_labels: [__address__]
+        target_label: instance
+```
+
+## See Also
+
+- [Mooncake Store Deployment Guide](mooncake-store-deployment-guide) — master startup flags including `--metrics_port`
+- [Mooncake Store HA Hot Standby](ha-hot-standby) — monitoring standby lag

docs/source/deployment/mooncake-store-deployment-guide.md

@@ -153,6 +153,126 @@ This sets the log level for yalantinglibs (including coro_rpc and coro_http) to
 
 Available log levels: trace, debug, info, warn (or warning), error, and critical.
 
+## S3 Snapshot Backend
+
+When `--snapshot_backend_type=s3` is set, Mooncake Master stores snapshots in an S3-compatible object store.
+
+### Required Environment Variables
+
+| Variable | Description |
+|----------|-------------|
+| `AWS_ACCESS_KEY_ID` | S3 access key (or IAM role credential) |
+| `AWS_SECRET_ACCESS_KEY` | S3 secret key |
+| `AWS_DEFAULT_REGION` | AWS region (e.g., `us-east-1`) |
+| `MOONCAKE_SNAPSHOT_S3_BUCKET` | Target S3 bucket name |
+| `MOONCAKE_SNAPSHOT_S3_PREFIX` | (Optional) Key prefix inside the bucket |
+
+For IAM role-based authentication (recommended in production), omit `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` — the SDK will use the instance profile automatically.
+
+### Bucket Naming Conventions
+
+- Create a dedicated bucket for Mooncake snapshots (e.g., `my-org-mooncake-snapshots`).
+- Enable versioning on the bucket for additional safety.
+- Set a lifecycle rule to expire old snapshot objects beyond the `--snapshot_retention_count` limit.
+
+> **Warning:** The S3 bucket is a managed directory. Do not store other data under the same key prefix as Mooncake snapshots — old snapshots are deleted automatically during cleanup.
+
+### Example Startup Command (S3 backend)
+
+```bash
+export AWS_ACCESS_KEY_ID=AKIAIOSFODNN7EXAMPLE
+export AWS_SECRET_ACCESS_KEY=wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY
+export AWS_DEFAULT_REGION=us-east-1
+export MOONCAKE_SNAPSHOT_S3_BUCKET=my-org-mooncake-snapshots
+export MOONCAKE_SNAPSHOT_S3_PREFIX=prod-cluster/
+
+mooncake_master \
+    --rpc_port=50051 \
+    --enable_snapshot=true \
+    --snapshot_backend_type=s3 \
+    --snapshot_interval_seconds=300 \
+    --snapshot_retention_count=3 \
+    --enable_snapshot_restore=true
+```
+
+### Restore Procedure
+
+On master restart with `--enable_snapshot_restore=true`, the master:
+1. Lists snapshots in the configured S3 bucket/prefix.
+2. Downloads the latest snapshot.
+3. Applies the snapshot to restore in-memory metadata.
+4. Resumes serving client RPCs.
+
+If `--snapshot_backup_dir` is also set, the downloaded snapshot is additionally saved locally as a fallback.
+
+---
+
+## Redis HA Backend
+
+Mooncake Store HA mode can use **Redis** instead of etcd for distributed leader election and cluster coordination.
+
+### Build Requirement
+
+Redis HA support must be enabled at compile time:
+
+```bash
+cmake .. \
+    -DSTORE_USE_REDIS=ON \
+    -DCMAKE_BUILD_TYPE=RelWithDebInfo
+make -j$(nproc)
+```
+
+This adds the `STORE_USE_REDIS` compile definition and links `hiredis`.
+
+### Connection Configuration
+
+| Environment Variable | Description |
+|----------------------|-------------|
+| `MC_REDIS_PASSWORD` | Redis AUTH password (omit if no auth is configured) |
+| `MC_REDIS_DB_INDEX` | Redis database index (default: `0`) |
+
+The Redis server URL is passed via `--etcd_endpoints` using the `redis://` scheme:
+
+```bash
+mooncake_master \
+    --rpc_port=50051 \
+    --enable_ha=true \
+    --etcd_endpoints="redis://10.0.0.10:6379" \
+    --cluster_id=prod-cluster \
+    --client_ttl=15
+```
+
+For a Redis Cluster or Sentinel setup:
+
+```bash
+# Redis Cluster (semicolon-separated nodes)
+--etcd_endpoints="redis://10.0.0.10:6379;redis://10.0.0.11:6379;redis://10.0.0.12:6379"
+```
+
+### Key Hash-Tag Conventions
+
+Mooncake uses Redis hash tags to ensure that all keys for a given cluster land on the same Redis cluster slot. The tag is derived from `--cluster_id` and sanitised via `SanitizeHashTagComponent()` (which strips characters that are illegal inside `{}`). For example:
+
+- `--cluster_id=prod-cluster` → Redis keys use `{prod-cluster}` as the hash tag.
+- All leader election keys, oplog entries, and session heartbeats share this tag.
+
+Do not use the same Redis instance / keyspace for other applications without ensuring their keys use different hash tags.
+
+### ConnectRedis() Helper
+
+The internal `ha::backends::redis::ConnectRedis()` helper reads `MC_REDIS_PASSWORD` and `MC_REDIS_DB_INDEX` from the environment automatically. No additional code changes are needed; configure these variables before starting the master.
+
+### Redis vs etcd
+
+| Feature | Redis | etcd |
+|---------|-------|------|
+| Build flag | `STORE_USE_REDIS=ON` | `STORE_USE_ETCD=ON` |
+| URL scheme | `redis://` | `etcd://` or bare `host:port` |
+| Cluster support | Redis Cluster / Sentinel | etcd cluster |
+| Recommended for | Environments already running Redis | New deployments |
+
+---
+
 ## Quick Tips
 
 - Scale `--rpc_thread_num` with available CPU cores and workload.
@@ -167,4 +287,7 @@ Available log levels: trace, debug, info, warn (or warning), error, and critical
 :maxdepth: 1
 
 ssd-offload
+ha-hot-standby
+monitoring
+multi-tier-storage
 :::

docs/source/deployment/multi-tier-storage.md

@@ -0,0 +1,140 @@
+# Multi-Tier Storage
+
+Mooncake Store supports a **multi-tier storage hierarchy** that extends the in-memory KV cache with a persistent distributed filesystem (DFS) layer. This enables workloads that cannot fit entirely in GPU/CPU memory to spill objects to a fast DFS backend (e.g., 3FS or any POSIX-compatible filesystem).
+
+## Tier Overview
+
+| Tier | Storage Medium | Access Path | Typical Latency |
+|------|---------------|-------------|----------------|
+| **G1** | GPU VRAM (on each client node) | Direct VRAM read/write | < 1 µs |
+| **G2** | CPU DRAM (on each client node) | RDMA or local memcpy | 1–10 µs |
+| **G3** | DFS (shared filesystem) | POSIX / 3FS USRBIO | 100 µs–ms |
+
+Objects flow down the hierarchy as memory pressure increases (G1 → G2 → G3) and are promoted back up on a cache hit.
+
+## Architecture
+
+```
+  Client Node 1                         Shared DFS
+ ┌──────────────────┐                  ┌──────────────┐
+ │ G1: GPU VRAM     │                  │ G3: 3FS /    │
+ │ G2: CPU DRAM     │ ── POSIX/USRBIO──│    NFS /     │
+ └──────────────────┘                  │    GPFS …    │
+                                        └──────────────┘
+  Client Node 2
+ ┌──────────────────┐
+ │ G1: GPU VRAM     │
+ │ G2: CPU DRAM     │
+ └──────────────────┘
+
+   ↑ All clients share G3 via the DFS mount point
+```
+
+The master (`mooncake_master`) manages the G3 segment as a special **DFS segment** registered at startup. Client nodes write KV objects to the DFS when instructed by the master's eviction policy.
+
+## Configuration
+
+### Master Flags
+
+The following flags enable and configure G3 DFS storage (from [Mooncake Store Deployment Guide](mooncake-store-deployment-guide)):
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--root_fs_dir` | — (empty, G3 disabled) | DFS mount directory. When set, the master registers a DFS-backed segment at startup. |
+| `--global_file_segment_size` | `INT64_MAX` | Maximum bytes the DFS segment may occupy. Set to the available DFS capacity. |
+
+### Client Flags
+
+Client nodes do not need additional flags for DFS; the master instructs clients to write to the DFS segment automatically during eviction. Ensure the DFS mount is accessible from every client node at the path specified in `--root_fs_dir`.
+
+## Deployment Example
+
+### Step 1: Mount the DFS on All Nodes
+
+```bash
+# Example: mount 3FS (see docs/source/getting_started/plugin-usage/3FS-USRBIO-Plugin.md)
+mount -t 3fs <3fs_server>:/kvcache /mnt/3fs
+
+# Example: mount NFS
+mount -t nfs <nfs_server>:/kvcache /mnt/kvcache
+```
+
+### Step 2: Start the Master with G3 Enabled
+
+```bash
+mooncake_master \
+    --rpc_port=50051 \
+    --enable_http_metadata_server=true \
+    --http_metadata_server_port=8080 \
+    --root_fs_dir=/mnt/3fs/mooncake \
+    --global_file_segment_size=107374182400   # 100 GB
+```
+
+The master will create the DFS segment directory under `--root_fs_dir` and register it as a G3 segment.
+
+### Step 3: Start Client Nodes (unchanged)
+
+```bash
+mooncake_client \
+    --master_server_address=<master_ip>:50051 \
+    --host=<client_ip> \
+    --protocol=rdma \
+    --device_names=mlx5_0 \
+    --global_segment_size="32GB" \
+    --metadata_server="P2PHANDSHAKE"
+```
+
+### Step 4: Connect the Application
+
+```python
+from mooncake.store import MooncakeDistributedStore
+
+store = MooncakeDistributedStore()
+store.setup(
+    local_hostname="<client_ip>",
+    metadata_server="P2PHANDSHAKE",
+    global_segment_size=32 * 1024 ** 3,
+    local_buffer_size=4 * 1024 ** 3,
+    protocol="rdma",
+    device_name="mlx5_0",
+    master_server_address="<master_ip>:50051",
+)
+```
+
+G3 eviction and promotion are fully transparent to the application.
+
+## 3FS USRBIO Integration
+
+For maximum DFS throughput, Mooncake supports the **3FS USRBIO** (User-space Block IO) interface, which bypasses the kernel page cache and achieves near-NVMe throughput from user space.
+
+To enable:
+1. Build Mooncake with `USE_3FS=ON` and ensure the 3FS USRBIO plugin is installed (see [3FS USRBIO Plugin Guide](../getting_started/plugin-usage/3FS-USRBIO-Plugin)).
+2. Mount 3FS with USRBIO support enabled.
+3. Pass `--root_fs_dir` pointing to the 3FS mount.
+
+When USRBIO is not available, Mooncake falls back to standard POSIX `read`/`write` calls.
+
+## Eviction Policy
+
+The master controls which objects are evicted from G1/G2 to G3 using the same watermark-based policy as SSD offload:
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--eviction_high_watermark_ratio` | `0.95` | Memory usage fraction that triggers eviction |
+| `--eviction_ratio` | `0.05` | Fraction of objects evicted per cycle |
+| `--allow_evict_soft_pinned_objects` | `true` | Whether soft-pinned objects can be evicted to G3 |
+
+Hard-pinned objects are never evicted to G3.
+
+## Performance Tips
+
+- **Place the DFS close to clients**: Use a high-bandwidth interconnect (e.g., InfiniBand / 100GbE) between clients and the DFS storage nodes.
+- **Use 3FS USRBIO** for best throughput when writing large KV cache tensors.
+- **Size `--global_file_segment_size`** conservatively: set it to 80–90 % of actual available DFS capacity to leave room for snapshots and other data.
+- **Monitor G3 usage** via the `/metrics` endpoint: `mooncake_master_dfs_segment_used_bytes`.
+
+## See Also
+
+- [SSD Offload](ssd-offload) — local NVMe offload (single-node)
+- [3FS USRBIO Plugin](../getting_started/plugin-usage/3FS-USRBIO-Plugin) — 3FS setup and USRBIO configuration
+- [Mooncake Store Deployment Guide](mooncake-store-deployment-guide) — all master flags

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/mooncake-pg.md

@@ -0,0 +1,135 @@
+# Mooncake Process Group (mooncake-pg)
+
+## Overview
+
+`mooncake-pg` is a custom PyTorch distributed process-group backend built on top of the Mooncake Transfer Engine. It provides `torch.distributed` collective and point-to-point (P2P) communication primitives that exploit high-speed interconnects (RDMA, NVLink, etc.) while remaining fully compatible with the standard `torch.distributed` API.
+
+It was designed for large-scale expert-parallelism and disaggregated inference scenarios such as those described in the [Kimi K2 deployment blog post](https://lmsys.org/blog/2025-07-20-k2-large-scale-ep/), where Mooncake replaced NCCL for all-to-all expert routing across 128 H200 GPUs.
+
+## Architecture
+
+```
+┌──────────────────────────────────────────────────────────────────┐
+│                       Python / torch.distributed                  │
+│  dist.send / dist.recv / dist.broadcast / dist.all_reduce / …    │
+└───────────────────────────────┬──────────────────────────────────┘
+                                │
+                    ┌───────────▼────────────┐
+                    │    MooncakeBackend     │  (c10d::Backend subclass)
+                    │  (mooncake_backend.h)  │
+                    └───┬──────────┬─────────┘
+                        │          │
+          ┌─────────────▼──┐  ┌───▼────────────────┐
+          │ ConnectionPoller│  │  MooncakeWorker     │
+          │(connection_     │  │  (mooncake_worker.  │
+          │  poller.h)      │  │   cu / .cuh)        │
+          └────────────────┘  └───────┬─────────────┘
+                                      │
+                           ┌──────────▼──────────┐
+                           │    P2PProxy          │
+                           │    (p2p_proxy.h)     │
+                           └──────────┬──────────┘
+                                      │
+                           ┌──────────▼──────────┐
+                           │  TransferEngine      │
+                           │  (RDMA / NVLink /    │
+                           │   TCP …)             │
+                           └─────────────────────┘
+```
+
+### MooncakeBackend
+
+`MooncakeBackend` is registered with PyTorch as two backends:
+- `"mooncake"` — for CUDA tensors.
+- `"mooncake-cpu"` — for CPU tensors.
+
+It subclasses `c10d::Backend` and overrides all collective operations: `send`, `recv`, `broadcast`, `allreduce`, `allgather`, `alltoall`, `scatter`, `reduce`, `gather`, `barrier`, and `batch_isend_irecv`.
+
+Under the hood it maintains:
+- A singleton `TransferEngine` instance shared across all backend instances within the same process, initialized with the RDMA/NVLink device configuration.
+- A `MooncakeWorkerManager` (one per backend instance) that owns a per-GPU CUDA worker thread queue.
+- A `P2PProxy` that serialises and routes send/recv payloads through the Transfer Engine.
+- A `ConnectionPoller` that continuously polls Transfer Engine for completed transfers and signals waiting Work objects.
+
+`MooncakeBackendOptions` carries an `activeRanks_` tensor (a boolean mask) that enables **elastic group membership**: ranks can be added dynamically via `extendGroupSizeTo()`, queried with `getActiveRanks()`, or recovered after a failure with `recoverRanks()`.
+
+### ConnectionPoller
+
+`ConnectionPoller` runs on a dedicated thread. It polls `TransferEngine::getTransferStatus()` for pending batch IDs and sets the `std::atomic<bool> completed` flag on the associated `MooncakeP2PWork` object, which unblocks `Work::wait()` in the calling Python thread.
+
+### MooncakeWorker (CUDA Worker)
+
+Each GPU rank spawns a `MooncakeWorker` CUDA thread (via a CUDA stream / host thread) to execute device-side work items — primarily memory copies involving GPU tensors (`cudaMemcpyAsync`, slicing, casting). This separation keeps the Python GIL thread free during GPU-side operations.
+
+### P2PProxy
+
+`P2PProxy` translates high-level send/recv requests (tensor + rank + tag) into one or more `TransferRequest` entries that the Transfer Engine can process. It handles:
+- Buffer registration (delegated to `TransferEngine::registerLocalMemory`).
+- Multi-slice splitting for large tensors.
+- Tag-based demultiplexing when multiple P2P streams are in flight.
+
+## Supported Operations
+
+| Operation | Supported | Notes |
+|-----------|-----------|-------|
+| `send` / `recv` | ✅ | Single-tensor P2P |
+| `batch_isend_irecv` | ✅ | Async batch P2P |
+| `broadcast` | ✅ | Rank 0 → all |
+| `allreduce` | ✅ | SUM only |
+| `allgather` | ✅ | |
+| `allgather_into_tensor` | ✅ | |
+| `alltoall` | ✅ | Equal-size exchange |
+| `alltoall_base` | ✅ | Variable-size exchange |
+| `scatter` | ✅ | |
+| `gather` | ✅ | |
+| `reduce` | ✅ | SUM only |
+| `barrier` | ✅ | Via allreduce on a dummy tensor |
+| Sparse tensors | ❌ | Not supported |
+| Non-SUM reduce ops | ❌ | Only SUM is implemented |
+
+## Elastic Group Membership
+
+`MooncakeBackend` exposes several extension APIs beyond the standard `c10d::Backend` interface:
+
+| Function | Description |
+|----------|-------------|
+| `getActiveRanks()` | Returns a boolean `torch.Tensor` of size `world_size` |
+| `getNumSyncedRanks()` | Number of currently synced (alive) ranks |
+| `extendGroupSizeTo(size)` | Grow the process group without restart |
+| `getPeerState(ranks)` | Check liveness of a list of ranks |
+| `recoverRanks(ranks)` | Re-admit previously failed ranks |
+| `getPreferredHca(location)` | Query optimal RDMA HCA for a memory location |
+
+These enable fault-tolerant and elastic training scenarios where nodes may join, leave, or fail during a run.
+
+## Build
+
+`mooncake-pg` requires a CUDA-capable build. It is an optional component enabled via the main CMake configuration:
+
+```bash
+mkdir build && cd build
+cmake .. \
+    -DWITH_PG=ON \
+    -DUSE_CUDA=ON \
+    -DCMAKE_BUILD_TYPE=RelWithDebInfo
+make -j$(nproc) mooncake_pg
+```
+
+The build produces `mooncake_pg.cpython-*.so`, which can be installed with pip via the wheel:
+
+```bash
+pip install -e ../mooncake-wheel --no-build-isolation
+```
+
+## Use Cases
+
+- **Expert parallelism (MoE)**: All-to-all expert token routing between GPU ranks using RDMA/NVLink instead of NCCL, as demonstrated in the Kimi K2 deployment.
+- **Pipeline parallelism**: Low-latency activation transfers between pipeline stages.
+- **Disaggregated prefill/decode**: Offloading KV-cache transfers from NCCL to Mooncake's topology-aware engine.
+- **Elastic training**: Dynamic rank management without process restart.
+
+## See Also
+
+- [Mooncake PG Usage Guide](../getting_started/examples/mooncake-pg-usage)
+- [Transfer Engine Architecture](transfer-engine/index)
+- [Kimi K2 deployment blog post](https://lmsys.org/blog/2025-07-20-k2-large-scale-ep/)

docs/source/design/mooncake-store.md

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

docs/source/design/ssd-offload.md

@@ -154,7 +154,7 @@ A single pre-allocated file (`kv_cache.data`) is shared by all objects. Space wi
 
 ## Eviction (BucketStorageBackend)
 
-When `MOONCAKE_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/barex_transport.md

@@ -0,0 +1,81 @@
+# Barex Transport
+
+The Barex transport (`BarexTransport`) is a **bare-metal RDMA extension** that provides an alternative RDMA data path with different queue pair management and flow-control characteristics compared to the standard `RdmaTransport`.
+
+## Overview
+
+While the standard `RdmaTransport` targets high-throughput, multi-NIC environments with endpoint pooling and topology-aware path selection, `BarexTransport` is designed for scenarios that require:
+
+- **Dedicated queue pairs per connection** rather than shared endpoint pools.
+- **Fine-grained flow control** via a countdown-latch mechanism that gates completion acknowledgements.
+- **Simplified connection lifecycle** for bare-metal or HPC environments where connections are long-lived and the overhead of dynamic endpoint management is undesirable.
+
+Both transports use the same ibverbs interface (`infiniband/verbs.h`) and can be compiled into the same binary; the choice is made at runtime via the protocol string.
+
+## Hardware Requirements
+
+- RDMA-capable NIC (InfiniBand or RoCE), same as `RdmaTransport`.
+- `libibverbs` installed on the host.
+
+## Build
+
+```bash
+cmake .. \
+    -DUSE_BAREX=ON \
+    -DCMAKE_BUILD_TYPE=RelWithDebInfo
+make -j$(nproc)
+```
+
+`USE_BAREX` links `barex_transport` into the `transfer_engine` shared library. It can be combined with `USE_CUDA`, `USE_MNNVL`, and other feature flags.
+
+## Usage
+
+### Protocol String
+
+Use `"barex"` as the protocol string:
+
+```python
+from mooncake.engine import TransferEngine
+
+te = TransferEngine()
+te.initialize("node1:12345", "P2PHANDSHAKE", "barex", "mlx5_0")
+```
+
+### When to Use Barex Instead of rdma
+
+| Scenario | Recommended Transport |
+|----------|-----------------------|
+| High-throughput, many concurrent connections, topology-aware routing | `rdma` |
+| Long-lived dedicated connections, HPC / bare-metal, fixed topology | `barex` |
+| AWS EFA | `efa` |
+
+## Internal Design
+
+### BarexContext
+
+Each `BarexContext` object manages the ibverbs resources for a single local RDMA NIC:
+- Protection Domain (`ibv_pd`)
+- Completion Queue (`ibv_cq`)
+- Memory Regions (`ibv_mr`) for each registered buffer
+
+### Queue Pair Management
+
+Unlike `RdmaTransport` which uses endpoint pooling with the SIEVE eviction algorithm, `BarexTransport` allocates a dedicated Queue Pair (QP) per connection pair. QPs are transitioned to the `RTS` (Ready to Send) state during the connection handshake and remain open for the lifetime of the Transfer Engine instance.
+
+### CountDownLatch
+
+`BarexTransport` uses a `CountDownLatch` synchronisation primitive to track in-flight send operations. When a batch of transfer requests is submitted, the latch count is set to the number of outstanding operations. Each CQ completion decrements the count; the initiating thread blocks until the count reaches zero.
+
+This approach simplifies correctness at the cost of some throughput compared to the fully asynchronous pipeline in `RdmaTransport`.
+
+## Limitations
+
+- Not designed for environments with many dynamically appearing / disappearing peers (the static QP model does not scale to thousands of short-lived connections).
+- Does not support topology-aware multi-NIC path selection.
+- Requires `USE_BAREX=ON` at compile time; cannot be selected at runtime if not compiled in.
+
+## See Also
+
+- [Transfer Engine Architecture](index) — standard RDMA transport
+- [Supported Protocols](../../getting_started/supported-protocols)
+- [EFA Transport](efa_transport) — alternative for AWS EFA

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

@@ -0,0 +1,97 @@
+# CXL Transport
+
+The CXL transport (`CxlTransport`) enables memory transfers over **Compute Express Link (CXL)**, a high-speed interconnect that allows CPUs and accelerators to share a unified memory pool across the PCIe bus.
+
+## Overview
+
+CXL memory pooling disaggregates physical DRAM from compute nodes: a CXL memory expander can be mounted into the address space of one or more host CPUs, making remote memory appear as local DRAM. `CxlTransport` leverages this to move data between CXL-attached memory regions across nodes without involving a network NIC.
+
+Typical use cases include:
+- **Memory disaggregation**: Offload KV cache tensors to a large, shared CXL memory pool that multiple inference nodes can access.
+- **Bandwidth aggregation**: Use CXL memory expanders to increase the total memory bandwidth available to a single node.
+
+## Hardware Requirements
+
+- CXL 2.0 or later capable host CPU (e.g., Intel Xeon Scalable 4th Gen / AMD EPYC 9004).
+- CXL memory expander (e.g., Samsung CMM-D, Micron CZ120, Ayar Labs).
+- CXL device must be mounted and visible as a NUMA node (verify with `numactl -H`).
+
+## Build
+
+```bash
+cmake .. \
+    -DUSE_CXL=ON \
+    -DCMAKE_BUILD_TYPE=RelWithDebInfo
+make -j$(nproc)
+```
+
+`USE_CXL` compiles `CxlTransport` into the `transfer_engine` library. No additional runtime library is required beyond the standard Linux kernel CXL driver (enabled by default in Linux 5.15+).
+
+## Usage
+
+### Protocol String
+
+Use `"cxl"` as the protocol string:
+
+```python
+from mooncake.engine import TransferEngine
+
+te = TransferEngine()
+te.initialize("localhost:12345", "P2PHANDSHAKE", "cxl", "")
+```
+
+### Memory Registration
+
+CXL memory regions are registered with a `cpu:N` location tag where `N` is the NUMA node corresponding to the CXL device:
+
+```bash
+# Find the CXL memory NUMA node
+numactl -H | grep -A5 "node distances"
+```
+
+```python
+import ctypes
+
+# Allocate on CXL NUMA node (e.g., node 2)
+buf = ctypes.create_string_buffer(256 * 1024 * 1024)
+te.register_memory(ctypes.addressof(buf), len(buf), "cpu:2")
+```
+
+### Base Address
+
+`CxlTransport` exposes a `getCxlBaseAddr()` method that returns the virtual base address of the mapped CXL region. This is useful when constructing scatter-gather lists that reference CXL memory directly.
+
+### CXL Segment Allocation
+
+`CxlTransport` maintains an internal segment table with a fixed number of local segments (controlled by `allocateLocalSegmentID()`). Each registered CXL buffer occupies one slot. The segment limit is set at compile time; contact the Mooncake team if you need to increase it.
+
+## Limitations
+
+- CXL memory regions must be on the same host or connected through a CXL switch; this transport does not use a network NIC.
+- CXL 1.1 devices (non-pooled) appear as standard DRAM NUMA nodes — no special transport is needed for them.
+- CXL 2.0 pooling (shared memory across multiple hosts) requires a CXL switch and OS-level support (Linux 6.6+).
+
+## Troubleshooting
+
+### CXL device not visible as NUMA node
+
+```bash
+ls /sys/bus/cxl/devices/
+daxctl list
+```
+
+If the device appears as a DAX device but not a NUMA node, use `daxctl reconfigure-device` to online it as system RAM.
+
+### Permission denied on CXL device file
+
+```bash
+sudo chmod 660 /dev/dax0.0
+sudo chown root:$(id -gn) /dev/dax0.0
+```
+
+## See Also
+
+- [Transfer Engine Architecture](index)
+- [Supported Protocols](../../getting_started/supported-protocols)
+- [SSD Offload](../../deployment/ssd-offload) — multi-tier storage with NVMe
+- [Multi-Tier Storage](../../deployment/multi-tier-storage) — G1/G2/G3 tier configuration

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

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

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

@@ -0,0 +1,81 @@
+# HIP Transport (AMD ROCm)
+
+The HIP transport (`HipTransport`) enables GPU-to-GPU data transfers on **AMD ROCm** platforms, using either IPC handles or Shareable handles for intra-node transfers between AMD GPUs.
+
+## Overview
+
+`HipTransport` is the AMD equivalent of the CUDA-based transfer paths. It is designed for **intra-node** GPU communication: moving data between AMD GPU VRAM buffers on the same host without routing through system memory or a network NIC.
+
+The transport uses the HIP runtime (`hip_runtime.h`) and supports two handle types for mapping peer GPU memory:
+
+| Handle Type | Use Case |
+|-------------|----------|
+| **IPC handle** (`hipIpcMemHandle_t`) | Mapping GPU memory from another process on the same host |
+| **Shareable handle** | Used when IPC handles are unavailable or for cross-device mappings |
+
+## Hardware Requirements
+
+- AMD GPU with ROCm 5.0 or later (e.g., MI200 / MI300 series).
+- Peer access must be supported between the source and destination GPUs (`hipDeviceCanAccessPeer`).
+
+## Build
+
+```bash
+cmake .. \
+    -DUSE_HIP=ON \
+    -DCMAKE_BUILD_TYPE=RelWithDebInfo
+make -j$(nproc)
+```
+
+The `USE_HIP` flag enables HIPification of the Transfer Engine sources (`hipify_files`) and links `hip::host` and the ROCm runtime library.
+
+> **Note:** `USE_HIP` and `USE_CUDA` are mutually exclusive. Do not enable both in the same build.
+
+## Usage
+
+### Protocol String
+
+Use `"hip"` as the protocol string:
+
+```python
+from mooncake.engine import TransferEngine
+
+te = TransferEngine()
+te.initialize("localhost:12345", "P2PHANDSHAKE", "hip", "")
+```
+
+### Memory Registration
+
+Register AMD GPU memory using the standard `TransferEngine::registerLocalMemory` API. Specify the location as `"cuda:N"` (the HIP transport reuses the same location tag convention as the CUDA transport, since HIP mirrors the CUDA device numbering):
+
+```python
+import ctypes, mooncake
+
+te = TransferEngine()
+te.initialize("localhost:12345", "P2PHANDSHAKE", "hip", "")
+
+# Allocate 256 MB on GPU 0 (HIP device 0)
+buf = mooncake.allocate_gpu_memory(256 * 1024 * 1024, device=0)
+te.register_memory(buf, 256 * 1024 * 1024, "cuda:0")
+```
+
+### Environment Variables
+
+`HipTransport` respects the same topology and configuration environment variables as the RDMA transport:
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `MC_MS_AUTO_DISC` | `1` | Auto-discover GPU topology |
+| `MC_MS_FILTERS` | — | NIC / device whitelist |
+
+## Limitations
+
+- Intra-node only: `HipTransport` does not support transfers across network links. Combine with `"rdma"` (via the standard RDMA transport) for cross-node transfers in a heterogeneous setup.
+- IPC handles require both processes to be on the same OS instance and the same AMD GPU driver version.
+- Inter-GPU peer access must be enabled at the OS / driver level.
+
+## See Also
+
+- [Supported Protocols](../../getting_started/supported-protocols) — protocol selection guide
+- [Transfer Engine Architecture](index)
+- [ROCm Documentation](https://rocm.docs.amd.com/)

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

@@ -330,6 +330,20 @@ ascend_transport
 heterogeneous_ascend
 :::
 
+## Accelerator and Specialized Transports
+
+:::{toctree}
+:maxdepth: 1
+
+nvlink_transport
+nvlink_intra_transport
+hip_transport
+cxl_transport
+barex_transport
+mlu_transport
+maca_transport
+:::
+
 ## Benchmark and Tuning Guide
 
 :::{toctree}

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/design/transfer-engine/maca_transport.md

@@ -0,0 +1,131 @@
+# MACA Transport (MetaX / Muxi)
+
+Mooncake Transfer Engine supports **MetaX (Muxi) MACA** GPU accelerators. MACA (MetaX Architecture for Computing Acceleration) is the GPU compute platform from MetaX Integrated Circuits.
+
+## Overview
+
+MACA support in Mooncake mirrors the CUDA code path: memory allocations on MetaX GPUs are registered with the Transfer Engine and transferred via the RDMA NIC that is closest to the device in the PCIe topology. MACA-specific runtime libraries replace the CUDA runtime.
+
+## Hardware Requirements
+
+- MetaX (Muxi) GPU accelerator.
+- MACA runtime installed at `MACA_HOME` (default `/opt/maca`).
+- RDMA-capable NIC in the same server.
+
+## Build
+
+```bash
+cmake .. \
+    -DUSE_MACA=ON \
+    -DCMAKE_BUILD_TYPE=RelWithDebInfo
+make -j$(nproc)
+```
+
+### Environment Variables (build-time)
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `MACA_HOME` | `/opt/maca` | Root of the MACA installation (headers and libraries) |
+
+CMake resolves `MACA_INCLUDE_DIR` as `${MACA_HOME}/include`.
+
+### Runtime Libraries
+
+By default, `USE_MACA=ON` links:
+
+```cmake
+mcruntime      # MACA compute runtime
+mxc-runtime64  # MetaX cross-platform runtime
+rt             # POSIX real-time library
+```
+
+Override the library list at configure time if your MACA installation uses different names:
+
+```bash
+cmake .. \
+    -DUSE_MACA=ON \
+    -DMACA_RUNTIME_LIBS="mcruntime;mxc-runtime64;rt"
+```
+
+## Usage
+
+### Protocol String
+
+Use `"rdma"` as the protocol string. MACA support extends the RDMA path with MACA-aware memory registration:
+
+```python
+from mooncake.engine import TransferEngine
+
+te = TransferEngine()
+te.initialize("node1:12345", "P2PHANDSHAKE", "rdma", "mlx5_0")
+```
+
+Register MACA GPU memory:
+
+```python
+# Allocate MACA GPU memory via mcMalloc (or framework allocator)
+maca_ptr = mc_malloc(256 * 1024 * 1024)
+
+# Register with Transfer Engine using "cuda:N" location tag
+# (MACA devices follow the same device-numbering convention as CUDA)
+te.register_memory(maca_ptr, 256 * 1024 * 1024, "cuda:0")
+```
+
+### Topology Discovery
+
+When `USE_MACA=ON`, the topology discovery module enumerates MetaX GPUs via the MACA runtime and maps them to their nearest RDMA NIC(s). The resulting `priority_matrix` entries use `"cuda:N"` location tags (shared convention with CUDA devices).
+
+## Environment Variables (run-time)
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `MC_MS_AUTO_DISC` | `1` | Auto-discover MACA GPU topology. Set `0` to disable. |
+| `MC_MS_FILTERS` | — | Comma-separated NIC whitelist |
+| `MACA_HOME` | `/opt/maca` | Used at runtime if the MACA shared libraries are not on `LD_LIBRARY_PATH` |
+
+Ensure the MACA runtime libraries are on `LD_LIBRARY_PATH`:
+
+```bash
+export LD_LIBRARY_PATH=$MACA_HOME/lib:$LD_LIBRARY_PATH
+```
+
+## Build with Both MACA and Other Features
+
+`USE_MACA=ON` can be combined with other flags such as `USE_EFA`, `USE_CXL`, or `USE_BAREX`. It cannot be combined with `USE_CUDA=ON` (both define the same CUDA-like runtime aliases).
+
+## Limitations
+
+- Cross-node MACA GPU transfers require an RDMA NIC; direct MACA peer-to-peer transfers between nodes are not supported.
+- MACA and CUDA cannot be enabled in the same build.
+
+## Troubleshooting
+
+### MACA runtime library not found
+
+```
+error while loading shared libraries: libmcruntime.so
+```
+
+Add the MACA library directory to `LD_LIBRARY_PATH`:
+
+```bash
+export LD_LIBRARY_PATH=/opt/maca/lib:$LD_LIBRARY_PATH
+```
+
+### Incorrect `MACA_HOME`
+
+```
+CMake Error: MACA_INCLUDE_DIR not found
+```
+
+Set `MACA_HOME` explicitly:
+
+```bash
+cmake .. -DUSE_MACA=ON -DMACA_HOME=/path/to/maca
+```
+
+## See Also
+
+- [Transfer Engine Architecture](index) — RDMA transport and topology-aware path selection
+- [Supported Protocols](../../getting_started/supported-protocols)
+- [Build Guide](../../getting_started/build) — `USE_MACA` flag in the build matrix

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

@@ -0,0 +1,121 @@
+# MLU Transport (Cambricon)
+
+Mooncake Transfer Engine supports **Cambricon MLU** (Machine Learning Unit) accelerators through an extension to the standard `RdmaTransport`. There is no separate `mlu` protocol string; MLU support adds MLU-aware memory registration and topology discovery on top of the normal RDMA data path.
+
+## Overview
+
+Cambricon MLU devices expose their on-device memory via the **DMA-BUF** kernel interface, the same mechanism used by NVIDIA GPUDirect RDMA. When `USE_MLU=ON` is enabled, Mooncake:
+
+1. Detects MLU devices during topology discovery.
+2. Registers MLU VRAM buffers with the RDMA NIC using DMA-BUF handles.
+3. Performs RDMA read/write operations directly to/from MLU VRAM — no intermediate CPU bounce buffer is needed.
+
+The result is a zero-copy transfer path between MLU device memory and remote DRAM / another MLU.
+
+## Hardware Requirements
+
+- Cambricon MLU accelerator (MLU370, MLU590, etc.).
+- Cambricon Neuware SDK installed (`neuware` package).
+- RDMA-capable NIC (InfiniBand or RoCE) in the same server.
+- Linux kernel with DMA-BUF support (5.6+).
+
+## Build
+
+```bash
+cmake .. \
+    -DUSE_MLU=ON \
+    -DUSE_CUDA=OFF \
+    -DCMAKE_BUILD_TYPE=RelWithDebInfo
+make -j$(nproc)
+```
+
+### Environment Variables (build-time)
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `NEUWARE_HOME` | `/usr/local/neuware` | Path to the Neuware installation |
+| `NEUWARE_ROOT` | — | Alternative Neuware root (checked if `NEUWARE_HOME` is unset) |
+
+The CMake scripts use these to locate `libcnrt`, `libcndrv`, and the MLU headers.
+
+### Runtime Libraries
+
+`USE_MLU=ON` links the following Neuware libraries into `transfer_engine`:
+
+| Library | Purpose |
+|---------|---------|
+| `cnrt` | Cambricon Runtime (memory allocation, device management) |
+| `cndrv` | Cambricon Driver (low-level device access, DMA-BUF export) |
+
+## Usage
+
+### Protocol String
+
+Use the standard `"rdma"` protocol string. MLU support is transparent at the API level:
+
+```python
+from mooncake.engine import TransferEngine
+
+te = TransferEngine()
+te.initialize("node1:12345", "P2PHANDSHAKE", "rdma", "mlx5_0")
+```
+
+To register MLU memory:
+
+```python
+import ctypes
+
+# Allocate MLU memory via cnrt (or use a framework allocator)
+mlu_ptr = cnrt_malloc(256 * 1024 * 1024)
+
+# Register with Transfer Engine — specify the MLU device location
+te.register_memory(mlu_ptr, 256 * 1024 * 1024, "mlu:0")
+```
+
+The location tag `"mlu:N"` tells the topology engine to associate the buffer with MLU device N and select an RDMA NIC that has direct PCIe connectivity to that device.
+
+### Topology Discovery
+
+When `USE_MLU=ON` and MLU devices are present, the Transfer Engine topology module:
+- Enumerates all MLU devices and their PCIe BDF addresses.
+- Maps each MLU to the nearest RDMA NIC(s) based on PCIe topology.
+- Populates the `priority_matrix` with `"mlu:N"` location entries.
+
+This ensures that DMA-BUF transfers use the NIC closest to the MLU, minimising PCIe switch hops.
+
+## Limitations
+
+- MLU support uses the RDMA data path; it does not add a dedicated protocol endpoint.
+- Cross-node MLU-to-MLU transfers require an RDMA NIC on each node.
+- Intra-node MLU-to-MLU transfers use the RDMA loopback path (not a direct NVLink / PCIe peer copy); direct peer-to-peer MLU copy is not yet implemented in Mooncake.
+- `USE_MLU=ON` and `USE_CUDA=ON` can coexist in the same build (the DMA-BUF registration code is independent), but has not been tested extensively in mixed-accelerator environments.
+
+## Troubleshooting
+
+### DMA-BUF export fails
+
+```
+cnDrvMemExportToDmaBuf: CNDRV_ERROR_NO_SUPPORT
+```
+
+Ensure the Cambricon driver version supports DMA-BUF export (driver ≥ 4.9) and that the kernel DMA-BUF interface is enabled:
+
+```bash
+zcat /proc/config.gz | grep DMA_BUF
+# CONFIG_DMA_BUF=y
+```
+
+### MLU not detected in topology
+
+```bash
+cnrt-diagnose   # should list MLU devices
+lspci | grep Cambricon
+```
+
+Verify that `NEUWARE_HOME` points to the correct Neuware installation.
+
+## See Also
+
+- [Transfer Engine Architecture](index) — RDMA transport and topology-aware path selection
+- [Supported Protocols](../../getting_started/supported-protocols) — `rdma` protocol usage with MLU
+- [Cambricon Neuware Documentation](https://developer.cambricon.com)

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

@@ -0,0 +1,68 @@
+# Intra-Node NVLink Transport
+
+The Intra-Node NVLink transport (`IntraNodeNvlinkTransport`) enables **zero-copy GPU-to-GPU data transfers within a single host** using the NVLink high-speed interconnect, without going through the PCIe bus or system memory.
+
+## Overview
+
+Modern NVIDIA servers (DGX A100/H100/H200) connect all GPUs on the same node via NVLink. For within-node transfers this provides substantially higher bandwidth and lower latency than PCIe-based copies or RDMA loopback.
+
+`IntraNodeNvlinkTransport` uses CUDA IPC handles and the **UBShmem Fabric Allocator** to map peer GPU memory directly into the local GPU's address space, then issues `cudaMemcpyAsync` on the NVLink path.
+
+## Hardware Requirements
+
+- NVIDIA GPUs with intra-node NVLink (e.g., A100 / H100 / H200 SXM variants in DGX systems).
+- CUDA 11.0+ with peer access support (`cudaDeviceEnablePeerAccess`).
+
+## Build
+
+```bash
+cmake .. \
+    -DUSE_INTRA_NVLINK=ON \
+    -DUSE_CUDA=ON \
+    -DCMAKE_BUILD_TYPE=RelWithDebInfo
+make -j$(nproc)
+```
+
+`USE_INTRA_NVLINK` enables the `IntraNodeNvlinkTransport` class in the `transfer_engine` library.
+
+## Usage
+
+### Protocol String
+
+Use `"nvlink_intra"` as the protocol string:
+
+```python
+from mooncake.engine import TransferEngine
+
+te = TransferEngine()
+te.initialize("localhost:12345", "P2PHANDSHAKE", "nvlink_intra", "cuda:0")
+```
+
+### Automatic Selection
+
+When `MooncakeDistributedStore` is used with `protocol="rdma"` and both the source and destination are on the same node, the topology-aware path selector may automatically prefer the intra-node NVLink path over RDMA when `USE_INTRA_NVLINK` is enabled and peer access is available.
+
+## Memory Registration
+
+`IntraNodeNvlinkTransport` uses the **UBShmem Fabric Allocator**, which:
+1. Creates a shareable CUDA IPC memory handle via `cudaIpcGetMemHandle`.
+2. Exports the handle to other CUDA contexts on the same host via a shared-memory segment.
+3. Maps the remote handle into the local GPU's address space with `cudaIpcOpenMemHandle`.
+
+Only memory in the VRAM of GPUs that are NVLink-connected can be mapped this way. CPU (DRAM) transfers fall through to the standard memcpy path.
+
+## Comparison with Inter-Node NVLink
+
+| Feature | Intra-Node (`nvlink_intra`) | Inter-Node (`nvlink`) |
+|---------|-----------------------------|-----------------------|
+| Scope | Single host, multiple GPUs | Multiple hosts (MNNVL fabric) |
+| Build flag | `USE_INTRA_NVLINK=ON` | `USE_MNNVL=ON` |
+| Memory mechanism | CUDA IPC handles | CUDA Virtual Memory Management + MNNVL fabric export |
+| Bandwidth | Full NVLink bandwidth (e.g., 900 GB/s total on H100 SXM) | Full MNNVL bandwidth |
+| Hardware needed | Any DGX / server with NVLink switch | MNNVL-capable nodes |
+
+## See Also
+
+- [NVLink Transport (MNNVL / Inter-Node)](nvlink_transport) — cross-node NVLink transfers
+- [Supported Protocols](../../getting_started/supported-protocols) — protocol selection guide
+- [Transfer Engine Architecture](index)

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

@@ -0,0 +1,83 @@
+# NVLink Transport (MNNVL / Inter-Node)
+
+The NVLink transport (`NvlinkTransport`) enables high-bandwidth, low-latency GPU-to-GPU data transfers across nodes using **NVIDIA Multi-Node NVLink (MNNVL)**. It bypasses the PCIe bus and the network stack entirely, delivering bandwidth that approaches the raw NVLink fabric speed.
+
+## Overview
+
+MNNVL is available on NVIDIA platforms that physically link GPUs across multiple nodes via NVLink cables (e.g., DGX SuperPOD / DGX GB200 NVL72). Within such a fabric, all GPUs share a unified address space exposed through the NVLink Allocator, making remote GPU memory directly addressable from any node in the fabric.
+
+`NvlinkTransport` registers GPU memory into the NVLink fabric via `cuMemCreate` / `cuMemMap` (CUDA Virtual Memory Management) and performs transfers through direct remote memory writes — no RDMA NIC is involved.
+
+## Hardware Requirements
+
+- NVIDIA MNNVL-capable hardware (e.g., H100 / H200 / B200 in a NVL72 configuration).
+- CUDA 12.0+ with NVLink fabric support (`cuMemGetAllocationGranularity` with `CU_MEM_ALLOC_GRANULARITY_RECOMMENDED`).
+- NVLink fabric initialised by the system firmware before the application starts.
+
+## Build
+
+Enable MNNVL support at CMake configure time:
+
+```bash
+cmake .. \
+    -DUSE_MNNVL=ON \
+    -DUSE_CUDA=ON \
+    -DCMAKE_BUILD_TYPE=RelWithDebInfo
+make -j$(nproc)
+```
+
+The `USE_MNNVL` flag compiles `NvlinkTransport` into the `transfer_engine` library.
+
+## Usage
+
+### Protocol String
+
+Use `"nvlink"` as the protocol string when initialising the Transfer Engine:
+
+```python
+from mooncake.engine import TransferEngine
+
+te = TransferEngine()
+te.initialize("node1:12345", "P2PHANDSHAKE", "nvlink", "")
+```
+
+### Environment Variables
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `MC_FORCE_MNNVL` | `false` | Force NVLink even when RDMA NICs are present. When RDMA HCAs are detected and `MC_FORCE_MNNVL` is not set, the engine prefers RDMA. Set `MC_FORCE_MNNVL=true` to override. |
+
+### When to Use
+
+- Use `"nvlink"` only when the deployment hardware has MNNVL connectivity and you want maximum GPU-to-GPU bandwidth between nodes.
+- On clusters without MNNVL hardware, fall back to `"rdma"` or `"tcp"`.
+- For **intra-node** NVLink transfers (within a single host), use `"nvlink_intra"` instead (see [Intra-Node NVLink Transport](nvlink_intra_transport)).
+
+## Memory Registration
+
+`NvlinkTransport` uses a custom **NVLink Allocator** that:
+1. Allocates GPU memory via CUDA Virtual Memory Management (`cuMemCreate`).
+2. Exports the allocation handle through the NVLink fabric (`cuMulticastAddMemory` / `cuMemExportToShareableHandle`).
+3. Registers the allocation with the Transfer Engine metadata service so remote peers can map it.
+
+All memory registered through `TransferEngine::registerLocalMemory` with a `cuda:N` location tag is automatically handled by the NVLink Allocator when `NvlinkTransport` is active.
+
+## Tuning Tips
+
+- **Buffer alignment**: Allocate buffers in multiples of the NVLink granularity returned by `cuMemGetAllocationGranularity`. Misaligned allocations will be rounded up internally.
+- **Pinning**: Keep GPU buffers pinned for the lifetime of a Transfer Engine session; frequent re-registration degrades performance.
+- **Topology**: MNNVL works best when the NVLink fabric is fully populated. Partial fabric configurations will fall back to slower paths for non-adjacent GPUs.
+
+## Test
+
+```bash
+./build/mooncake-transfer-engine/tests/nvlink_transport_test
+```
+
+The test suite requires MNNVL hardware. On systems without MNNVL, the test will skip automatically.
+
+## See Also
+
+- [Intra-Node NVLink Transport](nvlink_intra_transport) — within-node NVLink transfers
+- [Supported Protocols](../../getting_started/supported-protocols) — protocol selection guide
+- [Transfer Engine Architecture](index)

docs/source/getting_started/examples/mooncake-pg-usage.md

@@ -0,0 +1,253 @@
+# Mooncake Process Group — Usage Guide
+
+This guide shows how to build `mooncake-pg`, register it with PyTorch distributed, and use it as a drop-in replacement for NCCL in collective and point-to-point workloads.
+
+## Prerequisites
+
+| Requirement | Version |
+|-------------|---------|
+| CUDA | 11.8 or later |
+| PyTorch | 2.1 or later |
+| RDMA hardware | InfiniBand / RoCE / NVLink (TCP fallback available) |
+| Mooncake | Built with `-DWITH_PG=ON -DUSE_CUDA=ON` |
+
+## Installation
+
+### Step 1: Build Mooncake with PG support
+
+```bash
+git clone https://github.com/kvcache-ai/Mooncake.git
+cd Mooncake
+git submodule update --init --recursive
+
+mkdir build && cd build
+cmake .. \
+    -DWITH_PG=ON \
+    -DUSE_CUDA=ON \
+    -DCMAKE_BUILD_TYPE=RelWithDebInfo
+make -j$(nproc)
+```
+
+### Step 2: Install the Python wheel
+
+```bash
+# Copy built extension into the wheel directory
+cp mooncake-pg/mooncake_pg.cpython-*.so ../mooncake-wheel/mooncake/
+
+# Install with pip
+cd ..
+pip install -e mooncake-wheel --no-build-isolation
+```
+
+### Step 3: Verify the import
+
+```python
+import mooncake_pg   # no error = installed correctly
+import torch.distributed as dist
+```
+
+## Initializing the Process Group
+
+`mooncake-pg` registers two backends with PyTorch:
+- `"mooncake"` for CUDA tensors.
+- `"mooncake-cpu"` for CPU tensors.
+
+Use the standard `torch.distributed.init_process_group` call:
+
+```python
+import torch
+import torch.distributed as dist
+
+# Standard torchrun / mpirun launch sets MASTER_ADDR, MASTER_PORT,
+# RANK, WORLD_SIZE automatically.
+dist.init_process_group(
+    backend="mooncake",   # or "mooncake-cpu" for CPU-only
+    init_method="env://",
+)
+
+rank = dist.get_rank()
+world_size = dist.get_world_size()
+print(f"[rank {rank}/{world_size}] process group initialized")
+```
+
+The process group creation triggers `MooncakeBackend` construction, which:
+1. Initialises the shared `TransferEngine` singleton (first call only, subsequent calls reuse it).
+2. Registers local GPU memory buffers.
+3. Starts the `ConnectionPoller` background thread.
+
+## Example: All-Reduce
+
+```python
+import torch
+import torch.distributed as dist
+
+dist.init_process_group(backend="mooncake", init_method="env://")
+rank = dist.get_rank()
+
+# Create a tensor on the local GPU
+device = torch.device(f"cuda:{rank}")
+tensor = torch.ones(1024, device=device) * rank
+
+# In-place allreduce (SUM)
+dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+
+# On rank 0 with world_size=4: tensor == [0+1+2+3] * 1024 == 6144
+print(f"[rank {rank}] allreduce result[0] = {tensor[0].item()}")
+dist.destroy_process_group()
+```
+
+## Example: Point-to-Point (P2P)
+
+```python
+import torch
+import torch.distributed as dist
+
+dist.init_process_group(backend="mooncake", init_method="env://")
+rank = dist.get_rank()
+world_size = dist.get_world_size()
+
+device = torch.device(f"cuda:{rank}")
+N = 1024 * 1024  # 4 MB
+
+if rank == 0:
+    tensor = torch.arange(N, dtype=torch.float32, device=device)
+    dist.send(tensor, dst=1, tag=42)
+    print("[rank 0] sent tensor")
+elif rank == 1:
+    tensor = torch.empty(N, dtype=torch.float32, device=device)
+    dist.recv(tensor, src=0, tag=42)
+    print(f"[rank 1] received tensor[0]={tensor[0].item()}")
+
+dist.destroy_process_group()
+```
+
+## Example: All-to-All (Expert Parallelism)
+
+This pattern is the primary use case for `mooncake-pg` in MoE (Mixture-of-Experts) models, where tokens are routed to expert GPUs.
+
+```python
+import torch
+import torch.distributed as dist
+
+dist.init_process_group(backend="mooncake", init_method="env://")
+rank = dist.get_rank()
+world_size = dist.get_world_size()
+
+device = torch.device(f"cuda:{rank}")
+tokens_per_rank = 128
+hidden_dim = 4096
+
+# Each rank sends a different chunk to every other rank
+input_tensor = torch.randn(world_size * tokens_per_rank, hidden_dim, device=device)
+output_tensor = torch.empty_like(input_tensor)
+
+dist.all_to_all_single(output_tensor, input_tensor)
+print(f"[rank {rank}] all_to_all complete, output shape={output_tensor.shape}")
+
+dist.destroy_process_group()
+```
+
+## Example: Batch Async P2P (batch_isend_irecv)
+
+```python
+import torch
+import torch.distributed as dist
+
+dist.init_process_group(backend="mooncake", init_method="env://")
+rank = dist.get_rank()
+device = torch.device(f"cuda:{rank}")
+N = 1024
+
+peer = (rank + 1) % dist.get_world_size()
+send_tensor = torch.ones(N, device=device) * rank
+recv_tensor = torch.zeros(N, device=device)
+
+ops = [
+    dist.P2POp(dist.isend, send_tensor, peer),
+    dist.P2POp(dist.irecv, recv_tensor, peer),
+]
+works = dist.batch_isend_irecv(ops)
+for w in works:
+    w.wait()
+
+print(f"[rank {rank}] received from {peer}: {recv_tensor[0].item()}")
+dist.destroy_process_group()
+```
+
+## Launch Commands
+
+### torchrun (recommended)
+
+```bash
+# 4 GPUs on 1 node
+torchrun \
+    --nproc_per_node=4 \
+    --master_addr=127.0.0.1 \
+    --master_port=29500 \
+    your_script.py
+
+# 16 GPUs across 2 nodes (run on each node)
+torchrun \
+    --nproc_per_node=8 \
+    --nnodes=2 \
+    --node_rank=0 \          # 1 on second node
+    --master_addr=10.0.0.1 \
+    --master_port=29500 \
+    your_script.py
+```
+
+### mpirun
+
+```bash
+mpirun -np 8 \
+    -x MASTER_ADDR=10.0.0.1 \
+    -x MASTER_PORT=29500 \
+    python your_script.py
+```
+
+## Elastic Group Membership
+
+`mooncake-pg` exposes Python helpers for fault-tolerant setups via `mooncake_pg`:
+
+```python
+import mooncake_pg
+import torch.distributed as dist
+
+dist.init_process_group(backend="mooncake", init_method="env://")
+pg = dist.group.WORLD
+
+# Check which ranks are alive
+active = mooncake_pg.get_active_ranks(pg)           # torch.BoolTensor
+n_synced = mooncake_pg.get_num_synced_ranks(pg)      # int
+
+# Dynamically grow the group
+mooncake_pg.extend_group_size_to(pg, new_size=16)
+
+# Check specific peer liveness
+states = mooncake_pg.get_peer_state(pg, [2, 5, 7])  # list[bool]
+
+# Re-admit a recovered rank
+mooncake_pg.recover_ranks(pg, [5])
+```
+
+## Environment Variables
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `MOONCAKE_MASTER` | — | Mooncake master address (`ip:port`) |
+| `MOONCAKE_PROTOCOL` | `rdma` | Transfer protocol (`rdma`, `tcp`, `nvlink`, …) |
+| `MOONCAKE_DEVICE` | auto | RDMA/HCA device name (e.g., `mlx5_0`) |
+| `MC_METADATA_SERVER` | `P2PHANDSHAKE` | Metadata server URL |
+| `MC_FORCE_MNNVL` | `false` | Force MNNVL even when RDMA NICs are present |
+
+## Limitations
+
+- Only `SUM` is supported for reduce operations (`allreduce`, `reduce`).
+- Sparse tensors are not supported.
+- Each `send`/`recv` call must transfer a single tensor; use `batch_isend_irecv` for multiple tensors.
+
+## See Also
+
+- [Mooncake PG Design](../../design/mooncake-pg)
+- [Transfer Engine](../../design/transfer-engine/index)
+- [Supported Protocols](../supported-protocols)

docs/source/getting_started/examples/rl-training-integration.md

@@ -0,0 +1,156 @@
+# Reinforcement Learning Training Integration
+
+`mooncake-rl` demonstrates how to use `MooncakeDistributedStore` as a zero-copy data bus between **rollout engines** (inference workers) and **training engines** (gradient workers) in a distributed reinforcement-learning setup.
+
+The key idea is to decouple the rollout and training phases: rollout engines write experience batches into Mooncake Store under a well-known key, and training engines fetch those batches concurrently — without any explicit synchronisation barrier or shared memory.
+
+## Motivation
+
+In a typical RL-from-human-feedback (RLHF) or online RL pipeline:
+
+```
+ Rollout Engines                    Training Engines
+ (inference, e.g. SGLang/vLLM)      (gradient step, e.g. PyTorch)
+ ┌──────────────────┐                ┌─────────────────────┐
+ │  generate()      │ ─── RDMA ───▶  │  train(key)         │
+ │  put_tensor(key) │                │  get_tensor(key)    │
+ └──────────────────┘                └─────────────────────┘
+```
+
+Data flows over RDMA (or TCP fallback) via `MooncakeDistributedStore`, bypassing the CPU for GPU-to-GPU transfers and avoiding Python serialisation overhead.
+
+## Example: `rl_samples.py`
+
+The reference implementation lives in [`mooncake-rl/examples/rl_samples.py`](../../../../mooncake-rl/examples/rl_samples.py). It provides a minimal, runnable mock of the full RL loop:
+
+| Class | Role |
+|-------|------|
+| `RolloutEngine` | Generates random `(obs, action, reward)` samples |
+| `RolloutController` | Manages dataset state and connects to Mooncake Store |
+| `RolloutManager` | Orchestrates rollout engines; writes samples to the Store |
+| `TrainActor` | Consumes samples and performs a dummy training step |
+| `TrainGroup` | Coordinates training actors; reads samples from the Store |
+
+### Data Flow Walk-Through
+
+1. **RolloutManager.generate(rollout_id)** — each rollout engine produces a sample dict `{obs, action, reward}`. All samples for this rollout step are assembled into a list and written into the Store:
+
+   ```python
+   # RolloutController uses an RDMA-initialised Store client
+   self.rollout_client.put_tensor(str(rollout_id), rollout_samples)
+   ```
+
+2. **TrainGroup.train(rollout_id, key)** — the training side fetches those samples and distributes them across training actors:
+
+   ```python
+   samples = self.training_client.get_tensor(rollout_key)
+   for actor, sample in zip(self.actor_handlers, samples):
+       actor.train(sample)
+   ```
+
+3. **Training loop** — the top-level `train()` function alternates generation, training, evaluation, and checkpoint saving:
+
+   ```python
+   for rollout_id in range(start, num_rollout):
+       key = rollout_manager.generate(rollout_id)  # write to Store
+       actor_model.train(rollout_id, key)           # read from Store
+       actor_model.update_weights()                 # sync weights to rollout
+   ```
+
+## Running the Example
+
+### Prerequisites
+
+- A running Mooncake metadata server (e.g., the HTTP server or etcd).
+- A running `mooncake_master` (for distributed Store coordination).
+- Mooncake Python wheel installed (`pip install mooncake-transfer-engine`).
+
+### Start the Metadata Server
+
+```bash
+cd mooncake-transfer-engine/example/http-metadata-server-python
+pip install aiohttp
+python bootstrap_server.py &
+```
+
+### Start the Master
+
+```bash
+./build/mooncake_master \
+    --rpc_port=50051 \
+    --enable_http_metadata_server=true \
+    --http_metadata_server_host=0.0.0.0 \
+    --http_metadata_server_port=8080
+```
+
+### Run the RL Example
+
+```bash
+python mooncake-rl/examples/rl_samples.py \
+    --num_rollout=10 \
+    --num_train_actor=2 \
+    --num_rollout_actor=2 \
+    --save_interval=5 \
+    --eval_interval=5 \
+    --model_path=./checkpoints
+```
+
+## Connecting to a Real Metadata Server
+
+The example initialises two `MooncakeDistributedStore` clients — one on the rollout side and one on the training side — with different local hostnames and NIC device names:
+
+```python
+# Rollout client (on rollout node / process)
+self.rollout_client = MooncakeDistributedStore()
+self.rollout_client.setup(
+    "localhost:12346",                     # local_hostname:port
+    "http://localhost:8080/metadata",      # metadata_server URL
+    512 * 1024 * 1024,                     # global_segment_size (512 MB)
+    128 * 1024 * 1024,                     # local_buffer_size (128 MB)
+    "rdma",                                # protocol
+    "erdma_0",                             # NIC device (e.g. mlx5_0)
+    "localhost:50051",                     # mooncake_master address
+)
+
+# Training client (on training node / process)
+self.training_client = MooncakeDistributedStore()
+self.training_client.setup(
+    "localhost:12345",
+    "http://localhost:8080/metadata",
+    512 * 1024 * 1024,
+    128 * 1024 * 1024,
+    "rdma",
+    "erdma_1",
+    "localhost:50051",
+)
+```
+
+In a real deployment:
+- Replace `localhost` with actual hostnames or IPs.
+- Replace `erdma_0` / `erdma_1` with the NIC names from `ibv_devices`.
+- Use separate metadata server and master addresses accessible by both sets of nodes.
+
+## Key/Value Conventions
+
+By default, the example uses the rollout ID as the string key:
+
+```python
+key = str(rollout_id)               # e.g. "0", "1", "42"
+self.rollout_client.put_tensor(key, rollout_samples)
+```
+
+In production, use a structured key scheme to avoid collisions across concurrent runs:
+
+```python
+key = f"run:{run_id}/rollout:{rollout_id}"
+```
+
+## Integration with Real RL Frameworks
+
+This pattern is directly applicable to frameworks such as [THUDM/slime](https://github.com/THUDM/slime), [veRL](https://github.com/volcengine/verl), and [OpenRLHF](https://github.com/OpenRLHF/OpenRLHF). Replace the dummy `generate()` and `train()` implementations with real policy inference (SGLang / vLLM) and gradient update (PyTorch FSDP / DeepSpeed) calls, keeping the `put_tensor` / `get_tensor` calls as the data handoff point.
+
+## See Also
+
+- [MooncakeDistributedStore Python API](../../python-api-reference/mooncake-store)
+- [Mooncake Store Deployment Guide](../../deployment/mooncake-store-deployment-guide)
+- [SSD Offload](../../deployment/ssd-offload)

docs/source/getting_started/rust-bindings.md

@@ -0,0 +1,195 @@
+# Rust Bindings
+
+Mooncake provides Rust bindings for two components:
+
+| Crate | Path | Purpose |
+|-------|------|---------|
+| `mooncake_store` | `mooncake-store/rust/` | Distributed KV cache store client |
+| Transfer Engine (Rust bench) | `mooncake-transfer-engine/rust/` | Transfer Engine benchmark / integration |
+
+## Mooncake Store Rust Bindings
+
+### Overview
+
+`mooncake_store` is a Rust crate that wraps the Mooncake Store C ABI (`store_c.h`) using `bindgen`-generated bindings. It exposes a safe, idiomatic Rust API for storing and retrieving byte slices (KV cache objects) over RDMA.
+
+### Build Requirements
+
+The crate requires the compiled Mooncake Store shared library and its public headers. These are produced by the standard CMake build:
+
+```bash
+mkdir build && cd build
+cmake .. \
+    -DWITH_STORE_RUST=ON \   # ON by default
+    -DUSE_CUDA=ON \
+    -DCMAKE_BUILD_TYPE=RelWithDebInfo
+cmake --build . --target build_mooncake_store_rust
+```
+
+The CMake target sets the necessary Rust environment variables and runs `cargo build` automatically.
+
+If you want to build the crate independently (e.g., in your own workspace after a CMake install), set:
+
+```bash
+export MOONCAKE_STORE_LIB_DIR=/path/to/build/mooncake-store
+export MOONCAKE_STORE_INCLUDE_DIR=/path/to/Mooncake/mooncake-store/include
+cargo build
+```
+
+### Adding as a Dependency
+
+After building, you can depend on the crate from a local path:
+
+```toml
+# Cargo.toml
+[dependencies]
+mooncake_store = { path = "/path/to/Mooncake/mooncake-store/rust" }
+```
+
+### API Reference
+
+```rust
+use mooncake_store::{MooncakeStore, ReplicateConfig, StoreError};
+
+// 1. Create a store handle
+let store = MooncakeStore::new()?;
+
+// 2. Connect to a running Mooncake master
+store.setup(
+    "node1",                           // local_hostname
+    "http://10.0.0.1:8080/metadata",   // metadata_server
+    512 << 20,                         // global_segment_size (512 MiB)
+    128 << 20,                         // local_buffer_size  (128 MiB)
+    "rdma",                            // protocol: "tcp" | "rdma" | …
+    "mlx5_0",                          // device_name ("" for auto)
+    "10.0.0.1:50051",                  // mooncake_master address
+)?;
+
+// 3. Store a value
+store.put("my-key", b"hello, mooncake!", None)?;
+
+// 4. Store with replication options
+let config = ReplicateConfig {
+    replica_num: 2,
+    with_soft_pin: true,
+    with_hard_pin: false,
+    preferred_segments: vec!["seg-0".into()],
+};
+store.put("replicated-key", b"replicated value", Some(&config))?;
+
+// 5. Check existence
+let exists: bool = store.is_exist("my-key")?;
+
+// 6. Get size
+let size: u64 = store.get_size("my-key")?;
+
+// 7. Retrieve value
+let data: Vec<u8> = store.get("my-key")?;
+assert_eq!(data, b"hello, mooncake!");
+
+// 8. Remove
+store.remove("my-key", /*force=*/false)?;
+```
+
+### Error Handling
+
+All fallible operations return `Result<T, StoreError>`. The main variants are:
+
+| Variant | Meaning |
+|---------|---------|
+| `StoreError::OperationFailed(code)` | The underlying C library returned a non-zero status code |
+| `StoreError::SetupError(msg)` | Store setup / connection failed |
+| `StoreError::InvalidArgument(msg)` | A null pointer or invalid argument was passed |
+
+A missing key returns `StoreError::OperationFailed(code)` where `code` matches the Mooncake error table in [`mooncake-store/include/types.h`](../../../mooncake-store/include/types.h).
+
+### Running the Example
+
+```bash
+# Start metadata server
+cd mooncake-transfer-engine/example/http-metadata-server-python
+pip install aiohttp && python bootstrap_server.py &
+
+# Start mooncake_master
+./build/mooncake_master \
+    --enable_http_metadata_server=true \
+    --http_metadata_server_port=8080 \
+    --rpc_port=50051 &
+
+# Run the bundled example
+cd build
+cargo run --example basic_usage --manifest-path ../mooncake-store/rust/Cargo.toml
+```
+
+Or via CMake:
+
+```bash
+cd build
+cmake --build . --target build_mooncake_store_rust
+# The basic_usage binary is at:
+./mooncake-store/rust/target/debug/examples/basic_usage
+```
+
+---
+
+## Transfer Engine Rust Bindings
+
+### Overview
+
+The Transfer Engine Rust bindings live under `mooncake-transfer-engine/rust/`. They provide a Rust interface to `TransferEngine` and include a Rust port of the `transfer_engine_bench` benchmark tool.
+
+### Build
+
+```bash
+mkdir build && cd build
+cmake .. \
+    -DWITH_RUST_EXAMPLE=ON \
+    -DUSE_CUDA=ON \
+    -DCMAKE_BUILD_TYPE=RelWithDebInfo
+cmake --build . --target rust_transfer_engine_bench
+```
+
+### Running the Benchmark
+
+The Rust benchmark mirrors `transfer_engine_bench.cpp` in behaviour:
+
+```bash
+# Target node (receiver)
+./build/mooncake-transfer-engine/rust/target/release/transfer_engine_bench \
+    --mode=target \
+    --protocol=rdma \
+    --metadata_server=P2PHANDSHAKE
+
+# Initiator node (sender) — replace <target>:<port> from the target log
+./build/mooncake-transfer-engine/rust/target/release/transfer_engine_bench \
+    --mode=initiator \
+    --protocol=rdma \
+    --metadata_server=P2PHANDSHAKE \
+    --segment_id=<target>:<port> \
+    --operation=write \
+    --duration=10 \
+    --threads=8
+```
+
+---
+
+## CI Integration
+
+The Mooncake CI pipeline validates both Rust crates on every pull request:
+
+```yaml
+# Relevant CI steps (simplified from .github/workflows/ci.yml)
+- uses: dtolnay/rust-toolchain@stable
+- run: |
+    export MOONCAKE_STORE_LIB_DIR=$BUILD_DIR/mooncake-store
+    export MOONCAKE_STORE_INCLUDE_DIR=$SRC_DIR/mooncake-store/include
+    cargo check --manifest-path mooncake-store/rust/Cargo.toml
+```
+
+Both crates must pass `cargo check` (and `cargo clippy`) before merging.
+
+## See Also
+
+- [Mooncake Store Python API](../python-api-reference/mooncake-store)
+- [Transfer Engine C++ API](../design/transfer-engine/cpp-api)
+- [Mooncake Store Deployment Guide](../deployment/mooncake-store-deployment-guide)

docs/source/index.md

@@ -66,12 +66,15 @@ This repository also hosts its technical report and the open-sourced traces.
 getting_started/build
 getting_started/quick-start
 getting_started/supported-protocols
+getting_started/rust-bindings
 getting_started/plugin-usage/3FS-USRBIO-Plugin
 getting_started/examples/lmcache-integration
 getting_started/examples/lmdeploy-integration-v0.9
 getting_started/examples/sglang-integration-v1
 getting_started/examples/sglang-integration/index
 getting_started/examples/vllm-integration/index
+getting_started/examples/rl-training-integration
+getting_started/examples/mooncake-pg-usage
 :::
 
 % Making the most out of Mooncake
@@ -110,10 +113,12 @@ python-api-reference/ep-backend
 design/architecture
 design/mooncake-store
 design/p2p-store
+design/mooncake-pg
 design/transfer-engine/index
 design/tent/overview
 design/tent/tebench
 design/hicache-design
+design/conductor/indexer-api-design
 :::
 
 % Q&A for Mooncake
@@ -133,6 +138,9 @@ troubleshooting/troubleshooting
 :maxdepth: 2
 
 deployment/mooncake-store-deployment-guide
+deployment/ha-hot-standby
+deployment/monitoring
+deployment/multi-tier-storage
 :::
 
 % Community

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_backend.h

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

mooncake-pg/include/mooncake_worker.cuh

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

mooncake-pg/src/mooncake_backend.cpp

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

mooncake-pg/src/mooncake_worker.cu

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

mooncake-pg/src/mooncake_worker_thread.cpp

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

mooncake-store/go/build.sh

@@ -39,7 +39,7 @@ CGO_LDFLAGS="-L${BUILD_DIR}/mooncake-store/src"
 CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-store/src/cachelib_memory_allocator"
 CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-transfer-engine/src"
 CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-transfer-engine/src/common/base"
-CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-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.
@@ -117,11 +113,10 @@ int main(int argc, char *argv[]) {
         return -1;
     }
 
-    coro_rpc::coro_rpc_server server(FLAGS_threads, FLAGS_port, FLAGS_host);
+    coro_rpc::coro_rpc_server server(FLAGS_threads, FLAGS_port, "127.0.0.1");
     RegisterClientRpcService(server, *client_inst);
 
-    LOG(INFO) << "Starting real client service on " << FLAGS_host << ":"
-              << FLAGS_port;
+    LOG(INFO) << "Starting real client service on 127.0.0.1:" << FLAGS_port;
 
     return server.start();
 }

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/health_check_test.cpp

@@ -80,20 +80,6 @@ class HealthCheckTest : public ::testing::Test {
                                       16 * 1024 * 1024, FLAGS_protocol,
                                       rdma_devices, master_address_);
     }
-
-    bool WaitForHealthCode(
-        int expected_code,
-        std::chrono::milliseconds timeout = std::chrono::seconds(10),
-        std::chrono::milliseconds interval = std::chrono::milliseconds(100)) {
-        const auto deadline = std::chrono::steady_clock::now() + timeout;
-        while (std::chrono::steady_clock::now() < deadline) {
-            if (py_client_->health_check() == expected_code) {
-                return true;
-            }
-            std::this_thread::sleep_for(interval);
-        }
-        return py_client_->health_check() == expected_code;
-    }
 };
 
 // Test 1: health_check returns HC_NOT_INITIALIZED before setup
@@ -134,8 +120,9 @@ TEST_F(HealthCheckTest, ReturnsTwoWhenMasterDown) {
 
     // Stop master, wait for ping to fail
     master_.Stop();
-    EXPECT_TRUE(WaitForHealthCode(HC_MASTER_UNREACHABLE))
-        << "Timed out waiting for HC_MASTER_UNREACHABLE";
+    std::this_thread::sleep_for(std::chrono::seconds(3));
+
+    EXPECT_EQ(py_client_->health_check(), HC_MASTER_UNREACHABLE);
 
     py_client_->tearDownAll();
 }
@@ -167,8 +154,7 @@ TEST_F(HealthCheckTest, HttpReturns503WhenMasterDown) {
     EXPECT_EQ(py_client_->health_check(), HC_HEALTHY);
 
     master_.Stop();
-    ASSERT_TRUE(WaitForHealthCode(HC_MASTER_UNREACHABLE))
-        << "Timed out waiting for HC_MASTER_UNREACHABLE";
+    std::this_thread::sleep_for(std::chrono::seconds(3));
 
     auto resp = fetch_health(http_port);
     EXPECT_EQ(resp.http_status, 503);

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

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

mooncake-transfer-engine/example/CMakeLists.txt

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

mooncake-transfer-engine/example/batch_register_bench.py

@@ -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/efa_per_transfer_latency_bench.py

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

mooncake-transfer-engine/example/kvcache_prefix_bench.py

@@ -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,
@@ -472,21 +457,6 @@ static inline const std::string MakeNicPath(const std::string &server_name,
     return server_name + NIC_PATH_DELIM + nic_name;
 }
 
-// Strip the port from a nic_path to get a stable key for endpoint reuse.
-// "ip-172-31-45-191:15365@rdmap135s0" → "ip-172-31-45-191@rdmap135s0"
-// This allows the same physical peer to reuse endpoints across reconnections
-// (each run picks a random P2P handshake port).
-static inline std::string normalizeNicPath(const std::string &nic_path) {
-    std::string server_name = getServerNameFromNicPath(nic_path);
-    std::string nic_name = getNicNameFromNicPath(nic_path);
-    if (server_name.empty() || nic_name.empty()) return nic_path;
-    size_t colon = server_name.rfind(':');
-    if (colon != std::string::npos) {
-        server_name = server_name.substr(0, colon);
-    }
-    return server_name + NIC_PATH_DELIM + nic_name;
-}
-
 static inline bool overlap(const void *a, size_t a_len, const void *b,
                            size_t b_len) {
     return (a >= b && a < (char *)b + b_len) ||

mooncake-transfer-engine/include/config.h

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

mooncake-transfer-engine/include/multi_transport.h

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

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

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

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

@@ -52,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