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

---------

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

.claude/skills/mooncake-ci-local/SKILL.md

@@ -0,0 +1,266 @@
+---
+name: mooncake-ci-local
+description: Run Mooncake CI test suite locally — maps GitHub Actions CI steps to local commands. Use this skill whenever the user wants to run tests locally, reproduce a CI failure, check if their changes break tests, or run any subset of the CI test suite (C++ unit tests via ctest, Python integration tests, code format checks, or the full test pipeline). Trigger on phrases like "run tests", "run CI locally", "reproduce CI failure", "check my changes", "test before PR", "run ctest", "run python tests", "run all tests".
+---
+
+# Mooncake CI Local Test Runner
+
+You help users run the Mooncake CI test suite locally. The CI has three test layers. Map what the user wants to the right layer, check prerequisites, and run the tests.
+
+## CI Test Layers
+
+### Layer 1 — C++ Unit Tests (ctest)
+**CI equivalent:** `build` job in `ci.yml` — "Test (in build env) with coverage"
+
+**Prerequisite services:**
+```bash
+# 1. etcd (port 2379)
+etcd --advertise-client-urls http://127.0.0.1:2379 --listen-client-urls http://127.0.0.1:2379 &
+sleep 2
+etcdctl --endpoints=http://127.0.0.1:2379 endpoint health  # verify
+
+# 2. HTTP metadata server (port 8080)
+cd mooncake-transfer-engine/example/http-metadata-server-python
+pip install aiohttp
+python ./bootstrap_server.py &
+cd -
+```
+
+**Run:**
+```bash
+cd build
+export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
+MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 ctest -j --output-on-failure
+```
+
+**Run specific test:**
+```bash
+cd build
+MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 ctest -R <test_name_pattern> --output-on-failure
+# List all available tests: ctest -N
+```
+
+### Layer 2 — Python Integration Tests
+**CI equivalent:** `test-wheel-ubuntu` job — `run_tests.sh`
+
+**Prerequisite:** Mooncake wheel must be installed (either via `pip install` or via `make install` after build).
+
+**Check install:**
+```bash
+python -c "import mooncake; print('OK')"
+which mooncake_master  # must NOT be /usr/local/bin (must be from Python package)
+```
+
+**Run full suite:**
+```bash
+# Start metadata server first
+mooncake_http_metadata_server --port 8080 &
+sleep 1
+
+cd mooncake-wheel/tests
+MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 MC_FORCE_TCP=true \
+  bash ../../scripts/run_tests.sh
+```
+
+**Individual Python tests** (all require metadata server + mooncake_master on port 50051):
+```bash
+# Setup shared services
+mooncake_http_metadata_server --port 8080 &
+mooncake_master --default_kv_lease_ttl=500 &
+sleep 2
+
+cd mooncake-wheel/tests
+export MC_METADATA_SERVER=http://127.0.0.1:8080/metadata
+export DEFAULT_KV_LEASE_TTL=500
+export MC_FORCE_TCP=true
+
+# Pick any test:
+python test_distributed_object_store.py
+python test_replicated_distributed_object_store.py
+python test_put_get_tensor.py          # requires torch + numpy
+python test_safetensor_functions.py    # requires safetensors
+python test_dummy_client.py
+python test_cli.py
+python test_distributed_object_store_cxl.py  # requires CXL build
+```
+
+**Transfer engine tests specifically:**
+```bash
+cd mooncake-wheel/tests
+MC_METADATA_SERVER=http://127.0.0.1:8080/metadata MC_FORCE_TCP=true python transfer_engine_target.py &
+TARGET_PID=$!
+MC_METADATA_SERVER=http://127.0.0.1:8080/metadata MC_FORCE_TCP=true python transfer_engine_initiator_test.py
+kill $TARGET_PID
+```
+
+**Scripts-based tests** (from `test-wheel-ubuntu` job):
+```bash
+# Tensor API perf test
+export MOONCAKE_MASTER="127.0.0.1:50051"
+export MOONCAKE_TE_META_DATA_SERVER="http://127.0.0.1:8080/metadata"
+export MOONCAKE_PROTOCOL="tcp"
+export LOCAL_HOSTNAME="127.0.0.1"
+python scripts/test_tensor_api.py -n 1
+python scripts/test_async_store.py
+python scripts/test_copy_move_api.py
+```
+
+### Layer 3 — Static Checks (no services needed)
+**CI equivalent:** `clang-format` and `spell-check` jobs
+
+**Code format (changed files vs main):**
+```bash
+./scripts/code_format.sh --check --base origin/main
+# Auto-fix:
+./scripts/code_format.sh --base origin/main
+```
+
+**Spell check:**
+```bash
+# Requires typos tool: cargo install typos-cli
+typos
+```
+
+**Pre-commit (runs all hooks):**
+```bash
+pip install pre-commit
+pre-commit run --all-files
+# Or just on staged files:
+pre-commit run
+```
+
+## Build Configurations (from CI)
+
+If the user needs to build first, here are the CI-equivalent cmake flags:
+
+**Standard build with coverage (mirrors `build` job):**
+```bash
+mkdir build && cd build
+cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_CXL=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON -DENABLE_ASAN=ON -DCMAKE_BUILD_TYPE=Debug
+cmake --build .
+sudo cmake --install .
+```
+
+**All features ON (mirrors `build-flags` job):**
+```bash
+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
+cmake --build .
+sudo cmake --install .
+```
+
+**Transfer engine only:**
+```bash
+cd mooncake-transfer-engine
+mkdir build && cd build
+cmake -G Ninja .. -DUSE_ETCD=OFF -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON
+cmake --build .
+```
+
+## Workflow: Diagnosing and Running Tests
+
+### Step 1 — Understand what the user wants
+
+Ask (or infer from context):
+- All tests, or a specific subset?
+- Did a specific CI job fail? Which one?
+- Is the build already done, or do they need to build first?
+
+### Step 2 — Check and Fix Prerequisites
+
+**One-command setup** — this script checks all prerequisites and auto-fixes issues:
+
+```bash
+bash .claude/skills/mooncake-ci-local/scripts/check-prerequisites.sh
+```
+
+**What it checks:**
+1. ✓ Build directory exists
+2. ✓ mooncake package installed (auto-installs via cmake --install if missing)
+3. ✓ ctest available
+4. ✓ Restarts all services (etcd, metadata server) in clean state
+5. ✓ Verifies all services are healthy
+
+**If you need to build first:**
+```bash
+mkdir build && cd build
+cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CXL=ON -DSTORE_USE_ETCD=ON -DCMAKE_BUILD_TYPE=Debug
+cmake --build .
+sudo cmake --install .
+```
+
+**If script fails:**
+- Build issues: See "Build Configurations" section below
+- mooncake install fails: Try `pip install mooncake-wheel/dist/*.whl` manually
+- etcd install fails: Download from https://github.com/etcd-io/etcd/releases
+
+### Step 3 — Run and report
+
+Run the relevant test layer. On failure:
+1. Show the exact error message
+2. Check if it's a service/env issue (most common) vs a real test failure
+3. Suggest the fix (see common issues below)
+
+## Common Local Test Issues
+
+**"mooncake_master found in /usr/local/bin" error in run_tests.sh:**
+The test expects mooncake_master to come from the Python package, not a system install.
+```bash
+# Remove the system-installed binary:
+sudo rm /usr/local/bin/mooncake_master
+# Or use the wheel-installed one:
+pip install mooncake-wheel/dist/*.whl
+```
+
+**etcd port conflict:**
+```bash
+pkill etcd && sleep 1
+etcd --advertise-client-urls http://127.0.0.1:2379 --listen-client-urls http://127.0.0.1:2379 &
+```
+
+**Metadata server port conflict:**
+```bash
+pkill -f bootstrap_server.py
+pkill -f mooncake_http_metadata_server
+```
+
+**Tests hang (master not responding):**
+```bash
+pkill mooncake_master
+sleep 2
+mooncake_master --default_kv_lease_ttl=500 &
+sleep 1
+```
+
+**torch/numpy not installed for tensor tests:**
+```bash
+pip install torch numpy safetensors packaging
+```
+
+**ctest shows no tests found:**
+```bash
+# Rebuild with unit tests enabled:
+cd build
+cmake .. -DBUILD_UNIT_TESTS=ON
+cmake --build .
+```
+
+## Quick One-Liners
+
+```bash
+# Run ALL C++ tests (after building with etcd + metadata server running):
+# Note: full suite takes 5-15 minutes depending on hardware
+cd build && MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 ctest -j --output-on-failure
+
+# Run only fast tests (skip slow integration tests):
+cd build && MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 ctest -j --output-on-failure --exclude-regex "etcd|ha_test|redis"
+
+# Run ALL Python tests:
+mooncake_http_metadata_server --port 8080 & sleep 1 && cd mooncake-wheel/tests && MC_METADATA_SERVER=http://127.0.0.1:8080/metadata MC_FORCE_TCP=true bash ../../scripts/run_tests.sh
+
+# Check code format (changed files only):
+./scripts/code_format.sh --check --base origin/main
+
+# Full pre-commit check:
+pre-commit run --all-files
+```

.claude/skills/mooncake-ci-local/scripts/check-prerequisites.sh

@@ -0,0 +1,101 @@
+#!/bin/bash
+# Mooncake CI Local Test Prerequisites Check
+# Usage: bash check-prerequisites.sh
+# This script checks and auto-fixes all prerequisites for running Mooncake CI tests locally.
+
+set -e
+
+RED='\033[0;31m'
+GREEN='\033[0;32m'
+YELLOW='\033[1;33m'
+NC='\033[0m'
+
+echo "🔍 Checking Mooncake CI test prerequisites..."
+
+# 1. Check build directory
+if [ ! -f build/CMakeCache.txt ]; then
+  echo -e "${RED}✗ Build directory not found or not built${NC}"
+  echo "  → Run: mkdir build && cd build && cmake .. && cmake --build ."
+  exit 1
+fi
+echo -e "${GREEN}✓ Build exists${NC}"
+
+# 2. Check mooncake installation
+if ! python -c "import mooncake" 2>/dev/null; then
+  echo -e "${RED}✗ mooncake package not installed${NC}"
+  echo "  → Fixing: Installing mooncake package..."
+  cd build && sudo cmake --install . && cd - >/dev/null
+  if ! python -c "import mooncake" 2>/dev/null; then
+    echo "  → Alternative: pip install mooncake-wheel/dist/*.whl"
+    exit 1
+  fi
+  echo -e "${GREEN}✓ mooncake package installed${NC}"
+else
+  echo -e "${GREEN}✓ mooncake package already installed${NC}"
+fi
+
+# 3. Check ctest availability
+if ! command -v ctest &> /dev/null; then
+  echo -e "${RED}✗ ctest not found${NC}"
+  exit 1
+fi
+echo -e "${GREEN}✓ ctest available${NC}"
+
+# 4. Kill and restart services (safest approach for local testing)
+echo -e "\n${YELLOW}Cleaning up and restarting services...${NC}"
+pkill -f "^etcd" || true
+pkill -f bootstrap_server.py || true
+pkill -f mooncake_http_metadata_server || true
+sleep 1
+
+# 5. Start etcd
+if ! command -v etcd &> /dev/null; then
+  echo -e "${YELLOW}⚠ etcd not found, installing...${NC}"
+  ETCD_VER=v3.6.1
+  OS=$(uname -s | tr '[:upper:]' '[:lower:]')
+  ARCH=$(uname -m)
+  [ "$ARCH" = "x86_64" ] && ARCH="amd64"
+  DOWNLOAD_URL="https://github.com/etcd-io/etcd/releases/download/${ETCD_VER}/etcd-${ETCD_VER}-${OS}-${ARCH}.tar.gz"
+  echo "  Downloading from: $DOWNLOAD_URL"
+  cd /tmp
+  wget -q "$DOWNLOAD_URL" && tar xzf "etcd-${ETCD_VER}-${OS}-${ARCH}.tar.gz" && \
+    sudo mv "etcd-${ETCD_VER}-${OS}-${ARCH}"/etcd* /usr/local/bin/
+  cd - >/dev/null
+  echo -e "${GREEN}✓ etcd installed${NC}"
+fi
+
+etcd --advertise-client-urls http://127.0.0.1:2379 --listen-client-urls http://127.0.0.1:2379 >/dev/null 2>&1 &
+ETCD_PID=$!
+sleep 2
+if ! etcdctl --endpoints=http://127.0.0.1:2379 endpoint health &>/dev/null; then
+  echo -e "${RED}✗ etcd failed to start${NC}"
+  kill $ETCD_PID 2>/dev/null || true
+  exit 1
+fi
+echo -e "${GREEN}✓ etcd running (PID: $ETCD_PID)${NC}"
+
+# 6. Start HTTP metadata server
+if [ -f "mooncake-transfer-engine/example/http-metadata-server-python/bootstrap_server.py" ]; then
+  cd mooncake-transfer-engine/example/http-metadata-server-python
+  pip install -q aiohttp 2>/dev/null || true
+  python ./bootstrap_server.py >/dev/null 2>&1 &
+  METADATA_PID=$!
+  cd - >/dev/null
+  sleep 1
+  if curl -s http://127.0.0.1:8080/metadata > /dev/null 2>&1; then
+    echo -e "${GREEN}✓ HTTP Metadata server running (PID: $METADATA_PID)${NC}"
+  else
+    echo -e "${RED}✗ HTTP Metadata server failed to start${NC}"
+    kill $METADATA_PID $ETCD_PID 2>/dev/null || true
+    exit 1
+  fi
+else
+  echo -e "${YELLOW}⚠ Metadata server script not found, skipping${NC}"
+fi
+
+echo -e "\n${GREEN}✅ All prerequisites ready!${NC}"
+echo "Service PIDs: etcd=$ETCD_PID"
+[ -n "$METADATA_PID" ] && echo "Metadata server PID: $METADATA_PID"
+echo -e "\n${YELLOW}To kill services:${NC}"
+echo "  pkill -f '^etcd'"
+echo "  pkill -f bootstrap_server"

.claude/skills/mooncake-troubleshoot/SKILL.md

@@ -0,0 +1,366 @@
+---
+name: mooncake-troubleshoot
+description: Automatically diagnose Mooncake deployment and runtime issues. Checks services (mooncake_master, metadata server), RDMA devices, environment variables, connectivity, memory limits, and analyzes logs for common error patterns. Use when Mooncake deployment fails, services won't start, connections fail, or you encounter runtime errors like "Error from etcd client", "No matched device found", "Failed to register memory", "NO_AVAILABLE_HANDLE", or any RDMA/networking issues. Also use when user asks to troubleshoot, debug, diagnose, or fix Mooncake problems.
+---
+
+# Mooncake Deployment Troubleshooting
+
+You are a Mooncake deployment troubleshooting specialist. Your job is to systematically diagnose issues and provide actionable solutions based on the comprehensive troubleshooting knowledge from Mooncake documentation.
+
+## Diagnostic Strategy
+
+Run checks systematically, reporting findings as you go. Start with simple checks (services, connectivity) before diving into complex issues (RDMA, memory registration).
+
+### 1. Service Status Check
+
+Check if critical services are running:
+
+```bash
+# Check mooncake_master
+ps aux | grep mooncake_master | grep -v grep
+
+# Check port usage
+netstat -tuln | grep -E '(50051|8080|2379|9003)'
+
+# If using etcd
+ps aux | grep etcd | grep -v grep
+```
+
+**Common issues:**
+- `bind address already in use` → Port conflict, use different port with `--rpc_port`
+- Master not running → Check startup logs for errors
+
+### 2. Metadata Server Connectivity
+
+The metadata server is critical for node discovery and coordination.
+
+```bash
+# Test etcd connectivity
+curl -s http://127.0.0.1:2379/version
+
+# Or test custom metadata server
+curl -s $MC_METADATA_SERVER
+
+# Check for proxy interference
+echo "http_proxy: $http_proxy"
+echo "https_proxy: $https_proxy"
+```
+
+**Common issues:**
+- `Error from etcd client` → Metadata server unreachable
+  - **Fix:** Ensure etcd is bound to `0.0.0.0` not `127.0.0.1`:
+    ```bash
+    etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://<your_ip>:2379
+    ```
+  - **Fix:** Disable HTTP proxy:
+    ```bash
+    unset http_proxy https_proxy
+    ```
+
+### 3. Environment Variables Check
+
+Verify critical environment variables are set correctly:
+
+```bash
+# Display all MC_* variables
+env | grep ^MC_
+
+# Key variables to check:
+echo "MC_METADATA_SERVER: $MC_METADATA_SERVER"
+echo "MC_FORCE_TCP: $MC_FORCE_TCP"
+echo "MC_LOG_LEVEL: $MC_LOG_LEVEL"
+echo "MC_YLT_LOG_LEVEL: $MC_YLT_LOG_LEVEL"
+echo "MC_MS_AUTO_DISC: $MC_MS_AUTO_DISC"
+echo "MC_MS_FILTERS: $MC_MS_FILTERS"
+echo "MC_GID_INDEX: $MC_GID_INDEX"
+echo "MC_MTU: $MC_MTU"
+echo "MC_IB_PORT: $MC_IB_PORT"
+echo "MC_ENABLE_DEST_DEVICE_AFFINITY: $MC_ENABLE_DEST_DEVICE_AFFINITY"
+```
+
+**Key variables:**
+- `MC_METADATA_SERVER` - Metadata server URL (required)
+- `MC_FORCE_TCP=true` - Force TCP for testing without RDMA
+- `MC_LOG_LEVEL=0` - Enable verbose logging (0=INFO, 1=WARNING, 2=ERROR)
+- `MC_YLT_LOG_LEVEL=debug` - yalantinglibs log level
+- `MC_MS_AUTO_DISC=1` - Enable topology auto-discovery (default)
+- `MC_MS_FILTERS` - Filter specific RDMA devices (e.g., "mlx5_1,mlx5_2")
+- `MC_GID_INDEX` - RDMA GID index (set if GID is all zeros)
+- `MC_MTU` - RDMA MTU size
+- `MC_ENABLE_DEST_DEVICE_AFFINITY=1` - Reduce QP creation (fix "Failed to create QP")
+
+### 4. RDMA Device Check
+
+Only run if RDMA is being used (skip if `MC_FORCE_TCP=true`):
+
+```bash
+# List RDMA devices
+ibv_devices
+
+# Check device details and status
+ibv_devinfo
+
+# Check for ACTIVE ports
+ibv_devinfo | grep -A 10 "state:"
+
+# Check GID addresses (should NOT be all zeros)
+ibv_devinfo | grep -A 20 "GID"
+
+# Check peer memory modules
+lsmod | grep peer_mem
+lsmod | grep nvidia_peer_mem
+
+# Check QP count (if "Failed to create QP" error)
+rdma resource show qp
+```
+
+**Common issues:**
+- `No matched device found` → RDMA device name in config doesn't exist
+  - **Fix:** Use `ibv_devices` to get correct device names
+- `Device XXX port not active` → RDMA port not in ACTIVE state
+  - **Fix:** Check cable connections, verify with `ibv_devinfo | grep state`
+  - **Fix:** Try different port with `MC_IB_PORT` environment variable
+- GID all zeros → Wrong GID index
+  - **Fix:** Set `MC_GID_INDEX=1` (or 2, 3 depending on network)
+- `Failed to create QP: Cannot allocate memory` → Too many QPs created
+  - **Fix:** Set `MC_ENABLE_DEST_DEVICE_AFFINITY=1`
+
+### 5. Memory and Resource Limits
+
+Check system limits that affect RDMA memory registration:
+
+```bash
+# Check ulimits
+ulimit -a
+
+# Focus on max locked memory
+ulimit -l
+
+# Check RDMA device memory limits
+ibv_devinfo -v | grep max_mr_size
+
+# Check dmesg for memory errors
+dmesg -T | tail -50 | grep -i "out of mr size"
+```
+
+**Common issues:**
+- `Failed to register memory: Input/output error` → Memory registration limit exceeded
+  - **Diagnostic:** Check `max_mr_size` with `ibv_devinfo -v`
+  - **Fix:** Reduce memory allocation or split into smaller chunks
+- Cannot allocate memory → ulimit restriction
+  - **Fix:** Set unlimited locked memory:
+    ```bash
+    ulimit -l unlimited
+    ```
+  - **Permanent fix:** Add to `/etc/security/limits.conf`:
+    ```
+    * soft memlock unlimited
+    * hard memlock unlimited
+    ```
+
+### 6. Network Connectivity
+
+Test connectivity between nodes:
+
+```bash
+# Test basic RDMA connectivity
+ib_write_bw -d <device_name> -R
+
+# On peer node:
+ib_write_bw -d <device_name> -R <server_ip>
+
+# Test GPU Direct RDMA (if CUDA enabled)
+ib_write_bw -d <device_name> -R -x gdr
+
+# On peer node:
+ib_write_bw -d <device_name> -R -x gdr <server_ip>
+
+# Test DNS resolution
+nslookup <connectable_name>
+ping <connectable_name>
+```
+
+**Common issues:**
+- `connection refused` → Incorrect `connectable_name` or `rpc_port`
+  - **Fix:** Ensure `connectable_name` is NOT loopback (127.0.0.1/localhost)
+  - **Fix:** Use actual LAN/WAN IP or valid hostname
+- `Failed to exchange handshake` → RDMA connection setup failure
+  - **Fix:** Verify MTU matches: set `MC_MTU` environment variable
+  - **Fix:** Verify GID is valid (not all zeros)
+  - **Fix:** Test with `ib_send_bw` between nodes first
+
+### 7. Log Analysis
+
+Search logs for common error patterns and their meanings:
+
+**Metadata/Connectivity Errors:**
+- `Error from etcd client` → Cannot connect to metadata server
+- `ERR_METADATA` → Metadata server communication failed
+- `ERR_DNS` → Invalid `local_server_name` (not valid DNS/IP)
+
+**RDMA Errors:**
+- `No matched device found` → RDMA device name doesn't exist
+- `Device XXX port not active` → RDMA port not in ACTIVE state
+- `Failed to exchange handshake description` → RDMA handshake failed
+- `Failed to modify QP to RTR, check mtu, gid, peer lid, peer qp num` → MTU/GID mismatch
+- `Failed to register memory` → Memory registration limit exceeded
+- `Failed to create QP` → Too many QPs, enable `MC_ENABLE_DEST_DEVICE_AFFINITY=1`
+- `Worker: Process failed for slice` → Network instability
+- `work request flushed error` → Cascading error (find first error)
+
+**Store Errors:**
+- `NO_AVAILABLE_HANDLE` (-200) → Memory pool exhausted
+  - **Fix:** Increase `global_segment_size` in setup
+  - **Fix:** Check eviction is working (look for eviction logs)
+- `LEASE_EXPIRED` (-707) → Lease expired during transfer
+  - **Fix:** Increase `default_kv_lease_ttl` in master startup
+- `OBJECT_NOT_FOUND` (-704) → Object doesn't exist
+- `SEGMENT_NOT_FOUND` (-101) → No available segments
+- `Failed to get description of XXX` → Segment name mismatch
+  - **Fix:** Ensure segment name matches `local_hostname` from peer
+
+**Port/Service Errors:**
+- `bind address already in use` → Port conflict
+  - **Fix:** Use different port: `--rpc_port=50052`
+
+### 8. Configuration Validation
+
+Verify configuration is correct:
+
+```bash
+# Check connectable_name is not loopback
+hostname -I
+
+# Verify master startup flags
+ps aux | grep mooncake_master
+
+# Check if using correct protocol
+env | grep MC_FORCE_TCP
+```
+
+**Critical checks:**
+- `connectable_name` must be non-loopback IP or valid hostname
+- MTU and GID configurations must match network environment
+- RDMA device names must exist on the machine
+- Ports must not be in use by other services
+
+## Error Code Quick Reference
+
+### Transfer Engine Error Codes
+
+| Code | Name | Meaning | Fix |
+|------|------|---------|-----|
+| 0 | Success | Normal execution | - |
+| -12 | ERR_ADDRESS_NOT_REGISTERED | Memory not registered | Register memory before use |
+| -14 | ERR_DEVICE_NOT_FOUND | RDMA device not found | Check device name with `ibv_devices` |
+| -16 | ERR_DNS | Invalid local_server_name | Use valid IP/hostname |
+| -19 | ERR_REJECT_HANDSHAKE | Peer rejected handshake | Check peer logs for reason |
+| -20 | ERR_METADATA | Metadata server unreachable | Check etcd/HTTP server |
+
+### Store Error Codes
+
+| Code | Name | Meaning | Fix |
+|------|------|---------|-----|
+| 0 | Success | Operation successful | - |
+| -200 | NO_AVAILABLE_HANDLE | Memory pool exhausted | Increase segment size |
+| -707 | LEASE_EXPIRED | Lease expired | Increase lease TTL |
+| -704 | OBJECT_NOT_FOUND | Object doesn't exist | Check object key |
+| -101 | SEGMENT_NOT_FOUND | No available segments | Check segment registration |
+| -900 | RPC_FAIL | RPC failed | Check network/master |
+| -1000 | ETCD_OPERATION_ERROR | etcd operation failed | Check etcd status |
+
+## Output Format
+
+Provide a structured diagnostic report:
+
+```
+🔍 MOONCAKE DEPLOYMENT DIAGNOSTICS
+==================================
+
+✅ PASSED CHECKS:
+- Service status: mooncake_master running on port 50051
+- Metadata server: etcd accessible at http://127.0.0.1:2379
+- Environment: MC_METADATA_SERVER set correctly
+- [other passing checks]
+
+❌ FAILED CHECKS:
+- RDMA device: mlx5_0 port not ACTIVE (state: PORT_DOWN)
+- Memory limits: max locked memory is 64KB (too low)
+- [other failures with specific error messages]
+
+⚠️  WARNINGS:
+- GID index not set, may cause connection issues
+- HTTP proxy variables set, may interfere with metadata server
+- [other potential issues]
+
+🔧 RECOMMENDED FIXES:
+
+1. Fix RDMA port status:
+   - Check physical cable connections
+   - Verify driver configuration
+   - Command: ibv_devinfo | grep -A 10 "state:"
+
+2. Increase memory limits:
+   ulimit -l unlimited
+   # Or permanently in /etc/security/limits.conf:
+   * soft memlock unlimited
+   * hard memlock unlimited
+
+3. Set GID index:
+   export MC_GID_INDEX=1
+
+4. Disable HTTP proxy:
+   unset http_proxy https_proxy
+
+📋 SUMMARY:
+[Brief 2-3 sentence conclusion about deployment health and next steps]
+```
+
+## Troubleshooting Workflow
+
+1. **Start simple**: Check services and basic connectivity first
+2. **Read logs carefully**: First error is usually root cause (subsequent errors cascade)
+3. **Test incrementally**: Use `MC_FORCE_TCP=true` to isolate RDMA issues
+4. **Verify basics**: Check connectable_name, ports, env vars before deep diving
+5. **Use diagnostic tools**: ibv_devices, ibv_devinfo, ib_write_bw, curl
+6. **Reference documentation**: Check error codes and troubleshooting guide
+
+## Quick Fix Commands
+
+**Start metadata server properly:**
+```bash
+etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://<your_ip>:2379
+```
+
+**Enable verbose logging:**
+```bash
+export MC_LOG_LEVEL=0
+export MC_YLT_LOG_LEVEL=debug
+```
+
+**Force TCP mode for testing:**
+```bash
+export MC_FORCE_TCP=true
+```
+
+**Fix memory limits:**
+```bash
+ulimit -l unlimited
+```
+
+**Fix too many QPs:**
+```bash
+export MC_ENABLE_DEST_DEVICE_AFFINITY=1
+```
+
+**Fix GID issues:**
+```bash
+export MC_GID_INDEX=1  # or 2, 3 depending on network
+```
+
+**Use different port:**
+```bash
+mooncake_master --rpc_port=50052
+```
+
+Now execute the diagnostic checks systematically and provide the structured report.

.devcontainer/Dockerfile

@@ -23,11 +23,20 @@ RUN apt-get install -y libibverbs-dev \
     libhiredis-dev \
     libyaml-cpp-dev \
     libjemalloc-dev \
+    libzstd-dev \
+    libmsgpack-dev \
+    libgflags-dev \
     pkg-config \
     patchelf
 
-RUN wget https://go.dev/dl/go1.22.12.linux-amd64.tar.gz \
-    && tar -C /usr/local -xzf go1.22.12.linux-amd64.tar.gz
+RUN GO_VERSION="1.23.8" && \
+    ARCH=$(uname -m) && \
+    if [ "$ARCH" = "aarch64" ]; then GOARCH="arm64"; \
+    elif [ "$ARCH" = "x86_64" ]; then GOARCH="amd64"; \
+    else echo "Unsupported architecture: $ARCH" && exit 1; fi && \
+    wget https://go.dev/dl/go${GO_VERSION}.linux-${GOARCH}.tar.gz \
+    && tar -C /usr/local -xzf go${GO_VERSION}.linux-${GOARCH}.tar.gz \
+    && rm go${GO_VERSION}.linux-${GOARCH}.tar.gz
 
 RUN git clone https://github.com/alibaba/yalantinglibs.git \
     && cd yalantinglibs \

.github/CODEOWNERS

@@ -8,14 +8,16 @@
 
 .github @stmatengss @ykwd @Ann-1024 @luketong777
 /docs @ShangmingCai @stmatengss @ykwd
-/mooncake-ep @UNIDY2002 @ympcMark
-/mooncake-integration/ep @UNIDY2002 @ympcMark
-/mooncake-integration/transfer_engine @ShangmingCai @alogfans 
+/mooncake-ep @UNIDY2002 @ympcMark @yuechen-sys
+/mooncake-integration/transfer_engine @ShangmingCai @alogfans
 /mooncake-integration/store @ykwd @stmatengss
-/mooncake-pg @UNIDY2002 @ympcMark 
+/mooncake-pg @UNIDY2002 @ympcMark @yuechen-sys
 /mooncake-store @ykwd @stmatengss @XucSh @YiXR
+/mooncake-store/*/ha/ @Libotry @YiXR @00fish0
 /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,14 +6,25 @@ on:
   pull_request:
     branches: [ "main" ]
     types: [opened, synchronize, reopened, labeled]
+  workflow_dispatch: {}
 
+permissions:
+  contents: read
+
+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, 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:
@@ -24,6 +35,8 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
+      with:
+        persist-credentials: false
 
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v5
@@ -54,10 +67,10 @@ jobs:
         method: 'network'
         sub-packages: '["nvcc"]'
 
-    - name: Install coverage tools
+    - name: Install coverage tools and build utilities
       run: |
         sudo apt-get update
-        sudo apt-get install -y lcov gcovr
+        sudo apt-get install -y lcov gcovr ninja-build
 
     - name: Set up coverage compilation flags
       run: |
@@ -87,14 +100,14 @@ jobs:
         sudo bash -x dependencies.sh -y
         mkdir build
         cd build
-        cmake .. -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
       run: |
         cd build
-        make -j4
-        sudo make install
+        cmake --build .
+        sudo cmake --install .
       shell: bash
 
     - name: Build nvlink_allocator.so
@@ -112,12 +125,42 @@ jobs:
         python ./bootstrap_server.py &
       shell: bash
 
+    - name: Run Go store binding integration tests
+      run: |
+        $GITHUB_WORKSPACE/build/mooncake-store/src/mooncake_master \
+          --eviction_high_watermark_ratio=0.95 \
+          --cluster_id=ci_go_test_cluster \
+          --port 50051 &
+        MASTER_PID=$!
+        sleep 3
+        cd mooncake-store/go
+        export LD_LIBRARY_PATH=$GITHUB_WORKSPACE/build/mooncake-common:$GITHUB_WORKSPACE/build/mooncake-store/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base:$GITHUB_WORKSPACE/build/mooncake-common/etcd
+        export CGO_ENABLED=1
+        export CGO_CFLAGS="-I$GITHUB_WORKSPACE/mooncake-store/include -I$GITHUB_WORKSPACE/mooncake-transfer-engine/include"
+        export CGO_LDFLAGS="-L$GITHUB_WORKSPACE/build/mooncake-store/src -L$GITHUB_WORKSPACE/build/mooncake-store/src/cachelib_memory_allocator -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base -L$GITHUB_WORKSPACE/build/mooncake-common -L$GITHUB_WORKSPACE/build/mooncake-common/etcd -lmooncake_store -lcachelib_memory_allocator -ltransfer_engine -lbase -lasio -letcd_wrapper -lstdc++ -lnuma -lglog -lgflags -libverbs -ljsoncpp -lzstd -lcurl -luring -lasan -lm -lgcov"
+        # Link cudart if CUDA is available (needed for D2H staging in mooncake_store)
+        if [ -d /usr/local/cuda/lib64 ]; then export CGO_LDFLAGS="$CGO_LDFLAGS -L/usr/local/cuda/lib64 -lcudart"; fi
+        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
+
     - name: Test (in build env) with coverage
       run: |
         cd build
         export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
         ldconfig -v || echo "always continue"
-        MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 make test -j ARGS="-V"
+        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
@@ -193,14 +236,20 @@ jobs:
         path: mooncake-wheel/dist-py${{ steps.generate_tag_build.outputs.python_version_tag }}/*.whl
 
   build-musa:
+    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:
     - uses: actions/checkout@v4
+      with:
+        persist-credentials: false
 
     - name: Mark repository as safe
       run: git config --global --add safe.directory $GITHUB_WORKSPACE
@@ -209,22 +258,29 @@ jobs:
     - name: Configure project
       run: |
         apt update -y
+        apt install -y ninja-build
         bash -x dependencies.sh -y
         mkdir build
         cd build
-        cmake .. -DUSE_MUSA=ON -DUSE_MNNVL=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON -DUSE_CXL=ON -DUSE_TCP=ON -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF
+        cmake -G Ninja .. -DUSE_MUSA=ON -DUSE_MNNVL=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON -DUSE_CXL=ON -DUSE_TCP=ON -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF
       shell: bash
 
     - name: Build project
       run: |
         cd build
         source ~/.bashrc
-        make -j
-        make install
+        cmake --build .
+        cmake --install .
       shell: bash
 
   test-wheel-ubuntu:
-    needs: build-flags
+    needs: [spell-check, clang-format, build-flags]
+    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:
       matrix:
         ubuntu-version: [ubuntu-22.04, ubuntu-24.04]
@@ -232,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
@@ -287,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
@@ -337,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
@@ -347,15 +425,12 @@ jobs:
         kill $SERVER_PID 2>/dev/null || true
         wait $SERVER_PID 2>/dev/null || true
 
-    - name: Test Mooncake EP Backend (CPU Only)
+    - 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
-        # Disable these tests in CI as they fail occasionally.
-        # python -m unittest mooncake-wheel.tests.test_mooncake_backend_elastic
-        # python -m unittest mooncake-wheel.tests.test_mooncake_backend_p2p_cpu
       shell: bash
 
     - name: Test Safetensor Functions
@@ -366,10 +441,14 @@ jobs:
       shell: bash
 
   build-flags:
+    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:
@@ -377,12 +456,13 @@ jobs:
     env:
       CI: "true"
       BUILD_WITH_EP: "1"
-      EP_TORCH_VERSIONS: "2.9.0;2.9.1;2.10.0"
       TORCH_CUDA_ARCH_LIST: "8.0;9.0"
       SCCACHE_GHA_ENABLED: "true"
 
     steps:
     - uses: actions/checkout@v4
+      with:
+        persist-credentials: false
 
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v5
@@ -423,10 +503,14 @@ jobs:
     - name: Install dependencies
       run: |
         sudo apt update -y
+        sudo apt install -y ninja-build
         sudo bash -x dependencies.sh -y
         df -h
       shell: bash
 
+    - name: Install Rust toolchain
+      uses: dtolnay/rust-toolchain@stable
+
     - name: Build transfer engine only
       run: |
         cd mooncake-transfer-engine
@@ -434,9 +518,9 @@ 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 .. -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"
-        make -j4
-        sudo make install
+        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
       shell: bash
 
@@ -444,7 +528,7 @@ jobs:
       run: |
         mkdir build
         cd build
-        cmake .. -DUSE_ETCD=ON -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_STORE=ON -DWITH_P2P_STORE=ON -DWITH_EP=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
 
@@ -453,32 +537,39 @@ jobs:
         export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
         export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
         cd build
-        make -j4
-        sudo make install
+        cmake --build .
+        sudo cmake --install .
         df -h
       shell: bash
 
     - name: Configure project with unit tests and examples
       run: |
         cd build
-        cmake .. -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: |
         export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
         export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
         cd build
-        make -j4
-        sudo make install
+        cmake --build .
+        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 .. -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF -DUSE_HTTP=ON -DENABLE_SCCACHE=ON -DUSE_CXL=ON
+        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
@@ -486,8 +577,8 @@ jobs:
         export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
         export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
         cd build
-        make -j4
-        sudo make install
+        cmake --build .
+        sudo cmake --install .
       shell: bash
 
     - name: Build nvlink_allocator.so
@@ -520,13 +611,19 @@ jobs:
 
   build-docker:
     name: Build Docker Image
+    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:
     - uses: actions/checkout@v4
+      with:
+        persist-credentials: false
 
     - name: Set up Docker Buildx
       uses: docker/setup-buildx-action@v2
@@ -542,12 +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: Checkout Actions Repository
       uses: actions/checkout@v4
+      with:
+        persist-credentials: false
     - name: Spell Check Repo
       uses: crate-ci/typos@v1.30.2
 
@@ -555,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
@@ -563,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: |
@@ -597,3 +699,92 @@ jobs:
         echo "Comparing against: ${BASE_REF}"
         ./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

@@ -0,0 +1,364 @@
+name: 'CI Test on ASCEND Platform'
+
+on:
+  workflow_call:
+    inputs:
+      checkout_ref:
+        description: 'Git ref to checkout (PR head SHA for pull_request_target)'
+        required: false
+        type: string
+
+jobs:
+  build-and-test:
+    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
+                    --device /dev/davinci_manager --device /dev/devmm_svm --device /dev/hisi_hdc --ulimit nproc=65535:65535
+      env:
+        GITHUB_ACTIONS: "true"
+        LD_PRELOAD: "/usr/lib64/libjemalloc.so.2:"
+      volumes:
+        - /usr/local/dcmi:/usr/local/dcmi
+        - /usr/local/Ascend/driver/:/usr/local/Ascend/driver/
+        - /etc/ascend_install.info:/etc/ascend_install.info
+        - /etc/hccn.conf:/etc/hccn.conf
+
+    steps:
+    - name: Configure GitHub fetch defaults
+      shell: bash
+      run: |
+        git config --global protocol.version 2
+        git config --global http.version HTTP/1.1
+        git config --global http.lowSpeedLimit 1024
+        git config --global http.lowSpeedTime 30
+
+    - name: Checkout code
+      id: checkout_code
+      continue-on-error: true
+      uses: actions/checkout@v4
+      with:
+        ref: ${{ inputs.checkout_ref || github.sha }}
+        fetch-depth: 1
+        persist-credentials: false
+
+    - name: Retry checkout via GitHub mirror
+      if: steps.checkout_code.outcome == 'failure'
+      shell: bash
+      env:
+        ASCEND_GITHUB_MIRROR_URLS: ${{ vars.ASCEND_GITHUB_MIRROR_URLS }}
+        CHECKOUT_REF: ${{ inputs.checkout_ref || github.sha }}
+      run: |
+        set -euo pipefail
+
+        if [ -z "${ASCEND_GITHUB_MIRROR_URLS:-}" ]; then
+          echo "Checkout from GitHub failed and ASCEND_GITHUB_MIRROR_URLS is not set"
+          exit 1
+        fi
+
+        normalize_base() {
+          local base="$1"
+          base="${base#${base%%[![:space:]]*}}"
+          base="${base%${base##*[![:space:]]}}"
+          [ -n "$base" ] || return 1
+          [ "$base" != "https://github.com/" ] && base="${base%/}/"
+          printf '%s\n' "$base"
+        }
+
+        candidates=()
+        while IFS= read -r raw; do
+          base="$(normalize_base "$raw" || true)"
+          [ -n "$base" ] || continue
+          [ "$base" = "https://github.com/" ] && continue
+          candidates+=("$base")
+        done < <(printf '%s\n' "$ASCEND_GITHUB_MIRROR_URLS" | tr ',;' '\n')
+
+        if [ ${#candidates[@]} -eq 0 ]; then
+          echo "Checkout from GitHub failed and no valid mirror candidates were configured"
+          exit 1
+        fi
+
+        workdir="${GITHUB_WORKSPACE}"
+        git config --global --add safe.directory "$workdir"
+
+        for base in "${candidates[@]}"; do
+          mirror_url="${base}https://github.com/${GITHUB_REPOSITORY}.git"
+          echo "Retrying checkout with ${mirror_url}"
+
+          find "$workdir" -mindepth 1 -maxdepth 1 -exec rm -rf {} +
+          git init "$workdir"
+          git -C "$workdir" remote add origin "$mirror_url"
+
+          if git -C "$workdir" fetch --depth=1 origin "$CHECKOUT_REF" && \
+             git -C "$workdir" checkout --force --detach FETCH_HEAD; then
+            echo "Mirror checkout succeeded via ${base}"
+            exit 0
+          fi
+
+          echo "Mirror checkout failed via ${base}"
+          rm -rf "$workdir/.git"
+        done
+
+        echo "Direct GitHub checkout failed and all mirror retries failed"
+        exit 1
+
+    - name: Configure CMake
+      shell: bash
+      env:
+        ASCEND_GITHUB_MIRROR_URLS: ${{ vars.ASCEND_GITHUB_MIRROR_URLS }}
+      run: |
+        source /usr/local/Ascend/cann-9.0.0/set_env.sh
+        pwd
+
+        submodule_updated=false
+        if git submodule update --init --recursive; then
+          submodule_updated=true
+        elif [ -n "${ASCEND_GITHUB_MIRROR_URLS:-}" ]; then
+          normalize_base() {
+            local base="$1"
+            base="${base#${base%%[![:space:]]*}}"
+            base="${base%${base##*[![:space:]]}}"
+            [ -n "$base" ] || return 1
+            [ "$base" != "https://github.com/" ] && base="${base%/}/"
+            printf '%s\n' "$base"
+          }
+
+          while IFS= read -r raw; do
+            base="$(normalize_base "$raw" || true)"
+            [ -n "$base" ] || continue
+            [ "$base" = "https://github.com/" ] && continue
+
+            echo "Retrying submodule update with ${base}"
+            if git -c url."${base}https://github.com/".insteadOf=https://github.com/ \
+              submodule update --init --recursive; then
+              submodule_updated=true
+              break
+            fi
+          done < <(printf '%s\n' "$ASCEND_GITHUB_MIRROR_URLS" | tr ',;' '\n')
+        fi
+
+        if [ "$submodule_updated" != true ]; then
+          if [ ! -d "extern/pybind11" ] || [ -z "$(ls -A 'extern/pybind11' 2>/dev/null)" ]; then
+            echo "git submodule update failed (mirrors also exhausted), trying to cp pybind11..."
+            if [ -d "../pybind11" ]; then
+              cp -r ../pybind11 extern/
+            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
+        mkdir -p build
+        cd build
+
+        cmake .. \
+          -DUSE_ASCEND_DIRECT=ON \
+          -DBUILD_EXAMPLES=OFF \
+          -DBUILD_UNIT_TESTS=OFF
+
+    - name: Build
+      shell: bash
+      run: |
+        source /usr/local/Ascend/cann-9.0.0/set_env.sh
+        echo "Building..."
+        cd build
+        cmake --build . -j$(nproc)
+        cmake --install .
+        echo "Mooncake installed successfully."
+
+    - name: Run Hixl Mooncake Store Test
+      shell: bash
+      run: |
+          source /usr/local/Ascend/cann-9.0.0/set_env.sh
+          set -e
+          export ASCEND_PROCESS_LOG_PATH=/tmp/hixl-test-log/
+          export ASCEND_GLOBAL_LOG_LEVEL=3
+          echo "=== Cloning Hixl repository ==="
+          cd ..
+          rm -rf hixl
+          git clone https://gitcode.com/cann/hixl.git
+          cd hixl/examples/third_parties/mooncake_store/python/
+
+          export LD_LIBRARY_PATH=/usr/local/lib:${LD_LIBRARY_PATH}
+          echo "=== Starting Mooncake Master ==="
+
+          # Find mooncake_master binary
+          MOONCAKE_MASTER=$(find /usr/local/bin /usr/bin -name "mooncake_master" -type f 2>/dev/null | head -1)
+          if [ -z "$MOONCAKE_MASTER" ]; then
+              # Try finding in build directory
+              MOONCAKE_MASTER=$(find $GITHUB_WORKSPACE/build -name "mooncake_master" -type f 2>/dev/null | head -1)
+          fi
+
+          if [ -z "$MOONCAKE_MASTER" ]; then
+              echo "Error: mooncake_master binary not found"
+              exit 1
+          fi
+
+          echo "Found mooncake_master at: $MOONCAKE_MASTER"
+
+          # Start Mooncake master in background
+          $MOONCAKE_MASTER \
+            --enable_http_metadata_server=true \
+            --http_metadata_server_host=0.0.0.0 \
+            --http_metadata_server_port=8080 \
+            > /tmp/mooncake_master.log 2>&1 &
+          MASTER_PID=$!
+          echo "Mooncake Master started with PID: $MASTER_PID"
+
+          # Wait for master to be ready
+          echo "Waiting for Mooncake Master to initialize..."
+          sleep 5
+
+          # 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
+              exit 1
+          fi
+
+          echo "Mooncake Master is running"
+          echo "=== Running Hixl Mooncake Store Tests ==="
+          # List of test cases to run
+          TEST_CASES=(
+              "batch_put_get_sample.py"
+              "batch_put_get_multi_buffers_sample.py"
+          )
+
+          # List of test scenarios (HCCL_INTRA_ROCE_ENABLE settings)
+          TEST_SCENARIOS=(
+              "HCCL_INTRA_ROCE_ENABLE=1"
+              "HCCL_INTRA_ROCE_ENABLE_UNSET"
+          )
+
+          # Track test results
+          FAILED_TESTS=()
+          PASSED_TESTS=()
+
+          # Run each test scenario
+          for scenario in "${TEST_SCENARIOS[@]}"; do
+              echo ""
+              echo "========================================="
+              echo "Running scenario: $scenario"
+              echo "========================================="
+
+              # Configure environment variables for the current scenario
+              if [ "$scenario" = "HCCL_INTRA_ROCE_ENABLE=1" ]; then
+                  export HCCL_INTRA_ROCE_ENABLE=1
+                  unset ASCEND_BUFFER_POOL
+                  echo "HCCL_INTRA_ROCE_ENABLE is set to 1, ASCEND_BUFFER_POOL is unset"
+              else
+                  unset HCCL_INTRA_ROCE_ENABLE
+                  export ASCEND_BUFFER_POOL=4:8
+                  echo "HCCL_INTRA_ROCE_ENABLE is not set, ASCEND_BUFFER_POOL is set to 4:8"
+              fi
+
+              # Run each test case in the current scenario
+              for test_case in "${TEST_CASES[@]}"; do
+                  echo ""
+                  echo "-----------------------------------------"
+                  echo "Test: $test_case"
+                  echo "-----------------------------------------"
+
+                  if [ ! -f "$test_case" ]; then
+                      echo "Warning: Test file $test_case not found, skipping..."
+                      continue
+                  fi
+
+                  # Run the test with 2 devices in distributed mode
+                  # Run rank 0 on device 0
+                  python3 $test_case \
+                      --device_id=0 \
+                      --rank=0 \
+                      --world_size=2 \
+                      --distributed \
+                      2>&1 | tee "/tmp/hixl_test_${scenario//=/}_${test_case%.py}_rank0.log" &
+                  PID0=$!
+
+                  # Run rank 1 on device 1
+                  python3 $test_case \
+                      --device_id=2 \
+                      --rank=1 \
+                      --world_size=2 \
+                      --distributed \
+                      2>&1 | tee "/tmp/hixl_test_${scenario//=/}_${test_case%.py}_rank1.log" &
+                  PID1=$!
+
+                  # Wait for both processes to complete
+                  wait $PID0
+                  TEST_RESULT0=$?
+                  wait $PID1
+                  TEST_RESULT1=$?
+
+                  # Check test results
+                  if [ $TEST_RESULT0 -eq 0 ] && [ $TEST_RESULT1 -eq 0 ]; then
+                      echo "✓ $test_case PASSED (scenario: $scenario)"
+                      PASSED_TESTS+=("$scenario:$test_case")
+                  else
+                      echo "✗ $test_case FAILED (scenario: $scenario)"
+                      if [ $TEST_RESULT0 -ne 0 ]; then
+                          echo "  Rank 0 failed with code: $TEST_RESULT0"
+                      fi
+                      if [ $TEST_RESULT1 -ne 0 ]; then
+                          echo "  Rank 1 failed with code: $TEST_RESULT1"
+                      fi
+                      FAILED_TESTS+=("$scenario:$test_case")
+                  fi
+              done
+          done
+
+          echo ""
+          echo "========================================="
+          echo "Test Summary"
+          echo "========================================="
+          echo "Passed tests: ${#PASSED_TESTS[@]}"
+          for test in "${PASSED_TESTS[@]}"; do
+              echo "  ✓ $test"
+          done
+
+          echo ""
+          echo "Failed tests: ${#FAILED_TESTS[@]}"
+          for test in "${FAILED_TESTS[@]}"; do
+              echo "  ✗ $test"
+          done
+
+          # Cleanup: Stop Mooncake Master
+          echo ""
+          echo "Stopping Mooncake Master..."
+          kill $MASTER_PID 2>/dev/null || true
+          wait $MASTER_PID 2>/dev/null || true
+
+          # Exit with error if any tests failed
+          if [ ${#FAILED_TESTS[@]} -gt 0 ]; then
+              echo ""
+              echo "Some tests failed!"
+              exit 1
+          fi
+
+          echo ""
+          echo "All Hixl Mooncake Store tests completed successfully!"
+
+
+    - name: Test Summary
+      if: always()
+      shell: bash
+      run: |
+        echo "CI Test completed"
+
+    - name: Upload Test Logs
+      if: always()
+      uses: actions/upload-artifact@v4
+      with:
+        name: test-logs-${{ github.run_number }}
+        path: |
+          /tmp/hixl-test-log/*
+        retention-days: 30
+        if-no-files-found: warn

.github/workflows/ci_cu13.yml

@@ -1,18 +1,10 @@
 name: 'Build Wheel (CUDA 13)'
 
 on:
-  push:
-    branches: [ "main" ]
-  pull_request:
-    branches: [ "main" ]
-    types: [opened, synchronize, reopened, labeled]
+  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:
@@ -20,12 +12,13 @@ jobs:
     env:
       BUILD_WITH_EP: "1"
       CU13_BUILD: "1"
-      EP_TORCH_VERSIONS: "2.9.0;2.9.1;2.10.0"
       TORCH_CUDA_ARCH_LIST: "8.0;9.0"
       SCCACHE_GHA_ENABLED: "true"
 
     steps:
     - uses: actions/checkout@v4
+      with:
+        persist-credentials: false
 
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v5
@@ -66,6 +59,7 @@ jobs:
     - name: Install dependencies
       run: |
         sudo apt update -y
+        sudo apt install -y ninja-build
         sudo bash -x dependencies.sh -y
         df -h
       shell: bash
@@ -74,13 +68,14 @@ jobs:
       run: |
         mkdir build
         cd build
-        cmake .. \
+        cmake -G Ninja .. \
           -DUSE_ETCD=ON \
           -DUSE_REDIS=ON \
           -DUSE_HTTP=ON \
           -DWITH_STORE=ON \
           -DWITH_P2P_STORE=ON \
           -DWITH_EP=ON \
+          -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 \
@@ -96,8 +91,8 @@ jobs:
         export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
         export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
         cd build
-        make -j4
-        sudo make install
+        cmake --build .
+        sudo cmake --install .
         df -h
       shell: bash
 

.github/workflows/deploy.yml

@@ -4,6 +4,10 @@ on:
   # Runs on pushes targeting the default branch
   push:
     branches: ["main"]
+    paths:
+      - 'docs/**'
+      - 'requirements_docs.txt'
+      - '.github/workflows/deploy.yml'
   # Allows you to run this workflow manually from the Actions tab
   workflow_dispatch:
 
@@ -31,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,12 +1,16 @@
 name: 'Integration test (Linux)'
 
 on:
-  push:
-    branches: [ "main" ]
-  pull_request_target:
-    branches: [ "main" ]
-    types: [opened, synchronize, reopened, labeled]
-
+  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:
   test-sglang-integration:
@@ -17,18 +21,23 @@ jobs:
       - name: trigger T-one test
         if: ${{ env.tone_user_name != '' }}
         run: |
-          SHA="${{ github.event.pull_request.head.sha }}"
+          # 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_name }}" = "push" ]; then
             SHA="${{ github.sha }}"
+            PR_ID=""
           fi
+          echo "PR_ID=${PR_ID}"
           max_attempts=120
           attempt=1
           while [ $attempt -le $max_attempts ]; do
             echo "Attempt $attempt: Fetching artifact..."
-            if curl -L -fs -o artifact.json -H "Accept: application/vnd.github+json" -H "X-GitHub-Api-Version: 2022-11-28" https://api.github.com/repos/${{ github.repository }}/actions/artifacts; then
+            if curl -L -fs -o artifact.json -H "Accept: application/vnd.github+json" -H "X-GitHub-Api-Version: 2022-11-28" https://api.github.com/repos/${{ github.repository }}/actions/artifacts?per_page=100; then
               artifact_id=""
               if jq empty artifact.json >/dev/null 2>&1; then
-                artifact_id=$(jq -r ".artifacts[] | select(.name | contains(\"py312\") ) | select(.name | contains(\"cu130\") | not) | select(.workflow_run.head_sha == \"$SHA\" ) | .id" artifact.json | head -n 1)
+                artifact_id=$(jq -r ".artifacts[] | select(.name | contains(\"py312\") ) | select(.name | contains(\"mooncake\") ) | select(.name | contains(\"cu130\") | not) | select(.workflow_run.head_sha == \"$SHA\" ) | .id" artifact.json | head -n 1)
               else
                 echo "Failed to download artifact list. Retrying..."
               fi
@@ -53,9 +62,14 @@ jobs:
             echo "Failed to fetch artifacts after $max_attempts attempts"
             exit 1
           fi
+
+          ENV_INFO="ARTIFACT_ID=${artifact_id} GIT_REPO=${{ github.repository }}"
+          if [ -n "$PR_ID" ]; then
+            ENV_INFO="${ENV_INFO} PR_ID=${PR_ID}"
+          fi
           signature="${{ secrets.TONE_USER_NAME }}|${{ secrets.TONE_USER_TOKEN }}|$(python3 -c "import time;print(time.time())")"
           signature="$(python3 -c "import base64;print(base64.b64encode(\"$signature\".encode('utf-8')).decode('utf-8'))")"
-          curl -s -H 'Content-Type: application/json' -X POST -d "{\"workspace\":\"mooncake_test\",\"project\":\"mooncake-ci\",\"template\":\"mooncake-ci-test\",\"name\":\"mooncake-ci-${SHA}\",\"username\":\"${{ secrets.TONE_USER_NAME }}\",\"env_ifs\":\" \",\"env_info\":\"ARTIFACT_ID=${artifact_id} GIT_REPO=${{ github.repository }}\",\"signature\":\"$signature\"}" https://tone.openanolis.cn/api/job/create/ > job.json
+          curl -s -H 'Content-Type: application/json' -X POST -d "{\"workspace\":\"mooncake_test\",\"project\":\"mooncake-ci\",\"template\":\"mooncake-ci-test\",\"name\":\"mooncake-ci-${SHA}\",\"username\":\"${{ secrets.TONE_USER_NAME }}\",\"env_ifs\":\" \",\"env_info\":\"${ENV_INFO}\",\"signature\":\"$signature\"}" https://tone.openanolis.cn/api/job/create/ > job.json
           if [ "$(jq .code job.json)" == 200 ]; then
             echo "job created"
           else

.github/workflows/release-cuda13.yaml

@@ -18,7 +18,6 @@ jobs:
     env:
       BUILD_WITH_EP: "1"
       CU13_BUILD: "1"
-      EP_TORCH_VERSIONS: "2.9.0;2.9.1;2.10.0"
       TORCH_CUDA_ARCH_LIST: "8.0;9.0"
     steps:
       - name: Checkout source
@@ -66,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 -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

@@ -17,7 +17,6 @@ jobs:
         python-version: ['3.10', '3.11', '3.12', '3.13']
     env:
       BUILD_WITH_EP: "1"
-      EP_TORCH_VERSIONS: "2.9.0;2.9.1;2.10.0"
       TORCH_CUDA_ARCH_LIST: "8.0;9.0"
     steps:
       - name: Checkout source
@@ -65,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 -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

.gitmodules

@@ -2,3 +2,7 @@
 	path = extern/pybind11
 	url = https://github.com/pybind/pybind11.git
 	branch = stable
+[submodule "extern/yalantinglibs"]
+	path = extern/yalantinglibs
+	url = https://github.com/alibaba/yalantinglibs.git
+	branch = v0.5.7

.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

@@ -14,12 +14,14 @@ endif()
 
 option(WITH_TE "build mooncake transfer engine and sample code" ON)
 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
@@ -39,12 +41,25 @@ option(STORE_USE_ETCD "build mooncake store with etcd" OFF)
 if (STORE_USE_ETCD)
   add_compile_definitions(STORE_USE_ETCD)
 endif()
+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)
@@ -61,12 +76,103 @@ 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)
-  message(WARNING "Option `WITH_EP` is deprecated. Mooncake EP now builds with setuptools. Please set environment variable BUILD_WITH_EP=1 to enable.")
+  if (EP_USE_IDE)
+    message(WARNING "EP_USE_IDE enabled. DO NOT USE IN PRODUCTION!")
+    add_subdirectory(mooncake-ep)
+    include_directories(mooncake-ep/include)
+    add_subdirectory(mooncake-pg)
+    include_directories(mooncake-pg/include)
+  else ()
+    message(STATUS "WITH_EP enabled: building Mooncake EP and PG Python extensions")
+    find_package(CUDAToolkit REQUIRED)
+    message(STATUS "Detected CUDA version: ${CUDAToolkit_VERSION}")
+
+    # EP_TORCH_VERSIONS: semicolon-separated list of PyTorch versions to build for.
+    # Can be set via -DEP_TORCH_VERSIONS="2.9.1;2.8.0" or the EP_TORCH_VERSIONS env var.
+    # Empty means build with the currently-installed torch.
+    if(NOT EP_TORCH_VERSIONS)
+      set(EP_TORCH_VERSIONS "$ENV{EP_TORCH_VERSIONS}")
+    endif()
+    set(EP_TORCH_VERSIONS "${EP_TORCH_VERSIONS}" CACHE STRING
+        "PyTorch versions for EP/PG extensions, semicolon-separated (empty = use currently-installed torch)")
+
+    # TORCH_CUDA_ARCH_LIST forwarded to the torch CUDA extension build.
+    if(NOT TORCH_CUDA_ARCH_LIST)
+      set(TORCH_CUDA_ARCH_LIST "$ENV{TORCH_CUDA_ARCH_LIST}")
+    endif()
+    if(NOT TORCH_CUDA_ARCH_LIST)
+      set(TORCH_CUDA_ARCH_LIST "8.0;9.0")
+    endif()
+    set(TORCH_CUDA_ARCH_LIST "${TORCH_CUDA_ARCH_LIST}" CACHE STRING
+        "CUDA arch list for EP/PG extension builds (e.g. \"8.0;9.0\")")
+
+    # Staging directory: EP/PG .so files are placed here during make and later
+    # injected into the wheel AFTER auditwheel, so patchelf never touches the
+    # CUDA fatbins (which would cause cudaErrorInvalidKernelImage at runtime).
+    set(EP_PG_STAGING_DIR "${CMAKE_BINARY_DIR}/ep_pg_staging")
+
+    # Convert semicolon-separated lists to pipe-separated strings so they survive
+    # CMake's COMMAND list-splitting (semicolons are CMake list separators).
+    string(REPLACE ";" "|" _ep_torch_versions_pipe "${EP_TORCH_VERSIONS}")
+    string(REPLACE ";" "|" _torch_cuda_arch_list_pipe "${TORCH_CUDA_ARCH_LIST}")
+
+    add_custom_target(mooncake_ep_ext ALL
+      COMMAND ${CMAKE_COMMAND} -E make_directory "${EP_PG_STAGING_DIR}"
+      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}"
+        "-DENGINE_SO_PATH=$<TARGET_FILE:engine>"
+        -P "${CMAKE_CURRENT_SOURCE_DIR}/mooncake-ep/BuildEpExt.cmake"
+      COMMENT "Building Mooncake EP Python extension(s)"
+      DEPENDS engine
+      VERBATIM
+    )
+
+    add_custom_target(mooncake_pg_ext ALL
+      COMMAND ${CMAKE_COMMAND} -E make_directory "${EP_PG_STAGING_DIR}"
+      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}"
+        "-DENGINE_SO_PATH=$<TARGET_FILE:engine>"
+        -P "${CMAKE_CURRENT_SOURCE_DIR}/mooncake-pg/BuildPgExt.cmake"
+      COMMENT "Building Mooncake PG Python extension(s)"
+      DEPENDS engine mooncake_ep_ext
+      VERBATIM
+    )
+  endif ()
 endif()
 
 add_subdirectory(mooncake-integration)
 
+if (WITH_STORE_GO AND WITH_STORE)
+  add_custom_target(build_store_go DEPENDS mooncake_store transfer_engine)
+  add_custom_command(
+      TARGET build_store_go
+      COMMAND bash build.sh ${CMAKE_BINARY_DIR} ${CMAKE_CURRENT_BINARY_DIR} ${USE_ETCD} ${USE_REDIS} ${USE_HTTP} ${USE_ETCD_LEGACY}
+      WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/mooncake-store/go
+  )
+  set_property(TARGET build_store_go PROPERTY EXCLUDE_FROM_ALL FALSE)
+  message(STATUS "Mooncake Store Go bindings will be built")
+endif()
+
 if (WITH_P2P_STORE)
   add_subdirectory(mooncake-p2p-store)
   message(STATUS "P2P Store will be built")

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

@@ -15,6 +15,8 @@
   [![Docs](https://img.shields.io/badge/docs-live-brightgreen)](https://kvcache-ai.github.io/Mooncake/)
   [![PyPI](https://img.shields.io/pypi/v/mooncake-transfer-engine)](https://pypi.org/project/mooncake-transfer-engine)
   [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/mooncake-transfer-engine)](https://pypi.org/project/mooncake-transfer-engine)
+  [![CUDA <=12.9](https://img.shields.io/static/v1?label=CUDA&message=%3C%3D12.9&color=76B900)](https://pypi.org/project/mooncake-transfer-engine)
+  [![CUDA 13.0/13.1](https://img.shields.io/static/v1?label=CUDA&message=13.0%2F13.1&color=76B900)](https://pypi.org/project/mooncake-transfer-engine-cuda13)
   [![PyPI - Downloads](https://img.shields.io/pypi/dm/mooncake-transfer-engine)](https://pypi.org/project/mooncake-transfer-engine)
   [![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/kvcache-ai/Mooncake)
   [![GitHub commit activity](https://img.shields.io/github/commit-activity/w/kvcache-ai/Mooncake)](https://github.com/kvcache-ai/Mooncake/graphs/commit-activity)
@@ -29,6 +31,10 @@ This repository also hosts its technical report and the open-sourced traces.
 
 <h2 id="updates">🔄 Updates</h2>
 
+- **Mar 19, 2026**: [TorchSpec: Speculative Decoding Training at Scale](https://pytorch.org/blog/torchspec-speculative-decoding-training-at-scale) is [open sourced](https://github.com/torchspec-project/TorchSpec), using Mooncake to decouple inference and training via efficient hidden states management.
+- **Mar 5, 2026**: [LightX2V](https://github.com/ModelTC/LightX2V/pull/893) now supports disaggregated deployment based on Mooncake, enabling encoder/transformer service decoupling with Mooncake Transfer Engine for high-performance cross-device and cross-machine data transfer.
+- **Feb 25, 2026**: [SGLang](https://github.com/sgl-project/sglang) merged [Encoder Global Cache Manager](https://github.com/sgl-project/sglang/pull/16137), introducing a Mooncake-powered global multimodal embedding cache that enables cross-instance sharing of ViT embeddings to avoid redundant GPU computation.
+- **Feb 24, 2026**: [vLLM-Omni](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/disaggregated_inference/) introduces disaggregated inference connectors with support for both `MooncakeStoreConnector` and `MooncakeTransferEngineConnector` for multi-node omni-modality pipelines.
 - **Feb 12, 2026**: [Mooncake Joins PyTorch Ecosystem](https://pytorch.org/blog/mooncake-joins-pytorch-ecosystem/) We are thrilled to announce that Mooncake has officially joined the PyTorch Ecosystem!
 - **Jan 28, 2026**: [FlexKV](https://github.com/taco-project/FlexKV), a distributed KV store and cache system from Tencent and NVIDIA in collaboration with the community, now supports [distributed KVCache reuse](https://github.com/taco-project/FlexKV/blob/main/docs/dist_reuse/README_en.md) with the Mooncake Transfer Engine.
  - **Dec 27, 2025**: Collaboration with [ROLL](https://github.com/alibaba/ROLL)! Check out the paper [here](https://arxiv.org/abs/2512.22560).
@@ -91,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*.
@@ -172,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`:
@@ -195,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
@@ -223,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.
@@ -248,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>
 

benchmarks/storage_benchmark/storage_benchmark.py

@@ -0,0 +1,956 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Mooncake KVCache Storage Benchmark Tool
+"""
+
+import argparse
+import json
+import time
+import os
+import statistics
+import random
+import errno
+from pathlib import Path
+from typing import Dict, List, Optional
+from dataclasses import dataclass
+
+# ============================================================================
+# Constants
+# ============================================================================
+
+BLOCK_SIZE_TOKENS = 512  # Number of tokens per block
+DEFAULT_BYTES_PER_TOKEN = 2048  # 7B model FP16 (2KB per token)
+BLOCK_SIZE_BYTES = BLOCK_SIZE_TOKENS * DEFAULT_BYTES_PER_TOKEN  # 1MB per block
+MIN_LATENCY_MS = 0.001  # Minimum latency in milliseconds (1 microsecond)
+
+# Model KVCache sizes (bytes per token, based on LMCache calculator)
+# Source: https://lmcache.ai/kv_cache_calculator.html
+MODEL_BYTES_PER_TOKEN = {
+    "llama-3.1-405b": 327680,
+    "qwen3-32b": 81920,
+    "deepseek-v3": 1748992,
+    "glm-4.6": 157013,
+    "default": DEFAULT_BYTES_PER_TOKEN,
+}
+
+# ============================================================================
+# Data Structures
+# ============================================================================
+
+@dataclass
+class KVCacheRequest:
+    """KVCache request
+
+    Attributes:
+        timestamp: Request timestamp in milliseconds
+        hash_ids: List of block IDs (each ID corresponds to a 512-token block)
+        input_length: Input token count
+        output_length: Output token count
+    """
+    timestamp: float
+    hash_ids: List[int]
+    input_length: int
+    output_length: int
+
+# ============================================================================
+# Storage Layer: Offset Allocator
+# ============================================================================
+
+class OffsetAllocatorStorage:
+    """High-performance block storage based on Offset Allocator
+
+    Architecture:
+    -----------
+    1. Single large file stores all blocks (avoids file explosion)
+    2. Uses offset to manage file space (similar to Mooncake's OffsetAllocator)
+    3. hash_id -> offset mapping stored in memory (fast lookup)
+
+    Block Organization:
+    -----------
+    Each block corresponds to 512 tokens, fixed size 1MB:
+    - hash_id[0] -> block_0 (tokens [0...511])     -> offset 0
+    - hash_id[1] -> block_1 (tokens [512...1023])  -> offset 1
+    - hash_id[i] -> block_i (tokens [i*512...(i+1)*512-1]) -> offset i
+
+    Performance Advantages:
+    -----------
+    - Only one file, no file explosion
+    - Offset reuse, reduces memory allocation
+    - pread/pwrite, thread-safe, no seek needed
+    - Keep fd open, reduces open/close overhead
+    - Metadata in memory, O(1) lookup
+
+    Attributes:
+        storage_dir: Storage directory path
+        block_size_bytes: Block size in bytes
+        max_blocks: Maximum number of blocks
+        hash_id_to_offset: hash_id -> offset mapping
+        free_offsets: List of reusable offsets
+        next_offset: Next allocatable offset
+    """
+
+    def __init__(self, storage_dir: str, bytes_per_token: int = DEFAULT_BYTES_PER_TOKEN,
+                 max_blocks: int = 100000, block_size_tokens: int = 512,
+                 fsync_mode: str = 'batch', fsync_batch_size: int = 100):
+        """Initialize Offset Allocator storage
+
+        Args:
+            storage_dir: Storage directory path
+            bytes_per_token: Bytes per token
+            max_blocks: Maximum number of blocks (determines file size)
+            block_size_tokens: Number of tokens per block
+            fsync_mode: When to fsync ('batch', 'always', 'end', 'none')
+            fsync_batch_size: Number of writes between fsync in batch mode
+        """
+        self.storage_dir = Path(storage_dir)
+        self.bytes_per_token = bytes_per_token
+        self.block_size_tokens = block_size_tokens
+        self.block_size_bytes = self.block_size_tokens * self.bytes_per_token
+        self.max_blocks = max_blocks
+
+        # Fsync configuration
+        self.fsync_mode = fsync_mode
+        self.fsync_batch_size = fsync_batch_size
+        self.pending_sync_count = 0
+
+        # Create storage directory
+        self.storage_dir.mkdir(parents=True, exist_ok=True)
+
+        # Single large file
+        self.storage_file = self.storage_dir / "kvcache_storage.bin"
+        self.file_size = self.max_blocks * self.block_size_bytes
+
+        # Initialize storage file
+        if not self.storage_file.exists():
+            self._init_storage_file()
+
+        # hash_id -> offset mapping (metadata, in memory)
+        self.hash_id_to_offset: Dict[int, int] = {}
+
+        # Offset allocator (free list)
+        self.free_offsets: List[int] = []
+        self.next_offset = 0
+
+        # File descriptor (keep open, avoid repeated open/close)
+        self.fd = None
+
+        # Pre-allocated data buffer with pattern to avoid SSD compression artifacts
+        # Using a repeating pattern that looks like realistic data (not all zeros)
+        # Pattern: 64-byte repeated sequence mixed with some variation
+        pattern = bytes([(i & 0xFF) for i in range(256)])  # 0-255 byte pattern
+        pattern_repeats = (self.block_size_bytes // len(pattern)) + 1
+        self._data_buffer = (pattern * pattern_repeats)[:self.block_size_bytes]
+
+        # Statistics
+        self.stats = {
+            'read_count': 0,
+            'write_count': 0,
+            'read_bytes': 0,
+            'write_bytes': 0,
+            'read_latencies_ms': [],
+            'write_latencies_ms': [],
+            'sync_count': 0,  # Number of fsync operations performed
+        }
+
+    # ========================================================================
+    # Internal Methods
+    # ========================================================================
+
+    def _init_storage_file(self):
+        """Initialize storage file (pre-allocate space)
+
+        Create sparse file to avoid actual disk space usage until data is written
+        """
+        with open(self.storage_file, 'wb') as f:
+            f.seek(self.file_size - 1)
+            f.write(b'\0')
+            f.flush()
+            os.fsync(f.fileno())
+
+    def _get_fd(self):
+        """Get file descriptor (lazy open)
+
+        Returns:
+            int: File descriptor
+        """
+        if self.fd is None:
+            # Use O_RDWR | O_CREAT, no O_DIRECT (Python compatibility)
+            self.fd = os.open(self.storage_file, os.O_RDWR | os.O_CREAT)
+        return self.fd
+
+    def _allocate_offset(self) -> int:
+        """Allocate a new offset
+
+        Prioritize reusing freed offsets, otherwise allocate new offset
+
+        Returns:
+            int: Allocated offset
+        """
+        if self.free_offsets:
+            return self.free_offsets.pop()
+        offset = self.next_offset
+        self.next_offset += 1
+        return offset
+
+    def _free_offset(self, offset: int):
+        """Free offset for reuse
+
+        Args:
+            offset: Offset to free
+        """
+        self.free_offsets.append(offset)
+
+    # ========================================================================
+    # Public Interface
+    # ========================================================================
+
+    def block_exists(self, hash_id: int) -> bool:
+        """Check if block exists
+
+        Args:
+            hash_id: Unique block identifier
+
+        Returns:
+            bool: Whether block exists
+        """
+        return hash_id in self.hash_id_to_offset
+
+    def read_block(self, hash_id: int) -> float:
+        """Read block using pread
+
+        Args:
+            hash_id: Unique block identifier
+
+        Returns:
+            float: Read latency in milliseconds, or 0 if block doesn't exist
+        """
+        if hash_id not in self.hash_id_to_offset:
+            return 0.0  # Block doesn't exist, no latency to measure
+
+        offset = self.hash_id_to_offset[hash_id]
+        file_offset = offset * self.block_size_bytes
+
+        start = time.perf_counter()
+
+        try:
+            fd = self._get_fd()
+            data = os.pread(fd, self.block_size_bytes, file_offset)
+            latency_ms = (time.perf_counter() - start) * 1000.0
+
+            self.stats['read_count'] += 1
+            self.stats['read_bytes'] += len(data)
+            self.stats['read_latencies_ms'].append(latency_ms)
+            return latency_ms
+        except OSError as e:
+            print(f"Error reading block {hash_id} at offset {file_offset}: {e}")
+            return 0.0  # Error case, don't pollute stats
+
+    def write_block(self, hash_id: int) -> float:
+        """Write block using pwrite
+
+        Args:
+            hash_id: Unique block identifier
+
+        Returns:
+            float: Write latency in milliseconds
+        """
+        # Allocate offset
+        offset = self._allocate_offset()
+        file_offset = offset * self.block_size_bytes
+
+        # Use pre-allocated buffer (much faster than os.urandom)
+        data = self._data_buffer
+
+        start = time.perf_counter()
+
+        try:
+            fd = self._get_fd()
+            written = os.pwrite(fd, data, file_offset)
+
+            write_done = time.perf_counter()
+
+            # Conditional fsync based on mode
+            if self.fsync_mode == 'always':
+                # Include fsync in latency measurement
+                os.fsync(fd)
+                self.stats['sync_count'] += 1
+                self.pending_sync_count = 0
+                latency_ms = (time.perf_counter() - start) * 1000.0
+                # Evict from page cache AFTER fsync to ensure reads measure actual SSD performance
+                os.posix_fadvise(fd, file_offset, self.block_size_bytes, os.POSIX_FADV_DONTNEED)
+            elif self.fsync_mode == 'batch':
+                # For batch mode, only measure write time (fsync is deferred)
+                self.pending_sync_count += 1
+                if self.pending_sync_count >= self.fsync_batch_size:
+                    os.fsync(fd)
+                    self.stats['sync_count'] += 1
+                    self.pending_sync_count = 0
+                latency_ms = (write_done - start) * 1000.0  # Only write time
+                # Evict from page cache after each write
+                os.posix_fadvise(fd, file_offset, self.block_size_bytes, os.POSIX_FADV_DONTNEED)
+            elif self.fsync_mode == 'none':
+                latency_ms = (write_done - start) * 1000.0
+                # Evict from page cache even when not syncing
+                os.posix_fadvise(fd, file_offset, self.block_size_bytes, os.POSIX_FADV_DONTNEED)
+            else:  # 'end' mode
+                latency_ms = (write_done - start) * 1000.0
+                # Evict from page cache (fsync will happen at the end)
+                os.posix_fadvise(fd, file_offset, self.block_size_bytes, os.POSIX_FADV_DONTNEED)
+
+            # Update mapping
+            self.hash_id_to_offset[hash_id] = offset
+
+            self.stats['write_count'] += 1
+            self.stats['write_bytes'] += written
+            self.stats['write_latencies_ms'].append(latency_ms)
+            return latency_ms
+        except OSError as e:
+            if e.errno == errno.ENOSPC:
+                print(f"Error: Disk full when writing block {hash_id} at offset {file_offset}")
+            else:
+                print(f"Error writing block {hash_id} at offset {file_offset}: {e}")
+            return 0.0  # Error case, don't pollute stats
+
+    def __enter__(self):
+        """Context manager entry"""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit - ensures cleanup"""
+        # Perform final fsync before closing for 'end' and 'batch' modes
+        self._finalize_sync()
+        self.close(force_sync=False)  # Already synced above
+        return False
+
+    def _finalize_sync(self):
+        """Perform final fsync before closing (for 'end' mode and pending batch writes)"""
+        if self.fd is not None:
+            if self.fsync_mode == 'end':
+                try:
+                    os.fsync(self.fd)
+                    self.stats['sync_count'] += 1
+                except OSError:
+                    pass
+            elif self.fsync_mode == 'batch' and self.pending_sync_count > 0:
+                # Flush remaining pending writes
+                try:
+                    os.fsync(self.fd)
+                    self.stats['sync_count'] += 1
+                    self.pending_sync_count = 0
+                except OSError:
+                    pass
+
+    def close(self, force_sync: bool = True):
+        """Close file
+
+        Args:
+            force_sync: Whether to force fsync before closing
+        """
+        # For backward compatibility with non-context-manager usage
+        if force_sync:
+            self._finalize_sync()
+
+        if self.fd is not None:
+            os.close(self.fd)
+            self.fd = None
+
+    def get_stats(self) -> Dict:
+        """Get statistics
+
+        Returns:
+            Dict: Dictionary containing read/write statistics
+        """
+        def calc_stats(latencies):
+            """Calculate latency statistics"""
+            if not latencies:
+                return {'avg_ms': 0, 'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0}
+            return {
+                'avg_ms': statistics.mean(latencies),
+                **calc_percentiles(latencies),
+            }
+
+        return {
+            'read': {
+                'count': self.stats['read_count'],
+                'mb': self.stats['read_bytes'] / 1024 / 1024,
+                **calc_stats(self.stats['read_latencies_ms'])
+            },
+            'write': {
+                'count': self.stats['write_count'],
+                'mb': self.stats['write_bytes'] / 1024 / 1024,
+                **calc_stats(self.stats['write_latencies_ms'])
+            },
+            'sync_count': self.stats['sync_count'],
+            'total_blocks': len(self.hash_id_to_offset),
+            'free_blocks': len(self.free_offsets),
+        }
+
+
+# ============================================================================
+# Benchmark Layer
+# ============================================================================
+
+class StorageBenchmark:
+    """KVCache storage benchmark
+
+    Based on Mooncake OffsetAllocator + vLLM PagedAttention implementation:
+
+    Example:
+    -----
+    Request A: [1, 2, 4]
+    -> hash_id 1 -> not exist, write block_1 (offset=0, 1MB)
+    -> hash_id 2 -> not exist, write block_2 (offset=1, 1MB)
+    -> hash_id 4 -> not exist, write block_4 (offset=2, 1MB)
+
+    Request B: [1, 2, 4, 6]
+    -> hash_id 1 -> exists, read block_1 (offset=0) ✓ prefix reuse
+    -> hash_id 2 -> exists, read block_2 (offset=1) ✓ prefix reuse
+    -> hash_id 4 -> exists, read block_4 (offset=2) ✓ prefix reuse
+    -> hash_id 6 -> not exist, write block_6 (offset=3, 1MB)
+
+    Performance Advantages:
+    ---------
+    - Single file operation, no file explosion
+    - Offset reuse, reduces memory allocation
+    - pread/pwrite, thread-safe
+    """
+
+    def __init__(self, storage_dir: str, bytes_per_token: int = DEFAULT_BYTES_PER_TOKEN,
+                 max_blocks: int = 100000, block_size_tokens: int = 512,
+                 fsync_mode: str = 'batch', fsync_batch_size: int = 100):
+        """Initialize benchmark
+
+        Args:
+            storage_dir: Storage directory
+            bytes_per_token: Bytes per token
+            max_blocks: Maximum number of blocks
+            block_size_tokens: Number of tokens per block
+            fsync_mode: When to fsync ('batch', 'always', 'end', 'none')
+            fsync_batch_size: Number of writes between fsync in batch mode
+        """
+        self.storage = OffsetAllocatorStorage(
+            storage_dir, bytes_per_token, max_blocks,
+            block_size_tokens, fsync_mode, fsync_batch_size
+        )
+        self.bytes_per_token = bytes_per_token
+        self.block_size_tokens = block_size_tokens
+
+        # Statistics
+        self.stats = {
+            'total_requests': 0,
+            'total_blocks': 0,
+            'read_blocks': 0,
+            'write_blocks': 0,
+            'prefix_hit_blocks': 0,  # Number of prefix hit blocks
+            'request_latencies_ms': [],
+        }
+
+    def process_request(self, req: KVCacheRequest) -> float:
+        """Process a KVCache request
+
+        Based on vLLM's prefix caching mechanism:
+        - Each hash_id corresponds to an independent block
+        - Prefix reuse achieved through hash_id matching
+
+        Args:
+            req: KVCache request
+
+        Returns:
+            float: Request latency in milliseconds
+        """
+        self.stats['total_requests'] += 1
+        self.stats['total_blocks'] += len(req.hash_ids)
+
+        start_time = time.perf_counter()
+        total_latency = 0.0
+
+        # Process each hash_id (in order)
+        for hash_id in req.hash_ids:
+            if self.storage.block_exists(hash_id):
+                # Block exists, read (reuse cached block)
+                total_latency += self.storage.read_block(hash_id)
+                self.stats['read_blocks'] += 1
+                self.stats['prefix_hit_blocks'] += 1  # Count all cache hits as prefix reuse
+            else:
+                # Block doesn't exist, write (new block)
+                total_latency += self.storage.write_block(hash_id)
+                self.stats['write_blocks'] += 1
+
+        latency_ms = total_latency if total_latency > 0 else MIN_LATENCY_MS
+        self.stats['request_latencies_ms'].append(latency_ms)
+
+        return latency_ms
+
+    def get_stats(self) -> Dict:
+        """Get statistics
+
+        Returns:
+            Dict: Statistics dictionary
+        """
+        storage_stats = self.storage.get_stats()
+
+        request_latencies = self.stats['request_latencies_ms']
+
+        if request_latencies:
+            latency_stats = {
+                'avg_ms': statistics.mean(request_latencies),
+                **calc_percentiles(request_latencies),
+            }
+        else:
+            latency_stats = {'avg_ms': 0, 'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0}
+
+        total_blocks = self.stats['total_blocks']
+        read_blocks = self.stats['read_blocks']
+        write_blocks = self.stats['write_blocks']
+
+        return {
+            'total_requests': self.stats['total_requests'],
+            'total_blocks': total_blocks,
+            'read_blocks': read_blocks,
+            'write_blocks': write_blocks,
+            'prefix_hit_blocks': self.stats['prefix_hit_blocks'],
+            'block_hit_rate': read_blocks / total_blocks if total_blocks > 0 else 0,
+            'write_ratio': write_blocks / total_blocks if total_blocks > 0 else 0,
+            'tokens_per_block': self.block_size_tokens,  # Configurable block size in tokens
+            'latency': latency_stats,
+            'storage': storage_stats,
+        }
+
+    def __enter__(self):
+        """Context manager entry"""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit - ensures cleanup"""
+        self.close()
+        return False
+
+    def close(self, force_sync: bool = True):
+        """Close storage
+
+        Args:
+            force_sync: Whether to force final sync before closing
+        """
+        self.storage.close(force_sync=force_sync)
+
+
+# ============================================================================
+# Utility Functions
+# ============================================================================
+
+def calc_percentiles(data: List[float]) -> Dict[str, float]:
+    """Calculate latency percentiles
+
+    Uses linear interpolation for accurate percentile calculation.
+    This is more accurate than statistics.quantiles() for small datasets.
+
+    Args:
+        data: List of latency values in milliseconds
+
+    Returns:
+        Dict containing p50, p95, p99 percentiles
+    """
+    if not data:
+        return {'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0}
+
+    # Sort data for percentile calculation
+    sorted_data = sorted(data)
+    n = len(sorted_data)
+
+    def get_percentile(p: float) -> float:
+        """Get percentile using linear interpolation
+
+        Args:
+            p: Percentile (0-100)
+
+        Returns:
+            Value at percentile
+        """
+        index = (n - 1) * p / 100
+        lower = int(index)
+        upper = min(lower + 1, n - 1)
+
+        if lower == upper:
+            return sorted_data[lower]
+
+        # Linear interpolation
+        weight = index - lower
+        return sorted_data[lower] * (1 - weight) + sorted_data[upper] * weight
+
+    return {
+        'p50_ms': get_percentile(50),
+        'p95_ms': get_percentile(95),
+        'p99_ms': get_percentile(99),
+    }
+
+
+# ============================================================================
+# Trace Loader
+# ============================================================================
+
+class TraceLoader:
+    """Load KVCache trace"""
+
+    def __init__(self, trace_path: str):
+        """Initialize trace loader
+
+        Args:
+            trace_path: Trace file path
+        """
+        self.trace_path = trace_path
+        self.requests = []
+        self._load_trace()
+
+    def _load_trace(self):
+        """Load trace file with error handling"""
+        line_num = 0
+        try:
+            with open(self.trace_path, 'r') as f:
+                for line in f:
+                    line_num += 1
+                    line = line.strip()
+                    if not line:
+                        continue
+                    try:
+                        req = json.loads(line)
+                        # Validate required fields
+                        if not all(k in req for k in ['timestamp', 'hash_ids', 'input_length', 'output_length']):
+                            print(f"Warning: Line {line_num} missing required fields, skipping")
+                            continue
+                        if not isinstance(req['hash_ids'], list):
+                            print(f"Warning: Line {line_num} has invalid hash_ids (not a list), skipping")
+                            continue
+                        self.requests.append(KVCacheRequest(
+                            timestamp=float(req['timestamp']),
+                            hash_ids=req['hash_ids'],
+                            input_length=int(req['input_length']),
+                            output_length=int(req['output_length'])
+                        ))
+                    except (json.JSONDecodeError, ValueError, KeyError) as e:
+                        print(f"Warning: Line {line_num} has invalid format: {e}, skipping")
+                        continue
+        except FileNotFoundError:
+            raise FileNotFoundError(f"Trace file not found: {self.trace_path}")
+        except OSError as e:
+            raise OSError(f"Error reading trace file {self.trace_path}: {e}")
+
+    def get_requests(self) -> List[KVCacheRequest]:
+        """Get request list
+
+        Returns:
+            List[KVCacheRequest]: Request list
+        """
+        return self.requests
+
+
+# ============================================================================
+# Benchmark Runner
+# ============================================================================
+
+def run_benchmark(trace_path: str, storage_dir: str, bytes_per_token: int = DEFAULT_BYTES_PER_TOKEN,
+                   max_requests: Optional[int] = None, max_blocks: int = 100000,
+                   replay_timestamps: bool = False, time_scale: float = 1.0,
+                   block_size_tokens: int = 512,
+                   fsync_mode: str = 'batch', fsync_batch_size: int = 100) -> Dict:
+    """Run benchmark
+
+    Args:
+        trace_path: Trace file path
+        storage_dir: Storage directory
+        bytes_per_token: Bytes per token
+        max_requests: Maximum number of requests (None = all)
+        max_blocks: Maximum number of blocks
+        replay_timestamps: Whether to replay timestamps from trace (simulate realistic timing)
+        time_scale: Time scaling factor (1.0=real-time, 0.1=10x speed, 10.0=0.1x speed)
+        block_size_tokens: Number of tokens per block
+        fsync_mode: When to fsync ('batch', 'always', 'end', 'none')
+        fsync_batch_size: Number of writes between fsync in batch mode
+
+    Returns:
+        Dict: Benchmark results
+    """
+    block_size_bytes = block_size_tokens * bytes_per_token
+
+    print(f"\n{'='*80}")
+    print(f"Running: {Path(trace_path).name}")
+    print(f"Architecture: Offset Allocator (Mooncake style)")
+    print(f"Block size: {block_size_tokens} tokens/block ({block_size_bytes:,} bytes)")
+    print(f"Storage: Single large file with offset-based block management")
+    print(f"Bytes per token: {bytes_per_token}")
+    print(f"Max blocks: {max_blocks}")
+    print(f"Fsync mode: {fsync_mode}" + (f" (batch_size={fsync_batch_size})" if fsync_mode == 'batch' else ''))
+    print(f"Timestamp replay: {'Enabled' if replay_timestamps else 'Disabled'}")
+    if replay_timestamps:
+        scale_desc = 'real-time' if time_scale == 1.0 else f'{1/time_scale:.1f}x speed' if time_scale < 1.0 else f'{time_scale}x slower'
+        print(f"Time scale: {time_scale}x ({scale_desc})")
+    print(f"{'='*80}")
+
+    # Load trace
+    loader = TraceLoader(trace_path)
+    requests = loader.get_requests()
+
+    if max_requests:
+        requests = requests[:max_requests]
+
+    print(f"Loaded {len(requests)} requests")
+
+    # Show timestamp range
+    if replay_timestamps and requests:
+        timestamps = [req.timestamp for req in requests]
+        time_span_ms = max(timestamps) - min(timestamps)
+        print(f"Timestamp range: {min(timestamps):.1f} - {max(timestamps):.1f} ms (span: {time_span_ms:.1f} ms)")
+
+    # Create benchmark instance with context manager for cleanup
+    with StorageBenchmark(
+        storage_dir, bytes_per_token, max_blocks,
+        block_size_tokens, fsync_mode, fsync_batch_size
+    ) as benchmark:
+
+        # Run benchmark
+        start_time = time.perf_counter()
+        total_io_time = 0.0  # Actual I/O time (excluding sleep)
+        last_timestamp = None
+        base_time = time.time()  # Use wall time for replay synchronization
+
+        for i, req in enumerate(requests):
+            # Replay by timestamps
+            sleep_time = 0.0
+            if replay_timestamps and last_timestamp is not None:
+                # Calculate time interval from previous request
+                delta_ms = req.timestamp - last_timestamp
+                sleep_time = delta_ms / 1000.0 / time_scale  # Apply time scaling
+
+                if sleep_time > 0:
+                    time.sleep(sleep_time)
+
+            # Process request (measure I/O time)
+            req_start = time.perf_counter()
+            benchmark.process_request(req)
+            req_io_time = time.perf_counter() - req_start
+            total_io_time += req_io_time
+
+            # Record current request timestamp
+            last_timestamp = req.timestamp
+
+            # Progress output
+            if (i + 1) % 100 == 0:
+                if replay_timestamps:
+                    elapsed_wall_time = time.time() - base_time
+                    simulated_time = (req.timestamp - requests[0].timestamp) / 1000.0 / time_scale
+                    print(f"  Processed {i + 1}/{len(requests)}... (wall: {elapsed_wall_time:.1f}s, simulated: {simulated_time:.1f}s, io: {total_io_time:.1f}s)")
+                else:
+                    print(f"  Processed {i + 1}/{len(requests)}...")
+
+        elapsed = time.perf_counter() - start_time
+
+        # Perform final sync to include it in stats
+        benchmark.storage._finalize_sync()
+
+        # Get statistics (context manager will handle cleanup)
+        stats = benchmark.get_stats()
+
+    # Calculate actual I/O time (excluding sleep)
+    io_time = total_io_time if replay_timestamps else elapsed
+
+    return {
+        'trace_file': Path(trace_path).name,
+        'total_requests': len(requests),
+        'simulation_time_s': elapsed,
+        'io_time_s': io_time,  # Actual I/O time
+        'wall_time_s': elapsed,  # Wall time (including sleep)
+        'requests_per_second': len(requests) / io_time if io_time > 0 else 0,  # Based on I/O time
+        'timestamp_replay_enabled': replay_timestamps,
+        'time_scale': time_scale,
+        'bytes_per_token': bytes_per_token,
+        'block_size_tokens': block_size_tokens,
+        'fsync_mode': fsync_mode,
+        **stats,
+    }
+
+
+# ============================================================================
+# Result Output
+# ============================================================================
+
+def print_results(results: List[Dict]):
+    """Print benchmark results
+
+    Args:
+        results: List of benchmark results
+    """
+    for i, r in enumerate(results, 1):
+        print(f"\n{'='*80}")
+        print(f"  [{i}/{len(results)}] {r['trace_file']}")
+        print(f"{'='*80}")
+
+        print(f"\n[Performance Overview]")
+        print(f"  Total Requests:           {r['total_requests']:,}")
+        print(f"  Queries Per Second (QPS): {r['requests_per_second']:.2f}")
+        print(f"  Cache Hit Rate:           {r['block_hit_rate']:.2%}")
+        print(f"  Write Ratio:              {r['write_ratio']:.2%}")
+        print(f"  Total Blocks:             {r['total_blocks']:,}")
+        print(f"    Read Blocks:            {r['read_blocks']:,}")
+        print(f"    Write Blocks:           {r['write_blocks']:,}")
+        print(f"    Prefix Hits:            {r['prefix_hit_blocks']:,}")
+
+        print(f"\n[Latency Analysis]")
+        req_lat = r['latency']
+        print(f"  Request Latency (End-to-End): Avg={req_lat['avg_ms']:.2f}ms, P50={req_lat['p50_ms']:.2f}ms, P95={req_lat['p95_ms']:.2f}ms, P99={req_lat['p99_ms']:.2f}ms")
+        read_lat = r['storage']['read']
+        write_lat = r['storage']['write']
+        print(f"  Single I/O Operation (Per Block):")
+        print(f"    Read:  Avg={read_lat.get('avg_ms', 0):.3f}ms, P50={read_lat.get('p50_ms', 0):.3f}ms, P95={read_lat.get('p95_ms', 0):.3f}ms, P99={read_lat.get('p99_ms', 0):.3f}ms")
+        print(f"    Write: Avg={write_lat.get('avg_ms', 0):.3f}ms, P50={write_lat.get('p50_ms', 0):.3f}ms, P95={write_lat.get('p95_ms', 0):.3f}ms, P99={write_lat.get('p99_ms', 0):.3f}ms")
+
+        print(f"\n[I/O & Bandwidth]")
+        print(f"  Total Read I/O:      {r['storage']['read']['mb']:>10.1f} MB  ({r['storage']['read']['count']:,} ops)")
+        print(f"  Total Write I/O:     {r['storage']['write']['mb']:>10.1f} MB  ({r['storage']['write']['count']:,} ops)")
+        io_time = r['io_time_s']
+        bandwidth = (r['storage']['read']['mb'] + r['storage']['write']['mb']) / io_time
+        print(f"  Effective Bandwidth:   {bandwidth:>10.1f} MB/s")
+
+        print(f"\n[Storage Details]")
+        print(f"  Blocks in Use:       {r['storage']['total_blocks']:>10,}")
+        print(f"  Free Blocks:         {r['storage']['free_blocks']:>10,}")
+        print(f"  Tokens per Block:    {r['tokens_per_block']:>10,}")
+        print(f"  Block Size:          {r['tokens_per_block'] * r.get('bytes_per_token', 2048) / 1024 / 1024:>10.2f} MB")
+        if 'sync_count' in r['storage']:
+            print(f"  Fsync Operations:    {r['storage']['sync_count']:>10,}")
+
+        print(f"\n[Execution Time]")
+        if r.get('timestamp_replay_enabled'):
+            print(f"  Wall Time (Total):     {r['wall_time_s']:>10.2f} s")
+            print(f"  I/O Time (Actual):     {r['io_time_s']:>10.2f} s")
+            print(f"  Sleep Time (Replay):   {r['wall_time_s'] - r['io_time_s']:>10.2f} s")
+        else:
+            print(f"  Total Execution Time: {r['wall_time_s']:>10.2f} s")
+
+    print(f"\n{'='*80}\n")
+
+
+# ============================================================================
+# Main Program
+# ============================================================================
+
+def main():
+    """Main entry point"""
+    parser = argparse.ArgumentParser(
+        description='Mooncake KVCache Storage Benchmark',
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Quick test (100 requests)
+  python storage_benchmark.py --scenario=toolagent --max-requests=100
+
+  # Test with large model preset (Llama-3.1-405B)
+  python storage_benchmark.py --scenario=toolagent --model=llama-3.1-405b --max-requests=100
+
+  # Test with Deepseek V3 (extra large model)
+  python storage_benchmark.py --scenario=toolagent --model=deepseek-v3 --max-requests=100
+
+  # Realistic replay (with timestamps, 10x speed)
+  python storage_benchmark.py --scenario=toolagent --max-requests=1000 \\
+      --replay-timestamps --time-scale=0.1
+
+  # All scenarios with custom bytes_per_token
+  python storage_benchmark.py --scenario=all --bytes-per-token=512
+
+  # Test with different block sizes and fsync modes
+  python storage_benchmark.py --scenario=toolagent --block-size-tokens=256 --fsync-mode=always
+
+  # Test with custom fsync batch size
+  python storage_benchmark.py --scenario=toolagent --fsync-mode=batch --fsync-batch-size=50
+
+Performance Tuning:
+  --fsync-mode=batch (default): Balance between performance and safety
+  --fsync-mode=always: Safest but slowest, measures full persistence cost
+  --fsync-mode=end: Fastest, only measures write I/O (not persistence)
+  --fsync-mode=none: Testing only, no durability guarantees
+
+Available model presets:
+  llama-3.1-405b, qwen3-32b, deepseek-v3, glm-4.6, default
+
+For more information: tools/STORAGE_BENCHMARK_README.md
+        """
+    )
+
+    parser.add_argument('--trace-dir', type=str, default='../../FAST25-release/traces',
+                       help='Trace files directory')
+    parser.add_argument('--scenario', type=str, choices=['conversation', 'synthetic', 'toolagent', 'all'],
+                       default='toolagent', help='Test scenario')
+    parser.add_argument('--storage-dir', type=str, default='/tmp/mooncake_bench',
+                       help='Storage directory')
+    parser.add_argument('--model', type=str, choices=list(MODEL_BYTES_PER_TOKEN.keys()),
+                       default='default',
+                       help=f'Model preset (overrides --bytes-per-token). Available: {", ".join(MODEL_BYTES_PER_TOKEN.keys())}')
+    parser.add_argument('--bytes-per-token', type=int, default=DEFAULT_BYTES_PER_TOKEN,
+                       help='Bytes per token (default %d, overridden by --model if specified)' % DEFAULT_BYTES_PER_TOKEN)
+    parser.add_argument('--max-requests', type=int, default=None,
+                       help='Maximum number of requests (default: unlimited)')
+    parser.add_argument('--max-blocks', type=int, default=100000,
+                       help='Maximum number of blocks in storage file (determines file size)')
+    parser.add_argument('--replay-timestamps', action='store_true',
+                       help='Enable timestamp replay (simulate realistic request timing)')
+    parser.add_argument('--time-scale', type=float, default=1.0,
+                       help='Time scaling factor (1.0=real-time, 0.1=10x speed, 10.0=0.1x speed)')
+    parser.add_argument('--block-size-tokens', type=int, default=512,
+                       help='Number of tokens per block (default: 512)')
+    parser.add_argument('--fsync-mode', type=str, choices=['batch', 'always', 'end', 'none'],
+                       default='batch',
+                       help='When to fsync: batch=every N writes (default), always=after each write, end=only at close, none=never')
+    parser.add_argument('--fsync-batch-size', type=int, default=100,
+                       help='Number of writes between fsync in batch mode (default: 100)')
+
+    args = parser.parse_args()
+
+    # Print benchmark header
+    print(f"\n{'='*80}")
+    print(f"{'Mooncake KVCache Storage Benchmark':^80}")
+    print(f"{'='*80}")
+
+    # Determine bytes_per_token (model preset takes precedence)
+    bytes_per_token = MODEL_BYTES_PER_TOKEN.get(args.model, args.bytes_per_token)
+    if args.model != 'default':
+        print(f"Using model preset: {args.model} ({bytes_per_token} bytes/token, ~{bytes_per_token/1024:.1f} KB/token)")
+    else:
+        print(f"Using custom bytes_per_token: {bytes_per_token}")
+
+    # Determine test scenarios
+    scenarios = ['conversation', 'synthetic', 'toolagent'] if args.scenario == 'all' else [args.scenario]
+    trace_files = {
+        'conversation': 'conversation_trace.jsonl',
+        'synthetic': 'synthetic_trace.jsonl',
+        'toolagent': 'toolagent_trace.jsonl'
+    }
+
+    # Run benchmarks
+    results = []
+
+    for scenario in scenarios:
+        trace_path = Path(args.trace_dir) / trace_files[scenario]
+        if trace_path.exists():
+            result = run_benchmark(
+                str(trace_path),
+                str(Path(args.storage_dir) / scenario),
+                bytes_per_token,
+                args.max_requests,
+                args.max_blocks,
+                args.replay_timestamps,
+                args.time_scale,
+                args.block_size_tokens,
+                args.fsync_mode,
+                args.fsync_batch_size
+            )
+            results.append(result)
+        else:
+            print(f"Warning: Trace file not found: {trace_path}")
+
+    # Print results
+    if results:
+        print_results(results)
+
+
+if __name__ == '__main__':
+    main()

dependencies.sh

@@ -23,8 +23,7 @@ NC="\033[0m" # No Color
 # Configuration
 REPO_ROOT=`pwd`
 GITHUB_PROXY=${GITHUB_PROXY:-"https://github.com"}
-GOVER=1.23.8
-YALANTINGLIBS_VERSION=0.5.7
+GOVER=1.25.9
 
 # Function to print section headers
 print_section() {
@@ -76,8 +75,7 @@ echo -e "${YELLOW}Mooncake Dependencies Installer${NC}"
 echo -e "This script will install all required dependencies for Mooncake."
 echo -e "The following components will be installed:"
 echo -e "  - System packages (build tools, libraries)"
-echo -e "  - yalantinglibs"
-echo -e "  - Git submodules"
+echo -e "  - Git submodules (including pybind11 and yalantinglibs)"
 echo -e "  - Go $GOVER"
 echo
 
@@ -103,6 +101,7 @@ echo -e "${YELLOW}This may take a few minutes...${NC}"
 
 SYSTEM_PACKAGES="build-essential \
                   cmake \
+                  ninja-build \
                   git \
                   wget \
                   unzip \
@@ -127,6 +126,7 @@ SYSTEM_PACKAGES="build-essential \
                   libmsgpack-dev \
                   libzstd-dev \
                   libasio-dev \
+                  libxxhash-dev \
                   pkg-config \
                   patchelf \
                   libc6-dev \
@@ -136,48 +136,34 @@ apt-get install -y $SYSTEM_PACKAGES
 check_success "Failed to install system packages"
 print_success "System packages installed successfully"
 
-# Install yalantinglibs
+# Initialize and update git submodules
+print_section "Initializing Git Submodules"
+
+# Check if .gitmodules exists
+if [ -f "${REPO_ROOT}/.gitmodules" ]; then
+    echo "Enter repository root: ${REPO_ROOT}"
+    cd "${REPO_ROOT}"
+    check_success "Failed to change to repository root directory"
+
+    echo "Initializing git submodules..."
+    git submodule sync --recursive
+    check_success "Failed to sync git submodules"
+    git submodule update --init --recursive
+    check_success "Failed to initialize git submodules"
+
+    print_success "Git submodules initialized and updated successfully"
+else
+    echo -e "${YELLOW}No .gitmodules file found. Skipping...${NC}"
+    exit 1
+fi
+
+# Build and install yalantinglibs from submodule
 print_section "Installing yalantinglibs"
-
-# Check if thirdparties directory exists
-if [ ! -d "${REPO_ROOT}/thirdparties" ]; then
-    mkdir -p "${REPO_ROOT}/thirdparties"
-    check_success "Failed to create thirdparties directory"
-fi
-
-# Change to thirdparties directory
-cd "${REPO_ROOT}/thirdparties"
-check_success "Failed to change to thirdparties directory"
-
-# Check if yalantinglibs is already installed
-if [ -d "yalantinglibs-${YALANTINGLIBS_VERSION}" ]; then
-    echo -e "${YELLOW}yalantinglibs-${YALANTINGLIBS_VERSION} directory already exists. Removing for fresh install...${NC}"
-    rm -rf yalantinglibs-${YALANTINGLIBS_VERSION}
-    check_success "Failed to remove existing yalantinglibs directory"
-fi
-
-# Download yalantinglibs
-YALANTINGLIBS_ZIPFILE="yalantinglibs-${YALANTINGLIBS_VERSION}.zip"
-echo "Downloading yalantinglibs ${YALANTINGLIBS_VERSION} from ${GITHUB_PROXY}/alibaba/yalantinglibs/archive/refs/tags/${YALANTINGLIBS_VERSION}.zip"
-wget -q --show-progress -O ${YALANTINGLIBS_ZIPFILE} ${GITHUB_PROXY}/alibaba/yalantinglibs/archive/refs/tags/${YALANTINGLIBS_VERSION}.zip
-check_success "Failed to download yalantinglibs"
-
-# Extract yalantinglibs
-echo "Extracting yalantinglibs..."
-unzip -q ${YALANTINGLIBS_ZIPFILE}
-check_success "Failed to extract yalantinglibs"
-
-# Clean up downloaded ZIP file
-rm -f ${YALANTINGLIBS_ZIPFILE}
-check_success "Failed to clean up downloaded ZIP file"
-
-# Build and install yalantinglibs
-cd yalantinglibs-${YALANTINGLIBS_VERSION}
-check_success "Failed to change to yalantinglibs directory"
+cd "${REPO_ROOT}/extern/yalantinglibs"
+check_success "Failed to change to yalantinglibs submodule directory"
 
 mkdir -p build
 check_success "Failed to create build directory"
-
 cd build
 check_success "Failed to change to build directory"
 
@@ -194,32 +180,7 @@ cmake --install .
 check_success "Failed to install yalantinglibs"
 
 print_success "yalantinglibs installed successfully"
-
-# Initialize and update git submodules
-print_section "Initializing Git Submodules"
-
-# Check if .gitmodules exists
-if [ -f "${REPO_ROOT}/.gitmodules" ]; then
-    # Check if submodules are already initialized by looking for the .git directory in the first submodule
-    FIRST_SUBMODULE=$(grep "path" ${REPO_ROOT}/.gitmodules | head -1 | awk '{print $3}')
-
-    echo "Enter repository root: ${REPO_ROOT}"
-    cd "${REPO_ROOT}"
-    check_success "Failed to change to repository root directory"
-
-    if [ -d "${REPO_ROOT}/${FIRST_SUBMODULE}/.git" ] || [ -f "${REPO_ROOT}/${FIRST_SUBMODULE}/.git" ]; then
-        echo -e "${YELLOW}Git submodules already initialized. Skipping...${NC}"
-    else
-        echo "Initializing git submodules..."
-        git submodule update --init
-        check_success "Failed to initialize git submodules"
-
-        print_success "Git submodules initialized and updated successfully"
-    fi
-else
-    echo -e "${YELLOW}No .gitmodules file found. Skipping...${NC}"
-    exit 1
-fi
+cd "${REPO_ROOT}"
 
 print_section "Verifying essential build tools"
 
@@ -236,6 +197,8 @@ print_success "ldd found: $(ldd --version 2>&1 | head -1)"
 
 print_section "Installing Go $GOVER"
 
+USED_CN_MIRROR=false
+
 install_go() {
     ARCH=$(uname -m)
     if [ "$ARCH" = "aarch64" ]; then
@@ -246,18 +209,45 @@ install_go() {
         echo "Unsupported architecture: $ARCH"
         exit 1
     fi
-    # Download Go
-    echo "Downloading Go $GOVER..."
-    wget -q --show-progress https://go.dev/dl/go$GOVER.linux-$ARCH.tar.gz
-    check_success "Failed to download Go $GOVER"
+
+    GO_TARBALL="go$GOVER.linux-$ARCH.tar.gz"
+
+    # Try multiple download mirrors with fallback
+    GO_DOWNLOAD_URLS=(
+        "https://go.dev/dl/${GO_TARBALL}"
+        "https://golang.google.cn/dl/${GO_TARBALL}"
+        "https://mirrors.aliyun.com/golang/${GO_TARBALL}"
+    )
+
+    DOWNLOAD_SUCCESS=false
+    for url in "${GO_DOWNLOAD_URLS[@]}"; do
+        echo "Downloading Go $GOVER from ${url}..."
+        if wget -q --show-progress --timeout=30 --tries=2 -O "${GO_TARBALL}" "${url}"; then
+            DOWNLOAD_SUCCESS=true
+            # If the official source (go.dev) failed and we fell back to a CN mirror,
+            # it likely means the network has restricted access to international sites.
+            if [[ "$url" != "https://go.dev/dl/${GO_TARBALL}" ]]; then
+                USED_CN_MIRROR=true
+            fi
+            print_success "Downloaded Go $GOVER from ${url}"
+            break
+        else
+            echo -e "${YELLOW}Failed to download from ${url}, trying next mirror...${NC}"
+            rm -f "${GO_TARBALL}"
+        fi
+    done
+
+    if [ "$DOWNLOAD_SUCCESS" = false ]; then
+        print_error "Failed to download Go $GOVER from all mirrors"
+    fi
 
     # Install Go
     echo "Installing Go $GOVER..."
-    tar -C /usr/local -xzf go$GOVER.linux-$ARCH.tar.gz
+    tar -C /usr/local -xzf "${GO_TARBALL}"
     check_success "Failed to install Go $GOVER"
 
     # Clean up downloaded file
-    rm -f go$GOVER.linux-$ARCH.tar.gz
+    rm -f "${GO_TARBALL}"
     check_success "Failed to clean up Go installation file"
 
     print_success "Go $GOVER installed successfully"
@@ -283,6 +273,20 @@ if ! grep -q "export PATH=\$PATH:/usr/local/go/bin" ~/.bashrc; then
     echo -e "${YELLOW}Please run 'source ~/.bashrc' or start a new terminal to use Go${NC}"
 fi
 
+# Set GOPROXY only if Go download fell back to a CN mirror, indicating restricted
+# network access to international sites. Skip if user already configured GOPROXY.
+if [ "$USED_CN_MIRROR" = true ] && [ -z "$GOPROXY" ]; then
+    export GOPROXY=https://goproxy.cn,https://goproxy.io,direct
+    echo -e "${YELLOW}Detected restricted network (Go was downloaded from a CN mirror).${NC}"
+    echo -e "${YELLOW}GOPROXY set to: ${GOPROXY}${NC}"
+    if ! grep -q "export GOPROXY=" ~/.bashrc; then
+        echo 'export GOPROXY=https://goproxy.cn,https://goproxy.io,direct' >> ~/.bashrc
+        echo -e "${YELLOW}GOPROXY added to ~/.bashrc for future sessions${NC}"
+    fi
+elif [ -n "$GOPROXY" ]; then
+    echo -e "${GREEN}GOPROXY already set to: ${GOPROXY}${NC}"
+fi
+
 # Return to the repository root
 cd "${REPO_ROOT}"
 

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,17 +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 \
-        python3 \
-        python3-dev \
-        python3-pip \
-        python-is-python3 \
+        ninja-build \
+        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
@@ -44,16 +53,17 @@ RUN bash dependencies.sh -y
 # Configure & build Mooncake
 RUN mkdir -p build && \
     cd build && \
-    cmake .. \
+    cmake -G Ninja .. \
         -DBUILD_UNIT_TESTS=OFF \
         -DUSE_HTTP=ON \
         -DUSE_ETCD=ON \
         -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 . -j"$(nproc)"
+    cmake --build .
 
 # Build nvlink allocator to make wheel self-contained for CUDA paths
 RUN export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH && \
@@ -75,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 \
@@ -88,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/mooncake-store-deployment-guide.md

@@ -8,6 +8,7 @@ This page summarizes useful flags, environment variables, and HTTP endpoints to
   - `--rpc_port` (int, default 50051): RPC listen port.
   - `--rpc_thread_num` (int, default min(4, CPU cores)): RPC worker threads. If not set, uses `--max_threads` (default 4) capped by CPU cores.
   - `--rpc_address` (str, default `0.0.0.0`): RPC bind address.
+  - `--rpc_interface` (str, default empty): Network interface used to resolve the final RPC address. When set, Mooncake Master resolves the interface's current IPv4 address at startup and uses it as the final `rpc_address`. This overrides `--rpc_address`.
   - `--rpc_conn_timeout_seconds` (int, default `0`): RPC idle connection timeout; `0` disables.
   - `--rpc_enable_tcp_no_delay` (bool, default `true`): Enable TCP_NODELAY.
 
@@ -26,8 +27,8 @@ This page summarizes useful flags, environment variables, and HTTP endpoints to
     - `free_ratio_first`: Free-ratio-first strategy. Samples multiple candidates and selects those with highest free space ratio for better load balancing.
 
 - Eviction and TTLs
-  - `--default_kv_lease_ttl` (uint64, default `5000` ms): Default lease TTL for KV objects.
-  - `--default_kv_soft_pin_ttl` (uint64, default `1800000` ms): Soft pin TTL (30 minutes).
+  - `--default_kv_lease_ttl` (duration, default `5000` ms): Default lease TTL for KV objects. The default unit is milliseconds, so `5000` means `5000ms`. Duration strings such as `5000ms`, `5s`, `30m`, or `1h` are also supported.
+  - `--default_kv_soft_pin_ttl` (duration, default `1800000` ms): Soft pin TTL (30 minutes). The default unit is milliseconds, so `1800000` means `1800000ms`. Duration strings such as `1800000ms`, `30m`, or `1h` are also supported.
   - `--allow_evict_soft_pinned_objects` (bool, default `true`): Allow evicting soft-pinned objects.
   - `--eviction_ratio` (double, default `0.05`): Fraction evicted when hitting high watermark.
   - `--eviction_high_watermark_ratio` (double, default `0.95`): Usage ratio to trigger eviction.
@@ -74,6 +75,18 @@ mooncake_master \
   --enable_metric_reporting=true
 ```
 
+Example (resolve the master RPC address from a stable interface name in a container):
+
+```bash
+mooncake_master \
+  --rpc_interface=eth0 \
+  --enable_http_metadata_server=true \
+  --http_metadata_server_host=0.0.0.0 \
+  --http_metadata_server_port=8080
+```
+
+This resolves the current IPv4 address of `eth0` at startup and uses it as the final `rpc_address`.
+
 Example (use free-ratio-first allocation strategy for better load balancing):
 
 ```bash
@@ -92,6 +105,13 @@ mooncake_master \
   --config_path=mooncake-store/conf/master.yaml
 ```
 
+For config files, the equivalent setting is:
+
+```yaml
+rpc_interface: "eth0"
+rpc_port: 50051
+```
+
 ## Metrics Endpoints
 
 The master exposes Prometheus-style metrics over HTTP on `--metrics_port`:
@@ -138,3 +158,13 @@ Available log levels: trace, debug, info, warn (or warning), error, and critical
 - Scale `--rpc_thread_num` with available CPU cores and workload.
 - Start with default eviction settings; adjust `--eviction_high_watermark_ratio` and `--eviction_ratio` based on memory pressure and object churn.
 - Use `/metrics/summary` during bring-up; integrate `/metrics` with Prometheus/Grafana for production.
+
+
+---
+
+:::{toctree}
+:caption: Advanced Topics
+:maxdepth: 1
+
+ssd-offload
+:::

docs/source/deployment/ssd-offload.md

@@ -0,0 +1,280 @@
+# SSD Offload
+
+## Overview
+
+Mooncake Store supports offloading KV cache objects from distributed memory to local SSD. When memory pressure is high, the master instructs clients to persist selected objects to disk. On a cache miss, the client automatically falls back to reading from SSD.
+
+SSD offload is currently **only available in Real Client mode**. The real client is a standalone process that communicates with the application (e.g., SGLang) via RPC. All SSD reads and writes happen within this process.
+
+## Startup Steps
+
+### Step 1: Create the SSD storage directory
+
+```bash
+mkdir -p /nvme/mooncake_offload
+```
+
+### Step 2: Start the master
+
+```bash
+mooncake_master \
+    --rpc_port=50051 \
+    --enable-offload true
+```
+
+### Step 3: Start the real client with SSD offload enabled
+
+Use the `--enable_offload` flag to enable SSD offload, and set environment variables to specify the storage path and backend:
+
+```bash
+export MOONCAKE_OFFLOAD_FILE_STORAGE_PATH=/nvme/mooncake_offload
+export MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR=bucket_storage_backend
+
+mooncake_client \
+    --master_server_address=127.0.0.1:50051 \
+    --host=<machine IP> \
+    --protocol="rdma"  \
+    --device_names=<NIC name, e.g. eth0> \
+    --port=50052 \
+    --global_segment_size="4 GB" \
+    --enable_offload=true  \
+    --metadata_server="P2PHANDSHAKE"
+```
+
+> **Note:** On startup, the real client automatically scans existing SSD data and reports it to the master. No manual recovery is needed.
+
+### Step 4: Connect the application to the real client
+
+The application (e.g., SGLang) connects to the real client via the `MooncakeDistributedStore` Python SDK. SSD offload and fallback loading are handled transparently.
+
+```python
+from mooncake.store import MooncakeDistributedStore
+
+store = MooncakeDistributedStore()
+store.setup(
+    local_hostname="<machine IP>",
+    metadata_server="P2PHANDSHAKE",
+    global_segment_size=4 * 1024 * 1024 * 1024,   # 4 GB
+    local_buffer_size=512 * 1024 * 1024, #512MB
+    protocol="rdma",
+    device_name="eth0",
+    master_server_address="127.0.0.1:50051",
+)
+```
+
+---
+
+## Real Client Parameters
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--master_server_address` | `127.0.0.1:50051` | Master address |
+| `--host` | `0.0.0.0` | This machine's externally reachable IP |
+| `--port` | `50052` | Real client RPC listening port |
+| `--device_names` | ` ` | NIC name(s), e.g. `eth0` or `mlx5_0` |
+| `--protocol` | `tcp` | Transport protocol: `tcp` or `rdma` |
+| `--global_segment_size` | `4 GB` | Memory pool size allocated for this node |
+| `--enable_offload` | `false` | **Must be set to `true` to enable SSD offload** |
+| `--threads` | `1` | Number of RPC server threads |
+
+---
+
+## SSD Offload Configuration
+
+### Core settings
+
+| Environment Variable | Default | Description |
+|---|---|---|
+| `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` | `/data/file_storage` | Absolute path to the SSD storage directory |
+| `MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR` | `bucket_storage_backend` | Storage backend type (see below) |
+| `MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES` | `1342177280` (1.25 GB) | Client-side staging buffer size |
+| `MOONCAKE_OFFLOAD_SCANMETA_ITERATOR_KEYS_LIMIT` | `20000` | Max keys processed per iteration when scanning existing SSD metadata on startup |
+| `MOONCAKE_OFFLOAD_TOTAL_SIZE_LIMIT_BYTES` | `2199023255552` (2 TB) | Maximum disk usage |
+| `MOONCAKE_OFFLOAD_TOTAL_KEYS_LIMIT` | `10000000` | Maximum number of objects on disk |
+| `MOONCAKE_OFFLOAD_HEARTBEAT_INTERVAL_SECONDS` | `10` | Interval for offload heartbeat to master (seconds) |
+| `MOONCAKE_OFFLOAD_USE_URING` | `false` | Enable io_uring for async file I/O |
+
+### Bucket backend settings
+
+Applies when `MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR=bucket_storage_backend`.
+
+| Environment Variable | Default | Description |
+|---|---|---|
+| `MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES` | `268435456` (256 MB) | Max size per bucket |
+| `MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT` | `500` | Max keys per bucket |
+| `MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE` | `0` | Eviction threshold in bytes. When set to `0`, the backend uses **90% of the physical disk capacity** as the quota — it does not mean unlimited. Set an explicit value to control disk usage precisely. |
+| `MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY` | `none` | Eviction policy: `none` / `fifo` / `lru` |
+
+---
+
+## Storage Backends
+
+### `bucket_storage_backend` (recommended)
+
+Groups multiple objects into bucket files. Reduces filesystem overhead, supports efficient batch I/O, and supports FIFO and LRU eviction.
+
+**File layout:**
+```
+/nvme/mooncake_offload/
+├── 1710000000000-0.bucket   # data file (multiple KV pairs)
+├── 1710000000000-0.meta     # metadata file
+├── 1710000000001-0.bucket
+└── ...
+```
+
+Best for: general-purpose use, large-scale deployments.
+
+### `file_per_key_storage_backend`
+
+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`
+
+Pre-allocates a single large file and manages offset-based allocation within it. Highest concurrency via 1024-shard metadata.
+
+> **Warning:** This backend does **not** support metadata recovery on restart. On initialization, the data file is truncated and all in-memory metadata is cleared. Any previously offloaded objects become inaccessible after a process restart.
+
+**Capacity:** `MOONCAKE_OFFLOAD_TOTAL_SIZE_LIMIT_BYTES` is used directly as the pre-allocated file size (100%, no safety margin). Unlike `bucket_storage_backend`, there is no separate quota variable — this is the sole disk usage control. Set it below the physical disk capacity to avoid filling the disk; writes are rejected once usage reaches this limit.
+
+Best for: high-concurrency scenarios with many small objects where restart durability is not required.
+
+---
+
+## Eviction (Bucket Backend Only)
+
+When `MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE` is set, the backend automatically evicts buckets before writing new ones if total disk usage would exceed the limit.
+
+| Policy | Behavior |
+|--------|----------|
+| `none` | No eviction (default); writes fail when disk is full |
+| `fifo` | Evict the oldest bucket first |
+| `lru` | Evict the least recently read bucket first |
+
+Eviction is two-phase: the bucket is removed from metadata and master is notified first, then in-flight reads are drained before files are deleted.
+
+---
+
+## Example
+
+The following example starts a master and a real client on a single machine.
+
+### Environment
+
+- Machine IP: `192.168.1.10`
+- NIC: `eth0`
+- SSD mount point: `/nvme`
+- Memory pool size: 4 GB (smaller than the total data written, to trigger offload)
+
+### Start the master
+
+```bash
+mooncake_master \
+    --rpc_port=50051
+```
+
+### Start the real client (new terminal)
+
+```bash
+export MOONCAKE_OFFLOAD_FILE_STORAGE_PATH=/nvme/mooncake_offload
+export MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR=bucket_storage_backend
+export MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE=$((200 * 1024 * 1024 * 1024))  # 200 GB
+export MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY=lru
+
+mooncake_client \
+    --master_server_address="192.168.1.10:50051" \
+    --host="192.168.1.10" \
+    --device_names="eth0" \
+    --port=50052 \
+    --protocol="rdma" \
+    --global_segment_size="4GB" \
+    --enable_offload="true"
+```
+
+---
+
+## Notes
+
+- `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` must be an absolute path to an existing, writable directory. Symbolic links and paths containing `..` are rejected.
+- On real client restart, the backend automatically scans existing SSD files and reports them to the master, so previously offloaded objects remain accessible.
+- Eviction only notifies the master and deletes local files; objects replicated on other nodes are unaffected.
+- Each machine requires its own real client process. In multi-node deployments, ensure `--host` and `--port` are correctly set so nodes can reach each other.
+
+**2-node example:** suppose Node A (`192.168.1.10`) runs the master and Node B (`192.168.1.11`) is a second worker. Both real clients must point to the same master and advertise their own externally reachable IP:
+
+```bash
+# Node A — runs the master and its own real client
+mooncake_master --rpc_port=50051 --enable-offload true &
+
+export MOONCAKE_OFFLOAD_FILE_STORAGE_PATH=/nvme/mooncake_offload
+mooncake_client \
+    --master_server_address="192.168.1.10:50051" \
+    --host="192.168.1.10" \        # externally reachable IP of Node A
+    --device_names="eth0" \
+    --protocol="rdma" \
+    --metadata_server="P2PHANDSHAKE" \
+    --port=50052 \
+    --global_segment_size="4GB" \
+    --enable_offload="true"
+```
+
+```bash
+# Node B — real client only; points to the same master on Node A
+export MOONCAKE_OFFLOAD_FILE_STORAGE_PATH=/nvme/mooncake_offload
+mooncake_client \
+    --master_server_address="192.168.1.10:50051" \
+    --host="192.168.1.11" \        # externally reachable IP of Node B, NOT 127.0.0.1
+    --device_names="eth0" \
+    --protocol="rdma" \
+    --metadata_server="P2PHANDSHAKE" \
+    --port=50052 \
+    --global_segment_size="4GB" \
+    --enable_offload="true"
+```
+
+---
+
+## Troubleshooting
+
+### SSD offload is not triggering
+
+- Confirm `--enable_offload=true` is passed to `mooncake_client` and `--enable-offload true` is passed to `mooncake_master`.
+- Check that `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` points to an existing, writable directory. The client will fail silently if the path is invalid.
+- Verify memory pressure is actually high enough for the master to trigger offload. If the memory pool (`--global_segment_size`) is large relative to the data written, offload may never activate.
+
+### "Permission denied" or "No such file or directory" on the storage path
+
+- Ensure the directory exists before starting the client: `mkdir -p <path>`.
+- Confirm the process user has read/write access to the directory.
+- Symbolic links and paths containing `..` are rejected — use an absolute, canonical path.
+
+### "Failed to register buffer with UringFile" warning in logs
+
+This warning appears when `MOONCAKE_OFFLOAD_USE_URING=true` and the io_uring fixed-buffer registration fails. The most common cause is that `MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES` exceeds the process's locked-memory limit (`RLIMIT_MEMLOCK`). io_uring requires the registered buffer to be pinned in physical memory, which counts against this limit.
+
+Check the current limit:
+
+```bash
+ulimit -l        # in KB; "unlimited" means no cap
+```
+
+To raise it for the current session:
+
+```bash
+ulimit -l unlimited
+```
+
+To raise it permanently, add the following to `/etc/security/limits.conf`:
+
+```
+* soft memlock unlimited
+* hard memlock unlimited
+```
+
+Alternatively, reduce `MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES` to a value within the existing limit. Note that the warning does not abort startup — the client falls back to non-fixed-buffer I/O — but performance may be lower than expected.

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.
@@ -100,10 +101,30 @@ The data structure details of `ReplicateConfig` are as follows:
 struct ReplicateConfig {
     size_t replica_num{1};                    // Total number of replicas for the object
     bool with_soft_pin{false};               // Whether to enable soft pin mechanism for this object
+    bool with_hard_pin{false};               // Whether to enable hard pin (never evicted)
     std::string preferred_segment{};         // Preferred segment for allocation
 };
 ```
 
+### Upsert
+
+```C++
+tl::expected<void, ErrorCode> Upsert(const ObjectKey& key,
+                                     std::vector<Slice>& slices,
+                                     const ReplicateConfig& config);
+
+std::vector<tl::expected<void, ErrorCode>> BatchUpsert(
+    const std::vector<ObjectKey>& keys,
+    std::vector<std::vector<Slice>>& batched_slices,
+    const ReplicateConfig& config);
+```
+
+`Upsert` inserts `key` if it does not exist and updates the existing object if
+it does. It uses the same replication configuration model as `Put`, while
+allowing the store to reuse existing placement for in-place updates when the
+current layout permits it. `BatchUpsert` performs the same operation for
+multiple keys using a shared replication configuration.
+
 ### Remove
 
 ```C++
@@ -515,6 +536,40 @@ The Master Service handles object-related interfaces as follows:
 
 Before writing an object, the Client calls PutStart to request storage space allocation from the Master Service. After completing data writing, the Client calls PutEnd to notify the Master Service to mark the object write as completed.
 
+- Upsert
+
+```C++
+tl::expected<std::vector<Replica::Descriptor>, ErrorCode> UpsertStart(
+    const std::string& key,
+    const std::vector<size_t>& slice_lengths,
+    const ReplicateConfig& config);
+
+std::vector<tl::expected<std::vector<Replica::Descriptor>, ErrorCode>>
+BatchUpsertStart(const std::vector<std::string>& keys,
+                 const std::vector<std::vector<uint64_t>>& slice_lengths,
+                 const ReplicateConfig& config);
+
+tl::expected<void, ErrorCode> UpsertEnd(
+    const std::string& key, ReplicaType replica_type);
+
+std::vector<tl::expected<void, ErrorCode>> BatchUpsertEnd(
+    const std::vector<std::string>& keys);
+
+tl::expected<void, ErrorCode> UpsertRevoke(
+    const std::string& key, ReplicaType replica_type);
+
+std::vector<tl::expected<void, ErrorCode>> BatchUpsertRevoke(
+    const std::vector<std::string>& keys);
+```
+
+`UpsertStart` / `UpsertEnd` / `UpsertRevoke` mirror the existing put lifecycle
+but operate on insert-or-update semantics. If the key does not exist, the flow
+behaves like `PutStart`. If the key already exists, the Master may reuse the
+current allocation for an in-place update or allocate new space when the object
+layout changes. The batch variants provide the same control flow for multiple
+keys and are the lower-level primitives used by the high-level `BatchUpsert`
+path.
+
 - GetReplicaList
 
 ```C++
@@ -688,6 +743,18 @@ There are two startup parameters in `master_service` related to the soft pin mec
 
 Notably, soft pinned objects can still be removed using APIs such as `Remove` or `RemoveAll`.
 
+### Hard Pin
+
+For objects that must never be evicted under any circumstances (e.g., model weights, critical metadata), Mooncake Store provides a hard pin mechanism. Unlike soft pin, hard-pinned objects are permanently protected from eviction — they will never be selected as eviction candidates regardless of memory pressure.
+
+Hard pin is set at object creation time through the `with_hard_pin` field in `ReplicateConfig` and cannot be changed afterward. Hard-pinned objects can only be removed explicitly via `Remove` (with force) or `RemoveAll`.
+
+Key differences from soft pin:
+
+- Hard pin never expires. Soft pin status is removed after a configurable TTL if the object is not accessed.
+- Hard-pinned objects are completely skipped during eviction. Soft-pinned objects may still be evicted when no other candidates are available.
+- Hard pin is immutable once set. Soft pin status is automatically refreshed on access.
+
 ### Zombie Object Cleanup
 
 If a Client crashes or experiences a network failure after sending a `PutStart` request but before it can send the corresponding `PutEnd` or `PutRevoke` request to the Master, the object initiated by `PutStart` enters a "zombie" state—rendering it neither usable nor deletable. The existence of such "zombie objects" not only consumes storage space but also prevents subsequent `Put` operations on the same keys. To mitigate these issues, the Master records the start time of each `PutStart` request and employs two timeout thresholds—`put_start_discard_timeout` and `put_start_release_timeout`—to clean up zombie objects.
@@ -712,6 +779,7 @@ The preferred segment allocation feature is implemented through the `AllocationS
 struct ReplicateConfig {
     size_t replica_num{1};                    // Total number of replicas for the object
     bool with_soft_pin{false};               // Whether to enable soft pin mechanism for this object
+    bool with_hard_pin{false};               // Whether to enable hard pin (never evicted)
     std::string preferred_segment{};         // Preferred segment for allocation
 };
 ```
@@ -975,3 +1043,13 @@ When to bump the version:
 * **Major version (X.0.0)**: For breaking API changes, major architectural changes, or significant new features that affect backward compatibility
 * **Minor version (0.X.0)**: For new features, API additions, or notable improvements that maintain backward compatibility
 * **Patch version (0.0.X)**: For bug fixes, performance optimizations, or minor improvements that don't affect the API
+
+
+---
+
+:::{toctree}
+:caption: Related Design Docs
+:maxdepth: 1
+
+ssd-offload
+:::

docs/source/design/ssd-offload.md

@@ -0,0 +1,245 @@
+# SSD Offload Design
+
+## Overview
+
+Mooncake Store supports offloading KV cache objects from distributed memory to local SSD. This extends the effective cache capacity beyond DRAM limits at lower cost, while preserving the performance characteristics of the hot path through zero-copy RDMA-based memory transfers.
+
+SSD offload is implemented as a background subsystem within the **real client** process. It is transparent to the application: a `Put` that would otherwise be evicted from memory is persisted to disk, and a `Get` that finds no memory replica automatically falls back to reading from SSD.
+
+---
+
+## Architecture
+
+```
+┌─────────────────────────────────────────────────────────┐
+│                  Application (vLLM, etc.)               │
+└──────────────────────────┬──────────────────────────────┘
+                           │ MooncakeDistributedStore API
+                           ▼
+┌─────────────────────────────────────────────────────────┐
+│                      Real Client                        │
+│                                                         │
+│  ┌──────────────────────────────────────────────────┐   │
+│  │                  FileStorage                     │   │
+│  │  ┌────────────┐  ┌──────────────────────────┐    │   │
+│  │  │  Heartbeat │  │   ClientBuffer (staging) │    │   │
+│  │  │   Thread   │  └──────────────────────────┘    │   │
+│  │  └─────┬──────┘                                  │   │
+│  │        │ offload / load                          │   │
+│  │        ▼                                         │   │
+│  │  ┌─────────────────────────────────────────┐     │   │
+│  │  │        StorageBackendInterface          │     │   │
+│  │  │  ┌───────────┐ ┌──────────┐ ┌────────┐  │     │   │
+│  │  │  │  Bucket   │ │FilePerKey│ │Offset  │  │     │   │
+│  │  │  │  Backend  │ │ Backend  │ │Alloc.  │  │     │   │
+│  │  │  └───────────┘ └──────────┘ └────────┘  │     │   │
+│  │  └─────────────────────────────────────────┘     │   │
+│  └──────────────────────────────────────────────────┘   │
+│                                                         │
+│  ┌──────────────────────────────────────────────────┐   │
+│  │         In-memory distributed KV cache           │   │
+│  │              (Transfer Engine / RDMA)            │   │
+│  └──────────────────────────────────────────────────┘   │
+└─────────────────────────────────────────────────────────┘
+                           │
+                    Local SSD / NVMe
+```
+
+The key components are:
+
+- **FileStorage**: The top-level coordinator. It owns the storage backend, a staging buffer (`ClientBuffer`), and background threads for heartbeating and buffer garbage collection.
+- **StorageBackendInterface**: An abstract interface implemented by three backends (see below). Responsible for the actual on-disk layout and I/O.
+- **Heartbeat thread**: Periodically contacts the master. The master returns a list of objects to offload; the heartbeat thread writes them to SSD and notifies the master of completion.
+- **ClientBuffer**: A pre-registered, O_DIRECT-aligned staging area used for zero-copy reads from SSD back into application memory.
+
+---
+
+## Data Flow
+
+### Offload (memory → SSD)
+
+The offload path is driven entirely by the heartbeat thread inside `FileStorage`. No write path from the application is involved.
+
+```
+Heartbeat Thread          Master                  Local Memory Segment
+      │                     │                             │
+      │─OffloadObjectHB ───▶│                             │
+      │◀─ {key→size} map ───│  (objects to evict from     │
+      │                     │   memory to SSD)            │
+      │                     │                             │
+      │─ BatchQuery(keys) ───────────────────────────────▶│
+      │◀─ {key→Slice} ────────────────────────────────────│
+      │                     │                             │
+      │  [PrepareEviction: remove old buckets, notify master via BatchEvictDiskReplica]
+      │─ BatchEvictDiskReplica(evicted_keys) ────────────▶│ (master removes stale replicas)
+      │  [FinalizeEviction: delete evicted files]
+      │                     │                             │
+      │  BatchOffload(slices) → StorageBackend → SSD      │
+      │                     │                             │
+      │─ NotifyOffloadSuccess(keys, metadata) ───────────▶│
+      │                     │  (master adds LOCAL_DISK    │
+      │                     │   replica to object entry)  │
+```
+
+Step by step:
+
+1. **Heartbeat**: The heartbeat thread wakes up every `MOONCAKE_OFFLOAD_HEARTBEAT_INTERVAL_SECONDS` seconds and calls `client_->OffloadObjectHeartbeat(enable_offloading_, offloading_objects)`. The master replies with a map of `{key → size}` for objects it has selected to evict from memory.
+2. **Read slices from memory**: `FileStorage::OffloadObjects` groups the keys into buckets (for `BucketStorageBackend`) and calls `BatchQuerySegmentSlices` to obtain `{key → Slice}` from the local memory segment via `client_->BatchQuery`.
+3. **Eviction** (if capacity limit is set): Before writing, `PrepareEviction` removes old buckets from metadata under the exclusive lock and collects their keys. The `eviction_handler` callback calls `client_->BatchEvictDiskReplica` to notify the master in a single RPC. `FinalizeEviction` then deletes the corresponding files.
+4. **Write to SSD**: `StorageBackend::BatchOffload` serializes and writes the key-value data to disk.
+5. **Notify master**: On success, the `complete_handler` calls `client_->NotifyOffloadSuccess(keys, metadatas)`. The master adds a `LOCAL_DISK` replica entry (carrying the real client's RPC address as `transport_endpoint`) to the object's replica list.
+
+### Load (SSD → memory)
+
+The load path involves three parties: the **requesting client**, the **target client** that holds the SSD data, and the **Transfer Engine** for zero-copy data movement.
+
+```
+Requesting Client                 Target Client                    Master
+       │                                    │                         │
+       │─ BatchGet(keys) ──────────────────────────────────────────▶│
+       │◀─ QueryResult {replicas: [LOCAL_DISK(rpc_addr)]} ──────────│
+       │                                    │                         │
+       │  (no memory replica available)     │                         │
+       │─ batch_get_offload_object(keys, sizes) ───────────────────▶│
+       │                                    │                         │
+       │                    FileStorage::BatchGet                     │
+       │                    → StorageBackend::BatchLoad               │
+       │                    → read from SSD into ClientBuffer         │
+       │                                    │                         │
+       │◀─ BatchGetOffloadObjectResponse ───│                         │
+       │   {batch_id, pointers[], transfer_engine_addr, gc_ttl_ms}   │
+       │                                    │                         │
+       │─ Transfer Engine: BatchGetOffloadObject ──────────────────▶│
+       │  (RDMA/TCP: pull data from ClientBuffer into app memory)     │
+       │◀─ done ────────────────────────────│                         │
+       │                                    │                         │
+       │─ release_offload_buffer(batch_id) ────────────────────────▶│
+       │                                    │ (free ClientBuffer slot)│
+```
+
+Step by step:
+
+1. **Query master**: The requesting client calls `client_->BatchGet(keys, ...)` to query the master for replica locations. If the object has been offloaded, the master returns a `LOCAL_DISK` replica descriptor containing the target client's RPC address (`transport_endpoint`).
+2. **RPC to target client**: The requesting client calls `batch_get_offload_object(keys, sizes)` on the target client identified by `transport_endpoint`. The target client calls `FileStorage::BatchGet`, which allocates slots in `ClientBuffer` and reads the requested objects from SSD via `StorageBackend::BatchLoad`.
+3. **Response with buffer pointers**: The target client returns a `BatchGetOffloadObjectResponse` containing `batch_id`, a list of buffer `pointers` (addresses within `ClientBuffer`), the Transfer Engine address, and `gc_ttl_ms` (the buffer lease TTL).
+4. **Zero-copy transfer**: The requesting client invokes `client_->BatchGetOffloadObject(transfer_engine_addr, keys, pointers, slices)`, which uses the Transfer Engine (RDMA or TCP) to pull the data directly from the target client's `ClientBuffer` into the application's target memory (DRAM or VRAM). No intermediate copy is made on the requesting client side.
+5. **Release buffer**: After the transfer completes, the requesting client immediately calls `release_offload_buffer(batch_id)` on the target client to free the `ClientBuffer` slots. If the transfer takes longer than `gc_ttl_ms`, the buffer GC thread reclaims the slot automatically as a fallback.
+
+---
+
+## Storage Backends
+
+### BucketStorageBackend (default)
+
+Objects are grouped into **buckets** before being written to disk. Each bucket produces two files:
+
+- **`.bucket`** — binary data file containing serialized key-value records
+- **`.meta`** — metadata file describing the keys and byte offsets within the data file
+
+Bucket IDs are monotonically increasing timestamps with a sequence suffix, so `buckets_` (a `std::map<int64_t, BucketMetadata>`) is always ordered by creation time.
+
+**Grouping strategy** (`GroupOffloadingKeysByBucket`): objects are accumulated into a bucket until either `bucket_size_limit` (default 256 MB) or `bucket_keys_limit` (default 500) is reached. Objects that do not fill a complete bucket are held in `ungrouped_offloading_objects_` and retried on the next heartbeat.
+
+**In-flight read tracking**: A `BucketReadGuard` RAII object increments `BucketMetadata::inflight_reads_` on construction and decrements it on destruction. This allows safe deletion of bucket files even when concurrent reads are in progress.
+
+### StorageBackendAdaptor (FilePerKey)
+
+Each object is stored as an individual file. The file path is derived from the key via a two-level hash-sharded directory structure to avoid large flat directories. This backend is simple and easy to inspect but does not scale well to millions of objects.
+
+### OffsetAllocatorStorageBackend
+
+A single pre-allocated file (`kv_cache.data`) is shared by all objects. Space within the file is managed by an `OffsetAllocator`. Metadata is sharded across 1024 independent maps to reduce lock contention under high concurrency. Records follow the layout `[key_len: u32 | value_len: u32 | key | value]`.
+
+---
+
+## Eviction (BucketStorageBackend)
+
+When `MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE` is set, the backend evicts existing buckets to make room before writing a new one. Eviction is disabled by default (`BucketEvictionPolicy::NONE`).
+
+### Policies
+
+| Policy | Candidate selection |
+|--------|---------------------|
+| `FIFO` | `buckets_.begin()` — always the oldest bucket, since `buckets_` is ordered by bucket ID |
+| `LRU` | `std::min_element` over `BucketMetadata::last_access_ns_` — the bucket with the smallest last-read timestamp |
+
+`last_access_ns_` is an atomic `int64_t` updated on every `BatchLoad` with relaxed ordering. Buckets that have never been read have `last_access_ns_ == 0` and are therefore always evicted first under LRU, giving FIFO-among-unread semantics.
+
+### Two-phase eviction protocol
+
+Eviction is split into two phases to ensure that the master is notified before files are deleted, and that no in-flight reads are interrupted.
+
+**Phase 1 — `PrepareEviction(required_size)`** (called under exclusive lock):
+
+1. Repeatedly call `SelectEvictionCandidate()` until `total_size_ + required_size <= max_total_size`.
+2. For each selected bucket: remove it from `buckets_` and `object_bucket_map_`, subtract its size from `total_size_`.
+3. Collect all evicted keys and bucket metadata into a `PendingEviction` struct and return it — no file I/O at this point.
+
+**Between phases** — notify master:
+
+The caller invokes the `eviction_handler` callback with the full list of evicted keys. The handler calls `MasterClient::BatchEvictDiskReplica`, which sends a single RPC to the master to remove the disk replicas for all evicted keys atomically.
+
+**Phase 2 — `FinalizeEviction(pending)`** (called after master notification):
+
+For each evicted bucket:
+1. Spin-wait (with a 10-second timeout) until `inflight_reads_ == 0`.
+2. Evict any stale file-handle cache entries.
+3. Delete the `.bucket` and `.meta` files.
+
+This ordering guarantees:
+- The master never serves a stale disk-replica location for a file that has already been deleted.
+- Ongoing reads complete successfully before their files are removed.
+- Freed disk space is available for the incoming write before `WriteBucket` is called.
+
+---
+
+## io_uring File I/O
+
+When `MOONCAKE_OFFLOAD_USE_URING=true`, the storage backends replace POSIX `pread`/`pwrite` calls with an io_uring-based implementation (`UringFile`). The design prioritizes eliminating inter-thread lock contention, which was the dominant latency source in the previous global-ring approach.
+
+### Thread-local rings (`SharedUringRing`)
+
+Each thread owns exactly one `io_uring` ring, stored in `thread_local` storage. This means:
+
+- **No mutex between threads.** Each ring is accessed only by its owning thread, so concurrent I/O from multiple threads is fully parallel with zero synchronization overhead.
+- **Within-thread batching.** Multiple SQEs can be enqueued before calling `io_uring_submit_and_wait`, exposing NVMe queue depth > 1 within a single thread. `batch_read` exploits this to issue up to `QUEUE_DEPTH` (32) independent reads in one submission.
+- **File-descriptor registration is omitted.** The per-I/O `fdget()` overhead (~50 ns) is negligible compared to the lock contention (> 1 ms) the old global ring imposed, so `IOSQE_FIXED_FILE` is not used.
+
+Rings are initialized lazily on first use and destroyed when the thread exits. If ring initialization fails (e.g., kernel too old), the backend falls back gracefully to POSIX I/O.
+
+### Fixed-buffer registration
+
+The `ClientBuffer` (the staging buffer used for SSD reads) is registered with io_uring as a **fixed buffer** via `io_uring_register_buffers`. When a read destination falls within the registered region, the backend uses `io_uring_prep_read_fixed` instead of `io_uring_prep_read`, which avoids a per-I/O `mmap`/`munmap` in the kernel and reduces system-call overhead.
+
+Buffer registration is global but applied **lazily per thread**: `g_buf` stores the base address and length atomically; each thread-local ring calls `ensure_buf_registered()` on its first I/O and registers the buffer independently. This avoids a global barrier at startup.
+
+To prevent `io_uring`'s `FOLL_LONGTERM` page pinning from failing on systems with Transparent Huge Pages (THP) enabled, `MADV_NOHUGEPAGE` is applied to the buffer region before registration. This forces the kernel to back the range with 4 KB pages, making long-term pinning reliable regardless of system THP policy.
+
+### O_DIRECT and alignment
+
+`UringFile` supports an optional `O_DIRECT` mode. When enabled:
+
+- All file descriptors are opened with `O_DIRECT`.
+- Buffers, lengths, and offsets must be aligned to 4 KB (`ALIGNMENT_ = 4096`).
+- For unaligned writes (e.g., metadata serialized into a `std::string`), the backend allocates a temporary aligned bounce buffer via `posix_memalign`, copies the data, performs the aligned write, and frees the bounce buffer.
+- `read_aligned` and `write_aligned` are the primary I/O paths; they assert alignment constraints and delegate directly to the ring.
+
+### I/O operations
+
+| Method | Description |
+|--------|-------------|
+| `read` / `write` | Contiguous read or write, chunked into up to `QUEUE_DEPTH` SQEs per submission |
+| `read_aligned` / `write_aligned` | Same as above but with alignment preconditions for O_DIRECT |
+| `batch_read` | Submits multiple independent reads (different offsets) in batches of up to `QUEUE_DEPTH`, maximizing NVMe queue utilization |
+| `vector_read` / `vector_write` | Scatter/gather I/O: one SQE per `iovec`, submitted in batches |
+| `datasync` | Issues `IORING_FSYNC_DATASYNC` and waits for completion |
+
+### Integration with storage backends
+
+- **BucketStorageBackend**: uses `UringFile` for both bucket data files and metadata files when `use_uring_` is set. A file-handle cache (`file_cache_`) avoids repeated `open`/`close` for hot buckets. On eviction, the cache entry is explicitly removed before the file is deleted to prevent stale handles.
+- **OffsetAllocatorStorageBackend**: opens the single pre-allocated data file with `O_DIRECT` and `UringFile`, and uses `GetFileInstance()` to expose the file handle for external buffer registration.
+- **StorageBackendAdaptor** (FilePerKey): uses `UringFile` for reads when `use_uring_` is set; writes use POSIX paths.
+
+## Metadata Recovery on Restart
+
+On startup, `FileStorage::Init` calls `StorageBackend::ScanMeta`, which reads all on-disk metadata and invokes a callback for each discovered object. The callback calls `MasterClient::NotifyOffloadSuccess` to re-register the objects with the master. This restores the full disk-replica view without any application-level intervention.

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

@@ -2,6 +2,8 @@
 
 The source code path for Ascend Transport is `Mooncake/mooncake-transfer-engine/src/transport/ascend_transport`, which also includes automated build scripts and the README file.
 
+**ASCEND TRANSPORT is scheduled for deprecation, please use [ASCEND DIRECT TRANSPORT](./ascend_direct_transport.md) on ASCEND platform. **
+
 ## Overview
 
 Ascend Transport is a high-performance zero-copy NPU data transfer library with one-sided semantics, directly compatible with Mooncake Transfer Engine. To compile and use the Ascend Transport library, please set the `USE_ASCEND` flag to `"ON"` in the `mooncake-common/common.cmake` file.
@@ -142,7 +144,7 @@ Therefore, in testing:
    Watch the log produced by `mooncake-transfer-engine/src/transfer_engine.cpp`; you should see a line similar to  
    ```
    Transfer Engine RPC using <protocol> listening on <IP>:<actual-port>
-   ```  
+   ```
    Note the **actual port** the target node is listening on.
 
 2. **Edit the initiator’s launch command**:  

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,46 +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 \
-    libyaml-cpp-dev \
-    libgtest-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
@@ -76,12 +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-asio/libasio.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
@@ -102,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):
@@ -184,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
@@ -193,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
 
@@ -287,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.
 
@@ -303,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.
@@ -351,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,15 +187,25 @@ 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
 - `-DWITH_STORE=[ON|OFF]`: Build Mooncake Store component
 - `-DWITH_P2P_STORE=[ON|OFF]`: Enable Golang support and build P2P Store component, require go 1.23+
 - `-DWITH_WITH_RUST_EXAMPLE=[ON|OFF]`: Enable Rust support
+- `-DWITH_EP=[ON|OFF]`: Build the EP (Expert Parallelism) and PG Python extensions for CUDA. Requires CUDA toolkit and PyTorch. Use `-DEP_TORCH_VERSIONS="2.9.1"` (semicolon-separated) to build for specific PyTorch versions, or leave empty to use the currently-installed torch. The CUDA version is detected automatically. **Default: OFF.**
 - `-DUSE_REDIS=[ON|OFF]`: Enable Redis-based metadata service
 - `-DUSE_HTTP=[ON|OFF]`: Enable Http-based metadata service
 - `-DUSE_ETCD=[ON|OFF]`: Enable etcd-based metadata service, require go 1.23+

docs/source/getting_started/quick-start.md

@@ -200,6 +200,8 @@ mooncake_master \
 ```
 This exposes the metadata endpoint at `http://<host>:<port>/metadata`.
 
+If the master runs in a container and its IP is dynamic, set `--rpc_interface=<ifname>` such as `--rpc_interface=eth0`. Mooncake Master will resolve the current IPv4 address from that interface at startup instead of relying on a fixed `--rpc_address`.
+
 Optional: Use the free-ratio-first allocation strategy for better load balancing across segments with different sizes or utilization:
 
 ```bash
@@ -243,4 +245,4 @@ store.close()
 
 ### More Examples and Documentation
 
-Please refer to the [Mooncake Store Python API](../python-api-reference/mooncake-store.md), [Mooncake Store](../design/mooncake-store.md) and [Mooncake Store Deployment & Operations Guide](../deployment/mooncake-store-deployment-guide.md) for more examples and documentation.
+Please refer to the [Mooncake Store Python API](../python-api-reference/mooncake-store.md), [Mooncake Store](../design/mooncake-store.md) and [Mooncake Store Deployment & Operations Guide](../deployment/mooncake-store-deployment-guide.md) for more examples and documentation.

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

@@ -27,6 +27,7 @@ This repository also hosts its technical report and the open-sourced traces.
 
 <h2 id="updates">🔄 Updates</h2>
 
+ - **Mar 19, 2026**: [TorchSpec: Speculative Decoding Training at Scale](https://pytorch.org/blog/torchspec-speculative-decoding-training-at-scale) is [open sourced](https://github.com/torchspec-project/TorchSpec), using Mooncake to decouple inference and training via efficient hidden states management.
  - **Feb 12, 2026**: [Mooncake Joins PyTorch Ecosystem](https://pytorch.org/blog/mooncake-joins-pytorch-ecosystem/) We are thrilled to announce that Mooncake has officially joined the PyTorch Ecosystem!
  - **Jan 28, 2026**: [FlexKV](https://github.com/taco-project/FlexKV), a distributed KV store and cache system from Tencent and NVIDIA in collaboration with the community, now supports [distributed KVCache reuse](https://github.com/taco-project/FlexKV/blob/main/docs/dist_reuse/README_en.md) with the Mooncake Transfer Engine.
  - **Dec 23, 2025**: SGLang introduces [Encode-Prefill-Decode (EPD) Disaggregation](https://lmsys.org/blog/2026-01-12-epd/) with Mooncake as a transfer backend. This integration allows decoupling compute-intensive multimodal encoders (e.g., Vision Transformers) from language model nodes, utilizing Mooncake's RDMA engine for zero-copy transfer of large multimodal embeddings.
@@ -85,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
@@ -128,7 +130,7 @@ troubleshooting/troubleshooting
 
 :::{toctree}
 :caption: Deployment
-:maxdepth: 1
+:maxdepth: 2
 
 deployment/mooncake-store-deployment-guide
 :::

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/performance/storage-benchmark.md

@@ -0,0 +1,163 @@
+# Mooncake KVCache Storage Benchmark
+
+High-performance KVCache storage benchmark tool based on Mooncake Store architecture.
+
+## Overview
+
+Evaluates I/O performance of KVCache storage systems using:
+- Single large file (100GB) with offset-based block management
+- Prefix caching simulation with hash-based block lookup
+- Timestamp-based request replay for realistic testing
+- Comprehensive metrics: latency, bandwidth, hit rates
+
+## Test Flow
+
+1. **Load Traces**: Read request sequences from JSONL files (`FAST25-release/traces`)
+2. **Process Requests**: For each request, check hash_id prefix cache hits/misses
+3. **Perform I/O**: Read cached blocks from disk, write new blocks to storage
+4. **Collect Metrics**: Track latency, bandwidth, and cache hit rates
+
+## Quick Start
+
+```bash
+# Quick test (100 requests, no timestamp replay)
+python storage_benchmark.py --scenario=toolagent --max-requests=100
+
+# Test with large model preset (Llama-3.1-405B)
+python storage_benchmark.py --scenario=toolagent --model=llama-3.1-405b --max-requests=100
+
+# Test with Deepseek V3 (extra large model)
+python storage_benchmark.py --scenario=toolagent --model=deepseek-v3 --max-requests=100
+
+# Realistic replay (with timestamps, 10x speed)
+python storage_benchmark.py --scenario=toolagent --max-requests=1000 \
+    --replay-timestamps --time-scale=0.1
+
+# Test all scenarios with replay
+python storage_benchmark.py --scenario=all --time-scale=1.0
+```
+
+## Command-Line Options
+
+| Option | Description | Default |
+|--------|-------------|---------|
+| `--trace-dir` | Trace files directory | `../FAST25-release/traces` |
+| `--scenario` | Test scenario: `conversation`, `synthetic`, `toolagent`, `all` | `toolagent` |
+| `--storage-dir` | Storage directory | `/tmp/mooncake_bench` |
+| `--model` | Model preset (overrides `--bytes-per-token`) | `default` |
+| `--bytes-per-token` | Bytes per token (2048 for 7B FP16) | `2048` |
+| `--max-requests` | Maximum requests per scenario (unlimited if not specified) | `None` |
+| `--max-blocks` | Maximum number of blocks | `100000` |
+| `--replay-timestamps` | Enable timestamp replay | `False` |
+| `--time-scale` | Time scaling factor (1.0 = real-time, 0.1 = 10x faster) | `1.0` |
+
+## Model Presets
+
+The tool includes presets for popular LLM models with accurate KVCache sizes based on the [LMCache KVCache Calculator](https://lmcache.ai/kv_cache_calculator.html).
+
+| Model | Bytes/Token | Size | Notes |
+|-------|-------------|------|-------|
+| **Small Models (7B-13B)** |
+| `llama-3-8b` | 128 | 128 B/token | GQA optimized |
+| `mistral-7b` | 128 | 128 B/token | GQA optimized |
+| `qwen-14b` | 40 | 40 B/token | GQA optimized |
+| `gemma-7b` | 224 | 224 B/token | |
+| `llama-2-7b` | 512 | 512 B/token | |
+| `llama-2-13b` | 800 | 800 B/token | |
+| **Large Models (70B-405B)** |
+| `llama-2-70b` | 320 | 320 B/token | GQA optimized |
+| `llama-3-70b` | 320 | 320 B/token | GQA optimized |
+| `mixtral-8x7b` | 128 | 128 B/token | GQA optimized |
+| `mixtral-8x22b` | 224 | 224 B/token | GQA optimized |
+| `qwen-72b` | 320 | 320 B/token | GQA optimized |
+| `qwen-110b` | 320 | 320 B/token | GQA optimized |
+| `llama-3.1-405b` | 516018 | ~504 KB/token | Very large KVCache |
+| **Extra Large Models** |
+| `glm-4.6` | 156991 | ~153 KB/token | |
+| `deepseek-v3` | 1749384 | ~1.67 MB/token | Largest KVCache |
+| **Default** |
+| `default` | 2048 | 2 KB/token | Legacy 7B FP16 |
+
+**Usage**: `--model=llama-3.1-405b` (overrides `--bytes-per-token`)
+
+## Test Scenarios
+
+- **`conversation`**: Write-intensive workload (dialogue patterns)
+- **`synthetic`**: Read-intensive workload (cached patterns)
+- **`toolagent`**: Balanced read/write mix (tool use patterns)
+
+## Output Example
+
+```
+================================================================================
+Mooncake KVCache Storage Benchmark
+================================================================================
+Using model preset: llama-3.1-405b (516018 bytes/token, ~504.0 KB/token)
+
+[1/1] toolagent_trace.jsonl
+================================================================================
+
+[Performance Overview]
+  Total Requests:           100
+  Queries Per Second (QPS): 14.45
+  Cache Hit Rate:           24.27%
+  Write Ratio:              75.73%
+  Total Blocks:             1,949
+    Read Blocks:            473
+    Write Blocks:           1,476
+    Prefix Hits:            376
+
+[Latency Analysis]
+  Request Latency (End-to-End): Avg=69.18ms, P50=15.49ms, P95=239.99ms, P99=310.58ms
+  Single I/O Operation (Per Block):
+    Read:  Avg=14.572ms, P50=0.280ms, P95=0.280ms, P99=0.280ms
+    Write: Avg=5.120ms, P50=5.120ms, P95=5.120ms, P99=5.120ms
+
+[I/O & Bandwidth]
+  Total Read I/O:          473.0 MB  (473 ops)
+  Total Write I/O:       1476.0 MB  (1,476 ops)
+  Effective Bandwidth:    280.8 MB/s
+
+[Storage Details]
+  Blocks in Use:            1,476
+  Free Blocks:                  0
+  Tokens per Block:            512
+  Block Size:                   1.00 MB
+
+[Execution Time]
+  Total Execution Time:      8.42 s
+
+================================================================================
+```
+
+## Metrics
+
+| Metric | Description |
+|--------|-------------|
+| **QPS** | Queries per second (based on I/O time, excluding sleep) |
+| **Request Latency** | End-to-end latency for entire request (all I/O operations) |
+| **Single I/O Latency** | Latency for individual block read/write operations (512 tokens) |
+| **P50/P95/P99** | Latency percentiles (milliseconds) using linear interpolation |
+| **Hit Rate** | Cache hit ratio for blocks |
+| **Write Ratio** | Percentage of blocks that needed to be written |
+| **Bandwidth** | Effective throughput based on I/O time only |
+| **Prefix Hits** | Number of blocks served from prefix cache |
+
+**Note**: Request Latency measures the total time to process all blocks in a request, while Single I/O Latency measures the time for one block operation (512 tokens).
+
+## Trace Data Format
+
+```json
+{
+  "timestamp": 1234.567,
+  "hash_ids": [1, 2, 4, 7],
+  "input_length": 2048,
+  "output_length": 512
+}
+```
+
+Each `hash_id` corresponds to a 512-token block. The tool simulates prefix caching by checking if blocks are already in storage before writing.
+
+## Requirements
+
+- Python 3.10+

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)
@@ -629,6 +784,120 @@ result = store.put_batch(keys, values)
 
 ---
 
+#### upsert()
+
+Insert a new object if the key does not exist, or update the existing object in place when possible. They use the same replication configuration model as `put()`.
+
+Upsert binary data in the distributed storage.
+
+```python
+def upsert(self, key: str, value: bytes, config: ReplicateConfig = None) -> int
+```
+
+**Parameters:**
+- `key` (str): Unique object identifier
+- `value` (bytes): Binary data to insert or update
+- `config` (ReplicateConfig, optional): Replication configuration
+
+**Returns:**
+- `int`: Status code (0 = success, non-zero = error code)
+
+**Example:**
+```python
+config = ReplicateConfig()
+config.replica_num = 2
+
+rc = store.upsert("weights", b"new-bytes", config)
+if rc == 0:
+    print("Upsert succeeded")
+```
+
+#### upsert_from()
+
+Upsert object data directly from a pre-allocated buffer (zero-copy).
+
+```python
+def upsert_from(self, key: str, buffer_ptr: int, size: int, config: ReplicateConfig = None) -> int
+```
+
+**Parameters:**
+- `key` (str): Object identifier
+- `buffer_ptr` (int): Memory address of the source buffer
+- `size` (int): Number of bytes to insert or update
+- `config` (ReplicateConfig, optional): Replication configuration
+
+**Returns:**
+- `int`: Status code (0 = success, non-zero = error code)
+
+**Note:** This is the zero-copy counterpart of `upsert()`. As with
+`put_from()`, register the buffer before issuing the request.
+
+#### batch_upsert_from()
+
+Upsert multiple objects directly from pre-allocated buffers.
+
+```python
+def batch_upsert_from(self, keys: List[str], buffer_ptrs: List[int], sizes: List[int],
+                      config: ReplicateConfig = None) -> List[int]
+```
+
+**Parameters:**
+- `keys` (List[str]): List of object identifiers
+- `buffer_ptrs` (List[int]): List of source buffer addresses
+- `sizes` (List[int]): List of byte lengths for each buffer
+- `config` (ReplicateConfig, optional): Replication configuration shared by all objects
+
+**Returns:**
+- `List[int]`: List of status codes for each upsert
+
+#### upsert_parts()
+
+Upsert data from multiple buffer parts as a single object (insert or update).
+
+```python
+def upsert_parts(self, key: str, *parts, config: ReplicateConfig = None) -> int
+```
+
+**Parameters:**
+- `key` (str): Object identifier
+- `*parts`: Variable number of bytes-like objects to concatenate
+- `config` (ReplicateConfig, optional): Replication configuration
+
+**Returns:**
+- `int`: Status code (0 = success, non-zero = error code)
+
+**Example:**
+```python
+part1 = b"Hello, "
+part2 = b"World!"
+result = store.upsert_parts("greeting", part1, part2)
+```
+
+#### upsert_batch()
+
+Upsert multiple objects in a single batch operation.
+
+```python
+def upsert_batch(self, keys: List[str], values: List[bytes], config: ReplicateConfig = None) -> int
+```
+
+**Parameters:**
+- `keys` (List[str]): List of object identifiers
+- `values` (List[bytes]): List of binary data to insert or update
+- `config` (ReplicateConfig, optional): Replication configuration for all objects
+
+**Returns:**
+- `int`: Status code (0 = success, non-zero = error code)
+
+**Example:**
+```python
+keys = ["key1", "key2", "key3"]
+values = [b"value1", b"value2", b"value3"]
+result = store.upsert_batch(keys, values)
+```
+
+---
+
 #### get_batch()
 Retrieve multiple objects in a single batch operation.
 
@@ -721,6 +990,39 @@ print(f"Removed {count} objects")
 
 ---
 
+#### batch_remove()
+Remove multiple objects by their keys in a single batch operation.
+
+```python
+def batch_remove(self, keys: List[str], force: bool = False) -> List[int]
+```
+
+**Parameters:**
+- `keys` (List[str]): List of object identifiers to remove
+- `force` (bool): If True, skip lease and replication task checks (default: False)
+
+**Returns:**
+- `List[int]`: List of status codes for each key (0 = success, negative = error code)
+
+**Example:**
+```python
+# Remove multiple keys in one batch
+keys = ["key1", "key2", "key3", "key4", "key5"]
+results = store.batch_remove(keys)
+
+# Check results
+for key, result in zip(keys, results):
+    if result == 0:
+        print(f"✓ {key} removed successfully")
+    else:
+        print(f"✗ {key} failed with error code: {result}")
+
+# Force remove (bypass lease checks)
+results = store.batch_remove(keys, force=True)
+```
+
+---
+
 #### is_exist()
 Check if an object exists in the storage system.
 
@@ -1500,6 +1802,133 @@ def batch_pub_tensor(self, keys: List[str], tensors_list: List[torch.Tensor], co
 
 ---
 
+#### upsert_tensor()
+
+Insert a tensor if its key is missing, or update the existing tensor if the key already exists. The current tensor upsert helpers use the default `ReplicateConfig` and therefore do not take a `config` parameter.
+
+Upsert a PyTorch tensor into the store.
+
+```python
+def upsert_tensor(self, key: str, tensor: torch.Tensor) -> int
+```
+
+**Parameters:**
+- `key` (str): Object identifier
+- `tensor` (torch.Tensor): The PyTorch tensor to insert or update
+
+**Returns:**
+- `int`: Status code (0 = success, non-zero = error code)
+
+**Note:** This function requires `torch` to be installed and available in the environment.
+
+#### upsert_tensor_from()
+
+Upsert a tensor directly from a pre-allocated buffer. The buffer layout must be
+`[TensorMetadata][tensor data]`, matching the layout used by
+`get_tensor_into()`.
+
+```python
+def upsert_tensor_from(self, key: str, buffer_ptr: int, size: int) -> int
+```
+
+**Parameters:**
+- `key` (str): Object identifier
+- `buffer_ptr` (int): Buffer pointer containing serialized tensor metadata and payload
+- `size` (int): Actual serialized byte length of the tensor buffer
+
+**Returns:**
+- `int`: Status code (0 = success, non-zero = error code)
+
+**Note:** This function is not supported for dummy client.
+
+#### batch_upsert_tensor_from()
+
+Upsert multiple tensors directly from pre-allocated buffers. Each buffer must
+use layout `[TensorMetadata][tensor data]`.
+
+```python
+def batch_upsert_tensor_from(self, keys: List[str], buffer_ptrs: List[int], sizes: List[int]) -> List[int]
+```
+
+**Parameters:**
+- `keys` (List[str]): List of object identifiers
+- `buffer_ptrs` (List[int]): List of serialized tensor buffer pointers
+- `sizes` (List[int]): List of actual serialized byte lengths
+
+**Returns:**
+- `List[int]`: List of status codes for each tensor upsert
+
+#### batch_upsert_tensor()
+
+Upsert a batch of PyTorch tensors into the store (insert or update).
+
+```python
+def batch_upsert_tensor(self, keys: List[str], tensors_list: List[torch.Tensor]) -> List[int]
+```
+
+**Parameters:**
+- `keys` (List[str]): List of object identifiers
+- `tensors_list` (List[torch.Tensor]): List of tensors to insert or update
+
+**Returns:**
+- `List[int]`: List of status codes for each tensor operation.
+
+**Note:** This function requires `torch` to be installed and available in the environment. Not supported for dummy client.
+
+#### upsert_pub_tensor()
+
+Upsert a PyTorch tensor with configurable replication settings (insert or update).
+
+```python
+def upsert_pub_tensor(self, key: str, tensor: torch.Tensor, config: ReplicateConfig = None) -> int
+```
+
+**Parameters:**
+- `key` (str): Unique object identifier
+- `tensor` (torch.Tensor): PyTorch tensor to insert or update
+- `config` (ReplicateConfig, optional): Replication configuration
+
+**Returns:**
+- `int`: Status code (0 = success, non-zero = error code)
+
+**Note:** This function requires `torch` to be installed and available in the environment. Not supported for dummy client.
+
+**Example:**
+```python
+import torch
+from mooncake.store import ReplicateConfig
+
+tensor = torch.randn(100, 100)
+
+config = ReplicateConfig()
+config.replica_num = 2
+config.with_soft_pin = True
+
+result = store.upsert_pub_tensor("my_tensor", tensor, config)
+if result == 0:
+    print("Tensor upserted successfully")
+```
+
+#### batch_upsert_pub_tensor()
+
+Batch upsert PyTorch tensors with configurable replication settings (insert or update).
+
+```python
+def batch_upsert_pub_tensor(self, keys: List[str], tensors_list: List[torch.Tensor], config: ReplicateConfig = None) -> List[int]
+```
+
+**Parameters:**
+- `keys` (List[str]): List of object identifiers
+- `tensors_list` (List[torch.Tensor]): List of tensors to insert or update
+- `config` (ReplicateConfig, optional): Replication configuration
+
+**Returns:**
+- `List[int]`: List of status codes for each tensor operation.
+
+**Note:** This function requires `torch` to be installed and available in the environment. Not supported for dummy client.
+
+---
+
 ### PyTorch Tensor Operations (Zero Copy)
 
 These methods provide direct support for storing and retrieving PyTorch tensors. They automatically handle serialization and metadata, and include built-in support for **Tensor Parallelism (TP)** by automatically splitting and reconstructing tensor shards.
@@ -1535,8 +1964,8 @@ def batch_get_tensor_into(self, base_keys: List[str], buffer_ptrs: List[int], si
 **Parameters:**
 
   - `base_keys` (List[str]): List of base identifiers.
-  - `buffer_ptrs` (List[int]): List of the buffers pointer pre-allocated for tensor, and the buffers should be registered.
-  - `sizes` (List[int]): List of the size of buffers.
+  - `buffer_ptrs` (List[int]): List of buffer pointers pre-allocated for tensor; buffers should be registered.
+  - `sizes` (List[int]): List of buffer sizes.
 
 **Returns:**
 
@@ -1574,8 +2003,8 @@ def batch_get_tensor_with_tp_into(self, base_keys: List[str], buffer_ptrs: List[
 **Parameters:**
 
   - `base_keys` (List[str]): List of base identifiers.
-  - `buffer_ptrs` (List[int]): List of the buffers pointer pre-allocated for tensor, and the buffers should be registered.
-  - `sizes` (List[int]): List of the size of buffers.
+  - `buffer_ptrs` (List[int]): List of buffer pointers pre-allocated for tensor; buffers should be registered.
+  - `sizes` (List[int]): List of buffer sizes.
   - `tp_rank` (int): The tensor parallel rank to retrieve (default: 0).
   - `tp_size` (int): Total tensor parallel size (default: 1).
 
@@ -1583,6 +2012,84 @@ def batch_get_tensor_with_tp_into(self, base_keys: List[str], buffer_ptrs: List[
 
   - `List[torch.Tensor]`: List of retrieved tensors (or shards). Contains `None` for missing keys.
 
+#### put_tensor_from()
+
+Put a PyTorch tensor into the store directly from a pre-allocated buffer (zero-copy). The buffer must contain data in the same layout as produced by `get_tensor_into`: **\[TensorMetadata\]\[tensor data\]**. The buffer is only read during this call; no Python object references it.
+
+```python
+def put_tensor_from(self, key: str, buffer_ptr: int, size: int) -> int
+```
+
+**Parameters:**
+
+  - `key` (str): Object identifier for the tensor.
+  - `buffer_ptr` (int): The buffer pointer; the buffer should be registered. Layout must be \[TensorMetadata\]\[tensor data\].
+  - `size` (int): **Actual serialized byte length** of the data in the buffer (metadata + tensor bytes), not the buffer capacity.
+
+**Returns:**
+
+  - `int`: Status code (0 = success, non-zero = error code).
+
+#### batch_put_tensor_from()
+
+Put a batch of PyTorch tensors into the store directly from pre-allocated buffers (zero-copy). Each buffer must contain data in the layout **\[TensorMetadata\]\[tensor data\]**, same as `get_tensor_into`.
+
+```python
+def batch_put_tensor_from(self, keys: List[str], buffer_ptrs: List[int], sizes: List[int]) -> List[int]
+```
+
+**Parameters:**
+
+  - `keys` (List[str]): List of object identifiers.
+  - `buffer_ptrs` (List[int]): List of buffer pointers; buffers should be registered.
+  - `sizes` (List[int]): List of **actual serialized byte lengths** for each buffer (metadata + tensor bytes), not buffer capacities.
+
+**Returns:**
+
+  - `List[int]`: List of status codes for each tensor operation (0 = success, non-zero = error code).
+
+#### put_tensor_with_tp_from()
+
+Put a **full tensor** into the store directly from a pre-allocated buffer (zero-copy), for use with Tensor Parallelism. This is the zero-copy counterpart of `put_tensor_with_tp()`: the buffer must contain the complete tensor in layout **\[TensorMetadata\]\[tensor data\]**, and Mooncake will split it internally and store all shards under `key_tp_<rank>`.
+
+```python
+def put_tensor_with_tp_from(self, key: str, buffer_ptr: int, size: int, tp_rank: int = 0, tp_size: int = 1, split_dim: int = 0) -> int
+```
+
+**Parameters:**
+
+  - `key` (str): Base identifier for the tensor.
+  - `buffer_ptr` (int): The buffer pointer; the buffer should be registered.
+  - `size` (int): **Actual serialized byte length** of the full tensor in the buffer.
+  - `tp_rank` (int): Kept for signature compatibility with `put_tensor_with_tp()` (default: 0). It does **not** mean "only write one shard".
+  - `tp_size` (int): Total tensor parallel size (default: 1). If 1, equivalent to `put_tensor_from(key, buffer_ptr, size)`.
+  - `split_dim` (int): Dimension along which the full tensor is split before storing shards.
+
+**Returns:**
+
+  - `int`: Status code (0 = success, non-zero = error code).
+
+#### batch_put_tensor_with_tp_from()
+
+Put a batch of **full tensors** into the store directly from pre-allocated buffers (zero-copy). This is the zero-copy counterpart of `batch_put_tensor_with_tp()`: each buffer contains one full tensor in layout **\[TensorMetadata\]\[tensor data\]**, and Mooncake splits each tensor internally and stores all TP shards.
+
+```python
+def batch_put_tensor_with_tp_from(self, base_keys: List[str], buffer_ptrs: List[int], sizes: List[int], tp_rank: int = 0, tp_size: int = 1, split_dim: int = 0) -> List[int]
+```
+
+**Parameters:**
+
+  - `base_keys` (List[str]): List of base identifiers.
+  - `buffer_ptrs` (List[int]): List of buffer pointers; buffers should be registered.
+  - `sizes` (List[int]): List of **actual serialized byte lengths** for each full-tensor buffer.
+  - `tp_rank` (int): Kept for signature compatibility with `batch_put_tensor_with_tp()` (default: 0). It does **not** select a single shard to write.
+  - `tp_size` (int): Total tensor parallel size (default: 1). If 1, equivalent to `batch_put_tensor_from(base_keys, buffer_ptrs, sizes)`.
+  - `split_dim` (int): Dimension along which each full tensor is split before storing shards.
+
+**Returns:**
+
+  - `List[int]`: List of status codes for each tensor operation (0 = success, non-zero = error code).
+
 ---
 
 ### Batch Zero-Copy Operations

docs/source/troubleshooting/troubleshooting.md

@@ -63,7 +63,37 @@ Errors in this part usually indicate that the error occurred within the `mooncak
    **Solution:**
    Ensure that the total memory registration does not exceed the device's upper limit. You may need to reduce the amount of memory being registered or split large memory regions into smaller chunks that fit within the device's `max_mr_size` limit.
 
-5. If you encounter errors indicating inability to allocate memory space when requesting large memory regions, this may be due to ulimit restrictions. When the total memory requirement (number of registered RDMA devices × requested space) exceeds the ulimit, the system will display errors about failing to allocate space.
+5. If you encounter `Failed to register memory 0x...: Resource temporarily unavailable [11]` and kernel logs show `CREATE_MKEY failed, status no resources(0xf)`, this indicates that the RDMA NIC has exhausted its internal Memory Key (MKEY) resources, even though `ulimit -l` and `vm.max_map_count` may appear sufficient.
+
+   This typically happens when:
+   - Applications that use RDMA (e.g., SGLang with HiCache + Mooncake) have crashed or been killed multiple times without cleanly releasing RDMA resources.
+   - The leaked MKEY entries accumulate in the NIC firmware and are not reclaimed by the kernel, eventually hitting the hardware limit.
+   - Large memory regions (e.g., NSA indexer buffers at ~4.68 GB each across multiple TP ranks) amplify the problem since each registration consumes more internal NIC resources.
+
+   **Diagnostic Commands:**
+   ```bash
+   # Check current RDMA resource usage per device
+   rdma resource show
+
+   # Check kernel logs for CREATE_MKEY failures
+   dmesg | grep -i "CREATE_MKEY\|no resources\|mlx5_cmd_out_err"
+   # Example output:
+   # mlx5_core 0000:65:01.0: mlx5_cmd_out_err:829:(pid 3958462): CREATE_MKEY(0x200) op_mod(0x0) failed, status no resources(0xf), syndrome (0x2aac7c), err(-11)
+
+   # Ensure vm.max_map_count is large enough (default 65530 may be too small)
+   sysctl vm.max_map_count
+   ```
+
+   **Solutions:**
+   - **Reboot the node** to fully reset NIC firmware state and reclaim all leaked MKEY resources. This is the most reliable fix.
+   - Increase `vm.max_map_count` if it is at the default value: `sysctl -w vm.max_map_count=16777216`
+   - Ensure applications shut down cleanly (avoid `kill -9` when possible) so RDMA resources are properly deregistered.
+   - If rebooting is not feasible, try unloading and reloading the mlx5 kernel modules (may disrupt other services):
+     ```bash
+     modprobe -r mlx5_ib mlx5_core && modprobe mlx5_core mlx5_ib
+     ```
+
+6. If you encounter errors indicating inability to allocate memory space when requesting large memory regions, this may be due to ulimit restrictions. When the total memory requirement (number of registered RDMA devices × requested space) exceeds the ulimit, the system will display errors about failing to allocate space.
 
    **Diagnostic Commands:**
    - Use `ulimit -a` to check current limits, particularly the `max locked memory` value
@@ -80,7 +110,7 @@ Errors in this part usually indicate that the error occurred within the `mooncak
      *    hard    memlock    unlimited
      ```
 
-6. If the error `Failed to create QP: Cannot allocate memory` is displayed, it is typically caused by too many QP have been created, reaching the driver limit. You can use `rdma resource` to trace how many QP is created. One possible way to resolve this issue:
+7. If the error `Failed to create QP: Cannot allocate memory` is displayed, it is typically caused by too many QP have been created, reaching the driver limit. You can use `rdma resource` to trace how many QP is created. One possible way to resolve this issue:
    - Update Mooncake to version v0.3.5 or later
    - Set the environment variable `MC_ENABLE_DEST_DEVICE_AFFINITY=1` before starting the application
 

docs/source/zh_archive/ascend_transport.md

@@ -1,5 +1,8 @@
 # Ascend Transport
 Ascend Transport源代码路径为Mooncake/mooncake-transfer-engine/src/transport/ascend_transport，该路径下还包含自动化编译脚本、README文件。
+
+**Ascend Transport 已不再维护，昇腾平台推荐使用 [Ascend Direct Transport](./ascend_direct_transport.md). **
+
 ## 概述
 Ascend Transport是一个单边语义的高性能零拷贝NPU数据传输库，直接兼容Mooncake Transfer Engine。要编译使用Ascend Transport库，请在mooncake-common\common.cmake文件中将USE_ASCEND开关置于"ON"。
 

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
+```

docs/source/zh_archive/mooncake-store.md

@@ -779,6 +779,8 @@ Max threads: 4
 Master service listening on 0.0.0.0:50051
 ```
 
+如果 Master 运行在容器中，而容器 IP 可能动态变化，建议使用 `--rpc-interface=<网卡名>`（例如 `--rpc-interface=eth0`）而不是写死 `--rpc-address`。Master 会在启动时解析该网卡当前的 IPv4 地址，并将其作为最终的 `rpc_address` 使用。
+
 **高可用模式**:
 
 高可用模式依赖于 etcd 服务进行协调。如果 Transfer Engine 也使用 etcd 作为其元数据服务，那么 Mooncake Store 使用的 etcd 集群可以与 Transfer Engine 使用的集群共用，也可以是独立的。
@@ -788,6 +790,7 @@ Master service listening on 0.0.0.0:50051
 --enable-ha：启用高可用模式
 --etcd-endpoints：指定 etcd 服务的多个入口，使用分号 ';' 分隔
 --rpc-address：该实例的 RPC 地址。注意，这里填写的地址应当是客户端可访问的地址。
+--rpc-interface：按网卡名解析当前实例的 IPv4 地址。设置后会覆盖 --rpc-address，适合容器 IP 会变化的场景。
 ```
 
 例如:
@@ -798,6 +801,14 @@ Master service listening on 0.0.0.0:50051
     --rpc-address=10.0.0.1
 ```
 
+容器部署示例：
+```
+./build/mooncake-store/src/mooncake_master \
+    --enable-ha=true \
+    --etcd-endpoints="0.0.0.0:2379;0.0.0.0:2479;0.0.0.0:2579" \
+    --rpc-interface=eth0
+```
+
 ### 启动验证程序
 Mooncake Store 提供了多种验证程序，包括基于 C++ 和 Python 等接口形态。下面以 `stress_cluster_benchmark` 为例介绍一下如何运行。
 

extern/yalantinglibs

@@ -0,0 +1 @@
+Subproject commit 73dea196d23ad8fcd4914c6ef1238f390b9a1c48

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

@@ -40,6 +40,9 @@ add_definitions(-DCONFIG_ERDMA)
 
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
+# Memory-aware build parallelism (compile vs. link job pools)
+include(${CMAKE_CURRENT_LIST_DIR}/limit_jobs.cmake)
+
 option(ENABLE_SCCACHE "Whether to open sccache" OFF)
 if (ENABLE_SCCACHE)
   find_program(SCCACHE sccache REQUIRED)
@@ -57,7 +60,9 @@ 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)
 option(USE_NVMEOF "option for using NVMe over Fabric" OFF)
 option(USE_TCP "option for using TCP transport" ON)
@@ -69,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
@@ -102,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)
@@ -126,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)
@@ -143,6 +159,60 @@ 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)
+  add_compile_definitions(USE_MACA)
+  message(STATUS "MACA support is enabled")
+  include_directories(${MACA_INCLUDE_DIR})
+  if (EXISTS "${MACA_LIB_DIR}")
+    link_directories(${MACA_LIB_DIR})
+  endif()
+endif()
+
 if (USE_MUSA)
   add_compile_definitions(USE_MUSA)
   message(STATUS "MUSA support is enabled")

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/etcd_wrapper.go

@@ -34,11 +34,13 @@ import "C"
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"strings"
 	"sync"
 	"time"
 	"unsafe"
 
+	rpctypes "go.etcd.io/etcd/api/v3/v3rpc/rpctypes"
 	clientv3 "go.etcd.io/etcd/client/v3"
 )
 
@@ -64,11 +66,11 @@ var (
 	storeKeepAliveCtx   = make(map[int64]context.CancelFunc)
 	storeKeepAliveMutex sync.Mutex
 	// watch contexts for store
-	storeWatchCtx = make(map[string]context.CancelFunc)
-	storeWatchMutex    sync.Mutex
+	storeWatchCtx   = make(map[string]context.CancelFunc)
+	storeWatchMutex sync.Mutex
 	// etcd client for HA snapshot
-	snapshotClient  *clientv3.Client
-	snapshotMutex   sync.Mutex
+	snapshotClient *clientv3.Client
+	snapshotMutex  sync.Mutex
 	// watch contexts for prefix watch
 	storePrefixWatchCtx   = make(map[string]prefixWatchInfo)
 	storePrefixWatchMutex sync.Mutex
@@ -76,7 +78,7 @@ var (
 
 const (
 	// Snapshot client config (for GB-level snapshot files)
-	snapshotMaxMsgSize = 2000 * 1000 * 1000  // 2GB
+	snapshotMaxMsgSize = 2000 * 1000 * 1000 // 2GB
 	snapshotTimeout    = 60 * time.Second   // 1 minute for large files
 )
 
@@ -325,6 +327,25 @@ func EtcdStoreGrantLeaseWrapper(ttl int64, leaseId *int64, errMsg **C.char) int
 	return 0
 }
 
+//export EtcdStoreRevokeLeaseWrapper
+func EtcdStoreRevokeLeaseWrapper(leaseId int64, errMsg **C.char) int {
+	if storeClient == nil {
+		*errMsg = C.CString("etcd client not initialized")
+		return -1
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	_, err := storeClient.Revoke(ctx, clientv3.LeaseID(leaseId))
+	if err != nil {
+		if errors.Is(err, rpctypes.ErrLeaseNotFound) {
+			return 0
+		}
+		*errMsg = C.CString(err.Error())
+		return -1
+	}
+	return 0
+}
+
 //export EtcdStoreCreateWithLeaseWrapper
 func EtcdStoreCreateWithLeaseWrapper(key *C.char, keySize C.int, value *C.char, valueSize C.int,
 	leaseId int64, revisionId *int64, errMsg **C.char) int {
@@ -459,6 +480,14 @@ func cancelAndDeleteKeepAlive(leaseId int64) int {
 	return -1
 }
 
+func hasKeepAliveContext(leaseId int64) bool {
+	storeKeepAliveMutex.Lock()
+	defer storeKeepAliveMutex.Unlock()
+
+	_, exists := storeKeepAliveCtx[leaseId]
+	return exists
+}
+
 //export EtcdStoreKeepAliveWrapper
 func EtcdStoreKeepAliveWrapper(leaseId int64, errMsg **C.char) int {
 	if storeClient == nil {
@@ -518,6 +547,21 @@ func EtcdStoreCancelKeepAliveWrapper(leaseId int64, errMsg **C.char) int {
 	return 0
 }
 
+//export EtcdStoreWaitKeepAliveReadyWrapper
+func EtcdStoreWaitKeepAliveReadyWrapper(leaseId int64, timeoutMs int, errMsg **C.char) int {
+	deadline := time.Now().Add(time.Duration(timeoutMs) * time.Millisecond)
+	for {
+		if hasKeepAliveContext(leaseId) {
+			return 0
+		}
+		if timeoutMs <= 0 || !time.Now().Before(deadline) {
+			*errMsg = C.CString("keep alive context did not become ready before timeout")
+			return -1
+		}
+		time.Sleep(time.Millisecond)
+	}
+}
+
 //export EtcdStorePutWrapper
 func EtcdStorePutWrapper(key *C.char, keySize C.int, value *C.char, valueSize C.int, errMsg **C.char) int {
 	if storeClient == nil {

mooncake-common/etcd/go.mod

@@ -1,27 +1,28 @@
 module github.com/kvcache-ai/Mooncake/mooncake-common/etcd
 
-go 1.23.0
+go 1.25.0
 
-toolchain go1.23.7
+toolchain go1.25.9
 
-require go.etcd.io/etcd/client/v3 v3.5.21
+require (
+	go.etcd.io/etcd/api/v3 v3.5.21
+	go.etcd.io/etcd/client/v3 v3.5.21
+)
 
 require (
 	github.com/coreos/go-semver v0.3.0 // indirect
 	github.com/coreos/go-systemd/v22 v22.3.2 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
 	github.com/golang/protobuf v1.5.4 // indirect
-	go.etcd.io/etcd/api/v3 v3.5.21 // indirect
 	go.etcd.io/etcd/client/pkg/v3 v3.5.21 // indirect
 	go.uber.org/atomic v1.7.0 // indirect
 	go.uber.org/multierr v1.6.0 // indirect
 	go.uber.org/zap v1.17.0 // indirect
-	golang.org/x/net v0.38.0 // indirect
-	golang.org/x/sys v0.31.0 // indirect
-	golang.org/x/text v0.23.0 // indirect
-	google.golang.org/genproto v0.0.0-20230822172742-b8732ec3820d // indirect
-	google.golang.org/genproto/googleapis/api v0.0.0-20230822172742-b8732ec3820d // indirect
-	google.golang.org/genproto/googleapis/rpc v0.0.0-20230822172742-b8732ec3820d // indirect
-	google.golang.org/grpc v1.59.0 // indirect
-	google.golang.org/protobuf v1.33.0 // indirect
+	golang.org/x/net v0.48.0 // indirect
+	golang.org/x/sys v0.39.0 // indirect
+	golang.org/x/text v0.32.0 // indirect
+	google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217 // indirect
+	google.golang.org/grpc v1.79.3 // indirect
+	google.golang.org/protobuf v1.36.10 // indirect
 )

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/include/default_config.h

@@ -70,6 +70,18 @@ class DefaultConfig {
     void GetUInt64(const std::string& key, uint64_t* val,
                    uint64_t default_value = 0) const;
 
+    /**
+     * @brief GetDurationMs retrieves a duration value from the configuration
+     * and converts it to milliseconds.
+     * @param key The key to look up in the configuration
+     * @param val Pointer to store the retrieved value in milliseconds
+     * @param default_value Default value to return if the key is not found
+     * @note Duration strings may use ms, s, m, or h as suffixes. Bare numbers
+     * are interpreted as milliseconds.
+     */
+    void GetDurationMs(const std::string& key, uint64_t* val,
+                       uint64_t default_value = 0) const;
+
     /**
      * @brief GetDouble retrieves a double value from the configuration
      * @param key The key to look up in the configuration

mooncake-common/include/duration_utils.h

@@ -0,0 +1,95 @@
+#pragma once
+
+#include <cctype>
+#include <cstdint>
+#include <limits>
+#include <string>
+#include <string_view>
+
+namespace mooncake {
+
+inline std::string_view TrimAsciiWhitespace(std::string_view value) {
+    while (!value.empty() &&
+           std::isspace(static_cast<unsigned char>(value.front()))) {
+        value.remove_prefix(1);
+    }
+    while (!value.empty() &&
+           std::isspace(static_cast<unsigned char>(value.back()))) {
+        value.remove_suffix(1);
+    }
+    return value;
+}
+
+inline bool ParseDurationMs(std::string_view value, uint64_t* result,
+                            std::string* error = nullptr) {
+    auto set_error = [&](std::string message) {
+        if (error != nullptr) {
+            *error = std::move(message);
+        }
+        return false;
+    };
+
+    if (result == nullptr) {
+        return set_error("duration output pointer is null");
+    }
+
+    std::string_view trimmed = TrimAsciiWhitespace(value);
+    if (trimmed.empty()) {
+        return set_error(
+            "duration is empty; expected a non-negative integer optionally "
+            "followed by ms, s, m, or h");
+    }
+
+    size_t number_end = 0;
+    while (number_end < trimmed.size() &&
+           std::isdigit(static_cast<unsigned char>(trimmed[number_end]))) {
+        ++number_end;
+    }
+
+    if (number_end == 0) {
+        return set_error(
+            "duration must start with a non-negative integer and may use ms, "
+            "s, m, or h as the unit suffix");
+    }
+
+    uint64_t numeric_value = 0;
+    for (size_t i = 0; i < number_end; ++i) {
+        const uint64_t digit = static_cast<uint64_t>(trimmed[i] - '0');
+        if (numeric_value >
+            (std::numeric_limits<uint64_t>::max() - digit) / 10) {
+            return set_error("duration value is too large");
+        }
+        numeric_value = numeric_value * 10 + digit;
+    }
+
+    std::string_view suffix = TrimAsciiWhitespace(trimmed.substr(number_end));
+    std::string normalized_suffix;
+    normalized_suffix.reserve(suffix.size());
+    for (char ch : suffix) {
+        normalized_suffix.push_back(
+            static_cast<char>(std::tolower(static_cast<unsigned char>(ch))));
+    }
+
+    uint64_t multiplier = 1;
+    if (normalized_suffix.empty() || normalized_suffix == "ms") {
+        multiplier = 1;
+    } else if (normalized_suffix == "s") {
+        multiplier = 1000;
+    } else if (normalized_suffix == "m") {
+        multiplier = 60 * 1000;
+    } else if (normalized_suffix == "h") {
+        multiplier = 60 * 60 * 1000;
+    } else {
+        return set_error("unsupported duration unit '" + normalized_suffix +
+                         "'; supported units are ms, s, m, and h");
+    }
+
+    if (numeric_value > std::numeric_limits<uint64_t>::max() / multiplier) {
+        return set_error("duration value is too large after unit conversion");
+    }
+
+    *result = numeric_value * multiplier;
+    return true;
+}
+
+}  // namespace mooncake

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/limit_jobs.cmake

@@ -0,0 +1,86 @@
+# limit_jobs.cmake — Memory-aware build parallelism
+#
+# Auto-detects available memory and CPU count, calculates safe parallel job
+# limits for compilation and linking separately. With Ninja, creates job pools
+# so compilation uses many cores while memory-heavy linking is restricted.
+#
+# User overrides (cmake -D...):
+#   PARALLEL_COMPILE_JOBS  — override compile parallelism
+#   PARALLEL_LINK_JOBS     — override link parallelism
+#   MAX_COMPILER_MEMORY_MB — per-compile-job memory estimate (default: 1500)
+#   MAX_LINKER_MEMORY_MB   — per-link-job memory estimate (default: 4000)
+
+set(MAX_COMPILER_MEMORY_MB "1500" CACHE STRING
+    "Estimated peak memory per compile job in MB")
+set(MAX_LINKER_MEMORY_MB "4000" CACHE STRING
+    "Estimated peak memory per link job in MB")
+
+# Guard against invalid user input (division by zero)
+if(MAX_COMPILER_MEMORY_MB LESS_EQUAL 0)
+    message(WARNING "[limit_jobs] MAX_COMPILER_MEMORY_MB=${MAX_COMPILER_MEMORY_MB} "
+        "invalid, falling back to 1500")
+    set(MAX_COMPILER_MEMORY_MB 1500 CACHE STRING
+        "Estimated peak memory per compile job in MB" FORCE)
+endif()
+if(MAX_LINKER_MEMORY_MB LESS_EQUAL 0)
+    message(WARNING "[limit_jobs] MAX_LINKER_MEMORY_MB=${MAX_LINKER_MEMORY_MB} "
+        "invalid, falling back to 4000")
+    set(MAX_LINKER_MEMORY_MB 4000 CACHE STRING
+        "Estimated peak memory per link job in MB" FORCE)
+endif()
+
+# Detect system resources
+cmake_host_system_information(RESULT _available_mem_mb
+    QUERY AVAILABLE_PHYSICAL_MEMORY)
+cmake_host_system_information(RESULT _nproc
+    QUERY NUMBER_OF_LOGICAL_CORES)
+
+message(STATUS "[limit_jobs] Available memory: ${_available_mem_mb} MB, "
+    "CPU cores: ${_nproc}")
+
+# Calculate safe parallel jobs from memory
+math(EXPR _compile_jobs "${_available_mem_mb} / ${MAX_COMPILER_MEMORY_MB}")
+math(EXPR _link_jobs "${_available_mem_mb} / ${MAX_LINKER_MEMORY_MB}")
+
+# Clamp: [1, nproc]
+if(_compile_jobs LESS 1)
+    set(_compile_jobs 1)
+endif()
+if(_compile_jobs GREATER _nproc)
+    set(_compile_jobs ${_nproc})
+endif()
+if(_link_jobs LESS 1)
+    set(_link_jobs 1)
+endif()
+if(_link_jobs GREATER _nproc)
+    set(_link_jobs ${_nproc})
+endif()
+
+# Use auto-detected values unless user explicitly overrides with -D
+if(NOT DEFINED PARALLEL_COMPILE_JOBS)
+    set(PARALLEL_COMPILE_JOBS "${_compile_jobs}")
+endif()
+if(NOT DEFINED PARALLEL_LINK_JOBS)
+    set(PARALLEL_LINK_JOBS "${_link_jobs}")
+endif()
+
+message(STATUS "[limit_jobs] Compile jobs: ${PARALLEL_COMPILE_JOBS} "
+    "(${MAX_COMPILER_MEMORY_MB} MB/job), "
+    "Link jobs: ${PARALLEL_LINK_JOBS} (${MAX_LINKER_MEMORY_MB} MB/job)")
+
+# Apply to build system
+if(CMAKE_GENERATOR MATCHES "Ninja")
+    set_property(GLOBAL APPEND PROPERTY JOB_POOLS
+        compile_pool=${PARALLEL_COMPILE_JOBS}
+        link_pool=${PARALLEL_LINK_JOBS}
+    )
+    set(CMAKE_JOB_POOL_COMPILE "compile_pool" CACHE STRING "" FORCE)
+    set(CMAKE_JOB_POOL_LINK "link_pool" CACHE STRING "" FORCE)
+    message(STATUS "[limit_jobs] Ninja job pools: "
+        "compile=${PARALLEL_COMPILE_JOBS}, link=${PARALLEL_LINK_JOBS}")
+else()
+    message(STATUS "[limit_jobs] Hint: use -G Ninja for automatic "
+        "compile/link parallelism separation")
+    message(STATUS "[limit_jobs] With Make, recommend: "
+        "cmake --build . -j${PARALLEL_LINK_JOBS}")
+endif()

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-common/src/default_config.cpp

@@ -1,5 +1,7 @@
 #include "default_config.h"
 
+#include "duration_utils.h"
+
 #if __has_include(<jsoncpp/json/reader.h>)
 #include <jsoncpp/json/reader.h>
 #include <jsoncpp/json/value.h>  // Ubuntu
@@ -162,6 +164,71 @@ void DefaultConfig::GetUInt64(const std::string& key, uint64_t* val,
     }
 }
 
+void DefaultConfig::GetDurationMs(const std::string& key, uint64_t* val,
+                                  uint64_t default_value) const {
+    Node node;
+    if (!getValue(key, &node)) {
+        *val = default_value;
+        return;
+    }
+
+    if (type_ == ConfigType::YAML) {
+        std::string raw_value = node.yaml_node_.as<std::string>();
+        std::string error;
+        if (!ParseDurationMs(raw_value, val, &error)) {
+            throw std::runtime_error("Invalid duration for key '" + key +
+                                     "': " + error);
+        }
+        return;
+    }
+
+    if (type_ == ConfigType::JSON) {
+        if (node.json_value_.isString()) {
+            std::string error;
+            if (!ParseDurationMs(node.json_value_.asString(), val, &error)) {
+                throw std::runtime_error("Invalid duration for key '" + key +
+                                         "': " + error);
+            }
+            return;
+        }
+
+        if (node.json_value_.isUInt64()) {
+            *val = node.json_value_.asUInt64();
+            return;
+        }
+
+        if (node.json_value_.isUInt()) {
+            *val = static_cast<uint64_t>(node.json_value_.asUInt());
+            return;
+        }
+
+        if (node.json_value_.isInt64()) {
+            const int64_t numeric_value = node.json_value_.asInt64();
+            if (numeric_value < 0) {
+                throw std::runtime_error("Invalid duration for key '" + key +
+                                         "': value must be non-negative");
+            }
+            *val = static_cast<uint64_t>(numeric_value);
+            return;
+        }
+
+        if (node.json_value_.isInt()) {
+            const int numeric_value = node.json_value_.asInt();
+            if (numeric_value < 0) {
+                throw std::runtime_error("Invalid duration for key '" + key +
+                                         "': value must be non-negative");
+            }
+            *val = static_cast<uint64_t>(numeric_value);
+            return;
+        }
+
+        throw std::runtime_error("Invalid duration for key '" + key +
+                                 "': JSON value must be an integer or string");
+    }
+
+    *val = default_value;
+}
+
 void DefaultConfig::GetDouble(const std::string& key, double* val,
                               double default_value) const {
     Node node;

mooncake-common/tests/default_config_test.cpp

@@ -1,4 +1,5 @@
 #include "default_config.h"
+#include "duration_utils.h"
 
 #include <gtest/gtest.h>
 
@@ -77,6 +78,93 @@ TEST_F(DefaultConfigTest, LoadYamlSuccess) {
     ASSERT_EQ(config_data_.default_value, 10000);
 };
 
+TEST(DurationUtilsTest, ParseDurationMsSupportsLegacyMillisecondsAndUnits) {
+    uint64_t value = 0;
+
+    ASSERT_TRUE(ParseDurationMs("5000", &value));
+    ASSERT_EQ(value, 5000);
+
+    ASSERT_TRUE(ParseDurationMs("5000ms", &value));
+    ASSERT_EQ(value, 5000);
+
+    ASSERT_TRUE(ParseDurationMs("5s", &value));
+    ASSERT_EQ(value, 5000);
+
+    ASSERT_TRUE(ParseDurationMs("30m", &value));
+    ASSERT_EQ(value, 30 * 60 * 1000);
+
+    ASSERT_TRUE(ParseDurationMs("1H", &value));
+    ASSERT_EQ(value, 60 * 60 * 1000);
+
+    ASSERT_TRUE(ParseDurationMs(" 7 m ", &value));
+    ASSERT_EQ(value, 7 * 60 * 1000);
+}
+
+TEST(DurationUtilsTest, ParseDurationMsRejectsInvalidInput) {
+    uint64_t value = 0;
+
+    ASSERT_FALSE(ParseDurationMs("", &value));
+    ASSERT_FALSE(ParseDurationMs("abc", &value));
+    ASSERT_FALSE(ParseDurationMs("-1", &value));
+    ASSERT_FALSE(ParseDurationMs("1d", &value));
+    ASSERT_FALSE(ParseDurationMs("18446744073709551616", &value));
+    ASSERT_FALSE(ParseDurationMs("18446744073709552h", &value));
+}
+
+TEST_F(DefaultConfigTest, GetDurationMsFromJsonSupportsNumbersAndStrings) {
+    DefaultConfig config;
+    config.SetPath(path_ + "/../../mooncake-common/tests/test.json");
+    config.Load();
+
+    uint64_t legacy_ms = 0;
+    uint64_t seconds = 0;
+    uint64_t minutes = 0;
+    uint64_t hours = 0;
+    uint64_t whitespace = 0;
+    uint64_t missing_default = 0;
+
+    config.GetDurationMs("legacyDurationMs", &legacy_ms, 0);
+    config.GetDurationMs("durationSeconds", &seconds, 0);
+    config.GetDurationMs("durationMinutes", &minutes, 0);
+    config.GetDurationMs("durationHours", &hours, 0);
+    config.GetDurationMs("durationWhitespace", &whitespace, 0);
+    config.GetDurationMs("missingDuration", &missing_default, 1234);
+
+    ASSERT_EQ(legacy_ms, 5000);
+    ASSERT_EQ(seconds, 5000);
+    ASSERT_EQ(minutes, 30 * 60 * 1000);
+    ASSERT_EQ(hours, 60 * 60 * 1000);
+    ASSERT_EQ(whitespace, 7 * 60 * 1000);
+    ASSERT_EQ(missing_default, 1234);
+}
+
+TEST_F(DefaultConfigTest, GetDurationMsFromYamlSupportsNumbersAndStrings) {
+    DefaultConfig config;
+    config.SetPath(path_ + "/../../mooncake-common/tests/test.yaml");
+    config.Load();
+
+    uint64_t legacy_ms = 0;
+    uint64_t seconds = 0;
+    uint64_t minutes = 0;
+    uint64_t hours = 0;
+    uint64_t whitespace = 0;
+    uint64_t missing_default = 0;
+
+    config.GetDurationMs("legacyDurationMs", &legacy_ms, 0);
+    config.GetDurationMs("durationSeconds", &seconds, 0);
+    config.GetDurationMs("durationMinutes", &minutes, 0);
+    config.GetDurationMs("durationHours", &hours, 0);
+    config.GetDurationMs("durationWhitespace", &whitespace, 0);
+    config.GetDurationMs("missingDuration", &missing_default, 4321);
+
+    ASSERT_EQ(legacy_ms, 6000);
+    ASSERT_EQ(seconds, 6000);
+    ASSERT_EQ(minutes, 7 * 60 * 1000);
+    ASSERT_EQ(hours, 60 * 60 * 1000);
+    ASSERT_EQ(whitespace, 8 * 60 * 60 * 1000);
+    ASSERT_EQ(missing_default, 4321);
+}
+
 TEST_F(DefaultConfigTest, LoadInvalidFile) {
     DefaultConfig config;
     config.SetPath(path_ + "/invalid_file.txt");
@@ -87,4 +175,4 @@ TEST_F(DefaultConfigTest, LoadInvalidFile) {
 int main(int argc, char **argv) {
     ::testing::InitGoogleTest(&argc, argv);
     return RUN_ALL_TESTS();
-}
+}

mooncake-common/tests/test.json

@@ -3,8 +3,13 @@
     "testFloat": 4.15,
     "testString": "Hello, World",
     "testBoolean": false,
+    "legacyDurationMs": 5000,
+    "durationSeconds": "5s",
+    "durationMinutes": "30m",
+    "durationHours": "1H",
+    "durationWhitespace": " 7 m ",
     "testObject": {
         "nestedInteger": 1000,
         "nestedString": "Nested Hello, World"
     }
-}
+}

mooncake-common/tests/test.yaml

@@ -2,6 +2,11 @@ testInteger: 42
 testFloat: 3.14
 testString: "Hello, World!"
 testBoolean: true
+legacyDurationMs: 6000
+durationSeconds: 6s
+durationMinutes: "7m"
+durationHours: 1H
+durationWhitespace: " 8 h "
 testObject:
   nestedInteger: 100
   nestedString: "Nested Hello"

mooncake-ep/BuildEpExt.cmake

@@ -0,0 +1,93 @@
+# BuildEpExt.cmake - Build the Mooncake EP Python extension.
+#
+# Invoked at build time via cmake -P from the root CMakeLists.txt when
+# WITH_EP=ON.  Variables are passed with -D from the custom target:
+#
+#   SOURCE_DIR          - mooncake-ep source directory
+#   EP_CUDA_MAJOR       - CUDA major 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
+#   STAGING_DIR         - destination directory for the built .so files
+#   ENGINE_SO_PATH      - absolute path to the built engine.cpython-XYZ.so
+
+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}")
+endif()
+if(TORCH_CUDA_ARCH_LIST)
+  string(REPLACE "|" ";" TORCH_CUDA_ARCH_LIST "${TORCH_CUDA_ARCH_LIST}")
+endif()
+
+# ---------------------------------------------------------------------------
+# 1. Set up the build environment.
+# ---------------------------------------------------------------------------
+# Clear jobserver variables so that sub-processes started by setup.py do not
+# try to connect to the parent ninja's jobserver pipe FDs, which are not
+# inherited and cause: "ninja: error: Could not initialize jobserver: Invalid
+# file descriptors".
+set(ENV{MAKEFLAGS} "")
+set(ENV{MFLAGS} "")
+set(ENV{TORCH_CUDA_ARCH_LIST} "${TORCH_CUDA_ARCH_LIST}")
+
+# ---------------------------------------------------------------------------
+# 2. Ensure engine.so exists in mooncake-wheel/mooncake/ for setup.py linking.
+# ---------------------------------------------------------------------------
+# setup.py links against -l:engine.so in ../mooncake-wheel/mooncake/.
+# During the make phase only the versioned engine.cpython-XYZ.so exists in
+# the build tree; create a bare engine.so symlink so the linker can find it.
+set(_wheel_mooncake_dir "${SOURCE_DIR}/../mooncake-wheel/mooncake")
+set(_engine_symlink "${_wheel_mooncake_dir}/engine.so")
+if(ENGINE_SO_PATH AND NOT EXISTS "${_engine_symlink}")
+  message(STATUS "[EP] Creating engine.so symlink -> ${ENGINE_SO_PATH}")
+  execute_process(
+    COMMAND ${CMAKE_COMMAND} -E create_symlink "${ENGINE_SO_PATH}" "${_engine_symlink}"
+  )
+endif()
+
+# ---------------------------------------------------------------------------
+# 3. Build the EP Python extension.
+# ---------------------------------------------------------------------------
+if("${EP_TORCH_VERSIONS}" STREQUAL "")
+  message(STATUS "[EP] Building with currently-installed PyTorch")
+  execute_process(
+    COMMAND ${Python3_EXECUTABLE} setup.py build_ext --build-lib .
+    WORKING_DIRECTORY "${SOURCE_DIR}"
+    RESULT_VARIABLE _ret
+  )
+  if(NOT _ret EQUAL 0)
+    message(FATAL_ERROR "[EP] Extension build failed (exit code: ${_ret})")
+  endif()
+else()
+  message(STATUS "[EP] Building for PyTorch versions: ${EP_TORCH_VERSIONS}")
+  foreach(_version IN LISTS EP_TORCH_VERSIONS)
+    install_pytorch_wheel("${_version}" "${EP_CUDA_MAJOR}" "${EP_CUDA_MINOR}" "[EP]")
+
+    execute_process(
+      COMMAND ${Python3_EXECUTABLE} setup.py build_ext --build-lib . --force
+      WORKING_DIRECTORY "${SOURCE_DIR}"
+      RESULT_VARIABLE _ret
+    )
+    if(NOT _ret EQUAL 0)
+      message(FATAL_ERROR "[EP] Extension build failed for PyTorch ${_version}")
+    endif()
+  endforeach()
+endif()
+
+# ---------------------------------------------------------------------------
+# 4. Copy the built .so files to the staging directory.
+# ---------------------------------------------------------------------------
+file(MAKE_DIRECTORY "${STAGING_DIR}")
+file(GLOB _so_files "${SOURCE_DIR}/mooncake/*.so")
+foreach(_so IN LISTS _so_files)
+  get_filename_component(_fname "${_so}" NAME)
+  message(STATUS "[EP] Staging ${_fname} -> ${STAGING_DIR}")
+  file(COPY "${_so}" DESTINATION "${STAGING_DIR}" NO_SOURCE_PERMISSIONS)
+endforeach()
+
+message(STATUS "[EP] Mooncake EP extension build complete")

mooncake-ep/CMakeLists.txt

@@ -30,13 +30,4 @@ find_package(Torch REQUIRED)
 include_directories(${TORCH_INCLUDE_DIRS})
 
 include_directories(include)
-add_subdirectory(include)
 add_subdirectory(src)
-
-if (BUILD_UNIT_TESTS)
-  add_subdirectory(tests)
-endif()
-
-if (BUILD_EXAMPLES)
-  add_subdirectory(example)
-endif()

mooncake-ep/include/mooncake_ep_buffer.h

@@ -3,6 +3,7 @@
 
 #include <ATen/cuda/CUDAContext.h>
 #include <cuda_bf16.h>
+#include <cuda.h>
 #include <cuda_runtime.h>
 #include <fstream>
 #include <mooncake_ibgda/memheap.h>
@@ -71,7 +72,7 @@ struct MooncakeEpBuffer {
     void* gdr_buffer = nullptr;
 
     // IBGDA
-    static constexpr size_t CTRL_BUF_SIZE = 1024 * 1024 * 1024;  // 1024 MiB
+    static constexpr size_t CTRL_BUF_SIZE = 1024ULL * 1024 * 1024;  // 1024 MiB
     void* ctrl_buf = nullptr;
     // RDMA memory region for `gdr_buffer`. Must be nullptr when IBGDA init
     // fails.
@@ -85,6 +86,17 @@ struct MooncakeEpBuffer {
     bool is_roce_ = false;
     bool ibgda_disabled_ = false;
     int gid_index_ = -1;  // Dynamically discovered GID index
+    int USE_QP_COUNT = MAX_QP_COUNT;
+
+    mlx5dv_devx_umem* ctrl_buf_umem;
+    ibv_pd* pd;
+    mlx5dv_pd mpd;
+    memheap* ctrl_buf_heap;
+
+    // Fabric memory (MNNVL)
+    bool use_fabric_mem_ = false;
+    CUmemGenericAllocationHandle fabric_mem_handle_{};
+    size_t fabric_alloc_size_ = 0;
 
     // NVLink P2P
     int32_t* nvlink_available = nullptr;
@@ -156,16 +168,20 @@ struct MooncakeEpBuffer {
         return p2p_ipc_all_enabled_;
     }
 
+    void update_local_qpns();
+
     void sync_ib(const std::vector<int64_t>& remote_addrs,
                  const std::vector<int32_t>& remote_keys,
                  const std::vector<int32_t>& remote_qpns,
-                 const std::vector<int32_t>& remote_lids);
+                 const std::vector<int32_t>& remote_lids,
+                 const std::vector<int>& active_ranks_mask);
 
     void sync_roce(const std::vector<int64_t>& remote_addrs,
                    const std::vector<int32_t>& remote_keys,
                    const std::vector<int32_t>& remote_qpns,
                    const std::vector<int64_t>& subnet_prefixes,
-                   const std::vector<int64_t>& interface_ids);
+                   const std::vector<int64_t>& interface_ids,
+                   const std::vector<int>& active_ranks_mask);
 
     std::tuple<int64_t, int32_t> get_mr_info() {
         return {(int64_t)mr->addr, (int32_t)mr->rkey};
@@ -178,7 +194,7 @@ struct MooncakeEpBuffer {
 
     std::vector<int32_t> get_local_qpns() {
         std::vector<int32_t> local_qpns;
-        for (int i = 0; i < MAX_QP_COUNT; ++i) {
+        for (int i = 0; i < USE_QP_COUNT; ++i) {
             local_qpns.push_back((int32_t)qps[i]->qpn);
         }
         return local_qpns;
@@ -186,7 +202,7 @@ struct MooncakeEpBuffer {
 
     std::vector<int32_t> get_local_lids() {
         std::vector<int32_t> local_lids;
-        for (int i = 0; i < MAX_QP_COUNT; ++i) {
+        for (int i = 0; i < USE_QP_COUNT; ++i) {
             local_lids.push_back((int32_t)qps[i]->port_attr.lid);
         }
         return local_lids;
@@ -194,7 +210,8 @@ struct MooncakeEpBuffer {
 
     std::vector<int32_t> get_ipc_handle();
     void sync_nvlink_ipc_handles(
-        const std::vector<std::vector<int32_t>>& remote_handles);
+        const std::vector<std::vector<int32_t>>& remote_handles,
+        const std::vector<int>& active_ranks_mask);
 };
 
 inline size_t get_ep_buffer_size_hint(int num_max_dispatch_tokens_per_rank,

mooncake-ep/include/mooncake_ibgda/memheap.h

@@ -5,68 +5,153 @@
 #include <stdlib.h>
 #include <stdint.h>
 #include <stdalign.h>
+#include <stdbool.h>
+#include <errno.h>
 
 #include "os.h"
 
+#define MEMHEAP_MAX_ALLOCATIONS 1024
+
+struct memheap_allocation {
+    size_t offset;
+    size_t size;
+    bool used;
+};
+
 struct memheap {
     size_t size;
     pthread_mutex_t lock;
     size_t allocated;
+    struct memheap_allocation allocs[MEMHEAP_MAX_ALLOCATIONS];
+    int alloc_count;
 };
 
-static inline struct memheap *memheap_create(size_t size) {
-    struct memheap *heap = (struct memheap *)malloc(sizeof(struct memheap));
+static inline struct memheap* memheap_create(size_t size) {
+    struct memheap* heap = (struct memheap*)malloc(sizeof(struct memheap));
     if (!heap) {
         return NULL;
     }
     heap->size = size;
     heap->allocated = 0;
+    heap->alloc_count = 0;
     mutex_init(&heap->lock);
     return heap;
 }
 
-static inline void memheap_destroy(struct memheap *heap) {
+static inline void memheap_destroy(struct memheap* heap) {
     if (heap) {
         mutex_destroy(&heap->lock);
         free(heap);
     }
 }
 
-static inline size_t memheap_aligned_alloc(struct memheap *heap, size_t size,
+static inline size_t memheap_aligned_alloc(struct memheap* heap, size_t size,
                                            size_t align) {
     if (size == 0) {
-        return 0;  // No allocation for zero size
+        return (size_t)-1;  // No allocation for zero size
     }
     if (align == 0 || (align & (align - 1)) != 0) {
         errno = EINVAL;  // Invalid alignment
-        return -1;
+        return (size_t)-1;
     }
-    size_t ret = -1;
+
     mutex_lock(&heap->lock);
-    size_t offset = heap->allocated;
-    if (offset & (align - 1)) {
-        offset = (offset | (align - 1)) + 1;
+
+    size_t ret = (size_t)-1;
+
+    for (int i = 0; i < heap->alloc_count; i++) {
+        if (!heap->allocs[i].used) {
+            size_t offset = heap->allocs[i].offset;
+            size_t block_size = heap->allocs[i].size;
+
+            size_t aligned_offset = offset;
+            if (aligned_offset & (align - 1)) {
+                aligned_offset = (aligned_offset | (align - 1)) + 1;
+            }
+
+            if (aligned_offset + size <= offset + block_size) {
+                if (aligned_offset > offset) {
+                    int new_idx = heap->alloc_count;
+                    if (new_idx < MEMHEAP_MAX_ALLOCATIONS) {
+                        heap->allocs[new_idx].offset = offset;
+                        heap->allocs[new_idx].size = aligned_offset - offset;
+                        heap->allocs[new_idx].used = false;
+                        heap->alloc_count++;
+                    }
+                }
+
+                if (aligned_offset + size < offset + block_size) {
+                    int new_idx = heap->alloc_count;
+                    if (new_idx < MEMHEAP_MAX_ALLOCATIONS) {
+                        heap->allocs[new_idx].offset = aligned_offset + size;
+                        heap->allocs[new_idx].size =
+                            offset + block_size - (aligned_offset + size);
+                        heap->allocs[new_idx].used = false;
+                        heap->alloc_count++;
+                    }
+                }
+
+                heap->allocs[i].offset = aligned_offset;
+                heap->allocs[i].size = size;
+                heap->allocs[i].used = true;
+
+                ret = aligned_offset;
+                heap->allocated += size;
+                break;
+            }
+        }
     }
-    if (offset + size <= heap->size) {
-        ret = offset;
-        heap->allocated = offset + size;
-    } else {
-        errno = ENOMEM;  // Not enough memory
+
+    if (ret == (size_t)-1) {
+        size_t offset = heap->allocated;
+        if (offset & (align - 1)) {
+            offset = (offset | (align - 1)) + 1;
+        }
+        if (offset + size <= heap->size) {
+            ret = offset;
+
+            if (heap->alloc_count < MEMHEAP_MAX_ALLOCATIONS) {
+                heap->allocs[heap->alloc_count].offset = offset;
+                heap->allocs[heap->alloc_count].size = size;
+                heap->allocs[heap->alloc_count].used = true;
+                heap->alloc_count++;
+            }
+
+            heap->allocated = offset + size;
+        } else {
+            errno = ENOMEM;
+        }
     }
+
     mutex_unlock(&heap->lock);
     return ret;
 }
 
-static inline size_t memheap_alloc(struct memheap *heap, size_t size) {
+static inline size_t memheap_alloc(struct memheap* heap, size_t size) {
     size_t align = size & -size;
     if (align > alignof(max_align_t)) {
         align = alignof(max_align_t);
     }
+    if (align < 8) align = 8;
     return memheap_aligned_alloc(heap, size, align);
 }
 
-static inline void memheap_free(struct memheap *heap, size_t offset) {
-    // currently no-op
+static inline void memheap_free(struct memheap* heap, size_t offset) {
+    if (!heap || offset == (size_t)-1) {
+        return;
+    }
+
+    mutex_lock(&heap->lock);
+
+    for (int i = 0; i < heap->alloc_count; i++) {
+        if (heap->allocs[i].used && heap->allocs[i].offset == offset) {
+            heap->allocs[i].used = false;
+            heap->allocated -= heap->allocs[i].size;
+            break;
+        }
+    }
+
+    mutex_unlock(&heap->lock);
 }
 
-#endif
+#endif

mooncake-ep/include/mooncake_ibgda/mlx5gda.h

@@ -87,7 +87,7 @@ struct mlx5gda_qp *mlx5gda_create_rc_qp(struct mlx5dv_pd mpd, void *ctrl_buf,
                                         struct memheap *ctrl_buf_heap,
                                         struct ibv_pd *pd, int wqe,
                                         uint8_t port_num, cudaStream_t stream);
-void mlx5gda_destroy_qp(struct mlx5gda_qp *qp);
+void mlx5gda_destroy_qp(struct memheap *ctrl_buf_heap, struct mlx5gda_qp *qp);
 
 int mlx5gda_modify_rc_qp_rst2init(struct mlx5gda_qp *qp, uint16_t pkey_index);
 int mlx5gda_modify_rc_qp_init2rtr(struct mlx5gda_qp *qp,

mooncake-ep/setup.py

@@ -3,7 +3,7 @@ import re
 
 from setuptools import setup
 import torch
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME
 
 
 torch_version = re.match(r"\d+(?:\.\d+)*", torch.__version__).group()
@@ -13,6 +13,18 @@ module_name = "mooncake.ep" + version_suffix
 abi_flag = int(torch._C._GLIBCXX_USE_CXX11_ABI)
 current_dir = os.path.abspath(os.path.dirname(__file__))
 
+# Try to link against the CUDA driver stub library if it exists.
+cuda_libraries = ["ibverbs", "mlx5"]
+cuda_library_dirs = []
+
+if CUDA_HOME is not None:
+    cuda_stub_dir = os.path.join(CUDA_HOME, "lib64", "stubs")
+    cuda_stub_lib = os.path.join(cuda_stub_dir, "libcuda.so")
+    if os.path.exists(cuda_stub_lib):
+        cuda_libraries.insert(0, "cuda")
+        cuda_library_dirs.append(cuda_stub_dir)
+
+
 
 setup(
     name=module_name,
@@ -24,7 +36,7 @@ setup(
                 os.path.join(current_dir, "../mooncake-transfer-engine/include"),
             ],
             sources=[
-                "../mooncake-integration/ep/ep_py.cpp",
+                "src/ep_py.cpp",
                 "src/mooncake_ep_buffer.cpp",
                 "src/mooncake_ep_kernel.cu",
                 "src/mooncake_ibgda/mlx5gda.cpp",
@@ -33,7 +45,8 @@ setup(
                 "cxx": [f"-D_GLIBCXX_USE_CXX11_ABI={abi_flag}", "-std=c++20", "-O3", "-g0"],
                 "nvcc": [f"-D_GLIBCXX_USE_CXX11_ABI={abi_flag}", "-std=c++20", "-Xcompiler", "-O3", "-Xcompiler", "-g0"],
             },
-            libraries=["ibverbs", "mlx5"],
+            libraries=cuda_libraries,
+            library_dirs=cuda_library_dirs,
             extra_link_args=[
                 "-Wl,-rpath,$ORIGIN",
                 "-L" + os.path.join(current_dir, "../mooncake-wheel/mooncake"),

mooncake-ep/src/CMakeLists.txt

@@ -1,4 +1,4 @@
-add_library(mooncake_ep mooncake_backend.cpp mooncake_ep_buffer.cpp mooncake_ep_kernel.cu mooncake_worker.cu mooncake_worker_thread.cpp mooncake_ibgda/mlx5gda.cpp)
+add_library(mooncake_ep ep_py.cpp mooncake_ep_buffer.cpp mooncake_ep_kernel.cu mooncake_ibgda/mlx5gda.cpp)
 
 set_target_properties(mooncake_ep PROPERTIES POSITION_INDEPENDENT_CODE ON)
 target_link_libraries(mooncake_ep PUBLIC ${TORCH_LIBRARIES} transfer_engine ibverbs mlx5)

mooncake-ep/src/ep_py.cpp

@@ -24,6 +24,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         .def(py::init<int, int, int64_t, std::string>())
         .def("ibgda_disabled", &MooncakeEpBuffer::ibgda_disabled)
         .def("use_fast_path", &MooncakeEpBuffer::use_fast_path)
+        .def("update_local_qpns", &MooncakeEpBuffer::update_local_qpns)
         .def("is_roce", &MooncakeEpBuffer::is_roce)
         .def("sync_ib", &MooncakeEpBuffer::sync_ib)
         .def("sync_roce", &MooncakeEpBuffer::sync_roce)

mooncake-ep/src/mooncake_ep_buffer.cpp

@@ -1,8 +1,30 @@
 #include <mooncake_ep_buffer.h>
 #include <arpa/inet.h>
+#include <glog/logging.h>
 
 namespace mooncake {
 
+// Check if all GPUs support fabric memory handles (MNNVL).
+// Mirrors the check in nvlink_transport.cpp.
+static bool supportFabricMem() {
+    const char* nvlink_ipc = getenv("MC_USE_NVLINK_IPC");
+
+    bool fabric_enabled = nvlink_ipc && strcmp(nvlink_ipc, "0") == 0;
+    if (!fabric_enabled) return false;
+
+    int num_devices = 0;
+    cudaError_t err = cudaGetDeviceCount(&num_devices);
+    if (err != cudaSuccess || num_devices == 0) return false;
+
+    for (int dev = 0; dev < num_devices; ++dev) {
+        int supported = 0;
+        cuDeviceGetAttribute(
+            &supported, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, dev);
+        if (!supported) return false;
+    }
+    return true;
+}
+
 // Check if IPv6 address is an IPv4-mapped address (::ffff:x.x.x.x)
 static inline bool ipv6_addr_v4mapped(const struct in6_addr* a) {
     return ((a->s6_addr32[0] | a->s6_addr32[1]) == 0 &&
@@ -40,15 +62,100 @@ MooncakeEpBuffer::MooncakeEpBuffer(int rank, int num_ranks,
       num_ep_buffer_bytes(num_ep_buffer_bytes),
       device_name(std::move(device_name)),
       comm_stream(at::cuda::getStreamFromPool(true)) {
+    USE_QP_COUNT = MAX_QP_COUNT / num_ranks * num_ranks;
     // Get ranks
     CUDA_CHECK(cudaGetDevice(&device_id));
     CUDA_CHECK(cudaDeviceGetAttribute(&clock_rate_khz, cudaDevAttrClockRate,
                                       device_id));
-    CUDA_CHECK(cudaMalloc(&gdr_buffer, num_ep_buffer_bytes));
+
+    // Allocate gdr_buffer. On MNNVL clusters, use cuMemCreate with a fabric
+    // handle so the buffer is accessible cross-node via NVLink fabric.
+    // On IB clusters or single-node setups, fall back to cudaMalloc.
+    use_fabric_mem_ = supportFabricMem();
+    if (use_fabric_mem_) {
+        CUdevice cu_dev;
+        CUresult res = cuDeviceGet(&cu_dev, device_id);
+        if (res != CUDA_SUCCESS) {
+            LOG(ERROR) << "[EP] cuDeviceGet failed: " << res;
+            throw std::runtime_error("cuDeviceGet failed");
+        }
+
+        CUmemAllocationProp prop = {};
+        prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
+        prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
+        prop.location.id = cu_dev;
+        prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC;
+
+        int rdma_flag = 0;
+        cuDeviceGetAttribute(
+            &rdma_flag,
+            CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED,
+            cu_dev);
+        if (rdma_flag) prop.allocFlags.gpuDirectRDMACapable = 1;
+
+        size_t granularity = 0;
+        res = cuMemGetAllocationGranularity(&granularity, &prop,
+                                            CU_MEM_ALLOC_GRANULARITY_MINIMUM);
+        if (res != CUDA_SUCCESS) {
+            LOG(ERROR) << "[EP] cuMemGetAllocationGranularity failed: " << res;
+            throw std::runtime_error("cuMemGetAllocationGranularity failed");
+        }
+
+        fabric_alloc_size_ =
+            (num_ep_buffer_bytes + granularity - 1) & ~(granularity - 1);
+        if (fabric_alloc_size_ == 0) fabric_alloc_size_ = granularity;
+
+        res = cuMemCreate(&fabric_mem_handle_, fabric_alloc_size_, &prop, 0);
+        if (res != CUDA_SUCCESS) {
+            LOG(ERROR) << "[EP] cuMemCreate(FABRIC) failed: " << res;
+            throw std::runtime_error("cuMemCreate failed");
+        }
+
+        CUdeviceptr dptr = 0;
+        res = cuMemAddressReserve(&dptr, fabric_alloc_size_, granularity, 0, 0);
+        if (res != CUDA_SUCCESS) {
+            cuMemRelease(fabric_mem_handle_);
+            LOG(ERROR) << "[EP] cuMemAddressReserve failed: " << res;
+            throw std::runtime_error("cuMemAddressReserve failed");
+        }
+
+        res = cuMemMap(dptr, fabric_alloc_size_, 0, fabric_mem_handle_, 0);
+        if (res != CUDA_SUCCESS) {
+            cuMemAddressFree(dptr, fabric_alloc_size_);
+            cuMemRelease(fabric_mem_handle_);
+            LOG(ERROR) << "[EP] cuMemMap failed: " << res;
+            throw std::runtime_error("cuMemMap failed");
+        }
+
+        // Grant read/write access to all devices in the fabric clique
+        int device_count = 0;
+        cudaGetDeviceCount(&device_count);
+        std::vector<CUmemAccessDesc> access(device_count);
+        for (int i = 0; i < device_count; ++i) {
+            access[i].location.type = CU_MEM_LOCATION_TYPE_DEVICE;
+            access[i].location.id = i;
+            access[i].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
+        }
+        res = cuMemSetAccess(dptr, fabric_alloc_size_, access.data(),
+                             device_count);
+        if (res != CUDA_SUCCESS) {
+            cuMemUnmap(dptr, fabric_alloc_size_);
+            cuMemAddressFree(dptr, fabric_alloc_size_);
+            cuMemRelease(fabric_mem_handle_);
+            LOG(ERROR) << "[EP] cuMemSetAccess failed: " << res;
+            throw std::runtime_error("cuMemSetAccess failed");
+        }
+
+        gdr_buffer = reinterpret_cast<void*>(dptr);
+        LOG(INFO) << "[EP] Allocated " << fabric_alloc_size_
+                  << " bytes with fabric handle on GPU " << device_id;
+    } else {
+        CUDA_CHECK(cudaMalloc(&gdr_buffer, num_ep_buffer_bytes));
+    }
     CUDA_CHECK(cudaMalloc(&raddrs, num_ranks * sizeof(uint64_t)));
     CUDA_CHECK(cudaMalloc(&rkeys, num_ranks * sizeof(uint32_t)));
     CUDA_CHECK(
-        cudaMalloc(&qp_devctxs, MAX_QP_COUNT * sizeof(mlx5gda_qp_devctx)));
+        cudaMalloc(&qp_devctxs, USE_QP_COUNT * sizeof(mlx5gda_qp_devctx)));
 
     // Allocate NVLink P2P arrays
     CUDA_CHECK(cudaMalloc(&nvlink_available, num_ranks * sizeof(int32_t)));
@@ -71,7 +178,14 @@ MooncakeEpBuffer::MooncakeEpBuffer(int rank, int num_ranks,
 }
 
 MooncakeEpBuffer::~MooncakeEpBuffer() noexcept(false) {
-    cudaFree(gdr_buffer);
+    if (use_fabric_mem_) {
+        CUdeviceptr dptr = reinterpret_cast<CUdeviceptr>(gdr_buffer);
+        cuMemUnmap(dptr, fabric_alloc_size_);
+        cuMemAddressFree(dptr, fabric_alloc_size_);
+        cuMemRelease(fabric_mem_handle_);
+    } else {
+        cudaFree(gdr_buffer);
+    }
     cudaFree(raddrs);
     cudaFree(rkeys);
     cudaFree(qp_devctxs);
@@ -108,7 +222,7 @@ MooncakeEpBuffer::dispatch(const torch::Tensor& x,
                    x.size(0) <= num_max_dispatch_tokens_per_rank);
     EP_HOST_ASSERT(topk_idx.scalar_type() == torch::kInt64);
     EP_HOST_ASSERT(num_experts % num_ranks == 0);
-    EP_HOST_ASSERT(MAX_QP_COUNT % num_ranks == 0);
+    EP_HOST_ASSERT(USE_QP_COUNT % num_ranks == 0);
 
     auto num_tokens = static_cast<int>(x.size(0)),
          hidden = static_cast<int>(x.size(1));
@@ -381,12 +495,11 @@ int MooncakeEpBuffer::init_ibgda() {
     }
     ibv_free_device_list(dev_list);
 
-    ibv_pd* pd = ibv_alloc_pd(ctx);
+    pd = ibv_alloc_pd(ctx);
     if (!pd) {
         perror("Failed to allocate protection domain");
         return -1;
     }
-    mlx5dv_pd mpd;
     mlx5dv_obj dv_obj = {};
     dv_obj.pd.in = pd;
     dv_obj.pd.out = &mpd;
@@ -404,8 +517,8 @@ int MooncakeEpBuffer::init_ibgda() {
     // initialized as needed: CQ needs -1 (hardware requirement), DBR needs 0.
     // WQ doesn't need initialization as it's zeroed before each use.
     CUDA_CHECK(cudaMalloc(&ctrl_buf, CTRL_BUF_SIZE));
-    mlx5dv_devx_umem* ctrl_buf_umem = mlx5dv_devx_umem_reg(
-        ctx, ctrl_buf, CTRL_BUF_SIZE, IBV_ACCESS_LOCAL_WRITE);
+    ctrl_buf_umem = mlx5dv_devx_umem_reg(ctx, ctrl_buf, CTRL_BUF_SIZE,
+                                         IBV_ACCESS_LOCAL_WRITE);
     if (!ctrl_buf_umem) {
         perror("Failed to register control buffer as umem");
         fprintf(stderr,
@@ -419,14 +532,14 @@ int MooncakeEpBuffer::init_ibgda() {
         }
         return -1;
     }
-    memheap* ctrl_buf_heap = memheap_create(CTRL_BUF_SIZE);
+    ctrl_buf_heap = memheap_create(CTRL_BUF_SIZE);
     if (!ctrl_buf_heap) {
         perror("Failed to create memory heap");
         return -1;
     }
     // Individual regions (CQ, DBR) will be initialized as needed via async
     // memset.
-    for (int i = 0; i < MAX_QP_COUNT; ++i) {
+    for (int i = 0; i < USE_QP_COUNT; ++i) {
         mlx5gda_qp* qp =
             mlx5gda_create_rc_qp(mpd, ctrl_buf, ctrl_buf_umem, ctrl_buf_heap,
                                  pd, 16384, 1, comm_stream.stream());
@@ -458,11 +571,55 @@ int MooncakeEpBuffer::init_ibgda() {
     return 0;
 }
 
+void MooncakeEpBuffer::update_local_qpns() {
+    for (int i = 0; i < USE_QP_COUNT; ++i) {
+        if (qps[i]) {
+            mlx5gda_destroy_qp(ctrl_buf_heap, qps[i]);
+            qps[i] = nullptr;
+        }
+    }
+
+    for (int i = 0; i < USE_QP_COUNT; ++i) {
+        mlx5gda_qp* qp =
+            mlx5gda_create_rc_qp(mpd, ctrl_buf, ctrl_buf_umem, ctrl_buf_heap,
+                                 pd, 16384, 1, comm_stream.stream());
+        if (!qp) {
+            perror("Failed to recreate QP");
+            ibgda_disabled_ = true;
+            return;
+        }
+        is_roce_ = qp->port_attr.link_layer == IBV_LINK_LAYER_ETHERNET;
+        if (mlx5gda_modify_rc_qp_rst2init(qp, 0)) {
+            perror("Failed to mlx5gda_modify_rc_qp_rst2init");
+            ibgda_disabled_ = true;
+            return;
+        }
+        // Ensure all async memset operations are complete before accessing QP
+        // structures
+        CUDA_CHECK(cudaStreamSynchronize(comm_stream.stream()));
+
+        mlx5gda_qp_devctx qp_devctx = {
+            .qpn = qp->qpn,
+            .wqeid_mask = qp->num_wqebb - 1,
+            .wq = (mlx5gda_wqebb*)(ctrl_buf + qp->wq_offset),
+            .cq = (mlx5_cqe64*)(ctrl_buf + qp->send_cq->cq_offset),
+            .dbr = (mlx5gda_wq_dbr*)(ctrl_buf + qp->dbr_offset),
+            .bf = (char*)qp->uar->reg_addr,
+        };
+        cudaMemcpy(qp_devctxs + i * sizeof(mlx5gda_qp_devctx), &qp_devctx,
+                   sizeof(mlx5gda_qp_devctx), cudaMemcpyHostToDevice);
+        qps[i] = qp;
+    }
+}
+
 void MooncakeEpBuffer::sync_ib(const std::vector<int64_t>& remote_addrs,
                                const std::vector<int32_t>& remote_keys,
                                const std::vector<int32_t>& remote_qpns,
-                               const std::vector<int32_t>& remote_lids) {
-    for (int i = 0; i < MAX_QP_COUNT; ++i) {
+                               const std::vector<int32_t>& remote_lids,
+                               const std::vector<int>& active_ranks_mask) {
+    for (int i = 0; i < USE_QP_COUNT; ++i) {
+        int peer_rank = i * num_ranks / USE_QP_COUNT;
+        if (active_ranks_mask[peer_rank] == 0) continue;
         ibv_ah_attr ah_attr = {
             .dlid = (uint16_t)remote_lids[i],
             .port_num = 0,
@@ -478,6 +635,7 @@ void MooncakeEpBuffer::sync_ib(const std::vector<int64_t>& remote_addrs,
         }
     }
     for (int i = 0; i < num_ranks; ++i) {
+        if (active_ranks_mask[i] == 0) continue;
         uint64_t raddr =
             i == rank ? (uint64_t)mr->addr : (uint64_t)remote_addrs[i];
         cudaMemcpy(raddrs + i * sizeof(uint64_t), &raddr, sizeof(uint64_t),
@@ -492,13 +650,14 @@ void MooncakeEpBuffer::sync_roce(const std::vector<int64_t>& remote_addrs,
                                  const std::vector<int32_t>& remote_keys,
                                  const std::vector<int32_t>& remote_qpns,
                                  const std::vector<int64_t>& subnet_prefixes,
-                                 const std::vector<int64_t>& interface_ids) {
-    for (int i = 0; i < MAX_QP_COUNT; ++i) {
+                                 const std::vector<int64_t>& interface_ids,
+                                 const std::vector<int>& active_ranks_mask) {
+    for (int i = 0; i < USE_QP_COUNT; ++i) {
+        int peer_rank = i * num_ranks / USE_QP_COUNT;
+        if (active_ranks_mask[peer_rank] == 0) continue;
         ibv_gid remote_gid{};
-        remote_gid.global.subnet_prefix =
-            subnet_prefixes[i * num_ranks / MAX_QP_COUNT];
-        remote_gid.global.interface_id =
-            interface_ids[i * num_ranks / MAX_QP_COUNT];
+        remote_gid.global.subnet_prefix = subnet_prefixes[peer_rank];
+        remote_gid.global.interface_id = interface_ids[peer_rank];
         ibv_ah_attr ah_attr = {};
         ah_attr.is_global = 1;
         ah_attr.grh.dgid = remote_gid;
@@ -518,6 +677,7 @@ void MooncakeEpBuffer::sync_roce(const std::vector<int64_t>& remote_addrs,
         }
     }
     for (int i = 0; i < num_ranks; ++i) {
+        if (active_ranks_mask[i] == 0) continue;
         uint64_t raddr =
             i == rank ? (uint64_t)mr->addr : (uint64_t)remote_addrs[i];
         cudaMemcpy(raddrs + i * sizeof(uint64_t), &raddr, sizeof(uint64_t),
@@ -529,6 +689,12 @@ void MooncakeEpBuffer::sync_roce(const std::vector<int64_t>& remote_addrs,
 }
 
 std::vector<int32_t> MooncakeEpBuffer::get_ipc_handle() {
+    if (use_fabric_mem_) {
+        // Fabric memory is globally accessible via cuMemSetAccess — no IPC
+        // handle exchange needed. Return an empty vector so the caller knows
+        // to skip IPC for this rank.
+        return {};
+    }
     cudaIpcMemHandle_t handle;
     CUDA_CHECK(cudaIpcGetMemHandle(&handle, gdr_buffer));
     // Convert handle bytes to int32_t array
@@ -541,107 +707,114 @@ std::vector<int32_t> MooncakeEpBuffer::get_ipc_handle() {
 }
 
 void MooncakeEpBuffer::sync_nvlink_ipc_handles(
-    const std::vector<std::vector<int32_t>>& remote_handles) {
-    // We assume ranks are grouped by device_count (same node)
+    const std::vector<std::vector<int32_t>>& remote_handles,
+    const std::vector<int>& active_ranks_mask) {
     int device_count = 0;
     CUDA_CHECK(cudaGetDeviceCount(&device_count));
 
     std::vector<int32_t> nvlink_array(num_ranks, 0);
     nvlink_array[rank] = 1;
 
-    int node_id = rank / device_count;
-    int group_start = node_id * device_count;
-    int group_end = std::min(group_start + device_count, num_ranks);
+    if (use_fabric_mem_) {
+        // MNNVL: fabric addresses are globally visible across the clique.
+        // All ranks can directly access each other's gdr_buffer without IPC
+        // handle exchange — cuMemSetAccess already granted all devices
+        // read/write access during allocation.
+        for (int i = 0; i < num_ranks; ++i) {
+            if (active_ranks_mask[i] == 0) continue;
+            nvlink_array[i] = 1;
+            // Each rank's gdr_buffer is directly accessible; the remote
+            // addresses will be exchanged via the RDMA address sync path
+            // (sync_ib / sync_roce) or via a separate fabric address exchange.
+            // For local rank, point to our own buffer.
+            ipc_peer_ptrs_host[i] = (i == rank) ? gdr_buffer : nullptr;
+        }
+        p2p_ipc_all_enabled_ = true;
+        LOG(INFO) << "[EP] Fabric memory enabled, skipping IPC handle exchange";
+    } else {
+        // Non-MNNVL: use cudaIpc for intra-node P2P (original path)
+        int node_id = rank / device_count;
+        int group_start = node_id * device_count;
+        int group_end = std::min(group_start + device_count, num_ranks);
 
-    // Check peer access and enable it within the same node group
-    for (int dst_rank = group_start; dst_rank < group_end; ++dst_rank) {
-        if (dst_rank == rank) {
-            // Local rank - use local pointer
-            ipc_peer_ptrs_host[dst_rank] = gdr_buffer;
-            continue;
+        for (int dst_rank = group_start; dst_rank < group_end; ++dst_rank) {
+            if (active_ranks_mask[dst_rank] == 0) continue;
+            if (dst_rank == rank) {
+                ipc_peer_ptrs_host[dst_rank] = gdr_buffer;
+                continue;
+            }
+
+            int dst_device = dst_rank % device_count;
+            int can_access_peer = 0;
+            cudaError_t err = cudaDeviceCanAccessPeer(&can_access_peer,
+                                                      device_id, dst_device);
+            if (err == cudaSuccess && can_access_peer) {
+                cudaError_t peer_err =
+                    cudaDeviceEnablePeerAccess(dst_device, 0);
+                if (peer_err == cudaSuccess ||
+                    peer_err == cudaErrorPeerAccessAlreadyEnabled) {
+                    if (peer_err == cudaErrorPeerAccessAlreadyEnabled) {
+                        cudaGetLastError();
+                    }
+                    nvlink_array[dst_rank] = 1;
+
+                    if (dst_rank >= static_cast<int>(remote_handles.size())) {
+                        LOG(WARNING)
+                            << "[EP] Rank " << rank
+                            << " missing IPC handle for rank " << dst_rank;
+                        continue;
+                    }
+
+                    const size_t handle_size = sizeof(cudaIpcMemHandle_t);
+                    const size_t num_int32s =
+                        (handle_size + sizeof(int32_t) - 1) / sizeof(int32_t);
+                    const auto& handle_ints = remote_handles[dst_rank];
+                    if (handle_ints.size() < num_int32s) {
+                        LOG(WARNING)
+                            << "[EP] Rank " << rank
+                            << " invalid IPC handle size for rank " << dst_rank;
+                        continue;
+                    }
+
+                    cudaIpcMemHandle_t remote_handle;
+                    memcpy(&remote_handle, handle_ints.data(), handle_size);
+
+                    void* peer_ptr = nullptr;
+                    cudaError_t ipc_err =
+                        cudaIpcOpenMemHandle(&peer_ptr, remote_handle,
+                                             cudaIpcMemLazyEnablePeerAccess);
+                    if (ipc_err != cudaSuccess) {
+                        LOG(WARNING)
+                            << "[EP] Rank " << rank
+                            << " failed to open IPC handle for rank "
+                            << dst_rank << ": " << cudaGetErrorString(ipc_err);
+                        nvlink_array[dst_rank] = 0;
+                    } else {
+                        ipc_peer_ptrs_host[dst_rank] = peer_ptr;
+                    }
+                }
+            }
         }
 
-        int dst_device = dst_rank % device_count;
-        int can_access_peer = 0;
-        cudaError_t err =
-            cudaDeviceCanAccessPeer(&can_access_peer, device_id, dst_device);
-        if (err == cudaSuccess && can_access_peer) {
-            cudaError_t peer_err = cudaDeviceEnablePeerAccess(dst_device, 0);
-            if (peer_err == cudaSuccess ||
-                peer_err == cudaErrorPeerAccessAlreadyEnabled) {
-                // Clear sticky error on re-init so CUDA graph capture /
-                // dispatch later does not see
-                // cudaErrorPeerAccessAlreadyEnabled.
-                if (peer_err == cudaErrorPeerAccessAlreadyEnabled) {
-                    cudaGetLastError();
-                }
-                nvlink_array[dst_rank] = 1;
-
-                // Open IPC handle for this peer
-                if (dst_rank >= static_cast<int>(remote_handles.size())) {
-                    LOG(WARNING) << "[EP] Rank " << rank
-                                 << " missing IPC handle for rank " << dst_rank;
-                    continue;
-                }
-
-                const size_t handle_size = sizeof(cudaIpcMemHandle_t);
-                const size_t num_int32s =
-                    (handle_size + sizeof(int32_t) - 1) / sizeof(int32_t);
-                const auto& handle_ints = remote_handles[dst_rank];
-                if (handle_ints.size() < num_int32s) {
-                    LOG(WARNING)
-                        << "[EP] Rank " << rank
-                        << " invalid IPC handle size for rank " << dst_rank;
-                    continue;
-                }
-
-                cudaIpcMemHandle_t remote_handle;
-                memcpy(&remote_handle, handle_ints.data(), handle_size);
-
-                void* peer_ptr = nullptr;
-                cudaError_t ipc_err = cudaIpcOpenMemHandle(
-                    &peer_ptr, remote_handle, cudaIpcMemLazyEnablePeerAccess);
-                if (ipc_err != cudaSuccess) {
-                    LOG(WARNING)
-                        << "[EP] Rank " << rank
-                        << " failed to open IPC handle for rank " << dst_rank
-                        << ": " << cudaGetErrorString(ipc_err);
-                    nvlink_array[dst_rank] = 0;
-                } else {
-                    ipc_peer_ptrs_host[dst_rank] = peer_ptr;
-                }
+        p2p_ipc_all_enabled_ = true;
+        for (int i = 0; i < num_ranks; ++i) {
+            if (active_ranks_mask[i] == 0) continue;
+            if (nvlink_array[i] == 0 || ipc_peer_ptrs_host[i] == nullptr) {
+                p2p_ipc_all_enabled_ = false;
+                break;
+            }
+        }
+        if (p2p_ipc_all_enabled_ && num_ranks > 1) {
+            int first_node_id = 0 / device_count;
+            int last_node_id = (num_ranks - 1) / device_count;
+            if (first_node_id != last_node_id) {
+                p2p_ipc_all_enabled_ = false;
             }
         }
     }
 
-    // Check if P2P+IPC is available for ALL rank pairs.
-    // For P2P+IPC to be fully usable without IBGDA, every rank must be able to
-    // access every other rank via P2P+IPC. Since we only check within the same
-    // node group, all ranks must be in the same node group.
-    p2p_ipc_all_enabled_ = true;
-    for (int i = 0; i < num_ranks; ++i) {
-        // Must have P2P enabled and a valid peer pointer for every rank.
-        // Note: for local rank we set ipc_peer_ptrs_host[rank] = gdr_buffer.
-        if (nvlink_array[i] == 0 || ipc_peer_ptrs_host[i] == nullptr) {
-            p2p_ipc_all_enabled_ = false;
-            break;
-        }
-    }
-    // Verify all ranks are in the same node group (cross-node requires IBGDA)
-    if (p2p_ipc_all_enabled_ && num_ranks > 1) {
-        int first_node_id = 0 / device_count;
-        int last_node_id = (num_ranks - 1) / device_count;
-        if (first_node_id != last_node_id) {
-            // Ranks span multiple nodes, P2P only works within nodes
-            p2p_ipc_all_enabled_ = false;
-        }
-    }
-
-    // Copy NVLink availability to device memory
     CUDA_CHECK(cudaMemcpy(nvlink_available, nvlink_array.data(),
                           num_ranks * sizeof(int32_t), cudaMemcpyHostToDevice));
-
-    // Copy IPC pointers to device memory for kernel access
     CUDA_CHECK(cudaMemcpy(ipc_peer_ptrs, ipc_peer_ptrs_host,
                           num_ranks * sizeof(void*), cudaMemcpyHostToDevice));
 }

mooncake-ep/src/mooncake_ibgda/mlx5gda.cpp

@@ -29,13 +29,13 @@ constexpr T round_up_pow2(T n) {
 
 #define IBGDA_ROUND_UP_POW2_OR_0(_n) (((_n) == 0) ? 0 : round_up_pow2(_n))
 
-static void print_cuda_error(const char *msg) {
-    const char *err_str = cudaGetErrorString(cudaGetLastError());
+static void print_cuda_error(const char* msg) {
+    const char* err_str = cudaGetErrorString(cudaGetLastError());
     fprintf(stderr, "%s: %s\n", msg, err_str);
 }
 
-static struct mlx5dv_devx_uar *create_uar(struct ibv_context *ctx) {
-    struct mlx5dv_devx_uar *uar =
+static struct mlx5dv_devx_uar* create_uar(struct ibv_context* ctx) {
+    struct mlx5dv_devx_uar* uar =
         mlx5dv_devx_alloc_uar(ctx, MLX5DV_UAR_ALLOC_TYPE_BF);
     if (!uar) {
         errno = EIO;
@@ -52,7 +52,7 @@ static struct mlx5dv_devx_uar *create_uar(struct ibv_context *ctx) {
     return uar;
 }
 
-static void destroy_uar(struct mlx5dv_devx_uar *uar) {
+static void destroy_uar(struct mlx5dv_devx_uar* uar) {
     if (!uar) return;
     if (cudaHostUnregister(uar->reg_addr) != cudaSuccess) {
         print_cuda_error("Failed to unregister MMIO memory");
@@ -60,21 +60,21 @@ static void destroy_uar(struct mlx5dv_devx_uar *uar) {
     mlx5dv_devx_free_uar(uar);
 }
 
-struct mlx5gda_cq *mlx5gda_create_cq(void *ctrl_buf,
-                                     struct mlx5dv_devx_umem *ctrl_buf_umem,
-                                     struct memheap *ctrl_buf_heap,
-                                     struct ibv_pd *pd, int cqe,
+struct mlx5gda_cq* mlx5gda_create_cq(void* ctrl_buf,
+                                     struct mlx5dv_devx_umem* ctrl_buf_umem,
+                                     struct memheap* ctrl_buf_heap,
+                                     struct ibv_pd* pd, int cqe,
                                      cudaStream_t stream) {
-    struct mlx5gda_cq *cq = NULL;
-    struct mlx5dv_devx_uar *uar = NULL;
+    struct mlx5gda_cq* cq = NULL;
+    struct mlx5dv_devx_uar* uar = NULL;
     uint32_t eqn = 0;
     size_t cq_offset = -1;
     size_t dbr_offset = -1;
-    struct mlx5dv_devx_obj *mlx5_cq = NULL;
+    struct mlx5dv_devx_obj* mlx5_cq = NULL;
     uint32_t cqn = 0;
 
-    struct ibv_context *ctx = pd->context;
-    void *cq_context = NULL;
+    struct ibv_context* ctx = pd->context;
+    void* cq_context = NULL;
 
     if (cqe <= 0) {
         errno = EINVAL;
@@ -107,7 +107,7 @@ struct mlx5gda_cq *mlx5gda_create_cq(void *ctrl_buf,
         perror("Failed to allocate DBR memory");
         goto fail;
     }
-    cq = (struct mlx5gda_cq *)malloc(sizeof(struct mlx5gda_cq));
+    cq = (struct mlx5gda_cq*)malloc(sizeof(struct mlx5gda_cq));
     if (!cq) goto fail;
     if (mlx5dv_devx_query_eqn(ctx, 0, &eqn)) {
         perror("Failed to query EQN");
@@ -166,7 +166,7 @@ fail:
     return NULL;
 }
 
-void mlx5gda_destroy_cq(struct memheap *ctrl_buf_heap, struct mlx5gda_cq *cq) {
+void mlx5gda_destroy_cq(struct memheap* ctrl_buf_heap, struct mlx5gda_cq* cq) {
     if (!cq) return;
     if (cq->mcq) {
         mlx5dv_devx_obj_destroy(cq->mcq);
@@ -179,21 +179,21 @@ void mlx5gda_destroy_cq(struct memheap *ctrl_buf_heap, struct mlx5gda_cq *cq) {
     free(cq);
 }
 
-struct mlx5gda_qp *mlx5gda_create_rc_qp(struct mlx5dv_pd mpd, void *ctrl_buf,
-                                        struct mlx5dv_devx_umem *ctrl_buf_umem,
-                                        struct memheap *ctrl_buf_heap,
-                                        struct ibv_pd *pd, int wqe,
+struct mlx5gda_qp* mlx5gda_create_rc_qp(struct mlx5dv_pd mpd, void* ctrl_buf,
+                                        struct mlx5dv_devx_umem* ctrl_buf_umem,
+                                        struct memheap* ctrl_buf_heap,
+                                        struct ibv_pd* pd, int wqe,
                                         uint8_t port_num, cudaStream_t stream) {
-    struct mlx5gda_qp *qp = NULL;
-    struct mlx5gda_cq *send_cq = NULL;
-    struct mlx5dv_devx_uar *uar = NULL;
-    struct mlx5dv_devx_obj *mlx5_qp = NULL;
+    struct mlx5gda_qp* qp = NULL;
+    struct mlx5gda_cq* send_cq = NULL;
+    struct mlx5dv_devx_uar* uar = NULL;
+    struct mlx5dv_devx_obj* mlx5_qp = NULL;
     size_t wq_offset = -1;
     size_t dbr_offset = -1;
 
-    struct ibv_context *ctx = pd->context;
-    void *qp_context = NULL;
-    void *cap = NULL;
+    struct ibv_context* ctx = pd->context;
+    void* qp_context = NULL;
+    void* cap = NULL;
     uint32_t cqe_version = 0;
 
     if (wqe <= 0) {
@@ -208,7 +208,7 @@ struct mlx5gda_qp *mlx5gda_create_rc_qp(struct mlx5dv_pd mpd, void *ctrl_buf,
     uint8_t cmd_cap_in[DEVX_ST_SZ_BYTES(query_hca_cap_in)] = {0};
     uint8_t cmd_cap_out[DEVX_ST_SZ_BYTES(query_hca_cap_out)] = {0};
 
-    qp = (struct mlx5gda_qp *)calloc(1, sizeof(struct mlx5gda_qp));
+    qp = (struct mlx5gda_qp*)calloc(1, sizeof(struct mlx5gda_qp));
     if (!qp) {
         perror("Failed to allocate QP memory");
         goto fail;
@@ -347,7 +347,28 @@ fail:
     return NULL;
 }
 
-int mlx5gda_modify_rc_qp_rst2init(struct mlx5gda_qp *qp, uint16_t pkey_index) {
+void mlx5gda_destroy_qp(struct memheap* ctrl_buf_heap, struct mlx5gda_qp* qp) {
+    if (qp->mqp) {
+        mlx5dv_devx_obj_destroy(qp->mqp);
+    }
+    if (qp->uar) {
+        destroy_uar(qp->uar);
+    }
+    if (qp->send_cq) {
+        mlx5gda_destroy_cq(ctrl_buf_heap, qp->send_cq);
+    }
+    if (qp->wq_offset != -1) {
+        memheap_free(ctrl_buf_heap, qp->wq_offset);
+    }
+    if (qp->dbr_offset != -1) {
+        memheap_free(ctrl_buf_heap, qp->dbr_offset);
+    }
+    if (qp) {
+        free(qp);
+    }
+}
+
+int mlx5gda_modify_rc_qp_rst2init(struct mlx5gda_qp* qp, uint16_t pkey_index) {
     if (!qp || !qp->mqp) {
         errno = EINVAL;
         return -1;
@@ -358,7 +379,7 @@ int mlx5gda_modify_rc_qp_rst2init(struct mlx5gda_qp *qp, uint16_t pkey_index) {
     DEVX_SET(rst2init_qp_in, cmd_in, opcode, MLX5_CMD_OP_RST2INIT_QP);
     DEVX_SET(rst2init_qp_in, cmd_in, qpn, qp->qpn);
 
-    void *qpc = DEVX_ADDR_OF(rst2init_qp_in, cmd_in, qpc);
+    void* qpc = DEVX_ADDR_OF(rst2init_qp_in, cmd_in, qpc);
 
     DEVX_SET(qpc, qpc, rwe, 1);
     DEVX_SET(qpc, qpc, rre, 1);
@@ -381,7 +402,7 @@ int mlx5gda_modify_rc_qp_rst2init(struct mlx5gda_qp *qp, uint16_t pkey_index) {
     return ret;
 }
 
-int mlx5gda_modify_rc_qp_init2rtr(struct mlx5gda_qp *qp,
+int mlx5gda_modify_rc_qp_init2rtr(struct mlx5gda_qp* qp,
                                   struct ibv_ah_attr ah_attr,
                                   uint32_t remote_qpn, enum ibv_mtu mtu) {
     if (!qp || !qp->mqp) {
@@ -389,14 +410,14 @@ int mlx5gda_modify_rc_qp_init2rtr(struct mlx5gda_qp *qp,
         return -1;
     }
     int ret = 0;
-    struct ibv_ah *ah = NULL;
+    struct ibv_ah* ah = NULL;
     uint8_t cmd_in[DEVX_ST_SZ_BYTES(init2rtr_qp_in)] = {0};
     uint8_t cmd_out[DEVX_ST_SZ_BYTES(init2rtr_qp_out)] = {0};
 
     DEVX_SET(rst2init_qp_in, cmd_in, opcode, MLX5_CMD_OP_INIT2RTR_QP);
     DEVX_SET(rst2init_qp_in, cmd_in, qpn, qp->qpn);
 
-    void *qpc = DEVX_ADDR_OF(rst2init_qp_in, cmd_in, qpc);
+    void* qpc = DEVX_ADDR_OF(rst2init_qp_in, cmd_in, qpc);
 
     DEVX_SET(qpc, qpc, mtu, mtu);
     DEVX_SET(qpc, qpc, log_msg_max, 30);
@@ -445,7 +466,7 @@ cleanup:
     return ret;
 }
 
-int mlx5gda_modify_rc_qp_rtr2rts(struct mlx5gda_qp *qp) {
+int mlx5gda_modify_rc_qp_rtr2rts(struct mlx5gda_qp* qp) {
     if (!qp || !qp->mqp) {
         errno = EINVAL;
         return -1;
@@ -456,7 +477,7 @@ int mlx5gda_modify_rc_qp_rtr2rts(struct mlx5gda_qp *qp) {
     DEVX_SET(rst2init_qp_in, cmd_in, opcode, MLX5_CMD_OP_RTR2RTS_QP);
     DEVX_SET(rst2init_qp_in, cmd_in, qpn, qp->qpn);
 
-    void *qpc = DEVX_ADDR_OF(rst2init_qp_in, cmd_in, qpc);
+    void* qpc = DEVX_ADDR_OF(rst2init_qp_in, cmd_in, qpc);
 
     DEVX_SET(qpc, qpc, log_ack_req_freq, 0x0);  // Ack every packet
     DEVX_SET(qpc, qpc, log_sra_max, 1);         // log2(max_qp_rd_atomic)

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
@@ -24,6 +24,12 @@ include_directories("../mooncake-transfer-engine/include")
 
 find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
 
+execute_process(
+    COMMAND ${Python3_EXECUTABLE} -c "import sysconfig; print(sysconfig.get_config_var('EXT_SUFFIX'))"
+    OUTPUT_VARIABLE PYTHON_EXT_SUFFIX
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+)
+
 set(CMAKE_INSTALL_RPATH_USE_LINK_PATH TRUE)
 set(CMAKE_BUILD_WITH_INSTALL_RPATH TRUE)
 
@@ -188,3 +194,25 @@ if (WITH_TE)
     install(TARGETS engine DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})
     install(TARGETS asio_shared DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})
 endif()
+
+if (WITH_EP)
+    install(
+        DIRECTORY "${EP_PG_STAGING_DIR}/"
+        DESTINATION "${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}"
+        FILES_MATCHING PATTERN "*.so"
+    )
+    install(FILES
+        "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/ep.py"
+        "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_ep_buffer.py"
+        "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/pg.py"
+        DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}
+    )
+    # ep.so / pg.so link against engine.so by that exact bare name.
+    # Create a engine.so -> engine<EXT_SUFFIX> symlink so they can find it.
+    install(CODE "
+        execute_process(COMMAND ${CMAKE_COMMAND} -E create_symlink
+            \"engine${PYTHON_EXT_SUFFIX}\"
+            \"${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}/engine.so\"
+        )
+    ")
+endif()

mooncake-integration/store/store_py.cpp

@@ -256,11 +256,12 @@ class MooncakeStorePyWrapper {
 
     // Helper to initialize real client and register it
     std::shared_ptr<RealClient> init_real_client() {
+        auto &resource_tracker = ResourceTracker::getInstance();
         auto real_client = RealClient::create();
         use_dummy_client_ = false;
         store_ = real_client;
-        ResourceTracker::getInstance().registerInstance(
-            std::dynamic_pointer_cast<PyClient>(store_));
+        resource_tracker.registerInstance(
+            std::static_pointer_cast<PyClient>(store_));
         return real_client;
     }
 
@@ -700,7 +701,8 @@ class MooncakeStorePyWrapper {
         std::vector<std::string> all_chunk_keys;
         py::list all_chunks_list;
         std::vector<size_t> processed_indices;
-        std::vector<int> final_results(base_keys.size(), 0);
+        std::vector<int> final_results(base_keys.size(),
+                                       to_py_ret(ErrorCode::INVALID_PARAMS));
 
         try {
             // Chunking phase (GIL Held)
@@ -738,13 +740,18 @@ class MooncakeStorePyWrapper {
             // Aggregate results
             for (size_t i = 0; i < processed_indices.size(); ++i) {
                 size_t original_idx = processed_indices[i];
+                bool all_ok = true;
                 for (int j = 0; j < tp_size; ++j) {
                     int res = chunk_results[i * tp_size + j];
                     if (res != 0) {
                         final_results[original_idx] = res;  // First error wins
+                        all_ok = false;
                         break;
                     }
                 }
+                if (all_ok) {
+                    final_results[original_idx] = 0;
+                }
             }
 
         } catch (const std::exception &e) {
@@ -774,6 +781,416 @@ class MooncakeStorePyWrapper {
                                              tp_size, split_dim);
     }
 
+    // Zero-copy put from pre-allocated buffer (layout: [TensorMetadata][data])
+    int put_tensor_from(const std::string &key, uintptr_t buffer_ptr,
+                        size_t size) {
+        if (buffer_ptr == 0) {
+            LOG(ERROR) << "Buffer pointer cannot be null";
+            return to_py_ret(ErrorCode::INVALID_PARAMS);
+        }
+        void *buffer = reinterpret_cast<void *>(buffer_ptr);
+        if (!is_client_initialized()) {
+            LOG(ERROR) << "Client is not initialized";
+            return to_py_ret(ErrorCode::INVALID_PARAMS);
+        }
+        if (use_dummy_client_) {
+            LOG(ERROR) << "put_tensor_from is not supported for dummy client";
+            return to_py_ret(ErrorCode::INVALID_PARAMS);
+        }
+        if (size <= sizeof(TensorMetadata)) {
+            LOG(ERROR) << "Buffer size too small for tensor metadata";
+            return to_py_ret(ErrorCode::INVALID_PARAMS);
+        }
+        py::gil_scoped_release release_gil;
+        return store_->put_from(key, buffer, size, ReplicateConfig{});
+    }
+
+    std::vector<int> batch_put_tensor_from(
+        const std::vector<std::string> &keys,
+        const std::vector<uintptr_t> &buffer_ptrs,
+        const std::vector<size_t> &sizes) {
+        if (!is_client_initialized()) {
+            LOG(ERROR) << "Client is not initialized";
+            return std::vector<int>(keys.size(),
+                                    to_py_ret(ErrorCode::INVALID_PARAMS));
+        }
+        if (use_dummy_client_) {
+            LOG(ERROR)
+                << "batch_put_tensor_from is not supported for dummy client";
+            return std::vector<int>(keys.size(),
+                                    to_py_ret(ErrorCode::INVALID_PARAMS));
+        }
+        if (keys.empty()) {
+            return std::vector<int>();
+        }
+        if (keys.size() != buffer_ptrs.size() || keys.size() != sizes.size()) {
+            LOG(ERROR) << "Size mismatch: keys, buffer_ptrs, and sizes must "
+                          "have the same length";
+            return std::vector<int>(keys.size(),
+                                    to_py_ret(ErrorCode::INVALID_PARAMS));
+        }
+        for (size_t i = 0; i < sizes.size(); ++i) {
+            if (buffer_ptrs[i] == 0) {
+                LOG(ERROR) << "Buffer pointer at index " << i
+                           << " cannot be null";
+                return std::vector<int>(keys.size(),
+                                        to_py_ret(ErrorCode::INVALID_PARAMS));
+            }
+            if (sizes[i] <= sizeof(TensorMetadata)) {
+                LOG(ERROR) << "Buffer size at index " << i
+                           << " too small for tensor metadata";
+                return std::vector<int>(keys.size(),
+                                        to_py_ret(ErrorCode::INVALID_PARAMS));
+            }
+        }
+        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_gil;
+        return store_->batch_put_from(keys, buffers, sizes, ReplicateConfig{});
+    }
+
+    int put_tensor_with_tp_from(const std::string &key, uintptr_t buffer_ptr,
+                                size_t size, int tp_rank = 0, int tp_size = 1,
+                                int split_dim = 0) {
+        if (!is_client_initialized()) {
+            LOG(ERROR) << "Client is not initialized";
+            return to_py_ret(ErrorCode::INVALID_PARAMS);
+        }
+        if (use_dummy_client_) {
+            LOG(ERROR)
+                << "put_tensor_with_tp_from is not supported for dummy client";
+            return to_py_ret(ErrorCode::INVALID_PARAMS);
+        }
+        if (buffer_ptr == 0) {
+            LOG(ERROR) << "Buffer pointer cannot be null";
+            return to_py_ret(ErrorCode::INVALID_PARAMS);
+        }
+        if (size <= sizeof(TensorMetadata)) {
+            LOG(ERROR) << "Buffer size too small for tensor metadata";
+            return to_py_ret(ErrorCode::INVALID_PARAMS);
+        }
+        if (tp_size <= 1) {
+            return put_tensor_from(key, buffer_ptr, size);
+        }
+
+        pybind11::object tensor =
+            buffer_to_tensor(NULL, reinterpret_cast<char *>(buffer_ptr),
+                             static_cast<int64_t>(size));
+        if (tensor.is_none()) {
+            LOG(ERROR) << "Failed to decode full tensor buffer for "
+                          "put_tensor_with_tp_from";
+            return to_py_ret(ErrorCode::INVALID_PARAMS);
+        }
+        return put_tensor_with_tp_impl(key, tensor, ReplicateConfig{}, tp_rank,
+                                       tp_size, split_dim);
+    }
+
+    std::vector<int> batch_put_tensor_with_tp_from(
+        const std::vector<std::string> &base_keys,
+        const std::vector<uintptr_t> &buffer_ptrs,
+        const std::vector<size_t> &sizes, int tp_rank = 0, int tp_size = 1,
+        int split_dim = 0) {
+        if (!is_client_initialized()) {
+            LOG(ERROR) << "Client is not initialized";
+            return std::vector<int>(base_keys.size(),
+                                    to_py_ret(ErrorCode::INVALID_PARAMS));
+        }
+        if (use_dummy_client_) {
+            LOG(ERROR) << "batch_put_tensor_with_tp_from is not supported for "
+                          "dummy client";
+            return std::vector<int>(base_keys.size(),
+                                    to_py_ret(ErrorCode::INVALID_PARAMS));
+        }
+        if (tp_size <= 1) {
+            return batch_put_tensor_from(base_keys, buffer_ptrs, sizes);
+        }
+        if (base_keys.size() != buffer_ptrs.size() ||
+            base_keys.size() != sizes.size() || base_keys.empty()) {
+            if (!base_keys.empty()) {
+                LOG(ERROR) << "Size mismatch: base_keys, buffer_ptrs, and "
+                              "sizes must have the same length";
+            }
+            return std::vector<int>(base_keys.size(),
+                                    to_py_ret(ErrorCode::INVALID_PARAMS));
+        }
+
+        py::list tensors_list;
+        std::vector<size_t> processed_indices;
+        std::vector<int> final_results(base_keys.size(), 0);
+
+        for (size_t i = 0; i < base_keys.size(); ++i) {
+            if (buffer_ptrs[i] == 0) {
+                LOG(ERROR) << "Buffer pointer at index " << i
+                           << " cannot be null";
+                final_results[i] = to_py_ret(ErrorCode::INVALID_PARAMS);
+                continue;
+            }
+            if (sizes[i] <= sizeof(TensorMetadata)) {
+                LOG(ERROR) << "Buffer size at index " << i
+                           << " too small for tensor metadata";
+                final_results[i] = to_py_ret(ErrorCode::INVALID_PARAMS);
+                continue;
+            }
+
+            py::object tensor =
+                buffer_to_tensor(NULL, reinterpret_cast<char *>(buffer_ptrs[i]),
+                                 static_cast<int64_t>(sizes[i]));
+            if (tensor.is_none()) {
+                LOG(ERROR) << "Failed to decode full tensor buffer at index "
+                           << i;
+                final_results[i] = to_py_ret(ErrorCode::INVALID_PARAMS);
+                continue;
+            }
+            tensors_list.append(tensor);
+            processed_indices.push_back(i);
+        }
+
+        if (processed_indices.empty()) {
+            return final_results;
+        }
+
+        std::vector<std::string> valid_keys;
+        valid_keys.reserve(processed_indices.size());
+        for (size_t idx : processed_indices) {
+            valid_keys.push_back(base_keys[idx]);
+        }
+
+        std::vector<int> op_results = batch_put_tensor_with_tp_impl(
+            valid_keys, tensors_list, ReplicateConfig{}, tp_rank, tp_size,
+            split_dim);
+        for (size_t i = 0; i < processed_indices.size(); ++i) {
+            final_results[processed_indices[i]] = op_results[i];
+        }
+        return final_results;
+    }
+
+    // --- Upsert tensor methods ---
+
+    int upsert_tensor_impl(const std::string &key, pybind11::object tensor,
+                           const ReplicateConfig &config) {
+        auto info = extract_tensor_info(tensor, key);
+        if (!info.valid()) return to_py_ret(ErrorCode::INVALID_PARAMS);
+
+        std::vector<std::span<const char>> values;
+        values.emplace_back(reinterpret_cast<const char *>(&info.metadata),
+                            sizeof(TensorMetadata));
+        values.emplace_back(reinterpret_cast<const char *>(info.data_ptr),
+                            info.tensor_size);
+
+        py::gil_scoped_release release_gil;
+        int ret = store_->upsert_parts(key, values, config);
+        if (ret != 0)
+            LOG(ERROR) << "upsert_parts failed for key " << key << " with code "
+                       << ret;
+        return ret;
+    }
+
+    int upsert_tensor(const std::string &key, pybind11::object tensor) {
+        if (!is_client_initialized() || use_dummy_client_) {
+            LOG(ERROR) << "Client not initialized or Dummy client not "
+                          "supported for tensors";
+            return to_py_ret(ErrorCode::INVALID_PARAMS);
+        }
+        return upsert_tensor_impl(key, tensor, ReplicateConfig{});
+    }
+
+    int upsert_tensor_from(const std::string &key, uintptr_t buffer_ptr,
+                           size_t size) {
+        if (buffer_ptr == 0) {
+            LOG(ERROR) << "Buffer pointer cannot be null";
+            return to_py_ret(ErrorCode::INVALID_PARAMS);
+        }
+        void *buffer = reinterpret_cast<void *>(buffer_ptr);
+        if (!is_client_initialized()) {
+            LOG(ERROR) << "Client is not initialized";
+            return to_py_ret(ErrorCode::INVALID_PARAMS);
+        }
+        if (use_dummy_client_) {
+            LOG(ERROR)
+                << "upsert_tensor_from is not supported for dummy client";
+            return to_py_ret(ErrorCode::INVALID_PARAMS);
+        }
+        if (size <= sizeof(TensorMetadata)) {
+            LOG(ERROR) << "Buffer size too small for tensor metadata";
+            return to_py_ret(ErrorCode::INVALID_PARAMS);
+        }
+        py::gil_scoped_release release_gil;
+        return store_->upsert_from(key, buffer, size, ReplicateConfig{});
+    }
+
+    std::vector<int> batch_upsert_tensor_from(
+        const std::vector<std::string> &keys,
+        const std::vector<uintptr_t> &buffer_ptrs,
+        const std::vector<size_t> &sizes) {
+        if (!is_client_initialized()) {
+            LOG(ERROR) << "Client is not initialized";
+            return std::vector<int>(keys.size(),
+                                    to_py_ret(ErrorCode::INVALID_PARAMS));
+        }
+        if (use_dummy_client_) {
+            LOG(ERROR)
+                << "batch_upsert_tensor_from is not supported for dummy client";
+            return std::vector<int>(keys.size(),
+                                    to_py_ret(ErrorCode::INVALID_PARAMS));
+        }
+        if (keys.empty()) {
+            return std::vector<int>();
+        }
+        if (keys.size() != buffer_ptrs.size() || keys.size() != sizes.size()) {
+            LOG(ERROR) << "Size mismatch: keys, buffer_ptrs, and sizes must "
+                          "have the same length";
+            return std::vector<int>(keys.size(),
+                                    to_py_ret(ErrorCode::INVALID_PARAMS));
+        }
+        for (size_t i = 0; i < sizes.size(); ++i) {
+            if (buffer_ptrs[i] == 0) {
+                LOG(ERROR) << "Buffer pointer at index " << i
+                           << " cannot be null";
+                return std::vector<int>(keys.size(),
+                                        to_py_ret(ErrorCode::INVALID_PARAMS));
+            }
+            if (sizes[i] <= sizeof(TensorMetadata)) {
+                LOG(ERROR) << "Buffer size at index " << i
+                           << " too small for tensor metadata";
+                return std::vector<int>(keys.size(),
+                                        to_py_ret(ErrorCode::INVALID_PARAMS));
+            }
+        }
+        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_gil;
+        return store_->batch_upsert_from(keys, buffers, sizes,
+                                         ReplicateConfig{});
+    }
+
+    std::vector<int> batch_upsert_tensor_impl(
+        const std::vector<std::string> &keys,
+        const pybind11::list &tensors_list,
+        const ReplicateConfig &config = ReplicateConfig{}) {
+        std::vector<PyTensorInfo> infos(keys.size());
+        std::vector<int> results(keys.size(), 0);
+
+        // 1. Extract Metadata (GIL Held)
+        for (size_t i = 0; i < keys.size(); ++i) {
+            infos[i] = extract_tensor_info(tensors_list[i], keys[i]);
+            if (!infos[i].valid())
+                results[i] = to_py_ret(ErrorCode::INVALID_PARAMS);
+        }
+
+        // 2. Prepare Buffers and Execute (GIL Released)
+        {
+            py::gil_scoped_release release_gil;
+
+            std::vector<std::string> valid_keys;
+            std::vector<void *> buffer_ptrs;
+            std::vector<size_t> buffer_sizes;
+            std::vector<size_t> original_indices;
+
+            std::vector<std::unique_ptr<BufferHandle>> temp_allocations;
+
+            for (size_t i = 0; i < infos.size(); ++i) {
+                if (!infos[i].valid()) continue;
+
+                size_t total_size =
+                    sizeof(TensorMetadata) + infos[i].tensor_size;
+                auto alloc_result =
+                    store_->client_buffer_allocator_->allocate(total_size);
+
+                if (!alloc_result) {
+                    LOG(ERROR)
+                        << "Failed to allocate buffer for key: " << keys[i];
+                    results[i] = to_py_ret(ErrorCode::INVALID_PARAMS);
+                    continue;
+                }
+
+                // Copy Metadata & Data
+                char *dst = static_cast<char *>(alloc_result->ptr());
+                memcpy(dst, &infos[i].metadata, sizeof(TensorMetadata));
+                memcpy(dst + sizeof(TensorMetadata),
+                       reinterpret_cast<void *>(infos[i].data_ptr),
+                       infos[i].tensor_size);
+
+                valid_keys.push_back(keys[i]);
+                buffer_ptrs.push_back(alloc_result->ptr());
+                buffer_sizes.push_back(total_size);
+                original_indices.push_back(i);
+
+                temp_allocations.push_back(
+                    std::make_unique<BufferHandle>(std::move(*alloc_result)));
+            }
+
+            if (!valid_keys.empty()) {
+                std::vector<int> op_results = store_->batch_upsert_from(
+                    valid_keys, buffer_ptrs, buffer_sizes, config);
+                for (size_t i = 0; i < op_results.size(); ++i) {
+                    results[original_indices[i]] = op_results[i];
+                }
+            }
+        }
+
+        return results;
+    }
+
+    std::vector<int> batch_upsert_tensor(const std::vector<std::string> &keys,
+                                         const pybind11::list &tensors_list) {
+        if (!is_client_initialized() || use_dummy_client_)
+            return std::vector<int>(keys.size(),
+                                    to_py_ret(ErrorCode::INVALID_PARAMS));
+
+        if (keys.size() != tensors_list.size() || keys.empty()) {
+            if (!keys.empty()) LOG(ERROR) << "Size mismatch in batch_upsert";
+            return std::vector<int>(keys.size(),
+                                    to_py_ret(ErrorCode::INVALID_PARAMS));
+        }
+
+        return batch_upsert_tensor_impl(keys, tensors_list, ReplicateConfig{});
+    }
+
+    int upsert_pub_tensor(const std::string &key, pybind11::object tensor,
+                          const ReplicateConfig &config = ReplicateConfig{}) {
+        if (!is_client_initialized() || use_dummy_client_) {
+            LOG(ERROR) << "Client not initialized or Dummy client not "
+                          "supported for tensors";
+            return to_py_ret(ErrorCode::INVALID_PARAMS);
+        }
+
+        int validate_result = validate_replicate_config(config);
+        if (validate_result) return validate_result;
+
+        return upsert_tensor_impl(key, tensor, config);
+    }
+
+    std::vector<int> batch_upsert_pub_tensor(
+        const std::vector<std::string> &keys,
+        const pybind11::list &tensors_list,
+        const ReplicateConfig &config = ReplicateConfig{}) {
+        if (!is_client_initialized() || use_dummy_client_)
+            return std::vector<int>(keys.size(),
+                                    to_py_ret(ErrorCode::INVALID_PARAMS));
+
+        if (keys.size() != tensors_list.size() || keys.empty()) {
+            if (!keys.empty())
+                LOG(ERROR) << "Size mismatch in batch_upsert_pub";
+            return std::vector<int>(keys.size(),
+                                    to_py_ret(ErrorCode::INVALID_PARAMS));
+        }
+
+        int validate_result = validate_replicate_config(config);
+        if (validate_result)
+            return std::vector<int>(keys.size(),
+                                    to_py_ret(ErrorCode::INVALID_PARAMS));
+
+        return batch_upsert_tensor_impl(keys, tensors_list, config);
+    }
+
+    // --- End Upsert tensor methods ---
+
     int validate_replicate_config(
         const ReplicateConfig &config = ReplicateConfig{}) {
         if (!config.preferred_segments.empty() &&
@@ -980,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)
@@ -1144,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;
@@ -1155,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) {
@@ -1188,15 +1611,19 @@ 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,
                size_t local_buffer_size, const std::string &server_address) {
+                auto &resource_tracker = ResourceTracker::getInstance();
                 self.use_dummy_client_ = true;
                 self.store_ = std::make_shared<DummyClient>();
-                ResourceTracker::getInstance().registerInstance(
-                    std::dynamic_pointer_cast<PyClient>(self.store_));
+                resource_tracker.registerInstance(
+                    std::static_pointer_cast<PyClient>(self.store_));
                 auto [ip, port] = parseHostNameWithPort(server_address);
                 return self.store_->setup_dummy(
                     mem_pool_size, local_buffer_size, server_address,
@@ -1263,6 +1690,16 @@ PYBIND11_MODULE(store, m) {
             py::arg("force") = false,
             "Remove all objects from the store. If force=True, skip lease "
             "and replication task checks.")
+        .def(
+            "batch_remove",
+            [](MooncakeStorePyWrapper &self,
+               const std::vector<std::string> &keys, bool force) {
+                py::gil_scoped_release release;
+                return self.store_->batchRemove(keys, force);
+            },
+            py::arg("keys"), py::arg("force") = false,
+            "Batch remove objects by keys. Returns a list of status codes "
+            "(0=success, negative=error code) for each key.")
         .def("is_exist",
              [](MooncakeStorePyWrapper &self, const std::string &key) {
                  py::gil_scoped_release release;
@@ -1405,6 +1842,177 @@ PYBIND11_MODULE(store, m) {
             py::arg("tp_rank") = 0, py::arg("tp_size") = 1,
             "Get a batch of PyTorch tensor shards from the store directly into "
             "pre-allocated buffers for a given Tensor Parallel rank.")
+        .def("put_tensor_from", &MooncakeStorePyWrapper::put_tensor_from,
+             py::arg("key"), py::arg("buffer_ptr"), py::arg("size"),
+             "Put a tensor directly from a pre-allocated buffer. Buffer layout "
+             "must be [TensorMetadata][tensor data], same as get_tensor_into.")
+        .def("batch_put_tensor_from",
+             &MooncakeStorePyWrapper::batch_put_tensor_from, py::arg("keys"),
+             py::arg("buffer_ptrs"), py::arg("sizes"),
+             "Put tensors directly from pre-allocated buffers for multiple "
+             "keys. Each buffer layout: [TensorMetadata][tensor data].")
+        .def("put_tensor_with_tp_from",
+             &MooncakeStorePyWrapper::put_tensor_with_tp_from, py::arg("key"),
+             py::arg("buffer_ptr"), py::arg("size"), py::arg("tp_rank") = 0,
+             py::arg("tp_size") = 1, py::arg("split_dim") = 0,
+             "Put a full tensor directly from a pre-allocated buffer for "
+             "Tensor Parallelism. The buffer is split internally and stored "
+             "under key_tp_<rank> for all ranks.")
+        .def("batch_put_tensor_with_tp_from",
+             &MooncakeStorePyWrapper::batch_put_tensor_with_tp_from,
+             py::arg("base_keys"), py::arg("buffer_ptrs"), py::arg("sizes"),
+             py::arg("tp_rank") = 0, py::arg("tp_size") = 1,
+             py::arg("split_dim") = 0,
+             "Put a batch of full tensors directly from pre-allocated "
+             "buffers for Tensor Parallelism. Each buffer is internally split "
+             "and stored under key_tp_<rank> for all ranks.")
+        .def("upsert_tensor", &MooncakeStorePyWrapper::upsert_tensor,
+             py::arg("key"), py::arg("tensor"),
+             "Upsert a PyTorch tensor into the store (insert or update)")
+        .def("upsert_tensor_from", &MooncakeStorePyWrapper::upsert_tensor_from,
+             py::arg("key"), py::arg("buffer_ptr"), py::arg("size"),
+             "Upsert a tensor directly from a pre-allocated buffer. Buffer "
+             "layout must be [TensorMetadata][tensor data].")
+        .def("batch_upsert_tensor_from",
+             &MooncakeStorePyWrapper::batch_upsert_tensor_from, py::arg("keys"),
+             py::arg("buffer_ptrs"), py::arg("sizes"),
+             "Upsert tensors directly from pre-allocated buffers for "
+             "multiple keys. Each buffer layout: [TensorMetadata][tensor "
+             "data].")
+        .def("batch_upsert_tensor",
+             &MooncakeStorePyWrapper::batch_upsert_tensor, py::arg("keys"),
+             py::arg("tensors_list"),
+             "Upsert a batch of PyTorch tensors into the store (insert or "
+             "update)")
+        .def("upsert_pub_tensor", &MooncakeStorePyWrapper::upsert_pub_tensor,
+             py::arg("key"), py::arg("tensor"),
+             py::arg("config") = ReplicateConfig{},
+             "Upsert a PyTorch tensor with configurable replication settings")
+        .def("batch_upsert_pub_tensor",
+             &MooncakeStorePyWrapper::batch_upsert_pub_tensor, py::arg("keys"),
+             py::arg("tensors_list"), py::arg("config") = ReplicateConfig{},
+             "Batch upsert PyTorch tensors with configurable replication "
+             "settings")
+        .def(
+            "upsert_from",
+            [](MooncakeStorePyWrapper &self, const std::string &key,
+               uintptr_t buffer_ptr, size_t size,
+               const ReplicateConfig &config = ReplicateConfig{}) {
+                if (!self.is_client_initialized()) {
+                    LOG(ERROR) << "Client is not initialized";
+                    return to_py_ret(ErrorCode::INVALID_PARAMS);
+                }
+                void *buffer = reinterpret_cast<void *>(buffer_ptr);
+                py::gil_scoped_release release;
+                return self.store_->upsert_from(key, buffer, size, config);
+            },
+            py::arg("key"), py::arg("buffer_ptr"), py::arg("size"),
+            py::arg("config") = ReplicateConfig{},
+            "Upsert object data directly from a pre-allocated buffer")
+        .def(
+            "batch_upsert_from",
+            [](MooncakeStorePyWrapper &self,
+               const std::vector<std::string> &keys,
+               const std::vector<uintptr_t> &buffer_ptrs,
+               const std::vector<size_t> &sizes,
+               const ReplicateConfig &config = ReplicateConfig{}) {
+                if (!self.is_client_initialized()) {
+                    LOG(ERROR) << "Client is not initialized";
+                    return std::vector<int>(
+                        keys.size(), to_py_ret(ErrorCode::INVALID_PARAMS));
+                }
+                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_->batch_upsert_from(keys, buffers, sizes,
+                                                      config);
+            },
+            py::arg("keys"), py::arg("buffer_ptrs"), py::arg("sizes"),
+            py::arg("config") = ReplicateConfig{},
+            "Upsert object data directly from pre-allocated buffers for "
+            "multiple keys")
+        .def(
+            "upsert",
+            [](MooncakeStorePyWrapper &self, const std::string &key,
+               py::buffer buf,
+               const ReplicateConfig &config = ReplicateConfig{}) {
+                if (!self.is_client_initialized()) {
+                    LOG(ERROR) << "Client is not initialized";
+                    return to_py_ret(ErrorCode::INVALID_PARAMS);
+                }
+                py::buffer_info info = buf.request(/*writable=*/false);
+                py::gil_scoped_release release;
+                return self.store_->upsert(
+                    key,
+                    std::span<const char>(static_cast<char *>(info.ptr),
+                                          static_cast<size_t>(info.size)),
+                    config);
+            },
+            py::arg("key"), py::arg("value"),
+            py::arg("config") = ReplicateConfig{},
+            "Upsert raw bytes into the store (insert or update)")
+        .def(
+            "upsert_parts",
+            [](MooncakeStorePyWrapper &self, const std::string &key,
+               py::args parts,
+               const ReplicateConfig &config = ReplicateConfig{}) {
+                if (!self.is_client_initialized()) {
+                    LOG(ERROR) << "Client is not initialized";
+                    return to_py_ret(ErrorCode::INVALID_PARAMS);
+                }
+                std::vector<py::buffer_info> infos;
+                std::vector<std::span<const char>> spans;
+                infos.reserve(parts.size());
+                spans.reserve(parts.size());
+
+                for (auto &obj : parts) {
+                    py::buffer buf = py::reinterpret_borrow<py::buffer>(obj);
+                    infos.emplace_back(buf.request(false));
+                    const auto &info = infos.back();
+                    if (info.ndim != 1 || info.itemsize != 1)
+                        throw std::runtime_error(
+                            "parts must be 1-D bytes-like");
+
+                    spans.emplace_back(static_cast<const char *>(info.ptr),
+                                       static_cast<size_t>(info.size));
+                }
+
+                py::gil_scoped_release unlock;
+                return self.store_->upsert_parts(key, spans, config);
+            },
+            py::arg("key"), py::arg("config") = ReplicateConfig{},
+            "Upsert multiple byte parts as a single object (insert or update)")
+        .def(
+            "upsert_batch",
+            [](MooncakeStorePyWrapper &self,
+               const std::vector<std::string> &keys,
+               const std::vector<py::buffer> &buffers,
+               const ReplicateConfig &config = ReplicateConfig{}) {
+                if (!self.is_client_initialized()) {
+                    LOG(ERROR) << "Client is not initialized";
+                    return to_py_ret(ErrorCode::INVALID_PARAMS);
+                }
+                std::vector<py::buffer_info> infos;
+                std::vector<std::span<const char>> spans;
+                infos.reserve(buffers.size());
+                spans.reserve(buffers.size());
+
+                for (const auto &buf : buffers) {
+                    infos.emplace_back(buf.request(/*writable=*/false));
+                    const auto &info = infos.back();
+                    spans.emplace_back(static_cast<const char *>(info.ptr),
+                                       static_cast<size_t>(info.size));
+                }
+
+                py::gil_scoped_release release;
+                return self.store_->upsert_batch(keys, spans, config);
+            },
+            py::arg("keys"), py::arg("values"),
+            py::arg("config") = ReplicateConfig{},
+            "Batch upsert raw bytes for multiple keys (insert or update)")
         .def(
             "register_buffer",
             [](MooncakeStorePyWrapper &self, uintptr_t buffer_ptr,
@@ -1434,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,
@@ -1466,14 +2094,12 @@ PYBIND11_MODULE(store, m) {
             [](MooncakeStorePyWrapper &self, const std::string &key,
                uintptr_t buffer_ptr, size_t size,
                const ReplicateConfig &config = ReplicateConfig{}) {
-                // Put data directly from user-provided buffer
+                if (!self.is_client_initialized()) {
+                    LOG(ERROR) << "Client is not initialized";
+                    return to_py_ret(ErrorCode::INVALID_PARAMS);
+                }
                 void *buffer = reinterpret_cast<void *>(buffer_ptr);
                 py::gil_scoped_release release;
-                if (self.use_dummy_client_) {
-                    LOG(ERROR) << "put_from is not supported for dummy client "
-                                  "now";
-                    return -1;
-                }
                 return self.store_->put_from(key, buffer, size, config);
             },
             py::arg("key"), py::arg("buffer_ptr"), py::arg("size"),
@@ -1485,18 +2111,14 @@ PYBIND11_MODULE(store, m) {
                uintptr_t buffer_ptr, uintptr_t metadata_buffer_ptr, size_t size,
                size_t metadata_size,
                const ReplicateConfig &config = ReplicateConfig{}) {
-                // Put data directly from user-provided buffer with
-                // metadata
+                if (!self.is_client_initialized()) {
+                    LOG(ERROR) << "Client is not initialized";
+                    return to_py_ret(ErrorCode::INVALID_PARAMS);
+                }
                 void *buffer = reinterpret_cast<void *>(buffer_ptr);
                 void *metadata_buffer =
                     reinterpret_cast<void *>(metadata_buffer_ptr);
                 py::gil_scoped_release release;
-                if (self.use_dummy_client_) {
-                    LOG(ERROR)
-                        << "put_from_with_metadata is not supported for dummy "
-                           "client now";
-                    return -1;
-                }
                 return self.store_->put_from_with_metadata(
                     key, buffer, metadata_buffer, size, metadata_size, config);
             },
@@ -1604,13 +2226,11 @@ PYBIND11_MODULE(store, m) {
                const std::vector<std::vector<uintptr_t>> &all_buffer_ptrs,
                const std::vector<std::vector<size_t>> &all_sizes,
                const ReplicateConfig &config = ReplicateConfig{}) {
-                py::gil_scoped_release release;
-                if (self.use_dummy_client_) {
-                    LOG(ERROR)
-                        << "batch_put_from_multi_buffers is not supported for "
-                           "dummy client now";
+                if (!self.is_client_initialized()) {
+                    LOG(ERROR) << "Client is not initialized";
                     return std::vector<int>{};
                 }
+                py::gil_scoped_release release;
                 return self.store_->batch_put_from_multi_buffers(
                     keys, CastAddrs2Ptrs(all_buffer_ptrs), all_sizes, config);
             },
@@ -1627,12 +2247,6 @@ PYBIND11_MODULE(store, m) {
                const std::vector<std::vector<size_t>> &all_sizes,
                bool prefer_alloc_in_same_node = false) {
                 py::gil_scoped_release release;
-                if (self.use_dummy_client_) {
-                    LOG(ERROR)
-                        << "batch_get_into_multi_buffers is not supported for "
-                           "dummy client now";
-                    return std::vector<int>{};
-                }
                 return self.store_->batch_get_into_multi_buffers(
                     keys, CastAddrs2Ptrs(all_buffer_ptrs), all_sizes,
                     prefer_alloc_in_same_node);
@@ -1657,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,9 @@ 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());
             if (status.s == TransferStatusEnum::COMPLETED) {
@@ -637,14 +658,14 @@ int TransferEnginePy::getBatchTransferStatus(
                 LOG(INFO) << "Sync data transfer timeout";
             }
             auto current_ts = getCurrentTimeInNano();
-            if (current_ts - batch_desc->start_timestamp > entry.second) {
+            if (current_ts - start_timestamp > entry.second) {
                 LOG(INFO) << "Sync batch data transfer timeout after "
-                          << current_ts - batch_desc->start_timestamp << "ns";
+                          << current_ts - start_timestamp << "ns";
                 failed_or_timeout = true;
             }
         }
 
-        for (auto &remove_id : remove_ids) {
+        for (auto& remove_id : remove_ids) {
             timeout_table.erase(remove_id);
         }
 
@@ -652,7 +673,7 @@ int TransferEnginePy::getBatchTransferStatus(
     }
 
     if (failed_or_timeout) {
-        for (auto &entry : timeout_table) {
+        for (auto& entry : timeout_table) {
             engine_->freeBatchID(entry.first);
         }
     }
@@ -660,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) {
@@ -681,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;
 
@@ -716,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);
 }
@@ -725,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);
 }
 
@@ -766,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()) {
@@ -827,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;
@@ -862,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);
@@ -870,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);
@@ -887,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,
@@ -899,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) {
@@ -911,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);
 }
@@ -922,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);
@@ -941,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);
@@ -964,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});
     }
@@ -975,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
@@ -995,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,11 +33,15 @@ 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
+    EXT_LDFLAGS+=" -L/usr/local/cuda/lib64/stubs"
+fi
+
 if [ -d "/usr/local/cuda/lib64" ]; then
-    EXT_LDFLAGS+=" -L/usr/local/cuda/lib64 -lcudart"
+    EXT_LDFLAGS+=" -L/usr/local/cuda/lib64 -lcuda -lcudart"
 fi
 
 if [ -d "/opt/rocm/lib" ]; then

mooncake-pg/BuildPgExt.cmake

@@ -0,0 +1,94 @@
+# BuildPgExt.cmake - Build the Mooncake PG Python extension.
+#
+# Invoked at build time via cmake -P from the root CMakeLists.txt when
+# WITH_EP=ON.  Variables are passed with -D from the custom target:
+#
+#   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
+#   STAGING_DIR         - destination directory for the built .so files
+#   ENGINE_SO_PATH      - absolute path to the built engine.cpython-XYZ.so
+
+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}")
+endif()
+if(TORCH_CUDA_ARCH_LIST)
+  string(REPLACE "|" ";" TORCH_CUDA_ARCH_LIST "${TORCH_CUDA_ARCH_LIST}")
+endif()
+
+# ---------------------------------------------------------------------------
+# 1. Set up the build environment.
+# ---------------------------------------------------------------------------
+# Clear jobserver variables so that sub-processes started by setup.py do not
+# try to connect to the parent ninja's jobserver pipe FDs, which are not
+# inherited and cause: "ninja: error: Could not initialize jobserver: Invalid
+# file descriptors".
+set(ENV{MAKEFLAGS} "")
+set(ENV{MFLAGS} "")
+set(ENV{TORCH_CUDA_ARCH_LIST} "${TORCH_CUDA_ARCH_LIST}")
+
+# ---------------------------------------------------------------------------
+# 2. Ensure engine.so exists in mooncake-wheel/mooncake/ for setup.py linking.
+# ---------------------------------------------------------------------------
+# setup.py links against -l:engine.so in ../mooncake-wheel/mooncake/.
+# During the make phase only the versioned engine.cpython-XYZ.so exists in
+# the build tree; create a bare engine.so symlink so the linker can find it.
+set(_wheel_mooncake_dir "${SOURCE_DIR}/../mooncake-wheel/mooncake")
+set(_engine_symlink "${_wheel_mooncake_dir}/engine.so")
+if(ENGINE_SO_PATH AND NOT EXISTS "${_engine_symlink}")
+  message(STATUS "[PG] Creating engine.so symlink -> ${ENGINE_SO_PATH}")
+  execute_process(
+    COMMAND ${CMAKE_COMMAND} -E create_symlink "${ENGINE_SO_PATH}" "${_engine_symlink}"
+  )
+endif()
+
+# ---------------------------------------------------------------------------
+# 3. Build the PG Python extension.
+# ---------------------------------------------------------------------------
+if("${EP_TORCH_VERSIONS}" STREQUAL "")
+  message(STATUS "[PG] Building with currently-installed PyTorch")
+  execute_process(
+    COMMAND ${Python3_EXECUTABLE} setup.py build_ext --build-lib .
+    WORKING_DIRECTORY "${SOURCE_DIR}"
+    RESULT_VARIABLE _ret
+  )
+  if(NOT _ret EQUAL 0)
+    message(FATAL_ERROR "[PG] Extension build failed (exit code: ${_ret})")
+  endif()
+else()
+  message(STATUS "[PG] Building for PyTorch versions: ${EP_TORCH_VERSIONS}")
+  foreach(_version IN LISTS EP_TORCH_VERSIONS)
+    install_pytorch_wheel("${_version}" "${EP_CUDA_MAJOR}" "${EP_CUDA_MINOR}" "[PG]")
+
+    execute_process(
+      COMMAND ${Python3_EXECUTABLE} setup.py build_ext --build-lib . --force
+      WORKING_DIRECTORY "${SOURCE_DIR}"
+      RESULT_VARIABLE _ret
+    )
+    if(NOT _ret EQUAL 0)
+      message(FATAL_ERROR "[PG] Extension build failed for PyTorch ${_version}")
+    endif()
+  endforeach()
+endif()
+
+# ---------------------------------------------------------------------------
+# 4. Copy the built .so files to the staging directory.
+# ---------------------------------------------------------------------------
+file(MAKE_DIRECTORY "${STAGING_DIR}")
+file(GLOB _so_files "${SOURCE_DIR}/mooncake/*.so")
+foreach(_so IN LISTS _so_files)
+  get_filename_component(_fname "${_so}" NAME)
+  message(STATUS "[PG] Staging ${_fname} -> ${STAGING_DIR}")
+  file(COPY "${_so}" DESTINATION "${STAGING_DIR}" NO_SOURCE_PERMISSIONS)
+endforeach()
+
+message(STATUS "[PG] Mooncake PG extension build complete")

mooncake-pg/CMakeLists.txt

@@ -0,0 +1,33 @@
+cmake_minimum_required(VERSION 3.16)
+project(mooncake-pg)
+
+# Find PyTorch's CMake prefix path
+execute_process(
+    COMMAND ${PYTHON_EXECUTABLE} -c "import torch; print(torch.utils.cmake_prefix_path)"
+    OUTPUT_VARIABLE PYTORCH_CMAKE_PATH
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+)
+if(NOT PYTORCH_CMAKE_PATH)
+  message(WARNING "Could not find PyTorch CMake path! Please set Torch_DIR.")
+else ()
+  message(STATUS "Found PyTorch CMake path: ${PYTORCH_CMAKE_PATH}")
+  list(APPEND CMAKE_PREFIX_PATH "${PYTORCH_CMAKE_PATH}/Torch")
+endif()
+
+set(TORCH_CUDA_ARCH_LIST "8.0;9.0")
+
+find_package(CUDAToolkit REQUIRED)
+# https://discuss.pytorch.org/t/failed-to-find-nvtoolsext/179635/13
+if(NOT TARGET CUDA::nvToolsExt AND TARGET CUDA::nvtx3)
+  add_library(CUDA::nvToolsExt INTERFACE IMPORTED)
+  target_compile_definitions(
+      CUDA::nvToolsExt INTERFACE
+      TORCH_CUDA_USE_NVTX3
+  )
+  target_link_libraries(CUDA::nvToolsExt INTERFACE CUDA::nvtx3)
+endif()
+find_package(Torch REQUIRED)
+include_directories(${TORCH_INCLUDE_DIRS})
+
+include_directories(include)
+add_subdirectory(src)

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,28 +34,36 @@ 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_;
     int rank_;
 
-    // TODO: make it atomic and add `expandSize` to handle runtime scaling-up?
-    int size_;
+    std::atomic<int> groupSize_;
+
+    bool isDummy_;
+
+    // A mark tracking the group size for which all ranks
+    // in [0, establishedGroupSize_) have been successfully
+    // connected at least once (they may disconnect afterwards).
+    // Mainly used in `waitUntilNewRanksConnected()`.
+    std::atomic<int> establishedGroupSize_;
+
     uint64_t* local2global_rank_map_;
     c10::intrusive_ptr<::c10d::Store> store_;
 
@@ -68,16 +76,22 @@ class ConnectionContext {
 
     PeerConnection peerStates_[kMaxNumRanks];
 
+    // On MNNVL, warmup is skipped because CPU heap buffers aren't
+    // fabric-accessible for cross-node NVLink writes.
+    bool skip_warmup_;
+
     // warmup_send_region_ and warmup_recv_region_ are managed by
-    // ConnectionContext.
+    // ConnectionContext. nullptr when skip_warmup_ is true.
     int32_t* warmup_send_region_;
     int32_t* warmup_recv_region_;
 
     std::mutex backend_wakeup_mutex_;
     std::condition_variable backend_wakeup_cv_;
 
+    bool resource_abandoned_{false};
+
    public:
-    ConnectionContext(int backendIndex, int rank, int size,
+    ConnectionContext(int backendIndex, int rank, int size, bool isDummy,
                       uint64_t* local2global_rank_map,
                       c10::intrusive_ptr<::c10d::Store> store,
                       std::shared_ptr<TransferGroupMeta> meta,
@@ -88,14 +102,80 @@ class ConnectionContext {
     int32_t* warmup_send_region() const { return warmup_send_region_; }
     int32_t* warmup_recv_region() const { return warmup_recv_region_; }
 
-    int getTotalConnectedPeers() const {
-        return totalConnectedPeers_.load(std::memory_order_acquire);
-    }
-    bool isAllPeerConnected() const { return totalConnectedPeers_ == size_; }
+    /**
+     * @brief Get the total number of actively connected peers.
+     * @return The count of peers currently in the CONNECTED state.
+     */
+    int getTotalConnectedPeers() const;
 
+    /**
+     * @brief Expands the group to a new size.
+     *
+     * @note This is a non-blocking operation. Callers must invoke
+     *       `waitUntilNewRanksConnected()` prior to initiating any
+     *       subsequent communications (e.g., send, recv, putTaskCpu,
+     *       putTaskCuda) to ensure the new peers are ready.
+     *
+     * @param newGroupSize The target size for the extended group.
+     */
+    void extendGroupSizeTo(int newGroupSize);
+
+    /**
+     * @brief Checks whether all peers within the group have
+     *        established connections.
+     *
+     * @return True if all peers are fully connected.
+     */
+    bool isAllPeerConnected() const;
+
+    /**
+     * @brief Blocks until all peers in the group are connected.
+     *
+     * This method is primarily used during backend initialization.
+     * Upon completion, it set `establishedGroupSize_` to the
+     * current `groupSize_`.
+     */
     void waitUntilAllConnected();
+
+    void bootstrapLocalPeer(const std::string& localServerName,
+                            const SegmentInfo& localRankInfo);
+
+    /**
+     * @brief Blocks until all newly added ranks in the
+     *        extended group are connected.
+     *
+     * Specifically, it waits for pending ranks in the range
+     * `[establishedGroupSize_, groupSize_)` to reach the connected state.
+     * This should be called before starting new communications if
+     * `extendGroupSizeTo()` has been invoked.
+     * Upon completion, it set `establishedGroupSize_` to the
+     * current `groupSize_`.
+     */
+    void waitUntilNewRanksConnected();
+
     void shutdown();
 
+    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);
@@ -109,19 +189,25 @@ class ConnectionContext {
         return "extension_task_count_" + std::to_string(backendIndex) + "_" +
                std::to_string(rank);
     }
+    static std::string getExtensionActiveRanksStoreKey(int backendIndex,
+                                                       int rank) {
+        return "extension_active_ranks_" + std::to_string(backendIndex) + "_" +
+               std::to_string(rank);
+    }
 
    private:
     // 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() {
@@ -142,6 +228,7 @@ class ConnectionPoller {
 
    private:
     ConnectionPoller();
+    void ensureThreadStarted();
     void pollerLoop();
     bool processContext(const std::shared_ptr<ConnectionContext>& ctx);
     bool processPeer(const std::shared_ptr<ConnectionContext>& ctx,
@@ -150,6 +237,7 @@ class ConnectionPoller {
     std::mutex wakeup_mutex_;
     std::condition_variable wakeup_cv_;
     std::thread pollerThread_;
+    std::atomic<bool> pollerThreadStarted_{false};
 
     std::mutex contexts_mutex_;
     std::atomic<uint64_t> contexts_version_{0};

mooncake-pg/include/mooncake_backend.h

@@ -3,6 +3,8 @@
 
 #include <cstdint>
 #include <memory>
+#include <string>
+#include <vector>
 #include <mooncake_worker.cuh>
 #include <connection_poller.h>
 #include <p2p_proxy.h>
@@ -29,7 +31,18 @@ class MooncakeBackend final : public ::c10d::Backend {
         bool isExtension_ = false;
     };
 
-    MooncakeBackend(c10::intrusive_ptr<::c10d::Store> store, int rank, int size,
+    /**
+     * @brief Construct a Mooncake process-group backend instance.
+     *
+     * `distBackendOpts` contains the PyTorch process-group information for this
+     * backend instance. `options` contains Mooncake-specific settings and may
+     * be null when callers omit `pg_options`.
+     *
+     * @param distBackendOpts Process-group information supplied by PyTorch.
+     * @param options *Optional* Mooncake-specific backend options.
+     * @param isCpu Whether to initialize the CPU backend variant.
+     */
+    MooncakeBackend(c10d::DistributedBackendOptions distBackendOpts,
                     c10::intrusive_ptr<MooncakeBackendOptions> options,
                     bool isCpu = false);
 
@@ -37,6 +50,21 @@ class MooncakeBackend final : public ::c10d::Backend {
 
     const std::string getBackendName() const override;
 
+    /**
+     * @brief Return the stored Mooncake-specific backend options.
+     *
+     * PyTorch can use this to read Mooncake-specific options from an existing
+     * process group. This is used, for example, create sub-groups that inherit
+     * settings from the parent group.
+     *
+     * @return The stored backend options, or null when the backend was created
+     * without explicit Mooncake options.
+     */
+    c10::intrusive_ptr<::c10d::Backend::Options> getBackendOptions() override {
+        return c10::static_intrusive_pointer_cast<::c10d::Backend::Options>(
+            options_);
+    }
+
     // Point-to-point send/recv for torch.distributed P2POp/batch_isend_irecv.
     // Only single-tensor ops are supported.
     c10::intrusive_ptr<c10d::Work> send(std::vector<at::Tensor>& tensors,
@@ -97,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];
     }
@@ -122,11 +166,19 @@ class MooncakeBackend final : public ::c10d::Backend {
 
     void recoverRanks(const std::vector<int>& ranks);
 
+    void joinGroup();
+
    private:
+    void waitForExtensionState();
+    void publishLocalPeerMetadata();
+    void setLocalOnlyActiveRanks();
+    void syncActiveRanksTensor();
+
     static TransferEngine* engine_;
-    static MooncakeWorker* worker_;
+    std::shared_ptr<MooncakeWorker> worker_;
     static bool engineInitialized_;
     static int backendIndex_;
+    const c10::intrusive_ptr<MooncakeBackendOptions> options_;
     bool isCpu_{false};
     static std::string hostIp_;
     void* send_buffer_[2];
@@ -137,6 +189,7 @@ class MooncakeBackend final : public ::c10d::Backend {
     std::shared_ptr<TransferGroupMeta> meta_;
     bool isShutdown_{false};
     uint64_t local2global_rank_map_[kMaxNumRanks];
+    std::string localServerName_;
 
     // P2P async infrastructure
     // p2p_proxy_ is created in MooncakeBackend, but can live longer than
@@ -152,6 +205,7 @@ class MooncakeBackend final : public ::c10d::Backend {
     // Similar to p2p_proxy_, connection_ctx_ is created in MooncakeBackend, but
     // can live longer than MooncakeBackend.
     std::shared_ptr<ConnectionContext> connection_ctx_;
+    bool connectionPollerRegistered_{false};
 };
 
 }  // namespace mooncake

mooncake-pg/include/mooncake_worker.cuh

@@ -10,6 +10,13 @@
 #include <torch/csrc/distributed/c10d/Store.hpp>
 #include <transfer_engine.h>
 
+#include <memory>
+#include <atomic>
+#include <mutex>
+#include <thread>
+#include <unordered_map>
+#include <vector>
+
 namespace mooncake {
 
 static constexpr size_t kBufferSize = 1u << 24;
@@ -46,6 +53,7 @@ __global__ struct Task {
     size_t tensorSize;  // In bytes
     int64_t broadcastRoot;
     int bufferOffset;
+    uint64_t submitSequence = 0;
     BatchID batchID;
     void* transferGroupMeta;
 };
@@ -56,14 +64,24 @@ void launchReduceKernel(at::Tensor dst, size_t pos, size_t realSize, void* src,
 
 void launchReduceCpu(at::Tensor dst, size_t pos, size_t realSize, void* src,
                      size_t numRanks, c10d::ReduceOp op, bool* activeRanks);
+void preloadReduceKernels();
+
+class ConnectionContext;
+
+struct CudaTaskSubmissionToken {
+    size_t task_id;
+    uint64_t sequence;
+};
 
 class MooncakeWorker {
    public:
-    explicit MooncakeWorker();
+    explicit MooncakeWorker(int cuda_device_index = -1);
+    ~MooncakeWorker();
 
     c10::intrusive_ptr<c10d::Work> putTaskCpu(
         c10d::OpType opType, size_t tensorSize, int64_t broadcastRoot,
         const std::shared_ptr<TransferGroupMeta>& meta,
+        const std::shared_ptr<ConnectionContext>& connection_ctx,
         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)>&
@@ -72,22 +90,44 @@ class MooncakeWorker {
     c10::intrusive_ptr<c10d::Work> putTaskCuda(
         c10d::OpType opType, size_t tensorSize, int64_t broadcastRoot,
         const std::shared_ptr<TransferGroupMeta>& meta,
-        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 std::shared_ptr<ConnectionContext>& connection_ctx,
+        const at::cuda::CUDAStream& issue_stream,
+        const std::function<void(void* dst, size_t pos, size_t realSize,
+                                 const at::cuda::CUDAStream&)>& tensorToBuffer,
+        const std::function<void(void* src, size_t pos, size_t realSize,
+                                 const at::cuda::CUDAStream&)>& bufferToTensor);
 
-    void startWorker();
+    void Start();
 
-    void stopWorker() { 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();
+
     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};
+    int cuda_device_index_;
 
     Task *tasks_, *tasks_device_;
     bool hasCallback_[kNumTasks_]{};
@@ -95,6 +135,30 @@ 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 {
+   public:
+    static MooncakeWorkerManager& GetInstance() {
+        // leaky singleton to avoid destructor fiasco problem
+        static MooncakeWorkerManager* manager = new MooncakeWorkerManager;
+        return *manager;
+    }
+
+    std::shared_ptr<MooncakeWorker> GetCPUWorker();
+    std::shared_ptr<MooncakeWorker> GetCUDAWorker(int cuda_device_index);
+
+   private:
+    std::shared_ptr<MooncakeWorker> GetWorker(int worker_id);
+    static constexpr int CPUWorkerID = -1;
+    std::mutex manager_mutex_;
+    // Keep workers alive for the entire process lifetime because their
+    // detached threads must not outlive the MooncakeWorker object.
+    std::unordered_map<int, std::shared_ptr<MooncakeWorker>> workers_;
 };
 
 }  // namespace mooncake

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>
@@ -20,7 +21,6 @@ namespace mooncake {
 inline constexpr size_t kP2PBufferSize = 1u << 24;
 inline constexpr size_t kP2PNumSlots = 8;
 inline constexpr size_t kP2PSlotSize = kP2PBufferSize / kP2PNumSlots;
-inline constexpr size_t kP2PTotalBufferSize = kP2PBufferSize * kMaxNumRanks;
 
 struct alignas(64) AtomicHeadTail {
     uint32_t load(
@@ -58,6 +58,8 @@ struct P2PControlSlot {
 
 class P2PDeviceWorker;
 class P2PProxy {
+    static constexpr size_t kDrainTasksTimeoutMs = 5000;  // 5s
+
    public:
     friend class P2PDeviceWorker;
 
@@ -66,6 +68,7 @@ class P2PProxy {
         int rank = 0;
         int size = 0;
         int cuda_device_index = -1;
+        std::string location;
     };
 
     struct SendOp {
@@ -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,
@@ -231,7 +254,9 @@ class P2PProxy {
     int rank_ = 0;
     int size_ = 0;
     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/setup.py

@@ -3,7 +3,7 @@ import re
 
 from setuptools import setup
 import torch
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME
 
 
 torch_version = re.match(r"\d+(?:\.\d+)*", torch.__version__).group()
@@ -13,6 +13,18 @@ module_name = "mooncake.pg" + version_suffix
 abi_flag = int(torch._C._GLIBCXX_USE_CXX11_ABI)
 current_dir = os.path.abspath(os.path.dirname(__file__))
 
+# Link against the CUDA driver stub library if available.
+# Same approach as mooncake-ep/setup.py.
+cuda_libraries = ["ibverbs", "mlx5"]
+cuda_library_dirs = []
+
+if CUDA_HOME is not None:
+    cuda_stub_dir = os.path.join(CUDA_HOME, "lib64", "stubs")
+    cuda_stub_lib = os.path.join(cuda_stub_dir, "libcuda.so")
+    if os.path.exists(cuda_stub_lib):
+        cuda_libraries.insert(0, "cuda")
+        cuda_library_dirs.append(cuda_stub_dir)
+
 
 setup(
     name=module_name,
@@ -47,7 +59,8 @@ setup(
                     "-g0",
                 ],
             },
-            libraries=["ibverbs", "mlx5"],
+            libraries=cuda_libraries,
+            library_dirs=cuda_library_dirs,
             extra_link_args=[
                 "-Wl,-rpath,$ORIGIN",
                 "-L" + os.path.join(current_dir, "../mooncake-wheel/mooncake"),

mooncake-pg/src/CMakeLists.txt

@@ -0,0 +1,4 @@
+add_library(mooncake_pg connection_poller.cpp mooncake_backend.cpp mooncake_worker.cu mooncake_worker_thread.cpp p2p_proxy.cpp pg_py.cpp)
+
+set_target_properties(mooncake_pg PROPERTIES POSITION_INDEPENDENT_CODE ON)
+target_link_libraries(mooncake_pg PUBLIC ${TORCH_LIBRARIES} transfer_engine ibverbs mlx5)

mooncake-pg/src/connection_poller.cpp

@@ -1,6 +1,8 @@
 #include <c10/util/Exception.h>
 #include <connection_poller.h>
 #include <ATen/cuda/CUDAContext.h>
+#include <cuda.h>
+#include <cuda_alike.h>
 #include <cuda_runtime.h>
 #include <torch/torch.h>
 #include <atomic>
@@ -9,11 +11,38 @@
 #include <thread>
 #include <torch/csrc/distributed/c10d/Backend.hpp>
 #include <algorithm>
+#include <cstring>
 #include <limits>
+#include "memory_location.h"
 #include "mooncake_worker.cuh"
+#include "pg_utils.h"
 
 namespace mooncake {
+
+// Same check as nvlink_transport.cpp and mooncake_ep_buffer.cpp.
+// On MNNVL clusters all GPUs support fabric mem handles, meaning
+// NVLink transport can only access cuMemCreate(FABRIC) memory
+// cross-node -- CPU heap buffers are invisible to remote peers.
+static bool supportFabricMem() {
+    const char* nvlink_ipc = getenv("MC_USE_NVLINK_IPC");
+
+    bool fabric_enabled = nvlink_ipc && strcmp(nvlink_ipc, "0") == 0;
+    if (!fabric_enabled) return false;
+
+    int num_devices = 0;
+    cudaError_t err = cudaGetDeviceCount(&num_devices);
+    if (err != cudaSuccess || num_devices == 0) return false;
+
+    for (int dev = 0; dev < num_devices; ++dev) {
+        int supported = 0;
+        cuDeviceGetAttribute(
+            &supported, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, dev);
+        if (!supported) return false;
+    }
+    return true;
+}
 ConnectionContext::ConnectionContext(int backendIndex, int rank, int size,
+                                     bool isDummy,
                                      uint64_t* local2global_rank_map,
                                      c10::intrusive_ptr<::c10d::Store> store,
                                      std::shared_ptr<TransferGroupMeta> meta,
@@ -21,13 +50,25 @@ ConnectionContext::ConnectionContext(int backendIndex, int rank, int size,
                                      TransferEngine* engine)
     : backendIndex_(backendIndex),
       rank_(rank),
-      size_(size),
+      groupSize_(size),
+      isDummy_(isDummy),
+      establishedGroupSize_(0),
       local2global_rank_map_(local2global_rank_map),
       store_(std::move(store)),
       meta_(std::move(meta)),
       p2p_proxy_(std::move(p2p_proxy)),
-      engine_(engine) {
-    warmup_send_region_ = new int32_t[kMaxNumRanks];
+      engine_(engine),
+      skip_warmup_(supportFabricMem()) {
+    if (skip_warmup_) {
+        // On MNNVL clusters, CPU heap buffers aren't fabric-accessible so
+        // remote NVLink writes to them will fail. The fabric topology already
+        // guarantees connectivity, so we skip the warmup handshake entirely.
+        warmup_send_region_ = nullptr;
+        warmup_recv_region_ = nullptr;
+        return;
+    }
+
+    warmup_send_region_ = new int32_t[kMaxNumRanks]{};
     warmup_send_region_[0] = 1;
     int rc = engine_->registerLocalMemory(
         warmup_send_region_, kMaxNumRanks * sizeof(int32_t), kWildcardLocation);
@@ -40,24 +81,116 @@ ConnectionContext::ConnectionContext(int backendIndex, int rank, int size,
 }
 
 ConnectionContext::~ConnectionContext() {
-    for (int i = 0; i < size_; ++i) {
+    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());
         }
     }
 
-    engine_->unregisterLocalMemory(warmup_send_region_);
-    engine_->unregisterLocalMemory(warmup_recv_region_);
-    delete[] warmup_send_region_;
-    delete[] warmup_recv_region_;
+    if (warmup_send_region_) {
+        engine_->unregisterLocalMemory(warmup_send_region_);
+        delete[] warmup_send_region_;
+    }
+    if (warmup_recv_region_) {
+        engine_->unregisterLocalMemory(warmup_recv_region_);
+        delete[] warmup_recv_region_;
+    }
+}
+
+int ConnectionContext::getTotalConnectedPeers() const {
+    return totalConnectedPeers_.load(std::memory_order_acquire);
+}
+
+void ConnectionContext::extendGroupSizeTo(int newGroupSize) {
+    const int oldGroupSize = groupSize_.load(std::memory_order_acquire);
+    if (newGroupSize == oldGroupSize) return;
+
+    TORCH_CHECK(
+        newGroupSize >= 0 && static_cast<size_t>(newGroupSize) < kMaxNumRanks,
+        "Size out of range");
+    TORCH_CHECK(newGroupSize >= oldGroupSize, "newGroupSize < oldGroupSize");
+
+    // Reset local peer state for newly added ranks
+    for (int i = oldGroupSize; i < newGroupSize; ++i) {
+        meta_->peerConnected[i] = false;
+    }
+
+    groupSize_.store(newGroupSize, std::memory_order_release);
+}
+
+bool ConnectionContext::isAllPeerConnected() const {
+    return totalConnectedPeers_ == groupSize_;
 }
 
 void ConnectionContext::waitUntilAllConnected() {
+    if (isAllPeerConnected()) return;
+
     std::unique_lock<std::mutex> lock(backend_wakeup_mutex_);
     backend_wakeup_cv_.wait(lock, [this]() {
         return isAllPeerConnected() ||
                isShutdown_.load(std::memory_order_acquire);
     });
+    establishedGroupSize_.store(groupSize_, std::memory_order_release);
+}
+
+void ConnectionContext::waitUntilNewRanksConnected() {
+    if (isDummy_) {
+        return;
+    }
+    const int targetGroupSize = groupSize_.load(std::memory_order_acquire);
+    const int established =
+        establishedGroupSize_.load(std::memory_order_acquire);
+    if (established >= targetGroupSize) {
+        return;
+    }
+
+    std::unique_lock<std::mutex> lock(backend_wakeup_mutex_);
+    backend_wakeup_cv_.wait(lock, [this, targetGroupSize, established]() {
+        if (isShutdown_.load(std::memory_order_acquire)) {
+            return true;
+        }
+
+        for (int i = established; i < targetGroupSize; ++i) {
+            if (!meta_->peerConnected[i]) {
+                return false;
+            }
+        }
+        return true;
+    });
+
+    establishedGroupSize_.store(targetGroupSize, std::memory_order_release);
+}
+
+void ConnectionContext::bootstrapLocalPeer(const std::string& localServerName,
+                                           const SegmentInfo& localRankInfo) {
+    auto& peerState = peerStates_[rank_];
+    if (peerState.state == PeerConnectionState::CONNECTED) {
+        return;
+    }
+
+    auto segment_id = engine_->openSegment(localServerName);
+    meta_->segmentIDs[rank_] = segment_id;
+    peerState.segmentId = segment_id;
+    memcpy(&meta_->segmentInfos[rank_], &localRankInfo, sizeof(SegmentInfo));
+
+    meta_->peerConnected[rank_] = true;
+    ConnectionPoller::GetInstance()
+        .global_peerConnected_[local2global_rank_map_[rank_]] = true;
+    peerState.state = PeerConnectionState::CONNECTED;
+
+    {
+        std::lock_guard<std::mutex> lock(backend_wakeup_mutex_);
+        totalConnectedPeers_.store(1, std::memory_order_release);
+        if (isAllPeerConnected()) {
+            backend_wakeup_cv_.notify_all();
+        }
+    }
 }
 
 void ConnectionContext::shutdown() {
@@ -77,7 +210,7 @@ bool ConnectionContext::poll() {
     bool did_work = false;
 
     // Poll all peers sequentially.
-    for (int pollingRank = 0; pollingRank < size_; ++pollingRank) {
+    for (int pollingRank = 0; pollingRank < groupSize_; ++pollingRank) {
         did_work |= pollPeer(pollingRank);
     }
 
@@ -138,7 +271,19 @@ bool ConnectionContext::pollPeer(int pollingRank) {
             memcpy(&meta_->segmentInfos[pollingRank], buffer_data.data(),
                    sizeof(SegmentInfo));
 
-            if (pollingRank <= rank_) {
+            if (skip_warmup_) {
+                // MNNVL: fabric guarantees connectivity, skip warmup write
+                // since CPU heap buffers aren't fabric-accessible anyway.
+                meta_->peerConnected[pollingRank] = true;
+                global_peerConnected_[globalPollingRank] = true;
+                peerState.state = PeerConnectionState::CONNECTED;
+                {
+                    std::lock_guard<std::mutex> lock(backend_wakeup_mutex_);
+                    totalConnectedPeers_.fetch_add(1,
+                                                   std::memory_order_release);
+                    backend_wakeup_cv_.notify_all();
+                }
+            } else if (pollingRank <= rank_) {
                 // Send a warmup request to establish connections
                 auto batchID = engine_->allocateBatchID(1);
                 engine_->submitTransfer(
@@ -179,7 +324,7 @@ bool ConnectionContext::pollPeer(int pollingRank) {
                     std::lock_guard<std::mutex> lock(backend_wakeup_mutex_);
                     totalConnectedPeers_.fetch_add(1,
                                                    std::memory_order_release);
-                    if (isAllPeerConnected()) backend_wakeup_cv_.notify_all();
+                    backend_wakeup_cv_.notify_all();
                 }
                 state_changed = true;
             } else if (status.s == TransferStatusEnum::FAILED) {
@@ -206,7 +351,7 @@ bool ConnectionContext::pollPeer(int pollingRank) {
                     std::lock_guard<std::mutex> lock(backend_wakeup_mutex_);
                     totalConnectedPeers_.fetch_add(1,
                                                    std::memory_order_release);
-                    if (isAllPeerConnected()) backend_wakeup_cv_.notify_all();
+                    backend_wakeup_cv_.notify_all();
                 }
                 state_changed = true;
             }
@@ -245,13 +390,24 @@ bool ConnectionContext::pollPeer(int pollingRank) {
             // reports a failure. We must set both to false here.
             global_peerConnected_[globalPollingRank] = false;
             meta_->peerConnected[pollingRank] = false;
+            meta_->activeRanks[pollingRank] = false;
+            meta_->activeRanksTensor[pollingRank] = 0;
 
             // Reset store
-            store_->deleteKey(
-                getServerNameStoreKey(backendIndex_, pollingRank));
-            store_->deleteKey(getBufferStoreKey(backendIndex_, pollingRank));
-            store_->deleteKey(
-                getExtensionTaskCountStoreKey(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*>(
@@ -278,33 +434,66 @@ 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;
 }
 
-ConnectionPoller::ConnectionPoller() {
+ConnectionPoller::ConnectionPoller() = default;
+
+void ConnectionPoller::ensureThreadStarted() {
+    bool expected = false;
+    if (!pollerThreadStarted_.compare_exchange_strong(
+            expected, true, std::memory_order_acq_rel)) {
+        return;
+    }
     pollerThread_ = std::thread([this] { pollerLoop(); });
 }
 
 void ConnectionPoller::registerContext(
     const std::shared_ptr<ConnectionContext>& ctx) {
+    ensureThreadStarted();
     {
         std::lock_guard<std::mutex> lock(contexts_mutex_);
         contexts_.push_back(ctx);
@@ -316,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),
@@ -379,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))