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

---------

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

.github/CODEOWNERS

@@ -17,6 +17,7 @@
 /mooncake-transfer-engine @alogfans @doujiang24 @chestnut-Q
 /mooncake-transfer-engine/*/transport/hip_transport/ @alogfans @amd-arozanov
 /mooncake-transfer-engine/*/transport/ascend_transport/ @alogfans @ascend-direct-dev
+/mooncake-transfer-engine/*/transport/efa_transport/ @alogfans @whn09
 /mooncake-wheel @ShangmingCai @stmatengss
 /scripts/tone_tests @luketong777
 /scripts/ascend/ @ascend-direct-dev @VNightMare @MingYang119

.github/workflows/ci.yml

@@ -6,19 +6,25 @@ on:
   pull_request:
     branches: [ "main" ]
     types: [opened, synchronize, reopened, labeled]
+  workflow_dispatch: {}
 
 permissions:
   contents: read
-  actions: write
-  statuses: write
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref_name }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: true
 
 jobs:
   build:
-    needs: [spell-check, clang-format]
+    needs: [spell-check, clang-format, check-paths]
     if: >-
-      github.event_name == 'push' ||
-      github.event.action == 'opened' ||
-      contains(github.event.pull_request.labels.*.name, 'run-ci')
+      (needs.check-paths.outputs.should-run-downstream == 'true' ||
+       github.event_name == 'workflow_dispatch') &&
+      (github.event_name == 'push' ||
+       github.event_name == 'workflow_dispatch' ||
+       github.event.action == 'opened' ||
+       contains(github.event.pull_request.labels.*.name, 'run-ci'))
     runs-on: ubuntu-22.04
     strategy:
       matrix:
@@ -28,21 +34,9 @@ jobs:
       SCCACHE_GHA_ENABLED: "true"
 
     steps:
-    - name: Cancel workflow if checks failed
-      if: ${{ needs.spell-check.result != 'success' || needs.clang-format.result != 'success' }}
-      uses: actions/github-script@v7
-      with:
-        script: |
-          console.log('Cancelling workflow run due to spell-check or clang-format failure');
-          await github.rest.actions.cancelWorkflowRun({
-            owner: context.repo.owner,
-            repo: context.repo.repo,
-            run_id: context.runId
-          });
-          // Wait for cancel to propagate
-          await new Promise(resolve => setTimeout(resolve, 5000));
-          process.exit(1);
     - uses: actions/checkout@v4
+      with:
+        persist-credentials: false
 
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v5
@@ -106,7 +100,7 @@ jobs:
         sudo bash -x dependencies.sh -y
         mkdir build
         cd build
-        cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_CXL=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON -DENABLE_ASAN=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Debug
+        cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_CXL=ON -DUSE_UB=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON -DENABLE_ASAN=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Debug
       shell: bash
 
     - name: Build project
@@ -140,10 +134,12 @@ jobs:
         MASTER_PID=$!
         sleep 3
         cd mooncake-store/go
-        export LD_LIBRARY_PATH=$GITHUB_WORKSPACE/build/mooncake-asio:$GITHUB_WORKSPACE/build/mooncake-store/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base:$GITHUB_WORKSPACE/build/mooncake-common/etcd
+        export LD_LIBRARY_PATH=$GITHUB_WORKSPACE/build/mooncake-common:$GITHUB_WORKSPACE/build/mooncake-store/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base:$GITHUB_WORKSPACE/build/mooncake-common/etcd
         export CGO_ENABLED=1
         export CGO_CFLAGS="-I$GITHUB_WORKSPACE/mooncake-store/include -I$GITHUB_WORKSPACE/mooncake-transfer-engine/include"
-        export CGO_LDFLAGS="-L$GITHUB_WORKSPACE/build/mooncake-store/src -L$GITHUB_WORKSPACE/build/mooncake-store/src/cachelib_memory_allocator -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base -L$GITHUB_WORKSPACE/build/mooncake-asio -L$GITHUB_WORKSPACE/build/mooncake-common/etcd -lmooncake_store -lcachelib_memory_allocator -ltransfer_engine -lbase -lasio -letcd_wrapper -lstdc++ -lnuma -lglog -lgflags -libverbs -ljsoncpp -lzstd -lcurl -luring -lasan -lm -lgcov"
+        export CGO_LDFLAGS="-L$GITHUB_WORKSPACE/build/mooncake-store/src -L$GITHUB_WORKSPACE/build/mooncake-store/src/cachelib_memory_allocator -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base -L$GITHUB_WORKSPACE/build/mooncake-common -L$GITHUB_WORKSPACE/build/mooncake-common/etcd -lmooncake_store -lcachelib_memory_allocator -ltransfer_engine -lbase -lasio -letcd_wrapper -lstdc++ -lnuma -lglog -lgflags -libverbs -ljsoncpp -lzstd -lcurl -luring -lasan -lm -lgcov"
+        # Link cudart if CUDA is available (needed for D2H staging in mooncake_store)
+        if [ -d /usr/local/cuda/lib64 ]; then export CGO_LDFLAGS="$CGO_LDFLAGS -L/usr/local/cuda/lib64 -lcudart"; fi
         ASAN_OPTIONS=detect_leaks=0:verify_asan_link_order=0 MC_METADATA_SERVER=http://127.0.0.1:8080/metadata go test -v ./tests/...
         kill $MASTER_PID 2>/dev/null || true
       shell: bash
@@ -156,6 +152,17 @@ jobs:
         MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 ctest -j --output-on-failure
       shell: bash
 
+    - name: Drain HTTP E2E test
+      if: matrix.python-version == '3.12'
+      run: |
+        cd build
+        export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
+        # Keep the sanitizer gate on the C++ integration test. The Python
+        # drain script is manual/nightly only because pybind + ASan teardown in
+        # a Python host process is not stable.
+        DEFAULT_KV_LEASE_TTL=500 ./mooncake-store/tests/task_integration_test --gtest_filter='TaskExecutorIntegrationTest.DrainJobCompleteFlow'
+      shell: bash
+
     - name: Generate coverage report
       id: coverage
       run: |
@@ -229,28 +236,20 @@ jobs:
         path: mooncake-wheel/dist-py${{ steps.generate_tag_build.outputs.python_version_tag }}/*.whl
 
   build-musa:
-    needs: [spell-check, clang-format]
+    needs: [spell-check, clang-format, check-paths]
     if: >-
-      github.event_name == 'push' ||
-      github.event.action == 'opened' ||
-      contains(github.event.pull_request.labels.*.name, 'run-ci')
+      (needs.check-paths.outputs.should-run-downstream == 'true' ||
+       github.event_name == 'workflow_dispatch') &&
+      (github.event_name == 'push' ||
+       github.event_name == 'workflow_dispatch' ||
+       github.event.action == 'opened' ||
+       contains(github.event.pull_request.labels.*.name, 'run-ci'))
     runs-on: ubuntu-22.04
     container: mthreads/musa:rc4.3.0-devel-ubuntu22.04-amd64
     steps:
-    - name: Cancel workflow if checks failed
-      if: ${{ needs.spell-check.result != 'success' || needs.clang-format.result != 'success' }}
-      uses: actions/github-script@v7
-      with:
-        script: |
-          console.log('Cancelling workflow run due to spell-check or clang-format failure');
-          await github.rest.actions.cancelWorkflowRun({
-            owner: context.repo.owner,
-            repo: context.repo.repo,
-            run_id: context.runId
-          });
-          await new Promise(resolve => setTimeout(resolve, 5000));
-          process.exit(1);
     - uses: actions/checkout@v4
+      with:
+        persist-credentials: false
 
     - name: Mark repository as safe
       run: git config --global --add safe.directory $GITHUB_WORKSPACE
@@ -279,6 +278,7 @@ jobs:
     if: >-
       needs.build-flags.result == 'success' &&
       (github.event_name == 'push' ||
+       github.event_name == 'workflow_dispatch' ||
        github.event.action == 'opened' ||
        contains(github.event.pull_request.labels.*.name, 'run-ci'))
     strategy:
@@ -288,6 +288,8 @@ jobs:
     runs-on: ${{ matrix.ubuntu-version }}
     steps:
     - uses: actions/checkout@v4
+      with:
+        persist-credentials: false
 
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v5
@@ -343,6 +345,14 @@ jobs:
 
     - name: Run tests with ssd
       run: |
+        # Reserve port 50052 (mooncake_client RPC port) so the kernel never
+        # auto-allocates it as ephemeral source port for other outbound
+        # connections in the test suite. Without this, a random Python test
+        # connection can pick src_port=50052, leave a TIME_WAIT on
+        # <eth0_ip>:50052 for 60s, and block mooncake_client's bind to
+        # 0.0.0.0:50052 even with SO_REUSEADDR (Linux only relaxes
+        # TIME_WAIT+bind conflict for same-IP or loopback).
+        sudo sysctl -w net.ipv4.ip_local_reserved_ports=50052
         source test_env/bin/activate
         MC_STORE_MEMCPY=false TEST_SSD_OFFLOAD_IN_EVICT=true ./scripts/run_tests.sh
         rm -rf /tmp/mooncake_test_ssd
@@ -393,6 +403,18 @@ jobs:
         python scripts/test_copy_move_api.py
       shell: bash
 
+    - name: Run Python Drain HTTP E2E Test (CI check)
+      env:
+        MOONCAKE_MASTER: "127.0.0.1:50051"
+        MOONCAKE_TE_META_DATA_SERVER: "http://127.0.0.1:8080/metadata"
+        MOONCAKE_PROTOCOL: "tcp"
+        LOCAL_HOSTNAME: "127.0.0.1"
+      run: |
+        source test_env/bin/activate
+        export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
+        python scripts/test_drain_http_api.py --timeout-sec 90
+      shell: bash
+
     - name: Run RPC Communicator Bandwidth Test
       run: |
         source test_env/bin/activate
@@ -403,6 +425,14 @@ jobs:
         kill $SERVER_PID 2>/dev/null || true
         wait $SERVER_PID 2>/dev/null || true
 
+    - name: Test Mooncake PyTorch Backend (CPU Only)
+      env:
+        MC_FORCE_TCP: "true"
+      run: |
+        source test_env/bin/activate
+        python -m unittest mooncake-wheel.tests.test_mooncake_backend_cpu
+      shell: bash
+
     - name: Test Safetensor Functions
       run: |
         source test_env/bin/activate
@@ -411,11 +441,14 @@ jobs:
       shell: bash
 
   build-flags:
-    needs: [spell-check, clang-format]
+    needs: [spell-check, clang-format, check-paths]
     if: >-
-      github.event_name == 'push' ||
-      github.event.action == 'opened' ||
-      contains(github.event.pull_request.labels.*.name, 'run-ci')
+      (needs.check-paths.outputs.should-run-downstream == 'true' ||
+       github.event_name == 'workflow_dispatch') &&
+      (github.event_name == 'push' ||
+       github.event_name == 'workflow_dispatch' ||
+       github.event.action == 'opened' ||
+       contains(github.event.pull_request.labels.*.name, 'run-ci'))
     runs-on: ubuntu-22.04
     strategy:
       matrix:
@@ -427,20 +460,9 @@ jobs:
       SCCACHE_GHA_ENABLED: "true"
 
     steps:
-    - name: Cancel workflow if checks failed
-      if: ${{ needs.spell-check.result != 'success' || needs.clang-format.result != 'success' }}
-      uses: actions/github-script@v7
-      with:
-        script: |
-          console.log('Cancelling workflow run due to spell-check or clang-format failure');
-          await github.rest.actions.cancelWorkflowRun({
-            owner: context.repo.owner,
-            repo: context.repo.repo,
-            run_id: context.runId
-          });
-          await new Promise(resolve => setTimeout(resolve, 5000));
-          process.exit(1);
     - uses: actions/checkout@v4
+      with:
+        persist-credentials: false
 
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v5
@@ -486,6 +508,9 @@ jobs:
         df -h
       shell: bash
 
+    - name: Install Rust toolchain
+      uses: dtolnay/rust-toolchain@stable
+
     - name: Build transfer engine only
       run: |
         cd mooncake-transfer-engine
@@ -493,7 +518,7 @@ jobs:
         cd build
         export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH
         export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
-        cmake -G Ninja .. -DUSE_ETCD=OFF -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=OFF -DUSE_MNNVL=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
+        cmake -G Ninja .. -DUSE_ETCD=OFF -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=OFF -DUSE_MNNVL=OFF -DUSE_UB=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
         cmake --build .
         sudo cmake --install .
         df -h
@@ -503,7 +528,7 @@ jobs:
       run: |
         mkdir build
         cd build
-        cmake -G Ninja .. -DUSE_ETCD=ON -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_STORE=ON -DWITH_P2P_STORE=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=ON -DUSE_MNNVL=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
+        cmake -G Ninja .. -DUSE_ETCD=ON -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_STORE=ON -DWITH_P2P_STORE=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=ON -DUSE_MNNVL=OFF -DUSE_UB=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
       shell: bash
       # TODO: lack USE_NVMEOF,USE_MNNVL
 
@@ -520,9 +545,8 @@ jobs:
     - name: Configure project with unit tests and examples
       run: |
         cd build
-        cmake -G Ninja .. -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON
+        cmake -G Ninja .. -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DWITH_STORE_RUST=ON -DENABLE_SCCACHE=ON
       shell: bash
-      # TODO: lack WITH_RUST_EXAMPLE
 
     - name: Build project with unit tests and examples
       run: |
@@ -533,11 +557,19 @@ jobs:
         sudo cmake --install .
       shell: bash
 
+    - name: Check Mooncake Store Rust bindings and example
+      run: |
+        cd mooncake-store/rust
+        MOONCAKE_STORE_LIB_DIR=$GITHUB_WORKSPACE/build/mooncake-store/src \
+        MOONCAKE_STORE_INCLUDE_DIR=$GITHUB_WORKSPACE/mooncake-store/include \
+          cargo check --example basic_usage --tests
+      shell: bash
+
     - name: Configure project
       run: |
         cd build
         rm -r */tests
-        cmake -G Ninja .. -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF -DUSE_HTTP=ON -DENABLE_SCCACHE=ON -DUSE_CXL=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.0;2.9.1;2.10.0"
+        cmake -G Ninja .. -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF -DUSE_HTTP=ON -DENABLE_SCCACHE=ON -DUSE_CXL=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.0;2.9.1;2.10.0;2.11.0"
       shell: bash
 
     - name: Build project
@@ -579,27 +611,19 @@ jobs:
 
   build-docker:
     name: Build Docker Image
-    needs: [spell-check, clang-format]
+    needs: [spell-check, clang-format, check-paths]
     if: >-
-      github.event_name == 'push' ||
-      github.event.action == 'opened' ||
-      contains(github.event.pull_request.labels.*.name, 'run-ci')
+      (needs.check-paths.outputs.should-run-downstream == 'true' ||
+       github.event_name == 'workflow_dispatch') &&
+      (github.event_name == 'push' ||
+       github.event_name == 'workflow_dispatch' ||
+       github.event.action == 'opened' ||
+       contains(github.event.pull_request.labels.*.name, 'run-ci'))
     runs-on: ubuntu-22.04
     steps:
-    - name: Cancel workflow if checks failed
-      if: ${{ needs.spell-check.result != 'success' || needs.clang-format.result != 'success' }}
-      uses: actions/github-script@v7
-      with:
-        script: |
-          console.log('Cancelling workflow run due to spell-check or clang-format failure');
-          await github.rest.actions.cancelWorkflowRun({
-            owner: context.repo.owner,
-            repo: context.repo.repo,
-            run_id: context.runId
-          });
-          await new Promise(resolve => setTimeout(resolve, 5000));
-          process.exit(1);
     - uses: actions/checkout@v4
+      with:
+        persist-credentials: false
 
     - name: Set up Docker Buildx
       uses: docker/setup-buildx-action@v2
@@ -615,33 +639,15 @@ jobs:
     name: Spell Check with Typos
     if: >-
       github.event_name == 'push' ||
+      github.event_name == 'workflow_dispatch' ||
       github.event.action == 'opened' ||
       contains(github.event.pull_request.labels.*.name, 'run-ci')
     runs-on: ubuntu-22.04
     steps:
-    - name: Set pending status for downstream tests
-      if: github.event_name == 'pull_request'
-      continue-on-error: true
-      uses: actions/github-script@v7
-      with:
-        script: |
-          const contexts = [
-            'CI Test on ASCEND / build-and-test',
-            'Integration test / test-sglang-integration'
-          ];
-          for (const ctx of contexts) {
-            await github.rest.repos.createCommitStatus({
-              owner: context.repo.owner,
-              repo: context.repo.repo,
-              sha: context.payload.pull_request.head.sha,
-              state: 'pending',
-              context: ctx,
-              description: 'Waiting for Build & Test to complete...',
-              target_url: 'https://github.com/' + context.repo.owner + '/' + context.repo.repo + '/actions/runs/' + context.runId
-            });
-          }
     - name: Checkout Actions Repository
       uses: actions/checkout@v4
+      with:
+        persist-credentials: false
     - name: Spell Check Repo
       uses: crate-ci/typos@v1.30.2
 
@@ -649,6 +655,7 @@ jobs:
     name: Check code format
     if: >-
       github.event_name == 'push' ||
+      github.event_name == 'workflow_dispatch' ||
       github.event.action == 'opened' ||
       contains(github.event.pull_request.labels.*.name, 'run-ci')
     runs-on: ubuntu-22.04
@@ -657,6 +664,7 @@ jobs:
       uses: actions/checkout@v4
       with:
         fetch-depth: 0  # Need full history for branch comparison
+        persist-credentials: false
 
     - name: Install clang-format 20
       run: |
@@ -692,3 +700,91 @@ jobs:
         ./scripts/code_format.sh --check --base "${BASE_REF}"
       shell: bash
 
+
+  check-paths:
+    if: >-
+      github.event_name == 'push' ||
+      github.event_name == 'workflow_dispatch' ||
+      github.event.action == 'opened' ||
+      contains(github.event.pull_request.labels.*.name, 'run-ci')
+    runs-on: ubuntu-latest
+    outputs:
+      should-run-downstream: ${{ steps.dispatch-override.outputs.src || steps.filter.outputs.src }}
+    steps:
+      # workflow_dispatch has no PR/push diff context — skip paths-filter and default to true
+      - name: Default to true for workflow_dispatch
+        id: dispatch-override
+        if: github.event_name == 'workflow_dispatch'
+        run: echo "src=true" >> $GITHUB_OUTPUT
+      - uses: actions/checkout@v4
+        if: github.event_name != 'workflow_dispatch'
+        with:
+          fetch-depth: 2
+          persist-credentials: false
+      - uses: dorny/paths-filter@v3
+        if: github.event_name != 'workflow_dispatch'
+        id: filter
+        with:
+          filters: |
+            src:
+              - 'mooncake-*/**'
+              - 'extern/**'
+              - 'CMakeLists.txt'
+              - 'dependencies.sh'
+              - 'scripts/**'
+              - '.github/workflows/**'
+
+  build-wheel-cu13:
+    needs: [spell-check, clang-format, check-paths]
+    if: >-
+      (needs.check-paths.outputs.should-run-downstream == 'true' ||
+       github.event_name == 'workflow_dispatch') &&
+      (github.event_name == 'push' ||
+       github.event_name == 'workflow_dispatch' ||
+       github.event.action == 'opened' ||
+       contains(github.event.pull_request.labels.*.name, 'run-ci'))
+    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()
+    needs:
+      - spell-check
+      - clang-format
+      - build
+      - build-musa
+      - build-flags
+      - build-docker
+      - test-wheel-ubuntu
+      - build-wheel-cu13
+      - ascend-test
+      - integration-test
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check required job results
+        run: |
+          failing=$(echo "$NEEDS_JSON" | jq -r '
+            to_entries[] |
+            select(.value.result != "success" and .value.result != "skipped") |
+            "\(.key): \(.value.result)"')
+          if [ -n "$failing" ]; then
+            echo "::error::The following jobs failed or were cancelled:"
+            echo "$failing"
+            exit 1
+          fi
+          echo "All checks passed or were acceptably skipped."
+        env:
+          NEEDS_JSON: ${{ toJSON(needs) }}

.github/workflows/ci_ascend.yml

@@ -1,49 +1,22 @@
 name: 'CI Test on ASCEND Platform'
 
 on:
-  workflow_run:
-    workflows: ["Build & Test (Linux)"]
-    types:
-      - completed
-
-permissions:
-  statuses: write
+  workflow_call:
+    inputs:
+      checkout_ref:
+        description: 'Git ref to checkout (PR head SHA for pull_request_target)'
+        required: false
+        type: string
 
 jobs:
-  check-paths:
-    runs-on: ubuntu-latest
-    outputs:
-      should-run: ${{ steps.filter.outputs.changed }}
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ github.event.workflow_run.head_sha }}
-          fetch-depth: 0
-
-      - name: Check changed files
-        id: filter
-        uses: dorny/paths-filter@v3
-        with:
-          base: ${{ github.event.workflow_run.event == 'pull_request' && github.event.workflow_run.pull_requests[0].base.ref || 'main' }}
-          ref: ${{ github.event.workflow_run.head_sha }}
-          filters: |
-            changed:
-              - 'mooncake-*/**'
-              - 'extern/**'
-              - 'CMakeLists.txt'
-              - 'scripts/**'
-              - '.github/workflows/**'
-
   build-and-test:
-    needs: check-paths
-    if: ${{ github.event.workflow_run.conclusion == 'success' && github.repository == 'kvcache-ai/Mooncake' && needs.check-paths.outputs.should-run == 'true' }}
+    if: github.repository == 'kvcache-ai/Mooncake'
     runs-on: self-hosted
 
     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"
@@ -55,30 +28,132 @@ 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
-        if ! git submodule update --init --recursive; then
-            if [ ! -d "extern/pybind11" ] || [ -z "$(ls -A 'extern/pybind11' 2>/dev/null)" ]; then
-                echo "git submodule update failed, try to cp pybind11..."
-                if [ -d "../pybind11" ]; then
-                    cp -r ../pybind11 extern/
-                else
-                    echo "Error: ../pybind11 does not exist. Cannot copy pybind11."
-                    exit 1
-                fi
-            else
-                echo "Detected that extern/pybind11 already exists, continuing execution...."
+
+        submodule_updated=false
+        if git submodule update --init --recursive; then
+          submodule_updated=true
+        elif [ -n "${ASCEND_GITHUB_MIRROR_URLS:-}" ]; then
+          normalize_base() {
+            local base="$1"
+            base="${base#${base%%[![:space:]]*}}"
+            base="${base%${base##*[![:space:]]}}"
+            [ -n "$base" ] || return 1
+            [ "$base" != "https://github.com/" ] && base="${base%/}/"
+            printf '%s\n' "$base"
+          }
+
+          while IFS= read -r raw; do
+            base="$(normalize_base "$raw" || true)"
+            [ -n "$base" ] || continue
+            [ "$base" = "https://github.com/" ] && continue
+
+            echo "Retrying submodule update with ${base}"
+            if git -c url."${base}https://github.com/".insteadOf=https://github.com/ \
+              submodule update --init --recursive; then
+              submodule_updated=true
+              break
             fi
+          done < <(printf '%s\n' "$ASCEND_GITHUB_MIRROR_URLS" | tr ',;' '\n')
         fi
+
+        if [ "$submodule_updated" != true ]; then
+          if [ ! -d "extern/pybind11" ] || [ -z "$(ls -A 'extern/pybind11' 2>/dev/null)" ]; then
+            echo "git submodule update failed (mirrors also exhausted), trying to cp pybind11..."
+            if [ -d "../pybind11" ]; then
+              cp -r ../pybind11 extern/
+            else
+              echo "Error: ../pybind11 does not exist. Cannot copy pybind11."
+              exit 1
+            fi
+          else
+            echo "Detected that extern/pybind11 already exists, continuing execution...."
+          fi
+        fi
+
         bash scripts/ascend/dependencies_ascend_installation.sh
         echo "Configuring CMake..."
         rm -rf build
@@ -146,7 +221,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
 
@@ -270,7 +345,7 @@ jobs:
 
           echo ""
           echo "All Hixl Mooncake Store tests completed successfully!"
-          
+
 
     - name: Test Summary
       if: always()
@@ -287,50 +362,3 @@ jobs:
           /tmp/hixl-test-log/*
         retention-days: 30
         if-no-files-found: warn
-
-  report-status:
-    needs: [check-paths, build-and-test]
-    if: always()
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/github-script@v7
-        with:
-          script: |
-            const conclusion = context.payload.workflow_run.conclusion;
-            const shouldRun = '${{ needs.check-paths.outputs.should-run }}';
-            const testResult = '${{ needs.build-and-test.result }}';
-            const isTargetRepo = '${{ github.repository }}' === 'kvcache-ai/Mooncake';
-
-            let state, description;
-            if (conclusion === 'cancelled') {
-              state = 'failure';
-              description = 'CI cancelled (format/spell check failed)';
-            } else if (conclusion !== 'success') {
-              state = 'failure';
-              description = 'CI ' + conclusion + ', Ascend test skipped';
-            } else if (!isTargetRepo) {
-              state = 'success';
-              description = 'Skipped (fork repository)';
-            } else if (shouldRun !== 'true') {
-              state = 'success';
-              description = 'Skipped (no relevant file changes)';
-            } else if (testResult === 'success') {
-              state = 'success';
-              description = 'Ascend test passed';
-            } else if (testResult === 'skipped') {
-              state = 'success';
-              description = 'Skipped';
-            } else {
-              state = 'failure';
-              description = `Ascend test ${testResult}`;
-            }
-
-            await github.rest.repos.createCommitStatus({
-              owner: context.repo.owner,
-              repo: context.repo.repo,
-              sha: context.payload.workflow_run.head_sha,
-              state: state,
-              context: 'CI Test on ASCEND / build-and-test',
-              description: description,
-              target_url: `https://github.com/${context.repo.owner}/${context.repo.repo}/actions/runs/${context.runId}`
-            });

.github/workflows/ci_cu13.yml

@@ -1,32 +1,10 @@
 name: 'Build Wheel (CUDA 13)'
 
 on:
-  push:
-    branches: [ "main" ]
-    paths:
-      - 'mooncake-*/**'
-      - 'extern/**'
-      - 'CMakeLists.txt'
-      - 'dependencies.sh'
-      - 'scripts/**'
-      - '.github/workflows/ci_cu13.yml'
-  pull_request:
-    branches: [ "main" ]
-    types: [opened, synchronize, reopened, labeled]
-    paths:
-      - 'mooncake-*/**'
-      - 'extern/**'
-      - 'CMakeLists.txt'
-      - 'dependencies.sh'
-      - 'scripts/**'
-      - '.github/workflows/ci_cu13.yml'
+  workflow_call: {}
 
 jobs:
   build-wheel-cu13:
-    if: >-
-      github.event_name == 'push' ||
-      github.event.action == 'opened' ||
-      contains(github.event.pull_request.labels.*.name, 'run-ci')
     runs-on: ubuntu-22.04
     strategy:
       matrix:
@@ -39,6 +17,8 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
+      with:
+        persist-credentials: false
 
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v5
@@ -95,7 +75,7 @@ jobs:
           -DWITH_STORE=ON \
           -DWITH_P2P_STORE=ON \
           -DWITH_EP=ON \
-          -DEP_TORCH_VERSIONS="2.9.0;2.9.1;2.10.0" \
+          -DEP_TORCH_VERSIONS="2.9.0;2.9.1;2.10.0;2.11.0" \
           -DWITH_METRICS=ON \
           -DBUILD_UNIT_TESTS=OFF \
           -DBUILD_EXAMPLES=ON \

.github/workflows/deploy.yml

@@ -35,7 +35,9 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v4
-        
+        with:
+          persist-credentials: false
+
       - name: Setup Python
         uses: actions/setup-python@v4
         with:

.github/workflows/e2e-ci.yml

@@ -0,0 +1,94 @@
+name: E2E CI
+
+on:
+  pull_request_target:
+    branches: ["main"]
+    types: [labeled]
+  workflow_dispatch:
+    inputs:
+      pr_number:
+        description: 'PR number that triggered this'
+        required: false
+        type: string
+      pr_sha:
+        description: 'PR head SHA to checkout'
+        required: false
+        type: string
+      triggered_by:
+        description: 'User who triggered this'
+        required: false
+        type: string
+
+permissions:
+  contents: read
+  pull-requests: write
+
+concurrency:
+  group: e2e-ci-${{ github.event.pull_request.number || inputs.pr_number || github.sha }}
+  cancel-in-progress: true
+
+jobs:
+  ascend-test:
+    if: >
+      github.event_name == 'workflow_dispatch' ||
+      github.event.label.name == 'run-e2e-ci'
+    uses: ./.github/workflows/ci_ascend.yml
+    with:
+      checkout_ref: ${{ inputs.pr_sha || github.event.pull_request.head.sha }}
+    secrets: inherit
+
+  integration-test:
+    if: >
+      github.event_name == 'workflow_dispatch' ||
+      github.event.label.name == 'run-e2e-ci'
+    uses: ./.github/workflows/integration-test.yml
+    with:
+      pr_sha: ${{ inputs.pr_sha || github.event.pull_request.head.sha }}
+      pr_number: ${{ inputs.pr_number || github.event.pull_request.number }}
+    secrets: inherit
+
+  e2e-gate:
+    name: E2E Gate
+    if: >
+      always() &&
+      (github.event_name == 'workflow_dispatch' ||
+       github.event.label.name == 'run-e2e-ci')
+    needs:
+      - ascend-test
+      - integration-test
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check E2E results
+        run: |
+          echo "PR: #${{ inputs.pr_number || github.event.pull_request.number }}"
+          echo "SHA: ${{ inputs.pr_sha || github.event.pull_request.head.sha }}"
+          failing=$(echo "$NEEDS_JSON" | jq -r '
+            to_entries[] |
+            select(.value.result != "success" and .value.result != "skipped") |
+            "\(.key): \(.value.result)"')
+          if [ -n "$failing" ]; then
+            echo "::error::The following E2E jobs failed:"
+            echo "$failing"
+            exit 1
+          fi
+          echo "All E2E checks passed."
+        env:
+          NEEDS_JSON: ${{ toJSON(needs) }}
+
+  cleanup-label:
+    name: Cleanup E2E Label
+    if: >
+      always() &&
+      github.event_name != 'workflow_dispatch' &&
+      github.event.label.name == 'run-e2e-ci'
+    needs:
+      - e2e-gate
+    runs-on: ubuntu-latest
+    steps:
+      - name: Remove run-e2e-ci label
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: |
+          gh pr edit ${{ github.event.pull_request.number }} \
+            --repo ${{ github.repository }} \
+            --remove-label "run-e2e-ci" 2>/dev/null || true

.github/workflows/integration-test.yml

@@ -1,43 +1,19 @@
 name: 'Integration test (Linux)'
 
 on:
-  workflow_run:
-    workflows: ["Build & Test (Linux)"]
-    types:
-      - completed
-
-permissions:
-  statuses: write
+  workflow_call:
+    inputs:
+      pr_sha:
+        description: 'PR head SHA (passed from parent workflow for workflow_dispatch)'
+        required: false
+        type: string
+      pr_number:
+        description: 'PR number (passed from parent workflow for workflow_dispatch)'
+        required: false
+        type: string
 
 jobs:
-  check-paths:
-    runs-on: ubuntu-latest
-    outputs:
-      should-run: ${{ steps.filter.outputs.changed }}
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ github.event.workflow_run.head_sha }}
-          fetch-depth: 0
-
-      - name: Check changed files
-        id: filter
-        uses: dorny/paths-filter@v3
-        with:
-          base: ${{ github.event.workflow_run.event == 'pull_request' && github.event.workflow_run.pull_requests[0].base.ref || 'main' }}
-          ref: ${{ github.event.workflow_run.head_sha }}
-          filters: |
-            changed:
-              - 'mooncake-*/**'
-              - 'extern/**'
-              - 'CMakeLists.txt'
-              - 'scripts/**'
-              - '.github/workflows/**'
-
   test-sglang-integration:
-    needs: check-paths
-    if: ${{ github.event.workflow_run.conclusion == 'success' && needs.check-paths.outputs.should-run == 'true' }}
     runs-on: ubuntu-latest
     env:
       tone_user_name: ${{ secrets.TONE_USER_NAME }}
@@ -45,10 +21,11 @@ jobs:
       - name: trigger T-one test
         if: ${{ env.tone_user_name != '' }}
         run: |
-          SHA="${{ github.event.workflow_run.head_sha }}"
-          PR_ID="${{ github.event.workflow_run.pull_requests[0].number }}"
+          # Priority: explicit inputs > PR event context > push SHA
+          SHA="${{ inputs.pr_sha || github.event.pull_request.head.sha || github.sha }}"
+          PR_ID="${{ inputs.pr_number || github.event.pull_request.number }}"
 
-          if [ "${{ github.event.workflow_run.event }}" = "push" ]; then
+          if [ "${{ github.event_name }}" = "push" ]; then
             SHA="${{ github.sha }}"
             PR_ID=""
           fi
@@ -150,46 +127,3 @@ jobs:
             sleep 30
           done
         shell: bash
-
-  report-status:
-    needs: [check-paths, test-sglang-integration]
-    if: always()
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/github-script@v7
-        with:
-          script: |
-            const conclusion = context.payload.workflow_run.conclusion;
-            const shouldRun = '${{ needs.check-paths.outputs.should-run }}';
-            const testResult = '${{ needs.test-sglang-integration.result }}';
-
-            let state, description;
-            if (conclusion === 'cancelled') {
-              state = 'failure';
-              description = 'CI cancelled (format/spell check failed)';
-            } else if (conclusion !== 'success') {
-              state = 'failure';
-              description = 'CI ' + conclusion + ', integration test skipped';
-            } else if (shouldRun !== 'true') {
-              state = 'success';
-              description = 'Skipped (no relevant file changes)';
-            } else if (testResult === 'success') {
-              state = 'success';
-              description = 'Integration test passed';
-            } else if (testResult === 'skipped') {
-              state = 'success';
-              description = 'Skipped';
-            } else {
-              state = 'failure';
-              description = `Integration test ${testResult}`;
-            }
-
-            await github.rest.repos.createCommitStatus({
-              owner: context.repo.owner,
-              repo: context.repo.repo,
-              sha: context.payload.workflow_run.head_sha,
-              state: state,
-              context: 'Integration test / test-sglang-integration',
-              description: description,
-              target_url: `https://github.com/${context.repo.owner}/${context.repo.repo}/actions/runs/${context.runId}`
-            });

.github/workflows/release-cuda13.yaml

@@ -65,7 +65,7 @@ jobs:
           sudo bash -x dependencies.sh -y
           mkdir build
           cd build
-          cmake .. -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.0;2.9.1;2.10.0" -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
+          cmake .. -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.0;2.9.1;2.10.0;2.11.0" -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
         shell: bash
 
       - name: Build project

.github/workflows/release.yaml

@@ -64,7 +64,7 @@ jobs:
           sudo bash -x dependencies.sh -y
           mkdir build
           cd build
-          cmake .. -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.0;2.9.1;2.10.0" -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
+          cmake .. -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.0;2.9.1;2.10.0;2.11.0" -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
         shell: bash
 
       - name: Build project

.gitignore

@@ -5,6 +5,7 @@ build_ofed4
 old
 local_test
 go.sum
+!mooncake-common/etcd/go.sum
 *.so
 bin
 mod

.pre-commit-config.yaml

@@ -23,6 +23,16 @@ repos:
       - id: check-added-large-files
         args: ['--maxkb=1024']
 
+  - repo: local
+    hooks:
+      - id: mooncake-code-format
+        name: Run Mooncake code format script
+        entry: ./scripts/code_format.sh
+        language: system
+        pass_filenames: false
+        always_run: true
+        require_serial: true
+
   - repo: https://github.com/astral-sh/ruff-pre-commit
     rev: v0.6.9
     hooks:
@@ -37,7 +47,7 @@ repos:
     hooks:
       - id: codespell
         exclude: '^(extern/|FAST25-release/)'
-        args: ['--ignore-words-list=te,mooncake,KVCache']
+        args: ['--ignore-words-list=te,mooncake,KVCache,cann']
 
   - repo: https://github.com/pre-commit/mirrors-clang-format
     rev: v20.1.8

.typos.toml

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

CMakeLists.txt

@@ -17,10 +17,11 @@ option(WITH_STORE "build mooncake store library and sample code" ON)
 option(WITH_STORE_GO "build Go bindings for mooncake store" OFF)
 option(WITH_P2P_STORE "build p2p store library and sample code" OFF)
 option(WITH_RUST_EXAMPLE "build the Rust interface and sample code for the transfer engine" OFF)
+option(WITH_STORE_RUST "build the Rust bindings for the Mooncake Store" ON)
 option(WITH_EP "build mooncake with expert parallelism support" OFF)
 
+include(${CMAKE_CURRENT_SOURCE_DIR}/mooncake-common/SetupPython.cmake)
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extern/pybind11)
-set(PYTHON_EXECUTABLE "python3")
 execute_process(
     COMMAND ${PYTHON_EXECUTABLE} -c "import sys; print(sys.path[-1])"
     OUTPUT_VARIABLE PYTHON_SYS_PATH
@@ -44,12 +45,21 @@ option(STORE_USE_REDIS "build mooncake store with redis" OFF)
 if (STORE_USE_REDIS)
   add_compile_definitions(STORE_USE_REDIS)
 endif()
+option(STORE_USE_K8S_LEASE "build mooncake store with K8s Lease leader election" OFF)
+if (STORE_USE_K8S_LEASE)
+  if (STORE_USE_ETCD)
+    message(FATAL_ERROR "STORE_USE_K8S_LEASE and STORE_USE_ETCD cannot be enabled together because both build Go c-shared HA backends.")
+  endif()
+  if (USE_ETCD AND NOT USE_ETCD_LEGACY)
+    message(FATAL_ERROR "STORE_USE_K8S_LEASE cannot be enabled with non-legacy USE_ETCD because both build Go c-shared libraries in the same process.")
+  endif()
+  add_compile_definitions(STORE_USE_K8S_LEASE)
+endif()
 
 option(STORE_USE_JEMALLOC "Use jemalloc in mooncake store master" OFF)
 
-# Define ASIO macros before adding mooncake-asio subdirectory
+# Define ASIO macros before building targets that include ASIO headers.
 add_compile_definitions(ASIO_SEPARATE_COMPILATION ASIO_DYN_LINK)
-add_subdirectory(mooncake-asio)
 
 add_subdirectory(mooncake-common)
 include_directories(mooncake-common/etcd)
@@ -66,6 +76,14 @@ if (WITH_STORE)
   include_directories(mooncake-store/include)
 endif()
 
+if (WITH_STORE_RUST)
+  if (NOT WITH_STORE)
+    message(FATAL_ERROR "WITH_STORE_RUST=ON requires WITH_STORE=ON")
+  endif()
+  message(STATUS "Mooncake Store Rust bindings will be built")
+  add_subdirectory(mooncake-store/rust)
+endif()
+
 option(EP_USE_IDE "Enable intelligent indexing for IDEs" OFF)
 if (WITH_EP)
   if (EP_USE_IDE)
@@ -113,6 +131,7 @@ if (WITH_EP)
       COMMAND ${CMAKE_COMMAND}
         "-DSOURCE_DIR=${CMAKE_CURRENT_SOURCE_DIR}/mooncake-ep"
         "-DEP_CUDA_MAJOR=${CUDAToolkit_VERSION_MAJOR}"
+        "-DEP_CUDA_MINOR=${CUDAToolkit_VERSION_MINOR}"
         "-DEP_TORCH_VERSIONS=${_ep_torch_versions_pipe}"
         "-DTORCH_CUDA_ARCH_LIST=${_torch_cuda_arch_list_pipe}"
         "-DSTAGING_DIR=${EP_PG_STAGING_DIR}"
@@ -128,6 +147,7 @@ if (WITH_EP)
       COMMAND ${CMAKE_COMMAND}
         "-DSOURCE_DIR=${CMAKE_CURRENT_SOURCE_DIR}/mooncake-pg"
         "-DEP_CUDA_MAJOR=${CUDAToolkit_VERSION_MAJOR}"
+        "-DEP_CUDA_MINOR=${CUDAToolkit_VERSION_MINOR}"
         "-DEP_TORCH_VERSIONS=${_ep_torch_versions_pipe}"
         "-DTORCH_CUDA_ARCH_LIST=${_torch_cuda_arch_list_pipe}"
         "-DSTAGING_DIR=${EP_PG_STAGING_DIR}"

CONTRIBUTING.md

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

README.md

@@ -97,7 +97,7 @@ Mooncake establishes a full-stack, Tensor-oriented AI infrastructure where Tenso
 
 ### Use Transfer Engine Standalone ([Guide](https://kvcache-ai.github.io/Mooncake/design/transfer-engine/index.html))
 
-Transfer Engine is a high-performance data transfer framework. Transfer Engine provides a unified interface to transfer data from DRAM, VRAM or NVMe, while the technical details related to hardware are hidden. Transfer Engine supports multiple communication protocols including TCP, RDMA (InfiniBand/RoCEv2/eRDMA/NVIDIA GPUDirect), NVMe over Fabric (NVMe-of), NVLink, HIP, CXL, and Ascend. For a complete list of supported protocols and configuration guide, see the [Supported Protocols Documentation](https://kvcache-ai.github.io/Mooncake/getting_started/supported-protocols.html).
+Transfer Engine is a high-performance data transfer framework. Transfer Engine provides a unified interface to transfer data from DRAM, VRAM or NVMe, while the technical details related to hardware are hidden. Transfer Engine supports multiple communication protocols including TCP, RDMA (InfiniBand/RoCEv2/eRDMA/NVIDIA GPUDirect), NVMe over Fabric (NVMe-of), NVLink, HIP, CXL, and Ascend. When built with the corresponding runtime, Transfer Engine can also detect and route accelerator memory on CUDA, MUSA, HIP, and Cambricon MLU devices. For a complete list of supported protocols and configuration guide, see the [Supported Protocols Documentation](https://kvcache-ai.github.io/Mooncake/getting_started/supported-protocols.html).
 
 #### Highlights
 - **Efficient use of multiple RDMA NIC devices.** Transfer Engine supports the use of multiple RDMA NIC devices to achieve the *aggregation of transfer bandwidth*.
@@ -178,6 +178,7 @@ The following need to be installed before running any component of Mooncake:
 - RDMA Driver & SDK, such as Mellanox OFED.
 - Python 3.10, virtual environment is recommended.
 - CUDA 12.1 and above, including NVIDIA GPUDirect Storage Support, if the package is built with `-DUSE_CUDA` (disabled by default). *You may install them from [here](https://developer.nvidia.com/cuda-downloads)*.
+- Cambricon Neuware, if the package is built with `-DUSE_MLU`. By default Mooncake looks for Neuware under `NEUWARE_HOME` or `/usr/local/neuware`.
 
 ### Use Python package
 The simplest way to use Mooncake Transfer Engine is using `pip`:
@@ -201,6 +202,7 @@ pip install mooncake-transfer-engine-non-cuda
 > [!IMPORTANT]
 > - The CUDA version (`mooncake-transfer-engine`) includes Mooncake-EP and GPU topology detection, requiring CUDA 12.1+.
 > - The non-CUDA version (`mooncake-transfer-engine-non-cuda`) is for environments without CUDA dependencies.
+> - MLU support is currently available through source builds with `-DUSE_MLU=ON`; there is no dedicated prebuilt MLU wheel yet.
 > - If users encounter problems such as missing `lib*.so`, they should uninstall the package they installed and build the binaries manually.
 
 ### Use Docker image
@@ -229,6 +231,7 @@ The following are additional dependencies for building Mooncake:
 - Build essentials, including gcc, g++ (9.4+) and cmake (3.16+).
 - Go 1.20+, if you want to build with `-DWITH_P2P_STORE`, `-DUSE_ETCD` (enabled by default to use etcd as metadata servers), or `-DSTORE_USE_ETCD` (use etcd for the failover of the store master).
 - CUDA 12.1 and above, including NVIDIA GPUDirect Storage Support, if the package is built with `-DUSE_CUDA`. *This is NOT included in the `dependencies.sh` script. You may install them from [here](https://developer.nvidia.com/cuda-downloads)*.
+- Cambricon Neuware, if you want to build with `-DUSE_MLU`. *This is NOT included in the `dependencies.sh` script.* Mooncake resolves it from `NEUWARE_HOME` or `/usr/local/neuware` by default, and also supports overriding `MLU_INCLUDE_DIR` / `MLU_LIB_DIR` during CMake configure.
 - [Optional] Rust Toolchain, if you want to build with `-DWITH_RUST_EXAMPLE`. *This is NOT included in the `dependencies.sh` script.*
 - [Optional] `hiredis`, if you want to build with `-DUSE_REDIS` to use Redis instead of etcd as metadata servers.
 - [Optional] `curl`, if you want to build with `-DUSE_HTTP` to use HTTP instead of etcd as metadata servers.
@@ -254,6 +257,14 @@ The build and installation steps are as follows:
    sudo make install # optional, make it ready to be used by vLLM/SGLang
    ```
 
+For Cambricon MLU builds, configure CMake with `-DUSE_MLU=ON`. For example:
+```bash
+mkdir build
+cd build
+cmake .. -DUSE_MLU=ON -DNEUWARE_ROOT=/usr/local/neuware
+make -j
+```
+
 
 <h2 id="milestones"> 🛣️ Incoming Milestones</h2>
 

dependencies.sh

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

docker/mooncake.Dockerfile

@@ -12,6 +12,7 @@ ENV DEBIAN_FRONTEND=noninteractive \
     PYTHONUNBUFFERED=1
 
 ARG PYTHON_VERSION=3.10
+ARG PYPA_INDEX_URL=https://bootstrap.pypa.io
 ARG CMAKE_BUILD_TYPE=Release
 ARG EP_TORCH_VERSIONS="2.9.1"
 ARG TORCH_CUDA_ARCH_LIST="8.0;9.0"
@@ -22,18 +23,25 @@ ENV PYTHON_VERSION=${PYTHON_VERSION} \
     TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} \
     PATH="/usr/local/go/bin:${PATH}"
 
-# Install base build utilities and python bindings
+# Install base build utilities and the requested Python version via deadsnakes PPA
 RUN apt-get update && \
     apt-get install -y --no-install-recommends \
         ca-certificates \
         curl \
         git \
         ninja-build \
-        python3 \
-        python3-dev \
-        python3-pip \
-        python-is-python3 \
+        software-properties-common \
         pkg-config && \
+    add-apt-repository -y ppa:deadsnakes/ppa && \
+    apt-get update && \
+    apt-get install -y --no-install-recommends \
+        python${PYTHON_VERSION} \
+        python${PYTHON_VERSION}-dev \
+        python${PYTHON_VERSION}-venv && \
+    curl -sS ${PYPA_INDEX_URL}/get-pip.py | python${PYTHON_VERSION} && \
+    update-alternatives --install /usr/bin/python  python  /usr/bin/python${PYTHON_VERSION} 1 && \
+    update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 && \
+    apt-get purge -y --auto-remove software-properties-common && \
     rm -rf /var/lib/apt/lists/*
 
 WORKDIR /workspace
@@ -52,6 +60,7 @@ RUN mkdir -p build && \
         -DUSE_CUDA=ON \
         -DWITH_EP=ON \
         -DSTORE_USE_ETCD=ON \
+        -DPython3_EXECUTABLE=/usr/bin/python${PYTHON_VERSION} \
         -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE} && \
     export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH && \
     cmake --build .
@@ -76,11 +85,17 @@ ENV DEBIAN_FRONTEND=noninteractive \
     PYTHONUNBUFFERED=1 \
     PIP_NO_CACHE_DIR=1
 
-# Install runtime dependencies required by Mooncake
+# Inherit build-args so the runtime stage installs the matching interpreter
+ARG PYTHON_VERSION=3.10
+ARG PYPA_INDEX_URL=https://bootstrap.pypa.io
+ENV PYTHON_VERSION=${PYTHON_VERSION}
+
+# Install runtime dependencies and the requested Python version
 RUN apt-get update && \
     apt-get install -y --no-install-recommends \
-        python3 \
-        python3-pip \
+        ca-certificates \
+        curl \
+        software-properties-common \
         ibverbs-providers \
         rdma-core \
         libibverbs1 \
@@ -89,10 +104,18 @@ RUN apt-get update && \
         liburing2 \
         libyaml-0-2 \
         libcurl4 && \
+    add-apt-repository -y ppa:deadsnakes/ppa && \
+    apt-get update && \
+    apt-get install -y --no-install-recommends \
+        python${PYTHON_VERSION} && \
+    curl -sS ${PYPA_INDEX_URL}/get-pip.py | python${PYTHON_VERSION} && \
+    update-alternatives --install /usr/bin/python  python  /usr/bin/python${PYTHON_VERSION} 1 && \
+    update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 && \
+    apt-get purge -y --auto-remove software-properties-common curl && \
     rm -rf /var/lib/apt/lists/*
 
 # Copy wheels produced in builder stage and install them via pip
 COPY --from=builder /workspace/mooncake-wheel/dist /tmp/mooncake-wheel
-RUN python3 -m pip install --no-cache-dir /tmp/mooncake-wheel/*.whl && rm -rf /tmp/mooncake-wheel /root/.cache/pip
+RUN python${PYTHON_VERSION} -m pip install --no-cache-dir /tmp/mooncake-wheel/*.whl && rm -rf /tmp/mooncake-wheel /root/.cache/pip
 
 CMD ["/bin/bash"]

docs/source/deployment/ssd-offload.md

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

docs/source/design/mooncake-store.md

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

docs/source/design/ssd-offload.md

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

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

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

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

@@ -0,0 +1,293 @@
+# Kunpeng UB Transport for Mooncake
+
+This document describes how to build and use Mooncake with Kunpeng UB (Unified Bus) transport support using URMA (Unified Remote Memory Access).
+
+## Overview
+
+UB (Unified Bus) is a transport protocol at the same abstraction layer as RDMA, CXL, NVLink, and TCP, providing a flexible transport solution that can be selected at the application layer. Currently, UB protocol has two open-source implementations:
+
+- **URMA (Unified Remote Memory Access)**: Provides a unified programming abstraction and core semantic layer for upper-layer applications. It offers unified APIs and semantic interfaces for remote shared memory access and operations, leveraging the low-latency, high-bandwidth characteristics of the UB protocol.
+  - URMA open-source repository: https://atomgit.com/openeuler/umdk
+
+- **OBMM (Ownership Based Memory Management)**: A kernel memory management system for supernode environments, supporting cross-node physical memory sharing. It provides efficient remote memory access capabilities through a kernel module (obmm.ko) and a user-space library (libobmm.so).
+  - OBMM open-source repository: https://atomgit.com/openeuler/obmm
+
+## Prerequisites
+
+### 1. Hardware and Operating System
+
+- **Hardware Platform**: Kunpeng 950 CPU with native UB interconnect architecture
+- **OS Version**: openEuler 24.03 (LTS-SP3) [Download link](https://www.openeuler.openatom.cn/zh/download/#openEuler%2024.03%20LTS%20SP3)
+
+### 2. URMA Dependencies
+
+Install UMDK (URMA development package):
+
+```bash
+# Install via yum
+yum install umdk-urma-devel
+
+# Or build from source
+git clone https://atomgit.com/openeuler/umdk.git
+cd umdk
+mkdir build && cd build
+cmake ..
+make -j$(nproc)
+sudo make install
+```
+
+### 3. Build Dependencies
+
+```bash
+# Ubuntu/Debian
+sudo apt-get update
+sudo apt-get install -y \
+    build-essential \
+    cmake \
+    git \
+    libgflags-dev \
+    libgoogle-glog-dev \
+    libjsoncpp-dev \
+    libnuma-dev \
+    libibverbs-dev \
+    libboost-all-dev \
+    libcurl4-openssl-dev \
+    libgtest-dev \
+    libmsgpack-dev \
+    libxxhash-dev \
+    libyaml-cpp-dev \
+    pybind11-dev \
+    python3-dev
+
+# Install yalantinglibs (required)
+cd /tmp
+git clone https://github.com/alibaba/yalantinglibs.git
+cd yalantinglibs
+mkdir build && cd build
+cmake .. -DCMAKE_INSTALL_PREFIX=/usr/local
+make -j$(nproc)
+sudo make install
+```
+
+## Building Mooncake with UB Support
+
+### 1. Clone the Repository
+
+```bash
+git clone https://github.com/kvcache-ai/Mooncake.git
+cd Mooncake
+git submodule update --init --recursive
+```
+
+### 2. Build with UB Enabled
+
+```bash
+mkdir build && cd build
+
+cmake .. \
+    -DUSE_UB=ON \
+    -DURMA_INCLUDE_DIR=/usr/include \
+    -DURMA_LIBRARY=/usr/lib64/liburma.so \
+    -DCMAKE_BUILD_TYPE=RelWithDebInfo
+
+make -j$(nproc)
+```
+
+### 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/
+
+# Install with pip
+pip install -e ../mooncake-wheel --no-build-isolation
+```
+
+## Verification
+
+### Check UB Transport Registration
+
+```bash
+# Check if UB transport is registered
+./mooncake_server --list-transports
+# Expected output: rdma, tcp, nvlink, ub
+```
+
+### Test UB Transport Initialization
+
+```python
+from mooncake.engine import TransferEngine
+
+te = TransferEngine()
+result = te.initialize('127.0.0.1', 'P2PHANDSHAKE', 'ub', '')
+print(f'Initialize result: {result}')  # Should be 0
+
+# You should see logs like:
+# URMA module init success
+# found 1 devices.
+# device_name : urma0 EID : 01:02:03:04:05:06:07:08:09:0a:0b:0c:0d:0e:0f:10
+```
+
+## Usage
+
+### Single Node Benchmark Test
+
+```bash
+# Terminal 1: Target (receiver)
+./transfer_engine_bench \
+    --mode=target \
+    --protocol=ub \
+    --device_name=urma0 \
+    --local_server_name=127.0.0.1 \
+    --metadata_server=P2PHANDSHAKE
+
+# Terminal 2: Initiator (sender)
+./transfer_engine_bench \
+    --mode=initiator \
+    --protocol=ub \
+    --device_name=urma0 \
+    --metadata_server=P2PHANDSHAKE \
+    --segment_size=8388608 \
+    --batch_size=1 \
+    --segment_id=127.0.0.1:$PORT
+```
+
+### Multi-device Benchmark Test
+
+```bash
+# Auto-discovery of multiple URMA devices
+./transfer_engine_bench \
+    --protocol=ub \
+    --device_name=urma0,urma1,urma2,urma3
+```
+
+## Unit Tests
+
+Run the UB transport unit tests:
+
+```bash
+./build/mooncake-transfer-engine/tests/ub_transport_test
+```
+
+The test suite includes:
+
+| Test | Description |
+|------|-------------|
+| `MultiWrite` | Multiple write operations |
+| `MultipleRead` | Multiple read operations with data integrity check |
+
+You can also run all unit tests via CTest:
+
+```bash
+cd build && ctest --output-on-failure
+```
+
+Environment variables for test configuration:
+
+```bash
+export MC_METADATA_SERVER=P2PHANDSHAKE     # default
+export MC_LOCAL_SERVER_NAME=127.0.0.1:12345  # default
+```
+
+## Technical Details
+
+### UB Transport Architecture
+
+```
+┌─────────────────────────────────────────────────────┐
+│                   UbTransport                       │
+├─────────────────────────────────────────────────────┤
+│  UrmaContext (per device)                          │
+│  ├── urma_device   (URMA device handle)            │
+│  ├── urma_context  (URMA context)                 │
+│  ├── urma_jfce     (URMA jetty factory create)     │
+│  ├── urma_jfc      (URMA jetty factory send)       │
+│  └── urma_jfr      (URMA jetty factory receive)    │
+├─────────────────────────────────────────────────────┤
+│  UrmaEndpoint (per connection)                     │
+│  ├── urma_jetty    (URMA jetty for communication)  │
+│  ├── local_jetty   (local jetty ID)                │
+│  └── remote_jetty  (remote jetty ID)               │
+└─────────────────────────────────────────────────────┘
+```
+
+### Key Components
+
+1. **UbTransport**: The main transport class that manages URMA resources and endpoints
+2. **UrmaContext**: Represents a URMA device context, handling device initialization and resource management
+3. **UrmaEndpoint**: Represents a connection to a remote peer, handling data transfer operations
+4. **mock_urma_api.cpp**: Mock implementation of URMA API for testing without real URMA hardware
+
+### Protocol Advantages
+
+- **Optimized for Kunpeng**: URMA is specifically optimized for Kunpeng chip on-chip interconnect
+- **RDMA-like Semantics**: Provides similar memory semantics to RDMA
+- **High Performance**: Leverages UB's low-latency, high-bandwidth characteristics
+- **Unified Abstraction**: Offers a unified programming model for remote memory access
+
+## Troubleshooting
+
+### No URMA devices found
+
+```
+UbTransport: No URMA devices found
+```
+
+Solution: Verify URMA is properly installed and devices are available:
+```bash
+# Check URMA installation
+ls /usr/lib64/liburma.so
+ls /usr/include/ub/umdk/urma/urma_api.h
+
+# Check for URMA devices
+urma_admin -l
+```
+
+### URMA initialization failed
+
+```
+URMA module init failed
+```
+
+Solution: Ensure the URMA kernel module is loaded and the device is properly configured:
+```bash
+# Load URMA module
+sudo modprobe urma
+
+# Check module status
+sudo lsmod | grep urma
+
+# Check device status
+urma_admin -l
+```
+
+### Device port inactive
+
+```
+Device urma0 port not active
+```
+
+Solution: Ensure the UB port is properly configured and active:
+```bash
+# Check port status
+urma_admin -p urma0
+```
+
+### Missing liburma.so
+
+```
+cannot find -lurma
+```
+
+Solution: Verify URMA library is installed and in the library path:
+```bash
+export LD_LIBRARY_PATH=/usr/lib64:$LD_LIBRARY_PATH
+```
+
+## Conclusion
+
+Kunpeng UB Transport provides a high-performance, optimized transport solution for Mooncake on Kunpeng 950 CPU platforms. By leveraging the UB protocol's low-latency and high-bandwidth characteristics, it offers comparable performance to RDMA while being specifically tailored for Kunpeng chip architectures.
+
+With proper configuration and tuning, UB Transport can significantly improve the performance of distributed AI workloads, particularly for scenarios involving large-scale parameter transfers and distributed training.

docs/source/getting_started/build.md

@@ -18,6 +18,7 @@ pip install mooncake-transfer-engine-non-cuda
 📦 **Package Details**: [https://pypi.org/project/mooncake-transfer-engine-non-cuda/](https://pypi.org/project/mooncake-transfer-engine-non-cuda/)
 
 > **Note**: The CUDA version includes Mooncake-EP and GPU topology detection, requiring CUDA 12.1+. The non-CUDA version is for environments without CUDA dependencies.
+> **Note**: MLU support is currently source-build only. If you need Cambricon MLU memory support, install Neuware and build with `-DUSE_MLU=ON`.
 
 ## Automatic
 
@@ -112,8 +113,43 @@ pip install mooncake-transfer-engine-non-cuda
     ```bash
     export LIBRARY_PATH=$LIBRARY_PATH:/usr/local/musa/lib
     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/musa/lib
+    ```
 
-4. Install yalantinglibs
+4. If you want to compile Cambricon MLU support, first install the Cambricon Neuware SDK. After that:
+    1) Export `NEUWARE_HOME` or pass `-DNEUWARE_ROOT=/path/to/neuware` to CMake
+    2) Configure `LIBRARY_PATH` and `LD_LIBRARY_PATH` to ensure linking of `cnrt`, `cndrv`, and other Neuware libraries during compilation:
+    ```bash
+    export NEUWARE_HOME=/usr/local/neuware
+    export LIBRARY_PATH=$LIBRARY_PATH:${NEUWARE_HOME}/lib64
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${NEUWARE_HOME}/lib64
+    ```
+
+    If your Neuware installation lives outside the default include/library layout, you can also pass:
+    ```bash
+    cmake .. -DUSE_MLU=ON \
+      -DMLU_INCLUDE_DIR=/path/to/neuware/include \
+      -DMLU_LIB_DIR=/path/to/neuware/lib64
+    ```
+
+    For Cambricon MLU builds, enable the MLU backend explicitly:
+    ```bash
+    cmake .. -DUSE_MLU=ON -DNEUWARE_ROOT=${NEUWARE_HOME:-/usr/local/neuware}
+    make -j
+    ```
+
+5. If you want to compile MetaX (Muxi) MACA support (e.g. C500), install the MACA SDK so headers and libraries are available under `MACA_ROOT` (defaults to `MACA_HOME` env var if set, otherwise `/opt/maca`). SDK layouts vary; include both `lib` and `lib64` in runtime paths when needed:
+    ```bash
+    export MACA_HOME=/opt/maca
+    export LIBRARY_PATH=$LIBRARY_PATH:${MACA_HOME}/lib:${MACA_HOME}/lib64
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${MACA_HOME}/lib:${MACA_HOME}/lib64
+    ```
+    Build with `-DUSE_MACA=ON`. Optional overrides:
+    - `-DMACA_ROOT=/path/to/maca`
+    - `-DMACA_INCLUDE_DIR=/path/to/maca/include`
+    - `-DMACA_LIB_DIR=/path/to/maca/lib64`
+    - `-DMACA_RUNTIME_LIBS="mcruntime;mxc-runtime64;rt"` (semicolon-separated CMake list)
+
+6. Install yalantinglibs
     ```bash
     git clone https://github.com/alibaba/yalantinglibs.git
     cd yalantinglibs
@@ -123,7 +159,7 @@ pip install mooncake-transfer-engine-non-cuda
     make install
     ```
 
-5. In the root directory of this project, run the following commands:
+7. In the root directory of this project, run the following commands:
    ```bash
    mkdir build
    cd build
@@ -131,7 +167,7 @@ pip install mooncake-transfer-engine-non-cuda
    make -j
    ```
 
-6. Install Mooncake python package and mooncake_master executable
+8. Install Mooncake python package and mooncake_master executable
    ```bash
    make install
    ```
@@ -151,9 +187,18 @@ cd /Mooncake-main/build/mooncake-transfer-engine/example
 ## Advanced Compile Options
 The following options can be used during `cmake ..` to specify whether to compile certain components of Mooncake.
 - `-DUSE_CUDA=[ON|OFF]`: Enable GPU memory support (GPUDirect RDMA, NVMe-oF, and GPU-aware TCP transport). **Default: OFF.** Required when transferring GPU memory (e.g., KV cache in vLLM disaggregated serving), even when using TCP protocol.
-- `-DUSE_MNNVL=[ON|OFF]`: Enable Multi-Node NVLink transport support, default is OFF. **Note:** `-DUSE_CUDA` is required when `-DUSE_MNNVL` is on.
+- `-DUSE_MNNVL=[ON|OFF]`: Enable Multi-Node NVLink transport support, default is OFF. **Note:** `-DUSE_CUDA` is required when `-DUSE_MNNVL` is on (not used when building with `-DUSE_MUSA=ON`, `-DUSE_HIP=ON`, or `-DUSE_MACA=ON`).
 - `-DUSE_MUSA=[ON|OFF]`: Enable Moore Threads GPU support via MUSA
+- `-DUSE_MACA=[ON|OFF]`: Enable MetaX (Muxi) GPU support via MACA.
+- `-DMACA_ROOT=/path/to/maca`: Override the MACA SDK root (`MACA_HOME` env var is also honored; default `/opt/maca`).
+- `-DMACA_INCLUDE_DIR=/path/to/include`: Override MACA include directory when `-DUSE_MACA=ON`.
+- `-DMACA_LIB_DIR=/path/to/lib64`: Override MACA library directory when `-DUSE_MACA=ON`.
+- `-DMACA_RUNTIME_LIBS="mcruntime;mxc-runtime64;rt"`: Override MACA runtime libraries linked by `transfer_engine`.
 - `-DUSE_HIP=[ON|OFF]`: Enable AMD GPU support via HIP/ROCm
+- `-DUSE_MLU=[ON|OFF]`: Enable Cambricon MLU memory support via Neuware. **Default: OFF.** Supports MLU memory detection, topology discovery, and RDMA registration for Transfer Engine.
+- `-DNEUWARE_ROOT=/path/to/neuware`: Override the default Neuware SDK root used when `-DUSE_MLU=ON`. If unset, Mooncake uses `NEUWARE_HOME` or `/usr/local/neuware`.
+- `-DMLU_INCLUDE_DIR=/path/to/include`: Override the Neuware include directory when `-DUSE_MLU=ON`.
+- `-DMLU_LIB_DIR=/path/to/lib64`: Override the Neuware library directory when `-DUSE_MLU=ON`.
 - `-DUSE_EFA=[ON|OFF]`: Enable AWS Elastic Fabric Adapter transport via libfabric. **Default: OFF.** See [EFA Transport](../design/transfer-engine/efa_transport.md) for details.
 - `-DUSE_INTRA_NVLINK=[ON|OFF]`: Enable intranode nvlink transport
 - `-DUSE_CXL=[ON|OFF]`: Enable CXL support

docs/source/getting_started/supported-protocols.md

@@ -56,7 +56,7 @@ export MOONCAKE_PROTOCOL="tcp"
 
 ### RDMA (Recommended for Production)
 
-**Description:** Remote Direct Memory Access protocol providing high-performance, low-latency data transfer with minimal CPU overhead. Supports GPUDirect RDMA for zero-copy GPU memory transfers.
+**Description:** Remote Direct Memory Access protocol providing high-performance, low-latency data transfer with minimal CPU overhead. Supports accelerator-aware memory registration, including NVIDIA GPUDirect RDMA for CUDA buffers and Cambricon MLU buffers when built with Neuware.
 
 **Hardware Support:**
 - InfiniBand
@@ -64,6 +64,7 @@ export MOONCAKE_PROTOCOL="tcp"
 - eRDMA (Elastic RDMA)
 - NVIDIA GPUDirect RDMA
 - Non-NVIDAI GPUDirect RDMA (e.g., Intel E810 RDMA NIC)
+- Cambricon MLU memory via Neuware (`-DUSE_MLU=ON`)
 
 **Use When:**
 - High-performance networking is required
@@ -72,6 +73,8 @@ export MOONCAKE_PROTOCOL="tcp"
 
 **Note:** If no RDMA HCA (Host Channel Adapter) is detected on the system, the Transfer Engine will automatically fall back to TCP protocol for compatibility.
 
+**MLU Note:** Cambricon MLU support uses the standard `rdma` data path. There is no separate `mlu` protocol string. To enable MLU memory detection, topology discovery, and DMA-BUF based registration, build Transfer Engine with `-DUSE_MLU=ON` and make Neuware available through `NEUWARE_HOME` or `NEUWARE_ROOT`.
+
 **Configuration:**
 ```python
 # Python API - With specific device
@@ -321,6 +324,7 @@ export MOONCAKE_LOCAL_HOSTNAME="node1"
 | Cloud Environments | tcp or rdma (if available) | Check cloud provider support |
 | Multi-tier Storage | rdma + nvmeof | Combine protocols for different layers |
 | AMD GPU Clusters | rdma + hip | Use HIP for local GPU communication |
+| Cambricon MLU Clusters | rdma | Build with `-DUSE_MLU=ON`; MLU uses the normal RDMA protocol |
 | Ascend NPU Clusters | rdma + ascend | Use Ascend for NPU-specific operations |
 
 ## Troubleshooting

docs/source/index.md

@@ -86,6 +86,7 @@ performance/vllm-benchmark-results-v1
 performance/sglang-hicache-benchmark-results-v1
 performance/vllm-v1-support-benchmark
 performance/allocator-benchmark-result
+performance/ssd-offload-benchmark-results
 :::
 
 % API Documentation

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

@@ -0,0 +1,157 @@
+# Mooncake SSD Offload Benchmark
+
+This benchmark measures the performance benefit of Mooncake's SSD offload feature in multi-turn conversation scenarios. In the test, multiple clients send requests concurrently, each simulating a multi-round dialogue where every new round appends the previous context.
+
+We compare four storage configurations for the KV cache:
+
+* **GPU only**: KV cache resides entirely in GPU memory.
+* **(HiCache L1) + L2**: KV cache spans GPU and host memory via HiCache's two-level hierarchy.
+* **(HiCache L1 + L2) + Mooncake**: KV cache is further extended into an 80GB Mooncake distributed memory pool.
+* **(HiCache L1 + L2) + Mooncake + SSD**: On top of the above, SSD offload is enabled so that evicted cache entries are written to local NVMe storage rather than discarded.
+
+The benchmark targets the prefill stage and reports two primary metrics: Time-To-First-Token (TTFT) and input token throughput.
+
+## Benchmark Result
+
+![overall performance](../image/ssd_offload_overall.png)
+
+The figure above summarizes the end-to-end results on a single DGX node (8 × A100-SXM4-40GB, dual RDMA NICs). Enabling SSD offload cuts average TTFT by **57%** relative to GPU only and by **34%** relative to Mooncake without SSD, while delivering a **2.4×** improvement in input token throughput.
+
+![per-turn performance](../image/ssd_offload_per_turn.png)
+
+To better understand where the gains come from, we break down TTFT and cache hit rate by conversation round. The output length is fixed to 1 token so that decode overhead does not obscure prefill differences.
+
+During the first six rounds the 80GB memory pool is large enough, so `+ Mooncake` and `+ Mooncake + SSD` behave identically — both sustain hit rates above 80%.
+
+The divergence appears in round 7. Once the accumulated KV cache exceeds memory capacity, `+ Mooncake` must evict entries and its hit rate plunges from 83% to 36%, pushing TTFT from 6s to 16s. With SSD offload, those evicted entries survive on disk and remain retrievable; the hit rate stays above 84% through round 8, and TTFT remains at 9.4s — roughly half the latency of Mooncake without SSD.
+
+Note that a slight increase in TTFT is visible in round 8 with SSD offload (9.4s vs 7.4s in round 7), reflecting the additional latency of reading evicted entries from NVMe storage rather than RDMA memory. This overhead is modest compared to the alternative of re-computing evicted KV cache from scratch.
+
+This demonstrates that SSD offload turns local NVMe drives into a cost-effective extension of the cache hierarchy. In production, where long conversations and high concurrency are common, this prevents the sharp performance cliff that occurs when DRAM-based caching alone is exhausted.
+
+## Benchmark Setup
+
+### DGX Server
+
+**Experimental Environment**
+
+- GPU: 8 × NVIDIA A100-SXM4-40GB
+- Network: Dual RDMA NICs (ibp12s0, ibp75s0), InfiniBand 4X HDR 200 Gb/s each
+- Storage: 5 × Samsung NVMe SSDs in RAID0 — 3 × PM1733 3.84TB (PCIe Gen4, 7,000 MB/s seq read each) + 2 × PM983 1.92TB (PCIe Gen3, 3,000 MB/s seq read each). Aggregate theoretical sequential read bandwidth: ~27 GB/s. Mounted at /mnt/data (~14TB usable), used as the SSD offload target.
+- Model: Qwen3-8B
+
+**Benchmark Script:**
+
+We used SGLang's [multiturn benchmark](https://github.com/sgl-project/sglang/blob/main/benchmark/hicache/bench_multiturn.py) for the evaluation.
+
+```bash
+python3 benchmark/hicache/bench_multiturn.py \
+    --model-path $MODEL_PATH \
+    --host 127.0.0.1 \
+    --port 8189 \
+    --disable-random-sample \
+    --output-length 1 \
+    --request-length 4096 \
+    --num-clients 20 \
+    --num-rounds 10 \
+    --max-parallel 4 \
+    --request-rate 16 \
+    --ready-queue-policy random \
+    --disable-auto-run \
+    --enable-round-barrier
+```
+
+**GPU Only:**
+
+```bash
+python3 -m sglang.launch_server \
+    --model-path $MODEL_PATH \
+    --tp 1 \
+    --page-size 64 \
+    --attention-backend triton
+```
+
+**HiCache L1 + L2:**
+
+```bash
+python3 -m sglang.launch_server \
+    --model-path $MODEL_PATH \
+    --tp 1 \
+    --page-size 64 \
+    --attention-backend triton \
+    --enable-hierarchical-cache \
+    --hicache-ratio 2
+```
+
+**L1 + L2 + Mooncake:**
+
+Mooncake master and client must be started before launching the SGLang server.
+
+```bash
+# Start Mooncake master
+mooncake_master \
+    -http_metadata_server_port=8081 \
+    -metrics_port=9004 \
+    -logtostderr
+
+# Start Mooncake client (requires root)
+# Total Distributed Memory Pool: 80GB
+mooncake_client \
+    --host=127.0.0.1 \
+    --global_segment_size=80GB \
+    --master_server_address=localhost:50051 \
+    --metadata_server=P2PHANDSHAKE \
+    --protocol=rdma \
+    --device_names=ibp12s0,ibp75s0 \
+    --port=50052 \
+    --logtostderr
+```
+
+```bash
+MOONCAKE_MASTER="127.0.0.1:50051" \
+MOONCAKE_GLOBAL_SEGMENT_SIZE=0 \
+MOONCAKE_PROTOCOL="rdma" \
+MOONCAKE_DEVICE="ibp12s0,ibp75s0" \
+python3 -m sglang.launch_server \
+    --model-path $MODEL_PATH \
+    --tp 1 \
+    --page-size 64 \
+    --attention-backend triton \
+    --enable-hierarchical-cache \
+    --hicache-ratio 2 \
+    --hicache-storage-prefetch-policy wait_complete \
+    --hicache-mem-layout page_first_direct \
+    --hicache-storage-backend mooncake
+```
+
+**L1 + L2 + Mooncake + SSD:**
+
+Compared to the previous configuration, the only change is enabling SSD offload on both master and client. A 20GB local buffer absorbs write bursts before flushing to SSD.
+
+```bash
+# Start Mooncake master with offload enabled
+mooncake_master \
+    -enable_offload=true \
+    -http_metadata_server_port=8081 \
+    -metrics_port=9004 \
+    -logtostderr
+
+# Start Mooncake client with offload enabled (requires root)
+# Total Distributed Memory Pool: 80GB
+# 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_client \
+    --host=127.0.0.1 \
+    --global_segment_size=80GB \
+    --master_server_address=localhost:50051 \
+    --metadata_server=P2PHANDSHAKE \
+    --protocol=rdma \
+    --device_names=ibp12s0,ibp75s0 \
+    --enable_offload=true \
+    --port=50052 \
+    --logtostderr
+```
+
+The SGLang server launch command is identical to `L1 + L2 + Mooncake`.

docs/source/python-api-reference/mooncake-store.md

@@ -264,6 +264,151 @@ def get_into(self, key: str, buffer_ptr: int, size: int) -> int
 
 **Returns:** Number of bytes read, or negative on error
 
+#### get_into_ranges()
+Retrieve multiple byte ranges from multiple objects into registered buffers (zero-copy).
+
+```python
+def get_into_ranges(self, buffer_ptrs: List[int], all_keys: List[List[str]], all_dst_offsets: List[List[List[int]]], all_src_offsets: List[List[List[int]]], all_sizes: List[List[List[int]]]) -> List[List[List[int]]]
+```
+
+This API is **buffer-major** and supports **multiple fragments per key**.
+
+Think of the input shape as:
+- `buffer_ptrs[i]`: the `i`-th destination buffer
+- `all_keys[i][j]`: the `j`-th key that writes into buffer `i`
+- `all_dst_offsets[i][j][k]`: destination offset of fragment `k` for key `j` in buffer `i`
+- `all_src_offsets[i][j][k]`: source offset of fragment `k` inside key `j` for buffer `i`
+- `all_sizes[i][j][k]`: byte size of fragment `k`
+
+For each triple `(i, j, k)`, Mooncake reads the source range
+`[all_src_offsets[i][j][k], all_src_offsets[i][j][k] + all_sizes[i][j][k])`
+from object `all_keys[i][j]`, then writes it into destination buffer
+`buffer_ptrs[i]` at offset `all_dst_offsets[i][j][k]`.
+
+This lets one buffer gather interleaved fragments from multiple keys, and lets one key contribute multiple disjoint fragments to the same buffer in a single call.
+
+**Parameters:**
+- `buffer_ptrs`: Memory addresses of pre-allocated destination buffers. Every buffer must be registered with `register_buffer()` before calling this API.
+- `all_keys`: For each buffer, the ordered list of source object keys to read from.
+- `all_dst_offsets`: For each buffer and key, the destination offsets of that key's fragments.
+- `all_src_offsets`: For each buffer and key, the source offsets of that key's fragments inside the object.
+- `all_sizes`: For each buffer and key, the byte lengths of that key's fragments.
+
+**Shape rules:**
+- `len(buffer_ptrs) == len(all_keys) == len(all_dst_offsets) == len(all_src_offsets) == len(all_sizes)`
+- For each buffer `i`, `len(all_keys[i]) == len(all_dst_offsets[i]) == len(all_src_offsets[i]) == len(all_sizes[i])`
+- For each `(buffer i, key j)`, `len(all_dst_offsets[i][j]) == len(all_src_offsets[i][j]) == len(all_sizes[i][j])`
+
+If a top-level shape or per-key fragment shape does not match, the corresponding result entries are negative error codes.
+
+**Returns:** A nested list of per-buffer, per-key, per-fragment results. `results[i][j][k]` is the number of bytes read for fragment `k`, or a negative value on error.
+
+A successful call can still contain per-fragment failures. For example, if one key is missing but another key in the same buffer is valid, the missing key's fragment result will be negative while the valid fragment can still succeed.
+
+**Typical scenarios:**
+- **Partial read from one object:** You only need a slice of a large value, such as a header, metadata block, or a small subrange of a tensor shard. In this case, use one buffer, one key, and one or more fragments under that key.
+- **Stitch multiple fragments from one object into one buffer:** You need several non-contiguous ranges from the same object and want to pack them into one destination buffer. In this case, keep a single key entry and place multiple fragments under that key.
+- **Stitch data from multiple objects into one buffer:** You want to assemble one logical payload from several keys. In this case, use one destination buffer and list multiple keys under that buffer, with each key contributing one or more fragments.
+- **Fill multiple output buffers in one call:** You have several destination buffers, each with its own read plan. In this case, each top-level entry in `buffer_ptrs` and the parallel nested arrays describes one independent destination buffer.
+
+**How to use it for partial reads:**
+If you only want part of an object, do not call `get_into()` with the full object buffer size. Instead:
+1. Allocate and register a destination buffer sized for the bytes you actually want to materialize.
+2. Put that buffer pointer into `buffer_ptrs`.
+3. Put the source key into `all_keys`.
+4. Set `all_src_offsets` to the start offsets of the object ranges you want.
+5. Set `all_sizes` to the lengths of those ranges.
+6. Set `all_dst_offsets` to where those ranges should land in your destination buffer.
+
+A useful way to think about the arguments is:
+- `buffer_ptrs` answers **where does the data land**
+- `all_keys` answers **which object does it come from**
+- `all_src_offsets` and `all_sizes` answer **which bytes should be read**
+- `all_dst_offsets` answers **where each fragment should be placed in the destination buffer**
+
+If you are extracting a single contiguous slice from one object, the minimal shape is:
+
+```python
+results = store.get_into_ranges(
+    [buffer_ptr],
+    [["my_key"]],
+    [[[0]]],
+    [[[src_offset]]],
+    [[[size]]],
+)
+```
+
+This means:
+- one destination buffer
+- one source key for that buffer
+- one fragment for that key
+- read `size` bytes from `my_key[src_offset:src_offset + size]`
+- write them into `buffer_ptr[0:size]`
+
+If you want to read several disjoint ranges from the same object and pack them together, keep the same key and add more fragments under it. For example:
+
+```python
+results = store.get_into_ranges(
+    [buffer_ptr],
+    [["my_key"]],
+    [[[0, 16, 40]]],
+    [[[128, 4096, 8192]]],
+    [[[8, 12, 4]]],
+)
+```
+
+This reads three fragments from `my_key` and places them into the same destination buffer at offsets `0`, `16`, and `40`. This pattern is useful when you want to assemble only the needed pieces of a large object without reading the whole value.
+
+If you want to assemble one output buffer from multiple objects, keep one top-level buffer entry and add multiple keys under it. Each key can still contribute one or more fragments. For example, you might put a header from `meta_key` at the front of the buffer, then place a payload slice from `data_key` after it.
+
+**Usage example:**
+
+```python
+import ctypes
+
+buffer_size = 32
+buffer0 = (ctypes.c_ubyte * buffer_size)()
+buffer1 = (ctypes.c_ubyte * buffer_size)()
+buffer_ptr0 = ctypes.addressof(buffer0)
+buffer_ptr1 = ctypes.addressof(buffer1)
+
+store.register_buffer(buffer_ptr0, buffer_size)
+store.register_buffer(buffer_ptr1, buffer_size)
+
+# Buffer 0 reads:
+# - from key1: two fragments -> src[1:5] -> dst[0:4], src[30:33] -> dst[20:23]
+# - from key2: one fragment  -> src[2:7] -> dst[8:13]
+# Buffer 1 reads:
+# - from key2: one fragment  -> src[0:6] -> dst[4:10]
+# - from key1: one fragment  -> src[10:14] -> dst[16:20]
+results = store.get_into_ranges(
+    [buffer_ptr0, buffer_ptr1],
+    [["key1", "key2"], ["key2", "key1"]],
+    [[[0, 20], [8]], [[4], [16]]],
+    [[[1, 30], [2]], [[0], [10]]],
+    [[[4, 3], [5]], [[6], [4]]],
+)
+
+# results == [
+#   [[4, 3], [5]],
+#   [[6], [4]],
+# ]
+```
+
+In the example above:
+- `results[0][0][0] == 4`: buffer 0, key 0 (`"key1"`), fragment 0 succeeded with 4 bytes
+- `results[0][0][1] == 3`: buffer 0, key 0 (`"key1"`), fragment 1 succeeded with 3 bytes
+- `results[0][1][0] == 5`: buffer 0, key 1 (`"key2"`), fragment 0 succeeded with 5 bytes
+
+**Common pitfalls:**
+- Do not flatten all fragments for a buffer into one list. Fragments must be grouped under their corresponding key.
+- `all_dst_offsets`, `all_src_offsets`, and `all_sizes` are 3D, but `all_keys` is 2D.
+- Buffer overflow is checked against the registered destination buffer size.
+- Source overflow is checked against the source object's size.
+- Full-object `get_into()` and ranged `get_into_ranges()` are different APIs; use `get_into()` when you want the whole object into one buffer.
+
+**Current limitation:** true ranged items currently require the selected source replica to be memory-backed. Whole-object reads still follow the normal full-read path, but partial reads through `get_into_ranges()` do not support non-memory replicas.
+
 ---
 
 ## ReplicateConfig Configuration
@@ -301,6 +446,16 @@ config = ReplicateConfig()
 config.with_soft_pin = True  # Keep this object in memory longer
 ```
 
+#### with_hard_pin
+**Type:** `bool`
+**Default:** `False`
+**Description:** Enables hard pinning for the stored object. Hard pinned objects will not be evicted. This grants user to manually control the life time of stored objects.
+
+```python
+config = ReplicateConfig()
+config.with_hard_pin = True  # Keep this object in memory that will not be evicted
+```
+
 #### preferred_segment
 **Type:** `str`
 **Default:** `""` (empty string)

docs/source/zh_archive/build.md

@@ -110,7 +110,41 @@
     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/musa/lib
     ```
 
-4. 安装 yalantinglibs
+4. 若需编译寒武纪 MLU 支持，请先安装寒武纪 Neuware SDK。之后：
+    1) 导出 `NEUWARE_HOME`，或在 CMake 中传入 `-DNEUWARE_ROOT=/path/to/neuware`
+    2) 配置 `LIBRARY_PATH` 与 `LD_LIBRARY_PATH`，确保编译时能链接 `cnrt`、`cndrv` 等 Neuware 库：
+    ```bash
+    export NEUWARE_HOME=/usr/local/neuware
+    export LIBRARY_PATH=$LIBRARY_PATH:${NEUWARE_HOME}/lib64
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${NEUWARE_HOME}/lib64
+    ```
+
+    若 Neuware 安装路径与默认头文件/库布局不一致，还可显式指定：
+    ```bash
+    cmake .. -DUSE_MLU=ON \
+      -DMLU_INCLUDE_DIR=/path/to/neuware/include \
+      -DMLU_LIB_DIR=/path/to/neuware/lib64
+    ```
+
+    启用 MLU 后端示例：
+    ```bash
+    cmake .. -DUSE_MLU=ON -DNEUWARE_ROOT=${NEUWARE_HOME:-/usr/local/neuware}
+    make -j
+    ```
+
+5. 若需编译沐曦 MetaX MACA 支持（如 C500），请安装 MACA SDK，使头文件与库位于 `MACA_ROOT`（优先取 `MACA_HOME` 环境变量，未设置时默认 `/opt/maca`）。不同安装包可能把库放在 `lib` 或 `lib64`，建议在环境变量中同时加入两者，避免链接或运行时找不到共享库：
+    ```bash
+    export MACA_HOME=/opt/maca
+    export LIBRARY_PATH=$LIBRARY_PATH:${MACA_HOME}/lib:${MACA_HOME}/lib64
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${MACA_HOME}/lib:${MACA_HOME}/lib64
+    ```
+    使用 `-DUSE_MACA=ON` 配置构建。可选覆盖项：
+    - `-DMACA_ROOT=/path/to/maca`
+    - `-DMACA_INCLUDE_DIR=/path/to/maca/include`
+    - `-DMACA_LIB_DIR=/path/to/maca/lib64`
+    - `-DMACA_RUNTIME_LIBS="mcruntime;mxc-runtime64;rt"`（分号分隔的 CMake 列表）
+
+6. 安装 yalantinglibs
     ```bash
     git clone https://github.com/alibaba/yalantinglibs.git
     cd yalantinglibs
@@ -120,7 +154,7 @@
     make install
     ```
 
-5. 进入项目根目录，运行下列命令进行编译
+7. 进入项目根目录，运行下列命令进行编译
    ```bash
    mkdir build
    cd build
@@ -128,7 +162,7 @@
    make -j
    ```
 
-6. 安装 Mooncake python 包和 mooncake_master 可执行文件
+8. 安装 Mooncake python 包和 mooncake_master 可执行文件
    ```bash
    make install
    ```
@@ -137,6 +171,14 @@
 在执行 `cmake ..` 期间可以使用下列选项指定是否编译 Mooncake 的某些组件。
 - `-DUSE_CUDA=[ON|OFF]`: 启用 GPU Direct RDMA 及 NVMe-of 支持
 - `-DUSE_MUSA=[ON|OFF]`: 通过 MUSA 启用对摩尔线程 GPU 的支持
+- `-DUSE_MACA=[ON|OFF]`: 通过 MACA 启用对沐曦 MetaX GPU 的支持。
+- `-DMACA_ROOT=/path/to/maca`: 覆盖 MACA SDK 根路径（也支持 `MACA_HOME` 环境变量，默认 `/opt/maca`）。
+- `-DMACA_INCLUDE_DIR=/path/to/include`: 在 `-DUSE_MACA=ON` 时覆盖 MACA 头文件目录。
+- `-DMACA_LIB_DIR=/path/to/lib64`: 在 `-DUSE_MACA=ON` 时覆盖 MACA 库目录。
+- `-DMACA_RUNTIME_LIBS="mcruntime;mxc-runtime64;rt"`: 覆盖 `transfer_engine` 链接的 MACA 运行时库列表。
+- `-DUSE_MLU=[ON|OFF]`: 通过 Neuware 启用寒武纪 MLU 显存支持。默认 OFF；支持 MLU 显存探测、拓扑发现及 Transfer Engine 的 RDMA 注册。
+- `-DNEUWARE_ROOT=/path/to/neuware`: 在 `-DUSE_MLU=ON` 时覆盖默认 Neuware SDK 根路径；未设置时使用 `NEUWARE_HOME` 或 `/usr/local/neuware`。
+- `-DMLU_INCLUDE_DIR=/path/to/include` / `-DMLU_LIB_DIR=/path/to/lib64`: 在 `-DUSE_MLU=ON` 时覆盖 Neuware 头文件与库目录。
 - `-DUSE_HIP=[ON|OFF]`: 通过 HIP/ROCm 启用对 AMD GPU 的支持
 - `-DUSE_CXL=[ON|OFF]`: 启用 CXL 支持
 - `-DWITH_STORE=[ON|OFF]`: 编译 Mooncake Store 组件

docs/source/zh_archive/kunpeng_ub_transport.md

@@ -0,0 +1,90 @@
+# Kunpeng UB Transport 
+Kunpeng UbTransport源代码路径为Mooncake/mooncake-transfer-engine/src/transport/kunpneg_transport，该路径下有UB协议的Transport对接代码和实现逻辑。
+
+## 概述
+UB（Unified Bus，统一总线） 是与RDMA、CXL、NVLink 和TCP处于同一抽象层的传输协议，属于可在应用层灵活选择的传输方案。目前 UB 协议有两个开源实现：URMA（远程内存访问语义）和 OBMM（Load/Store 语义）。
+
+URMA（Unified Remote Memory Access，统一远程内存访问）是UB协议为上层应用提供的统一编程抽象与核心语义层。它基于 UB 协议低延迟、高带宽的底层特性，为远程共享内存的访问与操作提供统一的 API 和语义接口。
+
+URMA 开源代码仓库：https://atomgit.com/openeuler/umdk
+
+OBMM (Ownership Based Memory Management) 是面向超节点环境的内核内存管理系统，支持跨节点的物理内存共享。该系统通过内核模块 (obmm.ko) 和用户态库 (libobmm.so) 提供高效的远程内存访问能力。
+
+OBMM 开源代码仓库：https://atomgit.com/openeuler/obmm
+
+## 新增依赖
+Kunpeng UbTransport在Mooncake本身依赖的基础上，新增了一部分URMA和OBMM的依赖：
+
+- **硬件平台**: 支持原生UB互联架构的鲲鹏950 CPU
+- **OS版本**: openEuler 24.03 (LTS-SP3) [下载链接](https://www.openeuler.openatom.cn/zh/download/#openEuler%2024.03%20LTS%20SP3)
+- **URMA依赖**: UMDK: `yum install umdk-urma-devel` 或从[源码](https://atomgit.com/openeuler/umdk)构建。
+- **协议优势**: URMA 提供类似 RDMA 的内存语义，针对鲲鹏芯片片上互联进行了优化
+
+---
+
+## 构建与编译
+
+**前置条件**
+
+- openEuler 24.03 (LTS-SP3) [下载链接](https://www.openeuler.openatom.cn/zh/download/#openEuler%2024.03%20LTS%20SP3)
+- 已安装 UMDK: `yum install umdk-urma-devel` 或从[源码](https://atomgit.com/openeuler/umdk)构建
+
+**CMake 配置**
+
+```bash
+# 克隆 Mooncake 仓库
+git clone https://github.com/kvcache-ai/Mooncake.git 
+cd Mooncake
+
+# 启用 UB 传输层进行配置
+mkdir build && cd build
+cmake .. -DUSE_UB=ON \
+         -DURMA_INCLUDE_DIR=/usr/include \
+         -DURMA_LIBRARY=/usr/lib64/liburma.so
+
+# 编译
+make -j$(nproc)
+```
+
+**验证**
+
+```bash
+# 检查 UB 传输层是否已注册
+./mooncake_server --list-transports
+# 预期输出: rdma, tcp, nvlink, ub
+```
+
+---
+
+## 运行与测试
+
+**单节点基准测试**
+
+```bash
+# 终端 1: 目标端（Target）
+./transfer_engine_bench \
+    --mode=target \
+    --protocol=ub \
+    --device_name=urma0 \
+    --local_server_name=127.0.0.1 \
+    --metadata_server=P2PHANDSHAKE
+
+# 终端 2: 发起端（Initiator）
+./transfer_engine_bench \
+    --mode=initiator \
+    --protocol=ub \
+    --device_name=urma0 \
+    --metadata_server=P2PHANDSHAKE \
+    --segment_size=8388608 \
+    --batch_size=1\
+    --segment_id=127.0.0.1:$PORT
+```
+
+**多设备基准测试**
+
+```bash
+# 自动发现多个 URMA 设备
+./transfer_engine_bench \
+    --protocol=ub \
+    --device_name=urma0,urma1,urma2,urma3
+```

mooncake-asio/CMakeLists.txt

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

mooncake-common/CMakeLists.txt

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

mooncake-common/FindUrma.cmake

@@ -0,0 +1,20 @@
+include(FetchContent)
+
+# UMDK 头文件库
+FetchContent_Declare(
+        urma
+        GIT_REPOSITORY https://atomgit.com/openeuler/umdk.git
+        GIT_TAG        v25.12.0
+)
+
+FetchContent_MakeAvailable(urma)
+
+# 输出实际路径，确认位置
+message(STATUS "URMA source dir: ${urma_SOURCE_DIR}")
+message(STATUS "URMA binary dir: ${urma_BINARY_DIR}")
+
+# 假设 UMDK 头文件在其 include 目录下
+set(urma_INCLUDE_DIR ${urma_SOURCE_DIR}/src/urma/lib/urma/core/include)
+
+# 添加到需要的目标
+message(STATUS "urma_INCLUDE_DIR: ${urma_INCLUDE_DIR}")

mooncake-common/SetupPyTorchEnv.cmake

@@ -0,0 +1,64 @@
+# SetupPyTorchEnv.cmake
+#
+# This file provides helper functions for building Mooncake Pytorch extensions
+# and is meant to be included by BuildEpExt.cmake and BuildPgExt.cmake.
+
+# Ensure we have the correct Python interpreter (respects active virtualenvs)
+find_package(Python3 REQUIRED COMPONENTS Interpreter)
+
+# Install PyTorch for a specific version with proper CUDA compatibility handling.
+#
+# Usage:
+#   install_pytorch_wheel("<VERSION>" <CUDA_MAJOR> <CUDA_MINOR> "<MODULE_PREFIX>")
+#
+# Example:
+#   install_pytorch_wheel("2.11.0" 12 8 "[EP]")
+function(install_pytorch_wheel _version _cuda_major _cuda_minor _module_prefix)
+  message(STATUS "${_module_prefix} Installing PyTorch ${_version} via pip...")
+
+  set(_cu_tag "")
+
+  # Determine the specific CUDA tag for PyTorch wheels
+  if(_cuda_major GREATER_EQUAL 13)
+    # TODO: Fix when we need to support more CUDA 13 versions or when the CI env is fixed.
+    set(_cu_tag "cu130")
+
+  elseif(_cuda_major EQUAL 12 AND _version VERSION_GREATER_EQUAL "2.11.0")
+    # PyTorch 2.11.0+ defaults to CUDA 13.
+    # We must explicitly point to CUDA 12 wheels for these newer versions.
+    if(_cuda_minor GREATER_EQUAL 8)
+      set(_cu_tag "cu128")
+    elseif(_cuda_minor GREATER_EQUAL 6)
+      set(_cu_tag "cu126")
+    else()
+      message(FATAL_ERROR
+        "${_module_prefix} Can't find a matching PyTorch wheel for version ${_version} "
+        "with CUDA ${_cuda_major}.${_cuda_minor}"
+      )
+    endif()
+  endif()
+
+  # Construct pip command using the absolute path to the Python executable
+  set(_pip_cmd ${Python3_EXECUTABLE} -m pip install "torch==${_version}")
+
+  if(_cu_tag)
+    set(_index_url "https://download.pytorch.org/whl/${_cu_tag}")
+    message(STATUS "${_module_prefix} Using specific CUDA wheel: ${_index_url}")
+    list(APPEND _pip_cmd --index-url "${_index_url}")
+  else()
+    message(STATUS "${_module_prefix} Using default PyPI wheels for PyTorch ${_version}")
+  endif()
+
+  # Execute pip install
+  execute_process(
+    COMMAND ${_pip_cmd}
+    RESULT_VARIABLE _ret
+  )
+
+  if(NOT _ret EQUAL 0)
+    message(FATAL_ERROR "${_module_prefix} Failed to install PyTorch ${_version}."
+            " Command run: '${_pip_cmd}'")
+  endif()
+
+  message(STATUS "${_module_prefix} PyTorch ${_version} is ready.")
+endfunction()

mooncake-common/SetupPython.cmake

@@ -0,0 +1,10 @@
+# SetupPython.cmake — resolve the Python interpreter for execute_process() calls.
+#
+# Honour -DPython3_EXECUTABLE=... when provided (e.g. Docker builds that
+# install a non-system Python via deadsnakes), otherwise fall back to the
+# default "python3" on PATH.  Sets PYTHON_EXECUTABLE for legacy callers.
+
+if(NOT Python3_EXECUTABLE)
+    set(Python3_EXECUTABLE "python3")
+endif()
+set(PYTHON_EXECUTABLE "${Python3_EXECUTABLE}")

mooncake-common/common.cmake

@@ -60,6 +60,7 @@ option(BUILD_EXAMPLES "Build examples" ON)
 
 option(BUILD_UNIT_TESTS "Build unit tests" ON)
 option(USE_CUDA "option for enabling gpu features for NVIDIA GPU" OFF)
+option(USE_MLU "option for enabling Cambricon MLU features" OFF)
 option(USE_MUSA "option for enabling gpu features for MTHREADS GPU" OFF)
 option(USE_MACA "option for enabling gpu features for MUXI GPU with MACA" OFF)
 option(USE_HIP "option for enabling gpu features for AMD GPU" OFF)
@@ -73,6 +74,13 @@ option(USE_ASCEND_HETEROGENEOUS "option for transferring between ascend npu and
 option(USE_MNNVL "option for using Multi-Node NVLink transport" OFF)
 option(USE_CXL "option for using CXL protocol" OFF)
 option(USE_EFA "option for using AWS EFA transport" OFF)
+option(USE_UB "option for using UB protocol transport" OFF)
+
+if (USE_UB)
+  add_compile_definitions(USE_UB)
+  message(STATUS "ub transport is enabled")
+  include(${CMAKE_CURRENT_LIST_DIR}/FindUrma.cmake)
+endif()
 
 if (USE_EFA)
   # Find libfabric headers and library; default to AWS EFA installer path
@@ -106,7 +114,11 @@ option(WITH_NVIDIA_PEERMEM "disable to support RDMA without nvidia-peermem. If W
 option(USE_EVENT_DRIVEN_COMPLETION "option for using event-driven completion (store & transfer engine)" OFF)
 
 option(USE_TENT "option for building Mooncake TENT" OFF)
-
+option(ENABLE_MULTI_PROTOCOL "option for enabling multi-protocol support in transfer engine" OFF)
+if (ENABLE_MULTI_PROTOCOL)
+    add_compile_definitions(ENABLE_MULTI_PROTOCOL)
+    message(STATUS "Multi-protocol support is enabled")
+endif()
 option(USE_LRU_MASTER "option for using LRU in master service" OFF)
 option(USE_INTRA_NVLINK "option for using IntraNode nvlink transport" OFF)
 set(LRU_MAX_CAPACITY 1000)
@@ -130,7 +142,7 @@ if (USE_NVMEOF)
 endif()
 
 if (USE_MNNVL)
-  if (NOT USE_HIP AND NOT USE_MUSA)
+  if (NOT USE_HIP AND NOT USE_MUSA AND NOT USE_MACA)
     set(USE_CUDA ON)
   endif()
   add_compile_definitions(USE_MNNVL)
@@ -147,21 +159,58 @@ if (USE_CUDA)
   )
 endif()
 
+if (NOT DEFINED NEUWARE_ROOT OR NEUWARE_ROOT STREQUAL "")
+  if (DEFINED ENV{NEUWARE_HOME} AND NOT "$ENV{NEUWARE_HOME}" STREQUAL "")
+    set(NEUWARE_ROOT "$ENV{NEUWARE_HOME}" CACHE PATH "Path to Cambricon Neuware SDK" FORCE)
+  else()
+    set(NEUWARE_ROOT "/usr/local/neuware" CACHE PATH "Path to Cambricon Neuware SDK" FORCE)
+  endif()
+endif()
+
+if (NOT DEFINED MLU_INCLUDE_DIR OR MLU_INCLUDE_DIR STREQUAL "")
+  set(MLU_INCLUDE_DIR "${NEUWARE_ROOT}/include")
+endif()
+
+if (NOT DEFINED MLU_LIB_DIR OR MLU_LIB_DIR STREQUAL "")
+  set(MLU_LIB_DIR "${NEUWARE_ROOT}/lib64")
+endif()
+
+if (NOT DEFINED MACA_ROOT OR MACA_ROOT STREQUAL "")
+  if (DEFINED ENV{MACA_HOME} AND NOT "$ENV{MACA_HOME}" STREQUAL "")
+    set(MACA_ROOT "$ENV{MACA_HOME}" CACHE PATH "Path to MACA SDK" FORCE)
+  else()
+    set(MACA_ROOT "/opt/maca" CACHE PATH "Path to MACA SDK" FORCE)
+  endif()
+endif()
+
+if (NOT DEFINED MACA_INCLUDE_DIR OR MACA_INCLUDE_DIR STREQUAL "")
+  set(MACA_INCLUDE_DIR "${MACA_ROOT}/include")
+endif()
+
+if (NOT DEFINED MACA_LIB_DIR OR MACA_LIB_DIR STREQUAL "")
+  if (EXISTS "${MACA_ROOT}/lib64")
+    set(MACA_LIB_DIR "${MACA_ROOT}/lib64")
+  else()
+    set(MACA_LIB_DIR "${MACA_ROOT}/lib")
+  endif()
+endif()
+
+if (USE_MLU)
+  add_compile_definitions(USE_MLU)
+  message(STATUS "MLU support is enabled")
+  include_directories(${MLU_INCLUDE_DIR})
+  if (EXISTS "${MLU_LIB_DIR}")
+    link_directories(${MLU_LIB_DIR})
+  endif()
+endif()
+
 if (USE_MACA)
-  # MACA toolchain is CUDA-compatible in first-stage porting.
-  # Reuse CUDA code paths to get a runnable baseline quickly.
   add_compile_definitions(USE_MACA)
   message(STATUS "MACA support is enabled")
-  if(DEFINED ENV{MACA_HOME})
-    set(MACA_HOME $ENV{MACA_HOME})
-  else()
-    set(MACA_HOME /opt/maca)
+  include_directories(${MACA_INCLUDE_DIR})
+  if (EXISTS "${MACA_LIB_DIR}")
+    link_directories(${MACA_LIB_DIR})
   endif()
-  include_directories(${MACA_HOME}/include)
-  link_directories(
-    ${MACA_HOME}/lib
-    ${MACA_HOME}/lib64
-  )
 endif()
 
 if (USE_MUSA)

mooncake-common/etcd/CMakeLists.txt

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

mooncake-common/etcd/go.mod

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

mooncake-common/etcd/go.sum

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

mooncake-common/k8s-lease/CMakeLists.txt

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

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

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

mooncake-common/k8s-lease/go.mod

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

mooncake-common/k8s-lease/k8s_lease_integration_test.go

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

mooncake-common/k8s-lease/k8s_lease_wrapper.go

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

mooncake-common/k8s-lease/k8s_lease_wrapper_test.go

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

mooncake-common/src/CMakeLists.txt

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

mooncake-ep/BuildEpExt.cmake

@@ -13,6 +13,9 @@
 
 cmake_minimum_required(VERSION 3.16)
 
+# Include common build utilities.
+include("${SOURCE_DIR}/../mooncake-common/SetupPyTorchEnv.cmake")
+
 # Restore pipe-separated strings back to CMake semicolon-separated lists.
 if(EP_TORCH_VERSIONS)
   string(REPLACE "|" ";" EP_TORCH_VERSIONS "${EP_TORCH_VERSIONS}")
@@ -53,7 +56,7 @@ endif()
 if("${EP_TORCH_VERSIONS}" STREQUAL "")
   message(STATUS "[EP] Building with currently-installed PyTorch")
   execute_process(
-    COMMAND python setup.py build_ext --build-lib .
+    COMMAND ${Python3_EXECUTABLE} setup.py build_ext --build-lib .
     WORKING_DIRECTORY "${SOURCE_DIR}"
     RESULT_VARIABLE _ret
   )
@@ -63,26 +66,10 @@ if("${EP_TORCH_VERSIONS}" STREQUAL "")
 else()
   message(STATUS "[EP] Building for PyTorch versions: ${EP_TORCH_VERSIONS}")
   foreach(_version IN LISTS EP_TORCH_VERSIONS)
-    message(STATUS "[EP] Installing PyTorch ${_version}")
-    if(EP_CUDA_MAJOR GREATER_EQUAL 13)
-      # TODO: Fix when we need to support more CUDA 13 versions or when the CI
-      #       env is fixed.
-      execute_process(
-        COMMAND pip install "torch==${_version}" --index-url https://download.pytorch.org/whl/cu130
-        RESULT_VARIABLE _ret
-      )
-    else()
-      execute_process(
-        COMMAND pip install "torch==${_version}"
-        RESULT_VARIABLE _ret
-      )
-    endif()
-    if(NOT _ret EQUAL 0)
-      message(FATAL_ERROR "[EP] Failed to install PyTorch ${_version}")
-    endif()
+    install_pytorch_wheel("${_version}" "${EP_CUDA_MAJOR}" "${EP_CUDA_MINOR}" "[EP]")
 
     execute_process(
-      COMMAND python setup.py build_ext --build-lib . --force
+      COMMAND ${Python3_EXECUTABLE} setup.py build_ext --build-lib . --force
       WORKING_DIRECTORY "${SOURCE_DIR}"
       RESULT_VARIABLE _ret
     )

mooncake-integration/CMakeLists.txt

@@ -1,5 +1,5 @@
 file(GLOB SOURCES "*.cpp")
-set(PYTHON_EXECUTABLE "python3")
+include(${CMAKE_CURRENT_LIST_DIR}/../mooncake-common/SetupPython.cmake)
 execute_process(
     COMMAND ${PYTHON_EXECUTABLE} -c "import sys; print([s for s in sys.path if 'packages' in s][0])"
     OUTPUT_VARIABLE PYTHON_SYS_PATH

mooncake-integration/store/store_py.cpp

@@ -1397,6 +1397,7 @@ PYBIND11_MODULE(store, m) {
         .def(py::init<>())
         .def_readwrite("replica_num", &ReplicateConfig::replica_num)
         .def_readwrite("with_soft_pin", &ReplicateConfig::with_soft_pin)
+        .def_readwrite("with_hard_pin", &ReplicateConfig::with_hard_pin)
         .def_readwrite("preferred_segments",
                        &ReplicateConfig::preferred_segments)
         .def_readwrite("preferred_segment", &ReplicateConfig::preferred_segment)
@@ -1561,7 +1562,9 @@ PYBIND11_MODULE(store, m) {
                const std::string &protocol = "tcp",
                const std::string &rdma_devices = "",
                const std::string &master_server_addr = "127.0.0.1:50051",
-               const py::object &engine = py::none()) {
+               const py::object &engine = py::none(),
+               bool enable_ssd_offload = false,
+               const std::string &ssd_offload_path = "") {
                 auto real_client = self.init_real_client();
                 std::shared_ptr<mooncake::TransferEngine> transfer_engine =
                     nullptr;
@@ -1572,12 +1575,15 @@ PYBIND11_MODULE(store, m) {
                 return real_client->setup_real(
                     local_hostname, metadata_server, global_segment_size,
                     local_buffer_size, protocol, rdma_devices,
-                    master_server_addr, transfer_engine, "");
+                    master_server_addr, transfer_engine, "", enable_ssd_offload,
+                    ssd_offload_path);
             },
             py::arg("local_hostname"), py::arg("metadata_server"),
             py::arg("global_segment_size"), py::arg("local_buffer_size"),
             py::arg("protocol"), py::arg("rdma_devices"),
-            py::arg("master_server_addr"), py::arg("engine") = py::none())
+            py::arg("master_server_addr"), py::arg("engine") = py::none(),
+            py::arg("enable_ssd_offload") = false,
+            py::arg("ssd_offload_path") = "")
         .def(
             "setup",
             [](MooncakeStorePyWrapper &self, const py::dict &config_dict) {
@@ -1605,7 +1611,10 @@ PYBIND11_MODULE(store, m) {
             "  protocol: Transfer protocol (default 'tcp').\n"
             "  rdma_devices: RDMA device list.\n"
             "  master_server_addr: Master server address.\n"
-            "  ipc_socket_path: IPC socket path.")
+            "  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).")
         .def(
             "setup_dummy",
             [](MooncakeStorePyWrapper &self, size_t mem_pool_size,
@@ -2033,15 +2042,35 @@ PYBIND11_MODULE(store, m) {
                 // Get data directly into user-provided buffer
                 void *buffer = reinterpret_cast<void *>(buffer_ptr);
                 py::gil_scoped_release release;
-                if (self.use_dummy_client_) {
-                    LOG(ERROR) << "get_into is not supported for dummy client "
-                                  "now";
-                    return (int64_t)-1;
-                }
                 return self.store_->get_into(key, buffer, size);
             },
             py::arg("key"), py::arg("buffer_ptr"), py::arg("size"),
             "Get object data directly into a pre-allocated buffer")
+        .def(
+            "get_into_ranges",
+            [](MooncakeStorePyWrapper &self,
+               const std::vector<uintptr_t> &buffer_ptrs,
+               const std::vector<std::vector<std::string>> &all_keys,
+               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) {
+                std::vector<void *> buffers;
+                buffers.reserve(buffer_ptrs.size());
+                for (uintptr_t ptr : buffer_ptrs) {
+                    buffers.push_back(reinterpret_cast<void *>(ptr));
+                }
+                py::gil_scoped_release release;
+                return self.store_->get_into_ranges(buffers, all_keys,
+                                                    all_dst_offsets,
+                                                    all_src_offsets, all_sizes);
+            },
+            py::arg("buffer_ptrs"), py::arg("all_keys"),
+            py::arg("all_dst_offsets"), py::arg("all_src_offsets"),
+            py::arg("all_sizes"),
+            "Get multiple byte ranges from multiple objects into multiple "
+            "pre-allocated buffers")
         .def(
             "batch_get_into",
             [](MooncakeStorePyWrapper &self,
@@ -2242,6 +2271,20 @@ PYBIND11_MODULE(store, m) {
                 return self.store_->batch_get_replica_desc(keys);
             },
             py::arg("keys"))
+        .def(
+            "batch_replica_clear",
+            [](MooncakeStorePyWrapper &self,
+               const std::vector<std::string> &keys,
+               const std::string &segment_name) {
+                if (!self.is_client_initialized()) {
+                    LOG(ERROR) << "Client is not initialized";
+                    return std::vector<std::string>{};
+                }
+                py::gil_scoped_release release;
+                return self.store_->batch_replica_clear(keys, segment_name);
+            },
+            py::arg("keys"), py::arg("segment_name") = "",
+            "Clear replicas for the given keys. Requires lease to be expired.")
         .def(
             "create_copy_task",
             [](MooncakeStorePyWrapper &self, const std::string &key,

mooncake-integration/transfer_engine/transfer_engine_py.cpp

@@ -21,6 +21,10 @@
 #include <pybind11/stl.h>
 #include "transport/rpc_communicator/rpc_interface.h"
 
+#ifdef USE_HIP
+#include "transport/hip_transport/hip_transport.h"
+#endif
+
 #ifdef USE_MNNVL
 #include "transport/nvlink_transport/nvlink_transport.h"
 #endif
@@ -33,12 +37,12 @@
 #include <cuda_runtime.h>
 #endif
 
-static void *(*allocateMemory)(size_t) = nullptr;
-static void (*freeMemory)(void *) = nullptr;
+static void* (*allocateMemory)(size_t) = nullptr;
+static void (*freeMemory)(void*) = nullptr;
 static std::string g_protocol;
 
 //  Handle allocateMemory function pointer based on protocol
-void initMemoryAllocator(const char *protocol) {
+void initMemoryAllocator(const char* protocol) {
     if (allocateMemory != nullptr) {
         LOG(WARNING) << "Memory allocator already initialized with: "
                      << g_protocol;
@@ -47,23 +51,35 @@ void initMemoryAllocator(const char *protocol) {
     g_protocol = protocol;
     if (strcmp(protocol, "nvlink") == 0) {
 #ifdef USE_MNNVL
-        allocateMemory = [](size_t s) -> void * {
+        allocateMemory = [](size_t s) -> void* {
             return mooncake::NvlinkTransport::allocatePinnedLocalMemory(s);
         };
-        freeMemory = [](void *p) {
+        freeMemory = [](void* p) {
             mooncake::NvlinkTransport::freePinnedLocalMemory(p);
         };
         LOG(INFO) << "Selected MNNVL (NVLink) memory allocator";
 #else
         LOG(ERROR) << "Protocol 'nvlink' requires -DUSE_MNNVL=ON";
+#endif
+    } else if (strcmp(protocol, "hip") == 0) {
+#ifdef USE_HIP
+        allocateMemory = [](size_t s) -> void* {
+            return mooncake::HipTransport::allocatePinnedLocalMemory(s);
+        };
+        freeMemory = [](void* p) {
+            mooncake::HipTransport::freePinnedLocalMemory(p);
+        };
+        LOG(INFO) << "Selected HIP memory allocator";
+#else
+        LOG(ERROR) << "Protocol 'hip' requires -DUSE_HIP=ON";
 #endif
     } else if (strcmp(protocol, "nvlink_intra") == 0) {
 #ifdef USE_INTRA_NVLINK
-        allocateMemory = [](size_t s) -> void * {
+        allocateMemory = [](size_t s) -> void* {
             return mooncake::IntraNodeNvlinkTransport::
                 allocatePinnedLocalMemory(s);
         };
-        freeMemory = [](void *p) {
+        freeMemory = [](void* p) {
             mooncake::IntraNodeNvlinkTransport::freePinnedLocalMemory(p);
         };
         LOG(INFO) << "Selected Intra-NVLink memory allocator";
@@ -71,7 +87,6 @@ void initMemoryAllocator(const char *protocol) {
         LOG(ERROR) << "Protocol 'nvlink_intra' requires -DUSE_INTRA_NVLINK=ON";
 #endif
     } else {
-        // default fallback
         allocateMemory = malloc;
         freeMemory = free;
         LOG(WARNING) << "Using default malloc/free for protocol: " << protocol;
@@ -89,16 +104,16 @@ TransferEnginePy::TransferEnginePy() {
 }
 
 TransferEnginePy::~TransferEnginePy() {
-    for (auto &handle : handle_map_) engine_->closeSegment(handle.second);
+    for (auto& handle : handle_map_) engine_->closeSegment(handle.second);
     handle_map_.clear();
     engine_.reset();
-    for (auto &buffer : buffer_list_) freeMemory(buffer);
+    for (auto& buffer : buffer_list_) freeMemory(buffer);
     buffer_list_.clear();
-    for (auto &buffer : large_buffer_list_) freeMemory(buffer);
+    for (auto& buffer : large_buffer_list_) freeMemory(buffer);
     large_buffer_list_.clear();
 }
 
-std::vector<std::string> buildDeviceFilter(const std::string &device_names) {
+std::vector<std::string> buildDeviceFilter(const std::string& device_names) {
     std::stringstream ss(device_names);
     std::string item;
     std::vector<std::string> tokens;
@@ -109,7 +124,7 @@ std::vector<std::string> buildDeviceFilter(const std::string &device_names) {
 }
 
 std::pair<std::string, std::string> parseConnectionString(
-    const std::string &conn_string) {
+    const std::string& conn_string) {
     std::pair<std::string, std::string> result;
     std::string proto = "etcd";
     std::string domain;
@@ -130,8 +145,8 @@ std::pair<std::string, std::string> parseConnectionString(
     return result;
 }
 
-std::string buildConnString(const std::string &metadata_type,
-                            const std::string &metadata_server) {
+std::string buildConnString(const std::string& metadata_type,
+                            const std::string& metadata_server) {
     if (metadata_server == P2PHANDSHAKE) {
         return P2PHANDSHAKE;
     }
@@ -142,10 +157,10 @@ std::string buildConnString(const std::string &metadata_type,
     return conn_string;
 }
 
-int TransferEnginePy::initialize(const char *local_hostname,
-                                 const char *metadata_server,
-                                 const char *protocol,
-                                 const char *device_name) {
+int TransferEnginePy::initialize(const char* local_hostname,
+                                 const char* metadata_server,
+                                 const char* protocol,
+                                 const char* device_name) {
     initMemoryAllocator(protocol);
 
     auto conn_string = parseConnectionString(metadata_server);
@@ -153,11 +168,17 @@ int TransferEnginePy::initialize(const char *local_hostname,
                          device_name, conn_string.first.c_str());
 }
 
-int TransferEnginePy::initializeExt(const char *local_hostname,
-                                    const char *metadata_server,
-                                    const char *protocol,
-                                    const char *device_name,
-                                    const char *metadata_type) {
+int TransferEnginePy::initializeExt(const char* local_hostname,
+                                    const char* metadata_server,
+                                    const char* protocol,
+                                    const char* device_name,
+                                    const char* metadata_type) {
+    if (strcmp(protocol, "xgmi") == 0) {
+        LOG(ERROR) << "Protocol 'xgmi' is not exposed in the Python API. "
+                   << "Use 'hip' instead.";
+        return -1;
+    }
+
     std::string proto = protocol ? std::string(protocol) : "";
     std::string conn_string = buildConnString(metadata_type, metadata_server);
 
@@ -226,7 +247,7 @@ int TransferEnginePy::initializeExt(const char *local_hostname,
 
 int TransferEnginePy::getRpcPort() { return engine_->getRpcPort(); }
 
-char *TransferEnginePy::allocateRawBuffer(size_t capacity) {
+char* TransferEnginePy::allocateRawBuffer(size_t capacity) {
     auto buffer = allocateMemory(capacity);
     if (!buffer) return nullptr;
     int ret = engine_->registerLocalMemory(buffer, capacity, kWildcardLocation);
@@ -234,7 +255,7 @@ char *TransferEnginePy::allocateRawBuffer(size_t capacity) {
         freeMemory(buffer);
         return nullptr;
     }
-    return (char *)buffer;
+    return (char*)buffer;
 }
 
 int TransferEnginePy::findClassId(size_t size) {
@@ -258,7 +279,7 @@ int TransferEnginePy::doBuddyAllocate(int class_id) {
         if (ret) return ret;
     }
     assert(!free_list_[class_id + 1].empty());
-    char *buffer = free_list_[class_id + 1].top();
+    char* buffer = free_list_[class_id + 1].top();
     free_list_[class_id + 1].pop();
     free_list_[class_id].push(buffer);
     free_list_[class_id].push(buffer + kSlabSizeKB[class_id] * 1024);
@@ -269,21 +290,21 @@ uintptr_t TransferEnginePy::allocateManagedBuffer(size_t length) {
     std::lock_guard<std::mutex> guard(mutex_);
     int class_id = findClassId(length);
     if (class_id < 0) {
-        char *buffer = allocateRawBuffer(length);
+        char* buffer = allocateRawBuffer(length);
         if (buffer) large_buffer_list_.insert(buffer);
         return (uintptr_t)buffer;
     }
     if (free_list_[class_id].empty())
         if (doBuddyAllocate(class_id)) return 0;
     assert(!free_list_[class_id].empty());
-    char *buffer = free_list_[class_id].top();
+    char* buffer = free_list_[class_id].top();
     free_list_[class_id].pop();
     return (uintptr_t)buffer;
 }
 
 int TransferEnginePy::freeManagedBuffer(uintptr_t buffer_addr, size_t length) {
     std::lock_guard<std::mutex> guard(mutex_);
-    auto buffer = (char *)buffer_addr;
+    auto buffer = (char*)buffer_addr;
     int class_id = findClassId(length);
     if (class_id < 0) {
         large_buffer_list_.erase(buffer);
@@ -295,7 +316,7 @@ int TransferEnginePy::freeManagedBuffer(uintptr_t buffer_addr, size_t length) {
     return 0;
 }
 
-int TransferEnginePy::transferSyncWrite(const char *target_hostname,
+int TransferEnginePy::transferSyncWrite(const char* target_hostname,
                                         uintptr_t buffer,
                                         uintptr_t peer_buffer_address,
                                         size_t length) {
@@ -303,7 +324,7 @@ int TransferEnginePy::transferSyncWrite(const char *target_hostname,
                         TransferOpcode::WRITE);
 }
 
-int TransferEnginePy::transferSyncRead(const char *target_hostname,
+int TransferEnginePy::transferSyncRead(const char* target_hostname,
                                        uintptr_t buffer,
                                        uintptr_t peer_buffer_address,
                                        size_t length) {
@@ -312,40 +333,40 @@ int TransferEnginePy::transferSyncRead(const char *target_hostname,
 }
 
 int TransferEnginePy::batchTransferSyncWrite(
-    const char *target_hostname, std::vector<uintptr_t> buffers,
+    const char* target_hostname, std::vector<uintptr_t> buffers,
     std::vector<uintptr_t> peer_buffer_addresses, std::vector<size_t> lengths) {
     return batchTransferSync(target_hostname, buffers, peer_buffer_addresses,
                              lengths, TransferOpcode::WRITE);
 }
 
 int TransferEnginePy::batchTransferSyncRead(
-    const char *target_hostname, std::vector<uintptr_t> buffers,
+    const char* target_hostname, std::vector<uintptr_t> buffers,
     std::vector<uintptr_t> peer_buffer_addresses, std::vector<size_t> lengths) {
     return batchTransferSync(target_hostname, buffers, peer_buffer_addresses,
                              lengths, TransferOpcode::READ);
 }
 
 batch_id_t TransferEnginePy::batchTransferAsyncWrite(
-    const char *target_hostname, const std::vector<uintptr_t> &buffers,
-    const std::vector<uintptr_t> &peer_buffer_addresses,
-    const std::vector<size_t> &lengths) {
+    const char* target_hostname, const std::vector<uintptr_t>& buffers,
+    const std::vector<uintptr_t>& peer_buffer_addresses,
+    const std::vector<size_t>& lengths) {
     return batchTransferAsync(target_hostname, buffers, peer_buffer_addresses,
                               lengths, TransferOpcode::WRITE);
 }
 
 batch_id_t TransferEnginePy::batchTransferAsyncRead(
-    const char *target_hostname, const std::vector<uintptr_t> &buffers,
-    const std::vector<uintptr_t> &peer_buffer_addresses,
-    const std::vector<size_t> &lengths) {
+    const char* target_hostname, const std::vector<uintptr_t>& buffers,
+    const std::vector<uintptr_t>& peer_buffer_addresses,
+    const std::vector<size_t>& lengths) {
     return batchTransferAsync(target_hostname, buffers, peer_buffer_addresses,
                               lengths, TransferOpcode::READ);
 }
 
-int TransferEnginePy::transferSync(const char *target_hostname,
+int TransferEnginePy::transferSync(const char* target_hostname,
                                    uintptr_t buffer,
                                    uintptr_t peer_buffer_address, size_t length,
                                    TransferOpcode opcode,
-                                   TransferNotify *notify) {
+                                   TransferNotify* notify) {
     pybind11::gil_scoped_release release;
     Transport::SegmentHandle handle;
     {
@@ -379,7 +400,7 @@ int TransferEnginePy::transferSync(const char *target_hostname,
             entry.opcode = TransferRequest::READ;
         }
         entry.length = length;
-        entry.source = (void *)buffer;
+        entry.source = (void*)buffer;
         entry.target_id = handle;
         entry.target_offset = peer_buffer_address;
         entry.advise_retry_cnt = retry;
@@ -425,9 +446,8 @@ int TransferEnginePy::transferSync(const char *target_hostname,
             if (current_ts - start_ts > timeout) {
                 LOG(INFO) << "Sync data transfer timeout after "
                           << current_ts - start_ts << "ns, local buffer "
-                          << (void *)buffer << " remote buffer "
-                          << (void *)peer_buffer_address << " length "
-                          << length;
+                          << (void*)buffer << " remote buffer "
+                          << (void*)peer_buffer_address << " length " << length;
                 return -1;
             }
         }
@@ -436,9 +456,9 @@ int TransferEnginePy::transferSync(const char *target_hostname,
 }
 
 int TransferEnginePy::batchTransferSync(
-    const char *target_hostname, std::vector<uintptr_t> buffers,
+    const char* target_hostname, std::vector<uintptr_t> buffers,
     std::vector<uintptr_t> peer_buffer_addresses, std::vector<size_t> lengths,
-    TransferOpcode opcode, TransferNotify *notify) {
+    TransferOpcode opcode, TransferNotify* notify) {
     pybind11::gil_scoped_release release;
     Transport::SegmentHandle handle;
     {
@@ -472,7 +492,7 @@ int TransferEnginePy::batchTransferSync(
             entry.opcode = TransferRequest::READ;
         }
         entry.length = lengths[i];
-        entry.source = (void *)buffers[i];
+        entry.source = (void*)buffers[i];
         entry.target_id = handle;
         entry.target_offset = peer_buffer_addresses[i];
         entry.advise_retry_cnt = 0;
@@ -539,9 +559,9 @@ int TransferEnginePy::batchTransferSync(
 }
 
 batch_id_t TransferEnginePy::batchTransferAsync(
-    const char *target_hostname, const std::vector<uintptr_t> &buffers,
-    const std::vector<uintptr_t> &peer_buffer_addresses,
-    const std::vector<size_t> &lengths, TransferOpcode opcode) {
+    const char* target_hostname, const std::vector<uintptr_t>& buffers,
+    const std::vector<uintptr_t>& peer_buffer_addresses,
+    const std::vector<size_t>& lengths, TransferOpcode opcode) {
     pybind11::gil_scoped_release release;
     Transport::SegmentHandle handle;
     {
@@ -574,7 +594,7 @@ batch_id_t TransferEnginePy::batchTransferAsync(
             entry.opcode = TransferRequest::READ;
         }
         entry.length = lengths[i];
-        entry.source = (void *)buffers[i];
+        entry.source = (void*)buffers[i];
         entry.target_id = handle;
         entry.target_offset = peer_buffer_addresses[i];
         entry.advise_retry_cnt = 0;
@@ -583,7 +603,7 @@ batch_id_t TransferEnginePy::batchTransferAsync(
 
     for (int retry = 0; retry < max_retry; ++retry) {
         batch_id = engine_->allocateBatchID(batch_size);
-        auto batch_desc = reinterpret_cast<BatchDesc *>(batch_id);
+        auto batch_desc = reinterpret_cast<BatchDesc*>(batch_id);
 
         auto start_ts = getCurrentTimeInNano();
         batch_desc->start_timestamp = start_ts;
@@ -601,18 +621,18 @@ batch_id_t TransferEnginePy::batchTransferAsync(
 }
 
 int TransferEnginePy::getBatchTransferStatus(
-    const std::vector<batch_id_t> &batch_ids) {
+    const std::vector<batch_id_t>& batch_ids) {
     pybind11::gil_scoped_release release;
     TransferStatus status;
     std::unordered_map<batch_id_t, int64_t> timeout_table{};
-    for (auto &batch_id : batch_ids) {
+    for (auto& batch_id : batch_ids) {
         int64_t total_length = 0;
-        auto batch_desc = reinterpret_cast<BatchDesc *>(batch_id);
+        auto batch_desc = reinterpret_cast<BatchDesc*>(batch_id);
         const size_t task_count = batch_desc->task_list.size();
 
         for (size_t task_id = 0; task_id < task_count; task_id++) {
-            auto &task = batch_desc->task_list[task_id];
-            for (auto &slice : task.slice_list) {
+            auto& task = batch_desc->task_list[task_id];
+            for (auto& slice : task.slice_list) {
                 total_length += slice->length;
             }
         }
@@ -623,8 +643,8 @@ int TransferEnginePy::getBatchTransferStatus(
     bool failed_or_timeout = false;
     std::unordered_set<batch_id_t> remove_ids{};
     while (!timeout_table.empty() && !failed_or_timeout) {
-        for (auto &entry : timeout_table) {
-            auto batch_desc = reinterpret_cast<BatchDesc *>(entry.first);
+        for (auto& entry : timeout_table) {
+            auto batch_desc = reinterpret_cast<BatchDesc*>(entry.first);
             auto start_timestamp = batch_desc->start_timestamp;
             Status s = engine_->getBatchTransferStatus(entry.first, status);
             LOG_ASSERT(s.ok());
@@ -645,7 +665,7 @@ int TransferEnginePy::getBatchTransferStatus(
             }
         }
 
-        for (auto &remove_id : remove_ids) {
+        for (auto& remove_id : remove_ids) {
             timeout_table.erase(remove_id);
         }
 
@@ -653,7 +673,7 @@ int TransferEnginePy::getBatchTransferStatus(
     }
 
     if (failed_or_timeout) {
-        for (auto &entry : timeout_table) {
+        for (auto& entry : timeout_table) {
             engine_->freeBatchID(entry.first);
         }
     }
@@ -661,7 +681,7 @@ int TransferEnginePy::getBatchTransferStatus(
     return failed_or_timeout ? -1 : 0;
 }
 
-batch_id_t TransferEnginePy::transferSubmitWrite(const char *target_hostname,
+batch_id_t TransferEnginePy::transferSubmitWrite(const char* target_hostname,
                                                  uintptr_t buffer,
                                                  uintptr_t peer_buffer_address,
                                                  size_t length) {
@@ -682,7 +702,7 @@ batch_id_t TransferEnginePy::transferSubmitWrite(const char *target_hostname,
     TransferRequest entry;
     entry.opcode = TransferRequest::WRITE;
     entry.length = length;
-    entry.source = (void *)buffer;
+    entry.source = (void*)buffer;
     entry.target_id = handle;
     entry.target_offset = peer_buffer_address;
 
@@ -717,7 +737,7 @@ int TransferEnginePy::batchRegisterMemory(
     std::vector<BufferEntry> buffers;
     for (size_t i = 0; i < batch_size; i++) {
         buffers.push_back(
-            BufferEntry{(void *)buffer_addresses[i], capacities[i]});
+            BufferEntry{(void*)buffer_addresses[i], capacities[i]});
     }
     return engine_->registerLocalMemoryBatch(buffers, kWildcardLocation);
 }
@@ -726,20 +746,20 @@ int TransferEnginePy::batchUnregisterMemory(
     std::vector<uintptr_t> buffer_addresses) {
     pybind11::gil_scoped_release release;
     auto batch_size = buffer_addresses.size();
-    std::vector<void *> buffers;
+    std::vector<void*> buffers;
     for (size_t i = 0; i < batch_size; i++) {
-        buffers.push_back(reinterpret_cast<char *>(buffer_addresses[i]));
+        buffers.push_back(reinterpret_cast<char*>(buffer_addresses[i]));
     }
     return engine_->unregisterLocalMemoryBatch(buffers);
 }
 
 int TransferEnginePy::registerMemory(uintptr_t buffer_addr, size_t capacity) {
-    char *buffer = reinterpret_cast<char *>(buffer_addr);
+    char* buffer = reinterpret_cast<char*>(buffer_addr);
     return engine_->registerLocalMemory(buffer, capacity);
 }
 
 int TransferEnginePy::unregisterMemory(uintptr_t buffer_addr) {
-    char *buffer = reinterpret_cast<char *>(buffer_addr);
+    char* buffer = reinterpret_cast<char*>(buffer_addr);
     return engine_->unregisterLocalMemory(buffer);
 }
 
@@ -767,8 +787,8 @@ struct TransferOnCudaContext {
  *
  * @param data Pointer to a TransferOnCudaContext object.
  */
-void CUDART_CB transfer_on_cuda_callback(void *data) {
-    auto *ctx = reinterpret_cast<TransferOnCudaContext *>(data);
+void CUDART_CB transfer_on_cuda_callback(void* data) {
+    auto* ctx = reinterpret_cast<TransferOnCudaContext*>(data);
 
     auto status = ctx->engine->submitTransfer(ctx->batch_id, ctx->requests);
     if (!status.ok()) {
@@ -828,9 +848,9 @@ error_exit:
  * @param stream_ptr Handle to a CUDA stream (cudaStream_t as uintptr_t).
  */
 void TransferEnginePy::batchTransferOnCuda(
-    const char *target_hostname, const std::vector<uintptr_t> &buffers,
-    const std::vector<uintptr_t> &peer_buffer_addresses,
-    const std::vector<size_t> &lengths, TransferOpcode opcode,
+    const char* target_hostname, const std::vector<uintptr_t>& buffers,
+    const std::vector<uintptr_t>& peer_buffer_addresses,
+    const std::vector<size_t>& lengths, TransferOpcode opcode,
     uintptr_t stream_ptr) {
     pybind11::gil_scoped_release release;
     Transport::SegmentHandle handle;
@@ -863,7 +883,7 @@ void TransferEnginePy::batchTransferOnCuda(
                            ? TransferRequest::WRITE
                            : TransferRequest::READ;
         entry.length = lengths[i];
-        entry.source = (void *)buffers[i];
+        entry.source = (void*)buffers[i];
         entry.target_id = handle;
         entry.target_offset = peer_buffer_addresses[i];
         entries.push_back(entry);
@@ -871,7 +891,7 @@ void TransferEnginePy::batchTransferOnCuda(
     }
 
     auto batch_id = engine_->allocateBatchID(batch_size);
-    auto *ctx = new TransferOnCudaContext{engine_, batch_id, std::move(entries),
+    auto* ctx = new TransferOnCudaContext{engine_, batch_id, std::move(entries),
                                           total_bytes};
 
     cudaStream_t stream = reinterpret_cast<cudaStream_t>(stream_ptr);
@@ -888,7 +908,7 @@ void TransferEnginePy::batchTransferOnCuda(
 /**
  * @brief Async WRITE transfer triggered by a CUDA stream.
  */
-void TransferEnginePy::transferWriteOnCuda(const char *target_hostname,
+void TransferEnginePy::transferWriteOnCuda(const char* target_hostname,
                                            uintptr_t buffer,
                                            uintptr_t peer_buffer_address,
                                            size_t length,
@@ -900,7 +920,7 @@ void TransferEnginePy::transferWriteOnCuda(const char *target_hostname,
 /**
  * @brief Async READ transfer triggered by a CUDA stream.
  */
-void TransferEnginePy::transferReadOnCuda(const char *target_hostname,
+void TransferEnginePy::transferReadOnCuda(const char* target_hostname,
                                           uintptr_t buffer,
                                           uintptr_t peer_buffer_address,
                                           size_t length, uintptr_t stream_ptr) {
@@ -912,9 +932,9 @@ void TransferEnginePy::transferReadOnCuda(const char *target_hostname,
  * @brief Batch async WRITE transfer triggered by a CUDA stream.
  */
 void TransferEnginePy::batchTransferWriteOnCuda(
-    const char *target_hostname, const std::vector<uintptr_t> &buffers,
-    const std::vector<uintptr_t> &peer_buffer_addresses,
-    const std::vector<size_t> &lengths, uintptr_t stream_ptr) {
+    const char* target_hostname, const std::vector<uintptr_t>& buffers,
+    const std::vector<uintptr_t>& peer_buffer_addresses,
+    const std::vector<size_t>& lengths, uintptr_t stream_ptr) {
     batchTransferOnCuda(target_hostname, buffers, peer_buffer_addresses,
                         lengths, TransferOpcode::WRITE, stream_ptr);
 }
@@ -923,16 +943,16 @@ void TransferEnginePy::batchTransferWriteOnCuda(
  * @brief Batch async READ transfer triggered by a CUDA stream.
  */
 void TransferEnginePy::batchTransferReadOnCuda(
-    const char *target_hostname, const std::vector<uintptr_t> &buffers,
-    const std::vector<uintptr_t> &peer_buffer_addresses,
-    const std::vector<size_t> &lengths, uintptr_t stream_ptr) {
+    const char* target_hostname, const std::vector<uintptr_t>& buffers,
+    const std::vector<uintptr_t>& peer_buffer_addresses,
+    const std::vector<size_t>& lengths, uintptr_t stream_ptr) {
     batchTransferOnCuda(target_hostname, buffers, peer_buffer_addresses,
                         lengths, TransferOpcode::READ, stream_ptr);
 }
 #endif
 
 uintptr_t TransferEnginePy::getFirstBufferAddress(
-    const std::string &segment_name) {
+    const std::string& segment_name) {
     Transport::SegmentHandle segment_id =
         engine_->openSegment(segment_name.c_str());
     auto segment_desc = engine_->getMetadata()->getSegmentDescByID(segment_id);
@@ -942,7 +962,7 @@ uintptr_t TransferEnginePy::getFirstBufferAddress(
     return segment_desc->buffers[0].addr;
 }
 
-std::string TransferEnginePy::getLocalTopology(const char *device_name) {
+std::string TransferEnginePy::getLocalTopology(const char* device_name) {
     pybind11::gil_scoped_release release;
     auto device_name_safe = device_name ? std::string(device_name) : "";
     auto device_filter = buildDeviceFilter(device_name_safe);
@@ -965,7 +985,7 @@ std::vector<TransferEnginePy::TransferNotify> TransferEnginePy::getNotifies() {
         return result;
     }
 
-    for (const auto &notify : notifies) {
+    for (const auto& notify : notifies) {
         result.emplace_back(
             TransferEnginePy::TransferNotify{notify.name, notify.notify_msg});
     }
@@ -976,7 +996,7 @@ std::vector<TransferEnginePy::TransferNotify> TransferEnginePy::getNotifies() {
 namespace py = pybind11;
 
 // Implementation of coro_rpc_interface binding function
-void bind_coro_rpc_interface(py::module_ &m) {
+void bind_coro_rpc_interface(py::module_& m) {
     // Note: RpcInterface, ReceivedData and ReceivedTensor are already
     // registered by bind_rpc_interface() so we don't register them again here
     // to avoid duplicate type registration errors. The factory functions are
@@ -996,7 +1016,7 @@ PYBIND11_MODULE(engine, m) {
 
     py::class_<TransferEnginePy::TransferNotify>(m, "TransferNotify")
         .def(py::init<>())
-        .def(py::init<const std::string &, const std::string &>(),
+        .def(py::init<const std::string&, const std::string&>(),
              py::arg("name"), py::arg("msg"))
         .def_readwrite("name", &TransferEnginePy::TransferNotify::name)
         .def_readwrite("msg", &TransferEnginePy::TransferNotify::msg);

mooncake-p2p-store/build.sh

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

mooncake-pg/BuildPgExt.cmake

@@ -5,6 +5,7 @@
 #
 #   SOURCE_DIR          - mooncake-pg source directory
 #   EP_CUDA_MAJOR       - CUDA major version (integer)
+#   EP_CUDA_MINOR       - CUDA minor version (integer)
 #   EP_TORCH_VERSIONS   - pipe-separated (|) PyTorch versions to build for
 #                         (empty = use the currently-installed torch)
 #   TORCH_CUDA_ARCH_LIST - pipe-separated CUDA arch list forwarded to torch
@@ -13,6 +14,9 @@
 
 cmake_minimum_required(VERSION 3.16)
 
+# Include common build utilities.
+include("${SOURCE_DIR}/../mooncake-common/SetupPyTorchEnv.cmake")
+
 # Restore pipe-separated strings back to CMake semicolon-separated lists.
 if(EP_TORCH_VERSIONS)
   string(REPLACE "|" ";" EP_TORCH_VERSIONS "${EP_TORCH_VERSIONS}")
@@ -53,7 +57,7 @@ endif()
 if("${EP_TORCH_VERSIONS}" STREQUAL "")
   message(STATUS "[PG] Building with currently-installed PyTorch")
   execute_process(
-    COMMAND python setup.py build_ext --build-lib .
+    COMMAND ${Python3_EXECUTABLE} setup.py build_ext --build-lib .
     WORKING_DIRECTORY "${SOURCE_DIR}"
     RESULT_VARIABLE _ret
   )
@@ -63,26 +67,10 @@ if("${EP_TORCH_VERSIONS}" STREQUAL "")
 else()
   message(STATUS "[PG] Building for PyTorch versions: ${EP_TORCH_VERSIONS}")
   foreach(_version IN LISTS EP_TORCH_VERSIONS)
-    message(STATUS "[PG] Installing PyTorch ${_version}")
-    if(EP_CUDA_MAJOR GREATER_EQUAL 13)
-      # TODO: Fix when we need to support more CUDA 13 versions or when the CI
-      #       env is fixed.
-      execute_process(
-        COMMAND pip install "torch==${_version}" --index-url https://download.pytorch.org/whl/cu130
-        RESULT_VARIABLE _ret
-      )
-    else()
-      execute_process(
-        COMMAND pip install "torch==${_version}"
-        RESULT_VARIABLE _ret
-      )
-    endif()
-    if(NOT _ret EQUAL 0)
-      message(FATAL_ERROR "[PG] Failed to install PyTorch ${_version}")
-    endif()
+    install_pytorch_wheel("${_version}" "${EP_CUDA_MAJOR}" "${EP_CUDA_MINOR}" "[PG]")
 
     execute_process(
-      COMMAND python setup.py build_ext --build-lib . --force
+      COMMAND ${Python3_EXECUTABLE} setup.py build_ext --build-lib . --force
       WORKING_DIRECTORY "${SOURCE_DIR}"
       RESULT_VARIABLE _ret
     )

mooncake-pg/benchmark/README.md

@@ -16,6 +16,8 @@ python mooncake-pg/benchmark/pgbench.py \
   --collective all_reduce --backend mooncake --device cuda -g 8 -b 8 -e 128M -f 2
 ```
 
+Set `MOONCAKE_PGTEST_DEVICE_FILTERS=mlx5_1,mlx5_2,...` to explicitly set HCA whitelist.
+
 ## Notes
 - Single-node only (v1). Use `-g/--ngpus` as local ranks when spawning.
 - Dtypes align with nccl-tests; use `-d all` to sweep supported types.

mooncake-pg/benchmark/p2p_regular_k_bench.py

@@ -10,7 +10,11 @@ import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
 
-from pgbench_utils import parse_size, resolve_dtype
+from pgbench_utils import (
+    configure_mooncake_device_filter,
+    parse_size,
+    resolve_dtype,
+)
 
 
 def _parse_args() -> argparse.Namespace:
@@ -53,9 +57,9 @@ def _init_backend_device(args: argparse.Namespace) -> None:
 
     if args.backend in ("mooncake", "mooncake-cpu"):
         try:
-            import mooncake.pg as pg  # noqa: F401
+            import mooncake.pg as pg
 
-            pg.set_device_filter(["mlx5_1", "mlx5_2", "mlx5_3", "mlx5_4"])
+            configure_mooncake_device_filter(pg)
         except Exception as exc:
             raise RuntimeError(
                 "Failed to import mooncake.pg; ensure PYTHONPATH includes mooncake-pg"

mooncake-pg/benchmark/pgbench.py

@@ -9,11 +9,11 @@ from typing import List, Optional, Tuple
 import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
-import mooncake.pg as pg
 
 from pgbench_utils import (
     busbw_factor,
     compute_counts,
+    configure_mooncake_device_filter,
     format_header,
     format_result_line,
     list_supported_dtypes,
@@ -22,8 +22,6 @@ from pgbench_utils import (
     resolve_reduce_op,
 )
 
-pg.set_device_filter(["mlx5_1", "mlx5_2", "mlx5_3", "mlx5_4"])
-
 COLLECTIVES = {
     "all_reduce",
     "all_gather",
@@ -448,7 +446,9 @@ def _run_worker(local_rank: int, args: argparse.Namespace) -> None:
     backend = args.backend
     if backend in ("mooncake", "mooncake-cpu"):
         try:
-            import mooncake.pg as pg  # noqa: F401
+            import mooncake.pg as pg
+
+            configure_mooncake_device_filter(pg)
         except (
             Exception
         ) as exc:  # pragma: no cover - import-time failure should be explicit

mooncake-pg/benchmark/pgbench_utils.py

@@ -1,13 +1,15 @@
 from __future__ import annotations
 
 import datetime
+import os
 import re
-from typing import List, Optional, Tuple
+from typing import List, Optional, Sequence, Tuple
 
 import torch
 import torch.distributed as dist
 
 
+PGTEST_DEVICE_FILTER_ENV_VAR = "MOONCAKE_PGTEST_DEVICE_FILTERS"
 NCCL_DTYPE_ORDER = [
     "int8",
     "uint8",
@@ -26,6 +28,32 @@ NCCL_DTYPE_ORDER = [
 _SIZE_RE = re.compile(r"^(\d+)([KkMmGgTt])?[Bb]?$")
 
 
+def parse_device_filters(raw: str | None) -> list[str] | None:
+    if raw is None:
+        return None
+    filters = [item.strip() for item in raw.split(",") if item.strip()]
+    return filters or None
+
+
+def resolve_pgtest_device_filters(
+    device_filters: Sequence[str] | None = None,
+) -> list[str] | None:
+    if device_filters is not None:
+        resolved = [item.strip() for item in device_filters if item.strip()]
+        return resolved or None
+    return parse_device_filters(os.getenv(PGTEST_DEVICE_FILTER_ENV_VAR))
+
+
+def configure_mooncake_device_filter(
+    pg_module,
+    device_filters: Sequence[str] | None = None,
+) -> list[str] | None:
+    resolved = resolve_pgtest_device_filters(device_filters)
+    if resolved is not None:
+        pg_module.set_device_filter(resolved)
+    return resolved
+
+
 def parse_size(value: object) -> int:
     if isinstance(value, int):
         return value

mooncake-pg/include/connection_poller.h

@@ -25,8 +25,8 @@ enum class PeerConnectionState {
 };
 
 struct PeerConnection {
-    static constexpr size_t CHECK_STORE_INITIAL_BACKOFF_MS = 8;
-    static constexpr size_t CHECK_STORE_MAX_BACKOFF_MS = 1024;
+    static constexpr size_t kCheckStoreInitialBackoffMs = 8;
+    static constexpr size_t kCheckStoreMaxBackoffMs = 1024;
 
     PeerConnectionState state{PeerConnectionState::WAITING_STORE};
     std::optional<BatchID> warmupBatchId{std::nullopt};
@@ -34,21 +34,21 @@ struct PeerConnection {
 
     // Back off to avoid frequently checking store.
     std::chrono::steady_clock::time_point last_check_store;
-    size_t check_store_backoff_ms{CHECK_STORE_INITIAL_BACKOFF_MS};
+    size_t check_store_backoff_ms{kCheckStoreInitialBackoffMs};
 
     void increaseCheckStoreBackoff() {
         check_store_backoff_ms =
             (std::min)(check_store_backoff_ms * 2,
-                       PeerConnection::CHECK_STORE_MAX_BACKOFF_MS);
+                       PeerConnection::kCheckStoreMaxBackoffMs);
     }
 
     void resetCheckStoreBackoff() {
-        check_store_backoff_ms = CHECK_STORE_INITIAL_BACKOFF_MS;
+        check_store_backoff_ms = kCheckStoreInitialBackoffMs;
     }
 };
 
 class ConnectionContext {
-   private:
+    static constexpr size_t kDrainPollerTimeoutMs = 5000;  // 5s
     friend class ConnectionPoller;
 
     int backendIndex_;
@@ -88,6 +88,8 @@ class ConnectionContext {
     std::mutex backend_wakeup_mutex_;
     std::condition_variable backend_wakeup_cv_;
 
+    bool resource_abandoned_{false};
+
    public:
     ConnectionContext(int backendIndex, int rank, int size, bool isDummy,
                       uint64_t* local2global_rank_map,
@@ -155,6 +157,25 @@ class ConnectionContext {
 
     void setDummy(bool isDummy) { isDummy_ = isDummy; }
 
+    /**
+     * @brief Waits for the poller to stop all peer connections gracefully.
+     *
+     * Blocks until all peer connections have transitioned to the EXPIRING state
+     * or the timeout expires. Used during shutdown to ensure no pending
+     * transfers are active before resource cleanup.
+     *
+     * @return True if all peers stopped within the timeout; false otherwise.
+     */
+    bool drainPoller() const;
+
+    /**
+     * @brief Abandons resources instead of releasing them properly.
+     *
+     * When a hung operation prevents clean shutdown, this method marks
+     * resources as abandoned to prevent crashes during cleanup.
+     */
+    void abandonResources();
+
     static std::string getServerNameStoreKey(int backendIndex, int rank) {
         return "server_name_" + std::to_string(backendIndex) + "_" +
                std::to_string(rank);
@@ -178,14 +199,15 @@ class ConnectionContext {
     // For ConnectionManager
     bool poll();
     bool tryStop();
+    bool isStopped() const;
 
     // Internal helpers
     bool pollPeer(int pollingRank);
 };
 
 class ConnectionPoller {
-    static constexpr size_t CONNECTING_IDLE_SLEEP_MS = 50;
-    static constexpr size_t ALL_CONNECTED_IDLE_SLEEP_MS = 200;
+    static constexpr size_t kConnectingIdleSleepMs = 50;
+    static constexpr size_t kAllConnectedIdleSleepMs = 200;
 
    public:
     static ConnectionPoller& GetInstance() {

mooncake-pg/include/mooncake_backend.h

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

mooncake-pg/include/mooncake_worker.cuh

@@ -11,8 +11,11 @@
 #include <transfer_engine.h>
 
 #include <memory>
+#include <atomic>
 #include <mutex>
+#include <thread>
 #include <unordered_map>
+#include <vector>
 
 namespace mooncake {
 
@@ -50,6 +53,7 @@ __global__ struct Task {
     size_t tensorSize;  // In bytes
     int64_t broadcastRoot;
     int bufferOffset;
+    uint64_t submitSequence = 0;
     BatchID batchID;
     void* transferGroupMeta;
 };
@@ -63,9 +67,16 @@ 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);
+    ~MooncakeWorker();
 
     c10::intrusive_ptr<c10d::Work> putTaskCpu(
         c10d::OpType opType, size_t tensorSize, int64_t broadcastRoot,
@@ -80,15 +91,31 @@ class MooncakeWorker {
         c10d::OpType opType, size_t tensorSize, int64_t broadcastRoot,
         const std::shared_ptr<TransferGroupMeta>& meta,
         const std::shared_ptr<ConnectionContext>& connection_ctx,
-        const at::cuda::CUDAStream& stream,
-        const std::function<void(void* dst, size_t pos, size_t realSize)>&
-            tensorToBuffer,
-        const std::function<void(void* src, size_t pos, size_t realSize)>&
-            bufferToTensor);
+        const at::cuda::CUDAStream& issue_stream,
+        const std::function<void(void* dst, size_t pos, size_t realSize,
+                                 const at::cuda::CUDAStream&)>& tensorToBuffer,
+        const std::function<void(void* src, size_t pos, size_t realSize,
+                                 const at::cuda::CUDAStream&)>& bufferToTensor);
 
     void Start();
 
-    void Stop() { running_ = false; }
+    /**
+     * @brief Waits for all active collective tasks for the given backend to
+     * complete.
+     *
+     * Used during graceful shutdown to ensure no pending collective operations
+     * are active before releasing resources. Blocks until all tasks complete
+     * or the timeout expires.
+     *
+     * @param meta The transfer group metadata identifying the backend.
+     * @return True if all tasks completed within the timeout; false if timed
+     * out.
+     */
+    bool drainTasks(const TransferGroupMeta* meta) const;
+
+    bool waitUntilTasksSubmitted(
+        const std::vector<CudaTaskSubmissionToken>& tasks,
+        std::chrono::milliseconds timeout) const;
 
    private:
     void startWorker();
@@ -96,6 +123,7 @@ class MooncakeWorker {
     static constexpr size_t kNumTasks_ = 4;
 
     static constexpr size_t kPingTimeoutMicroseconds_ = 100;
+    static constexpr size_t kDrainTasksTimeoutMs = 5000;  // 5s
 
     bool running_ = false;
     std::atomic<bool> started_{false};
@@ -107,6 +135,10 @@ 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_;
 };
 
 class MooncakeWorkerManager {

mooncake-pg/include/p2p_proxy.h

@@ -5,6 +5,7 @@
 #include <torch/torch.h>
 #include <array>
 #include <atomic>
+#include <condition_variable>
 #include <cstddef>
 #include <cstdint>
 #include <deque>
@@ -57,6 +58,8 @@ struct P2PControlSlot {
 
 class P2PDeviceWorker;
 class P2PProxy {
+    static constexpr size_t kDrainTasksTimeoutMs = 5000;  // 5s
+
    public:
     friend class P2PDeviceWorker;
 
@@ -101,6 +104,26 @@ class P2PProxy {
 
     void ResetPeerState(int peer_rank);
 
+    /**
+     * @brief Waits for all active P2P send and receive tasks to complete.
+     *
+     * Used during graceful shutdown to ensure no pending P2P operations
+     * are active before releasing resources. Blocks until all tasks complete
+     * or the timeout expires.
+     *
+     * @return True if all tasks completed within the timeout; false if timed
+     * out.
+     */
+    bool DrainTasks() const;
+
+    /**
+     * @brief Abandons resources instead of releasing them properly.
+     *
+     * When a hung operation prevents clean shutdown, this method marks
+     * resources as abandoned to prevent crashes during destructor.
+     */
+    void AbandonResources();
+
    private:
     enum class TransferState {
         kDataCopy,
@@ -233,6 +256,7 @@ class P2PProxy {
     int cuda_device_index_ = -1;
     std::string location_;
     P2PResources resources_;
+    bool resource_abandoned_{false};
 
     std::queue<SendOpContext> send_queue_;
     std::mutex send_queue_mutex_;

mooncake-pg/include/pg_utils.h

@@ -0,0 +1,186 @@
+#ifndef MOONCAKE_PG_UTILS_H
+#define MOONCAKE_PG_UTILS_H
+
+#pragma once
+
+#include <chrono>
+#include <thread>
+#include <algorithm>
+#include <cstdint>
+
+// For PAUSE macro
+#include <transfer_engine.h>
+
+namespace mooncake {
+
+/**
+ * @brief Configuration parameters for the BackoffWaiter.
+ *
+ * Defines the thresholds and durations for the multi-stage backoff strategy:
+ * Spin -> Thread Yield -> Exponential Sleep.
+ */
+struct BackoffWaiterConfig {
+    /**
+     * @brief The maximum number of iterations to perform CPU
+     * busy-waiting (spinning). During this phase, the thread uses PAUSE
+     * to minimize latency.
+     */
+    uint32_t spin_limit = 200;
+
+    /**
+     * @brief The maximum number of times to yield.
+     * This phase occurs after spinning is exhausted, reducing CPU consumption
+     * while still maintaining relatively high responsiveness.
+     */
+    uint32_t yield_limit = 50;
+
+    /**
+     * @brief The initial sleep duration once the waiter enters the sleep phase.
+     */
+    std::chrono::microseconds init_sleep{10};
+
+    /**
+     * @brief The maximum allowed sleep duration. The sleep time will double
+     * exponentially up to this cap to prevent excessive overhead during waits.
+     */
+    std::chrono::microseconds max_sleep{100000};  // 100ms
+
+    /**
+     * @brief Creates a configuration that skips spinning and yielding, using
+     * only sleep-based backoff.
+     *
+     * @param init_sleep The initial sleep duration.
+     * @param max_sleep The maximum sleep duration cap.
+     * @return BackoffWaiterConfig instance using sleep-only strategy.
+     */
+    static BackoffWaiterConfig sleepOnly(
+        std::chrono::microseconds init_sleep,
+        std::chrono::microseconds max_sleep) noexcept {
+        return {0, 0, init_sleep, max_sleep};
+    }
+
+    /**
+     * @brief Creates a configuration that uses a fixed sleep duration.
+     *
+     * This disables spinning, yielding, and exponential backoff.
+     * The thread will sleep for a constant duration on each wait iteration.
+     *
+     * @param sleep_time The constant sleep duration to use.
+     * @return BackoffWaiterConfig instance with constant sleep behavior.
+     */
+    static BackoffWaiterConfig constantSleep(
+        std::chrono::microseconds sleep_time) noexcept {
+        return {0, 0, sleep_time, sleep_time};
+    }
+};
+
+/**
+ * @brief A waiting utility with adaptive backoff strategy.
+ *
+ * This class provides a three-stage backoff mechanism (Spin -> Yield -> Sleep)
+ * designed for efficiently polling asynchronous operations. It minimizes
+ * latency for fast operations by using CPU spinning initially, then
+ * progressively reduces CPU usage through thread yielding and exponential sleep
+ * backoff for long-running waits.
+ */
+class BackoffWaiter {
+   public:
+    explicit BackoffWaiter(
+        const BackoffWaiterConfig& cfg = BackoffWaiterConfig{})
+        : config_(cfg), current_sleep_(cfg.init_sleep) {}
+
+    void reset() noexcept {
+        spin_count_ = 0;
+        yield_count_ = 0;
+        current_sleep_ = config_.init_sleep;
+    }
+
+    /**
+     * @brief Advances the waiter to the next backoff state.
+     *
+     * @par Example Usage:
+     * Manually calling `step()` is useful for custom waiting
+     * where `wait()` or `wait_for()` is not applicable.
+     *
+     * @code
+     * std::atomic<bool> ready_flag{false};
+     * mooncake::BackoffWaiter waiter;
+     *
+     * // Wait indefinitely until the flag is set by another thread
+     * while (!ready_flag.load(std::memory_order_acquire)) {
+     *     // Perform some custom logic here if needed...
+     *     waiter.step();
+     * }
+     *
+     * // Reset the state if you plan to reuse this waiter instance later
+     * waiter.reset();
+     * @endcode
+     */
+    void step() {
+        if (spin_count_ < config_.spin_limit) {
+            PAUSE();
+            ++spin_count_;
+        } else if (yield_count_ < config_.yield_limit) {
+            std::this_thread::yield();
+            ++yield_count_;
+        } else {
+            std::this_thread::sleep_for(current_sleep_);
+            current_sleep_ = std::min(current_sleep_ * 2, config_.max_sleep);
+        }
+    }
+
+    /**
+     * @brief Blocks the current thread until the predicate is satisfied
+     * or the timeout expires.
+     *
+     * Repeatedly evaluates the given predicate. If the predicate returns false,
+     * it progresses the backoff state using step().
+     *
+     * @tparam Predicate A callable that returns a boolean condition.
+     * @tparam Rep An arithmetic type representing the number of ticks.
+     * @tparam Period A std::ratio representing the tick period.
+     * @param timeout The maximum duration to wait before giving up.
+     * @param pred The condition to wait for.
+     * @return true if the predicate evaluated to true within the timeout.
+     * @return false if the timeout expired before the predicate was satisfied.
+     */
+    template <typename Predicate, typename Rep, typename Period>
+    [[nodiscard]] bool wait_for(std::chrono::duration<Rep, Period> timeout,
+                                Predicate pred) {
+        reset();
+
+        auto start = std::chrono::steady_clock::now();
+        while (!pred()) {
+            if (std::chrono::steady_clock::now() - start > timeout) {
+                return false;
+            }
+            step();
+        }
+        return true;
+    }
+
+    /**
+     * @brief Blocks the current thread indefinitely until the predicate is
+     * satisfied.
+     *
+     * @tparam Predicate A callable that returns a boolean condition.
+     * @param pred The condition to wait for.
+     */
+    template <typename Predicate>
+    void wait(Predicate pred) {
+        reset();
+
+        while (!pred()) {
+            step();
+        }
+    }
+
+   private:
+    BackoffWaiterConfig config_;
+    uint32_t spin_count_{0};
+    uint32_t yield_count_{0};
+    std::chrono::microseconds current_sleep_;
+};
+}  // namespace mooncake
+
+#endif

mooncake-pg/src/connection_poller.cpp

@@ -15,6 +15,7 @@
 #include <limits>
 #include "memory_location.h"
 #include "mooncake_worker.cuh"
+#include "pg_utils.h"
 
 namespace mooncake {
 
@@ -80,6 +81,12 @@ ConnectionContext::ConnectionContext(int backendIndex, int rank, int size,
 }
 
 ConnectionContext::~ConnectionContext() {
+    if (resource_abandoned_) {
+        LOG(WARNING) << "Resource leak in ConnectionContext: cleanup skipped "
+                        "due to hung operations.";
+        return;
+    }
+
     for (int i = 0; i < groupSize_; ++i) {
         if (peerStates_[i].segmentId.has_value()) {
             engine_->closeSegment(peerStates_[i].segmentId.value());
@@ -387,13 +394,20 @@ bool ConnectionContext::pollPeer(int pollingRank) {
             meta_->activeRanksTensor[pollingRank] = 0;
 
             // Reset store
-            store_->deleteKey(
-                getServerNameStoreKey(backendIndex_, pollingRank));
-            store_->deleteKey(getBufferStoreKey(backendIndex_, pollingRank));
-            store_->deleteKey(
-                getExtensionTaskCountStoreKey(backendIndex_, pollingRank));
-            store_->deleteKey(
-                getExtensionActiveRanksStoreKey(backendIndex_, pollingRank));
+            try {
+                store_->deleteKey(
+                    getServerNameStoreKey(backendIndex_, pollingRank));
+                store_->deleteKey(
+                    getBufferStoreKey(backendIndex_, pollingRank));
+                store_->deleteKey(
+                    getExtensionTaskCountStoreKey(backendIndex_, pollingRank));
+                store_->deleteKey(getExtensionActiveRanksStoreKey(backendIndex_,
+                                                                  pollingRank));
+            } catch (const std::exception& e) {
+                LOG(WARNING) << "Rank " << rank_
+                             << " got an exception when deleteKey for peer "
+                             << pollingRank << ": " << e.what();
+            }
 
             // Reset warmup region
             *reinterpret_cast<volatile int32_t*>(
@@ -420,22 +434,47 @@ bool ConnectionContext::pollPeer(int pollingRank) {
     return state_changed;
 }
 
+bool ConnectionContext::isStopped() const {
+    for (auto& peerState : peerStates_) {
+        if (peerState.state != PeerConnectionState::EXPIRING) {
+            return false;
+        }
+    }
+    return true;
+}
+
+bool ConnectionContext::drainPoller() const {
+    BackoffWaiter waiter;
+    return waiter.wait_for(std::chrono::milliseconds(kDrainPollerTimeoutMs),
+                           [this] { return isStopped(); });
+}
+
+void ConnectionContext::abandonResources() { resource_abandoned_ = true; }
+
 bool ConnectionContext::tryStop() {
     bool stopped = true;
     for (auto& peerState : peerStates_) {
-        if (peerState.state == PeerConnectionState::WAITING_WARMUP_TRANSFER) {
-            TransferStatus status;
-            engine_->getTransferStatus(peerState.warmupBatchId.value(), 0,
-                                       status);
+        if (peerState.state == PeerConnectionState::EXPIRING) {
+            continue;
+        }
 
-            if (status.s == TransferStatusEnum::COMPLETED ||
-                status.s == TransferStatusEnum::FAILED) {
-                engine_->freeBatchID(peerState.warmupBatchId.value());
-                peerState.warmupBatchId = std::nullopt;
-                peerState.state = PeerConnectionState::EXPIRING;
-            } else {
-                stopped = false;
-            }
+        if (peerState.state != PeerConnectionState::WAITING_WARMUP_TRANSFER) {
+            peerState.state = PeerConnectionState::EXPIRING;
+            continue;
+        }
+
+        // For WAITING_WARMUP_TRANSFER, wait for the existing transfer to
+        // complete so that we can safely release the registered memory.
+        TransferStatus status;
+        engine_->getTransferStatus(peerState.warmupBatchId.value(), 0, status);
+
+        if (status.s == TransferStatusEnum::COMPLETED ||
+            status.s == TransferStatusEnum::FAILED) {
+            engine_->freeBatchID(peerState.warmupBatchId.value());
+            peerState.warmupBatchId = std::nullopt;
+            peerState.state = PeerConnectionState::EXPIRING;
+        } else {
+            stopped = false;
         }
     }
     return stopped;
@@ -466,6 +505,7 @@ void ConnectionPoller::registerContext(
 void ConnectionPoller::removeContext(
     const std::shared_ptr<ConnectionContext>& ctx) {
     TORCH_CHECK(ctx->isShutdown_, "connection context hasn't shutdown.");
+
     {
         std::lock_guard<std::mutex> lock(contexts_mutex_);
         contexts_.erase(std::remove(contexts_.begin(), contexts_.end(), ctx),
@@ -529,8 +569,8 @@ void ConnectionPoller::pollerLoop() {
         if (did_work) continue;
 
         std::unique_lock<std::mutex> lock(wakeup_mutex_);
-        auto sleep_ms = all_connected ? ALL_CONNECTED_IDLE_SLEEP_MS
-                                      : CONNECTING_IDLE_SLEEP_MS;
+        auto sleep_ms =
+            all_connected ? kAllConnectedIdleSleepMs : kConnectingIdleSleepMs;
         wakeup_cv_.wait_for(lock, std::chrono::milliseconds(sleep_ms), [&]() {
             if (local_version !=
                 contexts_version_.load(std::memory_order_acquire))

mooncake-pg/src/mooncake_backend.cpp

@@ -12,6 +12,7 @@
 #include "connection_poller.h"
 #include "memory_location.h"
 #include "mooncake_worker.cuh"
+#include "pg_utils.h"
 
 namespace mooncake {
 
@@ -68,18 +69,18 @@ class MooncakeP2PWork : public ::c10d::Work {
             return true;
         }
 
-        auto start = std::chrono::steady_clock::now();
-        while (!completed_->load(std::memory_order_acquire)) {
-            auto now = std::chrono::steady_clock::now();
-            auto elapsed =
-                std::chrono::duration_cast<std::chrono::milliseconds>(now -
-                                                                      start);
-            if (timeout.count() > 0 && elapsed >= timeout) {
-                return false;
-            }
-            std::this_thread::sleep_for(std::chrono::microseconds(10));
+        BackoffWaiterConfig cfg{};
+        cfg.max_sleep = std::chrono::microseconds(10);
+        BackoffWaiter waiter(cfg);
+
+        if (timeout.count() > 0) {
+            return waiter.wait_for(timeout, [this] {
+                return completed_->load(std::memory_order_acquire);
+            });
         }
 
+        waiter.wait(
+            [this] { return completed_->load(std::memory_order_acquire); });
         return true;
     }
 
@@ -102,17 +103,18 @@ MooncakeBackend::MooncakeBackend(
     const int size = distBackendOpts.group_size;
     const auto& globalRanks = distBackendOpts.global_ranks_in_group;
 
-    // Get device data
-    std::string location;
-    int deviceCount = 0;
-    cudaError_t err = cudaGetDeviceCount(&deviceCount);
-    if (err != cudaSuccess || deviceCount == 0) {
-        location = kWildcardLocation;
-    } else {
-        int deviceId_;
-        err = cudaGetDevice(&deviceId_);
-        TORCH_CHECK(!err, c10::str("Failed to get device id"));
-        location = GPU_PREFIX + std::to_string(deviceId_);
+    // Memory location for device specific buffers
+    // always kWildcardLocation for cpu backend
+    std::string location = kWildcardLocation;
+    if (!isCpu) {
+        int deviceCount = 0;
+        cudaError_t err = cudaGetDeviceCount(&deviceCount);
+        if (err == cudaSuccess && deviceCount != 0) {
+            int deviceId_;
+            err = cudaGetDevice(&deviceId_);
+            TORCH_CHECK(!err, c10::str("Failed to get device id"));
+            location = GPU_PREFIX + std::to_string(deviceId_);
+        }
     }
 
     // Initialize transfer engine
@@ -389,15 +391,18 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::broadcast(
         return worker_->putTaskCuda(
             c10d::OpType::BROADCAST, tensorSize, root, meta_, connection_ctx_,
             stream,
-            [=](void* dst, size_t pos, size_t realSize) {
+            [=](void* dst, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 if (isRoot) {
                     cudaMemcpyAsync(dst, (char*)tensor.data_ptr() + pos,
-                                    realSize, cudaMemcpyDeviceToDevice, stream);
+                                    realSize, cudaMemcpyDeviceToDevice,
+                                    enq_stream);
                 }
             },
-            [=](void* src, size_t pos, size_t realSize) {
+            [=](void* src, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 cudaMemcpyAsync((char*)tensor.data_ptr() + pos, src, realSize,
-                                cudaMemcpyDeviceToDevice, stream);
+                                cudaMemcpyDeviceToDevice, enq_stream);
             });
     }
 }
@@ -425,16 +430,18 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::allreduce(
         return worker_->putTaskCuda(
             c10d::OpType::ALLREDUCE, tensorSize, 0, meta_, connection_ctx_,
             stream,
-            [=](void* dst, size_t pos, size_t realSize) {
+            [=](void* dst, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 cudaMemcpyAsync(dst, (char*)tensor.data_ptr() + pos, realSize,
-                                cudaMemcpyDeviceToDevice, stream);
+                                cudaMemcpyDeviceToDevice, enq_stream);
             },
-            [=, this](void* src, size_t pos, size_t realSize) {
+            [=, this](void* src, size_t pos, size_t realSize,
+                      const at::cuda::CUDAStream& enq_stream) {
                 cudaMemsetAsync((char*)tensor.data_ptr() + pos, 0, realSize,
-                                stream);
+                                enq_stream);
                 launchReduceKernel(tensor, pos, realSize, src, meta_->size,
                                    opts.reduceOp, meta_->activeRanksDevice,
-                                   stream);
+                                   enq_stream);
             });
     }
 }
@@ -465,15 +472,17 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::allgather(
         return worker_->putTaskCuda(
             c10d::OpType::ALLGATHER, tensorSize, 0, meta_, connection_ctx_,
             stream,
-            [=](void* dst, size_t pos, size_t realSize) {
+            [=](void* dst, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 cudaMemcpyAsync(dst, (char*)inputTensor.data_ptr() + pos,
-                                realSize, cudaMemcpyDeviceToDevice, stream);
+                                realSize, cudaMemcpyDeviceToDevice, enq_stream);
             },
-            [=](void* src, size_t pos, size_t realSize) {
+            [=](void* src, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 for (const auto j : c10::irange(outputTensors_.size())) {
                     cudaMemcpyAsync((char*)outputTensors_[j].data_ptr() + pos,
                                     (char*)src + j * realSize, realSize,
-                                    cudaMemcpyDeviceToDevice, stream);
+                                    cudaMemcpyDeviceToDevice, enq_stream);
                 }
             });
     }
@@ -504,16 +513,18 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::_allgather_base(
         return worker_->putTaskCuda(
             c10d::OpType::_ALLGATHER_BASE, tensorSize, 0, meta_,
             connection_ctx_, stream,
-            [=](void* dst, size_t pos, size_t realSize) {
+            [=](void* dst, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 cudaMemcpyAsync(dst, (char*)inputBuffer.data_ptr() + pos,
-                                realSize, cudaMemcpyDeviceToDevice, stream);
+                                realSize, cudaMemcpyDeviceToDevice, enq_stream);
             },
-            [=, this](void* src, size_t pos, size_t realSize) {
+            [=, this](void* src, size_t pos, size_t realSize,
+                      const at::cuda::CUDAStream& enq_stream) {
                 for (const auto j : c10::irange(meta_->size)) {
                     cudaMemcpyAsync(
                         (char*)outputBuffer.data_ptr() + j * tensorSize + pos,
                         (char*)src + j * realSize, realSize,
-                        cudaMemcpyDeviceToDevice, stream);
+                        cudaMemcpyDeviceToDevice, enq_stream);
                 }
             });
     }
@@ -546,20 +557,22 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::_reduce_scatter_base(
         return worker_->putTaskCuda(
             c10d::OpType::_REDUCE_SCATTER_BASE, tensorSize, 0, meta_,
             connection_ctx_, stream,
-            [=, this](void* dst, size_t pos, size_t realSize) {
+            [=, this](void* dst, size_t pos, size_t realSize,
+                      const at::cuda::CUDAStream& enq_stream) {
                 for (const auto j : c10::irange(meta_->size)) {
                     cudaMemcpyAsync(
                         (char*)dst + j * realSize,
                         (char*)inputBuffer.data_ptr() + j * tensorSize + pos,
-                        realSize, cudaMemcpyDeviceToDevice, stream);
+                        realSize, cudaMemcpyDeviceToDevice, enq_stream);
                 }
             },
-            [=, this](void* src, size_t pos, size_t realSize) {
+            [=, this](void* src, size_t pos, size_t realSize,
+                      const at::cuda::CUDAStream& enq_stream) {
                 cudaMemsetAsync((char*)outputBuffer.data_ptr() + pos, 0,
-                                realSize, stream);
+                                realSize, enq_stream);
                 launchReduceKernel(outputBuffer, pos, realSize, src,
                                    meta_->size, opts.reduceOp,
-                                   meta_->activeRanksDevice, stream);
+                                   meta_->activeRanksDevice, enq_stream);
             });
     }
 }
@@ -590,18 +603,21 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::alltoall(
         return worker_->putTaskCuda(
             c10d::OpType::ALLTOALL, tensorSize, 0, meta_, connection_ctx_,
             stream,
-            [=](void* dst, size_t pos, size_t realSize) {
+            [=](void* dst, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 for (const auto j : c10::irange(inputTensors.size())) {
                     cudaMemcpyAsync((char*)dst + j * realSize,
                                     (char*)inputTensors[j].data_ptr() + pos,
-                                    realSize, cudaMemcpyDeviceToDevice, stream);
+                                    realSize, cudaMemcpyDeviceToDevice,
+                                    enq_stream);
                 }
             },
-            [=](void* src, size_t pos, size_t realSize) {
+            [=](void* src, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 for (const auto j : c10::irange(outputTensors.size())) {
                     cudaMemcpyAsync((char*)outputTensors[j].data_ptr() + pos,
                                     (char*)src + j * realSize, realSize,
-                                    cudaMemcpyDeviceToDevice, stream);
+                                    cudaMemcpyDeviceToDevice, enq_stream);
                 }
             });
     }
@@ -617,11 +633,12 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::barrier(
             [=](void*, size_t, size_t) {});
     } else {
         auto device_index = at::cuda::current_device();
-        auto stream = c10::cuda::getDefaultCUDAStream(device_index);
+        auto stream = at::cuda::getCurrentCUDAStream(device_index);
         return worker_->putTaskCuda(
             c10d::OpType::BARRIER, kBarrierDummyTensorSize, 0, meta_,
-            connection_ctx_, stream, [=](void*, size_t, size_t) {},
-            [=](void*, size_t, size_t) {});
+            connection_ctx_, stream,
+            [=](void*, size_t, size_t, const at::cuda::CUDAStream&) {},
+            [=](void*, size_t, size_t, const at::cuda::CUDAStream&) {});
     }
 }
 
@@ -651,17 +668,19 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::reduce(
         return worker_->putTaskCuda(
             c10d::OpType::REDUCE, tensorSize, root, meta_, connection_ctx_,
             stream,
-            [=](void* dst, size_t pos, size_t realSize) {
+            [=](void* dst, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 cudaMemcpyAsync(dst, (char*)tensor.data_ptr() + pos, realSize,
-                                cudaMemcpyDeviceToDevice, stream);
+                                cudaMemcpyDeviceToDevice, enq_stream);
             },
-            [=, this](void* src, size_t pos, size_t realSize) {
+            [=, this](void* src, size_t pos, size_t realSize,
+                      const at::cuda::CUDAStream& enq_stream) {
                 if (isRoot) {
                     cudaMemsetAsync((char*)tensor.data_ptr() + pos, 0, realSize,
-                                    stream);
+                                    enq_stream);
                     launchReduceKernel(tensor, pos, realSize, src, meta_->size,
                                        opts.reduceOp, meta_->activeRanksDevice,
-                                       stream);
+                                       enq_stream);
                 }
             });
     }
@@ -699,18 +718,20 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::gather(
         return worker_->putTaskCuda(
             c10d::OpType::GATHER, tensorSize, root, meta_, connection_ctx_,
             stream,
-            [=](void* dst, size_t pos, size_t realSize) {
+            [=](void* dst, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 cudaMemcpyAsync(dst, (char*)inputTensor.data_ptr() + pos,
-                                realSize, cudaMemcpyDeviceToDevice, stream);
+                                realSize, cudaMemcpyDeviceToDevice, enq_stream);
             },
-            [=](void* src, size_t pos, size_t realSize) {
+            [=](void* src, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 if (isRoot) {
                     auto outputTensors_ = outputTensors.back();
                     for (const auto j : c10::irange(outputTensors_.size())) {
                         cudaMemcpyAsync(
                             (char*)outputTensors_[j].data_ptr() + pos,
                             (char*)src + j * realSize, realSize,
-                            cudaMemcpyDeviceToDevice, stream);
+                            cudaMemcpyDeviceToDevice, enq_stream);
                     }
                 }
             });
@@ -751,20 +772,22 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::scatter(
         return worker_->putTaskCuda(
             c10d::OpType::SCATTER, tensorSize, root, meta_, connection_ctx_,
             stream,
-            [=](void* dst, size_t pos, size_t realSize) {
+            [=](void* dst, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 if (isRoot) {
                     auto inputTensors_ = inputTensors.back();
                     for (const auto j : c10::irange(inputTensors_.size())) {
                         cudaMemcpyAsync(
                             (char*)dst + j * realSize,
                             (char*)inputTensors_[j].data_ptr() + pos, realSize,
-                            cudaMemcpyDeviceToDevice, stream);
+                            cudaMemcpyDeviceToDevice, enq_stream);
                     }
                 }
             },
-            [=](void* src, size_t pos, size_t realSize) {
+            [=](void* src, size_t pos, size_t realSize,
+                const at::cuda::CUDAStream& enq_stream) {
                 cudaMemcpyAsync((char*)outputTensor.data_ptr() + pos, src,
-                                realSize, cudaMemcpyDeviceToDevice, stream);
+                                realSize, cudaMemcpyDeviceToDevice, enq_stream);
             });
     }
 }
@@ -775,28 +798,52 @@ void MooncakeBackend::shutdown() {
     }
     isShutdown_ = true;
 
-    p2p_device_worker_->removeProxy(p2p_proxy_);
-    p2p_proxy_.reset();
+    // If we encounter any hung operations, don't release resources
+    // to avoid potential crash. Instead, we allow those resources to leak
+    // and rely on the OS to reclaim them later.
+    bool has_hung_operation = false;
 
+    // Phase 1: Drain P2P tasks
+    p2p_device_worker_->removeProxy(p2p_proxy_);
+    has_hung_operation |= !p2p_proxy_->DrainTasks();
+
+    // Phase 2: Drain collective tasks for this backend
+    has_hung_operation |= !worker_->drainTasks(meta_.get());
+
+    // Phase 3: Drain warm-up transfers for connection poller
     connection_ctx_->shutdown();
     if (connectionPollerRegistered_) {
         ConnectionPoller::GetInstance().removeContext(connection_ctx_);
+        has_hung_operation |= !connection_ctx_->drainPoller();
         connectionPollerRegistered_ = false;
     }
 
-    for (size_t i = 0; i < 2; i++) {
-        engine_->unregisterLocalMemory(cpu_sync_send_region_[i]);
-        engine_->unregisterLocalMemory(cpu_sync_recv_region_[i]);
-        engine_->unregisterLocalMemory(send_buffer_[i]);
-        engine_->unregisterLocalMemory(recv_buffer_[i]);
-        delete[] cpu_sync_send_region_[i];
-        delete[] cpu_sync_recv_region_[i];
-        if (isCpu_) {
-            free(send_buffer_[i]);
-            free(recv_buffer_[i]);
-        } else {
-            cudaFree(send_buffer_[i]);
-            cudaFree(recv_buffer_[i]);
+    // Phase 4: CUDA synchronization
+    if (!isCpu_ && !has_hung_operation) {
+        cudaDeviceSynchronize();
+    }
+
+    // Phase 5: Release resources if no hung operations
+    if (has_hung_operation) {
+        p2p_proxy_->AbandonResources();
+        connection_ctx_->abandonResources();
+    }
+
+    if (!has_hung_operation) {
+        for (size_t i = 0; i < 2; i++) {
+            engine_->unregisterLocalMemory(cpu_sync_send_region_[i]);
+            engine_->unregisterLocalMemory(cpu_sync_recv_region_[i]);
+            engine_->unregisterLocalMemory(send_buffer_[i]);
+            engine_->unregisterLocalMemory(recv_buffer_[i]);
+            delete[] cpu_sync_send_region_[i];
+            delete[] cpu_sync_recv_region_[i];
+            if (isCpu_) {
+                free(send_buffer_[i]);
+                free(recv_buffer_[i]);
+            } else {
+                cudaFree(send_buffer_[i]);
+                cudaFree(recv_buffer_[i]);
+            }
         }
     }
 }
@@ -854,21 +901,20 @@ void MooncakeBackend::waitForExtensionState() {
     auto active_ranks_key = ConnectionContext::getExtensionActiveRanksStoreKey(
         meta_->backendIndex, rank_);
 
-    while (true) {
-        if (meta_->store->check({task_count_key, active_ranks_key})) {
-            auto task_count_data = meta_->store->get(task_count_key);
-            std::string task_count(task_count_data.begin(),
-                                   task_count_data.end());
-            meta_->taskCount = std::stoi(task_count);
+    BackoffWaiter waiter(
+        BackoffWaiterConfig::constantSleep(std::chrono::milliseconds(50)));
 
-            auto active_ranks = meta_->store->get(active_ranks_key);
-            deserializeActiveRanks(active_ranks, meta_->activeRanks,
-                                   meta_->size);
-            syncActiveRanksTensor();
-            return;
-        }
-        std::this_thread::sleep_for(std::chrono::milliseconds(50));
-    }
+    waiter.wait([&] {
+        return meta_->store->check({task_count_key, active_ranks_key});
+    });
+
+    auto task_count_data = meta_->store->get(task_count_key);
+    std::string task_count(task_count_data.begin(), task_count_data.end());
+    meta_->taskCount = std::stoi(task_count);
+
+    auto active_ranks = meta_->store->get(active_ranks_key);
+    deserializeActiveRanks(active_ranks, meta_->activeRanks, meta_->size);
+    syncActiveRanksTensor();
 }
 
 int MooncakeBackend::getNumSyncedRanks() {

mooncake-pg/src/mooncake_worker.cu

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

mooncake-pg/src/mooncake_worker_thread.cpp

@@ -3,6 +3,7 @@
 #include <mooncake_worker.cuh>
 #include <glog/logging.h>
 #include <transfer_engine.h>
+#include "pg_utils.h"
 
 namespace mooncake {
 
@@ -22,9 +23,51 @@ void MooncakeWorker::Start() {
     }
 }
 
+bool MooncakeWorker::drainTasks(const TransferGroupMeta* meta) const {
+    BackoffWaiter waiter;
+    return waiter.wait_for(
+        std::chrono::milliseconds(kDrainTasksTimeoutMs), [this, meta] {
+            for (size_t i = 0; i < kNumTasks_; ++i) {
+                if (tasks_[i].active && tasks_[i].transferGroupMeta == meta)
+                    return false;
+            }
+            return true;
+        });
+}
+
+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;
-    std::thread([this] {
+    worker_thread_ = std::thread([this] {
         if (cuda_device_index_ >= 0) {
             cudaSetDevice(cuda_device_index_);
         }
@@ -32,7 +75,6 @@ void MooncakeWorker::startWorker() {
         using clock = std::chrono::high_resolution_clock;
         clock::time_point activeTime[kNumTasks_];
         size_t rankToTaskId[kNumTasks_][kMaxNumRanks];
-        TransferMetadata::NotifyDesc msg{"ping", "ping"};
         while (running_) {
             PAUSE();
             for (size_t i = 0; i < kNumTasks_; ++i) {
@@ -49,7 +91,10 @@ void MooncakeWorker::startWorker() {
                                      group->rank != task.broadcastRoot) ||
                                     task.opType == c10d::OpType::BARRIER;
                 if (task_status[i].load(std::memory_order_acquire) == IDLE) {
+                    const auto submit_sequence = task.submitSequence;
                     if (skipTransfer) {
+                        submitted_task_sequence_[i].store(
+                            submit_sequence, std::memory_order_release);
                         task_status[i].store(TRANSFERRED_1,
                                              std::memory_order_release);
                         continue;
@@ -121,6 +166,8 @@ void MooncakeWorker::startWorker() {
                     task.batchID =
                         group->engine->allocateBatchID(entries.size());
                     group->engine->submitTransfer(task.batchID, entries);
+                    submitted_task_sequence_[i].store(
+                        submit_sequence, std::memory_order_release);
                     activeTime[i] = clock::now();
                     task_status[i].store(TRANSFERRED_1,
                                          std::memory_order_release);
@@ -146,8 +193,9 @@ void MooncakeWorker::startWorker() {
                                 if (status.s == TransferStatusEnum::FAILED ||
                                     (j != group->rank &&
                                      diff.count() > kPingTimeoutMicroseconds_ &&
-                                     group->engine->sendNotifyByID(
-                                         group->segmentIDs[j], msg))) {
+                                     group->engine->probePeerAliveByID(
+                                         group->segmentIDs[j]) !=
+                                         PeerLiveness::Alive)) {
                                     LOG(ERROR)
                                         << "Rank " << group->rank
                                         << " marking peer " << j
@@ -235,8 +283,9 @@ void MooncakeWorker::startWorker() {
                             if (status.s == TransferStatusEnum::FAILED ||
                                 (j != group->rank &&
                                  diff.count() > kPingTimeoutMicroseconds_ &&
-                                 group->engine->sendNotifyByID(
-                                     group->segmentIDs[j], msg))) {
+                                 group->engine->probePeerAliveByID(
+                                     group->segmentIDs[j]) !=
+                                     PeerLiveness::Alive)) {
                                 LOG(ERROR) << "Rank " << group->rank
                                            << " marking peer " << j
                                            << " as broken during syncing op "
@@ -281,7 +330,7 @@ void MooncakeWorker::startWorker() {
                 }
             }
         }
-    }).detach();
+    });
 }
 
 std::shared_ptr<MooncakeWorker> MooncakeWorkerManager::GetWorker(

mooncake-pg/src/p2p_proxy.cpp

@@ -10,6 +10,7 @@
 #include <limits>
 #include <thread>
 #include "memory_location.h"
+#include "pg_utils.h"
 
 namespace mooncake {
 
@@ -54,7 +55,15 @@ P2PProxy::P2PProxy(TransferEngine* engine, const Options& options)
     AllocateResources();
 }
 
-P2PProxy::~P2PProxy() { ReleaseResources(); }
+P2PProxy::~P2PProxy() {
+    if (resource_abandoned_) {
+        LOG(WARNING) << "Resource leak in P2PProxy: cleanup skipped due to "
+                        "hung operations.";
+        return;
+    }
+
+    ReleaseResources();
+}
 
 void P2PProxy::BindMeta(const std::shared_ptr<TransferGroupMeta>& meta) {
     meta_ = meta;
@@ -214,6 +223,9 @@ void P2PProxy::PerformRecvReset(int peer_rank) {
 }
 
 void P2PProxy::ReleaseResources() {
+    TORCH_CHECK(!resource_abandoned_,
+                "Should not release abandoned resources.");
+
     SetCudaDeviceIfNeeded(is_cpu_, cuda_device_index_,
                           "P2PProxy ReleaseResources cudaSetDevice failed");
 
@@ -287,6 +299,8 @@ void P2PProxy::ReleaseResources() {
     }
 }
 
+void P2PProxy::AbandonResources() { resource_abandoned_ = true; }
+
 void P2PProxy::EnqueueSend(SendOp op) {
     op.tensor_ =
         op.tensor_.is_contiguous() ? op.tensor_ : op.tensor_.contiguous();
@@ -899,6 +913,13 @@ bool P2PProxy::HasActiveRecvWork() const {
     return active_recv_tasks_.load(std::memory_order_acquire) > 0;
 }
 
+bool P2PProxy::DrainTasks() const {
+    BackoffWaiter waiter;
+    return waiter.wait_for(
+        std::chrono::milliseconds(kDrainTasksTimeoutMs),
+        [this] { return !HasActiveSendWork() && !HasActiveRecvWork(); });
+}
+
 void P2PDeviceWorker::Start() {
     bool expected_send = false;
     if (send_worker_running_.compare_exchange_strong(expected_send, true)) {

mooncake-pg/tests/README.md

@@ -0,0 +1,29 @@
+# PG Tests
+
+## Environment Variables
+
+- `MOONCAKE_PGTEST_DEVICE_FILTERS`
+  Comma-separated NIC / IB device names passed to
+  `pg.set_device_filter(...)`. Leave unset to use the backend default device
+  selection.
+
+- `MOONCAKE_PGTEST_MASTER_ADDR`
+  Rendezvous address used for the local test process group. Defaults to
+  `127.0.0.1`.
+
+- `MOONCAKE_PGTEST_MASTER_PORT`
+  Rendezvous port used for the local test process group. If unset, each test
+  allocates a free local port automatically.
+
+## Usage
+
+```bash
+# Run all test cases
+python -m unittest discover -s mooncake-pg/tests -v
+
+# Run CUDA test cases
+python -m unittest discover -s mooncake-pg/tests -k CUDA -v
+
+# Run CPU-only test cases
+python -m unittest discover -s mooncake-pg/tests -k CPU -v
+```

mooncake-pg/tests/pg_test_topology.py

@@ -0,0 +1,244 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterable
+
+import torch.distributed as dist
+
+
+@dataclass(frozen=True)
+class GroupSpec:
+    name: str
+    family: str
+    ranks: tuple[int, ...]
+
+
+@dataclass(frozen=True)
+class TopologySpec:
+    world_size: int
+    ordered_groups: tuple[GroupSpec, ...]
+    p2p_lanes: tuple[tuple[int, int], ...] = ()
+
+    def membership_for_rank(self, rank: int) -> dict[str, list[list[int]]]:
+        membership: dict[str, list[list[int]]] = {}
+        for group in self.ordered_groups:
+            if rank in group.ranks:
+                membership.setdefault(group.family, []).append(list(group.ranks))
+        return membership
+
+
+def find_local_group_spec(
+    topology: TopologySpec,
+    rank: int,
+    families: str | Iterable[str],
+) -> GroupSpec | None:
+    if isinstance(families, str):
+        families = (families,)
+    family_set = set(families)
+    for group_spec in topology.ordered_groups:
+        if group_spec.family in family_set and rank in group_spec.ranks:
+            return group_spec
+    return None
+
+
+def find_local_named_group(
+    topology: TopologySpec,
+    groups: dict[str, object],
+    rank: int,
+    families: str | Iterable[str],
+) -> tuple[GroupSpec | None, object | None]:
+    group_spec = find_local_group_spec(topology, rank, families)
+    if group_spec is None:
+        return None, None
+    return group_spec, groups.get(group_spec.name)
+
+def create_named_groups(
+    topology: TopologySpec,
+    *,
+    backend: str = "mooncake-cpu",
+) -> dict[str, object]:
+    groups: dict[str, object] = {}
+    rank = dist.get_rank()
+    for group_spec in topology.ordered_groups:
+        group = dist.new_group(ranks=list(group_spec.ranks), backend=backend)
+        if rank in group_spec.ranks:
+            groups[group_spec.name] = group
+    return groups
+
+
+def destroy_named_groups(groups: dict[str, object]) -> None:
+    for name in reversed(list(groups.keys())):
+        try:
+            dist.destroy_process_group(groups[name])
+        except Exception:
+            pass
+
+
+def build_tp_only_topology() -> TopologySpec:
+    return TopologySpec(
+        world_size=4,
+        ordered_groups=(
+            GroupSpec("tp0", "tp", (0, 1, 2, 3)),
+        ),
+    )
+
+
+def build_tp_only_topology_2() -> TopologySpec:
+    return TopologySpec(
+        world_size=2,
+        ordered_groups=(
+            GroupSpec("tp0", "tp", (0, 1)),
+        ),
+    )
+
+
+def build_dp_tp_topology() -> TopologySpec:
+    ordered_groups = (
+        GroupSpec("tp0", "tp", (0, 1, 2, 3)),
+        GroupSpec("tp1", "tp", (4, 5, 6, 7)),
+        GroupSpec("dp0", "dp", (0, 4)),
+        GroupSpec("dp1", "dp", (1, 5)),
+        GroupSpec("dp2", "dp", (2, 6)),
+        GroupSpec("dp3", "dp", (3, 7)),
+    )
+    return TopologySpec(world_size=8, ordered_groups=ordered_groups)
+
+
+def build_dp_tp_topology_4() -> TopologySpec:
+    ordered_groups = (
+        GroupSpec("tp0", "tp", (0, 1)),
+        GroupSpec("tp1", "tp", (2, 3)),
+        GroupSpec("dp0", "dp", (0, 2)),
+        GroupSpec("dp1", "dp", (1, 3)),
+    )
+    return TopologySpec(world_size=4, ordered_groups=ordered_groups)
+
+
+def build_tp_pp_topology() -> TopologySpec:
+    ordered_groups = (
+        GroupSpec("tp_stage0", "tp", (0, 1)),
+        GroupSpec("tp_stage1", "tp", (2, 3)),
+        GroupSpec("pp_lane0", "pp", (0, 2)),
+        GroupSpec("pp_lane1", "pp", (1, 3)),
+    )
+    return TopologySpec(
+        world_size=4,
+        ordered_groups=ordered_groups,
+        p2p_lanes=((0, 2), (1, 3)),
+    )
+
+
+def build_tp_pp_topology_2() -> TopologySpec:
+    ordered_groups = (
+        GroupSpec("tp_stage0", "tp", (0, 1)),
+        GroupSpec("pp_lane0", "pp", (0, 1)),
+    )
+    return TopologySpec(
+        world_size=2,
+        ordered_groups=ordered_groups,
+        p2p_lanes=((0, 1),),
+    )
+
+
+def build_dp_tp_ep_topology() -> TopologySpec:
+    ordered_groups = (
+        GroupSpec("tp_dp0_ep0", "tp", (0, 2)),
+        GroupSpec("tp_dp0_ep1", "tp", (1, 3)),
+        GroupSpec("tp_dp1_ep0", "tp", (4, 6)),
+        GroupSpec("tp_dp1_ep1", "tp", (5, 7)),
+        GroupSpec("dp_tp0_ep0", "dp", (0, 4)),
+        GroupSpec("dp_tp0_ep1", "dp", (1, 5)),
+        GroupSpec("dp_tp1_ep0", "dp", (2, 6)),
+        GroupSpec("dp_tp1_ep1", "dp", (3, 7)),
+        GroupSpec("ep_dp0_tp0", "ep", (0, 1)),
+        GroupSpec("ep_dp0_tp1", "ep", (2, 3)),
+        GroupSpec("ep_dp1_tp0", "ep", (4, 5)),
+        GroupSpec("ep_dp1_tp1", "ep", (6, 7)),
+    )
+    return TopologySpec(world_size=8, ordered_groups=ordered_groups)
+
+
+def build_dp_tp_ep_topology_4() -> TopologySpec:
+    ordered_groups = (
+        GroupSpec("tp_dp0_ep0", "tp", (0, 2)),
+        GroupSpec("tp_dp0_ep1", "tp", (1, 3)),
+        GroupSpec("dp_tp0_ep0", "dp", (0, 1)),
+        GroupSpec("dp_tp1_ep0", "dp", (2, 3)),
+        GroupSpec("ep_dp0_tp0", "ep", (0, 1)),
+        GroupSpec("ep_dp1_tp0", "ep", (2, 3)),
+    )
+    return TopologySpec(world_size=4, ordered_groups=ordered_groups)
+
+
+def build_prefill_decode_topology() -> TopologySpec:
+    ordered_groups = (
+        GroupSpec("prefill", "service", (0, 1)),
+        GroupSpec("decode", "service", (2, 3)),
+        GroupSpec("lane0", "lane", (0, 2)),
+        GroupSpec("lane1", "lane", (1, 3)),
+    )
+    return TopologySpec(
+        world_size=4,
+        ordered_groups=ordered_groups,
+        p2p_lanes=((0, 2), (1, 3)),
+    )
+
+
+def build_prefill_decode_topology_2() -> TopologySpec:
+    ordered_groups = (
+        GroupSpec("prefill", "service", (0,)),
+        GroupSpec("decode", "service", (1,)),
+        GroupSpec("lane0", "lane", (0, 1)),
+    )
+    return TopologySpec(
+        world_size=2,
+        ordered_groups=ordered_groups,
+        p2p_lanes=((0, 1),),
+    )
+
+
+def build_many_group_smoke_topology() -> TopologySpec:
+    ordered_groups = (
+        GroupSpec("tp0", "tp", (0, 1, 2, 3)),
+        GroupSpec("tp1", "tp", (4, 5, 6, 7)),
+        GroupSpec("dp0", "dp", (0, 4)),
+        GroupSpec("dp1", "dp", (1, 5)),
+        GroupSpec("dp2", "dp", (2, 6)),
+        GroupSpec("dp3", "dp", (3, 7)),
+        GroupSpec("ep0", "ep", (0, 1)),
+        GroupSpec("ep1", "ep", (2, 3)),
+        GroupSpec("ep2", "ep", (4, 5)),
+        GroupSpec("ep3", "ep", (6, 7)),
+        GroupSpec("lane0", "pp", (0, 6)),
+        GroupSpec("lane1", "pp", (1, 7)),
+        GroupSpec("lane2", "pp", (2, 4)),
+        GroupSpec("lane3", "pp", (3, 5)),
+    )
+    return TopologySpec(
+        world_size=8,
+        ordered_groups=ordered_groups,
+        p2p_lanes=((0, 6), (1, 7), (2, 4), (3, 5)),
+    )
+
+
+def build_many_group_smoke_topology_4() -> TopologySpec:
+    ordered_groups = (
+        GroupSpec("tp0", "tp", (0, 1)),
+        GroupSpec("tp1", "tp", (2, 3)),
+        GroupSpec("dp0", "dp", (0, 2)),
+        GroupSpec("dp1", "dp", (1, 3)),
+        GroupSpec("ep0", "ep", (0, 3)),
+        GroupSpec("ep1", "ep", (1, 2)),
+    )
+    return TopologySpec(world_size=4, ordered_groups=ordered_groups)
+
+
+def build_many_group_smoke_topology_2() -> TopologySpec:
+    ordered_groups = (
+        GroupSpec("tp0", "tp", (0, 1)),
+        GroupSpec("dp0", "dp", (0, 1)),
+    )
+    return TopologySpec(
+        world_size=2,
+        ordered_groups=ordered_groups,
+    )

mooncake-pg/tests/pg_test_utils.py

@@ -0,0 +1,579 @@
+import os
+import signal
+import socket
+import time
+import unittest
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Sequence
+
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from mooncake import pg
+
+
+DEVICE_FILTER_ENV_VAR = "MOONCAKE_PGTEST_DEVICE_FILTERS"
+MASTER_ADDR_ENV_VAR = "MOONCAKE_PGTEST_MASTER_ADDR"
+MASTER_PORT_ENV_VAR = "MOONCAKE_PGTEST_MASTER_PORT"
+DEFAULT_MASTER_ADDR = "127.0.0.1"
+DEFAULT_WAIT_TIMEOUT_S = 30.0
+DEFAULT_SPAWN_TIMEOUT_S = 30.0
+
+
+def parse_device_filters(raw: str | None) -> list[str] | None:
+    if raw is None:
+        return None
+    filters = [item.strip() for item in raw.split(",") if item.strip()]
+    return filters or None
+
+
+def resolve_device_filters(
+    device_filters: Sequence[str] | None = None,
+) -> list[str] | None:
+    if device_filters is not None:
+        resolved = [item.strip() for item in device_filters if item.strip()]
+        return resolved or None
+    return parse_device_filters(os.getenv(DEVICE_FILTER_ENV_VAR))
+
+
+def configure_mooncake_device_filter(
+    device_filters: Sequence[str] | None = None,
+) -> list[str] | None:
+    resolved = resolve_device_filters(device_filters)
+    if resolved is not None:
+        pg.set_device_filter(resolved)
+    return resolved
+
+
+def resolve_env_value(env_var: str, default: str) -> str:
+    value = os.getenv(env_var)
+    if value:
+        return value
+    return default
+
+
+def find_free_local_port() -> str:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        sock.bind((DEFAULT_MASTER_ADDR, 0))
+        return str(sock.getsockname()[1])
+
+
+@contextmanager
+def temporary_env(updates: dict[str, str]):
+    previous = {key: os.environ.get(key) for key in updates}
+    os.environ.update(updates)
+    try:
+        yield
+    finally:
+        for key, value in previous.items():
+            if value is None:
+                os.environ.pop(key, None)
+            else:
+                os.environ[key] = value
+
+
+def cuda_runtime_available(min_devices: int = 1) -> bool:
+    if not torch.cuda.is_available():
+        return False
+    return torch.cuda.device_count() >= min_devices
+
+
+def require_test_device(rank: int, device_type: str) -> torch.device:
+    if device_type == "cuda":
+        device_count = torch.cuda.device_count()
+        if device_count <= 0:
+            raise RuntimeError(
+                "CUDA backend requested but no CUDA devices are available"
+            )
+        if rank >= device_count:
+            raise RuntimeError(
+                f"rank {rank} requires a dedicated CUDA device but only {device_count} are visible"
+            )
+        torch.cuda.set_device(rank)
+        return torch.device("cuda", rank)
+    if device_type != "cpu":
+        raise ValueError(f"unsupported device_type: {device_type}")
+    return torch.device("cpu")
+
+
+def mooncake_backend_options(
+    world_size: int,
+    device_type: str,
+    *,
+    active_value: int = 0,
+    is_extension: bool = False,
+) -> pg.MooncakeBackendOptions:
+    device = torch.device(device_type)
+    active_ranks = torch.full(
+        (world_size,),
+        int(active_value),
+        dtype=torch.int32,
+        device=device,
+    )
+    if is_extension:
+        return pg.MooncakeBackendOptions(active_ranks, True)
+    return pg.MooncakeBackendOptions(active_ranks)
+
+
+def mooncake_cpu_options(world_size: int) -> pg.MooncakeBackendOptions:
+    return mooncake_backend_options(world_size, "cpu", active_value=0)
+
+
+def mooncake_extension_cpu_options(world_size: int) -> pg.MooncakeBackendOptions:
+    return mooncake_backend_options(
+        world_size,
+        "cpu",
+        active_value=1,
+        is_extension=True,
+    )
+
+
+def init_mooncake_group(
+    rank: int,
+    world_size: int,
+    *,
+    backend_name: str,
+    device_type: str,
+    device_filters: Sequence[str] | None = None,
+    use_pg_options: bool = True,
+    is_extension: bool = False,
+    active_value: int | None = None,
+) -> torch.device:
+    device = require_test_device(rank, device_type)
+    configure_mooncake_device_filter(device_filters)
+    kwargs = {
+        "backend": backend_name,
+        "rank": rank,
+        "world_size": world_size,
+    }
+    if device_type == "cuda":
+        kwargs["device_id"] = device
+    if use_pg_options:
+        resolved_active_value = (
+            1 if is_extension else 0 if active_value is None else active_value
+        )
+        kwargs["pg_options"] = mooncake_backend_options(
+            world_size,
+            device_type,
+            active_value=resolved_active_value,
+            is_extension=is_extension,
+        )
+    dist.init_process_group(**kwargs)
+    return device
+
+
+def init_mooncake_cpu_group(
+    rank: int,
+    world_size: int,
+    *,
+    device_filters: Sequence[str] | None = None,
+    use_pg_options: bool = True,
+) -> None:
+    init_mooncake_group(
+        rank,
+        world_size,
+        backend_name="mooncake-cpu",
+        device_type="cpu",
+        device_filters=device_filters,
+        use_pg_options=use_pg_options,
+    )
+
+
+def get_mooncake_backend(group=None, device_type: str = "cpu"):
+    if group is None:
+        group = dist.group.WORLD
+    return group._get_backend(torch.device(device_type))
+
+
+@dataclass(slots=True)
+class MooncakePGWorkerContext:
+    proc_rank: int
+    world_size: int
+    result_map: object
+    device_filters: Sequence[str] | None
+    backend_name: str
+    device_type: str
+    _device: torch.device | None = None
+
+    @property
+    def rank(self) -> int:
+        return self.proc_rank
+
+    @property
+    def device(self) -> torch.device:
+        if self._device is None:
+            raise RuntimeError("worker device is unavailable before init_group()")
+        return self._device
+
+    def init_group(
+        self,
+        *,
+        rank: int | None = None,
+        world_size: int | None = None,
+        device_filters: Sequence[str] | None = None,
+        use_pg_options: bool = True,
+        is_extension: bool = False,
+        active_value: int | None = None,
+    ) -> torch.device:
+        self._device = init_mooncake_group(
+            self.proc_rank if rank is None else rank,
+            self.world_size if world_size is None else world_size,
+            backend_name=self.backend_name,
+            device_type=self.device_type,
+            device_filters=self.device_filters
+            if device_filters is None
+            else device_filters,
+            use_pg_options=use_pg_options,
+            is_extension=is_extension,
+            active_value=active_value,
+        )
+        return self._device
+
+    def get_backend(self, group=None):
+        return get_mooncake_backend(group=group, device_type=self.device_type)
+
+    def record_result(self, payload: dict) -> None:
+        record_rank_result(
+            self.result_map,
+            self.proc_rank,
+            {"ok": True, "rank": self.proc_rank, **payload},
+        )
+
+    def record_error(self, exc: Exception) -> None:
+        record_rank_error(self.result_map, self.proc_rank, exc)
+
+    def synchronize(self) -> None:
+        if self.device_type == "cuda" and self._device is not None:
+            torch.cuda.synchronize(self._device)
+
+
+def destroy_process_groups(*groups) -> None:
+    for group in reversed([item for item in groups if item is not None]):
+        try:
+            dist.destroy_process_group(group)
+        except Exception:
+            pass
+    if dist.is_available() and dist.is_initialized():
+        try:
+            dist.destroy_process_group()
+        except Exception:
+            pass
+
+
+def record_rank_result(result_map, rank: int, payload: dict) -> None:
+    result_map[rank] = payload
+
+
+def record_rank_error(result_map, rank: int, exc: Exception) -> None:
+    record_rank_result(
+        result_map,
+        rank,
+        {
+            "ok": False,
+            "rank": rank,
+            "error_type": type(exc).__name__,
+            "error": str(exc),
+        },
+    )
+
+
+def collect_rank_results(result_map, count: int) -> list[dict]:
+    results = []
+    for rank in range(count):
+        if rank in result_map:
+            results.append(result_map[rank])
+        else:
+            results.append(
+                {
+                    "ok": False,
+                    "rank": rank,
+                    "error_type": "MissingResult",
+                    "error": f"Rank {rank} did not report a result (process may have crashed or timed out)",
+                }
+            )
+    return results
+
+
+def finalize_worker_results(result_map, count: int, *groups) -> None:
+    # Clean up process groups; result collection and timeouts are handled by the parent process.
+    destroy_process_groups(*groups)
+
+
+def _run_worker_with_finalizer(
+    rank: int,
+    world_size: int,
+    result_map,
+    device_filters,
+    worker,
+    args: tuple,
+) -> None:
+    try:
+        worker(rank, world_size, result_map, device_filters, *args)
+    except Exception as exc:  # noqa: BLE001
+        record_rank_error(result_map, rank, exc)
+    finally:
+        finalize_worker_results(result_map, world_size)
+
+
+def _run_backend_worker_with_finalizer(
+    rank: int,
+    world_size: int,
+    result_map,
+    device_filters,
+    result_count: int,
+    backend_name: str,
+    device_type: str,
+    worker,
+    args: tuple,
+) -> None:
+    ctx = MooncakePGWorkerContext(
+        proc_rank=rank,
+        world_size=world_size,
+        result_map=result_map,
+        device_filters=device_filters,
+        backend_name=backend_name,
+        device_type=device_type,
+    )
+    try:
+        worker(ctx, *args)
+    except Exception as exc:  # noqa: BLE001
+        record_rank_error(result_map, rank, exc)
+    finally:
+        finalize_worker_results(result_map, result_count)
+
+
+def wait_until(
+    predicate,
+    *,
+    timeout_s: float = DEFAULT_WAIT_TIMEOUT_S,
+    poll_interval_s: float = 0.05,
+    description: str = "condition",
+):
+    deadline = time.monotonic() + timeout_s
+    while time.monotonic() < deadline:
+        value = predicate()
+        if value:
+            return value
+        time.sleep(poll_interval_s)
+    raise TimeoutError(f"timed out waiting for {description}")
+
+
+def wait_for_spawn_context(ctx, timeout_s: float) -> None:
+    """Wait for spawn context with timeout; force kill if hung."""
+    deadline = time.monotonic() + timeout_s
+
+    # Phase 1: Normal wait
+    while time.monotonic() < deadline:
+        if not any(p.is_alive() for p in ctx.processes):
+            # All processes exited (success or failure)
+            return
+        time.sleep(0.1)
+
+    # Phase 2: Timeout - try graceful termination
+    for process in ctx.processes:
+        if process.is_alive():
+            try:
+                process.terminate()
+            except Exception:
+                pass
+
+    # Phase 3: Wait for terminations with hard timeout
+    term_deadline = time.monotonic() + 5.0
+    while time.monotonic() < term_deadline:
+        if not any(p.is_alive() for p in ctx.processes):
+            break
+        time.sleep(0.1)
+
+    # Phase 4: SIGKILL any survivors
+    for process in ctx.processes:
+        if process.is_alive():
+            try:
+                os.kill(process.pid, signal.SIGKILL)
+            except Exception:
+                pass
+
+    # Phase 5: Final wait (don't block forever)
+    final_deadline = time.monotonic() + 3.0
+    while time.monotonic() < final_deadline:
+        if not any(p.is_alive() for p in ctx.processes):
+            break
+        time.sleep(0.1)
+
+    raise AssertionError(f"Spawn timed out after {timeout_s} seconds")
+
+
+def spawn_and_collect(
+    worker,
+    world_size: int,
+    *args,
+    device_filters: Sequence[str] | None = None,
+    nprocs: int | None = None,
+    timeout_s: float | None = None,
+) -> list[dict]:
+    resolved_filters = resolve_device_filters(device_filters)
+    master_addr = resolve_env_value(MASTER_ADDR_ENV_VAR, DEFAULT_MASTER_ADDR)
+    default_master_port = find_free_local_port()
+    master_port = resolve_env_value(MASTER_PORT_ENV_VAR, default_master_port)
+    actual_nprocs = world_size if nprocs is None else nprocs
+    spawn_ctx = mp.get_context("spawn")
+
+    with spawn_ctx.Manager() as manager:
+        result_map = manager.dict()
+
+        with temporary_env(
+            {
+                "MASTER_ADDR": master_addr,
+                "MASTER_PORT": master_port,
+            }
+        ):
+            spawn_args = (
+                world_size,
+                result_map,
+                resolved_filters,
+                worker,
+                args,
+            )
+            if timeout_s is None:
+                mp.spawn(
+                    _run_worker_with_finalizer,
+                    args=spawn_args,
+                    nprocs=actual_nprocs,
+                    join=True,
+                )
+            else:
+                ctx = mp.spawn(
+                    _run_worker_with_finalizer,
+                    args=spawn_args,
+                    nprocs=actual_nprocs,
+                    join=False,
+                )
+                wait_for_spawn_context(ctx, timeout_s)
+
+        return collect_rank_results(result_map, actual_nprocs)
+
+
+class MultiProcessTestCase(unittest.TestCase):
+    world_size = 4
+    device_filters = None
+    spawn_timeout_s = DEFAULT_SPAWN_TIMEOUT_S
+
+    def spawn_and_collect(
+        self,
+        worker,
+        *args,
+        device_filters=None,
+        nprocs: int | None = None,
+        timeout_s: float | None = None,
+    ) -> list[dict]:
+        resolved_filters = (
+            self.device_filters if device_filters is None else device_filters
+        )
+        return spawn_and_collect(
+            worker,
+            self.world_size,
+            *args,
+            device_filters=resolved_filters,
+            nprocs=nprocs,
+            timeout_s=self.spawn_timeout_s if timeout_s is None else timeout_s,
+        )
+
+    def assert_all_ok(self, rows: list[dict]) -> None:
+        for row in rows:
+            if not row.get("ok", False):
+                self.fail(
+                    f"rank {row.get('rank', '?')} failed with "
+                    f"{row.get('error_type', 'UnknownError')}: {row.get('error', '')}"
+                )
+
+
+class BackendMultiProcessTestCase(MultiProcessTestCase):
+    backend_name: str | None = None
+    device_type: str | None = None
+
+    @classmethod
+    def configure_for_cuda_device_count(cls, device_count: int) -> None:
+        del device_count
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        super().setUpClass()
+        if cls.backend_name is None or cls.device_type is None:
+            raise RuntimeError(
+                f"{cls.__name__} must inherit a concrete Mooncake PG backend test base class"
+            )
+        if cls.device_type == "cuda":
+            device_count = torch.cuda.device_count() if torch.cuda.is_available() else 0
+            cls.configure_for_cuda_device_count(device_count)
+        if cls.device_type == "cuda" and not cuda_runtime_available(cls.world_size):
+            raise unittest.SkipTest(
+                f"{cls.__name__} requires {cls.world_size} visible CUDA devices"
+            )
+
+    def spawn_backend_and_collect(
+        self,
+        worker,
+        *args,
+        device_filters=None,
+        nprocs: int | None = None,
+        timeout_s: float | None = None,
+        world_size: int | None = None,
+    ) -> list[dict]:
+        if self.backend_name is None or self.device_type is None:
+            raise RuntimeError(
+                f"{type(self).__name__} must inherit a concrete Mooncake PG backend test base class"
+            )
+
+        resolved_filters = (
+            self.device_filters if device_filters is None else device_filters
+        )
+        resolved_world_size = (
+            self.world_size if world_size is None else world_size
+        )
+        master_addr = resolve_env_value(MASTER_ADDR_ENV_VAR, DEFAULT_MASTER_ADDR)
+        default_master_port = find_free_local_port()
+        master_port = resolve_env_value(MASTER_PORT_ENV_VAR, default_master_port)
+        actual_nprocs = resolved_world_size if nprocs is None else nprocs
+        spawn_ctx = mp.get_context("spawn")
+
+        with spawn_ctx.Manager() as manager:
+            result_map = manager.dict()
+
+            with temporary_env(
+                {
+                    "MASTER_ADDR": master_addr,
+                    "MASTER_PORT": master_port,
+                }
+            ):
+                spawn_args = (
+                    resolved_world_size,
+                    result_map,
+                    resolved_filters,
+                    actual_nprocs,
+                    self.backend_name,
+                    self.device_type,
+                    worker,
+                    args,
+                )
+                # Always use non-blocking spawn with timeout to avoid hangs
+                resolved_timeout = self.spawn_timeout_s if timeout_s is None else timeout_s
+                ctx = mp.spawn(
+                    _run_backend_worker_with_finalizer,
+                    args=spawn_args,
+                    nprocs=actual_nprocs,
+                    join=False,
+                )
+                wait_for_spawn_context(ctx, resolved_timeout)
+
+            return collect_rank_results(result_map, actual_nprocs)
+
+
+class MooncakePGCPUBackendTestCase(BackendMultiProcessTestCase):
+    backend_name = "mooncake-cpu"
+    device_type = "cpu"
+
+
+class MooncakePGCUDABackendTestCase(BackendMultiProcessTestCase):
+    backend_name = "mooncake"
+    device_type = "cuda"

mooncake-pg/tests/test_pg_collectives.py

@@ -0,0 +1,230 @@
+import unittest
+
+import torch
+import torch.distributed as dist
+
+from pg_test_utils import (
+    MooncakePGCPUBackendTestCase,
+    MooncakePGCUDABackendTestCase,
+    MooncakePGWorkerContext,
+    wait_until,
+)
+
+
+def _collective_payload(
+    ctx: MooncakePGWorkerContext,
+    case_name: str,
+    case_arg: str | None,
+) -> dict:
+    device = ctx.device
+    rank = ctx.rank
+    world_size = ctx.world_size
+    device_type = ctx.device_type
+
+    if case_name == "world_init_without_pg_options":
+        tensor = torch.tensor([rank + 1], dtype=torch.int32, device=device)
+        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+        return {"value": int(tensor.cpu().item())}
+
+    if case_name == "allreduce":
+        if case_arg == "sum":
+            tensor = torch.tensor([rank + 1], dtype=torch.int32, device=device)
+            op = dist.ReduceOp.SUM
+        elif case_arg == "min":
+            tensor = torch.tensor([rank + 10], dtype=torch.int32, device=device)
+            op = dist.ReduceOp.MIN
+        elif case_arg == "max":
+            tensor = torch.tensor([rank + 10], dtype=torch.int32, device=device)
+            op = dist.ReduceOp.MAX
+        elif case_arg == "product":
+            tensor = torch.tensor([2], dtype=torch.int32, device=device)
+            op = dist.ReduceOp.PRODUCT
+        else:
+            raise ValueError(f"unsupported allreduce case_arg: {case_arg}")
+        dist.all_reduce(tensor, op=op)
+        return {"value": int(tensor.cpu().item())}
+
+    if case_name == "broadcast":
+        tensor = torch.tensor([111 if rank == 0 else -1], dtype=torch.int32, device=device)
+        dist.broadcast(tensor, src=0)
+        return {"value": int(tensor.cpu().item())}
+
+    if case_name == "all_gather_into_tensor":
+        local = torch.tensor([rank], dtype=torch.int32, device=device)
+        gathered = torch.empty(world_size, dtype=torch.int32, device=device)
+        dist.all_gather_into_tensor(gathered, local)
+        return {"value": gathered.cpu().tolist()}
+
+    if case_name == "all_gather_list":
+        local = torch.tensor([rank], dtype=torch.int32, device=device)
+        gathered = [torch.empty_like(local) for _ in range(world_size)]
+        dist.all_gather(gathered, local)
+        return {"value": [int(t.cpu().item()) for t in gathered]}
+
+    if case_name == "reduce_scatter_sum":
+        input_buf = torch.arange(
+            rank * world_size,
+            (rank + 1) * world_size,
+            dtype=torch.int32,
+            device=device,
+        )
+        output = torch.empty(1, dtype=torch.int32, device=device)
+        if hasattr(dist, "reduce_scatter_tensor"):
+            dist.reduce_scatter_tensor(output, input_buf, op=dist.ReduceOp.SUM)
+        else:
+            dist.reduce_scatter(output, list(input_buf.chunk(world_size)), op=dist.ReduceOp.SUM)
+        return {"value": output.cpu().tolist()}
+
+    if case_name == "barrier":
+        dist.barrier()
+        return {"value": "ok"}
+
+    if case_name == "gather":
+        tensor = torch.tensor([rank], dtype=torch.int32, device=device)
+        if rank == 0:
+            gather_list = [torch.empty_like(tensor) for _ in range(world_size)]
+            dist.gather(tensor, gather_list, dst=0)
+            return {"value": [int(item.cpu().item()) for item in gather_list]}
+        dist.gather(tensor, dst=0)
+        return {"value": None}
+
+    if case_name == "scatter":
+        tensor = torch.zeros(1, dtype=torch.int32, device=device)
+        if rank == 0:
+            scatter_list = [
+                torch.tensor([peer], dtype=torch.int32, device=device)
+                for peer in range(world_size)
+            ]
+            dist.scatter(tensor, scatter_list, src=0)
+        else:
+            dist.scatter(tensor, src=0)
+        return {"value": int(tensor.cpu().item())}
+
+    if case_name == "reduce":
+        tensor = torch.tensor([1], dtype=torch.int32, device=device)
+        dist.reduce(tensor, dst=0, op=dist.ReduceOp.SUM)
+        return {"value": int(tensor.cpu().item()) if rank == 0 else None}
+
+    if case_name == "async_allreduce":
+        tensor = torch.tensor([rank + 1], dtype=torch.int32, device=device)
+        work = dist.all_reduce(tensor, op=dist.ReduceOp.SUM, async_op=True)
+        work.wait()
+        if device_type == "cuda":
+            torch.cuda.synchronize(device)
+        return {"value": int(tensor.cpu().item())}
+
+    raise ValueError(f"unsupported collective case_name: {case_name}")
+
+
+def _collective_worker(
+    ctx: MooncakePGWorkerContext,
+    case_name: str,
+    case_arg: str | None = None,
+) -> None:
+    ctx.init_group(use_pg_options=case_name != "world_init_without_pg_options")
+    payload = _collective_payload(ctx, case_name, case_arg)
+    ctx.record_result(payload)
+
+
+class _CollectiveTestMixin:
+    def test_world_init_without_pg_options(self) -> None:
+        rows = self.spawn_backend_and_collect(_collective_worker, "world_init_without_pg_options")
+        self.assert_all_ok(rows)
+
+        expected = sum(range(1, self.world_size + 1))
+        for row in rows:
+            self.assertEqual(row["value"], expected)
+
+    def test_allreduce_sum(self) -> None:
+        rows = self.spawn_backend_and_collect(_collective_worker, "allreduce", "sum")
+        self.assert_all_ok(rows)
+        expected = sum(range(1, self.world_size + 1))
+        for row in rows:
+            self.assertEqual(row["value"], expected)
+
+    def test_allreduce_min(self) -> None:
+        rows = self.spawn_backend_and_collect(_collective_worker, "allreduce", "min")
+        self.assert_all_ok(rows)
+        for row in rows:
+            self.assertEqual(row["value"], 10)
+
+    def test_allreduce_max(self) -> None:
+        rows = self.spawn_backend_and_collect(_collective_worker, "allreduce", "max")
+        self.assert_all_ok(rows)
+        expected = 10 + self.world_size - 1
+        for row in rows:
+            self.assertEqual(row["value"], expected)
+
+    def test_allreduce_product(self) -> None:
+        rows = self.spawn_backend_and_collect(_collective_worker, "allreduce", "product")
+        self.assert_all_ok(rows)
+        expected = 2 ** self.world_size
+        for row in rows:
+            self.assertEqual(row["value"], expected)
+
+    def test_broadcast(self) -> None:
+        rows = self.spawn_backend_and_collect(_collective_worker, "broadcast")
+        self.assert_all_ok(rows)
+        for row in rows:
+            self.assertEqual(row["value"], 111)
+
+    def test_all_gather_into_tensor(self) -> None:
+        rows = self.spawn_backend_and_collect(_collective_worker, "all_gather_into_tensor")
+        self.assert_all_ok(rows)
+        expected = list(range(self.world_size))
+        for row in rows:
+            self.assertEqual(row["value"], expected)
+
+    def test_all_gather_list(self) -> None:
+        rows = self.spawn_backend_and_collect(_collective_worker, "all_gather_list")
+        self.assert_all_ok(rows)
+        expected = list(range(self.world_size))
+        for row in rows:
+            self.assertEqual(row["value"], expected)
+
+    def test_reduce_scatter_sum(self) -> None:
+        rows = self.spawn_backend_and_collect(_collective_worker, "reduce_scatter_sum")
+        self.assert_all_ok(rows)
+        for row in rows:
+            rank = row["rank"]
+            expected = self.world_size * (self.world_size * (self.world_size - 1) // 2 + rank)
+            self.assertEqual(row["value"], [expected])
+
+    def test_barrier(self) -> None:
+        rows = self.spawn_backend_and_collect(_collective_worker, "barrier")
+        self.assert_all_ok(rows)
+
+    def test_gather(self) -> None:
+        rows = self.spawn_backend_and_collect(_collective_worker, "gather")
+        self.assert_all_ok(rows)
+        self.assertEqual(rows[0]["value"], list(range(self.world_size)))
+
+    def test_scatter(self) -> None:
+        rows = self.spawn_backend_and_collect(_collective_worker, "scatter")
+        self.assert_all_ok(rows)
+        for row in rows:
+            self.assertEqual(row["value"], row["rank"])
+
+    def test_reduce(self) -> None:
+        rows = self.spawn_backend_and_collect(_collective_worker, "reduce")
+        self.assert_all_ok(rows)
+        self.assertEqual(rows[0]["value"], self.world_size)
+
+    def test_async_allreduce_work_functional(self) -> None:
+        rows = self.spawn_backend_and_collect(_collective_worker, "async_allreduce")
+        self.assert_all_ok(rows)
+        expected = sum(range(1, self.world_size + 1))
+        for row in rows:
+            self.assertEqual(row["value"], expected)
+
+
+class TestMooncakePGCollectivesCPU(_CollectiveTestMixin, MooncakePGCPUBackendTestCase):
+    world_size = 4
+
+
+class TestMooncakePGCollectivesCUDA(_CollectiveTestMixin, MooncakePGCUDABackendTestCase):
+    world_size = 2
+
+
+if __name__ == "__main__":
+    unittest.main()

mooncake-pg/tests/test_pg_elastic.py

@@ -0,0 +1,280 @@
+import os
+import time
+import unittest
+
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+from mooncake import pg
+from pg_test_utils import (
+    MooncakePGCPUBackendTestCase,
+    MooncakePGCUDABackendTestCase,
+    MooncakePGWorkerContext,
+    wait_until,
+)
+
+
+BROKEN_RANK = 1
+
+
+def _extension_worker(
+    ctx: MooncakePGWorkerContext,
+    extend_event: mp.Event,
+    init_done_event: mp.Event,
+) -> None:
+    """Worker for testing extension mode - new ranks join existing group."""
+    initial_world_size = ctx.world_size - 1
+    extension_rank = ctx.world_size - 1
+
+    if ctx.proc_rank < initial_world_size:
+        # Original ranks
+        device = ctx.init_group(world_size=initial_world_size)
+        backend = ctx.get_backend()
+
+        # First collective
+        tensor = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device)
+        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+        baseline = int(tensor.cpu().item())
+
+        # Signal ready and extend
+        if ctx.proc_rank == 0:
+            extend_event.set()
+        pg.extend_group_size_to(backend, ctx.world_size)
+
+        # Wait for extension rank to complete init before collective
+        if not init_done_event.wait(timeout=30.0):
+            raise TimeoutError("timed out waiting for extension init")
+
+        # Final collective
+        final_tensor = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device)
+        dist.all_reduce(final_tensor, op=dist.ReduceOp.SUM)
+
+        ctx.record_result({
+            "role": "original",
+            "rank": ctx.proc_rank,
+            "baseline": baseline,
+        })
+    else:
+        # Extension rank
+        if not extend_event.wait(timeout=30.0):
+            raise TimeoutError("timed out waiting for extend_event")
+
+        device = ctx.init_group(
+            rank=extension_rank,
+            world_size=ctx.world_size,
+        )
+
+        # Signal init complete before collective
+        init_done_event.set()
+
+        # Final collective
+        final_tensor = torch.tensor([extension_rank + 1], dtype=torch.int32, device=device)
+        dist.all_reduce(final_tensor, op=dist.ReduceOp.SUM)
+
+        ctx.record_result({
+            "role": "extension",
+            "rank": extension_rank,
+        })
+
+
+def _fault_detection_worker(
+    ctx: MooncakePGWorkerContext,
+    broken_exited: mp.Event,
+) -> None:
+    """Worker for testing fault detection - survivors can continue without broken rank."""
+    device = ctx.init_group()
+
+    # Step 1: All ranks participate in first collective
+    tensor = torch.tensor([ctx.rank], dtype=torch.int32, device=device)
+    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+
+    if ctx.rank == BROKEN_RANK:
+        # Step 2: Broken rank exits after first collective
+        ctx.record_result({"role": "broken"})
+        broken_exited.set()
+        os._exit(0)
+
+    # Step 3: Survivors wait for broken rank to exit
+    broken_exited.wait()
+
+    # Step 4: Survivors run collective without broken rank
+    # This should not hang - verifies fault detection works
+    tensor = torch.tensor([ctx.rank], dtype=torch.int32, device=device)
+    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+
+    ctx.record_result({"role": "survivor"})
+
+
+def _replacement_recovery_worker(
+    ctx: MooncakePGWorkerContext,
+    broken_exited: mp.Event,
+    replacement_ready: mp.Event,
+    start_recovery: mp.Event,
+) -> None:
+    """Worker for testing replacement recovery."""
+    logical_rank = ctx.rank if ctx.proc_rank < ctx.world_size else BROKEN_RANK
+
+    if ctx.proc_rank < ctx.world_size:
+        # Original rank (0, 1, 2, or 3)
+        device = ctx.init_group(rank=logical_rank)
+
+        # First collective with all ranks
+        tensor = torch.tensor([logical_rank], dtype=torch.int32, device=device)
+        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+
+        if logical_rank == BROKEN_RANK:
+            # Broken rank exits
+            ctx.record_result({"role": "broken"})
+            broken_exited.set()
+            os._exit(0)
+
+        # Survivor ranks
+        broken_exited.wait()
+        backend = ctx.get_backend()
+
+        # Run collective without broken rank
+        tensor = torch.tensor([logical_rank], dtype=torch.int32, device=device)
+        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+
+        # Signal that we're ready for replacement
+        if logical_rank == 0:
+            start_recovery.set()
+
+        # Wait for replacement to be connected (metadata published)
+        # Use longer poll interval to avoid overloading the connection poller
+        wait_until(
+            lambda: pg.get_peer_state(backend, [BROKEN_RANK])[0],
+            timeout_s=30.0,
+            poll_interval_s=2.0,
+            description=f"rank {logical_rank} waiting for replacement to connect",
+        )
+
+        # Wait for replacement to be ready for join_group
+        replacement_ready.wait()
+
+        # All ranks call recover_ranks to include replacement
+        pg.recover_ranks(backend, [BROKEN_RANK])
+
+        # Final collective with all 4 ranks
+        tensor = torch.tensor([logical_rank], dtype=torch.int32, device=device)
+        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+
+        ctx.record_result({"role": "survivor"})
+    else:
+        # Replacement process (proc_rank = world_size)
+        # Wait for signal to start
+        start_recovery.wait()
+
+        # Replacement initializes with is_extension (local-only mode)
+        device = ctx.init_group(rank=logical_rank, is_extension=True)
+        backend = ctx.get_backend()
+
+        # Signal that we're initialized and ready for join_group
+        replacement_ready.set()
+
+        # join_group completes the connection and switches to global mode
+        pg.join_group(backend)
+
+        # Final collective with all 4 ranks
+        tensor = torch.tensor([logical_rank], dtype=torch.int32, device=device)
+        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+
+        ctx.record_result({"role": "replacement"})
+
+
+class _ElasticMixin:
+    world_size = 4
+    spawn_timeout_s = 30.0
+
+    def test_failed_rank(self) -> None:
+        """Test that survivors can continue collective after a rank fails."""
+        spawn_ctx = mp.get_context("spawn")
+        broken_exited = spawn_ctx.Event()
+
+        rows = self.spawn_backend_and_collect(
+            _fault_detection_worker,
+            broken_exited,
+            timeout_s=30.0,
+        )
+
+        # All survivors should complete
+        survivor_rows = [r for r in rows if r.get("role") == "survivor"]
+        self.assertEqual(len(survivor_rows), self.world_size - 1)
+
+        # Broken rank should have exited (may not have result)
+        broken_rows = [r for r in rows if r.get("role") == "broken"]
+        self.assertGreaterEqual(len(broken_rows), 1)
+
+    def test_recovery(self) -> None:
+        """Test that replacement can join and restore full collective."""
+        spawn_ctx = mp.get_context("spawn")
+        broken_exited = spawn_ctx.Event()
+        replacement_ready = spawn_ctx.Event()
+        start_recovery = spawn_ctx.Event()
+
+        rows = self.spawn_backend_and_collect(
+            _replacement_recovery_worker,
+            broken_exited,
+            replacement_ready,
+            start_recovery,
+            nprocs=self.world_size + 1,
+            timeout_s=30.0,
+        )
+
+        # Verify all participants completed
+        survivor_rows = [r for r in rows if r.get("role") == "survivor"]
+        replacement_rows = [r for r in rows if r.get("role") == "replacement"]
+        broken_rows = [r for r in rows if r.get("role") == "broken"]
+
+        self.assertEqual(len(survivor_rows), self.world_size - 1)
+        self.assertEqual(len(replacement_rows), 1)
+        self.assertGreaterEqual(len(broken_rows), 1)
+
+    def test_extension(self) -> None:
+        """Test extension mode allows new ranks to join existing group."""
+        spawn_ctx = mp.get_context("spawn")
+        extend_event = spawn_ctx.Event()
+        init_done_event = spawn_ctx.Event()
+
+        # Spawn world_size processes: (world_size - 1) original + 1 extension
+        rows = self.spawn_backend_and_collect(
+            _extension_worker,
+            extend_event,
+            init_done_event,
+            nprocs=self.world_size,
+            timeout_s=30.0,
+        )
+
+        # Verify all participants completed
+        original_rows = [r for r in rows if r.get("role") == "original"]
+        extension_rows = [r for r in rows if r.get("role") == "extension"]
+
+        # Original: world_size - 1 ranks, Extension: 1 rank
+        self.assertEqual(len(original_rows), self.world_size - 1)
+        self.assertEqual(len(extension_rows), 1)
+
+        # Verify baseline sum: 1+2+...+(world_size-1) = world_size*(world_size-1)/2
+        expected_baseline = (self.world_size - 1) * self.world_size // 2
+        for row in original_rows:
+            self.assertEqual(row.get("baseline"), expected_baseline)
+
+
+class TestMooncakePGElasticCPU(
+    _ElasticMixin, MooncakePGCPUBackendTestCase
+):
+    pass
+
+
+class TestMooncakePGElasticCUDA(
+    _ElasticMixin, MooncakePGCUDABackendTestCase
+):
+    @classmethod
+    def configure_for_cuda_device_count(cls, device_count: int) -> None:
+        if device_count < 2:
+            return
+        cls.world_size = min(device_count, 4)
+
+
+if __name__ == "__main__":
+    unittest.main()

mooncake-pg/tests/test_pg_inference_collectives.py

@@ -0,0 +1,202 @@
+import unittest
+
+import torch
+import torch.distributed as dist
+
+from pg_test_utils import (
+    MooncakePGCPUBackendTestCase,
+    MooncakePGWorkerContext,
+)
+from pg_test_topology import (
+    build_dp_tp_topology,
+    build_many_group_smoke_topology,
+    build_tp_pp_topology,
+    create_named_groups,
+    destroy_named_groups,
+    find_local_named_group,
+)
+
+
+def _local_group_by_prefix(
+    groups: dict[str, object], prefix: str
+) -> tuple[str | None, object | None]:
+    for name, group in groups.items():
+        if name.startswith(prefix):
+            return name, group
+    return None, None
+
+
+def _tp_allreduce_many_groups_worker(ctx: MooncakePGWorkerContext) -> None:
+    groups: dict[str, object] = {}
+    try:
+        topology = build_dp_tp_topology()
+        assert topology.world_size == ctx.world_size
+        device = ctx.init_group()
+        groups = create_named_groups(topology, backend=ctx.backend_name)
+        tp_spec, tp_group = find_local_named_group(topology, groups, ctx.rank, "tp")
+        if tp_spec is None or tp_group is None:
+            raise AssertionError(f"rank {ctx.rank} has no TP group")
+        group_ranks = list(tp_spec.ranks)
+        values = []
+        for iteration in range(3):
+            tensor = torch.tensor([ctx.rank + iteration + 1], dtype=torch.int32, device=device)
+            dist.all_reduce(tensor, group=tp_group, op=dist.ReduceOp.SUM)
+            values.append(int(tensor.item()))
+        ctx.record_result({"values": values, "group_ranks": group_ranks})
+    finally:
+        destroy_named_groups(groups)
+
+
+def _tp_allgather_many_groups_worker(ctx: MooncakePGWorkerContext) -> None:
+    groups: dict[str, object] = {}
+    try:
+        topology = build_dp_tp_topology()
+        assert topology.world_size == ctx.world_size
+        device = ctx.init_group()
+        groups = create_named_groups(topology, backend=ctx.backend_name)
+        tp_spec, tp_group = find_local_named_group(topology, groups, ctx.rank, "tp")
+        if tp_spec is None or tp_group is None:
+            raise AssertionError(f"rank {ctx.rank} has no TP group")
+        group_ranks = list(tp_spec.ranks)
+        values = []
+        for iteration in range(3):
+            local = torch.tensor([ctx.rank * 10 + iteration], dtype=torch.int32, device=device)
+            gathered = torch.empty(len(group_ranks), dtype=torch.int32, device=device)
+            dist.all_gather_into_tensor(gathered, local, group=tp_group)
+            values.append(gathered.tolist())
+        ctx.record_result({"values": values, "group_ranks": group_ranks})
+    finally:
+        destroy_named_groups(groups)
+
+
+def _pp_send_recv_smoke_worker(ctx: MooncakePGWorkerContext) -> None:
+    groups: dict[str, object] = {}
+    try:
+        topology = build_tp_pp_topology()
+        assert topology.world_size == ctx.world_size
+        device = ctx.init_group()
+        groups = create_named_groups(topology, backend=ctx.backend_name)
+
+        lane_spec, lane_group = find_local_named_group(topology, groups, ctx.rank, "pp")
+        if lane_group is None or lane_spec is None:
+            raise AssertionError(f"rank {ctx.rank} has no PP lane group")
+
+        if len(lane_spec.ranks) != 2:
+            raise AssertionError(f"PP lane {lane_spec.name} is not 2-rank: {lane_spec.ranks}")
+        local_index = lane_spec.ranks.index(ctx.rank)
+        peer_group_rank = 1 - local_index
+        src, dst = lane_spec.ranks
+        direct_value = None
+        batch_value = None
+        if local_index == 0:
+            direct = torch.tensor([src * 100 + dst], dtype=torch.int32, device=device)
+            dist.send(direct, group=lane_group, group_dst=peer_group_rank)
+            batch = torch.tensor([src * 1000 + dst], dtype=torch.int32, device=device)
+            requests = dist.batch_isend_irecv(
+                [dist.P2POp(dist.isend, batch, group=lane_group, group_peer=peer_group_rank)]
+            )
+            for request in requests:
+                request.wait()
+            direct_value = int(direct.item())
+            batch_value = int(batch.item())
+        else:
+            direct = torch.empty(1, dtype=torch.int32, device=device)
+            dist.recv(direct, group=lane_group, group_src=peer_group_rank)
+            batch = torch.empty(1, dtype=torch.int32, device=device)
+            requests = dist.batch_isend_irecv(
+                [dist.P2POp(dist.irecv, batch, group=lane_group, group_peer=peer_group_rank)]
+            )
+            for request in requests:
+                request.wait()
+            direct_value = int(direct.item())
+            batch_value = int(batch.item())
+
+        ctx.record_result({"direct": direct_value, "batch": batch_value})
+    finally:
+        destroy_named_groups(groups)
+
+
+def _overlapping_group_collective_worker(ctx: MooncakePGWorkerContext) -> None:
+    groups: dict[str, object] = {}
+    try:
+        topology = build_many_group_smoke_topology()
+        assert topology.world_size == ctx.world_size
+        device = ctx.init_group()
+        groups = create_named_groups(topology, backend=ctx.backend_name)
+
+        _, tp_group = _local_group_by_prefix(groups, "tp")
+        dp_name, dp_group = _local_group_by_prefix(groups, "dp")
+        ep_name, _ = _local_group_by_prefix(groups, "ep")
+
+        tp_tensor = torch.tensor([ctx.rank + 1], dtype=torch.int32, device=device)
+        dist.all_reduce(tp_tensor, group=tp_group, op=dist.ReduceOp.SUM)
+
+        dp_root = int(dp_name.removeprefix("dp"))
+        dp_tensor = torch.tensor(
+            [dp_root * 100 + 5 if ctx.rank == dp_root else -1], dtype=torch.int32, device=device
+        )
+        dist.broadcast(dp_tensor, src=dp_root, group=dp_group)
+
+        ctx.record_result({
+            "tp_value": int(tp_tensor.item()),
+            "dp_value": int(dp_tensor.item()),
+            "ep_group": list(next(spec.ranks for spec in topology.ordered_groups if spec.name == ep_name)),
+        })
+    finally:
+        destroy_named_groups(groups)
+
+
+class _InferenceCollectivesMixin:
+    world_size = 8
+
+    def test_tp_allreduce_many_groups(self) -> None:
+        rows = self.spawn_backend_and_collect(_tp_allreduce_many_groups_worker)
+        self.assert_all_ok(rows)
+
+        for row in rows:
+            ranks = row["group_ranks"]
+            expected = [sum(peer + iteration + 1 for peer in ranks) for iteration in range(3)]
+            self.assertEqual(row["values"], expected)
+
+    def test_tp_allgather_many_groups(self) -> None:
+        rows = self.spawn_backend_and_collect(_tp_allgather_many_groups_worker)
+        self.assert_all_ok(rows)
+
+        for row in rows:
+            ranks = row["group_ranks"]
+            expected = [[peer * 10 + iteration for peer in ranks] for iteration in range(3)]
+            self.assertEqual(row["values"], expected)
+
+    def test_overlapping_group_collective_traffic(self) -> None:
+        rows = self.spawn_backend_and_collect(_overlapping_group_collective_worker)
+        self.assert_all_ok(rows)
+
+
+class _InferenceP2PSmokeMixin:
+    world_size = 4
+
+    def test_pp_send_recv_smoke(self) -> None:
+        rows = self.spawn_backend_and_collect(_pp_send_recv_smoke_worker)
+        self.assert_all_ok(rows)
+
+        expected_direct = {0: 2, 1: 103, 2: 2, 3: 103}
+        expected_batch = {0: 2, 1: 1003, 2: 2, 3: 1003}
+        for row in rows:
+            self.assertEqual(row["direct"], expected_direct[row["rank"]])
+            self.assertEqual(row["batch"], expected_batch[row["rank"]])
+
+
+class TestMooncakePGInferenceCollectivesCPU(
+    _InferenceCollectivesMixin, MooncakePGCPUBackendTestCase
+):
+    pass
+
+
+class TestMooncakePGInferenceP2PSmokeCPU(
+    _InferenceP2PSmokeMixin, MooncakePGCPUBackendTestCase
+):
+    pass
+
+
+if __name__ == "__main__":
+    unittest.main()

mooncake-pg/tests/test_pg_inference_topologies.py

@@ -0,0 +1,213 @@
+import unittest
+
+import torch
+import torch.distributed as dist
+
+from pg_test_utils import (
+    MooncakePGCPUBackendTestCase,
+    MooncakePGWorkerContext,
+)
+from pg_test_topology import (
+    build_dp_tp_ep_topology,
+    build_dp_tp_topology,
+    build_prefill_decode_topology,
+    build_tp_only_topology,
+    build_tp_pp_topology,
+    create_named_groups,
+    destroy_named_groups,
+    find_local_named_group,
+)
+
+TOPOLOGY_BUILDERS = {
+    "tp_only": build_tp_only_topology,
+    "dp_tp": build_dp_tp_topology,
+    "tp_pp": build_tp_pp_topology,
+    "dp_tp_ep": build_dp_tp_ep_topology,
+    "prefill_decode": build_prefill_decode_topology,
+}
+
+
+def _run_lane_send_recv(
+    topology,
+    groups: dict[str, object],
+    rank: int,
+    device: torch.device,
+) -> dict[str, int | None]:
+    lane_spec, lane_group = find_local_named_group(
+        topology, groups, rank, ("pp", "lane")
+    )
+    if lane_group is None or lane_spec is None:
+        return {"lane_sent": None, "lane_recv": None}
+
+    if len(lane_spec.ranks) != 2:
+        raise AssertionError(f"lane {lane_spec.name} is not 2-rank: {lane_spec.ranks}")
+    local_index = lane_spec.ranks.index(rank)
+    peer_group_rank = 1 - local_index
+    src, dst = lane_spec.ranks
+    if local_index == 0:
+        value = src * 100 + dst
+        tensor = torch.tensor([value], dtype=torch.int32, device=device)
+        dist.send(tensor, group=lane_group, group_dst=peer_group_rank)
+        return {"lane_sent": value, "lane_recv": None}
+
+    tensor = torch.empty(1, dtype=torch.int32, device=device)
+    dist.recv(tensor, group=lane_group, group_src=peer_group_rank)
+    return {"lane_sent": None, "lane_recv": int(tensor.item())}
+
+
+def _topology_worker(
+    ctx: MooncakePGWorkerContext,
+    topology_name: str,
+) -> None:
+    groups: dict[str, object] = {}
+    try:
+        topology = TOPOLOGY_BUILDERS[topology_name]()
+        if topology.world_size != ctx.world_size:
+            raise AssertionError(
+                f"{topology_name} expects world_size={topology.world_size}, got {ctx.world_size}"
+            )
+
+        device = ctx.init_group()
+        groups = create_named_groups(topology, backend=ctx.backend_name)
+        membership = topology.membership_for_rank(ctx.rank)
+        payload: dict[str, object] = {
+            "membership": membership,
+        }
+
+        tp_spec, tp_group = find_local_named_group(topology, groups, ctx.rank, "tp")
+        if tp_group is not None:
+            tp_tensor = torch.tensor([ctx.rank + 1], dtype=torch.int32, device=device)
+            dist.all_reduce(tp_tensor, group=tp_group, op=dist.ReduceOp.SUM)
+            payload["tp_group"] = list(tp_spec.ranks)
+            payload["tp_value"] = int(tp_tensor.item())
+
+        dp_spec, dp_group = find_local_named_group(topology, groups, ctx.rank, "dp")
+        if dp_group is not None:
+            root_rank = min(dp_spec.ranks)
+            dp_tensor = torch.tensor(
+                [root_rank * 100 + 7 if ctx.rank == root_rank else -1],
+                dtype=torch.int32,
+                device=device,
+            )
+            dist.broadcast(dp_tensor, src=root_rank, group=dp_group)
+            payload["dp_group"] = list(dp_spec.ranks)
+            payload["dp_value"] = int(dp_tensor.item())
+
+        ep_spec, ep_group = find_local_named_group(topology, groups, ctx.rank, "ep")
+        if ep_group is not None:
+            payload["ep_group"] = list(ep_spec.ranks)
+
+        service_spec, service_group = find_local_named_group(
+            topology, groups, ctx.rank, "service"
+        )
+        if service_group is not None:
+            service_tensor = torch.tensor(
+                [ctx.rank + 1], dtype=torch.int32, device=device
+            )
+            dist.all_reduce(
+                service_tensor, group=service_group, op=dist.ReduceOp.SUM
+            )
+            payload["service_group"] = list(service_spec.ranks)
+            payload["service_value"] = int(service_tensor.item())
+
+        payload.update(_run_lane_send_recv(topology, groups, ctx.rank, device))
+        ctx.record_result(payload)
+    finally:
+        destroy_named_groups(groups)
+
+
+class _InferenceTopologiesMixin:
+    world_size = 4
+
+    def test_tp_only_topology(self) -> None:
+        rows = self.spawn_backend_and_collect(_topology_worker, "tp_only")
+        self.assert_all_ok(rows)
+
+        expected = sum(range(1, self.world_size + 1))
+        for row in rows:
+            self.assertEqual(row["tp_value"], expected)
+
+    def test_tp_pp_topology(self) -> None:
+        rows = self.spawn_backend_and_collect(_topology_worker, "tp_pp")
+        self.assert_all_ok(rows)
+
+        expected_tp = {
+            0: 3,
+            1: 3,
+            2: 7,
+            3: 7,
+        }
+        expected_lane_recv = {2: 2, 3: 103}
+        for row in rows:
+            self.assertEqual(row["tp_value"], expected_tp[row["rank"]])
+            if row["rank"] in expected_lane_recv:
+                self.assertEqual(row["lane_recv"], expected_lane_recv[row["rank"]])
+
+    def test_prefill_decode_topology(self) -> None:
+        rows = self.spawn_backend_and_collect(_topology_worker, "prefill_decode")
+        self.assert_all_ok(rows)
+
+        expected_service = {0: 3, 1: 3, 2: 7, 3: 7}
+        expected_lane_recv = {2: 2, 3: 103}
+        for row in rows:
+            self.assertEqual(row["service_value"], expected_service[row["rank"]])
+            if row["rank"] in expected_lane_recv:
+                self.assertEqual(row["lane_recv"], expected_lane_recv[row["rank"]])
+
+
+class _InferenceTopologiesLargeMixin:
+    world_size = 8
+
+    def test_dp_tp_topology(self) -> None:
+        rows = self.spawn_backend_and_collect(_topology_worker, "dp_tp")
+        self.assert_all_ok(rows)
+
+        expected_tp = {0: 10, 1: 10, 2: 10, 3: 10, 4: 26, 5: 26, 6: 26, 7: 26}
+        expected_dp = {
+            0: 7,
+            4: 7,
+            1: 107,
+            5: 107,
+            2: 207,
+            6: 207,
+            3: 307,
+            7: 307,
+        }
+        for row in rows:
+            self.assertEqual(row["tp_value"], expected_tp[row["rank"]])
+            self.assertEqual(row["dp_value"], expected_dp[row["rank"]])
+
+    def test_dp_tp_ep_topology(self) -> None:
+        rows = self.spawn_backend_and_collect(_topology_worker, "dp_tp_ep")
+        self.assert_all_ok(rows)
+
+        expected_tp = {0: 4, 2: 4, 1: 6, 3: 6, 4: 12, 6: 12, 5: 14, 7: 14}
+        expected_dp = {
+            0: 7,
+            4: 7,
+            1: 107,
+            5: 107,
+            2: 207,
+            6: 207,
+            3: 307,
+            7: 307,
+        }
+        for row in rows:
+            self.assertEqual(row["tp_value"], expected_tp[row["rank"]])
+            self.assertEqual(row["dp_value"], expected_dp[row["rank"]])
+
+
+class TestMooncakePGInferenceTopologiesCPU(
+    _InferenceTopologiesMixin, MooncakePGCPUBackendTestCase
+):
+    pass
+
+
+class TestMooncakePGInferenceTopologiesLargeCPU(
+    _InferenceTopologiesLargeMixin, MooncakePGCPUBackendTestCase
+):
+    pass
+
+
+if __name__ == "__main__":
+    unittest.main()

mooncake-pg/tests/test_pg_init_functional.py

@@ -0,0 +1,174 @@
+import unittest
+
+import torch
+import torch.distributed as dist
+
+from pg_test_utils import (
+    MooncakePGCPUBackendTestCase,
+    MooncakePGCUDABackendTestCase,
+    MooncakePGWorkerContext,
+)
+
+
+def _basic_init_worker(ctx: MooncakePGWorkerContext) -> None:
+    """Test basic init works and rank/world_size are correct."""
+    device = ctx.init_group()
+    # Verify rank and world_size are accessible
+    assert dist.get_rank() == ctx.rank, f"rank mismatch: {dist.get_rank()} != {ctx.rank}"
+    assert dist.get_world_size() == ctx.world_size, f"world_size mismatch"
+    # Simple collective to verify group works
+    tensor = torch.tensor([ctx.rank + 1], dtype=torch.int32, device=device)
+    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+    ctx.record_result({"sum": int(tensor.cpu().item())})
+
+
+def _null_init_worker(ctx: MooncakePGWorkerContext) -> None:
+    """Test init without pg_options (using defaults)."""
+    device = ctx.init_group(use_pg_options=False)
+    tensor = torch.tensor([ctx.rank + 1], dtype=torch.int32, device=device)
+    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+    ctx.record_result({"sum": int(tensor.cpu().item())})
+
+
+def _subgroup_create_destroy_worker(ctx: MooncakePGWorkerContext) -> None:
+    """Test subgroup creation and destruction."""
+    device = ctx.init_group()
+    world_size = ctx.world_size
+    rank = ctx.rank
+
+    # Create two subgroups: even ranks and odd ranks
+    # Note: new_group is collective, all ranks must call it
+    even_ranks = list(range(0, world_size, 2))
+    odd_ranks = list(range(1, world_size, 2))
+
+    even_group = None
+    odd_group = None
+    subgroup_sum = None
+
+    try:
+        # All ranks collectively create both groups
+        even_group = dist.new_group(ranks=even_ranks, backend=ctx.backend_name)
+        odd_group = dist.new_group(ranks=odd_ranks, backend=ctx.backend_name)
+
+        # Each rank uses its respective group
+        if rank in even_ranks and even_group is not None:
+            tensor = torch.tensor([rank], dtype=torch.int32, device=device)
+            dist.all_reduce(tensor, group=even_group, op=dist.ReduceOp.SUM)
+            subgroup_sum = int(tensor.cpu().item())
+        elif rank in odd_ranks and odd_group is not None:
+            tensor = torch.tensor([rank], dtype=torch.int32, device=device)
+            dist.all_reduce(tensor, group=odd_group, op=dist.ReduceOp.SUM)
+            subgroup_sum = int(tensor.cpu().item())
+
+        # Verify world group still works after subgroup operations
+        world_tensor = torch.tensor([1], dtype=torch.int32, device=device)
+        dist.all_reduce(world_tensor, op=dist.ReduceOp.SUM)
+
+        ctx.record_result({
+            "subgroup_sum": subgroup_sum,
+            "world_sum": int(world_tensor.cpu().item()),
+        })
+    finally:
+        for g in (even_group, odd_group):
+            if g is not None:
+                try:
+                    dist.destroy_process_group(g)
+                except Exception:
+                    pass
+
+
+def _destroy_and_reinit_worker(ctx: MooncakePGWorkerContext) -> None:
+    """Test destroy and re-init process group."""
+    # First init
+    device = ctx.init_group()
+    tensor1 = torch.tensor([ctx.rank + 1], dtype=torch.int32, device=device)
+    dist.all_reduce(tensor1, op=dist.ReduceOp.SUM)
+    sum1 = int(tensor1.cpu().item())
+
+    # Destroy WORLD group
+    try:
+        dist.destroy_process_group(dist.group.WORLD)
+    except Exception:
+        pass
+
+    # Re-init
+    device = ctx.init_group()
+    tensor2 = torch.tensor([ctx.rank + 1], dtype=torch.int32, device=device)
+    dist.all_reduce(tensor2, op=dist.ReduceOp.SUM)
+    sum2 = int(tensor2.cpu().item())
+
+    ctx.record_result({"sum1": sum1, "sum2": sum2})
+
+
+class _InitFunctionalMixin:
+    def test_basic_init(self) -> None:
+        """Test basic init works and rank/world_size are correct."""
+        rows = self.spawn_backend_and_collect(_basic_init_worker)
+        self.assert_all_ok(rows)
+        expected_sum = sum(range(1, self.world_size + 1))
+        for row in rows:
+            self.assertEqual(row["sum"], expected_sum)
+
+    def test_null_init(self) -> None:
+        """Test init without pg_options (using defaults)."""
+        rows = self.spawn_backend_and_collect(_null_init_worker)
+        self.assert_all_ok(rows)
+        expected_sum = sum(range(1, self.world_size + 1))
+        for row in rows:
+            self.assertEqual(row["sum"], expected_sum)
+
+    def test_subgroup_create_destroy(self) -> None:
+        """Test subgroup creation and destruction."""
+        if self.world_size < 4:
+            self.skipTest("subgroup test requires at least 4 ranks")
+
+        rows = self.spawn_backend_and_collect(_subgroup_create_destroy_worker)
+        self.assert_all_ok(rows)
+
+        # Verify world sum is correct
+        expected_world = self.world_size
+        for row in rows:
+            self.assertEqual(row["world_sum"], expected_world)
+
+        # Verify subgroup sums
+        even_ranks = list(range(0, self.world_size, 2))
+        odd_ranks = list(range(1, self.world_size, 2))
+        expected_even = sum(even_ranks)
+        expected_odd = sum(odd_ranks)
+
+        for row in rows:
+            if row["rank"] in even_ranks:
+                self.assertEqual(row["subgroup_sum"], expected_even)
+            else:
+                self.assertEqual(row["subgroup_sum"], expected_odd)
+
+    def test_destroy_and_reinit(self) -> None:
+        """Test destroy and re-init process group."""
+        rows = self.spawn_backend_and_collect(_destroy_and_reinit_worker)
+        self.assert_all_ok(rows)
+
+        expected = sum(range(1, self.world_size + 1))
+        for row in rows:
+            self.assertEqual(row["sum1"], expected)
+            self.assertEqual(row["sum2"], expected)
+
+
+class TestMooncakePGInitFunctionalCPU(_InitFunctionalMixin, MooncakePGCPUBackendTestCase):
+    world_size = 4
+
+
+class TestMooncakePGInitFunctionalCUDA(_InitFunctionalMixin, MooncakePGCUDABackendTestCase):
+    world_size = 4
+
+    @classmethod
+    def configure_for_cuda_device_count(cls, device_count: int) -> None:
+        if device_count < 2:
+            return
+        if device_count >= 4:
+            cls.world_size = 4
+            return
+        cls.world_size = 2
+
+
+if __name__ == "__main__":
+    unittest.main()

mooncake-pg/tests/test_pg_p2p.py

@@ -0,0 +1,182 @@
+import unittest
+
+import torch
+import torch.distributed as dist
+
+from pg_test_utils import (
+    MooncakePGCPUBackendTestCase,
+    MooncakePGCUDABackendTestCase,
+    MooncakePGWorkerContext,
+)
+
+
+def _ring_send_recv_worker(
+    ctx: MooncakePGWorkerContext,
+) -> None:
+    device = ctx.init_group()
+    send_tensor = torch.tensor([ctx.rank], dtype=torch.int64, device=device)
+    recv_tensor = torch.empty_like(send_tensor)
+    dst = (ctx.rank + 1) % ctx.world_size
+    src = (ctx.rank - 1 + ctx.world_size) % ctx.world_size
+    ops = [
+        dist.P2POp(op=dist.isend, tensor=send_tensor, peer=dst),
+        dist.P2POp(op=dist.irecv, tensor=recv_tensor, peer=src),
+    ]
+    works = dist.batch_isend_irecv(ops)
+    for work in works:
+        work.wait()
+    ctx.synchronize()
+    ctx.record_result({"value": int(recv_tensor.cpu().item())})
+
+
+def _direct_and_batch_send_recv_worker(
+    ctx: MooncakePGWorkerContext,
+) -> None:
+    if ctx.world_size != 2:
+        raise AssertionError("direct send/recv smoke expects world_size=2")
+    device = ctx.init_group()
+    peer = 1 - ctx.rank
+    if ctx.rank == 0:
+        direct = torch.tensor([12], dtype=torch.int32, device=device)
+        dist.send(direct, dst=peer)
+        batch = torch.tensor([1200], dtype=torch.int32, device=device)
+        requests = dist.batch_isend_irecv([dist.P2POp(dist.isend, batch, peer=peer)])
+        for request in requests:
+            request.wait()
+        value = {"direct": int(direct.cpu().item()), "batch": int(batch.cpu().item())}
+    else:
+        direct = torch.empty(1, dtype=torch.int32, device=device)
+        dist.recv(direct, src=peer)
+        batch = torch.empty(1, dtype=torch.int32, device=device)
+        requests = dist.batch_isend_irecv([dist.P2POp(dist.irecv, batch, peer=peer)])
+        for request in requests:
+            request.wait()
+        value = {"direct": int(direct.cpu().item()), "batch": int(batch.cpu().item())}
+    ctx.synchronize()
+    ctx.record_result(value)
+
+
+def _ordering_worker(
+    ctx: MooncakePGWorkerContext,
+) -> None:
+    device = ctx.init_group()
+    if ctx.rank >= 2:
+        ctx.record_result({"value": "skip"})
+        return
+
+    num_msgs = 4
+    if ctx.rank == 0:
+        send_tensors = [torch.tensor([i], dtype=torch.int64, device=device) for i in range(num_msgs)]
+        ops = [dist.P2POp(op=dist.isend, tensor=t, peer=1) for t in send_tensors]
+        works = dist.batch_isend_irecv(ops)
+        for work in works:
+            work.wait()
+        value = "ok"
+    else:
+        recv_tensors = [torch.empty(1, dtype=torch.int64, device=device) for _ in range(num_msgs)]
+        ops = [dist.P2POp(op=dist.irecv, tensor=t, peer=0) for t in recv_tensors]
+        works = dist.batch_isend_irecv(ops)
+        for work in works:
+            work.wait()
+        value = [int(t.cpu().item()) for t in recv_tensors]
+    ctx.synchronize()
+    ctx.record_result({"value": value})
+
+
+def _multiple_senders_worker(
+    ctx: MooncakePGWorkerContext,
+) -> None:
+    device = ctx.init_group()
+    if ctx.rank == 0:
+        send_tensor = torch.tensor([100], dtype=torch.int64, device=device)
+        recv_tensor = torch.empty(1, dtype=torch.int64, device=device)
+        ops = [
+            dist.P2POp(op=dist.isend, tensor=send_tensor, peer=1),
+            dist.P2POp(op=dist.irecv, tensor=recv_tensor, peer=1),
+        ]
+        works = dist.batch_isend_irecv(ops)
+        for work in works:
+            work.wait()
+        value = int(recv_tensor.cpu().item())
+    elif ctx.rank == 1:
+        recv_from_0 = torch.empty(1, dtype=torch.int64, device=device)
+        recv_from_2 = torch.empty(1, dtype=torch.int64, device=device)
+        send_tensor = torch.tensor([101], dtype=torch.int64, device=device)
+        ops = [
+            dist.P2POp(op=dist.irecv, tensor=recv_from_0, peer=0),
+            dist.P2POp(op=dist.irecv, tensor=recv_from_2, peer=2),
+            dist.P2POp(op=dist.isend, tensor=send_tensor, peer=0),
+        ]
+        works = dist.batch_isend_irecv(ops)
+        for work in works:
+            work.wait()
+        value = [int(recv_from_0.cpu().item()), int(recv_from_2.cpu().item())]
+    elif ctx.rank == 2:
+        send_tensor = torch.tensor([200], dtype=torch.int64, device=device)
+        works = dist.batch_isend_irecv([dist.P2POp(op=dist.isend, tensor=send_tensor, peer=1)])
+        for work in works:
+            work.wait()
+        value = "ok"
+    else:
+        value = "skip"
+    ctx.synchronize()
+    ctx.record_result({"value": value})
+
+
+class _P2PMixin:
+    world_size = 4
+
+    def test_ring_send_recv(self) -> None:
+        rows = self.spawn_backend_and_collect(_ring_send_recv_worker)
+        self.assert_all_ok(rows)
+        for row in rows:
+            expected = (row["rank"] - 1 + self.world_size) % self.world_size
+            self.assertEqual(row["value"], expected)
+
+    def test_direct_and_batch_send_recv(self) -> None:
+        rows = self.spawn_backend_and_collect(
+            _direct_and_batch_send_recv_worker,
+            world_size=2,
+            nprocs=2,
+        )
+        self.assert_all_ok(rows)
+        for row in rows:
+            self.assertEqual(row["direct"], 12)
+            self.assertEqual(row["batch"], 1200)
+
+    def test_ordering_between_two_ranks(self) -> None:
+        rows = self.spawn_backend_and_collect(_ordering_worker)
+        self.assert_all_ok(rows)
+        rank1 = next(row for row in rows if row["rank"] == 1)
+        self.assertEqual(rank1["value"], list(range(4)))
+
+    def test_multiple_senders_to_same_receiver(self) -> None:
+        if self.world_size < 3:
+            self.skipTest("multiple-sender P2P coverage requires at least 3 ranks")
+        rows = self.spawn_backend_and_collect(_multiple_senders_worker)
+        self.assert_all_ok(rows)
+        rank0 = next(row for row in rows if row["rank"] == 0)
+        rank1 = next(row for row in rows if row["rank"] == 1)
+        rank2 = next(row for row in rows if row["rank"] == 2)
+        self.assertEqual(rank0["value"], 101)
+        self.assertIn(100, rank1["value"])
+        self.assertIn(200, rank1["value"])
+        self.assertEqual(len(rank1["value"]), 2)
+        self.assertEqual(rank2["value"], "ok")
+
+
+class TestMooncakePGP2PCPU(_P2PMixin, MooncakePGCPUBackendTestCase):
+    pass
+
+
+class TestMooncakePGP2PCUDA(_P2PMixin, MooncakePGCUDABackendTestCase):
+
+    @classmethod
+    def configure_for_cuda_device_count(cls, device_count: int) -> None:
+        if device_count < 2:
+            return
+        cls.world_size = min(device_count, 4)
+
+
+if __name__ == "__main__":
+    unittest.main()

mooncake-store/go/build.sh

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

mooncake-store/go/mooncakestore/config.go

@@ -18,6 +18,7 @@ package mooncakestore
 type ReplicateConfig struct {
 	ReplicaNum        int
 	WithSoftPin       bool
+	WithHardPin       bool
 	PreferredSegments []string
 }
 
@@ -26,5 +27,6 @@ func DefaultReplicateConfig() ReplicateConfig {
 	return ReplicateConfig{
 		ReplicaNum:  1,
 		WithSoftPin: false,
+		WithHardPin: false,
 	}
 }

mooncake-store/go/mooncakestore/store.go

@@ -140,6 +140,9 @@ func (s *Store) toCConfig(cfg *ReplicateConfig) (C.mooncake_replicate_config_t,
 	if cfg.WithSoftPin {
 		cc.with_soft_pin = 1
 	}
+	if cfg.WithHardPin {
+		cc.with_hard_pin = 1
+	}
 
 	if len(cfg.PreferredSegments) > 0 {
 		cSegs = make([]*C.char, len(cfg.PreferredSegments))

mooncake-store/include/client_metric.h

@@ -1,8 +1,15 @@
 #pragma once
 
+#include <algorithm>
+#include <array>
 #include <atomic>
+#include <chrono>
+#include <iomanip>
+#include <mutex>
+#include <optional>
 #include <sstream>
 #include <thread>
+#include <unordered_set>
 #include <vector>
 #include <ylt/metric/counter.hpp>
 #include <ylt/metric/histogram.hpp>
@@ -44,6 +51,56 @@ const inline std::map<std::string, std::string> merge_labels(
     return merged_labels;
 }
 
+inline std::string format_metric_rate(double value, const char* suffix) {
+    const double KB = 1024.0;
+    const double MB = KB * 1024.0;
+    const double GB = MB * 1024.0;
+    const double TB = GB * 1024.0;
+
+    std::ostringstream oss;
+    oss << std::fixed << std::setprecision(2);
+    if (value >= TB) {
+        oss << value / TB << " T" << suffix;
+    } else if (value >= GB) {
+        oss << value / GB << " G" << suffix;
+    } else if (value >= MB) {
+        oss << value / MB << " M" << suffix;
+    } else if (value >= KB) {
+        oss << value / KB << " K" << suffix;
+    } else {
+        oss << value << " " << suffix;
+    }
+    return oss.str();
+}
+
+inline std::string format_metric_bandwidth(uint64_t total_bytes,
+                                           double elapsed_seconds) {
+    return format_metric_rate(total_bytes / elapsed_seconds, "B/s");
+}
+
+inline uint64_t elapsed_us_since(
+    std::chrono::steady_clock::time_point start_time) {
+    return static_cast<uint64_t>(
+        std::chrono::duration_cast<std::chrono::microseconds>(
+            std::chrono::steady_clock::now() - start_time)
+            .count());
+}
+
+template <typename Result, typename Operation, typename SuccessFn,
+          typename ObserveFn>
+Result execute_timed_operation(Operation&& operation, SuccessFn&& success_fn,
+                               ObserveFn&& observe_fn) {
+    const auto start_time = std::chrono::steady_clock::now();
+    Result result = std::forward<Operation>(operation)();
+    if (std::forward<SuccessFn>(success_fn)(result)) {
+        std::forward<ObserveFn>(observe_fn)(elapsed_us_since(start_time),
+                                            result);
+    }
+    return result;
+}
+
+enum class TransferOperationKind { kRead, kWrite };
+
 struct TransferMetric {
     TransferMetric(std::map<std::string, std::string> labels = {})
         : total_read_bytes("mooncake_transfer_read_bytes", "Total bytes read",
@@ -59,7 +116,8 @@ struct TransferMetric {
           get_latency_us("mooncake_transfer_get_latency",
                          "Get transfer latency (us)", kLatencyBucket, labels),
           put_latency_us("mooncake_transfer_put_latency",
-                         "Put transfer latency (us)", kLatencyBucket, labels) {}
+                         "Put transfer latency (us)", kLatencyBucket, labels),
+          start_time_(std::chrono::steady_clock::now()) {}
 
     ylt::metric::counter_t total_read_bytes;
     ylt::metric::counter_t total_write_bytes;
@@ -77,7 +135,7 @@ struct TransferMetric {
         put_latency_us.serialize(str);
     }
 
-    std::string summary_metrics() {
+    std::string summary_metrics(bool include_bandwidth = true) {
         std::stringstream ss;
         ss << "=== Transfer Metrics Summary ===\n";
 
@@ -86,6 +144,14 @@ struct TransferMetric {
         auto write_bytes = total_write_bytes.value();
         ss << "Total Read: " << byte_size_to_string(read_bytes) << "\n";
         ss << "Total Write: " << byte_size_to_string(write_bytes) << "\n";
+        if (include_bandwidth) {
+            ss << "Average Read Throughput: "
+               << format_metric_bandwidth(read_bytes, elapsed_seconds())
+               << "\n";
+            ss << "Average Write Throughput: "
+               << format_metric_bandwidth(write_bytes, elapsed_seconds())
+               << "\n";
+        }
 
         // Latency summaries
         ss << "\n=== Latency Summary (microseconds) ===\n";
@@ -100,6 +166,14 @@ struct TransferMetric {
     }
 
    private:
+    std::chrono::steady_clock::time_point start_time_;
+
+    double elapsed_seconds() const {
+        const auto elapsed = std::chrono::duration<double>(
+            std::chrono::steady_clock::now() - start_time_);
+        return std::max(elapsed.count(), 1e-9);
+    }
+
     std::string format_latency_summary(ylt::metric::histogram_t& hist) {
         // Access the internal sum and bucket counts
         auto sum_ptr =
@@ -272,9 +346,322 @@ 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 = {
+    50,       100,     200,                      // <200us (high-end NVMe)
+    500,      1000,    2000,    5000,    10000,  // 500us - 10ms
+    20000,    50000,   100000,  200000,          // 10ms - 200ms
+    500000,   1000000, 2000000, 5000000,         // 500ms - 5s
+    10000000, 30000000                           // 10s - 30s (3fs/nfs)
+};
+
+struct SsdMetric {
+    SsdMetric(std::map<std::string, std::string> labels = {})
+        : ssd_read_bytes("mooncake_ssd_read_bytes_total",
+                         "Total bytes read from SSD", labels),
+          ssd_write_bytes("mooncake_ssd_write_bytes_total",
+                          "Total bytes written to SSD", labels),
+          ssd_read_ops("mooncake_ssd_read_ops_total",
+                       "Total number of SSD read operations (key count)",
+                       labels),
+          ssd_write_ops("mooncake_ssd_write_ops_total",
+                        "Total number of SSD write operations (key count)",
+                        labels),
+          ssd_read_latency_us("mooncake_ssd_read_latency_us",
+                              "SSD BatchLoad latency per batch (us)",
+                              kSsdLatencyBucket, labels),
+          ssd_write_latency_us("mooncake_ssd_write_latency_us",
+                               "SSD BatchOffload latency per batch (us)",
+                               kSsdLatencyBucket, labels),
+          ssd_total_bytes("mooncake_ssd_total_bytes_total",
+                          "Total bytes read and written to SSD", labels),
+          ssd_total_ops("mooncake_ssd_total_ops_total",
+                        "Total number of SSD operations (key count)", labels),
+          ssd_total_latency_us("mooncake_ssd_total_latency_us",
+                               "SSD total latency per batch (us)",
+                               kSsdLatencyBucket, labels),
+          ssd_read_latency_summary("mooncake_ssd_read_latency_summary_us",
+                                   "SSD read latency quantiles (us)",
+                                   {0.5, 0.9, 0.99}, labels),
+          ssd_write_latency_summary("mooncake_ssd_write_latency_summary_us",
+                                    "SSD write latency quantiles (us)",
+                                    {0.5, 0.9, 0.99}, labels),
+          ssd_total_latency_summary("mooncake_ssd_total_latency_summary_us",
+                                    "SSD total latency quantiles (us)",
+                                    {0.5, 0.9, 0.99}, labels),
+          start_time_(std::chrono::steady_clock::now()) {}
+
+    ylt::metric::counter_t ssd_read_bytes;
+    ylt::metric::counter_t ssd_write_bytes;
+    ylt::metric::counter_t ssd_read_ops;
+    ylt::metric::counter_t ssd_write_ops;
+    ylt::metric::histogram_t ssd_read_latency_us;
+    ylt::metric::histogram_t ssd_write_latency_us;
+    ylt::metric::counter_t ssd_total_bytes;
+    ylt::metric::counter_t ssd_total_ops;
+    ylt::metric::histogram_t ssd_total_latency_us;
+    ylt::metric::summary_t ssd_read_latency_summary;
+    ylt::metric::summary_t ssd_write_latency_summary;
+    ylt::metric::summary_t ssd_total_latency_summary;
+    std::chrono::steady_clock::time_point start_time_;
+
+    void serialize(std::string& str) {
+        ssd_read_bytes.serialize(str);
+        ssd_write_bytes.serialize(str);
+        ssd_read_ops.serialize(str);
+        ssd_write_ops.serialize(str);
+        ssd_read_latency_us.serialize(str);
+        ssd_write_latency_us.serialize(str);
+        ssd_total_bytes.serialize(str);
+        ssd_total_ops.serialize(str);
+        ssd_total_latency_us.serialize(str);
+        ssd_read_latency_summary.serialize(str);
+        ssd_write_latency_summary.serialize(str);
+        ssd_total_latency_summary.serialize(str);
+    }
+
+    std::string summary_metrics() {
+        std::stringstream ss;
+        ss << "=== SSD Metrics Summary ===" << "\n";
+
+        auto read_bytes = ssd_read_bytes.value();
+        auto write_bytes = ssd_write_bytes.value();
+        auto read_ops = ssd_read_ops.value();
+        auto write_ops = ssd_write_ops.value();
+
+        auto elapsed_s = std::chrono::duration<double>(
+                             std::chrono::steady_clock::now() - start_time_)
+                             .count();
+
+        ss << "SSD Read: " << byte_size_to_string(read_bytes)
+           << ", ops=" << read_ops;
+        if (elapsed_s > 0 && read_bytes > 0) {
+            ss << ", throughput="
+               << byte_size_to_string(
+                      static_cast<int64_t>(read_bytes / elapsed_s))
+               << "/s";
+            ss << ", IOPS=" << std::fixed << std::setprecision(1)
+               << (read_ops / elapsed_s);
+        }
+        ss << "\n";
+
+        ss << "SSD Write: " << byte_size_to_string(write_bytes)
+           << ", ops=" << write_ops;
+        if (elapsed_s > 0 && write_bytes > 0) {
+            ss << ", throughput="
+               << byte_size_to_string(
+                      static_cast<int64_t>(write_bytes / elapsed_s))
+               << "/s";
+            ss << ", IOPS=" << std::fixed << std::setprecision(1)
+               << (write_ops / elapsed_s);
+        }
+        ss << "\n";
+
+        auto total_bytes = ssd_total_bytes.value();
+        auto total_ops = ssd_total_ops.value();
+        ss << "SSD Total: " << byte_size_to_string(total_bytes)
+           << ", ops=" << total_ops;
+        if (elapsed_s > 0 && total_bytes > 0) {
+            ss << ", throughput="
+               << byte_size_to_string(
+                      static_cast<int64_t>(total_bytes / elapsed_s))
+               << "/s";
+            ss << ", IOPS=" << std::fixed << std::setprecision(1)
+               << (total_ops / elapsed_s);
+        }
+        ss << "\n";
+
+        ss << "\n" << "=== SSD Latency Summary (microseconds) ===" << "\n";
+        ss << "Read: " << format_summary_percentiles(ssd_read_latency_summary)
+           << "\n";
+        ss << "Write: " << format_summary_percentiles(ssd_write_latency_summary)
+           << "\n";
+        ss << "Total: " << format_summary_percentiles(ssd_total_latency_summary)
+           << "\n";
+
+        return ss.str();
+    }
+
+   private:
+    std::string format_summary_percentiles(ylt::metric::summary_t& summary) {
+        double sum = 0;
+        uint64_t count = 0;
+        auto rates = summary.get_rates(sum, count);
+
+        if (count == 0) {
+            return "No data";
+        }
+
+        std::stringstream ss;
+        ss << std::fixed << std::setprecision(1);
+        ss << "count=" << count;
+        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";
+        if (count > 0) {
+            ss << ", avg=" << (sum / count) << "us";
+        }
+        return ss.str();
+    }
+};
+
 struct ClientMetric {
     TransferMetric transfer_metric;
     MasterClientMetric master_client_metric;
+    TransferOperationMetric transfer_operation_metric;
+    SsdMetric ssd_metric;
 
     /**
      * @brief Creates a ClientMetric instance based on environment variables
@@ -288,25 +675,44 @@ struct ClientMetric {
      *   (default: 0, 0 = collect but don't report)
      */
     static std::unique_ptr<ClientMetric> Create(
-        const std::map<std::string, std::string>& labels = {});
+        const std::map<std::string, std::string>& labels = {},
+        bool master_rpc_metrics_enabled = true);
+
+    void ObserveTransferOperation(TransferOperationKind kind,
+                                  const std::string& op_name, uint64_t bytes,
+                                  uint64_t latency_us) {
+        transfer_operation_metric.Observe(kind, op_name, bytes, latency_us);
+    }
 
     void serialize(std::string& str);
     std::string summary_metrics();
 
     uint64_t GetReportingInterval() const { return metrics_interval_seconds_; }
 
-    explicit ClientMetric(
-        uint64_t interval_seconds = 0,
-        const std::map<std::string, std::string>& labels = {});
+    explicit ClientMetric(uint64_t interval_seconds = 0,
+                          const std::map<std::string, std::string>& labels = {},
+                          bool bandwidth_reporting_enabled = true,
+                          bool master_rpc_metrics_enabled = true);
     ~ClientMetric();
 
    private:
+    struct TransferSnapshot {
+        uint64_t read_bytes;
+        uint64_t write_bytes;
+        std::chrono::steady_clock::time_point timestamp;
+    };
+
     // Metrics reporting thread management
     std::jthread metrics_reporting_thread_;
     std::atomic<bool> should_stop_metrics_thread_{false};
     uint64_t metrics_interval_seconds_{0};
+    bool bandwidth_reporting_enabled_{true};
+    bool master_rpc_metrics_enabled_{true};
+    std::mutex snapshot_mutex_;
+    std::optional<TransferSnapshot> last_report_snapshot_;
 
     void StartMetricsReportingThread();
     void StopMetricsReportingThread();
+    std::string BuildBandwidthReport();
 };
-};  // namespace mooncake
+};  // namespace mooncake

mooncake-store/include/client_service.h

@@ -24,6 +24,7 @@
 #include "master_metric_manager.h"
 #include "count_min_sketch.h"
 #include "local_hot_cache.h"
+#include "pinned_buffer_pool.h"
 
 namespace mooncake {
 
@@ -61,6 +62,8 @@ class Client {
    public:
     ~Client();
 
+    const UUID& getClientId() const { return client_id_; }
+
     /**
      * @brief Creates and initializes a new Client instance
      * @param local_hostname Local host address (IP:Port)
@@ -165,6 +168,10 @@ class Client {
     tl::expected<void, ErrorCode> Get(const std::string& object_key,
                                       const QueryResult& query_result,
                                       std::vector<Slice>& slices);
+    tl::expected<void, ErrorCode> Get(const std::string& object_key,
+                                      const QueryResult& query_result,
+                                      std::vector<Slice>& slices,
+                                      uint64_t src_offset);
     /**
      * @brief Transfers data using pre-queried object information
      * @param object_keys Keys of the objects
@@ -440,6 +447,15 @@ class Client {
         return master_client_.CalcCacheStats();
     }
 
+    void ObserveTransferOperation(TransferOperationKind kind,
+                                  const std::string& op_name, uint64_t bytes,
+                                  uint64_t latency_us) {
+        if (metrics_ != nullptr) {
+            metrics_->ObserveTransferOperation(kind, op_name, bytes,
+                                               latency_us);
+        }
+    }
+
     // For Prometheus-style metrics
     tl::expected<std::string, ErrorCode> SerializeMetrics() {
         if (metrics_ == nullptr) {
@@ -450,6 +466,10 @@ class Client {
         return str;
     }
 
+    SsdMetric* GetSsdMetricPtr() {
+        return metrics_ ? &metrics_->ssd_metric : nullptr;
+    }
+
     [[nodiscard]] std::string GetTransportEndpoint() {
         return transfer_engine_->getLocalIpAndPort();
     }
@@ -531,10 +551,16 @@ class Client {
     ErrorCode TransferData(const Replica::Descriptor& replica_descriptor,
                            std::vector<Slice>& slices,
                            TransferRequest::OpCode op_code);
+    ErrorCode TransferReadInternal(
+        const Replica::Descriptor& replica_descriptor,
+        std::vector<Slice>& slices, uint64_t src_offset);
     ErrorCode TransferWrite(const Replica::Descriptor& replica_descriptor,
                             std::vector<Slice>& slices);
     ErrorCode TransferRead(const Replica::Descriptor& replica_descriptor,
                            std::vector<Slice>& slices);
+    ErrorCode TransferReadRange(const Replica::Descriptor& replica_descriptor,
+                                std::vector<Slice>& slices,
+                                uint64_t src_offset);
 
     /**
      * @brief Prepare and use the storage backend for persisting data
@@ -646,6 +672,9 @@ class Client {
     const std::string protocol_;
 
     // Client persistent thread pool for async operations
+    // Pinned host memory pool for GPU D2H staging (must outlive
+    // write_thread_pool_)
+    std::unique_ptr<PinnedBufferPool> pinned_buffer_pool_;
     ThreadPool write_thread_pool_;
     std::shared_ptr<StorageBackend> storage_backend_;
 

mooncake-store/include/dummy_client.h

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

mooncake-store/include/file_storage.h

@@ -3,13 +3,17 @@
 #include "client_service.h"
 #include "client_buffer.hpp"
 #include "storage_backend.h"
+#include "pinned_buffer_pool.h"
 
 namespace mooncake {
 
+struct SsdMetric;
+
 class FileStorage {
    public:
     FileStorage(const FileStorageConfig& config, std::shared_ptr<Client> client,
-                const std::string& local_rpc_addr);
+                const std::string& local_rpc_addr,
+                SsdMetric* ssd_metric = nullptr);
     ~FileStorage();
 
     tl::expected<void, ErrorCode> Init();
@@ -100,7 +104,10 @@ class FileStorage {
     void ClientBufferGCThreadFunc();
 
     std::shared_ptr<Client> client_;
+    SsdMetric* ssd_metric_{nullptr};
     std::string local_rpc_addr_;
+    // Pinned host memory pool for GPU D2H staging in OffloadObjects
+    std::unique_ptr<PinnedBufferPool> pinned_buffer_pool_;
     std::shared_ptr<StorageBackendInterface> storage_backend_;
     std::shared_ptr<ClientBufferAllocator> client_buffer_allocator_;
     mutable Mutex client_buffer_mutex_;

mooncake-store/include/gpu_staging_utils.h

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

mooncake-store/include/ha/oplog/etcd_oplog_change_notifier.h

@@ -0,0 +1,103 @@
+// mooncake-store/include/etcd_oplog_change_notifier.h
+#pragma once
+
+#include <atomic>
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <mutex>
+#include <string>
+#include <thread>
+#include <vector>
+
+#include "ha/oplog/etcd_oplog_store.h"
+#include "ha/oplog/oplog_change_notifier.h"
+#include "ha/oplog/oplog_manager.h"
+#include "types.h"
+
+namespace mooncake {
+
+// Forward declaration
+class EtcdOpLogChangeNotifier;
+
+// Shared control block for safe C-style watch callbacks.
+// Because the etcd Watch goroutine can deliver callbacks after Stop()
+// returns (or even after destruction), we cannot pass a raw `this`
+// pointer as the callback context.
+struct ChangeNotifierCallbackContext {
+    std::mutex mutex;
+    EtcdOpLogChangeNotifier* notifier{nullptr};
+
+    ChangeNotifierCallbackContext() = default;
+    ChangeNotifierCallbackContext(const ChangeNotifierCallbackContext&) =
+        delete;
+    ChangeNotifierCallbackContext& operator=(
+        const ChangeNotifierCallbackContext&) = delete;
+};
+
+// OpLogChangeNotifier implementation backed by etcd Watch.
+// Migrated from OpLogReplicator's watch logic.
+class EtcdOpLogChangeNotifier : public OpLogChangeNotifier {
+   public:
+    explicit EtcdOpLogChangeNotifier(const std::string& cluster_id,
+                                     EtcdOpLogStore* oplog_store);
+    ~EtcdOpLogChangeNotifier();
+
+    ErrorCode Start(uint64_t start_sequence_id, EntryCallback on_entry,
+                    ErrorCallback on_error) override;
+    void Stop() override;
+    bool IsHealthy() const override;
+
+   private:
+    // Read historical entries and return the etcd revision for watch resume.
+    bool ReadOpLogSince(uint64_t start_seq_id, std::vector<OpLogEntry>& entries,
+                        EtcdRevisionId& revision_id);
+
+    // C-style callback for etcd Watch goroutine.
+    static void WatchCallback(void* context, const char* key, size_t key_size,
+                              const char* value, size_t value_size,
+                              int event_type, int64_t mod_revision);
+
+    // Background thread running the watch loop.
+    void WatchLoop();
+
+    // Handle a single watch event (PUT/DELETE/BROKEN).
+    void HandleWatchEvent(const std::string& key, const std::string& value,
+                          int event_type, int64_t mod_revision);
+
+    // Reconnection with exponential backoff.
+    void TryReconnect();
+
+    // Sync missed entries after reconnection.
+    bool SyncMissedEntries();
+
+    // Read and deliver all entries since start_seq_id. Updates
+    // last_processed_sequence_id_ and next_watch_revision_.
+    // Returns the number of delivered entries, or -1 on read failure.
+    int64_t DeliverHistoricalEntries(uint64_t start_seq_id);
+
+    std::string cluster_id_;
+    std::string watch_prefix_;     // "/oplog/{cluster_id}/"
+    EtcdOpLogStore* oplog_store_;  // Not owned
+
+    EntryCallback on_entry_;
+    ErrorCallback on_error_;
+
+    std::atomic<bool> running_{false};
+    std::thread watch_thread_;
+    std::atomic<uint64_t> last_processed_sequence_id_{0};
+    std::atomic<int64_t> next_watch_revision_{0};
+
+    ChangeNotifierCallbackContext* callback_ctx_{nullptr};
+
+    std::atomic<int> consecutive_errors_{0};
+    std::atomic<int> reconnect_count_{0};
+    std::atomic<bool> watch_healthy_{false};
+
+    static constexpr int kMaxConsecutiveErrors = 10;
+    static constexpr int kReconnectDelayMs = 1000;
+    static constexpr int kMaxReconnectDelayMs = 30000;
+    static constexpr int kSyncBatchSize = 1000;
+};
+
+}  // namespace mooncake

mooncake-store/include/ha/oplog/etcd_oplog_store.h

@@ -10,7 +10,8 @@
 #include <thread>
 #include <vector>
 
-#include "oplog_manager.h"
+#include "ha/oplog/oplog_manager.h"
+#include "ha/oplog/oplog_store.h"
 #include "types.h"
 
 namespace mooncake {
@@ -25,7 +26,7 @@ namespace mooncake {
  * The latest sequence_id is also stored at:
  *   /oplog/{cluster_id}/latest
  */
-class EtcdOpLogStore {
+class EtcdOpLogStore : public OpLogStore {
    public:
     /**
      * @brief Constructor.
@@ -51,7 +52,7 @@ class EtcdOpLogStore {
      *        background threads if enabled.
      * @return: Error code.
      */
-    ErrorCode Init();
+    ErrorCode Init() override;
 
     /**
      * @brief Write an OpLog entry to etcd.
@@ -60,7 +61,7 @@ class EtcdOpLogStore {
      *              If false, buffer it and return immediately (Group Commit).
      * @return: Error code.
      */
-    ErrorCode WriteOpLog(const OpLogEntry& entry, bool sync = true);
+    ErrorCode WriteOpLog(const OpLogEntry& entry, bool sync = true) override;
 
     /**
      * @brief Read an OpLog entry from etcd by sequence_id.
@@ -68,7 +69,7 @@ class EtcdOpLogStore {
      * @param entry: Output param, the OpLog entry.
      * @return: Error code.
      */
-    ErrorCode ReadOpLog(uint64_t sequence_id, OpLogEntry& entry);
+    ErrorCode ReadOpLog(uint64_t sequence_id, OpLogEntry& entry) override;
 
     /**
      * @brief Read OpLog entries starting from a given sequence_id.
@@ -78,7 +79,7 @@ class EtcdOpLogStore {
      * @return: Error code.
      */
     ErrorCode ReadOpLogSince(uint64_t start_sequence_id, size_t limit,
-                             std::vector<OpLogEntry>& entries);
+                             std::vector<OpLogEntry>& entries) override;
 
     // Like ReadOpLogSince, but also returns the etcd revision for consistent
     // "read then watch(from revision+1)" startup.
@@ -90,22 +91,22 @@ class EtcdOpLogStore {
     /**
      * @brief Get the latest sequence_id from etcd.
      * @param sequence_id: Output param, the latest sequence_id.
-     * @return: Error code. ETCD_KEY_NOT_EXIST if no OpLog exists yet.
+     * @return: Error code. OPLOG_ENTRY_NOT_FOUND if no OpLog exists yet.
      */
-    ErrorCode GetLatestSequenceId(uint64_t& sequence_id);
+    ErrorCode GetLatestSequenceId(uint64_t& sequence_id) override;
 
     // Stronger (than `/latest`) best-effort query: return the maximum existing
     // sequence_id by scanning etcd keys under /oplog/{cluster_id}/ with
     // descending key order.
-    // Return ETCD_KEY_NOT_EXIST if no OpLog exists yet.
-    ErrorCode GetMaxSequenceId(uint64_t& sequence_id);
+    // Return OPLOG_ENTRY_NOT_FOUND if no OpLog exists yet.
+    ErrorCode GetMaxSequenceId(uint64_t& sequence_id) override;
 
     /**
      * @brief Update the latest sequence_id in etcd.
      * @param sequence_id: The latest sequence_id to update.
      * @return: Error code.
      */
-    ErrorCode UpdateLatestSequenceId(uint64_t sequence_id);
+    ErrorCode UpdateLatestSequenceId(uint64_t sequence_id) override;
 
     /**
      * @brief Record the sequence_id corresponding to a snapshot.
@@ -114,16 +115,16 @@ class EtcdOpLogStore {
      * @return: Error code.
      */
     ErrorCode RecordSnapshotSequenceId(const std::string& snapshot_id,
-                                       uint64_t sequence_id);
+                                       uint64_t sequence_id) override;
 
     /**
      * @brief Get the sequence_id for a given snapshot.
      * @param snapshot_id: The snapshot ID.
      * @param sequence_id: Output param, the sequence_id.
-     * @return: Error code. ETCD_KEY_NOT_EXIST if snapshot not found.
+     * @return: Error code. OPLOG_ENTRY_NOT_FOUND if snapshot not found.
      */
     ErrorCode GetSnapshotSequenceId(const std::string& snapshot_id,
-                                    uint64_t& sequence_id);
+                                    uint64_t& sequence_id) override;
 
     /**
      * @brief Clean up OpLog entries before a given sequence_id.
@@ -131,7 +132,11 @@ class EtcdOpLogStore {
      * before_sequence_id will be deleted.
      * @return: Error code.
      */
-    ErrorCode CleanupOpLogBefore(uint64_t before_sequence_id);
+    ErrorCode CleanupOpLogBefore(uint64_t before_sequence_id) override;
+
+    // Create an EtcdOpLogChangeNotifier backed by this store.
+    std::unique_ptr<OpLogChangeNotifier> CreateChangeNotifier(
+        const std::string& cluster_id) override;
 
    private:
     /**
@@ -162,22 +167,6 @@ class EtcdOpLogStore {
     // Best-effort: find the maximum existing OpLog sequence_id in etcd.
     std::optional<uint64_t> GetMaxSequenceIdInternal() const;
 
-    /**
-     * @brief Serialize an OpLogEntry to JSON string.
-     * @param entry: The OpLog entry to serialize.
-     * @return: The JSON string.
-     */
-    std::string SerializeOpLogEntry(const OpLogEntry& entry) const;
-
-    /**
-     * @brief Deserialize a JSON string to OpLogEntry.
-     * @param json_str: The JSON string.
-     * @param entry: Output param, the OpLog entry.
-     * @return: true if successful, false otherwise.
-     */
-    bool DeserializeOpLogEntry(const std::string& json_str,
-                               OpLogEntry& entry) const;
-
     /**
      * @brief Batch update thread function.
      * Periodically updates latest_sequence_id in etcd.

mooncake-store/include/ha/oplog/localfs_oplog_store.h

@@ -0,0 +1,146 @@
+// mooncake-store/include/localfs_oplog_store.h
+#pragma once
+
+#include <atomic>
+#include <condition_variable>
+#include <cstdint>
+#include <deque>
+#include <mutex>
+#include <string>
+#include <thread>
+#include <vector>
+
+#include "ha/oplog/oplog_manager.h"
+#include "ha/oplog/oplog_store.h"
+#include "types.h"
+
+namespace mooncake {
+
+class LocalFsOpLogStore : public OpLogStore {
+   public:
+    explicit LocalFsOpLogStore(const std::string& cluster_id,
+                               const std::string& root_dir,
+                               bool enable_batch_write,
+                               int poll_interval_ms = 1000);
+    ~LocalFsOpLogStore();
+
+    ErrorCode Init() override;
+    ErrorCode WriteOpLog(const OpLogEntry& entry, bool sync = true) override;
+    ErrorCode ReadOpLog(uint64_t sequence_id, OpLogEntry& entry) override;
+    ErrorCode ReadOpLogSince(uint64_t start_sequence_id, size_t limit,
+                             std::vector<OpLogEntry>& entries) override;
+    ErrorCode GetLatestSequenceId(uint64_t& sequence_id) override;
+    ErrorCode GetMaxSequenceId(uint64_t& sequence_id) override;
+    ErrorCode UpdateLatestSequenceId(uint64_t sequence_id) override;
+    ErrorCode RecordSnapshotSequenceId(const std::string& snapshot_id,
+                                       uint64_t sequence_id) override;
+    ErrorCode GetSnapshotSequenceId(const std::string& snapshot_id,
+                                    uint64_t& sequence_id) override;
+    ErrorCode CleanupOpLogBefore(uint64_t before_sequence_id) override;
+    std::unique_ptr<OpLogChangeNotifier> CreateChangeNotifier(
+        const std::string& cluster_id) override;
+
+   private:
+    // Segment file format: 32-byte header + length-prefixed entries
+    static constexpr char kSegmentMagic[4] = {'M', 'C', 'S', 'G'};
+    static constexpr uint32_t kSegmentVersion = 1;
+    static constexpr size_t kSegmentHeaderSize = 32;
+
+    // Segment header layout (all little-endian):
+    //   [0..3]   magic    "MCSG"
+    //   [4..7]   version  uint32_t
+    //   [8..15]  min_seq  uint64_t
+    //   [16..23] max_seq  uint64_t
+    //   [24..27] count    uint32_t
+    //   [28..31] reserved uint32_t (0)
+    struct SegmentHeader {
+        char magic[4];
+        uint32_t version;
+        uint64_t min_seq;
+        uint64_t max_seq;
+        uint32_t entry_count;
+        uint32_t reserved;
+    };
+    static_assert(sizeof(SegmentHeader) == kSegmentHeaderSize,
+                  "SegmentHeader must be 32 bytes");
+
+    // Segment file name info parsed from filename
+    struct SegmentInfo {
+        std::string filename;
+        uint64_t min_seq;
+        uint64_t max_seq;
+    };
+
+    // Group Commit batch entry
+    struct BatchEntry {
+        std::string serialized_value;
+        uint64_t sequence_id;
+        bool is_sync;
+    };
+
+    // Directory and path helpers
+    std::string SegmentsDir() const;
+    std::string SnapshotsDir() const;
+    std::string LatestFilePath() const;
+    std::string BuildSegmentFilename(uint64_t min_seq, uint64_t max_seq) const;
+    std::string BuildSnapshotPath(const std::string& snapshot_id) const;
+
+    // Segment I/O
+    ErrorCode WriteSegmentFile(const std::vector<BatchEntry>& entries);
+    ErrorCode ReadSegmentEntries(const std::string& filepath,
+                                 std::vector<OpLogEntry>& entries);
+    ErrorCode ReadSegmentHeader(const std::string& filepath,
+                                SegmentHeader& header);
+    std::vector<SegmentInfo> ListSegments() const;
+    static bool ParseSegmentFilename(const std::string& filename,
+                                     uint64_t& min_seq, uint64_t& max_seq);
+    ErrorCode NormalizeBatchEntries(std::vector<BatchEntry>& entries) const;
+    ErrorCode VerifyPersistedEntryMatches(uint64_t sequence_id,
+                                          const std::string& serialized_value);
+    ErrorCode RecoverPersistedState();
+
+    // Atomic file write: write to .tmp, fsync, rename
+    ErrorCode AtomicWriteFile(const std::string& target_path,
+                              const std::string& content);
+    ErrorCode AtomicWriteFile(const std::string& target_path, const void* data,
+                              size_t size);
+
+    // Cleanup temp files from previous crash
+    void CleanupTempFiles();
+
+    // Snapshot ID validation
+    static bool ValidateSnapshotId(const std::string& snapshot_id);
+
+    // Read a uint64 value from a single-value text file
+    ErrorCode ReadUint64FromFile(const std::string& filepath,
+                                 uint64_t& value) const;
+
+    // Batch write thread
+    void BatchWriteThread();
+    void FlushBatch();
+
+    // Members
+    std::string cluster_id_;
+    std::string root_dir_;
+    std::string cluster_dir_;  // root_dir_/cluster_id_
+    bool enable_batch_write_;
+    int poll_interval_ms_;
+
+    // Group Commit state
+    mutable std::mutex batch_mutex_;
+    std::deque<BatchEntry> pending_batch_;
+    std::condition_variable cv_batch_updated_;
+    std::condition_variable cv_sync_completed_;
+    std::atomic<bool> batch_write_running_{false};
+    std::thread batch_write_thread_;
+    std::atomic<uint64_t> last_persisted_seq_id_{0};
+
+    // Batch write configs
+    static constexpr size_t kBatchCountLimit = 100;
+    static constexpr int kBatchTimeoutMs = 100;
+    static constexpr int kSyncWaitTimeoutMs = 3000;
+    static constexpr int kFlushRetryCount = 3;
+    static constexpr int kFlushRetryIntervalMs = 50;
+};
+
+}  // namespace mooncake

mooncake-store/include/ha/oplog/oplog_applier.h

@@ -9,13 +9,13 @@
 #include <string>
 #include <vector>
 
-#include "oplog_manager.h"
+#include "ha/oplog/oplog_manager.h"
 #include "metadata_store.h"
 
 namespace mooncake {
 
 // Forward declaration
-class EtcdOpLogStore;
+class OpLogStore;
 
 /**
  * @brief Apply OpLog entries to Standby metadata store with ordering guarantee
@@ -28,11 +28,21 @@ class OpLogApplier {
     /**
      * @brief Constructor
      * @param metadata_store Metadata store to apply changes to
-     * @param cluster_id Cluster ID for accessing etcd OpLog (optional, for
-     * requesting missing OpLog)
+     * @param cluster_id Cluster ID (for validation only)
+     * @param oplog_store Optional OpLogStore for requesting missing OpLog
+     *                    entries (caller owns the pointer)
      */
     explicit OpLogApplier(MetadataStore* metadata_store,
-                          const std::string& cluster_id = std::string());
+                          const std::string& cluster_id = std::string(),
+                          OpLogStore* oplog_store = nullptr);
+
+    ~OpLogApplier() = default;
+
+    /**
+     * @brief Set or replace the OpLogStore used for requesting missing entries
+     * @param oplog_store OpLogStore pointer (caller owns the pointer)
+     */
+    void SetOpLogStore(OpLogStore* oplog_store) { oplog_store_ = oplog_store; }
 
     /**
      * @brief Apply a single OpLog entry (with ordering checks)
@@ -48,14 +58,6 @@ class OpLogApplier {
      */
     size_t ApplyOpLogEntries(const std::vector<OpLogEntry>& entries);
 
-    /**
-     * @brief Get the current sequence ID for a key (DEPRECATED)
-     * @param key Object key
-     * @return Always returns 0 - global sequence_id is used for ordering
-     * @deprecated Use global sequence_id for ordering
-     */
-    uint64_t GetKeySequenceId(const std::string& key) const;
-
     /**
      * @brief Get the expected global sequence ID
      * @return Expected global sequence ID
@@ -116,30 +118,17 @@ class OpLogApplier {
     void ApplyRemove(const OpLogEntry& entry);
 
     /**
-     * @brief Request missing OpLog entry from etcd
+     * @brief Request missing OpLog entry from the store
      * @param missing_seq_id Missing sequence ID
      * @return true if entry was found and applied, false otherwise
      */
     bool RequestMissingOpLog(uint64_t missing_seq_id);
 
-    /**
-     * @brief Schedule wait for missing entries
-     * @param missing_seq_id Missing sequence ID
-     */
-    void ScheduleWaitForMissingEntries(uint64_t missing_seq_id);
-
     MetadataStore* metadata_store_;
 
-    // EtcdOpLogStore for requesting missing OpLog entries (optional)
+    // OpLogStore for requesting missing OpLog entries (optional, not owned)
     std::string cluster_id_;
-    mutable std::mutex etcd_oplog_store_mutex_;
-    mutable std::unique_ptr<EtcdOpLogStore> etcd_oplog_store_;
-
-    /**
-     * @brief Get or create EtcdOpLogStore instance (lazy initialization)
-     * @return Pointer to EtcdOpLogStore, or nullptr if cluster_id is not set
-     */
-    EtcdOpLogStore* GetEtcdOpLogStore() const;
+    OpLogStore* oplog_store_{nullptr};
 
     // Note: key_sequence_map_ has been removed.
     // Global sequence_id is sufficient for ordering guarantee.

mooncake-store/include/ha/oplog/oplog_change_notifier.h

@@ -0,0 +1,40 @@
+// mooncake-store/include/oplog_change_notifier.h
+#pragma once
+
+#include <cstdint>
+#include <functional>
+
+#include "ha/oplog/oplog_manager.h"
+#include "types.h"
+
+namespace mooncake {
+
+// Abstract interface for watching OpLog changes.
+//
+// Delivery semantics:
+//   The notifier provides **at-most-once delivery**. Its internal cursor
+//   tracks *fetch progress* (where to read from next), NOT consumer
+//   processing progress. The cursor advances unconditionally after entries
+//   are dispatched to EntryCallback, regardless of whether the callback
+//   processed them successfully.
+//
+//   Consumers (e.g. OpLogApplier) must maintain their own cursor
+//   (expected_sequence_id_) and handle gaps, duplicates, and late arrivals
+//   independently.
+//
+// Implementations: EtcdOpLogChangeNotifier (push via Watch),
+//                  PollingOpLogChangeNotifier (poll via ReadOpLogSince)
+class OpLogChangeNotifier {
+   public:
+    virtual ~OpLogChangeNotifier() = default;
+
+    using EntryCallback = std::function<void(const OpLogEntry& entry)>;
+    using ErrorCallback = std::function<void(ErrorCode error)>;
+
+    virtual ErrorCode Start(uint64_t start_sequence_id, EntryCallback on_entry,
+                            ErrorCallback on_error) = 0;
+    virtual void Stop() = 0;
+    virtual bool IsHealthy() const = 0;
+};
+
+}  // namespace mooncake

mooncake-store/include/ha/oplog/oplog_manager.h

@@ -15,7 +15,7 @@
 namespace mooncake {
 
 // Forward declaration
-class EtcdOpLogStore;
+class OpLogStore;
 
 // Operation types for hot-standby replication.
 // This is a minimal subset that can be extended later.
@@ -49,16 +49,16 @@ struct OpLogEntry {
  *
  * This class is intentionally simple: it keeps a bounded deque of OpLogEntry
  * and provides append / get-since primitives. It can later be extended to
- * or to spill to disk if needed. In the new etcd-based design, OpLog will be
- * written to etcd.
+ * or to spill to disk if needed. OpLog entries are persisted to the
+ * configured OpLogStore backend (etcd, local filesystem, etc.).
  */
 class OpLogManager {
    public:
     OpLogManager();
 
-    // Set the EtcdOpLogStore for writing OpLog to etcd (optional).
+    // Set the OpLogStore for writing OpLog to persistent storage (optional).
     // If not set, OpLog will only be stored in memory buffer.
-    void SetEtcdOpLogStore(std::shared_ptr<EtcdOpLogStore> etcd_oplog_store);
+    void SetOpLogStore(std::shared_ptr<OpLogStore> oplog_store);
 
     // Append a new entry and return the assigned sequence_id.
     // This is a best-effort (async) path: the entry is buffered in memory
@@ -78,11 +78,11 @@ class OpLogManager {
     OpLogEntry AllocateEntry(OpType type, const std::string& key,
                              const std::string& payload = std::string());
 
-    // Persist an already-allocated entry to etcd using its sequence_id.
+    // Persist an already-allocated entry to the store using its sequence_id.
     // Does NOT modify sequence counters.
-    ErrorCode PersistEntryToEtcd(const OpLogEntry& entry) const;
+    ErrorCode PersistEntry(const OpLogEntry& entry) const;
 
-    // Append a new entry and durably persist it to etcd (if EtcdOpLogStore is
+    // Append a new entry and durably persist it to the store (if OpLogStore is
     // set).
     //
     // This is intended for operations that may free/reuse memory (e.g. REMOVE),
@@ -92,7 +92,7 @@ class OpLogManager {
     //
     // Design (updated for seq pre-allocation):
     // - sequence_id is allocated first and never reused.
-    // - If etcd write fails, caller may retry PersistEntryToEtcd with the same
+    // - If etcd write fails, caller may retry PersistEntry with the same
     //   entry (sequence_id fixed and "smaller" than later entries).
     tl::expected<uint64_t, ErrorCode> AppendAndPersist(
         OpType type, const std::string& key,
@@ -109,9 +109,13 @@ class OpLogManager {
     // Current number of entries in the buffer.
     size_t GetEntryCount() const;
 
+    // Clean up OpLog entries in etcd before a given sequence_id.
+    // Delegates to EtcdOpLogStore::CleanupOpLogBefore.
+    ErrorCode CleanupOpLogBefore(uint64_t before_sequence_id);
+
     // Verify checksum of an OpLogEntry payload.
     // Returns true if checksum matches, false otherwise.
-    // This is public so OpLogWatcher and OpLogApplier can validate entries.
+    // This is public so OpLogReplicator and OpLogApplier can validate entries.
     static bool VerifyChecksum(const OpLogEntry& entry);
 
     // Basic DoS protection for externally sourced OpLog entries (etcd watch /
@@ -140,8 +144,8 @@ class OpLogManager {
     // All operations are applied in sequence_id order, which ensures
     // consistency.
 
-    // Optional etcd OpLog store for persistent storage
-    std::shared_ptr<EtcdOpLogStore> etcd_oplog_store_;
+    // Optional OpLog store for persistent storage
+    std::shared_ptr<OpLogStore> oplog_store_;
 
     // Simple bounds to avoid unbounded memory growth.
     static constexpr size_t kMaxBufferEntries_ = 100000;

mooncake-store/include/ha/oplog/oplog_replicator.h

@@ -0,0 +1,87 @@
+#pragma once
+
+#include <atomic>
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <string>
+
+#include "ha/oplog/oplog_change_notifier.h"
+#include "ha/oplog/oplog_manager.h"
+#include "standby_state_machine.h"
+#include "types.h"
+
+namespace mooncake {
+
+// Forward declarations
+class OpLogApplier;
+
+// Callback type for state events
+using ReplicatorStateCallback = std::function<void(StandbyEvent)>;
+
+/**
+ * @brief Replicate OpLog entries from a remote source and apply them locally.
+ *
+ * Delegates watch/notification to an OpLogChangeNotifier and applies
+ * received entries via OpLogApplier. This class is a thin orchestration
+ * layer; the actual watch implementation lives in OpLogChangeNotifier.
+ */
+class OpLogReplicator {
+   public:
+    /**
+     * @brief Constructor
+     * @param notifier Change notifier that delivers OpLog entries
+     * @param applier OpLog applier to process entries
+     */
+    OpLogReplicator(OpLogChangeNotifier* notifier, OpLogApplier* applier);
+
+    ~OpLogReplicator();
+
+    /**
+     * @brief Start replication from the beginning.
+     */
+    void Start();
+
+    /**
+     * @brief Start from a known last-applied sequence_id.
+     */
+    bool StartFromSequenceId(uint64_t start_seq_id);
+
+    /**
+     * @brief Stop replication.
+     */
+    void Stop();
+
+    /**
+     * @brief Get the last processed sequence ID.
+     */
+    uint64_t GetLastProcessedSequenceId() const;
+
+    /**
+     * @brief Set callback for state events.
+     */
+    void SetStateCallback(ReplicatorStateCallback callback) {
+        state_callback_ = std::move(callback);
+    }
+
+    /**
+     * @brief Check if replication is healthy.
+     */
+    bool IsHealthy() const;
+
+   private:
+    void NotifyStateEvent(StandbyEvent event) {
+        if (state_callback_) {
+            state_callback_(event);
+        }
+    }
+
+    OpLogChangeNotifier* notifier_;
+    OpLogApplier* applier_;
+    std::atomic<uint64_t> last_processed_sequence_id_{0};
+    std::atomic<bool> running_{false};
+
+    ReplicatorStateCallback state_callback_;
+};
+
+}  // namespace mooncake

mooncake-store/include/ha/oplog/oplog_serializer.h

@@ -0,0 +1,18 @@
+// mooncake-store/include/oplog_serializer.h
+#pragma once
+
+#include <string>
+
+#include "ha/oplog/oplog_manager.h"
+
+namespace mooncake {
+
+// Serialize an OpLogEntry to JSON string (with base64-encoded payload).
+// Format is backend-agnostic; all storage backends should use this.
+std::string SerializeOpLogEntry(const OpLogEntry& entry);
+
+// Deserialize a JSON string to OpLogEntry.
+// Returns true on success, false on parse error or size validation failure.
+bool DeserializeOpLogEntry(const std::string& json_str, OpLogEntry& entry);
+
+}  // namespace mooncake

mooncake-store/include/ha/oplog/oplog_store.h

@@ -0,0 +1,64 @@
+// mooncake-store/include/oplog_store.h
+#pragma once
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "ha/oplog/oplog_change_notifier.h"
+#include "ha/oplog/oplog_manager.h"
+#include "types.h"
+
+namespace mooncake {
+
+// Normalize and validate cluster_id for OpLog key prefix construction.
+// Strips trailing slashes, then validates the remaining string.
+// Returns true if valid (or empty after normalization), false otherwise.
+inline bool NormalizeAndValidateClusterId(std::string& cluster_id) {
+    while (!cluster_id.empty() && cluster_id.back() == '/') {
+        cluster_id.pop_back();
+    }
+    return cluster_id.empty() || IsValidClusterIdComponent(cluster_id);
+}
+
+// Abstract interface for OpLog persistent storage.
+// Implementations: EtcdOpLogStore, (future) HdfsOpLogStore, etc.
+class OpLogStore {
+   public:
+    virtual ~OpLogStore() = default;
+    virtual ErrorCode Init() = 0;
+
+    // Write
+    virtual ErrorCode WriteOpLog(const OpLogEntry& entry, bool sync = true) = 0;
+
+    // Read
+    virtual ErrorCode ReadOpLog(uint64_t sequence_id, OpLogEntry& entry) = 0;
+    virtual ErrorCode ReadOpLogSince(uint64_t start_sequence_id, size_t limit,
+                                     std::vector<OpLogEntry>& entries) = 0;
+
+    // Sequence ID management
+    virtual ErrorCode GetLatestSequenceId(uint64_t& sequence_id) = 0;
+    virtual ErrorCode GetMaxSequenceId(uint64_t& sequence_id) = 0;
+    virtual ErrorCode UpdateLatestSequenceId(uint64_t sequence_id) = 0;
+
+    // Snapshot
+    virtual ErrorCode RecordSnapshotSequenceId(const std::string& snapshot_id,
+                                               uint64_t sequence_id) = 0;
+    virtual ErrorCode GetSnapshotSequenceId(const std::string& snapshot_id,
+                                            uint64_t& sequence_id) = 0;
+
+    // Cleanup
+    virtual ErrorCode CleanupOpLogBefore(uint64_t before_sequence_id) = 0;
+
+    // Create a change notifier for this store.
+    // Each backend provides its own notifier (e.g., etcd watch, polling).
+    // Returns nullptr if the backend does not support change notification.
+    virtual std::unique_ptr<OpLogChangeNotifier> CreateChangeNotifier(
+        const std::string& cluster_id) {
+        (void)cluster_id;
+        return nullptr;
+    }
+};
+
+}  // namespace mooncake

mooncake-store/include/ha/oplog/oplog_store_factory.h

@@ -0,0 +1,74 @@
+// mooncake-store/include/oplog_store_factory.h
+#pragma once
+
+#include <algorithm>
+#include <memory>
+#include <string>
+
+#include "ha/oplog/oplog_store.h"
+
+namespace mooncake {
+
+enum class OpLogStoreRole {
+    WRITER,  // Primary: enable batch_write + batch_update
+    READER,  // Standby: read-only
+};
+
+enum class OpLogStoreType {
+    ETCD,
+    LOCAL_FS,
+};
+
+#ifdef STORE_USE_ETCD
+static constexpr OpLogStoreType kDefaultOpLogStoreType = OpLogStoreType::ETCD;
+#else
+static constexpr OpLogStoreType kDefaultOpLogStoreType =
+    OpLogStoreType::LOCAL_FS;
+#endif
+
+// Parse string to OpLogStoreType (case-insensitive).
+// Parse string to OpLogStoreType (case-insensitive).
+// Returns kDefaultOpLogStoreType for unrecognized strings.
+inline OpLogStoreType ParseOpLogStoreType(const std::string& type_str) {
+    std::string lower = type_str;
+    std::transform(lower.begin(), lower.end(), lower.begin(), ::tolower);
+    if (lower == "localfs" || lower == "local_fs") {
+        return OpLogStoreType::LOCAL_FS;
+    }
+    if (lower == "etcd") {
+        return OpLogStoreType::ETCD;
+    }
+    return kDefaultOpLogStoreType;
+}
+
+inline std::string OpLogStoreTypeToString(OpLogStoreType type) {
+    switch (type) {
+        case OpLogStoreType::LOCAL_FS:
+            return "localfs";
+        case OpLogStoreType::ETCD:
+        default:
+            return "etcd";
+    }
+}
+
+// Default configuration values for LocalFS OpLog store
+static constexpr const char* kDefaultOpLogRootDir = "/tmp/mooncake_oplog";
+static constexpr int kDefaultOpLogPollIntervalMs = 1000;
+
+class OpLogStoreFactory {
+   public:
+    /**
+     * @brief Create and initialize an OpLogStore instance.
+     *
+     * The returned instance is fully initialized (Init() has already been
+     * called internally). Callers must NOT call Init() again.
+     * Returns nullptr if the requested backend is unavailable or
+     * initialization fails.
+     */
+    static std::unique_ptr<OpLogStore> Create(
+        OpLogStoreType type, const std::string& cluster_id, OpLogStoreRole role,
+        const std::string& oplog_root_dir = kDefaultOpLogRootDir,
+        int poll_interval_ms = kDefaultOpLogPollIntervalMs);
+};
+
+}  // namespace mooncake

mooncake-store/include/ha/oplog/polling_oplog_change_notifier.h

@@ -0,0 +1,45 @@
+// mooncake-store/include/polling_oplog_change_notifier.h
+#pragma once
+
+#include <atomic>
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+#include "ha/oplog/oplog_change_notifier.h"
+#include "ha/oplog/oplog_store.h"
+
+namespace mooncake {
+
+class PollingOpLogChangeNotifier : public OpLogChangeNotifier {
+   public:
+    PollingOpLogChangeNotifier(OpLogStore* store, int poll_interval_ms);
+    ~PollingOpLogChangeNotifier();
+
+    ErrorCode Start(uint64_t start_sequence_id, EntryCallback on_entry,
+                    ErrorCallback on_error) override;
+    void Stop() override;
+    bool IsHealthy() const override;
+
+   private:
+    void PollLoop();
+
+    OpLogStore* store_;  // Not owned
+    int poll_interval_ms_;
+
+    EntryCallback on_entry_;
+    ErrorCallback on_error_;
+
+    std::atomic<bool> running_{false};
+    std::thread poll_thread_;
+    std::atomic<uint64_t> last_sequence_id_{0};
+    std::atomic<bool> healthy_{false};
+
+    // For interruptible sleep on Stop()
+    std::mutex stop_mutex_;
+    std::condition_variable stop_cv_;
+
+    static constexpr size_t kPollBatchSize = 1000;
+};
+
+}  // namespace mooncake

mooncake-store/include/ha/oplog_store.h

@@ -1,28 +0,0 @@
-#pragma once
-
-#include <chrono>
-#include <cstddef>
-
-#include <ylt/util/tl/expected.hpp>
-
-#include "ha/ha_types.h"
-
-namespace mooncake {
-namespace ha {
-
-class OpLogStore {
-   public:
-    virtual ~OpLogStore() = default;
-
-    virtual tl::expected<OpLogSequenceId, ErrorCode> Append(
-        const OpLogAppendRequest& request) = 0;
-
-    virtual tl::expected<OpLogPollResult, ErrorCode> PollFrom(
-        OpLogSequenceId start_seq, size_t max_records,
-        std::chrono::milliseconds timeout) = 0;
-
-    virtual tl::expected<OpLogSequenceId, ErrorCode> GetLatestSequence() = 0;
-};
-
-}  // namespace ha
-}  // namespace mooncake

mooncake-store/include/hot_standby_service.h

@@ -13,9 +13,10 @@
 #include <vector>
 
 #include "metadata_store.h"
-#include "oplog_applier.h"
-#include "oplog_manager.h"
-#include "oplog_watcher.h"
+#include "ha/oplog/oplog_applier.h"
+#include "ha/oplog/oplog_manager.h"
+#include "ha/oplog/oplog_replicator.h"
+#include "ha/oplog/oplog_store_factory.h"
 #include "ha/snapshot/snapshot_provider.h"
 #include "standby_state_machine.h"
 #include "types.h"
@@ -45,6 +46,11 @@ struct HotStandbyConfig {
     // snapshot bootstrap phase and keeps the standby in a snapshot-only steady
     // state.
     bool enable_oplog_following{true};
+
+    // OpLog store configuration
+    OpLogStoreType oplog_store_type{kDefaultOpLogStoreType};
+    std::string oplog_store_root_dir{kDefaultOpLogRootDir};
+    int oplog_poll_interval_ms{kDefaultOpLogPollIntervalMs};
 };
 
 /**
@@ -85,12 +91,12 @@ class HotStandbyService {
      * following
      * @param primary_address Address of the Primary Master (not used with
      * OpLog backend-based sync)
-     * @param etcd_endpoints Comma-separated OpLog backend endpoints
+     * @param oplog_endpoints Comma-separated OpLog backend endpoints
      * @param cluster_id Cluster identifier for OpLog path
      * @return ErrorCode::OK on success
      */
     ErrorCode Start(const std::string& primary_address,
-                    const std::string& etcd_endpoints,
+                    const std::string& oplog_endpoints,
                     const std::string& cluster_id);
 
     /**
@@ -162,7 +168,7 @@ class HotStandbyService {
     }
 
     /**
-     * @brief Callback for OpLogWatcher state changes
+     * @brief Callback for OpLogReplicator state changes
      * @param event The event to process
      */
     void OnWatcherEvent(StandbyEvent event);
@@ -225,6 +231,7 @@ class HotStandbyService {
         bool Remove(const std::string& key) override;
         bool Exists(const std::string& key) const override;
         size_t GetKeyCount() const override;
+        void Clear();
 
         // Snapshot for promotion/restore.
         void Snapshot(
@@ -241,10 +248,12 @@ class HotStandbyService {
 
     // OpLog replication components
     std::unique_ptr<OpLogApplier> oplog_applier_;
-    std::unique_ptr<OpLogWatcher> oplog_watcher_;
+    std::shared_ptr<OpLogStore> watcher_oplog_store_;
+    std::unique_ptr<OpLogChangeNotifier> oplog_change_notifier_;
+    std::unique_ptr<OpLogReplicator> oplog_replicator_;
 
-    // Configuration for etcd-based OpLog sync
-    std::string etcd_endpoints_;
+    // Configuration for OpLog sync
+    std::string oplog_endpoints_;
     std::string cluster_id_;
 
     // Replication state

mooncake-store/include/master_client.h

@@ -3,11 +3,14 @@
 #include <csignal>
 #include <memory>
 #include <string>
+#include <type_traits>
 #include <vector>
+#include <variant>
 #include <cstdlib>
 #include <boost/functional/hash.hpp>
 #include <ylt/coro_rpc/coro_rpc_client.hpp>
 #include <ylt/coro_io/client_pool.hpp>
+#include <ylt/coro_io/ibverbs/ib_socket.hpp>
 
 #include "client_metric.h"
 #include "replica.h"
@@ -20,6 +23,29 @@ namespace mooncake {
 
 static const std::string kDefaultMasterAddress = "localhost:50051";
 
+namespace detail {
+
+template <typename Variant, typename T>
+struct variant_contains : std::false_type {};
+
+template <typename... Ts, typename T>
+struct variant_contains<std::variant<Ts...>, T>
+    : std::bool_constant<(std::is_same_v<Ts, T> || ...)> {};
+
+template <typename Variant, typename T>
+inline constexpr bool variant_contains_v =
+    variant_contains<std::decay_t<Variant>, T>::value;
+
+template <typename SocketConfigVariant>
+inline void MaybeEnableRdmaSocketConfig(SocketConfigVariant& socket_config) {
+    if constexpr (variant_contains_v<SocketConfigVariant,
+                                     coro_io::ib_socket_t::config_t>) {
+        socket_config = coro_io::ib_socket_t::config_t{};
+    }
+}
+
+}  // namespace detail
+
 /**
  * @brief Client for interacting with the mooncake master service
  */
@@ -37,8 +63,8 @@ class MasterClient {
         pool_conf.host_alive_detect_duration = std::chrono::seconds(0);
         const char* value = std::getenv("MC_RPC_PROTOCOL");
         if (value && std::string_view(value) == "rdma") {
-            pool_conf.client_config.socket_config =
-                coro_io::ib_socket_t::config_t{};
+            detail::MaybeEnableRdmaSocketConfig(
+                pool_conf.client_config.socket_config);
         }
         client_pools_ =
             std::make_shared<coro_io::client_pools<coro_rpc::coro_rpc_client>>(

mooncake-store/include/master_config.h

@@ -81,6 +81,10 @@ struct MasterConfig {
     std::string cxl_path;
     size_t cxl_size;
     bool enable_cxl = false;
+
+    // Offload-on-evict: defer LOCAL_DISK offload to eviction time
+    bool offload_on_evict = false;
+    bool offload_force_evict = false;
 };
 
 class MasterServiceSupervisorConfig {
@@ -140,6 +144,8 @@ class MasterServiceSupervisorConfig {
     std::string cxl_path = DEFAULT_CXL_PATH;
     size_t cxl_size = DEFAULT_CXL_SIZE;
     bool enable_cxl = false;
+    bool offload_on_evict = false;
+    bool offload_force_evict = false;
     MasterServiceSupervisorConfig() = default;
 
     // From MasterConfig
@@ -155,6 +161,8 @@ class MasterServiceSupervisorConfig {
         eviction_high_watermark_ratio = config.eviction_high_watermark_ratio;
         client_live_ttl_sec = config.client_live_ttl_sec;
         enable_offload = config.enable_offload;
+        offload_on_evict = config.offload_on_evict;
+        offload_force_evict = config.offload_force_evict;
         rpc_port = static_cast<int>(config.rpc_port);
         rpc_thread_num = static_cast<size_t>(config.rpc_thread_num);
 
@@ -267,6 +275,9 @@ 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;
     std::string root_fs_dir = DEFAULT_ROOT_FS_DIR;
@@ -321,7 +332,13 @@ 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()) {
+            ha_backend_connstring = config.etcd_endpoints;
+        }
         cluster_id = config.cluster_id;
         root_fs_dir = config.root_fs_dir;
         global_file_segment_size = config.global_file_segment_size;
@@ -395,7 +412,13 @@ 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()) {
+            ha_backend_connstring = config.etcd_endpoints;
+        }
         cluster_id = config.cluster_id;
         root_fs_dir = config.root_fs_dir;
         global_file_segment_size = config.global_file_segment_size;
@@ -445,6 +468,7 @@ class MasterServiceConfigBuilder {
     int64_t client_live_ttl_sec_ = DEFAULT_CLIENT_LIVE_TTL_SEC;
     bool enable_ha_ = false;
     bool enable_offload_ = false;
+    std::string ha_backend_type_ = "etcd";
     std::string ha_backend_connstring_;
     std::string cluster_id_ = DEFAULT_CLUSTER_ID;
     std::string root_fs_dir_ = DEFAULT_ROOT_FS_DIR;
@@ -527,6 +551,12 @@ class MasterServiceConfigBuilder {
         return *this;
     }
 
+    MasterServiceConfigBuilder& set_ha_backend_type(
+        const std::string& backend_type) {
+        ha_backend_type_ = backend_type;
+        return *this;
+    }
+
     MasterServiceConfigBuilder& set_ha_backend_connstring(
         const std::string& connstring) {
         ha_backend_connstring_ = connstring;
@@ -720,6 +750,9 @@ 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;
     std::string root_fs_dir = DEFAULT_ROOT_FS_DIR;
@@ -770,6 +803,9 @@ 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;
         root_fs_dir = config.root_fs_dir;
@@ -825,6 +861,7 @@ inline MasterServiceConfig MasterServiceConfigBuilder::build() const {
     config.client_live_ttl_sec = client_live_ttl_sec_;
     config.enable_ha = enable_ha_;
     config.enable_offload = enable_offload_;
+    config.ha_backend_type = ha_backend_type_;
     config.ha_backend_connstring = ha_backend_connstring_;
     config.cluster_id = cluster_id_;
     config.root_fs_dir = root_fs_dir_;